mcp: better handling for streamable context cancellation (#677)

After walking through our handling of streamable client context
cancellation (due to encountering a shutdown deadlock), I think I've
settled on a more coherent strategy for handling call cancellation:

- In our call handler, retire the request if the call exits due to
cancellation: the caller will never see the actual result anyway.
- In connectSSE, use the actual request context (the same context used
in Write) for the client request, so that it terminates when the context
is cancelled. Thread through the initialization context for the
standalone SSE request.

Also, a couple minor improvements:
- Use a detached context for the background context of the client
connection. We want to preserve context values (see #513), but it is not
right to cancel the connection after Connect has already returned, if
the context times out.
- Don't use Last-Event-ID != "" as the signal for whether the connectSSE
call is initial: if the standalone SSE stream disconnects without an
event ID, we'll still reconnect it, and don't want to do so without a
delay.

+ tests, updating the streamable client connection test harness to
accomodate the new aspects being exercised.

Fixes #662
This commit is contained in:
Robert Findley 2025-12-03 12:26:50 -05:00 • committed by GitHub
commit 272e0cd4f5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -361,19 +361,26 @@ func (c *Connection) Call(ctx context.Context, method string, params any) *Async
if err := c.write(ctx, call); err != nil {
// Sending failed. We will never get a response, so deliver a fake one if it
// wasn't already retired by the connection breaking.
c.updateInFlight(func(s *inFlightState) {
if s.outgoingCalls[ac.id] == ac {
delete(s.outgoingCalls, ac.id)
ac.retire(&Response{ID: id, Error: err})
} else {
// ac was already retired by the readIncoming goroutine:
// perhaps our write raced with the Read side of the connection breaking.
}
})
c.Retire(ac, err)
}
return ac
}
// Retire stops tracking the call, and reports err as its terminal error.
//
// Retire is safe to call multiple times: if the call is already no longer
// tracked, Retire is a no op.
func (c *Connection) Retire(ac *AsyncCall, err error) {
c.updateInFlight(func(s *inFlightState) {
if s.outgoingCalls[ac.id] == ac {
delete(s.outgoingCalls, ac.id)
ac.retire(&Response{ID: ac.id, Error: err})
} else {
// ac was already retired elsewhere.
}
})
}
// Async, signals that the current jsonrpc2 request may be handled
// asynchronously to subsequent requests, when ctx is the request context.
//
@ -437,6 +444,9 @@ func (ac *AsyncCall) IsReady() bool {
}
// retire processes the response to the call.
//
// It is an error to call retire more than once: retire is guarded by the
// connection's outgoingCalls map.
func (ac *AsyncCall) retire(response *Response) {
select {
case <-ac.ready:
@ -450,6 +460,9 @@ func (ac *AsyncCall) retire(response *Response) {
// Await waits for (and decodes) the results of a Call.
// The response will be unmarshaled from JSON into the result.
//
// If the call is cancelled due to context cancellation, the result is
// ctx.Err().
func (ac *AsyncCall) Await(ctx context.Context, result any) error {
select {
case <-ctx.Done():