From 5e069e9ca2249e3e82858e6d16f8955892775dcf Mon Sep 17 00:00:00 2001 From: kerthans <1255829396@qq.com> Date: Wed, 21 Oct 2026 10:41:01 +0800 Subject: [PATCH] mcp: don't fall back to initialize when discover fails authorization With an OAuthHandler on StreamableClientTransport, Client.Connect sends server/discover first. If that request is still unauthorized after OAuthHandler.Authorize ran, or Authorize itself failed, Connect treated the error like any other non-modern discover error and fell back to initialize. The initialize POST got the same 401/403, so the transport ran Authorize a second time, which for interactive handlers means a second consent prompt. A 401 or 403 is not a protocol era signal: a legacy server that requires authorization rejects initialize the same way. The streamable client now wraps authorization failures with an unexported sentinel, and Connect returns that error instead of falling back. A legacy server behind authorization still connects through initialize, since the retried discover reaches it with valid credentials and fails for a non-auth reason. Fixes #1423 --- mcp/client.go | 8 +- mcp/streamable.go | 10 ++- mcp/streamable_client_test.go | 146 +++++++++++++++++++++++++++++++++- 3 files changed, 159 insertions(+), 5 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 3b24b2e9..1af2ae0c 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -371,8 +371,14 @@ func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOp } } } + // A legacy server behind authorization would reject initialize the + // same way, so falling back would only repeat the authorization flow. + if errors.Is(err, errAuthorizationFailed) { + _ = cs.Close() + return nil, err + } // Per the spec, fall back to the legacy initialize handshake on any - // non-modern error from server/discover. + // other non-modern error from server/discover. break } // Use the latest legacy protocol version for the fallback initialize. diff --git a/mcp/streamable.go b/mcp/streamable.go index 819bc08f..c9bc112b 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -2604,7 +2604,7 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e // Wrap with ErrRejected so the jsonrpc2 connection doesn't set writeErr // and permanently break the connection. // Wrap the authorization error as well for client inspection. - return fmt.Errorf("%s: %w: %w", requestSummary, jsonrpc2.ErrRejected, err) + return fmt.Errorf("%s: %w: %w: %w", requestSummary, jsonrpc2.ErrRejected, errAuthorizationFailed, err) } // Retry the request after successful authorization. _, resp, err = doRequest() @@ -2614,6 +2614,9 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e } if err := c.checkResponse(ctx, requestSummary, resp); err != nil { + if (resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden) && c.oauthHandler != nil { + err = fmt.Errorf("%w: %w", err, errAuthorizationFailed) + } if (requestMethod == methodDiscover || requestMethod == methodSubscriptionsListen) && !errors.Is(err, jsonrpc2.ErrRejected) { // Wrap the discover or subscriptions/listen failure with ErrRejected so // the jsonrpc2 layer doesn't set writeErr, which would break the connection @@ -2836,6 +2839,11 @@ func (c *streamableClientConn) handleSSE(ctx context.Context, requestSummary str } } +// errAuthorizationFailed marks a request that is still unauthorized after the +// OAuthHandler ran, either because Authorize failed or because the server +// rejected the resulting credentials. +var errAuthorizationFailed = errors.New("authorization failed") + // checkResponse checks the status code of the provided response, and // translates it into an error if the request was unsuccessful. // diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index 928b65ed..f8871ebe 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -12,6 +12,7 @@ import ( "io" "net/http" "net/http/httptest" + "slices" "strings" "sync" "sync/atomic" @@ -1005,12 +1006,16 @@ func TestStreamableClientDisableStandaloneSSE(t *testing.T) { } type mockOAuthHandler struct { + mu sync.Mutex token *oauth2.Token authorizeErr error - authorizeCalled bool + authorizedToken *oauth2.Token // if set, replaces token after a successful Authorize + authorizeCalls int } func (h *mockOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) { + h.mu.Lock() + defer h.mu.Unlock() if h.token == nil { return nil, nil } @@ -1018,10 +1023,21 @@ func (h *mockOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, } func (h *mockOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error { - h.authorizeCalled = true + h.mu.Lock() + defer h.mu.Unlock() + h.authorizeCalls++ + if h.authorizeErr == nil && h.authorizedToken != nil { + h.token = h.authorizedToken + } return h.authorizeErr } +func (h *mockOAuthHandler) calls() int { + h.mu.Lock() + defer h.mu.Unlock() + return h.authorizeCalls +} + func TestStreamableClientOAuth_AuthorizationHeader(t *testing.T) { ctx := context.Background() token := &oauth2.Token{AccessToken: "test-token"} @@ -1106,11 +1122,135 @@ func TestStreamableClientOAuth_401(t *testing.T) { t.Fatalf("client.Connect() error does not contain 'Unauthorized': %v", err) } - if !oauthHandler.authorizeCalled { + if oauthHandler.calls() == 0 { t.Errorf("expected Authorize to be called") } } +// TestStreamableClientOAuth_DiscoverUnauthorized_NoFallback is a regression +// test for #1423: an authorization failure on server/discover must not fall +// back to initialize, which would run the authorization flow a second time. +func TestStreamableClientOAuth_DiscoverUnauthorized_NoFallback(t *testing.T) { + tests := []struct { + name string + authorizeErr error + wantErr string + }{ + {"authorize fails", errors.New("access_denied: user declined consent"), "access_denied"}, + {"token still rejected", nil, "Unauthorized"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var ( + mu sync.Mutex + methods []string + ) + httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost { + body, _ := io.ReadAll(r.Body) + if msg, err := jsonrpc.DecodeMessage(body); err == nil { + if req, ok := msg.(*jsonrpc.Request); ok { + mu.Lock() + methods = append(methods, req.Method) + mu.Unlock() + } + } + } + w.Header().Set("WWW-Authenticate", `Bearer error="invalid_token"`) + w.WriteHeader(http.StatusUnauthorized) + })) + defer httpServer.Close() + + oauthHandler := &mockOAuthHandler{ + authorizeErr: test.authorizeErr, + authorizedToken: &oauth2.Token{AccessToken: "rejected-token"}, + } + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + OAuthHandler: oauthHandler, + } + client := NewClient(testImpl, nil) + session, err := client.Connect(context.Background(), transport, nil) + if err == nil { + session.Close() + t.Fatal("client.Connect() succeeded unexpectedly") + } + if !strings.Contains(err.Error(), test.wantErr) { + t.Errorf("client.Connect() error = %v, want it to contain %q", err, test.wantErr) + } + if strings.Contains(err.Error(), methodInitialize) { + t.Errorf("client.Connect() error = %v, want no mention of %q", err, methodInitialize) + } + if test.authorizeErr != nil && !errors.Is(err, test.authorizeErr) { + t.Errorf("client.Connect() error = %v, want it to wrap %v", err, test.authorizeErr) + } + if got := oauthHandler.calls(); got != 1 { + t.Errorf("Authorize called %d times, want 1", got) + } + mu.Lock() + defer mu.Unlock() + if slices.Contains(methods, methodInitialize) { + t.Errorf("server received %v, want no %q", methods, methodInitialize) + } + }) + } +} + +// TestStreamableClientOAuth_LegacyServerAuthorizedFallback verifies that a +// legacy server behind authorization still connects through initialize, with +// a single authorization. +func TestStreamableClientOAuth_LegacyServerAuthorizedFallback(t *testing.T) { + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodInitialize, ""}: { + header: header{ + "Content-Type": "application/json", + sessionIDHeader: "123", + }, + responseFunc: func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ID: r.ID, Result: mustMarshal(initResult)}), http.StatusOK + }, + }, + {"POST", "123", notificationInitialized, ""}: { + status: http.StatusAccepted, + wantProtocolVersion: protocolVersion20251125, + }, + {"GET", "123", "", ""}: { + header: header{"Content-Type": "text/event-stream"}, + optional: true, + }, + {"DELETE", "123", "", ""}: {optional: true}, + }, + } + verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) { + if token != "test-token" { + return nil, auth.ErrInvalidToken + } + return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil + } + httpServer := httptest.NewServer(auth.RequireBearerToken(verifier, nil)(fake)) + defer httpServer.Close() + + oauthHandler := &mockOAuthHandler{authorizedToken: &oauth2.Token{AccessToken: "test-token"}} + transport := &StreamableClientTransport{ + Endpoint: httpServer.URL, + OAuthHandler: oauthHandler, + } + client := NewClient(testImpl, nil) + session, err := client.Connect(context.Background(), transport, nil) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + defer session.Close() + if got := session.InitializeResult().ProtocolVersion; got != protocolVersion20251125 { + t.Errorf("InitializeResult.ProtocolVersion = %q, want %q (initialize fallback)", got, protocolVersion20251125) + } + if got := oauthHandler.calls(); got != 1 { + t.Errorf("Authorize called %d times, want 1", got) + } +} + // 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