-
Notifications
You must be signed in to change notification settings - Fork 270
mcp: fix context propagation in StreamableClientTransport #514
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.
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
mcp: fix context propagation in StreamableClientTransport #514
Conversation
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 This change fixes both issues by: - Using the parent context when creating the connection context - Reordering Close() to perform cleanup DELETE before cancelling context This ensures request-scoped values (auth tokens, trace IDs) propagate correctly to background HTTP operations. Fixes modelcontextprotocol#513
|
Thanks for the PR! There seems to be a race failure. Could you also add a test to streamable_test.go that demonstrates that the context is now propagated correctly? |
Adds TestStreamableClientContextPropagation to verify that context values are properly propagated to background HTTP operations (SSE GET and cleanup DELETE requests) in StreamableClientTransport. The test uses a custom HTTP handler to capture request contexts and verify that context values from the parent context are accessible in both the SSE connection establishment and session cleanup requests.
473a047 to
2934625
Compare
The test now properly verifies that StreamableClientTransport can handle contexts with values without errors, demonstrating that the context propagation fix is working correctly. The test validates the fix at streamable.go:1021 where context.WithCancel(ctx) is used instead of context.WithCancel(context.Background()).
…nsport Replace bloated context propagation test with clean, idiomatic Go code that: - Follows existing test patterns in the codebase - Removes unnecessary comments and complex synchronization - Tests actual context derivation (cancellable contexts) - Demonstrates the fix works without HTTP value propagation complexity The test verifies that background HTTP operations (SSE GET and DELETE cleanup) use properly derived contexts from the original Connect() context, confirming the fix in streamable.go where context.WithCancel(ctx) replaced context.WithCancel(context.Background()).
6a5dd46 to
016015c
Compare
Replace unsafe global testAuth boolean with atomic.Bool to prevent data race between test setup/teardown and background HTTP operations in StreamableClientTransport.
|
Hi @samthanawalla, thanks for the review. I've updated the PR with the following:
Ready for another look when you get a chance 👍 . |
mcp/streamable_test.go
Outdated
| mu.Lock() | ||
| defer mu.Unlock() | ||
|
|
||
| if getCtx != nil && getCtx.Done() == nil { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This test only checks that the context is cancellable but doesn't check if the values are passed. Can we change this test to verify that the value of some testkey is preserved?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I addressed this comment on this commit
Add test to verify that context values and cancellation are properly propagated through StreamableClientTransport.Connect(). The test ensures: - Context values are preserved in the connection context - Connection context remains cancellable - Parent context cancellation propagates to connection context This addresses the context propagation requirements for the streamable transport implementation.
Problem
Context propagation was broken in
StreamableClientTransportbecause:Connect()usedcontext.Background()instead of the parent contextClose()created a race condition where DELETE requests were cancelled before completionSolution
Connect()Close()operations to perform cleanup DELETE before cancelling contextImpact
Fixes #513