Support hooking token refresh

[smlx]

I'm working on an MCP client and using the new OAuth2 functionality in the v1.5.0 pre-release has worked well for me!

However, I have found an area of the API that I'm having to hack around and I'm not sure if it is a missing feature or if I'm just missing something.

My client stores the oauth2.Token locally so it can load it on start for use by the mcp client. I want to hook the automatic refresh of the token so that I can store the updated access and refresh token in my client state on disk. What is the best way to achieve this?

Right now I'm using a custom AuthorizationCodeHandlerConfig.Client that intercepts the refresh request but it feels very hacky. Would appreciate any advice. Thanks!

edit: Similarly, I would like to be able to extract the ClientID, ClientSecret, and token Endpoint so that I can reconstruct a oauth2.Config object.
https://pkg.go.dev/golang.org/x/oauth2#Config

[jba]

I wonder if you can wrap the token in a ReuseTokenSource, using your own TokenSource as the second argument. That second argument's Token method is called to refresh, and you can implement that however you want.

[smlx]

Yes, that is what I ended up doing for intercepting token refreshes.

However it was much more difficult to extract the ClientID and TokenURL, both of which can be dynamically generated during the authorization code flow. Do you have any ideas for that?

So far what I have is the aforementioned AuthorizationCodeHandlerConfig.Client that inspects its requests and extracts the TokenURL from any POST that has grant_type=authorization_code. Like I said, quite hacky.

[maciej-kisiel]

Hi @smlx,

I initially included some elements in the client-side OAuth proposal (#785) that aimed to support use cases similar to yours, but we removed them since we didn't know if there's demand for them and if the APIs I came up with are optimal.

These elements would be:

• AuthorizationCodeHandler.SetTokenSource – ability to pre-set a token source to be used by the handler without running the authorization flow. If I understand correctly, you would need something like that to use the tokens you load from disk. How are you doing that right now without it?
• AuthorizationCodeHandlerConfig.TokenStore – ability to provide an implementation of an interface to automatically store tokens. The definition of the TokenStore was along these lines:

   // TokenStore is an interface than can be used by OAuthHandler implementations
   // to save tokens to a persistent storage.
   type TokenStore interface {
   	Save(context.Context, *oauth2.Token) error
   }

  But I was also wondering if/how to expose the additional data that is needed to build an *oauth2.Config needed to create a refreshable *oauth2.TokenSource. Maybe we could consider passing the resolved *oauth2.Config instance to the token store as well?
• auth.NewPersistentTokenSource – a function to enrich a token source with capability to save the tokens to the token store. Useful when manually pre-setting the token source in the handler e.g. from disk to get the same behavior. The proposed signature was func (ctx context.Context, wrapped oauth2.TokenSource, store TokenStore) oauth2.TokenSource.

Let me know what do you think about these proposals and whether they would solve your use case.

[smlx]

Hi, thanks for the detailed response! I've only recently stumbled into the MCP space so I was unaware of your proposal in #785.

AuthorizationCodeHandler.SetTokenSource – ability to pre-set a token source to be used by the handler without running the authorization flow. If I understand correctly, you would need something like that to use the tokens you load from disk. How are you doing that right now without it?

Basically I force a token refresh on initialization and return early from Authorize() if it works. It is functional, but far from optimal. Code is here.

I think your interface is nicer and avoids the unnecessary forced refresh. 🙂

AuthorizationCodeHandlerConfig.TokenStore – ability to provide an implementation of an interface to automatically store tokens.

If this was called every time the token was refreshed (which I understand from the NewPersistentTokenSource description it would be) then yeah I think this works.

But I was also wondering if/how to expose the additional data that is needed to build an *oauth2.Config needed to create a refreshable *oauth2.TokenSource. Maybe we could consider passing the resolved *oauth2.Config instance to the token store as well?

My understanding is that oauth.Config shouldn't really change after initialization in Authorize(), so I am not sure it makes sense to pass it to Save(), which may be called on every call to TokenSource.Token().

My idea for this was to add a callback to the AuthorizationCodeHandlerConfig. I have implemented this in a fork here. I admit I am not as familiar with either OAuth2 or this library as you are so I do not know if this is a bad idea or not but it seems to work in my brief testing.

auth.NewPersistentTokenSource – a function to enrich a token source with capability to save the tokens to the token store. Useful when manually pre-setting the token source in the handler e.g. from disk to get the same behavior. The proposed signature was func (ctx context.Context, wrapped oauth2.TokenSource, store TokenStore) oauth2.TokenSource.

Looking at #785 I see that this would call TokenStore.Save on every call to Token() which I think makes sense. My TokenStore would just have to compare the token in memory to the new one to decide whether or not to update the token on disk.

I see that you have thought about this a great deal more than I have and I hope that you are able to add some of the persistence ideas from #785 back to the API.

[smlx]

My first comment was rather rambling, so here is a summary of my persistence use case:

1. Load an oauth2.Config and oauth2.Token from disk, inject them into some interface in this SDK and, if the access token and/or refresh token are valid, have this state used without forcing the user through an interactive authorization flow.
2. Save the oauth2.Config generated in AuthorizationCodeHandler.Authorize() to disk after it has been successfully used to obtain a token.
3. Save the oauth2.Token to disk whenever it is refreshed.

[smlx]

My idea for this was to add a callback to the AuthorizationCodeHandlerConfig. I have implemented this in a fork here.

If you do happen to like this idea, let me know and I can open a PR. 🙂

[maciej-kisiel]

Let me advertise this topic on the Discord channel, maybe we can get more opinions.