Skip to content

fix(core,server): keep StdioServerTransport running when a handler is cancelled #1101

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gibkim wants to merge 2 commits into
base: main
Choose a base branch
from
+71 −21
Open

fix(core,server): keep StdioServerTransport running when a handler is cancelled #1101

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
00f90b7
fix(server): keep StdioServerTransport running when a handler is canc…
gibkim Oct 7, 2026
806578c
fix(core): answer a leaked CancellationException in the serial phase
Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions kotlin-sdk-core/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/Protocol.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ import kotlinx.coroutines.NonCancellable
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.currentCoroutineContext
import kotlinx.coroutines.ensureActive
import kotlinx.coroutines.job
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Semaphore
Expand Down Expand Up @@ -625,8 +626,9 @@ public abstract class Protocol(@PublishedApi internal val options: ProtocolOptio
)
} catch (e: CancellationException) {
// Intentional exception to the always-rethrow-CE convention, only at this boundary.
if (handlerJob == null) throw e // serial phase: inline dispatch is not response-suppressed
if (handlerJob.isCancelled) throw e // genuine peer or close cancel: suppress the response
// Serial phase: the handler runs inline, so a real close cancels the delivering coroutine itself.
if (handlerJob == null) currentCoroutineContext().ensureActive()
if (handlerJob?.isCancelled == true) throw e // genuine peer or close cancel: suppress the response
// CE escaped a live handler job (e.g. a leaked inner withTimeout). Answer
// INTERNAL_ERROR so the peer does not hang until its own timeout.
respondWithError(capturedTransport, request, e)
Expand Down
29 changes: 29 additions & 0 deletions ...rc/commonTest/kotlin/io/modelcontextprotocol/kotlin/sdk/shared/ProtocolConcurrencyTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,35 @@ class ProtocolConcurrencyTest {
error.error.code shouldBe RPCError.ErrorCode.INTERNAL_ERROR
}

@Test
fun `handler leaking CancellationException in the serial phase produces INTERNAL_ERROR`() = runTest {
val (protocol, transport) = connected(concurrent = false)
protocol.fallbackRequestHandler = { _, _ -> throw CancellationException("leaked") }

transport.deliver(JSONRPCRequest(id = RequestId(7L), method = "test/leak"))

val error = errorsOn(transport).single()
error.id shouldBe RequestId(7L)
error.error.code shouldBe RPCError.ErrorCode.INTERNAL_ERROR
}

@Test
fun `cancelling the delivering coroutine in the serial phase sends no response`() = runTest {
val (protocol, transport) = connected(concurrent = false)
protocol.fallbackRequestHandler = { _, _ ->
delay(10.seconds)
EmptyResult()
}

val delivery = launch { transport.deliver(JSONRPCRequest(id = RequestId(8L), method = "test/slow")) }
runCurrent()
delivery.cancel()
advanceUntilIdle()

delivery.isCancelled shouldBe true
transport.sentWithOptions shouldBe emptyList()
}

@Test
fun `suspension-free notification handlers observe strict arrival order`() = runTest {
val (protocol, transport) = connected()
Expand Down
7 changes: 6 additions & 1 deletion ...r/src/commonMain/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransport.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ClosedSendChannelException
import kotlinx.coroutines.currentCoroutineContext
import kotlinx.coroutines.ensureActive
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
Expand Down Expand Up @@ -218,7 +219,11 @@ public class StdioServerTransport private constructor(
try {
_onMessage(message)
} catch (e: CancellationException) {
throw e
// Only the transport's own cancellation stops the loop. A cancellation that
// escapes a handler is a failure of that handler, like any other exception.
currentCoroutineContext().ensureActive()
logger.error(e) { "Message handler was cancelled" }
_onError(e)

@jetbrains-air jetbrains-air Bot Oct 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Problem: keeping the loop alive is the right call, but it now leaves the request the cancelled handler was serving with no reply at all. In the concurrent dispatch phase Protocol.onRequest already turns a cancellation that escapes a live handler job into an INTERNAL_ERROR response (Protocol.kt:626-632, "so the peer does not hang until its own timeout"). The serial, pre-notifications/initialized phase takes the handlerJob == null branch and rethrows instead — and until this change that rethrow landed here and tore the transport down, so the peer at least observed a closed connection. With the transport surviving, nothing goes out.

Failure scenario: a client sends initialize, then tools/call before the server has processed notifications/initialized (so dispatch is still inline), and the tool handler leaks a cancellation — e.g. an inner withTimeout expiring. I ran this against the branch: the server answers initialize, answers a later unrelated ping, and never answers the tools/call. The client blocks on that call until its own request timeout. On main the same scenario closed the stream, so the client failed fast.

Suggested fix: make the serial branch in Protocol.onRequest answer instead of rethrowing, using the same liveness check as here so a real close() during an inline handler still propagates:

if (handlerJob == null) {
    currentCoroutineContext().ensureActive()
    respondWithError(capturedTransport, request, e)
    return@withContext
}

Produced by Air Automations. Name: Code review / Run: https://air.jetbrains.cloud/org/05cf1a7f-6ab5-713b-abd3-29d0c8a05e2d/automations/409567ec-d3a7-40ac-ba95-8518ec577ac2?run=9ec05c83-ec7c-4a82-afff-a07cf8545630

@gibkim gibkim Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 14c4d86, as suggested: the serial branch in Protocol.onRequest now calls ensureActive() and answers INTERNAL_ERROR instead of rethrowing. Added a core test for the leak (fails on the previous commit) and one that cancels the delivering coroutine and expects no response. This moves part of the fix into core, so I noted that in the PR description.

} catch (e: Throwable) {
logger.error(e) { "Error processing message" }
_onError(e)
Expand Down
50 changes: 32 additions & 18 deletions .../src/jvmTest/kotlin/io/modelcontextprotocol/kotlin/sdk/server/StdioServerTransportTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.asCoroutineDispatcher
import kotlinx.coroutines.cancel
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.job
import kotlinx.coroutines.launch
import kotlinx.coroutines.withTimeoutOrNull
import kotlinx.io.Buffer
Expand Down Expand Up @@ -225,24 +226,6 @@ class StdioServerTransportTest {
server.close()
}

@Test
fun `should not invoke onError for CancellationException in handler`() = runIntegrationTest {
val server = StdioServerTransport(input = bufferedInput, output = printOutput)
val capturedError = CompletableDeferred<Throwable>()
val closed = CompletableDeferred<Unit>()
server.onError { capturedError.complete(it) }
server.onClose { closed.complete(Unit) }
server.onMessage { throw CancellationException("cancelled") }
server.start()

inputWriter.write(serializeMessage(PingRequest().toJSON()))
inputWriter.flush()

// The cancellation stops the processor; the transport then shuts down without reporting an error.
closed.await()
capturedError.isCompleted shouldBe false
}

// endregion

@Test
Expand Down Expand Up @@ -442,6 +425,36 @@ class StdioServerTransportTest {
ex.code shouldBe RPCError.ErrorCode.CONNECTION_CLOSED
}

@Test
fun `should stop without onError when the scope is cancelled during a handler`() = runIntegrationTest {
val externalScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
val server = StdioServerTransport(input = bufferedInput, output = printOutput) {
scope = externalScope
}
val capturedErrors = mutableListOf<Throwable>()
val handlerEntered = CompletableDeferred<Unit>()
val closeCaptured = CompletableDeferred<Unit>()
server.onError { capturedErrors.add(it) }
server.onClose { closeCaptured.complete(Unit) }
server.onMessage {
handlerEntered.complete(Unit)
CompletableDeferred<Unit>().await()
}
server.start()

inputWriter.write(serializeMessage(PingRequest().toJSON()))
inputWriter.flush()
handlerEntered.await()
externalScope.cancel()

closeCaptured.await()
// onClose can fire before the processor has handled the cancellation, so wait for all pumps.
// Closing the input lets the blocking reader return.
inputWriter.close()
externalScope.coroutineContext.job.join()
capturedErrors shouldBe emptyList()
}

@Test
fun `should use ioDispatcher for reader and writer`() = runIntegrationTest {
val threadName = "stdio-test-io-thread"
Expand Down Expand Up @@ -541,6 +554,7 @@ class StdioServerTransportTest {
private fun handlerErrors() = listOf(
RuntimeException("handler failure"),
OutOfMemoryError("handler error"),
CancellationException("handler cancelled"),
)

/** A [RawSource] that immediately throws [throwable] on every read attempt. */
Expand Down
Loading