mcp: fix context propagation in StreamableClientTransport (#514)
## Problem Context propagation was broken in `StreamableClientTransport` because: 1. `Connect()` used `context.Background()` instead of the parent context 2. `Close()` created a race condition where DELETE requests were cancelled before completion ## Solution - Use parent context when creating the connection context in `Connect()` - Reorder `Close()` operations to perform cleanup DELETE before cancelling context ## Impact - Request-scoped values (auth tokens, trace IDs) now propagate correctly to background HTTP operations - Eliminates race condition in cleanup operations - Maintains proper Go context semantics Fixes #513
This commit is contained in:
parent
33ff851ef3
commit
5d64d61974
1 changed files with 65 additions and 10 deletions
|
|
@ -1018,7 +1018,7 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, er
|
|||
// Create a new cancellable context that will manage the connection's lifecycle.
|
||||
// This is crucial for cleanly shutting down the background SSE listener by
|
||||
// cancelling its blocking network operations, which prevents hangs on exit.
|
||||
connCtx, cancel := context.WithCancel(context.Background())
|
||||
connCtx, cancel := context.WithCancel(ctx)
|
||||
conn := &streamableClientConn{
|
||||
url: t.Endpoint,
|
||||
client: client,
|
||||
|
|
@ -1230,7 +1230,7 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e
|
|||
|
||||
// testAuth controls whether a fake Authorization header is added to outgoing requests.
|
||||
// TODO: replace with a better mechanism when client-side auth is in place.
|
||||
var testAuth = false
|
||||
var testAuth atomic.Bool
|
||||
|
||||
func (c *streamableClientConn) setMCPHeaders(req *http.Request) {
|
||||
c.mu.Lock()
|
||||
|
|
@ -1242,7 +1242,7 @@ func (c *streamableClientConn) setMCPHeaders(req *http.Request) {
|
|||
if c.sessionID != "" {
|
||||
req.Header.Set(sessionIDHeader, c.sessionID)
|
||||
}
|
||||
if testAuth {
|
||||
if testAuth.Load() {
|
||||
req.Header.Set("Authorization", "Bearer foo")
|
||||
}
|
||||
}
|
||||
|
|
@ -1394,14 +1394,10 @@ func (c *streamableClientConn) reconnect(lastEventID string) (*http.Response, er
|
|||
// Close implements the [Connection] interface.
|
||||
func (c *streamableClientConn) Close() error {
|
||||
c.closeOnce.Do(func() {
|
||||
// Cancel any hanging network requests.
|
||||
c.cancel()
|
||||
close(c.done)
|
||||
|
||||
if errors.Is(c.failure(), errSessionMissing) {
|
||||
// If the session is missing, no need to delete it.
|
||||
} else {
|
||||
req, err := http.NewRequest(http.MethodDelete, c.url, nil)
|
||||
req, err := http.NewRequestWithContext(c.ctx, http.MethodDelete, c.url, nil)
|
||||
if err != nil {
|
||||
c.closeErr = err
|
||||
} else {
|
||||
|
|
@ -1411,6 +1407,10 @@ func (c *streamableClientConn) Close() error {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Cancel any hanging network requests after cleanup.
|
||||
c.cancel()
|
||||
close(c.done)
|
||||
})
|
||||
return c.closeErr
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue