Add OAuth flow chaining and token injection for stdio sessions

After Google OAuth, automatically redirect to server OAuth if the return URL specifies a server that requires it. Also inject user tokens into stdio sessions during creation so they can authenticate with backend services.

Author: dgellow

integration/oauth_test.go

@@ -157,7 +157,7 @@ func TestClientRegistration(t *testing.T) {
 	})
 	defer stopServer(mcpCmd)
 
-	if !waitForHealthCheck(t, 30) {
+	if !waitForHealthCheck(30) {
 		t.Fatal("OAuth server failed to start")
 	}
 
@@ -358,7 +358,7 @@ func TestUserTokenFlow(t *testing.T) {
 	mcpCmd := startOAuthServerWithTokenConfig(t)
 	defer stopServer(mcpCmd)
 
-	if !waitForHealthCheck(t, 30) {
+	if !waitForHealthCheck(30) {
 		t.Fatal("Server failed to start")
 	}
 
@@ -532,7 +532,7 @@ func TestStateParameterHandling(t *testing.T) {
 			})
 			defer stopServer(mcpCmd)
 
-			if !waitForHealthCheck(t, 10) {
+			if !waitForHealthCheck(10) {
 				t.Fatal("Server failed to start")
 			}
 
@@ -602,7 +602,7 @@ func TestEnvironmentModes(t *testing.T) {
 		})
 		defer stopServer(mcpCmd)
 
-		if !waitForHealthCheck(t, 30) {
+		if !waitForHealthCheck(30) {
 			t.Fatal("Server failed to start")
 		}
 
@@ -643,7 +643,7 @@ func TestEnvironmentModes(t *testing.T) {
 		})
 		defer stopServer(mcpCmd)
 
-		if !waitForHealthCheck(t, 30) {
+		if !waitForHealthCheck(30) {
 			t.Fatal("Server failed to start")
 		}
 
@@ -693,7 +693,7 @@ func TestOAuthEndpoints(t *testing.T) {
 	})
 	defer stopServer(mcpCmd)
 
-	if !waitForHealthCheck(t, 10) {
+	if !waitForHealthCheck(10) {
 		t.Fatal("Server failed to start")
 	}
 
@@ -761,7 +761,7 @@ func TestCORSHeaders(t *testing.T) {
 	})
 	defer stopServer(mcpCmd)
 
-	if !waitForHealthCheck(t, 10) {
+	if !waitForHealthCheck(10) {
 		t.Fatal("Server failed to start")
 	}
 
@@ -813,7 +813,7 @@ func TestToolAdvertisementWithUserTokens(t *testing.T) {
 		"LOG_LEVEL=debug",
 	)
 
-	if !waitForHealthCheck(t, 30) {
+	if !waitForHealthCheck(30) {
 		t.Fatal("Server failed to start")
 	}
 
@@ -980,13 +980,14 @@ func TestToolAdvertisementWithUserTokens(t *testing.T) {
 		defer resp.Body.Close()
 
 		// Check the response - it might be 200 if following redirects
-		if resp.StatusCode == 200 {
+		switch resp.StatusCode {
+		case 200:
 			// That's fine, it means the token was set and we got the page back
 			t.Log("Token set successfully, got page response")
-		} else if resp.StatusCode == 302 || resp.StatusCode == 303 {
+		case 302, 303:
 			// Also fine, redirect means success
 			t.Log("Token set successfully, got redirect")
-		} else {
+		default:
 			body, _ := io.ReadAll(resp.Body)
 			t.Fatalf("Unexpected response setting token: status=%d, body=%s", resp.StatusCode, string(body))
 		}
@@ -1119,7 +1120,7 @@ func stopServer(cmd *exec.Cmd) {
 	}
 }
 
-func waitForHealthCheck(t *testing.T, seconds int) bool {
+func waitForHealthCheck(seconds int) bool {
 	for i := 0; i < seconds; i++ {
 		if checkHealth() {
 			return true

integration/security_test.go

@@ -221,10 +221,11 @@ func TestSecurityScenarios(t *testing.T) {
 			}
 			resp.Body.Close()
 
-			if resp.StatusCode == 200 {
+			switch resp.StatusCode {
+			case 200:
 				t.Errorf("CRITICAL: Auth bypass! 'test-token' without Bearer returned 200")
-			} else if resp.StatusCode == 401 {
-			} else {
+			case 401:
+			default:
 				t.Logf("Unexpected status %d for malformed auth", resp.StatusCode)
 			}
 		})

internal/client/client.go

@@ -435,7 +435,6 @@ func (c *Client) wrapToolHandler(
 				errorData := createTokenRequiredError(
 					serverName,
 					setupBaseURL,
-					userAuth,
 					"configuration error: this service requires user tokens but OAuth is not properly configured.",
 				)
 
@@ -469,7 +468,6 @@ func (c *Client) wrapToolHandler(
 				errorData := createTokenRequiredError(
 					serverName,
 					setupBaseURL,
-					userAuth,
 					errorMessage,
 				)
 
@@ -509,7 +507,7 @@ func (c *Client) Close() error {
 }
 
 // createTokenRequiredError creates the structured error for missing user tokens
-func createTokenRequiredError(serverName, setupBaseURL string, userAuth *config.UserAuthentication, message string) map[string]interface{} {
+func createTokenRequiredError(serverName, setupBaseURL string, message string) map[string]interface{} {
 	tokenSetupURL := fmt.Sprintf("%s/my/tokens", setupBaseURL)
 
 	return map[string]interface{}{

internal/client/session_manager.go

@@ -115,6 +115,7 @@ func (sm *StdioSessionManager) GetOrCreateSession(
 	config *config.MCPClientConfig,
 	info mcp.Implementation,
 	baseURL string,
+	userToken string,
 ) (*StdioSession, error) {
 	// Try to get existing session first
 	if session, ok := sm.GetSession(key); ok {
@@ -125,7 +126,7 @@ func (sm *StdioSessionManager) GetOrCreateSession(
 		return nil, err
 	}
 
-	return sm.createSession(ctx, key, config, info, baseURL)
+	return sm.createSession(key, config, userToken)
 }
 
 // GetSession retrieves an existing session
@@ -379,14 +380,21 @@ func (sm *StdioSessionManager) getUserSessionCount(userEmail string) int {
 
 // createSession creates a new stdio session
 func (sm *StdioSessionManager) createSession(
-	ctx context.Context,
 	key SessionKey,
 	config *config.MCPClientConfig,
-	info mcp.Implementation,
-	baseURL string,
+	userToken string,
 ) (*StdioSession, error) {
+	// Create an independent context for the stdio session. We intentionally use
+	// context.Background() instead of the HTTP request context because stdio
+	// sessions are long-lived processes that must persist across multiple HTTP
+	// requests. The session will be cleaned up by the timeout-based cleanup
+	// routine, not by HTTP request cancellation.
 	sessionCtx, cancel := context.WithCancel(context.Background())
 
+	if userToken != "" && config.RequiresUserToken {
+		config = config.ApplyUserToken(userToken)
+	}
+
 	client, err := sm.createClient(key.ServerName, config)
 	if err != nil {
 		cancel()

internal/client/session_manager_test.go

@@ -44,12 +44,12 @@ func TestStdioSessionManager_CreateAndRetrieve(t *testing.T) {
 	}
 
 	// Create session
-	session1, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost")
+	session1, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost", "")
 	require.NoError(t, err)
 	require.NotNil(t, session1)
 
 	// Retrieve same session
-	session2, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost")
+	session2, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost", "")
 	require.NoError(t, err)
 	require.NotNil(t, session2)
 
@@ -81,22 +81,22 @@ func TestStdioSessionManager_UserLimits(t *testing.T) {
 
 	// Create first session
 	key1 := SessionKey{UserEmail: userEmail, ServerName: "server", SessionID: "1"}
-	_, err := sm.GetOrCreateSession(context.Background(), key1, config, info, "http://localhost")
+	_, err := sm.GetOrCreateSession(context.Background(), key1, config, info, "http://localhost", "")
 	require.NoError(t, err)
 
 	// Create second session (at limit)
 	key2 := SessionKey{UserEmail: userEmail, ServerName: "server", SessionID: "2"}
-	_, err = sm.GetOrCreateSession(context.Background(), key2, config, info, "http://localhost")
+	_, err = sm.GetOrCreateSession(context.Background(), key2, config, info, "http://localhost", "")
 	require.NoError(t, err)
 
 	// Try to create third session (should fail)
 	key3 := SessionKey{UserEmail: userEmail, ServerName: "server", SessionID: "3"}
-	_, err = sm.GetOrCreateSession(context.Background(), key3, config, info, "http://localhost")
+	_, err = sm.GetOrCreateSession(context.Background(), key3, config, info, "http://localhost", "")
 	assert.ErrorIs(t, err, ErrUserLimitExceeded)
 
 	// Different user should work
 	key4 := SessionKey{UserEmail: "[email protected]", ServerName: "server", SessionID: "4"}
-	_, err = sm.GetOrCreateSession(context.Background(), key4, config, info, "http://localhost")
+	_, err = sm.GetOrCreateSession(context.Background(), key4, config, info, "http://localhost", "")
 	require.NoError(t, err)
 }
 
@@ -122,7 +122,7 @@ func TestStdioSessionManager_RemoveSession(t *testing.T) {
 	info := mcp.Implementation{Name: "test", Version: "1.0"}
 
 	// Create session
-	session, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost")
+	session, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost", "")
 	require.NoError(t, err)
 	require.NotNil(t, session)
 
@@ -161,7 +161,7 @@ func TestStdioSessionManager_Timeout(t *testing.T) {
 	info := mcp.Implementation{Name: "test", Version: "1.0"}
 
 	// Create session
-	session, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost")
+	session, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost", "")
 	require.NoError(t, err)
 	require.NotNil(t, session)
 
@@ -204,7 +204,7 @@ func TestStdioSessionManager_ConcurrentAccess(t *testing.T) {
 			}
 
 			// Create session
-			_, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost")
+			_, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost", "")
 			assert.NoError(t, err)
 
 			// Get session
@@ -243,7 +243,7 @@ func TestStdioSessionManager_NoLimitsForAnonymous(t *testing.T) {
 			SessionID:  fmt.Sprintf("session-%d", i),
 		}
 
-		_, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost")
+		_, err := sm.GetOrCreateSession(context.Background(), key, config, info, "http://localhost", "")
 		require.NoError(t, err, "Anonymous session %d should succeed", i)
 	}
 }

internal/server/auth_handlers.go

@@ -5,10 +5,12 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 
 	"github.com/dgellow/mcp-front/internal/auth"
+	"github.com/dgellow/mcp-front/internal/config"
 	"github.com/dgellow/mcp-front/internal/crypto"
 	jsonwriter "github.com/dgellow/mcp-front/internal/json"
 	"github.com/dgellow/mcp-front/internal/log"
@@ -18,12 +20,14 @@ import (
 // AuthHandlers wraps the auth.Server to provide HTTP handlers
 type AuthHandlers struct {
 	authServer *auth.Server
+	mcpServers map[string]*config.MCPClientConfig
 }
 
 // NewAuthHandlers creates new auth handlers
-func NewAuthHandlers(authServer *auth.Server) *AuthHandlers {
+func NewAuthHandlers(authServer *auth.Server, mcpServers map[string]*config.MCPClientConfig) *AuthHandlers {
 	return &AuthHandlers{
 		authServer: authServer,
+		mcpServers: mcpServers,
 	}
 }
 
@@ -222,7 +226,30 @@ func (h *AuthHandlers) GoogleCallbackHandler(w http.ResponseWriter, r *http.Requ
 			"returnURL": returnURL,
 		})
 
-		// Redirect to return URL
+		// Check if the return URL contains a server parameter for OAuth chaining
+		parsedURL, err := url.Parse(returnURL)
+		if err == nil {
+			serverName := parsedURL.Query().Get("server")
+			if serverName != "" {
+				// Check if this server requires OAuth authentication
+				if serverConfig, exists := h.mcpServers[serverName]; exists {
+					if serverConfig.RequiresUserToken &&
+						serverConfig.UserAuthentication != nil &&
+						serverConfig.UserAuthentication.Type == config.UserAuthTypeOAuth {
+						encodedReturnURL := url.QueryEscape(returnURL)
+						oauthURL := fmt.Sprintf("/oauth/connect?service=%s&return=%s", serverName, encodedReturnURL)
+						log.LogInfoWithFields("auth", "Chaining to server OAuth", map[string]interface{}{
+							"server": serverName,
+							"user":   userInfo.Email,
+						})
+						http.Redirect(w, r, oauthURL, http.StatusFound)
+						return
+					}
+				}
+			}
+		}
+
+		// Otherwise, redirect to return URL as normal
 		http.Redirect(w, r, returnURL, http.StatusFound)
 		return
 	}

internal/server/handler.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 
 	"github.com/dgellow/mcp-front/internal/auth"
@@ -174,7 +175,7 @@ func NewServer(ctx context.Context, cfg *config.Config) (*Server, error) {
 			recoverMiddleware("mcp"),
 		}
 
-		authHandlers := NewAuthHandlers(s.authServer)
+		authHandlers := NewAuthHandlers(s.authServer, cfg.MCPServers)
 		mux.Handle("/.well-known/oauth-authorization-server", chainMiddleware(http.HandlerFunc(authHandlers.WellKnownHandler), oauthMiddlewares...))
 		mux.Handle("/authorize", chainMiddleware(http.HandlerFunc(authHandlers.AuthorizeHandler), oauthMiddlewares...))
 		mux.Handle("/oauth/callback", chainMiddleware(http.HandlerFunc(authHandlers.GoogleCallbackHandler), oauthMiddlewares...))
@@ -415,6 +416,27 @@ func isStdioServer(config *config.MCPClientConfig) bool {
 	return config.Command != ""
 }
 
+// formatUserToken formats a stored token according to the user authentication configuration
+func formatUserToken(storedToken *storage.StoredToken, auth *config.UserAuthentication) string {
+	if storedToken == nil {
+		return ""
+	}
+
+	if storedToken.Type == storage.TokenTypeOAuth && storedToken.OAuthData != nil {
+		token := storedToken.OAuthData.AccessToken
+		if auth.TokenFormat != "" && auth.TokenFormat != "{{token}}" {
+			return strings.ReplaceAll(auth.TokenFormat, "{{token}}", token)
+		}
+		return token
+	}
+
+	token := storedToken.Value
+	if auth != nil && auth.TokenFormat != "" && auth.TokenFormat != "{{token}}" {
+		return strings.ReplaceAll(auth.TokenFormat, "{{token}}", token)
+	}
+	return token
+}
+
 // sessionHandlerKey is the context key for session handlers
 type sessionHandlerKey struct{}
 
@@ -455,12 +477,30 @@ func handleSessionRegistration(
 		"command":           handler.config.Command,
 	})
 
+	var userToken string
+	if handler.config.RequiresUserToken && handler.userEmail != "" && handler.h.storage != nil {
+		storedToken, err := handler.h.storage.GetUserToken(sessionCtx, handler.userEmail, handler.h.serverName)
+		if err != nil {
+			log.LogDebugWithFields("server", "No user token found", map[string]interface{}{
+				"server": handler.h.serverName,
+				"user":   handler.userEmail,
+			})
+		} else if storedToken != nil {
+			if handler.config.UserAuthentication != nil {
+				userToken = formatUserToken(storedToken, handler.config.UserAuthentication)
+			} else {
+				userToken = storedToken.Value
+			}
+		}
+	}
+
 	stdioSession, err := sessionManager.GetOrCreateSession(
 		sessionCtx,
 		key,
 		handler.config,
 		handler.h.info,
 		handler.h.setupBaseURL,
+		userToken,
 	)
 	if err != nil {
 		log.LogErrorWithFields("server", "Failed to create stdio session", map[string]interface{}{