You can close this window. Please check the proxy server logs to verify that onboarding completed successfully. No need to restart the proxy.
") else: self.send_response(400) self.send_header("Content-type", "text/html") self.end_headers() self.wfile.write(b"Please try again.
") def authenticate_user(request: Request): """Authenticate the user with multiple methods.""" # Check for API key in query parameters first (for Gemini client compatibility) api_key = request.query_params.get("key") if api_key and api_key == GEMINI_AUTH_PASSWORD: return "api_key_user" # Check for API key in x-goog-api-key header (Google SDK format) goog_api_key = request.headers.get("x-goog-api-key", "") if goog_api_key and goog_api_key == GEMINI_AUTH_PASSWORD: return "goog_api_key_user" # Check for API key in Authorization header (Bearer token format) auth_header = request.headers.get("authorization", "") if auth_header.startswith("Bearer "): bearer_token = auth_header[7:] if bearer_token == GEMINI_AUTH_PASSWORD: return "bearer_user" # Check for HTTP Basic Authentication if auth_header.startswith("Basic "): try: encoded_credentials = auth_header[6:] decoded_credentials = base64.b64decode(encoded_credentials).decode('utf-8') username, password = decoded_credentials.split(':', 1) if password == GEMINI_AUTH_PASSWORD: return username except Exception: pass # If none of the authentication methods work raise HTTPException( status_code=401, detail="Invalid authentication credentials. Use HTTP Basic Auth, Bearer token, 'key' query parameter, or 'x-goog-api-key' header.", headers={"WWW-Authenticate": "Basic"}, ) def save_credentials(creds, project_id=None): global credentials_from_env # Don't save credentials to file if they came from environment variable, # but still save project_id if provided and no file exists or file lacks project_id if credentials_from_env: if project_id and os.path.exists(CREDENTIAL_FILE): try: with open(CREDENTIAL_FILE, "r") as f: existing_data = json.load(f) # Only update project_id if it's missing from the file if "project_id" not in existing_data: existing_data["project_id"] = project_id with open(CREDENTIAL_FILE, "w") as f: json.dump(existing_data, f, indent=2) logging.info(f"Added project_id {project_id} to existing credential file") except Exception as e: logging.warning(f"Could not update project_id in credential file: {e}") return creds_data = { "client_id": CLIENT_ID, "client_secret": CLIENT_SECRET, "token": creds.token, "refresh_token": creds.refresh_token, "scopes": creds.scopes if creds.scopes else SCOPES, "token_uri": "https://oauth2.googleapis.com/token", } if creds.expiry: if creds.expiry.tzinfo is None: from datetime import timezone expiry_utc = creds.expiry.replace(tzinfo=timezone.utc) else: expiry_utc = creds.expiry # Keep the existing ISO format for backward compatibility, but ensure it's properly handled during loading creds_data["expiry"] = expiry_utc.isoformat() if project_id: creds_data["project_id"] = project_id elif os.path.exists(CREDENTIAL_FILE): try: with open(CREDENTIAL_FILE, "r") as f: existing_data = json.load(f) if "project_id" in existing_data: creds_data["project_id"] = existing_data["project_id"] except Exception: pass with open(CREDENTIAL_FILE, "w") as f: json.dump(creds_data, f, indent=2) def get_credentials(allow_oauth_flow=True): """Loads credentials matching gemini-cli OAuth2 flow.""" global credentials, credentials_from_env, user_project_id if credentials and credentials.token: return credentials # Check for credentials in environment variable (JSON string) env_creds_json = os.getenv("GEMINI_CREDENTIALS") if env_creds_json: # First, check if we have a refresh token - if so, we should always be able to load credentials try: raw_env_creds_data = json.loads(env_creds_json) # SAFEGUARD: If refresh_token exists, we should always load credentials successfully if "refresh_token" in raw_env_creds_data and raw_env_creds_data["refresh_token"]: logging.info("Environment refresh token found - ensuring credentials load successfully") try: creds_data = raw_env_creds_data.copy() # Handle different credential formats if "access_token" in creds_data and "token" not in creds_data: creds_data["token"] = creds_data["access_token"] if "scope" in creds_data and "scopes" not in creds_data: creds_data["scopes"] = creds_data["scope"].split() # Handle problematic expiry formats that cause parsing errors if "expiry" in creds_data: expiry_str = creds_data["expiry"] # If expiry has timezone info that causes parsing issues, try to fix it if isinstance(expiry_str, str) and ("+00:00" in expiry_str or "Z" in expiry_str): try: # Try to parse and reformat the expiry to a format Google Credentials can handle from datetime import datetime if "+00:00" in expiry_str: # Handle ISO format with timezone offset parsed_expiry = datetime.fromisoformat(expiry_str) elif expiry_str.endswith("Z"): # Handle ISO format with Z suffix parsed_expiry = datetime.fromisoformat(expiry_str.replace('Z', '+00:00')) else: parsed_expiry = datetime.fromisoformat(expiry_str) # Convert to UTC timestamp format that Google Credentials library expects import time timestamp = parsed_expiry.timestamp() creds_data["expiry"] = datetime.utcfromtimestamp(timestamp).strftime("%Y-%m-%dT%H:%M:%SZ") logging.info(f"Converted environment expiry format from '{expiry_str}' to '{creds_data['expiry']}'") except Exception as expiry_error: logging.warning(f"Could not parse environment expiry format '{expiry_str}': {expiry_error}, removing expiry field") # Remove problematic expiry field - credentials will be treated as expired but still loadable del creds_data["expiry"] credentials = Credentials.from_authorized_user_info(creds_data, SCOPES) credentials_from_env = True # Mark as environment credentials # Extract project_id from environment credentials if available if "project_id" in raw_env_creds_data: user_project_id = raw_env_creds_data["project_id"] logging.info(f"Extracted project_id from environment credentials: {user_project_id}") # Try to refresh if expired and refresh token exists if credentials.expired and credentials.refresh_token: try: logging.info("Environment credentials expired, attempting refresh...") credentials.refresh(GoogleAuthRequest()) logging.info("Environment credentials refreshed successfully") except Exception as refresh_error: logging.warning(f"Failed to refresh environment credentials: {refresh_error}") logging.info("Using existing environment credentials despite refresh failure") elif not credentials.expired: logging.info("Environment credentials are still valid, no refresh needed") elif not credentials.refresh_token: logging.warning("Environment credentials expired but no refresh token available") return credentials except Exception as parsing_error: # SAFEGUARD: Even if parsing fails, try to create minimal credentials with refresh token logging.warning(f"Failed to parse environment credentials normally: {parsing_error}") logging.info("Attempting to create minimal environment credentials with refresh token") try: minimal_creds_data = { "client_id": raw_env_creds_data.get("client_id", CLIENT_ID), "client_secret": raw_env_creds_data.get("client_secret", CLIENT_SECRET), "refresh_token": raw_env_creds_data["refresh_token"], "token_uri": "https://oauth2.googleapis.com/token", } credentials = Credentials.from_authorized_user_info(minimal_creds_data, SCOPES) credentials_from_env = True # Mark as environment credentials # Extract project_id from environment credentials if available if "project_id" in raw_env_creds_data: user_project_id = raw_env_creds_data["project_id"] logging.info(f"Extracted project_id from minimal environment credentials: {user_project_id}") # Force refresh since we don't have a valid token try: logging.info("Refreshing minimal environment credentials...") credentials.refresh(GoogleAuthRequest()) logging.info("Minimal environment credentials refreshed successfully") return credentials except Exception as refresh_error: logging.error(f"Failed to refresh minimal environment credentials: {refresh_error}") # Even if refresh fails, return the credentials - they might still work return credentials except Exception as minimal_error: logging.error(f"Failed to create minimal environment credentials: {minimal_error}") # Fall through to file-based credentials else: logging.warning("No refresh token found in environment credentials") # Fall through to file-based credentials except Exception as e: logging.error(f"Failed to parse environment credentials JSON: {e}") # Fall through to file-based credentials # Check for credentials file (CREDENTIAL_FILE now includes GOOGLE_APPLICATION_CREDENTIALS path if set) if os.path.exists(CREDENTIAL_FILE): # First, check if we have a refresh token - if so, we should always be able to load credentials try: with open(CREDENTIAL_FILE, "r") as f: raw_creds_data = json.load(f) # SAFEGUARD: If refresh_token exists, we should always load credentials successfully if "refresh_token" in raw_creds_data and raw_creds_data["refresh_token"]: logging.info("Refresh token found - ensuring credentials load successfully") try: creds_data = raw_creds_data.copy() # Handle different credential formats if "access_token" in creds_data and "token" not in creds_data: creds_data["token"] = creds_data["access_token"] if "scope" in creds_data and "scopes" not in creds_data: creds_data["scopes"] = creds_data["scope"].split() # Handle problematic expiry formats that cause parsing errors if "expiry" in creds_data: expiry_str = creds_data["expiry"] # If expiry has timezone info that causes parsing issues, try to fix it if isinstance(expiry_str, str) and ("+00:00" in expiry_str or "Z" in expiry_str): try: # Try to parse and reformat the expiry to a format Google Credentials can handle from datetime import datetime if "+00:00" in expiry_str: # Handle ISO format with timezone offset parsed_expiry = datetime.fromisoformat(expiry_str) elif expiry_str.endswith("Z"): # Handle ISO format with Z suffix parsed_expiry = datetime.fromisoformat(expiry_str.replace('Z', '+00:00')) else: parsed_expiry = datetime.fromisoformat(expiry_str) # Convert to UTC timestamp format that Google Credentials library expects import time timestamp = parsed_expiry.timestamp() creds_data["expiry"] = datetime.utcfromtimestamp(timestamp).strftime("%Y-%m-%dT%H:%M:%SZ") logging.info(f"Converted expiry format from '{expiry_str}' to '{creds_data['expiry']}'") except Exception as expiry_error: logging.warning(f"Could not parse expiry format '{expiry_str}': {expiry_error}, removing expiry field") # Remove problematic expiry field - credentials will be treated as expired but still loadable del creds_data["expiry"] credentials = Credentials.from_authorized_user_info(creds_data, SCOPES) # Mark as environment credentials if GOOGLE_APPLICATION_CREDENTIALS was used credentials_from_env = bool(os.getenv("GOOGLE_APPLICATION_CREDENTIALS")) # Try to refresh if expired and refresh token exists if credentials.expired and credentials.refresh_token: try: logging.info("File-based credentials expired, attempting refresh...") credentials.refresh(GoogleAuthRequest()) logging.info("File-based credentials refreshed successfully") save_credentials(credentials) except Exception as refresh_error: logging.warning(f"Failed to refresh file-based credentials: {refresh_error}") logging.info("Using existing file-based credentials despite refresh failure") elif not credentials.expired: logging.info("File-based credentials are still valid, no refresh needed") elif not credentials.refresh_token: logging.warning("File-based credentials expired but no refresh token available") return credentials except Exception as parsing_error: # SAFEGUARD: Even if parsing fails, try to create minimal credentials with refresh token logging.warning(f"Failed to parse credentials normally: {parsing_error}") logging.info("Attempting to create minimal credentials with refresh token") try: minimal_creds_data = { "client_id": raw_creds_data.get("client_id", CLIENT_ID), "client_secret": raw_creds_data.get("client_secret", CLIENT_SECRET), "refresh_token": raw_creds_data["refresh_token"], "token_uri": "https://oauth2.googleapis.com/token", } credentials = Credentials.from_authorized_user_info(minimal_creds_data, SCOPES) credentials_from_env = bool(os.getenv("GOOGLE_APPLICATION_CREDENTIALS")) # Force refresh since we don't have a valid token try: logging.info("Refreshing minimal credentials...") credentials.refresh(GoogleAuthRequest()) logging.info("Minimal credentials refreshed successfully") save_credentials(credentials) return credentials except Exception as refresh_error: logging.error(f"Failed to refresh minimal credentials: {refresh_error}") # Even if refresh fails, return the credentials - they might still work return credentials except Exception as minimal_error: logging.error(f"Failed to create minimal credentials: {minimal_error}") # Fall through to new login as last resort else: logging.warning("No refresh token found in credentials file") # Fall through to new login except Exception as e: logging.error(f"Failed to read credentials file {CREDENTIAL_FILE}: {e}") # Fall through to new login only if file is completely unreadable # Only start OAuth flow if explicitly allowed if not allow_oauth_flow: logging.info("OAuth flow not allowed - returning None (credentials will be required on first request)") return None client_config = { "installed": { "client_id": CLIENT_ID, "client_secret": CLIENT_SECRET, "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", } } flow = Flow.from_client_config( client_config, scopes=SCOPES, redirect_uri="http://localhost:8080" ) flow.oauth2session.scope = SCOPES auth_url, _ = flow.authorization_url( access_type="offline", prompt="consent", include_granted_scopes='true' ) print(f"\n{'='*80}") print(f"AUTHENTICATION REQUIRED") print(f"{'='*80}") print(f"Please open this URL in your browser to log in:") print(f"{auth_url}") print(f"{'='*80}\n") logging.info(f"Please open this URL in your browser to log in: {auth_url}") server = HTTPServer(("", 8080), _OAuthCallbackHandler) server.handle_request() auth_code = _OAuthCallbackHandler.auth_code if not auth_code: return None import oauthlib.oauth2.rfc6749.parameters original_validate = oauthlib.oauth2.rfc6749.parameters.validate_token_parameters def patched_validate(params): try: return original_validate(params) except Warning: pass oauthlib.oauth2.rfc6749.parameters.validate_token_parameters = patched_validate try: flow.fetch_token(code=auth_code) credentials = flow.credentials credentials_from_env = False # Mark as file-based credentials save_credentials(credentials) logging.info("Authentication successful! Credentials saved.") return credentials except Exception as e: logging.error(f"Authentication failed: {e}") return None finally: oauthlib.oauth2.rfc6749.parameters.validate_token_parameters = original_validate def onboard_user(creds, project_id): """Ensures the user is onboarded, matching gemini-cli setupUser behavior.""" global onboarding_complete if onboarding_complete: return if creds.expired and creds.refresh_token: try: creds.refresh(GoogleAuthRequest()) save_credentials(creds) except Exception as e: raise Exception(f"Failed to refresh credentials during onboarding: {str(e)}") headers = { "Authorization": f"Bearer {creds.token}", "Content-Type": "application/json", "User-Agent": get_user_agent(), } load_assist_payload = { "cloudaicompanionProject": project_id, "metadata": get_client_metadata(project_id), } try: import requests resp = requests.post( f"{CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist", data=json.dumps(load_assist_payload), headers=headers, ) resp.raise_for_status() load_data = resp.json() tier = None if load_data.get("currentTier"): tier = load_data["currentTier"] else: for allowed_tier in load_data.get("allowedTiers", []): if allowed_tier.get("isDefault"): tier = allowed_tier break if not tier: tier = { "name": "", "description": "", "id": "legacy-tier", "userDefinedCloudaicompanionProject": True, } if tier.get("userDefinedCloudaicompanionProject") and not project_id: raise ValueError("This account requires setting the GOOGLE_CLOUD_PROJECT env var.") if load_data.get("currentTier"): onboarding_complete = True return onboard_req_payload = { "tierId": tier.get("id"), "cloudaicompanionProject": project_id, "metadata": get_client_metadata(project_id), } while True: onboard_resp = requests.post( f"{CODE_ASSIST_ENDPOINT}/v1internal:onboardUser", data=json.dumps(onboard_req_payload), headers=headers, ) onboard_resp.raise_for_status() lro_data = onboard_resp.json() if lro_data.get("done"): onboarding_complete = True break time.sleep(5) except requests.exceptions.HTTPError as e: raise Exception(f"User onboarding failed. Please check your Google Cloud project permissions and try again. Error: {e.response.text if hasattr(e, 'response') else str(e)}") except Exception as e: raise Exception(f"User onboarding failed due to an unexpected error: {str(e)}") def get_user_project_id(creds): """Gets the user's project ID matching gemini-cli setupUser logic.""" global user_project_id # Priority 1: Check environment variable first (always check, even if user_project_id is set) env_project_id = os.getenv("GOOGLE_CLOUD_PROJECT") if env_project_id: logging.info(f"Using project ID from GOOGLE_CLOUD_PROJECT environment variable: {env_project_id}") user_project_id = env_project_id save_credentials(creds, user_project_id) return user_project_id # If we already have a cached project_id and no env var override, use it if user_project_id: logging.info(f"Using cached project ID: {user_project_id}") return user_project_id # Priority 2: Check cached project ID in credential file if os.path.exists(CREDENTIAL_FILE): try: with open(CREDENTIAL_FILE, "r") as f: creds_data = json.load(f) cached_project_id = creds_data.get("project_id") if cached_project_id: logging.info(f"Using cached project ID from credential file: {cached_project_id}") user_project_id = cached_project_id return user_project_id except Exception as e: logging.warning(f"Could not read project_id from credential file: {e}") # Priority 3: Make API call to discover project ID # Ensure we have valid credentials for the API call if creds.expired and creds.refresh_token: try: logging.info("Refreshing credentials before project ID discovery...") creds.refresh(GoogleAuthRequest()) save_credentials(creds) logging.info("Credentials refreshed successfully for project ID discovery") except Exception as e: logging.error(f"Failed to refresh credentials while getting project ID: {e}") # Continue with existing credentials - they might still work if not creds.token: raise Exception("No valid access token available for project ID discovery") headers = { "Authorization": f"Bearer {creds.token}", "Content-Type": "application/json", "User-Agent": get_user_agent(), } probe_payload = { "metadata": get_client_metadata(), } try: import requests logging.info("Attempting to discover project ID via API call...") resp = requests.post( f"{CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist", data=json.dumps(probe_payload), headers=headers, ) resp.raise_for_status() data = resp.json() discovered_project_id = data.get("cloudaicompanionProject") if not discovered_project_id: raise ValueError("Could not find 'cloudaicompanionProject' in loadCodeAssist response.") logging.info(f"Discovered project ID via API: {discovered_project_id}") user_project_id = discovered_project_id save_credentials(creds, user_project_id) return user_project_id except requests.exceptions.HTTPError as e: logging.error(f"HTTP error during project ID discovery: {e}") if hasattr(e, 'response') and e.response: logging.error(f"Response status: {e.response.status_code}, body: {e.response.text}") raise Exception(f"Failed to discover project ID via API: {e}") except Exception as e: logging.error(f"Unexpected error during project ID discovery: {e}") raise Exception(f"Failed to discover project ID: {e}")