fix(cli): complete OAuth auth flow for CLI mode

auth/google_auth.py

@@ -329,7 +329,7 @@ async def start_auth_flow(
     user_google_email: Optional[str],
     service_name: str,  # e.g., "Google Calendar", "Gmail" for user messages
     redirect_uri: str,  # Added redirect_uri as a required parameter
-) -> str:
+) -> tuple[str, str]:
     """
     Initiates the Google OAuth flow and returns an actionable message for the user.
 
@@ -339,7 +339,7 @@ async def start_auth_flow(
         redirect_uri: The URI Google will redirect to after authorization.
 
     Returns:
-        A formatted string containing guidance for the LLM/user.
+        Tuple of (formatted message string, auth_url).
 
     Raises:
         Exception: If the OAuth flow cannot be initiated.
@@ -378,7 +378,7 @@ async def start_auth_flow(
             state=oauth_state,
         )
 
-        auth_url, _ = flow.authorization_url(access_type="offline", prompt="consent")
+        auth_url, _ = flow.authorization_url(prompt="consent")
 
         session_id = None
         try:
@@ -422,7 +422,7 @@ async def start_auth_flow(
         message_lines.append(
             f"\nThe application will use the new credentials. If '{user_google_email}' was provided, it must match the authenticated account."
         )
-        return "\n".join(message_lines)
+        return "\n".join(message_lines), auth_url
 
     except FileNotFoundError as e:
         error_text = f"OAuth client credentials not found: {e}. Please either:\n1. Set environment variables: GOOGLE_OAUTH_CLIENT_ID and GOOGLE_OAUTH_CLIENT_SECRET\n2. Ensure '{CONFIG_CLIENT_SECRETS_PATH}' file exists"
@@ -939,14 +939,14 @@ async def get_authenticated_google_service(
             )
 
         # Generate auth URL and raise exception with it
-        auth_response = await start_auth_flow(
+        auth_response, auth_url = await start_auth_flow(
             user_google_email=user_google_email,
             service_name=f"Google {service_name.title()}",
             redirect_uri=redirect_uri,
         )
 
-        # Extract the auth URL from the response and raise with it
-        raise GoogleAuthenticationError(auth_response)
+        # Raise with both the message and the auth URL
+        raise GoogleAuthenticationError(auth_response, auth_url=auth_url)
 
     try:
         service = build(service_name, version, credentials=credentials)

auth/oauth_callback_server.py

@@ -43,6 +43,10 @@ def __init__(self, port: int = 8000, base_uri: str = "http://localhost"):
         self.server_thread = None
         self.is_running = False
 
+        # CLI auth completion signaling
+        self.auth_completed = threading.Event()
+        self.auth_result: Optional[dict] = None  # {"success": bool, "user_id": str|None, "error": str|None}
+
         # Setup the callback route
         self._setup_callback_route()
         # Setup attachment serving route
@@ -62,19 +66,25 @@ async def oauth_callback(request: Request):
                     f"Authentication failed: Google returned an error: {error}."
                 )
                 logger.error(error_message)
+                self.auth_result = {"success": False, "user_id": None, "error": error_message}
+                self.auth_completed.set()
                 return create_error_response(error_message)
 
             if not code:
                 error_message = (
                     "Authentication failed: No authorization code received from Google."
                 )
                 logger.error(error_message)
+                self.auth_result = {"success": False, "user_id": None, "error": error_message}
+                self.auth_completed.set()
                 return create_error_response(error_message)
 
             try:
                 # Check if we have credentials available (environment variables or file)
                 error_message = check_client_secrets()
                 if error_message:
+                    self.auth_result = {"success": False, "user_id": None, "error": error_message}
+                    self.auth_completed.set()
                     return create_server_error_response(error_message)
 
                 logger.info(
@@ -96,12 +106,18 @@ async def oauth_callback(request: Request):
                     f"OAuth callback: Successfully authenticated user: {verified_user_id}."
                 )
 
+                # Signal completion for CLI auth flow
+                self.auth_result = {"success": True, "user_id": verified_user_id, "error": None}
+                self.auth_completed.set()
+
                 # Return success page using shared template
                 return create_success_response(verified_user_id)
 
             except Exception as e:
                 error_message_detail = f"Error processing OAuth callback: {str(e)}"
                 logger.error(error_message_detail, exc_info=True)
+                self.auth_result = {"success": False, "user_id": None, "error": str(e)}
+                self.auth_completed.set()
                 return create_server_error_response(str(e))
 
     def _setup_attachment_route(self):
@@ -200,6 +216,22 @@ def run_server():
         logger.error(error_msg)
         return False, error_msg
 
+    def wait_for_auth(self, timeout: float = 300) -> Optional[dict]:
+        """
+        Block until OAuth callback is received or timeout.
+
+        Args:
+            timeout: Maximum seconds to wait (default 5 minutes)
+
+        Returns:
+            Auth result dict {"success": bool, "user_id": str|None, "error": str|None}
+            or None if timed out
+        """
+        completed = self.auth_completed.wait(timeout=timeout)
+        if completed:
+            return self.auth_result
+        return None
+
     def stop(self):
         """Stop the minimal OAuth server."""
         if not self.is_running:
@@ -279,6 +311,34 @@ def ensure_oauth_callback_available(
         return False, error_msg
 
 
+def set_cli_oauth_server(server: MinimalOAuthServer) -> None:
+    """Register a MinimalOAuthServer as the global instance (used by CLI mode)."""
+    global _minimal_oauth_server
+    _minimal_oauth_server = server
+
+
+def get_cli_oauth_port() -> int:
+    """
+    Find an available port for the CLI OAuth callback server.
+    Tries ports 8000-8009, falls back to OS-assigned port.
+
+    Returns:
+        Available port number
+    """
+    for port in range(8000, 8010):
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.bind(("localhost", port))
+                return port
+        except OSError:
+            continue
+
+    # Fallback: let OS assign a port
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("localhost", 0))
+        return s.getsockname()[1]
+
+
 def cleanup_oauth_callback_server():
     """Clean up the minimal OAuth server if it was started."""
     global _minimal_oauth_server

core/cli_handler.py

@@ -16,10 +16,16 @@
 import asyncio
 import json
 import logging
+import os
 import sys
 from typing import Any, Dict, List, Optional
 
-from auth.oauth_config import set_transport_mode
+from auth.oauth_config import set_transport_mode, reload_oauth_config
+from auth.oauth_callback_server import (
+    MinimalOAuthServer,
+    get_cli_oauth_port,
+    set_cli_oauth_server,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -264,6 +270,10 @@ async def run_tool(server, tool_name: str, args: Dict[str, Any]) -> str:
             f"Provided parameters: {list(call_args.keys())}"
         )
     except Exception as e:
+        # Let GoogleAuthenticationError propagate for CLI auth handling
+        from auth.google_auth import GoogleAuthenticationError
+        if isinstance(e, GoogleAuthenticationError):
+            raise
         logger.error(f"[CLI] Error executing {tool_name}: {e}", exc_info=True)
         return f"Error: {type(e).__name__}: {e}"
 
@@ -361,6 +371,74 @@ def read_stdin_args() -> Dict[str, Any]:
     return {}
 
 
+def _setup_cli_oauth_server() -> MinimalOAuthServer:
+    """
+    Pre-start a MinimalOAuthServer for CLI auth flow.
+
+    Finds an available port, configures the redirect URI, starts the server,
+    and registers it as the global OAuth server.
+
+    Returns:
+        The started MinimalOAuthServer instance
+    """
+    port = get_cli_oauth_port()
+    redirect_uri = f"http://localhost:{port}/oauth2callback"
+
+    # Set redirect URI in env so OAuthConfig picks it up
+    os.environ["GOOGLE_OAUTH_REDIRECT_URI"] = redirect_uri
+    reload_oauth_config()
+
+    oauth_server = MinimalOAuthServer(port=port, base_uri="http://localhost")
+    success, error_msg = oauth_server.start()
+    if not success:
+        raise RuntimeError(f"Failed to start OAuth server: {error_msg}")
+
+    set_cli_oauth_server(oauth_server)
+    logger.info(f"[CLI] OAuth callback server ready on port {port}")
+    return oauth_server
+
+
+async def _handle_cli_auth_flow(
+    server, tool_name: str, args: Dict[str, Any], auth_error, oauth_server: MinimalOAuthServer
+) -> str:
+    """
+    Handle OAuth authentication flow in CLI mode.
+
+    Prints the auth URL, waits for the browser callback, then retries the tool.
+
+    Args:
+        server: The FastMCP server instance
+        tool_name: Name of the tool to retry after auth
+        args: Tool arguments
+        auth_error: The GoogleAuthenticationError with auth_url
+        oauth_server: The running MinimalOAuthServer
+
+    Returns:
+        Tool result as string
+    """
+    print("\n" + "=" * 60, file=sys.stderr)
+    print("Authentication Required", file=sys.stderr)
+    print("=" * 60, file=sys.stderr)
+    print(f"\nOpen this URL in your browser:\n", file=sys.stderr)
+    print(f"  {auth_error.auth_url}\n", file=sys.stderr)
+    print("Waiting for authentication (timeout: 5 minutes)...", file=sys.stderr)
+
+    # Wait for the OAuth callback in a thread (blocking call)
+    result = await asyncio.to_thread(oauth_server.wait_for_auth, 300)
+
+    if result is None:
+        raise RuntimeError("Authentication timed out after 5 minutes.")
+
+    if not result["success"]:
+        raise RuntimeError(f"Authentication failed: {result['error']}")
+
+    print(f"\nAuthenticated as: {result['user_id']}", file=sys.stderr)
+    print("Retrying tool execution...\n", file=sys.stderr)
+
+    # Retry the tool
+    return await run_tool(server, tool_name, args)
+
+
 async def handle_cli_mode(server, cli_args: List[str]) -> int:
     """
     Main entry point for CLI mode.
