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:
Manuel Ibar 2025-09-23 13:27:39 -03:00 • committed by GitHub
commit 5d64d61974
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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
}