Fix stateless mode OAuth infinite authentication loop

- Implement email-based credential lookup in stateless mode
- Allow session-independent OAuth state validation
- Handle OAuth callbacks when sessions are terminated

Essential files only for WORKSPACE_MCP_STATELESS_MODE=true

Author: lucyliu542789

auth/google_auth.py

@@ -486,9 +486,19 @@ def handle_auth_callback(
         user_google_email = user_info["email"]
         logger.info(f"Identified user_google_email: {user_google_email}")
 
-        # Save the credentials
-        credential_store = get_credential_store()
-        credential_store.store_credential(user_google_email, credentials)
+        # In stateless mode, if session_id is None or invalid, create a new one
+        # This handles the case where the client terminated the session before OAuth completed
+        if is_stateless_mode() and not session_id:
+            import uuid
+            session_id = str(uuid.uuid4())
+            logger.info(f"Stateless mode: Created new session {session_id} for {user_google_email} after OAuth")
+
+        # Save the credentials to file (skip in stateless mode)
+        if not is_stateless_mode():
+            credential_store = get_credential_store()
+            credential_store.store_credential(user_google_email, credentials)
+        else:
+            logger.info(f"Stateless mode: Skipping file storage for {user_google_email}")
 
         # Always save to OAuth21SessionStore for centralized management
         store = get_oauth21_session_store()

auth/oauth21_session_store.py

@@ -262,21 +262,37 @@ def validate_and_consume_oauth_state(
                 raise ValueError("Invalid or expired OAuth state parameter")
 
             bound_session = state_info.get("session_id")
-            if bound_session and session_id and bound_session != session_id:
-                # Consume the state to prevent replay attempts
-                del self._oauth_states[state]
-                logger.error(
-                    "SECURITY: OAuth state session mismatch (expected %s, got %s)",
-                    bound_session,
-                    session_id,
-                )
-                raise ValueError("OAuth state does not match the initiating session")
+            
+            # In stateless mode, don't require session to match
+            # The state itself is cryptographically secure (32 bytes random)
+            # and is one-time use, which provides sufficient security
+            from auth.oauth_config import is_stateless_mode
+            if not is_stateless_mode():
+                # In non-stateless mode, enforce strict session binding
+                if bound_session and session_id and bound_session != session_id:
+                    # Consume the state to prevent replay attempts
+                    del self._oauth_states[state]
+                    logger.error(
+                        "SECURITY: OAuth state session mismatch (expected %s, got %s)",
+                        bound_session,
+                        session_id,
+                    )
+                    raise ValueError("OAuth state does not match the initiating session")
+            else:
+                # In stateless mode, log but don't fail on session mismatch
+                if bound_session and session_id and bound_session != session_id:
+                    logger.debug(
+                        "Stateless mode: OAuth state session mismatch (expected %s, got %s) - allowing",
+                        bound_session,
+                        session_id,
+                    )
 
             # State is valid – consume it to prevent reuse
             del self._oauth_states[state]
             logger.debug(
-                "Validated OAuth state %s",
+                "Validated OAuth state %s (stateless mode: %s)",
                 state[:8] if len(state) > 8 else state,
+                is_stateless_mode()
             )
             return state_info
 

auth/service_decorator.py

@@ -210,6 +210,8 @@ async def get_authenticated_google_service_oauth21(
     """
     OAuth 2.1 authentication using the session store with security validation.
     """
+    from auth.oauth_config import is_stateless_mode
+    
     provider = get_auth_provider()
     access_token = get_access_token()
 
@@ -254,6 +256,36 @@ async def get_authenticated_google_service_oauth21(
         return service, resolved_email
 
     store = get_oauth21_session_store()
+    
+    # In stateless mode with bearer token, use email-based lookup
+    # This allows credentials to work across different session IDs
+    if is_stateless_mode() and access_token:
+        token_email = None
+        if getattr(access_token, "claims", None):
+            token_email = access_token.claims.get("email")
+        
+        if token_email:
+            # Security: Validate token email matches requested email
+            if user_google_email and token_email != user_google_email:
+                raise GoogleAuthenticationError(
+                    f"Token email {token_email} does not match requested user {user_google_email}."
+                )
+            
+            # Direct email-based lookup - works across sessions!
+            logger.debug(f"[{tool_name}] Stateless mode: Using email-based credential lookup for {token_email}")
+            credentials = store.get_credentials(token_email)
+            
+            if credentials:
+                # Validate scopes
+                scopes_available = set(credentials.scopes or [])
+                if not all(scope in scopes_available for scope in required_scopes):
+                    raise GoogleAuthenticationError(
+                        f"OAuth 2.1 credentials lack required scopes. Need: {required_scopes}, Have: {sorted(scopes_available)}"
+                    )
+                
+                service = build(service_name, version, credentials=credentials)
+                logger.info(f"[{tool_name}] Stateless mode: Authenticated {service_name} for {token_email}")
+                return service, token_email
 
     # Use the validation method to ensure session can only access its own credentials
     credentials = store.get_credentials_with_validation(