mcp: do not re-prompt OAuth after cancelled Authorize (#885)
When `OAuthHandler.Authorize` returns a context-cancelled error, the cancellation notification that the call layer sends in response to the same ctx cancellation is routed back through the streamable client transport, re-triggering the `401` → `Authorize` path and prompting the user a second time for a request they have 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 (`TestStreamableClientOAuth_CancelledAuthorize_NoReprompt`) with a blocking `OAuthHandler` 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 --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
5f2cd8fd14
commit
db50910c70
3 changed files with 115 additions and 3 deletions
|
|
@ -1814,6 +1814,22 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e
|
||||||
|
|
||||||
if (resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden) && c.oauthHandler != nil {
|
if (resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden) && c.oauthHandler != nil {
|
||||||
if err := c.oauthHandler.Authorize(ctx, req, resp); err != nil {
|
if err := c.oauthHandler.Authorize(ctx, req, resp); err != nil {
|
||||||
|
// If the caller's context was cancelled while we were running the
|
||||||
|
// authorization flow, treat the connection as failed so subsequent
|
||||||
|
// operations on it (e.g. the cancellation notify the call layer
|
||||||
|
// sends in response to ctx cancellation) short-circuit instead of
|
||||||
|
// re-invoking the OAuth handler. Otherwise the user gets prompted
|
||||||
|
// to authorize a request they have already abandoned. See #882.
|
||||||
|
//
|
||||||
|
// We check ctx.Err() rather than the error returned by Authorize,
|
||||||
|
// because the handler is user-implemented and may return an error
|
||||||
|
// that does not wrap context.Canceled (e.g. a custom sentinel or
|
||||||
|
// a fmt.Errorf with %v). The context itself is the authoritative
|
||||||
|
// source for whether the caller abandoned the request.
|
||||||
|
ctxErr := ctx.Err()
|
||||||
|
if errors.Is(ctxErr, context.Canceled) || errors.Is(ctxErr, context.DeadlineExceeded) {
|
||||||
|
c.fail(fmt.Errorf("%s: authorization cancelled: %w", requestSummary, err))
|
||||||
|
}
|
||||||
// Wrap with ErrRejected so the jsonrpc2 connection doesn't set writeErr
|
// Wrap with ErrRejected so the jsonrpc2 connection doesn't set writeErr
|
||||||
// and permanently break the connection.
|
// and permanently break the connection.
|
||||||
// Wrap the authorization error as well for client inspection.
|
// Wrap the authorization error as well for client inspection.
|
||||||
|
|
|
||||||
|
|
@ -1017,6 +1017,94 @@ func TestStreamableClientOAuth_401(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// blockingCountingOAuthHandler is an OAuthHandler that blocks inside
|
||||||
|
// Authorize until the caller's context is cancelled, then returns a custom
|
||||||
|
// error that does NOT wrap context.Canceled. This mirrors real-world OAuth
|
||||||
|
// handlers that catch the cancellation internally and surface their own
|
||||||
|
// error type. The fix for #882 checks ctx.Err() directly rather than
|
||||||
|
// relying on the error from Authorize, so this must still trigger c.fail().
|
||||||
|
// It records how many times Authorize is invoked.
|
||||||
|
type blockingCountingOAuthHandler struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
callCount int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *blockingCountingOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *blockingCountingOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error {
|
||||||
|
h.mu.Lock()
|
||||||
|
h.callCount++
|
||||||
|
h.mu.Unlock()
|
||||||
|
// Block until the caller's context is cancelled, mirroring an
|
||||||
|
// interactive OAuth flow that the user has abandoned.
|
||||||
|
<-ctx.Done()
|
||||||
|
// Return a custom error that does not wrap context.Canceled, as a
|
||||||
|
// real-world handler might. The code under test must check ctx.Err()
|
||||||
|
// to detect the cancellation, not this error.
|
||||||
|
return fmt.Errorf("oauth flow interrupted")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *blockingCountingOAuthHandler) Calls() int {
|
||||||
|
h.mu.Lock()
|
||||||
|
defer h.mu.Unlock()
|
||||||
|
return h.callCount
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestStreamableClientOAuth_CancelledAuthorize_NoReprompt is a regression
|
||||||
|
// test for #882. When OAuthHandler.Authorize returns a context-cancelled
|
||||||
|
// error, the connection must enter a failed state so that the cancellation
|
||||||
|
// notification the call layer sends in response to ctx cancellation does
|
||||||
|
// not flow back through the same broken auth path and re-invoke Authorize.
|
||||||
|
func TestStreamableClientOAuth_CancelledAuthorize_NoReprompt(t *testing.T) {
|
||||||
|
handler := &blockingCountingOAuthHandler{}
|
||||||
|
|
||||||
|
fake := &fakeStreamableServer{
|
||||||
|
t: t,
|
||||||
|
responses: fakeResponses{
|
||||||
|
{"POST", "", methodInitialize, ""}: {
|
||||||
|
header: header{
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
sessionIDHeader: "123",
|
||||||
|
},
|
||||||
|
body: jsonBody(t, initResp),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) {
|
||||||
|
return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil
|
||||||
|
}
|
||||||
|
httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake))
|
||||||
|
t.Cleanup(httpServer.Close)
|
||||||
|
|
||||||
|
transport := &StreamableClientTransport{
|
||||||
|
Endpoint: httpServer.URL,
|
||||||
|
OAuthHandler: handler,
|
||||||
|
}
|
||||||
|
client := NewClient(testImpl, nil)
|
||||||
|
|
||||||
|
// Use a context with a tight deadline so the cancellation path runs
|
||||||
|
// while the auth flow is in progress.
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
_, err := client.Connect(ctx, transport, nil)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected client.Connect to fail")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Give the cancellation Notify path a moment to (try to) run.
|
||||||
|
time.Sleep(50 * time.Millisecond)
|
||||||
|
|
||||||
|
// Authorize should be invoked exactly once. The bug in #882 caused
|
||||||
|
// it to be invoked a second time when the call layer sent the
|
||||||
|
// cancellation notification through the same auth-broken connection.
|
||||||
|
if got := handler.Calls(); got != 1 {
|
||||||
|
t.Errorf("expected Authorize to be called exactly 1 time, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestTokenInfo(t *testing.T) {
|
func TestTokenInfo(t *testing.T) {
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -14,13 +14,20 @@ import (
|
||||||
"net"
|
"net"
|
||||||
"os"
|
"os"
|
||||||
"sync"
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
internaljson "github.com/modelcontextprotocol/go-sdk/internal/json"
|
internaljson "github.com/modelcontextprotocol/go-sdk/internal/json"
|
||||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||||
"github.com/modelcontextprotocol/go-sdk/internal/xcontext"
|
|
||||||
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
|
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// notifyCancellationTimeout bounds the cancellation notification we send to
|
||||||
|
// the peer when the caller's context is cancelled. The notification is
|
||||||
|
// best-effort: a degraded connection (e.g. an OAuth flow that has been
|
||||||
|
// abandoned) must not be able to block the caller's return path or
|
||||||
|
// re-trigger expensive recovery on its behalf. See issue #882.
|
||||||
|
const notifyCancellationTimeout = 5 * time.Second
|
||||||
|
|
||||||
// ErrConnectionClosed is returned when sending a message to a connection that
|
// ErrConnectionClosed is returned when sending a message to a connection that
|
||||||
// is closed or in the process of closing.
|
// is closed or in the process of closing.
|
||||||
var ErrConnectionClosed = errors.New("connection closed")
|
var ErrConnectionClosed = errors.New("connection closed")
|
||||||
|
|
@ -220,8 +227,9 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params
|
||||||
case errors.Is(err, jsonrpc2.ErrClientClosing), errors.Is(err, jsonrpc2.ErrServerClosing):
|
case errors.Is(err, jsonrpc2.ErrClientClosing), errors.Is(err, jsonrpc2.ErrServerClosing):
|
||||||
return fmt.Errorf("%w: calling %q: %v", ErrConnectionClosed, method, err)
|
return fmt.Errorf("%w: calling %q: %v", ErrConnectionClosed, method, err)
|
||||||
case ctx.Err() != nil:
|
case ctx.Err() != nil:
|
||||||
// Notify the peer of cancellation.
|
notifyCtx, cancelNotify := context.WithTimeout(context.WithoutCancel(ctx), notifyCancellationTimeout)
|
||||||
err := conn.Notify(xcontext.Detach(ctx), notificationCancelled, &CancelledParams{
|
defer cancelNotify()
|
||||||
|
err := conn.Notify(notifyCtx, notificationCancelled, &CancelledParams{
|
||||||
Reason: ctx.Err().Error(),
|
Reason: ctx.Err().Error(),
|
||||||
RequestID: call.ID().Raw(),
|
RequestID: call.ID().Raw(),
|
||||||
})
|
})
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue