From 1b2a4db1ed1451a50b50c62211119544ad3d3ea6 Mon Sep 17 00:00:00 2001 From: "Paul Gauthier (aider)" Date: Fri, 28 Mar 2025 17:13:42 -1000 Subject: [PATCH] feat: Add OpenRouter OAuth PKCE flow for authentication --- aider/onboarding.py | 257 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 251 insertions(+), 6 deletions(-) diff --git a/aider/onboarding.py b/aider/onboarding.py index b525acbcd..9179beaeb 100644 --- a/aider/onboarding.py +++ b/aider/onboarding.py @@ -1,11 +1,24 @@ +import base64 +import hashlib +import http.server import os +import secrets +import socketserver +import threading +import time +import webbrowser +from urllib.parse import parse_qs, urlparse + +import requests from aider import urls +from aider.utils import check_pip_install_extra def select_default_model(args, io, analytics): """ Selects a default model based on available API keys if no model is specified. + Offers OAuth flow for OpenRouter if no keys are found. Args: args: The command line arguments object. @@ -29,18 +42,250 @@ def select_default_model(args, io, analytics): ] selected_model = None + found_key_env_var = None for env_key, model_name in model_key_pairs: - if os.environ.get(env_key): + api_key_value = os.environ.get(env_key) + # Special check for Vertex AI project which isn't a key but acts like one for selection + is_vertex = env_key == "VERTEXAI_PROJECT" and api_key_value + if api_key_value and (not is_vertex or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")): selected_model = model_name + found_key_env_var = env_key io.tool_warning(f"Using {model_name} model with {env_key} environment variable.") # Track which API key was used for auto-selection analytics.event("auto_model_selection", api_key=env_key) break - if not selected_model: - io.tool_error("You need to specify a --model and an --api-key to use.") - io.offer_url(urls.models_and_keys, "Open documentation url for more info?") - analytics.event("auto_model_selection", api_key=None) + if selected_model: + return selected_model + + # No API keys found - Offer OpenRouter OAuth + io.tool_warning("No API key environment variables found (e.g., OPENAI_API_KEY, ANTHROPIC_API_KEY...).") + # Use confirm_ask which handles non-interactive cases + if io.confirm_ask( + "Authenticate with OpenRouter via browser to get an API key?", default="y", group="openrouter_oauth" + ): + analytics.event("oauth_flow_initiated", provider="openrouter") + openrouter_key = start_openrouter_oauth_flow(io, analytics) + if openrouter_key: + # Successfully got key via OAuth, use the default OpenRouter model + # Ensure OPENROUTER_API_KEY is now set in the environment for later use + os.environ["OPENROUTER_API_KEY"] = openrouter_key + selected_model = "openrouter/anthropic/claude-3.7-sonnet" # Default OR model + io.tool_warning(f"Using {selected_model} model via OpenRouter OAuth.") + # Track OAuth success leading to model selection + analytics.event("auto_model_selection", api_key="OPENROUTER_API_KEY_OAUTH") + return selected_model + else: + # OAuth failed or was cancelled by user implicitly (e.g., closing browser) + # Error messages are handled within start_openrouter_oauth_flow + io.tool_error("OpenRouter authentication did not complete successfully.") + # Fall through to the final error message + + # Final fallback if no key found and OAuth not attempted or failed/declined + io.tool_error( + "No model specified and no API key found or configured.\n" + "Please set an API key environment variable (e.g., OPENAI_API_KEY),\n" + "use the OpenRouter authentication flow (if offered),\n" + "or specify both --model and --api-key." + ) + io.offer_url(urls.models_and_keys, "Open documentation URL for more info?") + analytics.event("auto_model_selection", api_key=None) # Track failure + return None + + +# Helper function to find an available port +def find_available_port(start_port=8484, end_port=8584): + for port in range(start_port, end_port + 1): + try: + with socketserver.TCPServer(("localhost", port), None) as s: + return port + except OSError: + continue + return None + + +# PKCE code generation +def generate_pkce_codes(): + code_verifier = secrets.token_urlsafe(64) + hasher = hashlib.sha256() + hasher.update(code_verifier.encode("utf-8")) + code_challenge = base64.urlsafe_b64encode(hasher.digest()).rstrip(b"=").decode("utf-8") + return code_verifier, code_challenge + + +# Function to exchange the authorization code for an API key +def exchange_code_for_key(code, code_verifier, io): + try: + response = requests.post( + "https://openrouter.ai/api/v1/auth/keys", + headers={"Content-Type": "application/json"}, + json={ + "code": code, + "code_verifier": code_verifier, + "code_challenge_method": "S256", + }, + timeout=30, # Add a timeout + ) + response.raise_for_status() # Raise exception for bad status codes (4xx or 5xx) + data = response.json() + api_key = data.get("key") + if not api_key: + io.tool_error("Error: 'key' not found in OpenRouter response.") + io.tool_error(f"Response: {response.text}") + return None + return api_key + except requests.exceptions.Timeout: + io.tool_error("Error: Request to OpenRouter timed out during code exchange.") + return None + except requests.exceptions.HTTPError as e: + io.tool_error(f"Error exchanging code for OpenRouter key: {e.status_code} {e.response.reason}") + io.tool_error(f"Response: {e.response.text}") + return None + except requests.exceptions.RequestException as e: + io.tool_error(f"Error exchanging code for OpenRouter key: {e}") + return None + except Exception as e: + io.tool_error(f"Unexpected error during code exchange: {e}") return None - return selected_model + +# Function to start the OAuth flow +def start_openrouter_oauth_flow(io, analytics): + """Initiates the OpenRouter OAuth PKCE flow using a local server.""" + + # Check for requests library + if not check_pip_install_extra(io, "requests", "OpenRouter OAuth", "aider[oauth]"): + return None + + port = find_available_port() + if not port: + io.tool_error("Could not find an available port between 8484 and 8584.") + io.tool_error("Please ensure a port in this range is free, or configure manually.") + return None + + callback_url = f"http://localhost:{port}/callback" + auth_code = None + server_error = None + server_started = threading.Event() + shutdown_server = threading.Event() + + class OAuthCallbackHandler(http.server.SimpleHTTPRequestHandler): + def do_GET(self): + nonlocal auth_code, server_error + parsed_path = urlparse(self.path) + if parsed_path.path == "/callback": + query_params = parse_qs(parsed_path.query) + if "code" in query_params: + auth_code = query_params["code"][0] + self.send_response(200) + self.send_header("Content-type", "text/html") + self.end_headers() + self.wfile.write( + b"

Success!

" + b"

Aider has received the authentication code. " + b"You can close this browser tab.

" + ) + # Signal the main thread to shut down the server + shutdown_server.set() + else: + server_error = "Missing 'code' parameter in callback URL." + self.send_response(400) + self.send_header("Content-type", "text/html") + self.end_headers() + self.wfile.write( + b"

Error

" + b"

Missing 'code' parameter in callback URL.

" + b"

Please check the Aider terminal.

" + ) + shutdown_server.set() + else: + self.send_response(404) + self.end_headers() + self.wfile.write(b"Not Found") + + def log_message(self, format, *args): + # Suppress server logging to keep terminal clean + pass + + def run_server(): + nonlocal server_error + try: + with socketserver.TCPServer(("localhost", port), OAuthCallbackHandler) as httpd: + io.tool_output(f"Temporary server listening on {callback_url}", log_only=True) + server_started.set() # Signal that the server is ready + # Wait until shutdown is requested or timeout occurs (handled by main thread) + while not shutdown_server.is_set(): + httpd.handle_request() # Handle one request at a time + # Add a small sleep to prevent busy-waiting if needed, + # though handle_request should block appropriately. + time.sleep(0.1) + io.tool_output("Shutting down temporary server.", log_only=True) + except Exception as e: + server_error = f"Failed to start or run temporary server: {e}" + server_started.set() # Signal even if failed, error will be checked + shutdown_server.set() # Ensure shutdown logic proceeds + + server_thread = threading.Thread(target=run_server, daemon=True) + server_thread.start() + + # Wait briefly for the server to start, or for an error + if not server_started.wait(timeout=5): + io.tool_error("Temporary authentication server failed to start in time.") + shutdown_server.set() # Ensure thread exits if it eventually starts + server_thread.join(timeout=1) + return None + + # Check if server failed during startup + if server_error: + io.tool_error(server_error) + shutdown_server.set() # Ensure thread exits + server_thread.join(timeout=1) + return None + + # Generate codes and URL + code_verifier, code_challenge = generate_pkce_codes() + auth_url = f"https://openrouter.ai/auth?callback_url={callback_url}&code_challenge={code_challenge}&code_challenge_method=S256" + + io.tool_output("\nPlease open the following URL in your web browser to authorize Aider with OpenRouter:") + io.tool_output(auth_url) + io.tool_output("\nWaiting for authentication... (Timeout: 2 minutes)") + + try: + webbrowser.open(auth_url) + except Exception as e: + io.tool_warning(f"Could not automatically open browser: {e}") + io.tool_output("Please manually open the URL above.") + + # Wait for the callback to set the auth_code or for timeout/error + shutdown_server.wait(timeout=120) # 2 minute timeout + + # Join the server thread to ensure it's cleaned up + server_thread.join(timeout=1) + + if server_error: + io.tool_error(f"Authentication failed: {server_error}") + analytics.event("oauth_flow_failed", provider="openrouter", reason=server_error) + return None + + if not auth_code: + io.tool_error("Authentication timed out. No code received from OpenRouter.") + analytics.event("oauth_flow_failed", provider="openrouter", reason="timeout") + return None + + io.tool_output("Authentication code received. Exchanging for API key...") + analytics.event("oauth_flow_code_received", provider="openrouter") + + # Exchange code for key + api_key = exchange_code_for_key(auth_code, code_verifier, io) + + if api_key: + io.tool_output("Successfully obtained and configured OpenRouter API key.") + # Securely store this key? For now, set env var for the session. + os.environ["OPENROUTER_API_KEY"] = api_key + io.tool_warning("Set OPENROUTER_API_KEY environment variable for this session.") + analytics.event("oauth_flow_success", provider="openrouter") + return api_key + else: + io.tool_error("Failed to obtain OpenRouter API key from code.") + analytics.event("oauth_flow_failed", provider="openrouter", reason="code_exchange_failed") + return None