@@ -390,13 +468,26 @@ async def handle_cli_mode(server, cli_args: List[str]) -> int:
             return 0
 
         if parsed["command"] == "run":
+            # Pre-start OAuth server before running tools that may need auth
+            oauth_server = _setup_cli_oauth_server()
+
             # Merge stdin args with inline args (inline takes precedence)
             args = read_stdin_args()
             args.update(parsed["tool_args"])
 
-            result = await run_tool(server, parsed["tool_name"], args)
-            print(result)
-            return 0
+            try:
+                result = await run_tool(server, parsed["tool_name"], args)
+                print(result)
+                return 0
+            except Exception as e:
+                from auth.google_auth import GoogleAuthenticationError
+                if isinstance(e, GoogleAuthenticationError) and e.auth_url:
+                    result = await _handle_cli_auth_flow(
+                        server, parsed["tool_name"], args, e, oauth_server
+                    )
+                    print(result)
+                    return 0
+                raise
 
         # Unknown command
         print(f"Unknown command: {parsed['command']}")

core/server.py

@@ -565,7 +565,7 @@ async def start_google_auth(
         return f"**Authentication Error:** {error_message}"
 
     try:
-        auth_message = await start_auth_flow(
+        auth_message, _ = await start_auth_flow(
             user_google_email=user_google_email,
             service_name=service_name,
             redirect_uri=get_oauth_redirect_uri_for_current_mode(),

main.py

@@ -311,9 +311,12 @@ def main():
             safe_print(
                 f"   {tool_icons[tool]} {tool.title()} - Google {tool.title()} API integration"
             )
-        except ModuleNotFoundError as exc:
+        except Exception as exc:
             logger.error("Failed to import tool '%s': %s", tool, exc, exc_info=True)
-            safe_print(f"   ⚠️ Failed to load {tool.title()} tool module ({exc}).")
+            if _CLI_MODE:
+                print(f"   Warning: Failed to load {tool.title()} tool module ({exc}).", file=sys.stderr)
+            else:
+                safe_print(f"   ⚠️ Failed to load {tool.title()} tool module ({exc}).")
     safe_print("")
 
     # Filter tools based on tier configuration (if tier-based loading is enabled)