auth: support OAuth2 session persistence (#1058)

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.

Fixes: #901.

See #901 for the previous discussion behind this design.

---------

Co-authored-by: Guglielmo Colombo <guglielmoc@google.com>
Co-authored-by: Jonathan Amsterdam <jba@users.noreply.github.com>

auth/auth_example_test.go

@@ -0,0 +1,222 @@
+// 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,
+	}
+}
+
+// sessionStore is an in-memory OAuth 2.0 session store used to demonstrate
+// persistence. A production implementation would persist the session to disk or
+// a secret store.
+type sessionStore struct {
+	config *oauth2.Config
+	token  *oauth2.Token
+}
+
+// save persists the given OAuth 2.0 config and token.
+func (s *sessionStore) save(config *oauth2.Config, token *oauth2.Token) error {
+	fmt.Printf("Saving token: %s\n", token.AccessToken)
+	s.config = config
+	s.token = token
+	return nil
+}
+
+// restore loads a previously persisted OAuth 2.0 session, if one exists.
+func (s *sessionStore) restore() (*oauth2.Config, *oauth2.Token, error) {
+	if s.config != nil && s.token != nil {
+		fmt.Println("Restoring session.")
+	} else {
+		fmt.Println("No session found to restore.")
+	}
+	return s.config, s.token, nil
+}
+
+// newMockAuthServer returns an httptest.Server that simulates both an MCP
+// resource server requiring authorization and its OAuth authorization server.
+func newMockAuthServer() *httptest.Server {
+	return 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)
+		}
+	}))
+}
+
+// This example shows how OAuth2 session persistence might be implemented. It
+// connects twice using a shared in-memory session store: the first connection
+// authorizes and saves the session, and the second restores and reuses it.
+func Example_persistence() {
+	// Simulate an MCP server that requires authorization and an OAuth server.
+	mockServer := newMockAuthServer()
+	defer mockServer.Close()
+
+	// store persists the OAuth2 session across both connections.
+	store := &sessionStore{}
+
+	// connect performs a single client connection, restoring any saved session
+	// beforehand and saving any new session acquired during authorization.
+	connect := func() {
+		// Load the OAuth2 session if available, and use it.
+		var initialTS oauth2.TokenSource
+		cfg, tok, err := store.restore()
+		if err == nil && cfg != nil && tok != nil {
+			initialTS = NewSavingTokenSource(
+				cfg.TokenSource(context.Background(), tok),
+				cfg, tok, store.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.
+				_ = store.save(cfg, token)
+				return NewSavingTokenSource(
+					cfg.TokenSource(ctx, token), cfg, token, store.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)
+	}
+
+	// The first connection has no saved session, so it authorizes and saves.
+	fmt.Println("--- First connect ---")
+	connect()
+	// The second connection restores the saved session and reuses it, so no
+	// further authorization or save occurs.
+	fmt.Println("--- Second connect ---")
+	connect()
+
+	// Output:
+	// --- First connect ---
+	// No session found to restore.
+	// No token source found. Transport is calling Authorize()...
+	// Saving token: mock-token
+	// --- Second connect ---
+	// Restoring session.
+}

auth/authorization_code.go

@@ -14,6 +14,7 @@ import (
 	"net/url"
 	"slices"
 	"strings"
+	"sync"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/authutil"
 	"github.com/modelcontextprotocol/go-sdk/internal/util"
@@ -126,6 +127,24 @@ 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]. Implementations must use the
+	// provided context, which is properly configured for constructing a
+	// 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
@@ -133,6 +152,9 @@ type AuthorizationCodeHandlerConfig struct {
 type AuthorizationCodeHandler struct {
 	config *AuthorizationCodeHandlerConfig
 
+	// mu protects concurrent access to tokenSource and grantedScopes.
+	mu sync.RWMutex
+
 	// tokenSource is the token source to use for authorization.
 	tokenSource oauth2.TokenSource
 
@@ -143,6 +165,8 @@ type AuthorizationCodeHandler struct {
 var _ OAuthHandler = (*AuthorizationCodeHandler)(nil)
 
 func (h *AuthorizationCodeHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
 	return h.tokenSource, nil
 }
 
@@ -199,6 +223,7 @@ func NewAuthorizationCodeHandler(config *AuthorizationCodeHandlerConfig) (*Autho
 	}
 	return &AuthorizationCodeHandler{
 		config:        config,
+		tokenSource:   config.InitialTokenSource,
 		grantedScopes: make(map[string][]string),
 	}, nil
 }
@@ -304,7 +329,10 @@ func (h *AuthorizationCodeHandler) Authorize(ctx context.Context, req *http.Requ
 	// Accumulate scopes: union previously granted scopes with the newly
 	// challenged scopes so that step-up authorization does not lose
 	// permissions granted in earlier rounds (SEP-2350).
-	requestedScopes = authutil.UnionScopes(h.grantedScopes[asm.Issuer], requestedScopes)
+	h.mu.RLock()
+	granted := h.grantedScopes[asm.Issuer]
+	h.mu.RUnlock()
+	requestedScopes = authutil.UnionScopes(granted, requestedScopes)
 
 	cfg := &oauth2.Config{
 		ClientID:     resolvedClientConfig.clientID,
@@ -615,23 +643,42 @@ 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)
+	var ts oauth2.TokenSource
+	if h.config.NewTokenSource == nil {
+		ts = cfg.TokenSource(refreshCtx, token)
+	} else {
+		var err error
+		ts, err = h.config.NewTokenSource(refreshCtx, cfg, token)
+		if err != nil {
+			return fmt.Errorf("constructing token source failed: %w", err)
+		}
+	}
+	h.mu.Lock()
+	h.tokenSource = ts
+	h.mu.Unlock()
 	return nil
 }
 
 // updateGrantedScopes updates the granted scopes based on the token source and requested scopes.
 func (h *AuthorizationCodeHandler) updateGrantedScopes(issuer string, requestedScopes []string) error {
-	if h.tokenSource == nil {
+	h.mu.RLock()
+	ts := h.tokenSource
+	h.mu.RUnlock()
+
+	if ts == nil {
 		return nil
 	}
-	tok, err := h.tokenSource.Token()
+	tok, err := ts.Token()
 	if err != nil {
 		return err
 	}
+
+	h.mu.Lock()
 	if tokenScopes := authutil.ScopesFromToken(tok); tokenScopes == nil {
 		h.grantedScopes[issuer] = requestedScopes
 	} else {
 		h.grantedScopes[issuer] = tokenScopes
 	}
+	h.mu.Unlock()
 	return nil
 }

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)
+	}
+}