Skip to content

Commit 88305e8

Browse files
committed
mcp: do not re-prompt OAuth after cancelled Authorize
When OAuthHandler.Authorize returned a context-cancelled error, the cancellation notification the call layer sent in response to the same ctx cancellation was routed back through the streamable client transport, re-triggering the 401 -> Authorize path and prompting the user a second time for a request they had already abandoned. Fix in two layers: 1. mcp/transport.go: bound the cancellation Notify with a fresh context.Background()-derived timeout instead of reusing the cancelled caller context. This guarantees the caller never blocks on a doomed notify against a degraded connection. 2. mcp/streamable.go: when Authorize returns a context-cancelled error, mark the streamable client connection as failed via c.fail(). The existing failure check in Write() then short-circuits the cancellation notification without re-invoking the OAuth handler, matching the bug reporter's expectation that cancelling Authorize abandons the connection. Includes a regression test with a blocking OAuth handler that asserts Authorize is invoked exactly once. The test was verified to fail on unfixed code (got 2 calls) and pass with the fix. Fixes #882
1 parent 862d78a commit 88305e8

3 files changed

Lines changed: 112 additions & 3 deletions

File tree

mcp/streamable.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1801,6 +1801,15 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e
18011801

18021802
if (resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden) && c.oauthHandler != nil {
18031803
if err := c.oauthHandler.Authorize(ctx, req, resp); err != nil {
1804+
// If the caller's context was cancelled while we were running the
1805+
// authorization flow, treat the connection as failed so subsequent
1806+
// operations on it (e.g. the cancellation notify the call layer
1807+
// sends in response to ctx cancellation) short-circuit instead of
1808+
// re-invoking the OAuth handler. Otherwise the user gets prompted
1809+
// to authorize a request they have already abandoned. See #882.
1810+
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
1811+
c.fail(fmt.Errorf("%s: authorization cancelled: %w", requestSummary, err))
1812+
}
18041813
// Wrap with ErrRejected so the jsonrpc2 connection doesn't set writeErr
18051814
// and permanently break the connection.
18061815
// Wrap the authorization error as well for client inspection.

mcp/streamable_client_test.go

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1016,6 +1016,88 @@ func TestStreamableClientOAuth_401(t *testing.T) {
10161016
}
10171017
}
10181018

