[SSE][OAuth] Add OAuth support to SSE client #340
Conversation
|
""" WalkthroughThis change adds OAuth authentication support to the SSE transport, introduces a new OAuth SSE client constructor, and updates the Streamable HTTP client to use a renamed OAuth option. It also adds comprehensive tests for SSE OAuth integration, normalizes the "Bearer" token type, and refactors OAuth client initialization and retry logic in the example client. Changes
Possibly related PRs
Suggested labels
Suggested reviewers
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (1.64.8)Error: you are using a configuration file for golangci-lint v2 with golangci-lint v1: please use golangci-lint v2 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
examples/oauth_client/main.go (1)
106-170: 🛠️ Refactor suggestionConsider returning errors instead of calling
log.Fatalf.The function calls
log.Fatalfon errors, making it impossible for callers to handle errors differently. Consider returning an error to make the function more flexible and testable.-func maybeAuthorize(err error) { +func maybeAuthorize(err error) error { // Check if we need OAuth authorization if client.IsOAuthAuthorizationRequiredError(err) { fmt.Println("OAuth authorization required. Starting authorization flow...") // Get the OAuth handler from the error oauthHandler := client.GetOAuthHandler(err) // Start a local server to handle the OAuth callback callbackChan := make(chan map[string]string) server := startCallbackServer(callbackChan) defer server.Close() // Generate PKCE code verifier and challenge codeVerifier, err := client.GenerateCodeVerifier() if err != nil { - log.Fatalf("Failed to generate code verifier: %v", err) + return fmt.Errorf("failed to generate code verifier: %w", err) } codeChallenge := client.GenerateCodeChallenge(codeVerifier) // Generate state parameter state, err := client.GenerateState() if err != nil { - log.Fatalf("Failed to generate state: %v", err) + return fmt.Errorf("failed to generate state: %w", err) } err = oauthHandler.RegisterClient(context.Background(), "mcp-go-oauth-example") if err != nil { - log.Fatalf("Failed to register client: %v", err) + return fmt.Errorf("failed to register client: %w", err) } // Get the authorization URL authURL, err := oauthHandler.GetAuthorizationURL(context.Background(), state, codeChallenge) if err != nil { - log.Fatalf("Failed to get authorization URL: %v", err) + return fmt.Errorf("failed to get authorization URL: %w", err) } // Open the browser to the authorization URL fmt.Printf("Opening browser to: %s\n", authURL) openBrowser(authURL) // Wait for the callback fmt.Println("Waiting for authorization callback...") params := <-callbackChan // Verify state parameter if params["state"] != state { - log.Fatalf("State mismatch: expected %s, got %s", state, params["state"]) + return fmt.Errorf("state mismatch: expected %s, got %s", state, params["state"]) } // Exchange the authorization code for a token code := params["code"] if code == "" { - log.Fatalf("No authorization code received") + return fmt.Errorf("no authorization code received") } fmt.Println("Exchanging authorization code for token...") err = oauthHandler.ProcessAuthorizationResponse(context.Background(), code, state, codeVerifier) if err != nil { - log.Fatalf("Failed to process authorization response: %v", err) + return fmt.Errorf("failed to process authorization response: %w", err) } fmt.Println("Authorization successful!") + return nil } + return err }Then update the callers:
- maybeAuthorize(err) + if err := maybeAuthorize(err); err != nil { + log.Fatalf("Failed to authorize: %v", err) + }
🧹 Nitpick comments (3)
client/transport/oauth.go (1)
140-146: Good defensive programming for token type normalization.The normalization of lowercase "bearer" to "Bearer" improves compatibility with OAuth implementations that may return inconsistent casing. This follows the OAuth 2.0 specification which uses "Bearer" with a capital B.
Consider making this more comprehensive by handling other potential variations:
// Some auth implementations are strict about token type tokenType := token.TokenType -if tokenType == "bearer" { - tokenType = "Bearer" -} +if strings.ToLower(tokenType) == "bearer" { + tokenType = "Bearer" +}This would handle variations like "BEARER", "Bearer", "bearer", etc.
client/oauth.go (1)
38-39: Fix the function documentation.The comment incorrectly describes this as a "streamable-http-based" client when it should be "SSE-based".
-// NewOAuthStreamableHttpClient creates a new streamable-http-based MCP client with OAuth support. +// NewOAuthSSEClient creates a new SSE-based MCP client with OAuth support.examples/oauth_client/main.go (1)
69-88: Consider refactoring to reduce code duplication.The
InitializeRequeststructure is duplicated. Consider extracting it to a variable to improve maintainability.- result, err := c.Initialize(context.Background(), mcp.InitializeRequest{ - Params: struct { - ProtocolVersion string `json:"protocolVersion"` - Capabilities mcp.ClientCapabilities `json:"capabilities"` - ClientInfo mcp.Implementation `json:"clientInfo"` - }{ - ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, - ClientInfo: mcp.Implementation{ - Name: "mcp-go-oauth-example", - Version: "0.1.0", - }, - }, - }) + initRequest := mcp.InitializeRequest{ + Params: struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities mcp.ClientCapabilities `json:"capabilities"` + ClientInfo mcp.Implementation `json:"clientInfo"` + }{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{ + Name: "mcp-go-oauth-example", + Version: "0.1.0", + }, + }, + } + + result, err := c.Initialize(context.Background(), initRequest) if err != nil { maybeAuthorize(err) - result, err = c.Initialize(context.Background(), mcp.InitializeRequest{ - Params: struct { - ProtocolVersion string `json:"protocolVersion"` - Capabilities mcp.ClientCapabilities `json:"capabilities"` - ClientInfo mcp.Implementation `json:"clientInfo"` - }{ - ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, - ClientInfo: mcp.Implementation{ - Name: "mcp-go-oauth-example", - Version: "0.1.0", - }, - }, - }) + result, err = c.Initialize(context.Background(), initRequest) if err != nil { log.Fatalf("Failed to initialize client: %v", err) } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
client/oauth.go(2 hunks)client/transport/oauth.go(1 hunks)client/transport/sse.go(10 hunks)client/transport/sse_oauth_test.go(1 hunks)client/transport/streamable_http.go(1 hunks)client/transport/streamable_http_oauth_test.go(3 hunks)examples/oauth_client/main.go(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (5)
client/transport/streamable_http.go (2)
client/oauth.go (1)
OAuthConfig(11-11)client/transport/oauth.go (1)
OAuthConfig(18-34)
client/transport/streamable_http_oauth_test.go (3)
client/transport/streamable_http.go (2)
NewStreamableHTTP(84-110)WithHTTPOAuth(44-48)client/oauth.go (1)
OAuthConfig(11-11)client/transport/oauth.go (1)
OAuthConfig(18-34)
examples/oauth_client/main.go (1)
mcp/types.go (6)
InitializeRequest(389-398)Params(154-154)ClientCapabilities(427-437)Implementation(468-471)LATEST_PROTOCOL_VERSION(99-99)ListResourcesRequest(526-528)
client/oauth.go (4)
client/transport/streamable_http.go (1)
WithHTTPOAuth(44-48)client/transport/oauth.go (1)
OAuthConfig(18-34)client/transport/sse.go (3)
ClientOption(45-45)WithOAuth(65-69)NewSSE(73-99)client/client.go (2)
Client(16-25)NewClient(44-54)
client/transport/sse.go (4)
client/transport/oauth.go (3)
OAuthHandler(111-119)OAuthConfig(18-34)NewOAuthHandler(122-131)client/oauth.go (3)
OAuthConfig(11-11)OAuthAuthorizationRequiredError(61-61)GetOAuthHandler(70-76)client/client.go (1)
ClientOption(27-27)client/transport/streamable_http.go (2)
OAuthAuthorizationRequiredError(162-164)ErrOAuthAuthorizationRequired(159-159)
🔇 Additional comments (13)
client/transport/streamable_http.go (1)
43-44: Good naming improvement for API clarity.Renaming
WithOAuthtoWithHTTPOAuthclearly distinguishes this option from the SSE transport's OAuth option, improving the API's usability and reducing potential confusion.client/transport/streamable_http_oauth_test.go (1)
76-76: Test updates correctly reflect the API changes.All test function calls have been properly updated to use
WithHTTPOAuthinstead ofWithOAuth, maintaining consistency with the renamed function in the implementation.Also applies to: 166-166, 207-207
client/oauth.go (2)
29-29: Correctly updated to use the renamed function.The update to
transport.WithHTTPOAuthis consistent with the function rename in the transport layer.
40-49: Well-implemented SSE OAuth client constructor.The new function follows the same pattern as the HTTP OAuth client, with proper error handling and consistent API design. The use of
transport.WithOAuth(for SSE) vstransport.WithHTTPOAuth(for HTTP) maintains clear separation between the two transport types.examples/oauth_client/main.go (3)
46-50: LGTM! Clean retry pattern for OAuth authorization.The retry logic correctly handles the OAuth authorization flow by calling
maybeAuthorizeand retrying theStartoperation.
89-104: Good addition to demonstrate successful OAuth flow.The resource listing effectively shows that the client is properly authenticated and can make API calls.
132-136: Verify if client registration is always required.The code unconditionally calls
RegisterClient, which might fail if:
- The client is already registered
- The server doesn't support dynamic client registration
Consider checking if registration is needed or handling registration errors more gracefully.
client/transport/sse.go (5)
8-8: LGTM! Clean OAuth integration.The addition of the
errorsimport andoauthHandlerfield properly sets up OAuth support for the SSE transport.Also applies to: 41-43
65-70: LGTM! Well-structured client option.The
WithOAuthoption follows the established pattern and properly initializes the OAuth handler.
91-97: LGTM! Proper OAuth metadata discovery setup.The code correctly extracts the base URL and sets it for OAuth metadata discovery.
154-161: LGTM! Consistent 401 error handling.The code properly handles 401 Unauthorized responses by returning
OAuthAuthorizationRequiredErroracross all methods.Also applies to: 377-384, 478-485
506-515: LGTM! Useful OAuth status accessors.The
GetOAuthHandlerandIsOAuthEnabledmethods provide clean access to OAuth configuration status.client/transport/sse_oauth_test.go (1)
1-238: Excellent test coverage for OAuth functionality!The tests comprehensively cover:
- OAuth authorization flow with retry logic
- 401 Unauthorized error handling
- OAuth enablement status checking
- Authorization header verification
The test structure is clean and follows Go testing best practices.
|
Thank you @giridhar-murthy-glean! |
Description
This PR adds OAuth support to SSE clients. This is necessary to deal with MCP servers that exhibit the behavior described in the backwards compatibility section of the spec:
https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#backwards-compatibility
Fixes #<issue_number> (if applicable)
Type of Change
Checklist
MCP Spec Compliance
Additional Information
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Chores