auth: support OAuth2 session persistence

smlx · smlx · commit 7615caf33f0a · 2026-07-08T21:36:44.000+08:00
This change adds NewTokenSource and InitialTokenSource fields to
AuthorizationCodeHandlerConfig. NewTokenSource enables wrapping the
underlying token source to intercept token refreshes. InitialTokenSource
enables configuring the token source used by the
AuthorizationCodeHandler on initialization. Together these features
facilitate persisting OAuth2 sessions across restarts.

Includes a package example demonstrating the persistence pattern.
diff --git a/auth/auth_example_test.go b/auth/auth_example_test.go
@@ -0,0 +1,184 @@
+// Copyright 2026 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by the license
+// that can be found in the LICENSE file.
+
+package auth_test
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync"
+
+	"github.com/modelcontextprotocol/go-sdk/auth"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/modelcontextprotocol/go-sdk/oauthex"
+	"golang.org/x/oauth2"
+)
+
+// savingTokenSource is an oauth2.TokenSource that passes the oauth2 config and
+// token to the given saver function each time the access token value changes.
+type savingTokenSource struct {
+	mu          sync.Mutex
+	src         oauth2.TokenSource
+	saver       func(*oauth2.Config, *oauth2.Token) error
+	config      *oauth2.Config
+	accessToken string
+}
+
+func (s *savingTokenSource) Token() (*oauth2.Token, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	tok, err := s.src.Token()
+	if err != nil {
+		return nil, err
+	}
+	if s.accessToken != tok.AccessToken {
+		s.accessToken = tok.AccessToken
+		// This saver implementation always returns nil.
+		_ = s.saver(s.config, tok)
+	}
+	return tok, nil
+}
+
+// NewSavingTokenSource persists OAuth 2.0 sessions by intercepting token
+// changes from the wrapped oauth2.TokenSource. When this wrapper detects an
+// access token change, it calls the provided session saver with the
+// oauth2.Config and the new oauth2.Token.
+//
+// initial is an optional access token the caller may already hold (such as a
+// token loaded from storage). It initialises the wrapper's state so that if
+// the first wrapped.Token() call returns the same token, does not trigger a
+// redundant call to saver(). Pass nil when there is no existing token, in
+// which case the first token produced by wrapped.Token() is saved.
+func NewSavingTokenSource(wrapped oauth2.TokenSource, config *oauth2.Config, initial *oauth2.Token, saver func(*oauth2.Config, *oauth2.Token) error) oauth2.TokenSource {
+	if wrapped == nil {
+		return nil
+	}
+	if saver == nil {
+		return wrapped
+	}
+	var accessToken string
+	if initial != nil {
+		accessToken = initial.AccessToken
+	}
+	return &savingTokenSource{
+		src:         wrapped,
+		saver:       saver,
+		config:      config,
+		accessToken: accessToken,
+	}
+}
+
+// This example shows how OAuth2 session persistence might be implemented.
+func Example_persistence() {
+	var savedConfig *oauth2.Config
+	var savedToken *oauth2.Token
+
+	// Save OAuth2 session.
+	// This function might persist the session to disk or key storage.
+	save := func(config *oauth2.Config, token *oauth2.Token) error {
+		fmt.Printf("Saving token: %s\n", token.AccessToken)
+		savedConfig = config
+		savedToken = token
+		return nil
+	}
+
+	// Restore OAuth2 session.
+	// This function might load the session from disk or key storage.
+	restore := func() (*oauth2.Config, *oauth2.Token, error) {
+		if savedConfig != nil && savedToken != nil {
+			fmt.Println("Restoring session.")
+		} else {
+			fmt.Println("No session found to restore.")
+		}
+		return savedConfig, savedToken, nil
+	}
+
+	// Simulate an MCP server that requires authorization and an OAuth server.
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/.well-known/oauth-authorization-server":
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = fmt.Fprintf(w, `{"issuer": "http://%s", "authorization_endpoint": "http://%s/auth", "token_endpoint": "http://%s/token", "code_challenge_methods_supported": ["S256"]}`, r.Host, r.Host, r.Host)
+		case "/token":
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"access_token": "mock-token", "token_type": "bearer"}`))
+		default:
+			// The mock MCP endpoint returns 401 until the client presents a valid
+			// bearer token.
+			if r.Header.Get("Authorization") == "Bearer mock-token" {
+				// A real server would return a valid MCP message here. The empty
+				// response causes Connect to return an error after authorization,
+				// which is ignored for this example which demonstrates token
+				// persistence only.
+				w.WriteHeader(http.StatusOK)
+				return
+			}
+			w.Header().Set("WWW-Authenticate", "Bearer")
+			w.WriteHeader(http.StatusUnauthorized)
+		}
+	}))
+	defer mockServer.Close()
+
+	// Load the OAuth2 session if available, and use it.
+	var initialTS oauth2.TokenSource
+	cfg, tok, err := restore()
+	if err == nil && cfg != nil && tok != nil {
+		initialTS = NewSavingTokenSource(cfg.TokenSource(context.Background(), tok), cfg, tok, save)
+	}
+
+	// Configure and initialize the AuthorizationCodeHandler with a
+	// NewTokenSource that saves the session when the access token changes, and
+	// the InitialTokenSource set to the session loaded via restore().
+	config := &auth.AuthorizationCodeHandlerConfig{
+		RedirectURL: "http://localhost/callback",
+		PreregisteredClient: &oauthex.ClientCredentials{
+			ClientID: "example",
+		},
+		Client: mockServer.Client(),
+		AuthorizationCodeFetcher: func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) {
+			fmt.Println("No token source found. Transport is calling Authorize()...")
+			// Extract the generated state from the authorization URL
+			u, _ := url.Parse(args.URL)
+			state := u.Query().Get("state")
+			return &auth.AuthorizationResult{Code: "mock-code", State: state}, nil
+		},
+		NewTokenSource: func(ctx context.Context, cfg *oauth2.Config, token *oauth2.Token) (oauth2.TokenSource, error) {
+			// This save implementation always returns nil.
+			_ = save(cfg, token)
+			return NewSavingTokenSource(cfg.TokenSource(ctx, token), cfg, token, save), nil
+		},
+		InitialTokenSource: initialTS,
+	}
+	handler, err := auth.NewAuthorizationCodeHandler(config)
+	if err != nil {
+		fmt.Printf("Error creating handler: %v\n", err)
+		return
+	}
+
+	// Set the constructed handler on a transport.
+	transport := &mcp.StreamableClientTransport{
+		Endpoint:     mockServer.URL + "/sse",
+		OAuthHandler: handler,
+	}
+
+	// Create a client and attempt to connect using the configured transport.
+	// The transport will automatically:
+	// 1. Call TokenSource() to check for an existing session. This will return
+	//    InitialTokenSource, if set.
+	// 2. Try the MCP endpoint, and encounter a 401 response from the mock server.
+	// 3. Call Authorize() to perform the OAuth flow, which calls NewTokenSource.
+	//    In this example NewTokenSource saves the newly acquired session.
+	// 4. Retry the MCP endpoint with a valid bearer token and get a 200.
+	client := mcp.NewClient(&mcp.Implementation{Name: "example", Version: "1.0.0"}, nil)
+	// Response ignored: this example asserts authorization only.
+	_, _ = client.Connect(context.Background(), transport, nil)
+
+	// Output:
+	// No session found to restore.
+	// No token source found. Transport is calling Authorize()...
+	// Saving token: mock-token
+}
diff --git a/auth/authorization_code.go b/auth/authorization_code.go
@@ -126,6 +126,23 @@ type AuthorizationCodeHandlerConfig struct {
 	// https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#server-side-request-forgery-ssrf
 	// If not provided, http.DefaultClient will be used.
 	Client *http.Client
+
+	// NewTokenSource is an optional function that can be set to construct the
+	// token source that will be used by the [AuthorizationCodeHandler]. If
+	// non-nil, it is called after the authorization code is successfully
+	// exchanged for a token in [AuthorizationCodeHandler.Authorize]
+	// to obtain the [oauth2.TokenSource] returned by
+	// [AuthorizationCodeHandler.TokenSource]. The default is to call
+	// [oauth2.Config.TokenSource].
+	NewTokenSource func(context.Context, *oauth2.Config, *oauth2.Token) (oauth2.TokenSource, error)
+
+	// InitialTokenSource is an optional field that can be set to inject the
+	// token source that will be used by the [AuthorizationCodeHandler]. If
+	// non-nil, it is set as the token source that will be returned by
+	// [AuthorizationCodeHandler.TokenSource] during handler initialization.
+	// The default is nil, which means no token source has been set initially,
+	// and will trigger a call to [AuthorizationCodeHandler.Authorize].
+	InitialTokenSource oauth2.TokenSource
 }
 
 // AuthorizationCodeHandler is an implementation of [OAuthHandler] that uses
@@ -197,10 +214,14 @@ func NewAuthorizationCodeHandler(config *AuthorizationCodeHandlerConfig) (*Autho
 	if config.Client == nil {
 		config.Client = http.DefaultClient
 	}
-	return &AuthorizationCodeHandler{
+	h := &AuthorizationCodeHandler{
 		config:        config,
 		grantedScopes: make(map[string][]string),
-	}, nil
+	}
+	if config.InitialTokenSource != nil {
+		h.tokenSource = config.InitialTokenSource
+	}
+	return h, nil
 }
 
 func isNonRootHTTPSURL(u string) bool {
@@ -615,7 +636,15 @@ func (h *AuthorizationCodeHandler) exchangeAuthorizationCode(ctx context.Context
 	// completes. Use a background context that still carries the configured HTTP
 	// client so refreshes keep working for the life of the token source.
 	refreshCtx := context.WithValue(context.Background(), oauth2.HTTPClient, h.config.Client)
-	h.tokenSource = cfg.TokenSource(refreshCtx, token)
+	if h.config.NewTokenSource != nil {
+		ts, err := h.config.NewTokenSource(refreshCtx, cfg, token)
+		if err != nil {
+			return fmt.Errorf("construction of new token source failed: %v", err)
+		}
+		h.tokenSource = ts
+	} else {
+		h.tokenSource = cfg.TokenSource(refreshCtx, token)
+	}
 	return nil
 }
 
diff --git a/auth/authorization_code_test.go b/auth/authorization_code_test.go
@@ -1233,3 +1233,92 @@ func validConfig() *AuthorizationCodeHandlerConfig {
 		},
 	}
 }
+
+func TestNewTokenSource(t *testing.T) {
+	// mock the /token endpoint to successfully return an access token on code exchange
+	mockTS := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/token" {
+			w.Header().Set("Content-Type", "application/json")
+			w.Write([]byte(`{"access_token": "test_token", "token_type": "bearer"}`))
+			return
+		}
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer mockTS.Close()
+
+	// configure the handler and set NewTokenSource
+	var called bool
+	handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{
+		RedirectURL: "http://localhost/callback",
+		PreregisteredClient: &oauthex.ClientCredentials{
+			ClientID: "test_client",
+		},
+		AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) {
+			u, _ := url.Parse(args.URL)
+			return &AuthorizationResult{
+				Code:  "test_code",
+				State: u.Query().Get("state"),
+			}, nil
+		},
+		NewTokenSource: func(ctx context.Context, cfg *oauth2.Config, token *oauth2.Token) (oauth2.TokenSource, error) {
+			called = true
+			return oauth2.StaticTokenSource(token), nil
+		},
+		Client: mockTS.Client(),
+	})
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	// Simulate a 401 response from a resource server.
+	// The WWW-Authenticate: Bearer header triggers the authorization logic.
+	req := httptest.NewRequest(http.MethodGet, mockTS.URL, nil)
+	resp := &http.Response{
+		StatusCode: http.StatusUnauthorized,
+		Header:     make(http.Header),
+		Body:       http.NoBody,
+		Request:    req,
+	}
+	resp.Header.Set("WWW-Authenticate", "Bearer")
+
+	// Authorize and confirm NewTokenSource was called.
+	err = handler.Authorize(t.Context(), req, resp)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !called {
+		t.Error("expected NewTokenSource to be called")
+	}
+}
+
+func TestInitialTokenSource(t *testing.T) {
+	handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{
+		RedirectURL: "http://localhost:12345/callback",
+		PreregisteredClient: &oauthex.ClientCredentials{
+			ClientID: "test_client_id",
+		},
+		AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) {
+			return nil, nil
+		},
+		InitialTokenSource: oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "set_token"}),
+	})
+	if err != nil {
+		t.Fatalf("NewAuthorizationCodeHandler failed: %v", err)
+	}
+
+	ts, err := handler.TokenSource(t.Context())
+	if err != nil {
+		t.Fatalf("failed to get token source: %v", err)
+	}
+	if ts == nil {
+		t.Fatal("expected token source to be non-nil")
+	}
+
+	tok, err := ts.Token()
+	if err != nil {
+		t.Fatalf("failed to get Token: %v", err)
+	}
+	if tok.AccessToken != "set_token" {
+		t.Errorf("expected access token 'set_token', got '%s'", tok.AccessToken)
+	}
+}
