Merge pull request #2448 from pandego/fix/2417-mcp-oauth-refresh

fix(mcp): reuse discovered auth server for token refresh

Author: dgageot

pkg/tools/mcp/oauth.go

@@ -259,9 +259,10 @@ func (t *oauthTransport) getValidToken(ctx context.Context) *OAuthToken {
 	slog.Debug("Attempting silent token refresh", "url", t.baseURL)
 
 	o := &oauth{metadataClient: &http.Client{Timeout: 5 * time.Second}}
-	metadata, err := o.getAuthorizationServerMetadata(ctx, t.baseURL)
+	authServer := cmp.Or(token.AuthServer, t.baseURL)
+	metadata, err := o.getAuthorizationServerMetadata(ctx, authServer)
 	if err != nil {
-		slog.Debug("Failed to fetch auth server metadata for refresh", "error", err)
+		slog.Debug("Failed to fetch auth server metadata for refresh", "auth_server", authServer, "error", err)
 		return nil
 	}
 
@@ -273,6 +274,7 @@ func (t *oauthTransport) getValidToken(ctx context.Context) *OAuthToken {
 		t.mu.Unlock()
 		return nil
 	}
+	newToken.AuthServer = authServer
 
 	t.mu.Lock()
 	t.refreshFailedAt = time.Time{} // reset on success
@@ -443,6 +445,7 @@ func (t *oauthTransport) handleManagedOAuthFlow(ctx context.Context, authServer,
 
 	token.ClientID = clientID
 	token.ClientSecret = clientSecret
+	token.AuthServer = resourceMetadata.AuthorizationServers[0]
 
 	if err := t.tokenStore.StoreToken(t.baseURL, token); err != nil {
 		return fmt.Errorf("failed to store token: %w", err)
@@ -539,6 +542,7 @@ func (t *oauthTransport) handleUnmanagedOAuthFlow(ctx context.Context, authServe
 		token.ExpiresIn = int(expiresIn)
 		token.ExpiresAt = time.Now().Add(time.Duration(token.ExpiresIn) * time.Second)
 	}
+	token.AuthServer = resourceMetadata.AuthorizationServers[0]
 
 	if refreshToken, ok := tokenData["refresh_token"].(string); ok {
 		token.RefreshToken = refreshToken

pkg/tools/mcp/oauth_test.go

@@ -172,6 +172,80 @@ func TestGetValidToken_UsesStoredCredentialsForRefresh(t *testing.T) {
 	}
 }
 
+// TestGetValidToken_UsesStoredAuthServerForRefresh verifies that silent
+// refresh uses the discovered auth server rather than assuming the MCP
+// server URL also hosts the OAuth metadata.
+func TestGetValidToken_UsesStoredAuthServerForRefresh(t *testing.T) {
+	var refreshRequests int
+
+	authMux := http.NewServeMux()
+	authSrv := httptest.NewServer(authMux)
+	defer authSrv.Close()
+
+	authMux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"issuer":                 authSrv.URL,
+			"token_endpoint":         authSrv.URL + "/token",
+			"authorization_endpoint": authSrv.URL + "/authorize",
+		})
+	})
+	authMux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
+		refreshRequests++
+		if err := r.ParseForm(); err != nil {
+			t.Fatal(err)
+		}
+		if got := r.FormValue("refresh_token"); got != "old-rt" {
+			t.Fatalf("refresh_token = %q, want %q", got, "old-rt")
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"access_token":  "fresh-at",
+			"token_type":    "Bearer",
+			"expires_in":    3600,
+			"refresh_token": "fresh-rt",
+		})
+	})
+
+	mcpSrv := httptest.NewServer(http.NotFoundHandler())
+	defer mcpSrv.Close()
+
+	store := NewInMemoryTokenStore()
+	expiredToken := &OAuthToken{
+		AccessToken:  "old-at",
+		TokenType:    "Bearer",
+		RefreshToken: "old-rt",
+		ExpiresAt:    time.Now().Add(-1 * time.Hour),
+		ClientID:     "stored-cid",
+		ClientSecret: "stored-csec",
+		AuthServer:   authSrv.URL,
+	}
+	if err := store.StoreToken(mcpSrv.URL, expiredToken); err != nil {
+		t.Fatal(err)
+	}
+
+	transport := &oauthTransport{
+		base:       http.DefaultTransport,
+		tokenStore: store,
+		baseURL:    mcpSrv.URL,
+	}
+
+	got := transport.getValidToken(t.Context())
+	if got == nil {
+		t.Fatal("getValidToken returned nil, expected refreshed token")
+	}
+	if got.AccessToken != "fresh-at" {
+		t.Fatalf("AccessToken = %q, want %q", got.AccessToken, "fresh-at")
+	}
+	if got.AuthServer != authSrv.URL {
+		t.Fatalf("AuthServer = %q, want %q", got.AuthServer, authSrv.URL)
+	}
+	if refreshRequests != 1 {
+		t.Fatalf("refreshRequests = %d, want 1", refreshRequests)
+	}
+}
+
 // TestOAuthTokenClientCredentials_JSONRoundTrip verifies that ClientID and
 // ClientSecret survive JSON serialization (important for keyring storage).
 func TestOAuthTokenClientCredentials_JSONRoundTrip(t *testing.T) {
@@ -183,6 +257,7 @@ func TestOAuthTokenClientCredentials_JSONRoundTrip(t *testing.T) {
 		ExpiresAt:    time.Date(2025, 6, 1, 0, 0, 0, 0, time.UTC),
 		ClientID:     "cid",
 		ClientSecret: "csec",
+		AuthServer:   "https://auth.example.com",
 	}
 
 	data, err := json.Marshal(token)
@@ -201,6 +276,9 @@ func TestOAuthTokenClientCredentials_JSONRoundTrip(t *testing.T) {
 	if got.ClientSecret != "csec" {
 		t.Errorf("ClientSecret = %q, want %q", got.ClientSecret, "csec")
 	}
+	if got.AuthServer != "https://auth.example.com" {
+		t.Errorf("AuthServer = %q, want %q", got.AuthServer, "https://auth.example.com")
+	}
 }
 
 // TestOAuthTokenClientCredentials_OmittedWhenEmpty verifies the omitempty

pkg/tools/mcp/tokenstore.go

@@ -26,6 +26,7 @@ type OAuthToken struct {
 	ExpiresAt    time.Time `json:"expires_at"`
 	ClientID     string    `json:"client_id,omitempty"`
 	ClientSecret string    `json:"client_secret,omitempty"`
+	AuthServer   string    `json:"auth_server,omitempty"`
 }
 
 // IsExpired checks if the token is expired