1019+
// blockingCountingOAuthHandler is an OAuthHandler that blocks inside
1020+
// Authorize until the caller's context is cancelled, then returns ctx.Err().
1021+
// It records how many times Authorize is invoked. Used by the regression
1022+
// test for #882.
1023+
type blockingCountingOAuthHandler struct {
1024+
mu sync.Mutex
1025+
callCount int
1026+
}
1027+
1028+
func (h *blockingCountingOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
1029+
return nil, nil
1030+
}
1031+
1032+
func (h *blockingCountingOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error {
1033+
h.mu.Lock()
1034+
h.callCount++
1035+
h.mu.Unlock()
1036+
// Block until the caller's context is cancelled, mirroring an
1037+
// interactive OAuth flow that the user has abandoned.
1038+
<-ctx.Done()
1039+
return ctx.Err()
1040+
}
1041+
1042+
func (h *blockingCountingOAuthHandler) Calls() int {
1043+
h.mu.Lock()
1044+
defer h.mu.Unlock()
1045+
return h.callCount
1046+
}
1047+
1048+
// TestStreamableClientOAuth_CancelledAuthorize_NoReprompt is a regression
1049+
// test for #882. When OAuthHandler.Authorize returns a context-cancelled
1050+
// error, the connection must enter a failed state so that the cancellation
1051+
// notification the call layer sends in response to ctx cancellation does
1052+
// not flow back through the same broken auth path and re-invoke Authorize.
1053+
func TestStreamableClientOAuth_CancelledAuthorize_NoReprompt(t *testing.T) {
1054+
handler := &blockingCountingOAuthHandler{}
1055+
1056+
fake := &fakeStreamableServer{
1057+
t: t,
1058+
responses: fakeResponses{
1059+
{"POST", "", methodInitialize, ""}: {
1060+
header: header{
1061+
"Content-Type": "application/json",
1062+
sessionIDHeader: "123",
1063+
},
1064+
body: jsonBody(t, initResp),
1065+
},
1066+
},
1067+
}
1068+
verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) {
1069+
return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil
1070+
}
1071+
httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake))
1072+
t.Cleanup(httpServer.Close)
1073+
1074+
transport := &StreamableClientTransport{
1075+
Endpoint: httpServer.URL,
1076+
OAuthHandler: handler,
1077+
}
1078+
client := NewClient(testImpl, nil)
1079+
1080+
// Use a context with a tight deadline so the cancellation path runs
1081+
// while the auth flow is in progress.
1082+
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
1083+
defer cancel()
1084+
1085+
_, err := client.Connect(ctx, transport, nil)
1086+
if err == nil {
1087+
t.Fatal("expected client.Connect to fail")
1088+
}
1089+
1090+
// Give the cancellation Notify path a moment to (try to) run.
1091+
time.Sleep(50 * time.Millisecond)
1092+
1093+
// Authorize should be invoked exactly once. The bug in #882 caused
1094+
// it to be invoked a second time when the call layer sent the
1095+
// cancellation notification through the same auth-broken connection.
1096+
if got := handler.Calls(); got != 1 {
1097+
t.Errorf("expected Authorize to be called exactly 1 time, got %d", got)
1098+
}
1099+
}
1100+
10191101
func TestTokenInfo(t *testing.T) {
10201102
ctx := context.Background()
10211103

mcp/transport.go

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,20 @@ import (
1414
"net"
1515
"os"
1616
"sync"
17+
"time"
1718

1819
internaljson "github.com/modelcontextprotocol/go-sdk/internal/json"
1920
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
20-
"github.com/modelcontextprotocol/go-sdk/internal/xcontext"
2121
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
2222
)
2323

24+
// notifyCancellationTimeout bounds the cancellation notification we send to
25+
// the peer when the caller's context is cancelled. The notification is
26+
// best-effort: a degraded connection (e.g. an OAuth flow that has been
27+
// abandoned) must not be able to block the caller's return path or
28+
// re-trigger expensive recovery on its behalf. See issue #882.
29+
const notifyCancellationTimeout = 5 * time.Second
30+
2431
// ErrConnectionClosed is returned when sending a message to a connection that
2532
// is closed or in the process of closing.
2633
var ErrConnectionClosed = errors.New("connection closed")
@@ -216,8 +223,19 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params
216223
case errors.Is(err, jsonrpc2.ErrClientClosing), errors.Is(err, jsonrpc2.ErrServerClosing):
217224
return fmt.Errorf("%w: calling %q: %v", ErrConnectionClosed, method, err)
218225
case ctx.Err() != nil:
219-
// Notify the peer of cancellation.
220-
err := conn.Notify(xcontext.Detach(ctx), notificationCancelled, &CancelledParams{
226+
// Best-effort notify the peer of cancellation. We deliberately bound
227+
// this with a fresh, short-lived context derived from
228+
// context.Background() rather than reusing (or merely detaching) the
229+
// caller's already-cancelled context. The connection may be in a
230+
// degraded state — for example, the original failure may have come
231+
// from an OAuth flow whose handler context expired (see #882) — and
232+
// reusing that context would either return immediately with an error
233+
// or, worse, re-trigger expensive recovery (re-auth) on the caller's
234+
// return path. The bounded background context lets the notification
235+
// attempt to deliver but never blocks the caller indefinitely.
236+
notifyCtx, cancelNotify := context.WithTimeout(context.Background(), notifyCancellationTimeout)
237+
defer cancelNotify()
238+
err := conn.Notify(notifyCtx, notificationCancelled, &CancelledParams{
221239
Reason: ctx.Err().Error(),
222240
RequestID: call.ID().Raw(),
223241
})

0 commit comments

Comments
 (0)