Proposal: add StreamableClientTransport.OAuthHandler

[jba]

We should add this field to StreamableClientTransport:

OAuthHandler auth.OAuthHandler

The type is defined here.

I believe that gives us everything we need for auth. It does couple the streamable transport to OAuth, but so does the spec, and the TypeScript SDK.

Our SDK will call the handler when a request returns a 403. If the handler returns a TokenSource, the SDK will use it to set the Authorization header on subsequent calls.

If the handler can't return a TokenSource because user approval must occur out of band, it should return an error, which will be passed up to the caller. The caller can then do whatever it needs to obtain a token, and then begin again with with a new transport whose handler simply returns the token. Or one handler can suffice for both cases: it first looks for a token somewhere (a shared variable), and if it finds it, returns it. Otherwise, it conducts the OAuth flow and returns the resulting token.

Previously, we tried to do this entirely within an http.Client. That approach is cleaner, but is a bit messy to implement, and more important, has an unfixable bug (until golang/go#75814 is available) in which a RoundTripper inserts an Authorization header incorrectly.

We also proposed adding a ModifyRequest option. That is more general, and maybe we should still add it to support adding other headers and modifying the request in other ways. But for setting the Authorization header, I think this proposal suffices. If you have an access token, you can get the effect of ModifyRequest easily enough: wrap it in a TokenSource by creating a Token and passing it to a StaticTokenSource, and then your handler can just return that TokenSource.

[jba]

/cc @findleyr @neild @wagnerjt

[wagnerjt]

Heya @jba -- I took a moment to test the OAuthHandler you have by extending the auth-middleware example to test out a simple flow to see if I could:

1. Leverage the oauth2.StaticTokenSource on the apikey route
2. Use the oauthex package to create the PRM on the server and add the rest of the client logic leveraging the jwt route

So far I am still on the apikey route, and I am seeing the client code is failing to initialize and actually is returning nil for the client session as well as the error on .Connect. You can reproduce this using the auth middleware example as-is for the server and the following code for the client. Essentially, I can validate on 401s, the client is executing the oAuthHandler function as expected and adding the static token source. However, on the second attempt that occurs automatically, the server responds with a 400 and this code is returning the raw response and error as nil if it isn't a 401. I haven't dug more into the internal client connect to see where the client session and error aren't being handled in this scenario if you can help me here.

package main

import (
	"context"
	"fmt"
	"iter"
	"log"
	"net/http"

	"github.com/modelcontextprotocol/go-sdk/auth"
	"github.com/modelcontextprotocol/go-sdk/mcp"
	"golang.org/x/oauth2"
)

const (
	endpoint     = "http://localhost:8080/mcp/apikey"
	staticApiKey = "sk-1234567890abcdef"
)

func main() {
	ctx := context.Background()
	apiTokenSource := oauth2.StaticTokenSource(&oauth2.Token{
		AccessToken: staticApiKey,
		TokenType:   "Bearer",
	})
	// test static api token for apikey
	oAuthHandler := func(ctx context.Context, args auth.OAuthHandlerArgs) (oauth2.TokenSource, error) {
		return apiTokenSource, nil
	}

	httpTransport, err := auth.NewHTTPTransport(oAuthHandler, nil)
	if err != nil {
		log.Fatalf("NewHTTPTransport() failed: %v", err)
	}
	httpClient := &http.Client{Transport: httpTransport}

	transport := &mcp.StreamableClientTransport{
		Endpoint:   endpoint,
		HTTPClient: httpClient,
	}

	client := mcp.NewClient(&mcp.Implementation{Name: "mcp-client", Version: "v1.0.0"}, nil)
	cs, err := client.Connect(ctx, transport, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer cs.Close()

	t := cs.InitializeResult().Capabilities.Tools
	log.Println(t)
}

However, I did manage to get a successful client initialization and tool list + calls by explicitly adding the oAuth2 option in the httpClient transport like so:

httpTransport, err := auth.NewHTTPTransport(oAuthHandler, &auth.HTTPTransportOptions{
		Base: &oauth2.Transport{Base: http.DefaultTransport, Source: apiTokenSource},
	})

[wagnerjt]

Ah okay, my issue seems to be the same reported in #590. I'll work on getting the actual oauth flow example this week

[jba]

We just fixed #590. Can you try at head?

[wagnerjt]

Great! I just tested on the head (v1.1.0-pre.1) and it seems to be working with the initial example ☝️. For knowledge/docs when the error that returns when the ClientSession cannot connect due to auth, we receive calling "initialize": broken session: 401 Unauthorized which is sufficient!

I'll continue with the jwt example and share updates when I carve time for it