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:
parent
8adc1e3a3d
commit
272e0cd4f5
1 changed files with 284 additions and 68 deletions
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Reference in a new issue