Proposal: Client API support for saving and restoring OAuth2 sessions

[smlx]

Is your feature request related to a problem? Please describe.

I am adding MCP OAuth2 support to a client and don't want my users to be re-prompted for authorization if there is still a valid refresh token from their last session.

Currently, clients using AuthorizationCodeHandler and StreamableClientTransport do not have a built-in mechanism to persist and restore OAuth2 sessions across application restarts.

Describe the solution you'd like

I propose adding an API to support saving and restoring OAuth2 sessions.

The specific additions to the auth package API are:

1. Add a SessionSaver callback function to AuthorizationCodeHandlerConfig:

	// SessionSaver is an optional function that can be set to save oauth2
	// sessions to persistent storage. If it is set:
	//  - It is called immediately after the authorization code is successfully
	//    exchanged for a token in [AuthorizationCodeHandler.Authorize].
	//  - After a successful call to [AuthorizationCodeHandler.Authorize], the
	//    token source returned by [AuthorizationCodeHandler.TokenSource] will be
	//    a [NewSavingTokenSource].
	// To ensure SessionSaver is called on subsequent token refreshes, clients
	// must use the token source returned by [AuthorizationCodeHandler.TokenSource]
	// (for example, by returning it from your own [OAuthHandler.TokenSource]
	// implementation) after [AuthorizationCodeHandler.Authorize] returns
	// successfully.
	SessionSaver func(*oauth2.Config, *oauth2.Token)

2. Add a NewSavingTokenSource function:

// NewSavingTokenSource persists OAuth 2.0 sessions by intercepting token
// refreshes from the wrapped [oauth2.TokenSource]. When this wrapper detects
// an access token refresh, it calls the provided session saver with the
// [oauth2.Config] and the new [oauth2.Token].
//
// The initial token argument prevents saver from being called unnecessarily if
// a valid token already exists. If initial is invalid or nil, saver may be
// called once before the wrapped token is refreshed.
//
// Clients implementing [OAuthHandler.TokenSource] can use this constructor to
// wrap token sources loaded from storage to ensure that subsequent refreshes
// during the application's lifetime are intercepted.
func NewSavingTokenSource(wrapped oauth2.TokenSource, config *oauth2.Config, initialToken *oauth2.Token, saver func(*oauth2.Config, *oauth2.Token)) oauth2.TokenSource

This design is heavily inspired by the initial design in #785. The config argument was added so that the entire oauth2 session can be persisted. The context was removed as I didn't understand the need for it. I haven't found it necessary but am open to re-adding it if there is a good reason.

The two individual API additions handle saving the oauth2 session during its two phases:

1. After the authorization code exchange in Authorize().
2. Later when the access token expires and is automatically refreshed.

This design is more complex than I would like because it requires passing a session saver function in two places to manage these two phases of the token lifecycle. However I wasn't able to come up with a simpler abstraction within the constraints of the existing API.

Describe alternatives you've considered

Individual SaveConfig / SaveToken / RestoreConfig / RestoreToken methods/fields.

The lifecycle of Token and Config does not align. oauth2.Token will change on every refresh. oauth2.Config will change only during Authorize(). So the minimal work required during a token refresh would be to call SaveToken. However this design has several problems:

• What the client actually cares about is the "session". The Token is useless without the Config. So it is likely that clients will save Token and Config together as "session state".
• It is a larger API surface and relies on careful ordering of function calls.
• The client has to individually save Token and Config to their persistent storage, and then restore them individually using the two restore functions, but still ensure the associated Token and Config are restored together.

Additional context

See discussion about this topic on #881.

[jba]

Would a more general design be to add a constructor to AuthorizationCodeHandlerConfig?

NewTokenSource func(context.Context, *oauth2.Config, *oauth2.Token) (oauth2.TokenSource, error)

The only other code change would be changing this line in authorization_code.go:

h.tokenSource = cfg.TokenSource(clientCtx, token)

to this:

if h.config.NewTokenSource != nil {
    h.tokenSource, err = h.config.NewTokenSource(clientCtx, cfg, token)
    if err ...
} else {
    // the above
}

By the way, we add a context whenever there is a potential for a long-running operation, like I/O. That way the caller can cancel it.

[smlx]

Thank you for the thoughtful suggestion. That API design is indeed smaller and cleaner than the one I came up with.

I've implemented it and done some testing locally. It requires a bit more code on the client side but not an excessive amount. And the desired functionality is all there.

I came up with this documentation for the new API. Does it make sense?

// NewTokenSource is an optional function that can be set to construct the
// token source that will be used by the [AuthorizationCodeHandler]. If it is
// set:
//  - It is called immediately after the authorization code is successfully
//    exchanged for a token in [AuthorizationCodeHandler.Authorize].
//  - After a successful call to [AuthorizationCodeHandler.Authorize], the
//    token source returned by [AuthorizationCodeHandler.TokenSource] will be
//    the token source returned from [NewTokenSource].
// Clients can ensure the same token source is used throughout a
// [StreamableClientTransport] connection lifecycle by obtaining it from
// [AuthorizationCodeHandler.TokenSource] and returning it from the
// [OAuthHandler.TokenSource] implementation.
NewTokenSource func(context.Context, *oauth2.Config, *oauth2.Token) (oauth2.TokenSource, error)

[jba]

Looks good. We don't usually use bullets in quite that way. How about:

"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]."

Then a blank commented line "//" to separate paragraphs, then "Clients can ensure...".

