KTOR-8938 Inherit server coroutine context in WebSocket session (#5426)

* KTOR-8938 Add failing test for WebSocket context inheritance

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* KTOR-8938 Inherit server coroutine context in WebSocket session

The WebSocket session's coroutineContext was built solely from
the engine dispatcher, dropping custom CoroutineContext elements
from the server/call scope. Use call.coroutineContext as the base
so user-provided elements are preserved.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* KTOR-8938 Fix import ordering in WebSocketEngineSuite

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* KTOR-8938 Preserve original Job parent in WebSocket context fix

Strip the call's Job before merging its context, and restore
the original Job from the engine's processing coroutine. This
ensures custom context elements are inherited without changing
the WebSocket session's Job parent-child relationship.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* KTOR-8938 Use call coroutine context including Job for WebSocket session

Use call.coroutineContext directly (with its Job) so that the
WebSocket session is a proper child of the server's coroutine
scope. This ensures cancellation propagates correctly from
server to WebSocket sessions.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* KTOR-8938 Add test for WebSocket session cancellation on server stop

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* KTOR-8938 Make server-stop cancellation test open, skip on Tomcat

Tomcat's servlet-based WebSocket upgrade creates an upgradeJob
with no parent connection to the server scope, so server stop
does not propagate cancellation to WebSocket sessions. This is
a pre-existing Tomcat limitation unrelated to this fix.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* KTOR-8938 Fix WebSocket cancellation on server stop for servlet engines

Pass the servlet's SupervisorJob as parent for the upgrade handler's Job,
ensuring cancellation propagates when the servlet is destroyed. Use
onCancelling=true for WebConnection cleanup to break the deadlock between
blocked IO reads and Job completion.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* KTOR-8938 Skip server-stop cancellation test on JS/WASM

server.stop() calls runBlocking internally via runBlockingBridge,
which throws on JS and WASM platforms. Skip the test on those
targets with a runtime platform check.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* KTOR-8938 Narrow @OptIn(InternalCoroutinesApi::class) scope to init method

Move the annotation from class-level to the specific method that uses
invokeOnCompletion(onCancelling = true), in both servlet and servlet-jakarta
upgrade handlers.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* KTOR-8938 Use stopSuspend instead of stop in WebSocket test

- Replace server.stop() with server.stopSuspend() as suggested by reviewer
- Remove JS/WASM platform check since stopSuspend doesn't use runBlocking
- Remove open modifier from test function (no longer needed for overrides)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: e5l

ktor-server/ktor-server-plugins/ktor-server-websockets/common/src/io/ktor/server/websocket/WebSocketUpgrade.kt

@@ -95,7 +95,7 @@ public class WebSocketUpgrade(
             output,
             plugin.maxFrameSize,
             plugin.masking,
-            coroutineContext = engineContext + (coroutineContext[Job] ?: EmptyCoroutineContext),
+            coroutineContext = call.coroutineContext + engineContext,
             channelsConfig = plugin.channelsConfig
         )
 

ktor-server/ktor-server-servlet-jakarta/jvm/src/io/ktor/server/servlet/jakarta/AsyncServlet.kt

@@ -125,7 +125,9 @@ public open class AsyncServletApplicationResponse(
         (call.request as AsyncServletApplicationRequest).upgraded()
         completed = true
 
-        servletUpgradeImpl.performUpgrade(upgrade, servletRequest, servletResponse, engineContext, userContext)
+        val servletJob = coroutineContext[Job]?.parent
+        val upgradeEngineContext = if (servletJob != null) engineContext + servletJob else engineContext
+        servletUpgradeImpl.performUpgrade(upgrade, servletRequest, servletResponse, upgradeEngineContext, userContext)
     }
 
     @UseHttp2Push

ktor-server/ktor-server-servlet-jakarta/jvm/src/io/ktor/server/servlet/jakarta/ServletUpgrade.kt

@@ -84,13 +84,14 @@ public class ServletUpgradeHandler : HttpUpgradeHandler, CoroutineScope {
 
     override val coroutineContext: CoroutineContext get() = upgradeJob
 
+    @OptIn(InternalCoroutinesApi::class)
     override fun init(webConnection: WebConnection?) {
         if (webConnection == null) {
             throw IllegalArgumentException("Upgrade processing requires WebConnection instance")
         }
 
         upgradeJob = Job(up.engineContext[Job])
-        upgradeJob.invokeOnCompletion {
+        upgradeJob.invokeOnCompletion(onCancelling = true) {
             webConnection.close()
         }
 

ktor-server/ktor-server-servlet/jvm/src/io/ktor/server/servlet/AsyncServlet.kt

@@ -125,7 +125,9 @@ public open class AsyncServletApplicationResponse(
         (call.request as AsyncServletApplicationRequest).upgraded()
         completed = true
 
-        servletUpgradeImpl.performUpgrade(upgrade, servletRequest, servletResponse, engineContext, userContext)
+        val servletJob = coroutineContext[Job]?.parent
+        val upgradeEngineContext = if (servletJob != null) engineContext + servletJob else engineContext
+        servletUpgradeImpl.performUpgrade(upgrade, servletRequest, servletResponse, upgradeEngineContext, userContext)
     }
 
     @UseHttp2Push

ktor-server/ktor-server-servlet/jvm/src/io/ktor/server/servlet/ServletUpgrade.kt

@@ -84,13 +84,14 @@ public class ServletUpgradeHandler : HttpUpgradeHandler, CoroutineScope {
 
     override val coroutineContext: CoroutineContext get() = upgradeJob
 
+    @OptIn(InternalCoroutinesApi::class)
     override fun init(webConnection: WebConnection?) {
         if (webConnection == null) {
             throw IllegalArgumentException("Upgrade processing requires WebConnection instance")
         }
 
         upgradeJob = Job(up.engineContext[Job])
-        upgradeJob.invokeOnCompletion {
+        upgradeJob.invokeOnCompletion(onCancelling = true) {
             webConnection.close()
         }
 

ktor-server/ktor-server-test-suites/common/src/io/ktor/server/testing/suites/WebSocketEngineSuite.kt

@@ -21,6 +21,7 @@ import kotlinx.coroutines.*
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.channels.*
 import kotlinx.io.*
+import kotlin.coroutines.*
 import kotlin.random.*
 import kotlin.test.*
 import kotlin.time.Duration.Companion.milliseconds
@@ -653,6 +654,62 @@ abstract class WebSocketEngineSuite<TEngine : ApplicationEngine, TConfiguration
         checkFrame(second.await())
     }
 
+    @Test
+    fun testWebSocketSessionInheritsServerCoroutineContext() = runTest {
+        val result = CompletableDeferred<Boolean>()
+
+        val customElement = object : AbstractCoroutineContextElement(CustomTestElement) {}
+
+        createAndStartServer(parent = customElement) {
+            webSocket("/") {
+                val hasElement = coroutineContext[CustomTestElement] != null
+                result.complete(hasElement)
+            }
+        }
+
+        useSocket {
+            negotiateHttpWebSocket()
+
+            output.apply {
+                // close frame with code 1000
+                writeHex("0x88 0x02 0x03 0xe8")
+                flush()
+            }
+
+            assertCloseFrame()
+        }
+
+        assertTrue(result.await(), "WebSocket session should inherit custom coroutine context elements from server")
+    }
+
+    @Test
+    fun testWebSocketSessionCancelledOnServerStop() = runTest {
+        val sessionStarted = CompletableDeferred<Unit>()
+        val sessionCancelled = CompletableDeferred<Unit>()
+
+        createAndStartServer {
+            webSocket("/") {
+                sessionStarted.complete(Unit)
+                try {
+                    incoming.consumeEach {}
+                } finally {
+                    sessionCancelled.complete(Unit)
+                }
+            }
+        }
+
+        useSocket {
+            negotiateHttpWebSocket()
+
+            sessionStarted.await()
+            server!!.stopSuspend(0, 0)
+
+            withTimeout(5000) {
+                sessionCancelled.await()
+            }
+        }
+    }
+
     @Test
     fun testCorruptFrameWithBadOpcode() = runTest {
         createAndStartServer {
@@ -830,6 +887,8 @@ internal suspend fun ByteWriteChannel.writeFrameTest(frame: Frame, masking: Bool
 
 internal fun Boolean.flagAt(at: Int) = if (this) 1 shl at else 0
 
+private object CustomTestElement : CoroutineContext.Key<AbstractCoroutineContextElement>
+
 private fun Source.mask(maskKey: Int): Source = withMemory(4) { maskMemory ->
     maskMemory.storeIntAt(0, maskKey)
     buildPacket {