If you're up for it, it would be great to have an example that shows the full save/restore logic, like NewSavingTokenSource, so people can see how to do that. You can define the example in terms of two functions, save and restore, which you don't have to implement. Just provide dummy implementations in the example file so everything compiles.

[maciej-kisiel]

This would still require a SetTokenSource method on the handler to restore the OAuth session, correct?

Together, these two APIs should allow to save and restore tokens, but they require more work to achieve it. If we decide to go this way, I believe we should add a proper example showing how they are intended do be used. EDIT: I see Jonathan already proposed that, I guess I failed to read his comment until the end.

[smlx]

This would still require a SetTokenSource method on the handler to restore the OAuth session, correct?

No, I don't think so. Implementing an OAuthHandler that does the loading/restoring of the session on the transport is enough for the session restore to work.

[maciej-kisiel]

What do you mean by "Implementing an OAuthHandler"?

Maybe let me describe what I mean precisely: let's assume the user wants to use auth.AuthorizationCodeHandler as the OAuth handler (I believe the APIs we're designing here are for this type). When you have the OAuth data persisted, you probably want to feed it into the handler without triggering the whole authorization flow. How do you want to do that without SetTokenSource or something similar? NewTokenSource, if I understand the proposal correctly is supposed to only be called at the end of the authorization flow, so it doesn't address this scenario.

[smlx]

Your question makes me think that I've been using this API in a way that it was not designed for 😄

Here is the way I have been doing it: calling the AuthorizationCodeHandler manually inside my own OAuthHandler.

func (h *myOAuthHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error {
	// set up redirect listener
  // ...

	// configure the mcp go-sdk authz code handler
	authCfg := &auth.AuthorizationCodeHandlerConfig{
		// ...
		NewTokenSource: func(_ context.Context, cfg *oauth2.Config, tok *oauth2.Token) (oauth2.TokenSource, error) {
			if err := h.store.SaveOAuth2Session(cfg, tok); err != nil {
				slog.Warn("Failed to save MCP OAuth2 session", "name", h.name, "error", err)
			}
			return NewSavingTokenSource(cfg.TokenSource(ctx, tok), cfg, tok,
				func(cfg *oauth2.Config, tok *oauth2.Token) {
					if err := h.store.SaveOAuth2Session(cfg, tok); err != nil {
						slog.Warn("Failed to save MCP OAuth2 session", "name", h.name, "error", err)
					}
				}), nil
		},
	}

	// run the auth code flow
	sdkHandler, err := auth.NewAuthorizationCodeHandler(authCfg)
	if err != nil {
		return fmt.Errorf("couldn't create authz code handler: %v", err)
	}
	if err := sdkHandler.Authorize(authCtx, req, resp); err != nil {
		return fmt.Errorf("couldn't perform authz flow: %v", err)
	}
	// get the new token source, and save it to this OAuthHandler
	ts, err := sdkHandler.TokenSource(authCtx)
	if err != nil {
		return fmt.Errorf("couldn't get token source after authorize: %v", err)
	}
	h.tokenSource = ts

	return nil
}

func (h *myOAuthHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
	// Check if token source already available.
	if h.tokenSource != nil {
		return h.tokenSource, nil
	}
	// Otherwise try to load it from disk. If there isn't a state available just
	// return nil, as this will trigger a call to Authorize().
	cfg, tok, err := h.store.LoadOAuth2Session()
	if errors.Is(err, state.ErrNoStore) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("couldn't load MCP OAuth2 state: %v", err)
	}
	if cfg == nil || tok == nil {
		return nil, nil
	}
	h.tokenSource = NewSavingTokenSource(cfg.TokenSource(ctx, tok), cfg, tok,
		func(cfg *oauth2.Config, tok *oauth2.Token) {
			if err := h.store.SaveOAuth2Session(cfg, tok); err != nil {
				slog.Warn("Failed to save MCP OAuth2 token", "name", h.name, "error", err)
			}
		})
	return h.tokenSource, nil
}

I understand now from your question that AuthorizationCodeHandler is supposed to be configured directly as the transport OAuthHandler? And you want the save/restore functionality added directly to the AuthorizationCodeHandler? If so, I'll have to rethink the whole design of this session save/restore API.

[maciej-kisiel]

Yes, AuthorizationCodeHandler was generally designed to be used directly. The fact that you are wrapping it means that you have some needs that it was not addressing. Are there any additional such needs besides saving and restoration of tokens?

[smlx]

No. Saving/restoring of OAuth2 sessions (config+token) is the only missing functionality I have found at this stage.

[smlx]

I prototyped a potential SetTokenSource, and can confirm that it allows the AuthorizationCodeHandler to be used directly as a transport OAuthHandler with save/restore.

Here's what I came up with:

	// SetTokenSource is an optional function that can be set to inject the
	// token source that will be used by the [AuthorizationCodeHandler]. If
	// non-nil, it is called in [AuthorizationCodeHandler.TokenSource] to obtain
	// a token source if the handler does not already have one configured. The
	// default is to return nil if no token source has already been set, which
	// will trigger a call to [AuthorizationCodeHandler.Authorize].
	SetTokenSource func(context.Context) (oauth2.TokenSource, error)

func (h *AuthorizationCodeHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
	if h.tokenSource != nil {
		return h.tokenSource, nil
	}
	if h.config.SetTokenSource != nil {
		ts, err := h.config.SetTokenSource(ctx)
		if err != nil {
			return nil, fmt.Errorf("failed to set token source: %v", err)
		}
		h.tokenSource = ts
	}
	return h.tokenSource, nil
}