From d5a43fd8a402d9bc57e882127562051b09547b45 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 13:56:51 +0100 Subject: [PATCH 001/221] =?UTF-8?q?feat(providers):=20=E2=9C=A8=20add=20an?= =?UTF-8?q?tigravity=20provider=20and=20auth=20base?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a new Antigravity provider and authentication base to integrate with the Antigravity (internal Google) API. - Add providers/antigravity_auth_base.py: OAuth2 token management with env/file loading, atomic saves, refresh logic, backoff/queue tracking, interactive and headless browser auth flow, and helper utilities. - Add providers/antigravity_provider.py: request/response transformations (OpenAI → Gemini CLI → Antigravity), model aliasing, thinking/reasoning config mapping, tool response grouping, streaming & non-streaming handling, and base-URL fallback. - Update provider_factory.py and providers/__init__.py to register the new provider. - Bump project metadata in pyproject.toml (package name and version). BREAKING CHANGE: project packaging metadata updated — package name changed to "rotator_library" and version bumped to 0.95. Update any dependency or packaging references that relied on the previous name/version. --- src/rotator_library/provider_factory.py | 2 + src/rotator_library/providers/__init__.py | 2 + .../providers/antigravity_auth_base.py | 466 ++++++++++ .../providers/antigravity_provider.py | 869 ++++++++++++++++++ src/rotator_library/pyproject.toml | 4 +- 5 files changed, 1341 insertions(+), 2 deletions(-) create mode 100644 src/rotator_library/providers/antigravity_auth_base.py create mode 100644 src/rotator_library/providers/antigravity_provider.py diff --git a/src/rotator_library/provider_factory.py b/src/rotator_library/provider_factory.py index f53eabd0..f13d16aa 100644 --- a/src/rotator_library/provider_factory.py +++ b/src/rotator_library/provider_factory.py @@ -3,11 +3,13 @@ from .providers.gemini_auth_base import GeminiAuthBase from .providers.qwen_auth_base import QwenAuthBase from .providers.iflow_auth_base import IFlowAuthBase +from .providers.antigravity_auth_base import AntigravityAuthBase PROVIDER_MAP = { "gemini_cli": GeminiAuthBase, "qwen_code": QwenAuthBase, "iflow": IFlowAuthBase, + "antigravity": AntigravityAuthBase, } def get_provider_auth_class(provider_name: str): diff --git a/src/rotator_library/providers/__init__.py b/src/rotator_library/providers/__init__.py index 3541d11a..c6bee073 100644 --- a/src/rotator_library/providers/__init__.py +++ b/src/rotator_library/providers/__init__.py @@ -112,6 +112,8 @@ def _register_providers(): "chutes", "iflow", "qwen_code", + "gemini_cli", + "antigravity", ]: continue diff --git a/src/rotator_library/providers/antigravity_auth_base.py b/src/rotator_library/providers/antigravity_auth_base.py new file mode 100644 index 00000000..14b470f5 --- /dev/null +++ b/src/rotator_library/providers/antigravity_auth_base.py @@ -0,0 +1,466 @@ +# src/rotator_library/providers/antigravity_auth_base.py + +import os +import webbrowser +from typing import Union, Optional +import json +import time +import asyncio +import logging +from pathlib import Path +from typing import Dict, Any +import tempfile +import shutil + +import httpx +from rich.console import Console +from rich.panel import Panel +from rich.text import Text + +from ..utils.headless_detection import is_headless_environment + +lib_logger = logging.getLogger('rotator_library') + +# Antigravity OAuth credentials from CLIProxyAPI +CLIENT_ID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" +CLIENT_SECRET = "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" +TOKEN_URI = "https://oauth2.googleapis.com/token" +USER_INFO_URI = "https://www.googleapis.com/oauth2/v1/userinfo" +REFRESH_EXPIRY_BUFFER_SECONDS = 30 * 60 # 30 minutes buffer before expiry + +# Antigravity requires additional scopes +OAUTH_SCOPES = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + "https://www.googleapis.com/auth/cclog", # Antigravity-specific + "https://www.googleapis.com/auth/experimentsandconfigs" # Antigravity-specific +] + +console = Console() + +class AntigravityAuthBase: + """ + Base authentication class for Antigravity provider. + Handles OAuth2 flow, token management, and refresh logic. + + Based on GeminiAuthBase but uses Antigravity-specific OAuth credentials and scopes. + """ + + def __init__(self): + self._credentials_cache: Dict[str, Dict[str, Any]] = {} + self._refresh_locks: Dict[str, asyncio.Lock] = {} + self._locks_lock = asyncio.Lock() # Protects the locks dict from race conditions + # [BACKOFF TRACKING] Track consecutive failures per credential + self._refresh_failures: Dict[str, int] = {} # Track consecutive failures per credential + self._next_refresh_after: Dict[str, float] = {} # Track backoff timers (Unix timestamp) + + # [QUEUE SYSTEM] Sequential refresh processing + self._refresh_queue: asyncio.Queue = asyncio.Queue() + self._queued_credentials: set = set() # Track credentials already in queue + self._unavailable_credentials: set = set() # Mark credentials unavailable during re-auth + self._queue_tracking_lock = asyncio.Lock() # Protects queue sets + self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task + + def _load_from_env(self) -> Optional[Dict[str, Any]]: + """ + Load OAuth credentials from environment variables for stateless deployments. + + Expected environment variables: + - ANTIGRAVITY_ACCESS_TOKEN (required) + - ANTIGRAVITY_REFRESH_TOKEN (required) + - ANTIGRAVITY_EXPIRY_DATE (optional, defaults to 0) + - ANTIGRAVITY_CLIENT_ID (optional, uses default) + - ANTIGRAVITY_CLIENT_SECRET (optional, uses default) + - ANTIGRAVITY_TOKEN_URI (optional, uses default) + - ANTIGRAVITY_UNIVERSE_DOMAIN (optional, defaults to googleapis.com) + - ANTIGRAVITY_EMAIL (optional, defaults to "env-user") + + Returns: + Dict with credential structure if env vars present, None otherwise + """ + access_token = os.getenv("ANTIGRAVITY_ACCESS_TOKEN") + refresh_token = os.getenv("ANTIGRAVITY_REFRESH_TOKEN") + + # Both access and refresh tokens are required + if not (access_token and refresh_token): + return None + + lib_logger.debug("Loading Antigravity credentials from environment variables") + + # Parse expiry_date as float, default to 0 if not present + expiry_str = os.getenv("ANTIGRAVITY_EXPIRY_DATE", "0") + try: + expiry_date = float(expiry_str) + except ValueError: + lib_logger.warning(f"Invalid ANTIGRAVITY_EXPIRY_DATE value: {expiry_str}, using 0") + expiry_date = 0 + + creds = { + "access_token": access_token, + "refresh_token": refresh_token, + "expiry_date": expiry_date, + "client_id": os.getenv("ANTIGRAVITY_CLIENT_ID", CLIENT_ID), + "client_secret": os.getenv("ANTIGRAVITY_CLIENT_SECRET", CLIENT_SECRET), + "token_uri": os.getenv("ANTIGRAVITY_TOKEN_URI", TOKEN_URI), + "universe_domain": os.getenv("ANTIGRAVITY_UNIVERSE_DOMAIN", "googleapis.com"), + "_proxy_metadata": { + "email": os.getenv("ANTIGRAVITY_EMAIL", "env-user"), + "last_check_timestamp": time.time(), + "loaded_from_env": True # Flag to indicate env-based credentials + } + } + + return creds + + async def _load_credentials(self, path: str) -> Dict[str, Any]: + """ + Load credentials from a file. First attempts file-based load, + then falls back to environment variables if file not found. + + Args: + path: File path to load credentials from + + Returns: + Dict containing the credentials + + Raises: + ValueError: If credentials cannot be loaded from either source + """ + # If path is special marker "env", load from environment + if path == "env": + env_creds = self._load_from_env() + if env_creds: + lib_logger.debug("Using Antigravity credentials from environment variables") + return env_creds + raise ValueError("ANTIGRAVITY_ACCESS_TOKEN and ANTIGRAVITY_REFRESH_TOKEN environment variables not set") + + # Try loading from cache first + if path in self._credentials_cache: + cached_creds = self._credentials_cache[path] + lib_logger.debug(f"Using cached Antigravity credentials for: {Path(path).name}") + return cached_creds + + # Try loading from file + try: + with open(path, 'r') as f: + creds = json.load(f) + self._credentials_cache[path] = creds + lib_logger.debug(f"Loaded Antigravity credentials from file: {Path(path).name}") + return creds + except FileNotFoundError: + # Fall back to environment variables + lib_logger.debug(f"Credential file not found: {path}, attempting environment variables") + env_creds = self._load_from_env() + if env_creds: + lib_logger.debug("Using Antigravity credentials from environment variables as fallback") + # Cache with special path marker + self._credentials_cache[path] = env_creds + return env_creds + raise ValueError(f"Credential file not found: {path} and environment variables not set") + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON in credential file {path}: {e}") + + async def _save_credentials(self, path: str, creds: Dict[str, Any]) -> None: + """ + Save credentials to a file. Skip if credentials were loaded from environment. + + Args: + path: File path to save credentials to + creds: Credentials dictionary to save + """ + # Don't save environment-based credentials to file + if creds.get("_proxy_metadata", {}).get("loaded_from_env"): + lib_logger.debug("Skipping credential save (loaded from environment)") + return + + # Don't save if path is special marker + if path == "env": + return + + try: + # Ensure directory exists + Path(path).parent.mkdir(parents=True, exist_ok=True) + + # Write atomically using temp file + rename + temp_fd, temp_path = tempfile.mkstemp( + dir=Path(path).parent, + prefix='.tmp_', + suffix='.json' + ) + try: + with os.fdopen(temp_fd, 'w') as f: + json.dump(creds, f, indent=2) + shutil.move(temp_path, path) + lib_logger.debug(f"Saved Antigravity credentials to: {Path(path).name}") + except Exception: + # Clean up temp file on error + try: + os.unlink(temp_path) + except Exception: + pass + raise + except Exception as e: + lib_logger.warning(f"Failed to save Antigravity credentials to {path}: {e}") + + def _is_token_expired(self, creds: Dict[str, Any]) -> bool: + """ + Check if the access token is expired or close to expiry. + + Args: + creds: Credentials dict with expiry_date field (in milliseconds) + + Returns: + True if token is expired or within buffer time of expiry + """ + if 'expiry_date' not in creds: + return True + + # expiry_date is in milliseconds + expiry_timestamp = creds['expiry_date'] / 1000.0 + current_time = time.time() + + # Consider expired if within buffer time + return (expiry_timestamp - current_time) <= REFRESH_EXPIRY_BUFFER_SECONDS + + async def _refresh_token(self, path: str, creds: Dict[str, Any]) -> Dict[str, Any]: + """ + Refresh an expired OAuth token using the refresh token. + + Args: + path: Credential file path (for saving updated credentials) + creds: Current credentials dict with refresh_token + + Returns: + Updated credentials dict with fresh access token + + Raises: + ValueError: If refresh fails + """ + if 'refresh_token' not in creds: + raise ValueError("No refresh token available") + + lib_logger.debug(f"Refreshing Antigravity OAuth token for: {Path(path).name if path != 'env' else 'env'}") + + client_id = creds.get('client_id', CLIENT_ID) + client_secret = creds.get('client_secret', CLIENT_SECRET) + token_uri = creds.get('token_uri', TOKEN_URI) + + async with httpx.AsyncClient() as client: + try: + response = await client.post( + token_uri, + data={ + 'client_id': client_id, + 'client_secret': client_secret, + 'refresh_token': creds['refresh_token'], + 'grant_type': 'refresh_token' + }, + timeout=30.0 + ) + response.raise_for_status() + token_data = response.json() + + # Update credentials with new token + creds['access_token'] = token_data['access_token'] + creds['expiry_date'] = (time.time() + token_data['expires_in']) * 1000 + + # Update metadata + if '_proxy_metadata' not in creds: + creds['_proxy_metadata'] = {} + creds['_proxy_metadata']['last_check_timestamp'] = time.time() + + # Save updated credentials + await self._save_credentials(path, creds) + + # Update cache + self._credentials_cache[path] = creds + + # Reset failure count on success + self._refresh_failures[path] = 0 + + lib_logger.info(f"Successfully refreshed Antigravity OAuth token for: {Path(path).name if path != 'env' else 'env'}") + return creds + + except httpx.HTTPStatusError as e: + # Track failures for backoff + self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 + raise ValueError(f"Failed to refresh Antigravity token (HTTP {e.response.status_code}): {e.response.text}") + except Exception as e: + self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 + raise ValueError(f"Failed to refresh Antigravity token: {e}") + + async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + """ + Initialize or refresh an OAuth token. Handles the complete OAuth flow if needed. + + Args: + creds_or_path: Either a credentials dict or a file path string + + Returns: + Valid credentials dict with fresh access token + """ + path = creds_or_path if isinstance(creds_or_path, str) else None + + if isinstance(creds_or_path, dict): + display_name = creds_or_path.get("_proxy_metadata", {}).get("display_name", "in-memory object") + else: + display_name = Path(path).name if path and path != "env" else "env" + + lib_logger.debug(f"Initializing Antigravity token for '{display_name}'...") + + try: + creds = await self._load_credentials(creds_or_path) if path else creds_or_path + reason = "" + if not creds.get("refresh_token"): + reason = "refresh token is missing" + elif self._is_token_expired(creds): + reason = "token is expired" + + if reason: + if reason == "token is expired" and creds.get("refresh_token"): + try: + return await self._refresh_token(path, creds) + except Exception as e: + lib_logger.warning(f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login.") + + lib_logger.warning(f"Antigravity OAuth token for '{display_name}' needs setup: {reason}.") + + is_headless = is_headless_environment() + + auth_code_future = asyncio.get_event_loop().create_future() + server = None + + async def handle_callback(reader, writer): + try: + request_line_bytes = await reader.readline() + if not request_line_bytes: + return + path_str = request_line_bytes.decode('utf-8').strip().split(' ')[1] + # Consume headers + while await reader.readline() != b'\r\n': + pass + + from urllib.parse import urlparse, parse_qs + query_params = parse_qs(urlparse(path_str).query) + + writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n") + if 'code' in query_params: + if not auth_code_future.done(): + auth_code_future.set_result(query_params['code'][0]) + writer.write(b"

Authentication successful!

You can close this window.

") + else: + error = query_params.get('error', ['Unknown error'])[0] + if not auth_code_future.done(): + auth_code_future.set_exception(Exception(f"OAuth failed: {error}")) + writer.write(f"

Authentication Failed

Error: {error}. Please try again.

".encode()) + await writer.drain() + except Exception as e: + lib_logger.error(f"Error in OAuth callback handler: {e}") + finally: + writer.close() + + try: + server = await asyncio.start_server(handle_callback, '127.0.0.1', 8085) + + from urllib.parse import urlencode + auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode({ + "client_id": CLIENT_ID, + "redirect_uri": "http://localhost:8085/oauth2callback", + "scope": " ".join(OAUTH_SCOPES), + "access_type": "offline", + "response_type": "code", + "prompt": "consent" + }) + + if is_headless: + auth_panel_text = Text.from_markup( + "Running in headless environment (no GUI detected).\n" + "Please open the URL below in a browser on another machine to authorize:\n" + ) + else: + auth_panel_text = Text.from_markup( + "1. Your browser will now open to log in and authorize the application.\n" + "2. If it doesn't open automatically, please open the URL below manually." + ) + + console.print(Panel(auth_panel_text, title=f"Antigravity OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue")) + console.print(f"[bold]URL:[/bold] [link={auth_url}]{auth_url}[/link]\n") + + if not is_headless: + try: + webbrowser.open(auth_url) + lib_logger.info("Browser opened successfully for OAuth flow") + except Exception as e: + lib_logger.warning(f"Failed to open browser automatically: {e}. Please open the URL manually.") + + with console.status("[bold green]Waiting for you to complete authentication in the browser...[/bold green]", spinner="dots"): + auth_code = await asyncio.wait_for(auth_code_future, timeout=300) + except asyncio.TimeoutError: + raise Exception("OAuth flow timed out. Please try again.") + finally: + if server: + server.close() + await server.wait_closed() + + lib_logger.info(f"Attempting to exchange authorization code for tokens...") + async with httpx.AsyncClient() as client: + response = await client.post(TOKEN_URI, data={ + "code": auth_code.strip(), + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "redirect_uri": "http://localhost:8085/oauth2callback", + "grant_type": "authorization_code" + }) + response.raise_for_status() + token_data = response.json() + + creds = token_data.copy() + creds["expiry_date"] = (time.time() + creds.pop("expires_in")) * 1000 + creds["client_id"] = CLIENT_ID + creds["client_secret"] = CLIENT_SECRET + creds["token_uri"] = TOKEN_URI + creds["universe_domain"] = "googleapis.com" + + # Fetch user info + user_info_response = await client.get( + USER_INFO_URI, + headers={"Authorization": f"Bearer {creds['access_token']}"} + ) + user_info_response.raise_for_status() + user_info = user_info_response.json() + + creds["_proxy_metadata"] = { + "email": user_info.get("email"), + "last_check_timestamp": time.time() + } + + if path: + await self._save_credentials(path, creds) + + lib_logger.info(f"Antigravity OAuth initialized successfully for '{display_name}'.") + return creds + + lib_logger.info(f"Antigravity OAuth token at '{display_name}' is valid.") + return creds + except Exception as e: + raise ValueError(f"Failed to initialize Antigravity OAuth for '{display_name}': {e}") + + async def get_valid_token(self, credential_path: str) -> str: + """ + Get a valid access token, refreshing if necessary. + + Args: + credential_path: Path to credential file or "env" for environment variables + + Returns: + Valid access token string + + Raises: + ValueError: If token cannot be obtained + """ + try: + creds = await self.initialize_token(credential_path) + return creds['access_token'] + except Exception as e: + raise ValueError(f"Failed to get valid Antigravity token: {e}") diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py new file mode 100644 index 00000000..79e21516 --- /dev/null +++ b/src/rotator_library/providers/antigravity_provider.py @@ -0,0 +1,869 @@ +# src/rotator_library/providers/antigravity_provider.py + +import json +import httpx +import logging +import time +import asyncio +import random +import uuid +import copy +from typing import List, Dict, Any, AsyncGenerator, Union, Optional, Tuple +from .provider_interface import ProviderInterface +from .antigravity_auth_base import AntigravityAuthBase +from ..model_definitions import ModelDefinitions +import litellm +from litellm.exceptions import RateLimitError +from litellm.llms.vertex_ai.common_utils import _build_vertex_schema + +lib_logger = logging.getLogger('rotator_library') + +# Antigravity base URLs with fallback order +# Priority: daily (sandbox) → autopush (sandbox) → production +BASE_URLS = [ + "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal", + "https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal", + "https://cloudcode-pa.googleapis.com/v1internal" # Production fallback +] + +# Hardcoded models available via Antigravity +HARDCODED_MODELS = [ + "gemini-2.5-pro", + "gemini-2.5-flash", + "gemini-2.5-flash-lite", + "gemini-3-pro-preview", + "gemini-3-pro-image-preview", + "gemini-2.5-computer-use-preview-10-2025" +] + + +class AntigravityProvider(AntigravityAuthBase, ProviderInterface): + """ + Antigravity provider implementation for Gemini models. + + Antigravity is an experimental internal Google API that provides access to Gemini models + including Gemini 3 with thinking/reasoning capabilities. It wraps standard Gemini API + requests with additional metadata and uses sandbox endpoints. + + Key features: + - Model aliasing (gemini-3-pro-high ↔ gemini-3-pro-preview) + - Gemini 3 thinkingLevel support + - Thinking signature preservation for multi-turn conversations + - Sophisticated tool response grouping + - Base URL fallback (sandbox → production) + """ + skip_cost_calculation = True + + def __init__(self): + super().__init__() + self.model_definitions = ModelDefinitions() + self._current_base_url = BASE_URLS[0] # Start with daily sandbox + self._base_url_index = 0 + + # ============================================================================ + # MODEL ALIAS SYSTEM + # ============================================================================ + + def _model_name_to_alias(self, model_name: str) -> str: + """ + Convert internal Antigravity model names to public aliases. + + Args: + model_name: Internal model name + + Returns: + Public alias name, or empty string if model should be excluded + """ + alias_map = { + "rev19-uic3-1p": "gemini-2.5-computer-use-preview-10-2025", + "gemini-3-pro-image": "gemini-3-pro-image-preview", + "gemini-3-pro-high": "gemini-3-pro-preview", + "claude-sonnet-4-5": "gemini-claude-sonnet-4-5", + "claude-sonnet-4-5-thinking": "gemini-claude-sonnet-4-5-thinking", + } + + # Filter out excluded models (return empty string to skip) + excluded = [ + "chat_20706", "chat_23310", "gemini-2.5-flash-thinking", + "gemini-3-pro-low", "gemini-2.5-pro" + ] + if model_name in excluded: + return "" + + return alias_map.get(model_name, model_name) + + def _alias_to_model_name(self, alias: str) -> str: + """ + Convert public aliases to internal Antigravity model names. + + Args: + alias: Public alias name + + Returns: + Internal model name + """ + reverse_map = { + "gemini-2.5-computer-use-preview-10-2025": "rev19-uic3-1p", + "gemini-3-pro-image-preview": "gemini-3-pro-image", + "gemini-3-pro-preview": "gemini-3-pro-high", + "gemini-claude-sonnet-4-5": "claude-sonnet-4-5", + "gemini-claude-sonnet-4-5-thinking": "claude-sonnet-4-5-thinking", + } + return reverse_map.get(alias, alias) + + # ============================================================================ + # RANDOM ID GENERATION + # ============================================================================ + + @staticmethod + def generate_request_id() -> str: + """Generate Antigravity request ID: agent-{uuid}""" + return f"agent-{uuid.uuid4()}" + + @staticmethod + def generate_session_id() -> str: + """Generate Antigravity session ID: -{random_number}""" + # Generate random 19-digit number + n = random.randint(1_000_000_000_000_000_000, 9_999_999_999_999_999_999) + return f"-{n}" + + @staticmethod + def generate_project_id() -> str: + """Generate fake project ID: {adj}-{noun}-{random}""" + adjectives = ["useful", "bright", "swift", "calm", "bold"] + nouns = ["fuze", "wave", "spark", "flow", "core"] + adj = random.choice(adjectives) + noun = random.choice(nouns) + random_part = str(uuid.uuid4())[:5].lower() + return f"{adj}-{noun}-{random_part}" + + # ============================================================================ + # MESSAGE TRANSFORMATION (OpenAI → Gemini CLI format) + # ============================================================================ + + def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: + """ + Transform OpenAI messages to Gemini CLI format. + Reused from GeminiCliProvider with modifications for Antigravity. + + Returns: + Tuple of (system_instruction, gemini_contents) + """ + system_instruction = None + gemini_contents = [] + + # Make a copy to avoid modifying original + messages = copy.deepcopy(messages) + + # Separate system prompt from other messages + if messages and messages[0].get('role') == 'system': + system_prompt_content = messages.pop(0).get('content', '') + if system_prompt_content: + system_instruction = { + "role": "user", + "parts": [{"text": system_prompt_content}] + } + + # Build tool call ID to name mapping + tool_call_id_to_name = {} + for msg in messages: + if msg.get("role") == "assistant" and msg.get("tool_calls"): + for tool_call in msg["tool_calls"]: + if tool_call.get("type") == "function": + tool_call_id_to_name[tool_call["id"]] = tool_call["function"]["name"] + + # Convert each message + for msg in messages: + role = msg.get("role") + content = msg.get("content") + parts = [] + gemini_role = "model" if role == "assistant" else "tool" if role == "tool" else "user" + + if role == "user": + if isinstance(content, str): + # Simple text content + if content: + parts.append({"text": content}) + elif isinstance(content, list): + # Multi-part content (text, images, etc.) + for item in content: + if item.get("type") == "text": + text = item.get("text", "") + if text: + parts.append({"text": text}) + elif item.get("type") == "image_url": + # Handle image data URLs + image_url = item.get("image_url", {}).get("url", "") + if image_url.startswith("data:"): + try: + # Parse: data:image/png;base64,iVBORw0KG... + header, data = image_url.split(",", 1) + mime_type = header.split(":")[1].split(";")[0] + parts.append({ + "inlineData": { + "mimeType": mime_type, + "data": data + } + }) + except Exception as e: + lib_logger.warning(f"Failed to parse image data URL: {e}") + + elif role == "assistant": + if isinstance(content, str) and content: + parts.append({"text": content}) + if msg.get("tool_calls"): + for tool_call in msg["tool_calls"]: + if tool_call.get("type") == "function": + try: + args_dict = json.loads(tool_call["function"]["arguments"]) + except (json.JSONDecodeError, TypeError): + args_dict = {} + + # Add function call part with thoughtSignature + func_call_part = { + "functionCall": { + "name": tool_call["function"]["name"], + "args": args_dict + }, + "thoughtSignature": "skip_thought_signature_validator" + } + parts.append(func_call_part) + + elif role == "tool": + tool_call_id = msg.get("tool_call_id") + function_name = tool_call_id_to_name.get(tool_call_id) + if function_name: + # Wrap the tool response in a 'result' object + response_content = {"result": content} + parts.append({"functionResponse": {"name": function_name, "response": response_content}}) + + if parts: + gemini_contents.append({"role": gemini_role, "parts": parts}) + + # Ensure first message is from user + if not gemini_contents or gemini_contents[0]['role'] != 'user': + gemini_contents.insert(0, {"role": "user", "parts": [{"text": ""}]}) + + return system_instruction, gemini_contents + + # ============================================================================ + # THINKING/REASONING CONFIGURATION + # ============================================================================ + + def _map_reasoning_effort_to_thinking_config( + self, + reasoning_effort: Optional[str], + model: str + ) -> Optional[Dict[str, Any]]: + """ + Map OpenAI reasoning_effort to Gemini thinking configuration. + Handles Gemini 3 thinkingLevel vs other models thinkingBudget. + + Args: + reasoning_effort: OpenAI reasoning_effort value + model: Model name (public alias) + + Returns: + Dictionary with thinkingConfig or None + """ + internal_model = self._alias_to_model_name(model) + is_gemini_3 = internal_model.startswith("gemini-3-") + + # Default for gemini-3-pro-preview when no reasoning_effort specified + if not reasoning_effort: + if model == "gemini-3-pro-preview" or internal_model == "gemini-3-pro-high": + return { + "thinkingBudget": -1, + "include_thoughts": True + } + return None + + if reasoning_effort == "none": + return { + "thinkingBudget": 0, + "include_thoughts": False + } + + if reasoning_effort == "auto": + # Auto always uses thinkingBudget=-1, even for Gemini 3 + return { + "thinkingBudget": -1, + "include_thoughts": True + } + + if is_gemini_3: + # Gemini 3: Use thinkingLevel + level_map = { + "low": "low", + "medium": "high", # Medium not released yet, map to high + "high": "high" + } + level = level_map.get(reasoning_effort, "high") + return { + "thinkingLevel": level, + "include_thoughts": True + } + else: + # Non-Gemini-3: Use thinkingBudget with normalization + budget_map = { + "low": 1024, + "medium": 8192, + "high": 32768 + } + budget = budget_map.get(reasoning_effort, -1) + # TODO: Add model-specific normalization via model registry + return { + "thinkingBudget": budget, + "include_thoughts": True + } + + # ============================================================================ + # TOOL RESPONSE GROUPING + # ============================================================================ + + def _fix_tool_response_grouping(self, contents: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Group function calls with their responses for Antigravity compatibility. + + Converts linear format (function call, response, function call, response) + to grouped format (model with calls, function role with all responses). + + Args: + contents: List of Gemini content objects + + Returns: + List of grouped content objects + """ + new_contents = [] + pending_groups = [] # Groups awaiting responses + collected_responses = [] # Standalone responses to match + + for content in contents: + role = content.get("role") + parts = content.get("parts", []) + + # Check if this content has function responses + response_parts = [p for p in parts if "functionResponse" in p] + + if response_parts: + # Collect responses + collected_responses.extend(response_parts) + + # Try to satisfy pending groups + for i in range(len(pending_groups) - 1, -1, -1): + group = pending_groups[i] + if len(collected_responses) >= group["responses_needed"]: + # Take needed responses + group_responses = collected_responses[:group["responses_needed"]] + collected_responses = collected_responses[group["responses_needed"]:] + + # Create merged function response content + function_response_content = { + "parts": group_responses, + "role": "function" # Changed from tool + } + new_contents.append(function_response_content) + + # Remove satisfied group + pending_groups.pop(i) + break + + continue # Skip adding this content + + # If this is model content with function calls, create a group + if role == "model": + function_calls = [p for p in parts if "functionCall" in p] + + if function_calls: + # Add model content first + new_contents.append(content) + + # Create pending group + pending_groups.append({ + "model_content": content, + "function_calls": function_calls, + "responses_needed": len(function_calls) + }) + else: + # Regular model content without function calls + new_contents.append(content) + else: + # Non-model content (user, etc.) + new_contents.append(content) + + # Handle remaining pending groups + for group in pending_groups: + if len(collected_responses) >= group["responses_needed"]: + group_responses = collected_responses[:group["responses_needed"]] + collected_responses = collected_responses[group["responses_needed"]:] + + function_response_content = { + "parts": group_responses, + "role": "function" + } + new_contents.append(function_response_content) + + return new_contents + + # ============================================================================ + # ANTIGRAVITY REQUEST TRANSFORMATION + # ============================================================================ + + def _transform_to_antigravity_format( + self, + gemini_cli_payload: Dict[str, Any], + model: str + ) -> Dict[str, Any]: + """ + Transform Gemini CLI format to complete Antigravity format. + + Args: + gemini_cli_payload: Request in Gemini CLI format + model: Model name (public alias) + + Returns: + Complete Antigravity request payload + """ + internal_model = self._alias_to_model_name(model) + + # 1. Wrap in Antigravity envelope + antigravity_payload = { + "project": self.generate_project_id(), + "userAgent": "antigravity", + "requestId": self.generate_request_id(), + "model": internal_model, # Use internal name + "request": copy.deepcopy(gemini_cli_payload) + } + + # 2. Add session ID + antigravity_payload["request"]["sessionId"] = self.generate_session_id() + + # 3. Remove fields that Antigravity doesn't support + antigravity_payload["request"].pop("safetySettings", None) + if "generationConfig" in antigravity_payload["request"]: + antigravity_payload["request"]["generationConfig"].pop("maxOutputTokens", None) + + # 4. Set toolConfig mode + if "toolConfig" not in antigravity_payload["request"]: + antigravity_payload["request"]["toolConfig"] = {} + if "functionCallingConfig" not in antigravity_payload["request"]["toolConfig"]: + antigravity_payload["request"]["toolConfig"]["functionCallingConfig"] = {} + antigravity_payload["request"]["toolConfig"]["functionCallingConfig"]["mode"] = "VALIDATED" + + # 5. Handle Gemini 3 specific thinking logic + # For non-Gemini-3 models, convert thinkingLevel to thinkingBudget + if not internal_model.startswith("gemini-3-"): + gen_config = antigravity_payload["request"].get("generationConfig", {}) + thinking_config = gen_config.get("thinkingConfig", {}) + if "thinkingLevel" in thinking_config: + # Remove thinkingLevel for non-Gemini-3 models + del thinking_config["thinkingLevel"] + # Set thinkingBudget to -1 (auto/dynamic) + thinking_config["thinkingBudget"] = -1 + + # 6. Preserve/add thoughtSignature to ALL function calls in model role content + for content in antigravity_payload["request"].get("contents", []): + if content.get("role") == "model": + for part in content.get("parts", []): + # Add signature to function calls OR preserve if already exists + if "functionCall" in part and "thoughtSignature" not in part: + part["thoughtSignature"] = "skip_thought_signature_validator" + # If thoughtSignature already exists, preserve it (important for Gemini 3) + + # 7. Handle Claude models (special tool schema conversion) + if internal_model.startswith("claude-sonnet-"): + # For Claude models, convert parametersJsonSchema back to parameters + for tool in antigravity_payload["request"].get("tools", []): + for func_decl in tool.get("functionDeclarations", []): + if "parametersJsonSchema" in func_decl: + func_decl["parameters"] = func_decl.pop("parametersJsonSchema") + # Remove $schema if present + if "parameters" in func_decl and "$schema" in func_decl["parameters"]: + del func_decl["parameters"]["$schema"] + + return antigravity_payload + + #============================================================================ + # BASE URL FALLBACK LOGIC + # ============================================================================ + + def _get_current_base_url(self) -> str: + """Get the current base URL from the fallback list.""" + return self._current_base_url + + def _try_next_base_url(self) -> bool: + """ + Switch to the next base URL in the fallback list. + + Returns: + True if successfully switched to next URL, False if no more URLs available + """ + if self._base_url_index < len(BASE_URLS) - 1: + self._base_url_index += 1 + self._current_base_url = BASE_URLS[self._base_url_index] + lib_logger.info(f"Switching to fallback Antigravity base URL: {self._current_base_url}") + return True + return False + + def _reset_base_url(self): + """Reset to the primary base URL (daily sandbox).""" + self._base_url_index = 0 + self._current_base_url = BASE_URLS[0] + + # ============================================================================ + # RESPONSE TRANSFORMATION (Antigravity → OpenAI) + # ============================================================================ + + def _unwrap_antigravity_response(self, antigravity_response: Dict[str, Any]) -> Dict[str, Any]: + """ + Extract Gemini response from Antigravity envelope. + + Args: + antigravity_response: Response from Antigravity API + + Returns: + Gemini response (unwrapped) + """ + # For both streaming and non-streaming, response is in 'response' field + return antigravity_response.get("response", antigravity_response) + + def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> litellm.ModelResponse: + """ + Convert a single Gemini response chunk to OpenAI format. + Based on GeminiCliProvider logic. + + Args: + gemini_chunk: Gemini response chunk + model: Model name + + Returns: + OpenAI-format ModelResponse + """ + # Extract candidate + candidates = gemini_chunk.get("candidates", []) + if not candidates: + # Empty chunk, return minimal response + return litellm.ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=int(time.time()), + model=model, + choices=[] + ) + + candidate = candidates[0] + content_parts = candidate.get("content", {}).get("parts", []) + + # Extract text, tool calls, and thinking + text_content = "" + tool_calls = [] + + for part in content_parts: + # Extract text + if "text" in part: + text_content += part["text"] + + # Extract function calls (tool calls) + if "functionCall" in part: + func_call = part["functionCall"] + tool_calls.append({ + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": func_call.get("name", ""), + "arguments": json.dumps(func_call.get("args", {})) + } + }) + + # Build delta + delta = {} + if text_content: + delta["content"] = text_content + if tool_calls: + delta["tool_calls"] = tool_calls + + # Get finish reason + finish_reason = candidate.get("finishReason", "").lower() if candidate.get("finishReason") else None + if finish_reason == "stop": + finish_reason = "stop" + elif finish_reason == "max_tokens": + finish_reason = "length" + + # Build choice + choice = { + "index": 0, + "delta": delta, + "finish_reason": finish_reason + } + + # Extract usage (if present) + usage_metadata = gemini_chunk.get("usageMetadata", {}) + usage = None + if usage_metadata: + usage = { + "prompt_tokens": usage_metadata.get("promptTokenCount", 0), + "completion_tokens": usage_metadata.get("candidatesTokenCount", 0), + "total_tokens": usage_metadata.get("totalTokenCount", 0) + } + + return litellm.ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=int(time.time()), + model=model, + choices=[choice], + usage=usage + ) + + # ============================================================================ + # PROVIDER INTERFACE IMPLEMENTATION + # ============================================================================ + + def has_custom_logic(self) -> bool: + """Antigravity uses custom translation logic.""" + return True + + async def get_auth_header(self, credential_identifier: str) -> Dict[str, str]: + """ + Get OAuth authorization header for Antigravity. + + Args: + credential_identifier: Credential file path or "env" + + Returns: + Dict with Authorization header + """ + access_token = await self.get_valid_token(credential_identifier) + return {"Authorization": f"Bearer {access_token}"} + + async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: + """ + Fetch available models from Antigravity. + + For Antigravity, we use the fetchAvailableModels endpoint and apply + alias mapping to convert internal names to public names. + + Args: + api_key: Credential path (not a traditional API key) + client: HTTP client + + Returns: + List of public model names + """ + credential_path = api_key # For OAuth providers, this is the credential path + + try: + access_token = await self.get_valid_token(credential_path) + base_url = self._get_current_base_url() + + # Generate required IDs + project_id = self.generate_project_id() + request_id = self.generate_request_id() + + # Fetch models endpoint + url = f"{base_url}/fetchAvailableModels" + + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json" + } + + payload = { + "project": project_id, + "requestId": request_id, + "userAgent": "antigravity" + } + + lib_logger.debug(f"Fetching Antigravity models from: {url}") + + response = await client.post(url, json=payload, headers=headers, timeout=30.0) + response.raise_for_status() + + data = response.json() + + # Extract model names and apply aliasing + models = [] + if "models" in data: + for model_info in data["models"]: + internal_name = model_info.get("name", "").replace("models/", "") + if internal_name: + public_name = self._model_name_to_alias(internal_name) + if public_name: # Skip excluded models (empty string) + models.append(public_name) + + if models: + lib_logger.info(f"Discovered {len(models)} Antigravity models") + return models + else: + lib_logger.warning("No models returned from Antigravity, using hardcoded list") + return HARDCODED_MODELS + + except Exception as e: + lib_logger.warning(f"Failed to fetch Antigravity models: {e}, using hardcoded list") + return HARDCODED_MODELS + + async def acompletion( + self, + client: httpx.AsyncClient, + **kwargs + ) -> Union[litellm.ModelResponse, AsyncGenerator[litellm.ModelResponse, None]]: + """ + Handle completion requests for Antigravity. + + This is the main entry point that: + 1. Extracts the model and credential path + 2. Transforms OpenAI request → Gemini CLI → Antigravity format + 3. Makes the API call with fallback logic + 4. Transforms Antigravity response → Gemini → OpenAI format + + Args: + client: HTTP client + **kwargs: LiteLLM completion parameters + + Returns: + ModelResponse (non-streaming) or AsyncGenerator (streaming) + """ + # Extract key parameters + model = kwargs.get("model", "gemini-2.5-pro") + messages = kwargs.get("messages", []) + stream = kwargs.get("stream", False) + credential_path = kwargs.pop("credential_identifier", kwargs.get("api_key", "")) + tools = kwargs.get("tools") + reasoning_effort = kwargs.get("reasoning_effort") + temperature = kwargs.get("temperature") + top_p = kwargs.get("top_p") + max_tokens = kwargs.get("max_tokens") + + lib_logger.info(f"Antigravity completion: model={model}, stream={stream}, messages={len(messages)}") + + # Step 1: Transform messages (OpenAI → Gemini CLI) + system_instruction, gemini_contents = self._transform_messages(messages) + + # Apply tool response grouping + gemini_contents = self._fix_tool_response_grouping(gemini_contents) + + # Step 2: Build Gemini CLI payload + gemini_cli_payload = { + "contents": gemini_contents + } + + if system_instruction: + gemini_cli_payload["system_instruction"] = system_instruction + + # Add generation config + generation_config = {} + if temperature is not None: + generation_config["temperature"] = temperature + if top_p is not None: + generation_config["topP"] = top_p + + # Handle thinking config + thinking_config = self._map_reasoning_effort_to_thinking_config(reasoning_effort, model) + if thinking_config: + generation_config.setdefault("thinkingConfig", {}).update(thinking_config) + + if generation_config: + gemini_cli_payload["generationConfig"] = generation_config + + # Add tools + if tools: + gemini_tools = [] + for tool in tools: + if tool.get("type") == "function": + func = tool.get("function", {}) + schema = _build_vertex_schema(parameters=func.get("parameters", {})) + gemini_tools.append({ + "functionDeclarations": [{ + "name": func.get("name", ""), + "description": func.get("description", ""), + "parametersJsonSchema": schema + }] + }) + if gemini_tools: + gemini_cli_payload["tools"] = gemini_tools + + # Step 3: Transform to Antigravity format + antigravity_payload = self._transform_to_antigravity_format(gemini_cli_payload, model) + + # Step 4: Make API call + access_token = await self.get_valid_token(credential_path) + base_url = self._get_current_base_url() + + endpoint = ":streamGenerateContent" if stream else ":generateContent" + url = f"{base_url}{endpoint}" + + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json" + } + + lib_logger.debug(f"Antigravity request to: {url}") + + try: + if stream: + return self._handle_streaming(client, url, headers, antigravity_payload, model) + else: + return await self._handle_non_streaming(client, url, headers, antigravity_payload, model) + except Exception as e: + # Try fallback URL if available + if self._try_next_base_url(): + lib_logger.warning(f"Retrying Antigravity request with fallback URL: {e}") + base_url = self._get_current_base_url() + url = f"{base_url}{endpoint}" + + if stream: + return self._handle_streaming(client, url, headers, antigravity_payload, model) + else: + return await self._handle_non_streaming(client, url, headers, antigravity_payload, model) + else: + raise + + async def _handle_non_streaming( + self, + client: httpx.AsyncClient, + url: str, + headers: Dict[str, str], + payload: Dict[str, Any], + model: str + ) -> litellm.ModelResponse: + """Handle non-streaming completion.""" + response = await client.post(url, headers=headers, json=payload, timeout=120.0) + response.raise_for_status() + + antigravity_response = response.json() + + # Unwrap Antigravity envelope + gemini_response = self._unwrap_antigravity_response(antigravity_response) + + # Convert to OpenAI format + return self._gemini_to_openai_chunk(gemini_response, model) + + async def _handle_streaming( + self, + client: httpx.AsyncClient, + url: str, + headers: Dict[str, str], + payload: Dict[str, Any], + model: str + ) -> AsyncGenerator[litellm.ModelResponse, None]: + """Handle streaming completion.""" + async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response: + response.raise_for_status() + + async for line in response.aiter_lines(): + if line.startswith("data: "): + data_str = line[6:] + if data_str == "[DONE]": + break + + try: + antigravity_chunk = json.loads(data_str) + + # Unwrap Antigravity envelope + gemini_chunk = self._unwrap_antigravity_response(antigravity_chunk) + + # Convert to OpenAI format + openai_chunk = self._gemini_to_openai_chunk(gemini_chunk, model) + + yield openai_chunk + except json.JSONDecodeError: + lib_logger.warning(f"Failed to parse Antigravity chunk: {data_str[:100]}") + continue diff --git a/src/rotator_library/pyproject.toml b/src/rotator_library/pyproject.toml index a8dacd37..4cfa41a3 100644 --- a/src/rotator_library/pyproject.toml +++ b/src/rotator_library/pyproject.toml @@ -3,8 +3,8 @@ requires = ["setuptools>=61.0"] build-backend = "setuptools.build_meta" [project] -name = "rotating-api-key-client" -version = "0.9" +name = "rotator_library" +version = "0.95" authors = [ { name="Mirrowel", email="nuh@uh.com" }, ] From 34cb9f83a3f06fd7475c69882c99ff945b3d8fa5 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 14:17:31 +0100 Subject: [PATCH 002/221] =?UTF-8?q?feat(providers):=20=E2=9C=A8=20add=20Ge?= =?UTF-8?q?mini=203=20thoughtSignature=20handling=20and=20reasoning=5Fcont?= =?UTF-8?q?ent=20separation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Introduce Gemini 3 special mechanics in AntigravityProvider: - append a constant thoughtSignature into functionCall payloads to preserve Gemini reasoning continuity - filter out thoughtSignature parts from returned content to avoid exposing encrypted reasoning data - separate parts flagged with thought=true into a new reasoning_content field while keeping regular content in content - include thoughtsTokenCount in token accounting: prompt_tokens now includes reasoning tokens and reasoning_tokens are reported under completion_tokens_details.reasoning_tokens when present - Update comments, docstrings, and conversion logic to reflect Gemini 3 behavior - Rotate Antigravity OAuth client secret in AntigravityAuthBase --- .../providers/antigravity_auth_base.py | 2 +- .../providers/antigravity_provider.py | 59 ++++++++++++++++--- 2 files changed, 53 insertions(+), 8 deletions(-) diff --git a/src/rotator_library/providers/antigravity_auth_base.py b/src/rotator_library/providers/antigravity_auth_base.py index 14b470f5..df15dae9 100644 --- a/src/rotator_library/providers/antigravity_auth_base.py +++ b/src/rotator_library/providers/antigravity_auth_base.py @@ -23,7 +23,7 @@ # Antigravity OAuth credentials from CLIProxyAPI CLIENT_ID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" -CLIENT_SECRET = "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" +CLIENT_SECRET = "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" TOKEN_URI = "https://oauth2.googleapis.com/token" USER_INFO_URI = "https://www.googleapis.com/oauth2/v1/userinfo" REFRESH_EXPIRY_BUFFER_SECONDS = 30 * 60 # 30 minutes buffer before expiry diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 79e21516..d1833021 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -48,9 +48,19 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface): Key features: - Model aliasing (gemini-3-pro-high ↔ gemini-3-pro-preview) - Gemini 3 thinkingLevel support - - Thinking signature preservation for multi-turn conversations + - ThoughtSignature preservation for multi-turn conversations + - Reasoning content separation (thought=true parts) - Sophisticated tool response grouping - Base URL fallback (sandbox → production) + + Gemini 3 Special Mechanics: + 1. ThinkingLevel: Uses thinkingLevel (low/high) instead of thinkingBudget for Gemini 3 models + 2. ThoughtSignature: Function calls include thoughtSignature="skip_thought_signature_validator" + - This is a CONSTANT validation bypass flag, not a session key + - Preserved across conversation turns to maintain reasoning continuity + - Filtered from responses to prevent exposing encrypted internal data + 3. Reasoning Content: Text parts with thought=true flag are separated into reasoning_content + 4. Token Counting: thoughtsTokenCount is included in prompt_tokens and reported as reasoning_tokens """ skip_cost_calculation = True @@ -220,6 +230,9 @@ def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[ args_dict = {} # Add function call part with thoughtSignature + # ThoughtSignature is required for Gemini to process function calls correctly + # The constant "skip_thought_signature_validator" tells Gemini to bypass signature validation + # This is preserved across conversation turns to maintain reasoning continuity func_call_part = { "functionCall": { "name": tool_call["function"]["name"], @@ -530,7 +543,11 @@ def _unwrap_antigravity_response(self, antigravity_response: Dict[str, Any]) -> def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> litellm.ModelResponse: """ Convert a single Gemini response chunk to OpenAI format. - Based on GeminiCliProvider logic. + + Handles Gemini 3 special mechanics: + - Filters thoughtSignature parts (encrypted reasoning data) + - Separates reasoning content (thought=true) from regular content + - Includes thoughtsTokenCount in usage metadata Args: gemini_chunk: Gemini response chunk @@ -553,14 +570,27 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> l candidate = candidates[0] content_parts = candidate.get("content", {}).get("parts", []) - # Extract text, tool calls, and thinking + # Extract text, tool calls, and reasoning content text_content = "" + reasoning_content = "" tool_calls = [] for part in content_parts: - # Extract text + # CRITICAL: Skip parts with thoughtSignature (encrypted reasoning data) + # This prevents exposing internal Gemini reasoning signatures to clients + if "thoughtSignature" in part and part["thoughtSignature"]: + continue + + # Extract text - separate regular content from reasoning/thinking if "text" in part: - text_content += part["text"] + # Check for thought flag (Gemini 3 reasoning indicator) + thought = part.get("thought") + if thought is True or (isinstance(thought, str) and thought.lower() == 'true'): + # This is reasoning/thinking content + reasoning_content += part["text"] + else: + # Regular content + text_content += part["text"] # Extract function calls (tool calls) if "functionCall" in part: @@ -578,6 +608,9 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> l delta = {} if text_content: delta["content"] = text_content + if reasoning_content: + # OpenAI o1-style reasoning content field + delta["reasoning_content"] = reasoning_content if tool_calls: delta["tool_calls"] = tool_calls @@ -599,11 +632,23 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> l usage_metadata = gemini_chunk.get("usageMetadata", {}) usage = None if usage_metadata: + # Get token counts + prompt_tokens = usage_metadata.get("promptTokenCount", 0) + thoughts_tokens = usage_metadata.get("thoughtsTokenCount", 0) + completion_tokens = usage_metadata.get("candidatesTokenCount", 0) + + # OpenAI o1-style token counting: thoughts are included in prompt_tokens usage = { - "prompt_tokens": usage_metadata.get("promptTokenCount", 0), - "completion_tokens": usage_metadata.get("candidatesTokenCount", 0), + "prompt_tokens": prompt_tokens + thoughts_tokens, + "completion_tokens": completion_tokens, "total_tokens": usage_metadata.get("totalTokenCount", 0) } + + # Add reasoning tokens details if thinking was used + if thoughts_tokens > 0: + if "completion_tokens_details" not in usage: + usage["completion_tokens_details"] = {} + usage["completion_tokens_details"]["reasoning_tokens"] = thoughts_tokens return litellm.ModelResponse( id=f"chatcmpl-{uuid.uuid4()}", From 7c758a6f4939981d70b43f02ee5ae43a03db2802 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 14:34:31 +0100 Subject: [PATCH 003/221] feat(providers): add Antigravity file logging, reasoning mapping and token counting Add a per-request file logger and reasoning configuration mapping to the Antigravity provider and expose a token counting helper. - Introduce _AntigravityFileLogger to persist request payloads, streaming chunks, errors, and final responses under logs/antigravity_logs with timestamped directories. - Add optional enable_request_logging kwarg to completion flow to enable per-call file logging; wire logger through streaming and non-streaming handlers. - Log request payloads, raw response chunks, parse errors, and final unwrapped responses when enabled. - Add _map_reasoning_effort_to_thinking_config to map reasoning_effort ('low'|'medium'|'high'|'disable'|None) to Gemini thinkingConfig for gemini-2.5 and gemini-3 families (budgets/levels and include_thoughts). - Add count_tokens method that calls Antigravity :countTokens endpoint using transformed Gemini payloads and returns prompt/total token counts. - Add cautionary comment about Claude parametersJsonSchema handling requiring investigation. No behavioral breaking changes; new logging is opt-in via enable_request_logging and token counting is additive. --- .../providers/antigravity_provider.py | 271 +++++++++++++++++- 1 file changed, 266 insertions(+), 5 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index d1833021..5ab0db9d 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -8,6 +8,8 @@ import random import uuid import copy +from pathlib import Path +from datetime import datetime from typing import List, Dict, Any, AsyncGenerator, Union, Optional, Tuple from .provider_interface import ProviderInterface from .antigravity_auth_base import AntigravityAuthBase @@ -36,6 +38,64 @@ "gemini-2.5-computer-use-preview-10-2025" ] +# Logging configuration +LOGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "logs" +ANTIGRAVITY_LOGS_DIR = LOGS_DIR / "antigravity_logs" + + +class _AntigravityFileLogger: + """A simple file logger for a single Antigravity transaction.""" + def __init__(self, model_name: str, enabled: bool = True): + self.enabled = enabled + if not self.enabled: + return + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + request_id = str(uuid.uuid4()) + # Sanitize model name for directory + safe_model_name = model_name.replace('/', '_').replace(':', '_') + self.log_dir = ANTIGRAVITY_LOGS_DIR / f"{timestamp}_{safe_model_name}_{request_id}" + try: + self.log_dir.mkdir(parents=True, exist_ok=True) + except Exception as e: + lib_logger.error(f"Failed to create Antigravity log directory: {e}") + self.enabled = False + + def log_request(self, payload: Dict[str, Any]): + """Logs the request payload sent to Antigravity.""" + if not self.enabled: return + try: + with open(self.log_dir / "request_payload.json", "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, ensure_ascii=False) + except Exception as e: + lib_logger.error(f"_AntigravityFileLogger: Failed to write request: {e}") + + def log_response_chunk(self, chunk: str): + """Logs a raw chunk from the Antigravity response stream.""" + if not self.enabled: return + try: + with open(self.log_dir / "response_stream.log", "a", encoding="utf-8") as f: + f.write(chunk + "\n") + except Exception as e: + lib_logger.error(f"_AntigravityFileLogger: Failed to write response chunk: {e}") + + def log_error(self, error_message: str): + """Logs an error message.""" + if not self.enabled: return + try: + with open(self.log_dir / "error.log", "a", encoding="utf-8") as f: + f.write(f"[{datetime.utcnow().isoformat()}] {error_message}\n") + except Exception as e: + lib_logger.error(f"_AntigravityFileLogger: Failed to write error: {e}") + + def log_final_response(self, response_data: Dict[str, Any]): + """Logs the final, reassembled response.""" + if not self.enabled: return + try: + with open(self.log_dir / "final_response.json", "w", encoding="utf-8") as f: + json.dump(response_data, f, indent=2, ensure_ascii=False) + except Exception as e: + lib_logger.error(f"_AntigravityFileLogger: Failed to write final response: {e}") class AntigravityProvider(AntigravityAuthBase, ProviderInterface): """ @@ -418,6 +478,71 @@ def _fix_tool_response_grouping(self, contents: List[Dict[str, Any]]) -> List[Di return new_contents + + # ============================================================================ + # REASONING PARAMETER HANDLING + # ============================================================================ + + def _map_reasoning_effort_to_thinking_config( + self, + reasoning_effort: Optional[str], + model: str + ) -> Optional[Dict[str, Any]]: + """ + Map reasoning_effort parameter to thinkingConfig for Gemini models. + + This enables default thinking for Gemini 2.5 and 3 models, allowing + them to use internal reasoning/thinking capabilities. + + Args: + reasoning_effort: Optional reasoning effort level ('low', 'medium', 'high', 'disable', or None) + model: Model name (public alias) + + Returns: + thinkingConfig dict if applicable, None otherwise + """ + # Only apply to gemini-2.5 and gemini-3 model families + if "gemini-2.5" not in model and "gemini-3" not in model: + return None + + # If no reasoning_effort provided, enable default thinking (auto mode) + if reasoning_effort is None: + # For Gemini 3, use thinkingLevel + if "gemini-3" in model: + return {"thinkingLevel": 1, "include_thoughts": True} + # For Gemini 2.5, use thinkingBudget in auto mode (-1) + else: + return {"thinkingBudget": -1, "include_thoughts": True} + + # Handle explicit disable + if reasoning_effort == "disable": + if "gemini-3" in model: + return {"thinkingLevel": 0, "include_thoughts": False} + else: + return {"thinkingBudget": 0, "include_thoughts": False} + + # Map reasoning effort to budget for Gemini 2.5 + if "gemini-2.5" in model: + if "gemini-2.5-pro" in model: + budgets = {"low": 8192, "medium": 16384, "high": 32768} + elif "gemini-2.5-flash" in model: + budgets = {"low": 6144, "medium": 12288, "high": 24576} + else: + # Fallback for other gemini-2.5 models + budgets = {"low": 1024, "medium": 2048, "high": 4096} + + budget = budgets.get(reasoning_effort, -1) # -1 = auto for invalid values + # Note: Not dividing by 4 like Gemini CLI does, using full budget + return {"thinkingBudget": budget, "include_thoughts": True} + + # For Gemini 3, map to thinkingLevel + if "gemini-3" in model: + levels = {"low": 1, "medium": 2, "high": 3} + level = levels.get(reasoning_effort, 1) # Default to level 1 + return {"thinkingLevel": level, "include_thoughts": True} + + return None + # ============================================================================ # ANTIGRAVITY REQUEST TRANSFORMATION # ============================================================================ @@ -483,7 +608,21 @@ def _transform_to_antigravity_format( part["thoughtSignature"] = "skip_thought_signature_validator" # If thoughtSignature already exists, preserve it (important for Gemini 3) - # 7. Handle Claude models (special tool schema conversion) + # ======================================================================== + # IMPORTANT: CLAUDE SCHEMA HANDLING - REQUIRES INVESTIGATION + # ======================================================================== + # WARNING: This code block may be incorrect! + # + # INVESTIGATION REQUIRED BEFORE MAKING CHANGES: + # - Test Claude model access through Antigravity with tools + # - Verify whether parametersJsonSchema → parameters conversion is needed + # - The Go reference suggests Antigravity expects parametersJsonSchema for ALL models + # + # Current behavior: Converts parametersJsonSchema back to parameters for Claude models + # Potential issue: Antigravity may actually expect parametersJsonSchema for Claude too + # + # DO NOT MODIFY without first confirming actual API behavior! + # ======================================================================== if internal_model.startswith("claude-sonnet-"): # For Claude models, convert parametersJsonSchema back to parameters for tool in antigravity_payload["request"].get("tools", []): @@ -776,9 +915,16 @@ async def acompletion( temperature = kwargs.get("temperature") top_p = kwargs.get("top_p") max_tokens = kwargs.get("max_tokens") + enable_request_logging = kwargs.pop("enable_request_logging", False) lib_logger.info(f"Antigravity completion: model={model}, stream={stream}, messages={len(messages)}") + # Create file logger + file_logger = _AntigravityFileLogger( + model_name=model, + enabled=enable_request_logging + ) + # Step 1: Transform messages (OpenAI → Gemini CLI) system_instruction, gemini_contents = self._transform_messages(messages) @@ -828,6 +974,9 @@ async def acompletion( # Step 3: Transform to Antigravity format antigravity_payload = self._transform_to_antigravity_format(gemini_cli_payload, model) + # Log the request + file_logger.log_request(antigravity_payload) + # Step 4: Make API call access_token = await self.get_valid_token(credential_path) base_url = self._get_current_base_url() @@ -844,9 +993,9 @@ async def acompletion( try: if stream: - return self._handle_streaming(client, url, headers, antigravity_payload, model) + return self._handle_streaming(client, url, headers, antigravity_payload, model, file_logger) else: - return await self._handle_non_streaming(client, url, headers, antigravity_payload, model) + return await self._handle_non_streaming(client, url, headers, antigravity_payload, model, file_logger) except Exception as e: # Try fallback URL if available if self._try_next_base_url(): @@ -867,7 +1016,8 @@ async def _handle_non_streaming( url: str, headers: Dict[str, str], payload: Dict[str, Any], - model: str + model: str, + file_logger: Optional[_AntigravityFileLogger] = None ) -> litellm.ModelResponse: """Handle non-streaming completion.""" response = await client.post(url, headers=headers, json=payload, timeout=120.0) @@ -875,6 +1025,10 @@ async def _handle_non_streaming( antigravity_response = response.json() + # Log response + if file_logger: + file_logger.log_final_response(antigravity_response) + # Unwrap Antigravity envelope gemini_response = self._unwrap_antigravity_response(antigravity_response) @@ -887,13 +1041,18 @@ async def _handle_streaming( url: str, headers: Dict[str, str], payload: Dict[str, Any], - model: str + model: str, + file_logger: Optional[_AntigravityFileLogger] = None ) -> AsyncGenerator[litellm.ModelResponse, None]: """Handle streaming completion.""" async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response: response.raise_for_status() async for line in response.aiter_lines(): + # Log raw chunk + if file_logger: + file_logger.log_response_chunk(line) + if line.startswith("data: "): data_str = line[6:] if data_str == "[DONE]": @@ -910,5 +1069,107 @@ async def _handle_streaming( yield openai_chunk except json.JSONDecodeError: + if file_logger: + file_logger.log_error(f"Failed to parse chunk: {data_str[:100]}") lib_logger.warning(f"Failed to parse Antigravity chunk: {data_str[:100]}") continue + + # ============================================================================ + # TOKEN COUNTING + # ============================================================================ + + async def count_tokens( + self, + client: httpx.AsyncClient, + credential_path: str, + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + litellm_params: Optional[Dict[str, Any]] = None + ) -> Dict[str, int]: + """ + Counts tokens for the given prompt using the Antigravity :countTokens endpoint. + + Args: + client: The HTTP client to use + credential_path: Path to the credential file + model: Model name to use for token counting + messages: List of messages in OpenAI format + tools: Optional list of tool definitions + litellm_params: Optional additional parameters + + Returns: + Dict with 'prompt_tokens' and 'total_tokens' counts + """ + # Get auth token + access_token = await self.get_valid_token(credential_path) + + # Convert public alias to internal name + internal_model = self._alias_to_model_name(model) + + # Transform messages to Gemini format + system_instruction, contents = self._transform_messages(messages) + + # Build Gemini CLI payload + gemini_cli_payload = { + "contents": contents + } + + if system_instruction: + gemini_cli_payload["systemInstruction"] = system_instruction + + if tools: + # Transform tools to Gemini format + gemini_tools = [] + for tool in tools: + if tool.get("type") == "function": + func = tool.get("function", {}) + schema = _build_vertex_schema(parameters=func.get("parameters", {})) + gemini_tools.append({ + "functionDeclarations": [{ + "name": func.get("name", ""), + "description": func.get("description", ""), + "parametersJsonSchema": schema + }] + }) + if gemini_tools: + gemini_cli_payload["tools"] = gemini_tools + + # Wrap in Antigravity envelope + antigravity_payload = { + "project": self.generate_project_id(), + "userAgent": "antigravity", + "requestId": self.generate_request_id(), + "model": internal_model, + "request": gemini_cli_payload + } + + # Make the request + base_url = self._get_current_base_url() + url = f"{base_url}:countTokens" + + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json" + } + + try: + response = await client.post(url, headers=headers, json=antigravity_payload, timeout=30) + response.raise_for_status() + data = response.json() + + # Unwrap Antigravity response + unwrapped = self._unwrap_antigravity_response(data) + + # Extract token counts from response + total_tokens = unwrapped.get('totalTokens', 0) + + return { + 'prompt_tokens': total_tokens, + 'total_tokens': total_tokens, + } + + except httpx.HTTPStatusError as e: + lib_logger.error(f"Failed to count tokens: {e}") + # Return 0 on error rather than raising + return {'prompt_tokens': 0, 'total_tokens': 0} From 14953252ac4453d33bc9b9106747bb434ac52cf7 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 16:07:47 +0100 Subject: [PATCH 004/221] =?UTF-8?q?feat(providers):=20=E2=9C=A8=20support?= =?UTF-8?q?=20gemini=202.5/3=20reasoning=20configs=20and=20custom=20budget?= =?UTF-8?q?=20toggle?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce a consolidated mapping for reasoning effort targeted at Gemini 2.5 and Gemini 3 models: - Replace older duplicated logic with a single _map_reasoning_effort_to_thinking_config that detects gemini-2.5 vs gemini-3. - Gemini 2.5: map reasoning_effort to model-specific thinkingBudget values (pro/flash/fallback). Default auto = -1. Apply division by 4 unless kwargs['custom_reasoning_budget'] is True. - Gemini 3: use string thinkingLevel ("low" or "high"), default to "high" when unspecified and do not allow disabling thinking. - Return None for non-Gemini models to avoid changing other providers (e.g., Claude). - Propagate a new custom_reasoning_budget toggle from kwargs to the mapping call. - Add threading and os imports and remove the old obsolete mapping implementation. BREAKING CHANGE: Gemini 3 thinkingConfig format and defaults changed: - thinkingLevel is now a string ("low"/"high") instead of numeric levels. Update any code that inspects thinkingConfig thinkingLevel. - Default thinking behavior for Gemini 3 is now "high" when reasoning_effort is omitted. - The mapping function signature/behavior changed (added custom_reasoning_budget handling). If this method was called externally, update callers to pass the new parameter or rely on kwargs propagation. --- .../providers/antigravity_provider.py | 202 ++++++++---------- 1 file changed, 87 insertions(+), 115 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 5ab0db9d..af254600 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -8,6 +8,8 @@ import random import uuid import copy +import threading +import os from pathlib import Path from datetime import datetime from typing import List, Dict, Any, AsyncGenerator, Union, Optional, Tuple @@ -320,75 +322,102 @@ def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[ return system_instruction, gemini_contents # ============================================================================ - # THINKING/REASONING CONFIGURATION + # REASONING CONFIGURATION (GEMINI 2.5 & 3 ONLY) # ============================================================================ def _map_reasoning_effort_to_thinking_config( self, reasoning_effort: Optional[str], - model: str + model: str, + custom_reasoning_budget: bool = False ) -> Optional[Dict[str, Any]]: """ - Map OpenAI reasoning_effort to Gemini thinking configuration. - Handles Gemini 3 thinkingLevel vs other models thinkingBudget. + Map reasoning_effort to thinking configuration for Gemini 2.5 and 3 models. + + IMPORTANT: This function ONLY applies to Gemini 2.5 and 3 models. + For other models (e.g., Claude via Antigravity), it returns None. + + Gemini 2.5 and 3 use separate budgeting systems: + - Gemini 2.5: thinkingBudget (integer tokens, based on Gemini CLI logic) + - Gemini 3: thinkingLevel (string: "low" or "high") + + Default behavior (no reasoning_effort): + - Gemini 2.5: thinkingBudget=-1 (auto mode) + - Gemini 3: thinkingLevel="high" (always enabled at high level) Args: - reasoning_effort: OpenAI reasoning_effort value + reasoning_effort: Effort level ('low', 'medium', 'high', 'disable', or None) model: Model name (public alias) + custom_reasoning_budget: If True, use full budgets; if False, divide by 4 Returns: - Dictionary with thinkingConfig or None + Dict with thinkingConfig or None if not a Gemini 2.5/3 model """ internal_model = self._alias_to_model_name(model) + + # Detect model family - ONLY support gemini-2.5 and gemini-3 + # For other models (Claude, etc.), return None without filtering + is_gemini_25 = "gemini-2.5" in model is_gemini_3 = internal_model.startswith("gemini-3-") - # Default for gemini-3-pro-preview when no reasoning_effort specified - if not reasoning_effort: - if model == "gemini-3-pro-preview" or internal_model == "gemini-3-pro-high": - return { - "thinkingBudget": -1, - "include_thoughts": True - } + # Return None for unsupported models - no reasoning config changes + if not is_gemini_25 and not is_gemini_3: return None - if reasoning_effort == "none": - return { - "thinkingBudget": 0, - "include_thoughts": False - } - - if reasoning_effort == "auto": - # Auto always uses thinkingBudget=-1, even for Gemini 3 - return { - "thinkingBudget": -1, - "include_thoughts": True - } + # ======================================================================== + # GEMINI 2.5: Use Gemini CLI logic with thinkingBudget + # ======================================================================== + if is_gemini_25: + # Default: auto mode + if not reasoning_effort: + return {"thinkingBudget": -1, "include_thoughts": True} + + # Disable thinking + if reasoning_effort == "disable": + return {"thinkingBudget": 0, "include_thoughts": False} + + # Model-specific budgets (same as Gemini CLI) + if "gemini-2.5-pro" in model: + budgets = {"low": 8192, "medium": 16384, "high": 32768} + elif "gemini-2.5-flash" in model: + budgets = {"low": 6144, "medium": 12288, "high": 24576} + else: + # Fallback for other gemini-2.5 models + budgets = {"low": 1024, "medium": 2048, "high": 4096} + + budget = budgets.get(reasoning_effort, -1) # -1 for invalid/auto + + # Apply custom_reasoning_budget toggle + # If False (default), divide by 4 like Gemini CLI + if not custom_reasoning_budget: + budget = budget // 4 + + return {"thinkingBudget": budget, "include_thoughts": True} + # ======================================================================== + # GEMINI 3: Use STRING thinkingLevel ("low" or "high") + # ======================================================================== if is_gemini_3: - # Gemini 3: Use thinkingLevel - level_map = { - "low": "low", - "medium": "high", # Medium not released yet, map to high - "high": "high" - } - level = level_map.get(reasoning_effort, "high") - return { - "thinkingLevel": level, - "include_thoughts": True - } - else: - # Non-Gemini-3: Use thinkingBudget with normalization - budget_map = { - "low": 1024, - "medium": 8192, - "high": 32768 - } - budget = budget_map.get(reasoning_effort, -1) - # TODO: Add model-specific normalization via model registry - return { - "thinkingBudget": budget, - "include_thoughts": True - } + # Default: Always use "high" if not specified + # Gemini 3 cannot be disabled - always has thinking enabled + if not reasoning_effort: + return {"thinkingLevel": "high", "include_thoughts": True} + + # Map reasoning effort to string level + # Note: "disable" is ignored - Gemini 3 cannot disable thinking + if reasoning_effort == "low": + level = "low" + # Medium level not yet available - map to high + # When medium is released, uncomment the following line: + # elif reasoning_effort == "medium": + # level = "medium" + else: + # "medium", "high", "disable", or any invalid value → "high" + level = "high" + + return {"thinkingLevel": level, "include_thoughts": True} + + return None # ============================================================================ # TOOL RESPONSE GROUPING @@ -478,71 +507,6 @@ def _fix_tool_response_grouping(self, contents: List[Dict[str, Any]]) -> List[Di return new_contents - - # ============================================================================ - # REASONING PARAMETER HANDLING - # ============================================================================ - - def _map_reasoning_effort_to_thinking_config( - self, - reasoning_effort: Optional[str], - model: str - ) -> Optional[Dict[str, Any]]: - """ - Map reasoning_effort parameter to thinkingConfig for Gemini models. - - This enables default thinking for Gemini 2.5 and 3 models, allowing - them to use internal reasoning/thinking capabilities. - - Args: - reasoning_effort: Optional reasoning effort level ('low', 'medium', 'high', 'disable', or None) - model: Model name (public alias) - - Returns: - thinkingConfig dict if applicable, None otherwise - """ - # Only apply to gemini-2.5 and gemini-3 model families - if "gemini-2.5" not in model and "gemini-3" not in model: - return None - - # If no reasoning_effort provided, enable default thinking (auto mode) - if reasoning_effort is None: - # For Gemini 3, use thinkingLevel - if "gemini-3" in model: - return {"thinkingLevel": 1, "include_thoughts": True} - # For Gemini 2.5, use thinkingBudget in auto mode (-1) - else: - return {"thinkingBudget": -1, "include_thoughts": True} - - # Handle explicit disable - if reasoning_effort == "disable": - if "gemini-3" in model: - return {"thinkingLevel": 0, "include_thoughts": False} - else: - return {"thinkingBudget": 0, "include_thoughts": False} - - # Map reasoning effort to budget for Gemini 2.5 - if "gemini-2.5" in model: - if "gemini-2.5-pro" in model: - budgets = {"low": 8192, "medium": 16384, "high": 32768} - elif "gemini-2.5-flash" in model: - budgets = {"low": 6144, "medium": 12288, "high": 24576} - else: - # Fallback for other gemini-2.5 models - budgets = {"low": 1024, "medium": 2048, "high": 4096} - - budget = budgets.get(reasoning_effort, -1) # -1 = auto for invalid values - # Note: Not dividing by 4 like Gemini CLI does, using full budget - return {"thinkingBudget": budget, "include_thoughts": True} - - # For Gemini 3, map to thinkingLevel - if "gemini-3" in model: - levels = {"low": 1, "medium": 2, "high": 3} - level = levels.get(reasoning_effort, 1) # Default to level 1 - return {"thinkingLevel": level, "include_thoughts": True} - - return None - # ============================================================================ # ANTIGRAVITY REQUEST TRANSFORMATION # ============================================================================ @@ -946,8 +910,16 @@ async def acompletion( if top_p is not None: generation_config["topP"] = top_p + # Extract custom_reasoning_budget toggle + # Check kwargs first, then headers if not found + custom_reasoning_budget = kwargs.get("custom_reasoning_budget", False) + # Handle thinking config - thinking_config = self._map_reasoning_effort_to_thinking_config(reasoning_effort, model) + thinking_config = self._map_reasoning_effort_to_thinking_config( + reasoning_effort, + model, + custom_reasoning_budget + ) if thinking_config: generation_config.setdefault("thinkingConfig", {}).update(thinking_config) From ff827398a926e2e4246cb7fe2086de056c4497a4 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 16:08:38 +0100 Subject: [PATCH 005/221] =?UTF-8?q?feat(providers):=20=E2=9C=A8=20add=20se?= =?UTF-8?q?rver-side=20thoughtSignature=20cache=20and=20preserve=20thought?= =?UTF-8?q?Signature=20handling=20for=20Gemini=203?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Introduce ThoughtSignatureCache: TTL-based, thread-safe, auto-cleanup cache for mapping tool_call_id → thoughtSignature. - Integrate cache into AntigravityProvider and add env toggles: - ANTIGRAVITY_SIGNATURE_CACHE_TTL (default 3600s) - ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES (client passthrough) - ANTIGRAVITY_ENABLE_SIGNATURE_CACHE (server-side caching) - Update message transformation to accept model and implement a 3-tier thoughtSignature fallback: 1. client-provided signature 2. server-side cache 3. bypass constant ("skip_thought_signature_validator") with warning for Gemini 3 - Fix Gemini → OpenAI chunk conversion: - Stop dropping function calls that include signatures (skip only standalone signature parts). - Store signatures into server cache and optionally include them in responses when passthrough is enabled. - Robustly parse tool responses, map finish reasons, and include reasoning token counts in usage. - Improve tool response grouping and id generation; add informative logging for signature-preservation behavior --- .../providers/antigravity_provider.py | 332 +++++++++++++----- 1 file changed, 250 insertions(+), 82 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index af254600..c5a9c21c 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -99,6 +99,75 @@ def log_final_response(self, response_data: Dict[str, Any]): except Exception as e: lib_logger.error(f"_AntigravityFileLogger: Failed to write final response: {e}") +class ThoughtSignatureCache: + """ + Server-side cache for thoughtSignatures to maintain Gemini 3 conversation context. + + Maps tool_call_id → thoughtSignature to preserve encrypted reasoning signatures + across turns, even if clients don't support the thought_signature field. + + Features: + - TTL-based expiration to prevent memory growth + - Thread-safe for concurrent access + - Automatic cleanup of expired entries + """ + + def __init__(self, ttl_seconds: int = 3600): + """ + Initialize the signature cache. + + Args: + ttl_seconds: Time-to-live for cache entries in seconds (default: 1 hour) + """ + self._cache: Dict[str, Tuple[str, float]] = {} # {call_id: (signature, timestamp)} + self._ttl = ttl_seconds + self._lock = threading.Lock() + + def store(self, tool_call_id: str, signature: str): + """ + Store a signature for a tool call ID. + + Args: + tool_call_id: Unique identifier for the tool call + signature: Encrypted thoughtSignature from Antigravity API + """ + with self._lock: + self._cache[tool_call_id] = (signature, time.time()) + self._cleanup_expired() + + def retrieve(self, tool_call_id: str) -> Optional[str]: + """ + Retrieve signature for a tool call ID. + + Args: + tool_call_id: Unique identifier for the tool call + + Returns: + The signature if found and not expired, None otherwise + """ + with self._lock: + if tool_call_id not in self._cache: + return None + + signature, timestamp = self._cache[tool_call_id] + if time.time() - timestamp > self._ttl: + del self._cache[tool_call_id] + return None + + return signature + + def _cleanup_expired(self): + """Remove expired entries from cache.""" + now = time.time() + expired = [k for k, (_, ts) in self._cache.items() if now - ts > self._ttl] + for k in expired: + del self._cache[k] + + def clear(self): + """Clear all cached signatures.""" + with self._lock: + self._cache.clear() + class AntigravityProvider(AntigravityAuthBase, ProviderInterface): """ Antigravity provider implementation for Gemini models. @@ -131,6 +200,32 @@ def __init__(self): self.model_definitions = ModelDefinitions() self._current_base_url = BASE_URLS[0] # Start with daily sandbox self._base_url_index = 0 + + # Initialize thoughtSignature cache for Gemini 3 multi-turn conversations + cache_ttl = int(os.getenv("ANTIGRAVITY_SIGNATURE_CACHE_TTL", "3600")) + self._signature_cache = ThoughtSignatureCache(ttl_seconds=cache_ttl) + + # Check if client passthrough is enabled (default: TRUE for testing) + self._preserve_signatures_in_client = os.getenv( + "ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES", + "true" # Default ON for testing + ).lower() in ("true", "1", "yes") + + # Check if server-side cache is enabled (default: TRUE for testing) + self._enable_signature_cache = os.getenv( + "ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", + "true" # Default ON for testing + ).lower() in ("true", "1", "yes") + + if self._preserve_signatures_in_client: + lib_logger.info("Antigravity: thoughtSignature client passthrough ENABLED") + else: + lib_logger.info("Antigravity: thoughtSignature client passthrough DISABLED") + + if self._enable_signature_cache: + lib_logger.info(f"Antigravity: thoughtSignature server-side cache ENABLED (TTL: {cache_ttl}s)") + else: + lib_logger.info("Antigravity: thoughtSignature server-side cache DISABLED") # ============================================================================ # MODEL ALIAS SYSTEM @@ -183,6 +278,19 @@ def _alias_to_model_name(self, alias: str) -> str: } return reverse_map.get(alias, alias) + def _is_gemini_3_model(self, model: str) -> bool: + """ + Check if model is Gemini 3 (requires thoughtSignature preservation). + + Args: + model: Model name (public alias) + + Returns: + True if this is a Gemini 3 model + """ + internal_model = self._alias_to_model_name(model) + return internal_model.startswith("gemini-3-") or model.startswith("gemini-3-") + # ============================================================================ # RANDOM ID GENERATION # ============================================================================ @@ -213,11 +321,20 @@ def generate_project_id() -> str: # MESSAGE TRANSFORMATION (OpenAI → Gemini CLI format) # ============================================================================ - def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: + def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: """ Transform OpenAI messages to Gemini CLI format. Reused from GeminiCliProvider with modifications for Antigravity. + UPDATED: Now handles thoughtSignature preservation with 3-tier fallback: + 1. Use client-provided signature (if present) + 2. Fall back to server-side cache + 3. Use bypass constant as last resort + + Args: + messages: List of OpenAI-formatted messages + model: Model name for Gemini 3 detection + Returns: Tuple of (system_instruction, gemini_contents) """ @@ -244,7 +361,7 @@ def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[ if tool_call.get("type") == "function": tool_call_id_to_name[tool_call["id"]] = tool_call["function"]["name"] - # Convert each message + #Convert each message for msg in messages: role = msg.get("role") content = msg.get("content") @@ -291,34 +408,64 @@ def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[ except (json.JSONDecodeError, TypeError): args_dict = {} - # Add function call part with thoughtSignature - # ThoughtSignature is required for Gemini to process function calls correctly - # The constant "skip_thought_signature_validator" tells Gemini to bypass signature validation - # This is preserved across conversation turns to maintain reasoning continuity + tool_call_id = tool_call.get("id", "") + func_call_part = { "functionCall": { "name": tool_call["function"]["name"], "args": args_dict - }, - "thoughtSignature": "skip_thought_signature_validator" + } } + + # PRIORITY 1: Use client-provided signature if available + client_signature = tool_call.get("thought_signature") + + # PRIORITY 2: Fall back to server-side cache + if not client_signature and tool_call_id and self._enable_signature_cache: + client_signature = self._signature_cache.retrieve(tool_call_id) + if client_signature: + lib_logger.debug(f"Retrieved thoughtSignature from cache for {tool_call_id}") + + # PRIORITY 3: Use bypass constant as last resort + if client_signature: + func_call_part["thoughtSignature"] = client_signature + else: + func_call_part["thoughtSignature"] = "skip_thought_signature_validator" + + # WARNING: Missing signature for Gemini 3 + if self._is_gemini_3_model(model): + lib_logger.warning( + f"Gemini 3 tool call '{tool_call_id}' missing thoughtSignature. " + f"Client didn't provide it and cache lookup failed. " + f"Using bypass - reasoning quality may degrade." + ) + parts.append(func_call_part) elif role == "tool": - tool_call_id = msg.get("tool_call_id") - function_name = tool_call_id_to_name.get(tool_call_id) - if function_name: - # Wrap the tool response in a 'result' object - response_content = {"result": content} - parts.append({"functionResponse": {"name": function_name, "response": response_content}}) + # Tool responses grouped by function name + tool_call_id = msg.get("tool_call_id", "") + function_name = tool_call_id_to_name.get(tool_call_id, "unknown_function") + tool_content = msg.get("content", "{}") + + try: + response_data = json.loads(tool_content) + except (json.JSONDecodeError, TypeError): + response_data = {"result": tool_content} + + parts.append({ + "functionResponse": { + "name": function_name, + "response": response_data + } + }) if parts: - gemini_contents.append({"role": gemini_role, "parts": parts}) - - # Ensure first message is from user - if not gemini_contents or gemini_contents[0]['role'] != 'user': - gemini_contents.insert(0, {"role": "user", "parts": [{"text": ""}]}) - + gemini_contents.append({ + "role": gemini_role, + "parts": parts + }) + return system_instruction, gemini_contents # ============================================================================ @@ -643,106 +790,117 @@ def _unwrap_antigravity_response(self, antigravity_response: Dict[str, Any]) -> # For both streaming and non-streaming, response is in 'response' field return antigravity_response.get("response", antigravity_response) - def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> litellm.ModelResponse: + def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> Dict[str, Any]: """ - Convert a single Gemini response chunk to OpenAI format. + Convert a Gemini API response chunk to OpenAI format. - Handles Gemini 3 special mechanics: - - Filters thoughtSignature parts (encrypted reasoning data) - - Separates reasoning content (thought=true) from regular content - - Includes thoughtsTokenCount in usage metadata + UPDATED: Now preserves thoughtSignatures for Gemini 3 multi-turn conversations: + - Stores signatures in server-side cache (if enabled) + - Includes signatures in response (if client passthrough enabled) + - Filters standalone signature parts (no functionCall/text) Args: - gemini_chunk: Gemini response chunk - model: Model name + gemini_chunk: Gemini API response chunk + model: Model name for Gemini 3 detection Returns: - OpenAI-format ModelResponse + OpenAI-compatible response chunk """ - # Extract candidate + # Extract the main response structure candidates = gemini_chunk.get("candidates", []) if not candidates: - # Empty chunk, return minimal response - return litellm.ModelResponse( - id=f"chatcmpl-{uuid.uuid4()}", - created=int(time.time()), - model=model, - choices=[] - ) + return {} candidate = candidates[0] - content_parts = candidate.get("content", {}).get("parts", []) + content = candidate.get("content", {}) + content_parts = content.get("parts", []) - # Extract text, tool calls, and reasoning content + # Build delta components text_content = "" reasoning_content = "" tool_calls = [] for part in content_parts: - # CRITICAL: Skip parts with thoughtSignature (encrypted reasoning data) - # This prevents exposing internal Gemini reasoning signatures to clients - if "thoughtSignature" in part and part["thoughtSignature"]: - continue + has_function_call = "functionCall" in part + has_text = "text" in part + has_signature = "thoughtSignature" in part and part["thoughtSignature"] + + # FIXED: Only skip if ONLY signature (standalone encryption part) + # Previously this filtered out ALL function calls with signatures! + if has_signature and not has_function_call and not has_text: + continue # Skip standalone signature parts - # Extract text - separate regular content from reasoning/thinking - if "text" in part: - # Check for thought flag (Gemini 3 reasoning indicator) + # Process text content + if has_text: thought = part.get("thought") if thought is True or (isinstance(thought, str) and thought.lower() == 'true'): - # This is reasoning/thinking content reasoning_content += part["text"] else: - # Regular content text_content += part["text"] - # Extract function calls (tool calls) - if "functionCall" in part: + # Process function calls (NOW WORKS with signatures!) + if has_function_call: func_call = part["functionCall"] - tool_calls.append({ - "id": f"call_{uuid.uuid4().hex[:24]}", + tool_call_id = f"call_{uuid.uuid4().hex[:24]}" + + tool_call = { + "id": tool_call_id, "type": "function", "function": { "name": func_call.get("name", ""), "arguments": json.dumps(func_call.get("args", {})) } - }) + } + + # Store signature in server-side cache (if enabled and signature exists) + if has_signature and self._enable_signature_cache: + signature = part["thoughtSignature"] + self._signature_cache.store(tool_call_id, signature) + lib_logger.debug(f"Stored thoughtSignature in cache for {tool_call_id}") + + # Include in response if client passthrough enabled + if self._preserve_signatures_in_client: + tool_call["thought_signature"] = signature + + tool_calls.append(tool_call) # Build delta delta = {} if text_content: delta["content"] = text_content if reasoning_content: - # OpenAI o1-style reasoning content field delta["reasoning_content"] = reasoning_content if tool_calls: delta["tool_calls"] = tool_calls + delta["role"] = "assistant" + elif text_content or reasoning_content: + delta["role"] = "assistant" + + # Handle finish reason + finish_reason = candidate.get("finishReason") + if finish_reason: + # Map Gemini finish reasons to OpenAI + finish_reason_map = { + "STOP": "stop", + "MAX_TOKENS": "length", + "SAFETY": "content_filter", + "RECITATION": "content_filter", + "OTHER": "stop" + } + finish_reason = finish_reason_map.get(finish_reason, "stop") + if tool_calls: + finish_reason = "tool_calls" - # Get finish reason - finish_reason = candidate.get("finishReason", "").lower() if candidate.get("finishReason") else None - if finish_reason == "stop": - finish_reason = "stop" - elif finish_reason == "max_tokens": - finish_reason = "length" - - # Build choice - choice = { - "index": 0, - "delta": delta, - "finish_reason": finish_reason - } - - # Extract usage (if present) - usage_metadata = gemini_chunk.get("usageMetadata", {}) + # Build usage metadata usage = None + usage_metadata = gemini_chunk.get("usageMetadata", {}) if usage_metadata: - # Get token counts prompt_tokens = usage_metadata.get("promptTokenCount", 0) thoughts_tokens = usage_metadata.get("thoughtsTokenCount", 0) completion_tokens = usage_metadata.get("candidatesTokenCount", 0) - # OpenAI o1-style token counting: thoughts are included in prompt_tokens usage = { - "prompt_tokens": prompt_tokens + thoughts_tokens, + "prompt_tokens": prompt_tokens + thoughts_tokens, # Include thoughts in prompt "completion_tokens": completion_tokens, "total_tokens": usage_metadata.get("totalTokenCount", 0) } @@ -753,13 +911,23 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> l usage["completion_tokens_details"] = {} usage["completion_tokens_details"]["reasoning_tokens"] = thoughts_tokens - return litellm.ModelResponse( - id=f"chatcmpl-{uuid.uuid4()}", - created=int(time.time()), - model=model, - choices=[choice], - usage=usage - ) + # Build final response + response = { + "id": gemini_chunk.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [{ + "index": 0, + "delta": delta, + "finish_reason": finish_reason + }] + } + + if usage: + response["usage"] = usage + + return response # ============================================================================ # PROVIDER INTERFACE IMPLEMENTATION @@ -890,7 +1058,7 @@ async def acompletion( ) # Step 1: Transform messages (OpenAI → Gemini CLI) - system_instruction, gemini_contents = self._transform_messages(messages) + system_instruction, gemini_contents = self._transform_messages(messages, model=model) # Apply tool response grouping gemini_contents = self._fix_tool_response_grouping(gemini_contents) @@ -1080,7 +1248,7 @@ async def count_tokens( internal_model = self._alias_to_model_name(model) # Transform messages to Gemini format - system_instruction, contents = self._transform_messages(messages) + system_instruction, contents = self._transform_messages(messages, model=internal_model) # Build Gemini CLI payload gemini_cli_payload = { From 065d589302a6b090790a536aceae58355aed07ae Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 16:24:42 +0100 Subject: [PATCH 006/221] =?UTF-8?q?fix(providers):=20=F0=9F=90=9B=20ensure?= =?UTF-8?q?=20only=20first=20parallel=20tool=20call=20retains=20thoughtSig?= =?UTF-8?q?nature=20and=20decouple=20cache/passthrough?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enforce Gemini 3 behavior where only the first tool call in parallel receives a thoughtSignature. Previously caching and client passthrough were coupled and could result in multiple signatures being stored or passed. This change: - add a first_signature_seen flag to ensure only the first tool call gets the signature - store signature in server-side cache only when _enable_signature_cache is true - pass signature to the client only when _preserve_signatures_in_client is true - preserve logging when a signature is stored in cache --- .../providers/antigravity_provider.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index c5a9c21c..dae1ea60 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -820,6 +820,10 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D reasoning_content = "" tool_calls = [] + # Track if we've seen a signature yet (for parallel tool call handling) + # Per Gemini 3 spec: only FIRST tool call in parallel gets signature + first_signature_seen = False + for part in content_parts: has_function_call = "functionCall" in part has_text = "text" in part @@ -852,13 +856,19 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D } } - # Store signature in server-side cache (if enabled and signature exists) - if has_signature and self._enable_signature_cache: + # Handle thoughtSignature if present + # CRITICAL FIX: Cache and passthrough are INDEPENDENT toggles + if has_signature and not first_signature_seen: + # Only first tool call gets signature (parallel call handling) + first_signature_seen = True signature = part["thoughtSignature"] - self._signature_cache.store(tool_call_id, signature) - lib_logger.debug(f"Stored thoughtSignature in cache for {tool_call_id}") - # Include in response if client passthrough enabled + # Option 1: Store in server-side cache (if enabled) + if self._enable_signature_cache: + self._signature_cache.store(tool_call_id, signature) + lib_logger.debug(f"Stored thoughtSignature in cache for {tool_call_id}") + + # Option 2: Pass to client (if enabled) - INDEPENDENT of cache! if self._preserve_signatures_in_client: tool_call["thought_signature"] = signature From fc70523f3dc7f82c38c0e74c2ea8ab2304f52458 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 17:13:36 +0100 Subject: [PATCH 007/221] =?UTF-8?q?feat(providers):=20=E2=9C=A8=20add=20cl?= =?UTF-8?q?aude-sonnet-4-5=20models=20and=20remove=20unnecessary=20aliasin?= =?UTF-8?q?g?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add "claude-sonnet-4-5" and "claude-sonnet-4-5-thinking" to HARDCODED_MODELS and simplify the alias mappings by removing explicit alias entries for these Claude models since their public names match internal names. This ensures the provider recognizes the new Claude Sonnet variants and avoids incorrect alias translations. --- src/rotator_library/providers/antigravity_provider.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index dae1ea60..ed30d417 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -37,7 +37,9 @@ "gemini-2.5-flash-lite", "gemini-3-pro-preview", "gemini-3-pro-image-preview", - "gemini-2.5-computer-use-preview-10-2025" + "gemini-2.5-computer-use-preview-10-2025", + "claude-sonnet-4-5", + "claude-sonnet-4-5-thinking" ] # Logging configuration @@ -245,8 +247,7 @@ def _model_name_to_alias(self, model_name: str) -> str: "rev19-uic3-1p": "gemini-2.5-computer-use-preview-10-2025", "gemini-3-pro-image": "gemini-3-pro-image-preview", "gemini-3-pro-high": "gemini-3-pro-preview", - "claude-sonnet-4-5": "gemini-claude-sonnet-4-5", - "claude-sonnet-4-5-thinking": "gemini-claude-sonnet-4-5-thinking", + # Claude models: no aliasing needed (public name = internal name) } # Filter out excluded models (return empty string to skip) @@ -273,8 +274,7 @@ def _alias_to_model_name(self, alias: str) -> str: "gemini-2.5-computer-use-preview-10-2025": "rev19-uic3-1p", "gemini-3-pro-image-preview": "gemini-3-pro-image", "gemini-3-pro-preview": "gemini-3-pro-high", - "gemini-claude-sonnet-4-5": "claude-sonnet-4-5", - "gemini-claude-sonnet-4-5-thinking": "claude-sonnet-4-5-thinking", + # Claude models: no aliasing needed (public name = internal name) } return reverse_map.get(alias, alias) From 97f19509e5ed77370a79364915ebe15bd74675f2 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 17:54:39 +0100 Subject: [PATCH 008/221] feat(auth): extract GoogleOAuthBase and add antigravity provider - Add providers/google_oauth_base.py to centralize Google OAuth logic (auth flow, token refresh, env loading, atomic saves, backoff/retry, queueing, headless support, and validation). - Migrate GeminiAuthBase and AntigravityAuthBase to inherit from GoogleOAuthBase and expose provider-specific constants (CLIENT_ID, CLIENT_SECRET, OAUTH_SCOPES, ENV_PREFIX, CALLBACK_PORT, CALLBACK_PATH). - Register "antigravity" in DEFAULT_OAUTH_DIRS and mark it as OAuth-only in credential_tool; include a user-friendly display name for interactive flows. - Remove large duplicated OAuth implementations from provider-specific files and consolidate behavior to reduce maintenance surface and ensure consistent token handling. --- src/rotator_library/credential_manager.py | 1 + src/rotator_library/credential_tool.py | 8 +- .../providers/antigravity_auth_base.py | 476 +------------ .../providers/gemini_auth_base.py | 642 +---------------- .../providers/google_oauth_base.py | 653 ++++++++++++++++++ 5 files changed, 695 insertions(+), 1085 deletions(-) create mode 100644 src/rotator_library/providers/google_oauth_base.py diff --git a/src/rotator_library/credential_manager.py b/src/rotator_library/credential_manager.py index c5426d76..0678f7c2 100644 --- a/src/rotator_library/credential_manager.py +++ b/src/rotator_library/credential_manager.py @@ -14,6 +14,7 @@ "gemini_cli": Path.home() / ".gemini", "qwen_code": Path.home() / ".qwen", "iflow": Path.home() / ".iflow", + "antigravity": Path.home() / ".antigravity", # Add other providers like 'claude' here if they have a standard CLI path } diff --git a/src/rotator_library/credential_tool.py b/src/rotator_library/credential_tool.py index 82c8b05e..a1705a13 100644 --- a/src/rotator_library/credential_tool.py +++ b/src/rotator_library/credential_tool.py @@ -98,7 +98,7 @@ async def setup_api_key(): # Discover custom providers and add them to the list # Note: gemini_cli is OAuth-only, but qwen_code and iflow support both OAuth and API keys _, PROVIDER_PLUGINS = _ensure_providers_loaded() - oauth_only_providers = {'gemini_cli'} + oauth_only_providers = {'gemini_cli', 'antigravity'} discovered_providers = { p.replace('_', ' ').title(): p.upper() + "_API_KEY" for p in PROVIDER_PLUGINS.keys() @@ -195,7 +195,8 @@ async def setup_new_credential(provider_name: str): oauth_friendly_names = { "gemini_cli": "Gemini CLI (OAuth)", "qwen_code": "Qwen Code (OAuth - also supports API keys)", - "iflow": "iFlow (OAuth - also supports API keys)" + "iflow": "iFlow (OAuth - also supports API keys)", + "antigravity": "Antigravity (OAuth)" } display_name = oauth_friendly_names.get(provider_name, provider_name.replace('_', ' ').title()) @@ -578,7 +579,8 @@ async def main(clear_on_start=True): oauth_friendly_names = { "gemini_cli": "Gemini CLI (OAuth)", "qwen_code": "Qwen Code (OAuth - also supports API keys)", - "iflow": "iFlow (OAuth - also supports API keys)" + "iflow": "iFlow (OAuth - also supports API keys)", + "antigravity": "Antigravity (OAuth)", } provider_text = Text() diff --git a/src/rotator_library/providers/antigravity_auth_base.py b/src/rotator_library/providers/antigravity_auth_base.py index df15dae9..7240304e 100644 --- a/src/rotator_library/providers/antigravity_auth_base.py +++ b/src/rotator_library/providers/antigravity_auth_base.py @@ -1,466 +1,24 @@ # src/rotator_library/providers/antigravity_auth_base.py -import os -import webbrowser -from typing import Union, Optional -import json -import time -import asyncio -import logging -from pathlib import Path -from typing import Dict, Any -import tempfile -import shutil +from .google_oauth_base import GoogleOAuthBase -import httpx -from rich.console import Console -from rich.panel import Panel -from rich.text import Text - -from ..utils.headless_detection import is_headless_environment - -lib_logger = logging.getLogger('rotator_library') - -# Antigravity OAuth credentials from CLIProxyAPI -CLIENT_ID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" -CLIENT_SECRET = "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" -TOKEN_URI = "https://oauth2.googleapis.com/token" -USER_INFO_URI = "https://www.googleapis.com/oauth2/v1/userinfo" -REFRESH_EXPIRY_BUFFER_SECONDS = 30 * 60 # 30 minutes buffer before expiry - -# Antigravity requires additional scopes -OAUTH_SCOPES = [ - "https://www.googleapis.com/auth/cloud-platform", - "https://www.googleapis.com/auth/userinfo.email", - "https://www.googleapis.com/auth/userinfo.profile", - "https://www.googleapis.com/auth/cclog", # Antigravity-specific - "https://www.googleapis.com/auth/experimentsandconfigs" # Antigravity-specific -] - -console = Console() - -class AntigravityAuthBase: +class AntigravityAuthBase(GoogleOAuthBase): """ - Base authentication class for Antigravity provider. - Handles OAuth2 flow, token management, and refresh logic. + Antigravity OAuth2 authentication implementation. - Based on GeminiAuthBase but uses Antigravity-specific OAuth credentials and scopes. + Inherits all OAuth functionality from GoogleOAuthBase with Antigravity-specific configuration. + Uses Antigravity's OAuth credentials and includes additional scopes for cclog and experimentsandconfigs. """ - def __init__(self): - self._credentials_cache: Dict[str, Dict[str, Any]] = {} - self._refresh_locks: Dict[str, asyncio.Lock] = {} - self._locks_lock = asyncio.Lock() # Protects the locks dict from race conditions - # [BACKOFF TRACKING] Track consecutive failures per credential - self._refresh_failures: Dict[str, int] = {} # Track consecutive failures per credential - self._next_refresh_after: Dict[str, float] = {} # Track backoff timers (Unix timestamp) - - # [QUEUE SYSTEM] Sequential refresh processing - self._refresh_queue: asyncio.Queue = asyncio.Queue() - self._queued_credentials: set = set() # Track credentials already in queue - self._unavailable_credentials: set = set() # Mark credentials unavailable during re-auth - self._queue_tracking_lock = asyncio.Lock() # Protects queue sets - self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task - - def _load_from_env(self) -> Optional[Dict[str, Any]]: - """ - Load OAuth credentials from environment variables for stateless deployments. - - Expected environment variables: - - ANTIGRAVITY_ACCESS_TOKEN (required) - - ANTIGRAVITY_REFRESH_TOKEN (required) - - ANTIGRAVITY_EXPIRY_DATE (optional, defaults to 0) - - ANTIGRAVITY_CLIENT_ID (optional, uses default) - - ANTIGRAVITY_CLIENT_SECRET (optional, uses default) - - ANTIGRAVITY_TOKEN_URI (optional, uses default) - - ANTIGRAVITY_UNIVERSE_DOMAIN (optional, defaults to googleapis.com) - - ANTIGRAVITY_EMAIL (optional, defaults to "env-user") - - Returns: - Dict with credential structure if env vars present, None otherwise - """ - access_token = os.getenv("ANTIGRAVITY_ACCESS_TOKEN") - refresh_token = os.getenv("ANTIGRAVITY_REFRESH_TOKEN") - - # Both access and refresh tokens are required - if not (access_token and refresh_token): - return None - - lib_logger.debug("Loading Antigravity credentials from environment variables") - - # Parse expiry_date as float, default to 0 if not present - expiry_str = os.getenv("ANTIGRAVITY_EXPIRY_DATE", "0") - try: - expiry_date = float(expiry_str) - except ValueError: - lib_logger.warning(f"Invalid ANTIGRAVITY_EXPIRY_DATE value: {expiry_str}, using 0") - expiry_date = 0 - - creds = { - "access_token": access_token, - "refresh_token": refresh_token, - "expiry_date": expiry_date, - "client_id": os.getenv("ANTIGRAVITY_CLIENT_ID", CLIENT_ID), - "client_secret": os.getenv("ANTIGRAVITY_CLIENT_SECRET", CLIENT_SECRET), - "token_uri": os.getenv("ANTIGRAVITY_TOKEN_URI", TOKEN_URI), - "universe_domain": os.getenv("ANTIGRAVITY_UNIVERSE_DOMAIN", "googleapis.com"), - "_proxy_metadata": { - "email": os.getenv("ANTIGRAVITY_EMAIL", "env-user"), - "last_check_timestamp": time.time(), - "loaded_from_env": True # Flag to indicate env-based credentials - } - } - - return creds - - async def _load_credentials(self, path: str) -> Dict[str, Any]: - """ - Load credentials from a file. First attempts file-based load, - then falls back to environment variables if file not found. - - Args: - path: File path to load credentials from - - Returns: - Dict containing the credentials - - Raises: - ValueError: If credentials cannot be loaded from either source - """ - # If path is special marker "env", load from environment - if path == "env": - env_creds = self._load_from_env() - if env_creds: - lib_logger.debug("Using Antigravity credentials from environment variables") - return env_creds - raise ValueError("ANTIGRAVITY_ACCESS_TOKEN and ANTIGRAVITY_REFRESH_TOKEN environment variables not set") - - # Try loading from cache first - if path in self._credentials_cache: - cached_creds = self._credentials_cache[path] - lib_logger.debug(f"Using cached Antigravity credentials for: {Path(path).name}") - return cached_creds - - # Try loading from file - try: - with open(path, 'r') as f: - creds = json.load(f) - self._credentials_cache[path] = creds - lib_logger.debug(f"Loaded Antigravity credentials from file: {Path(path).name}") - return creds - except FileNotFoundError: - # Fall back to environment variables - lib_logger.debug(f"Credential file not found: {path}, attempting environment variables") - env_creds = self._load_from_env() - if env_creds: - lib_logger.debug("Using Antigravity credentials from environment variables as fallback") - # Cache with special path marker - self._credentials_cache[path] = env_creds - return env_creds - raise ValueError(f"Credential file not found: {path} and environment variables not set") - except json.JSONDecodeError as e: - raise ValueError(f"Invalid JSON in credential file {path}: {e}") - - async def _save_credentials(self, path: str, creds: Dict[str, Any]) -> None: - """ - Save credentials to a file. Skip if credentials were loaded from environment. - - Args: - path: File path to save credentials to - creds: Credentials dictionary to save - """ - # Don't save environment-based credentials to file - if creds.get("_proxy_metadata", {}).get("loaded_from_env"): - lib_logger.debug("Skipping credential save (loaded from environment)") - return - - # Don't save if path is special marker - if path == "env": - return - - try: - # Ensure directory exists - Path(path).parent.mkdir(parents=True, exist_ok=True) - - # Write atomically using temp file + rename - temp_fd, temp_path = tempfile.mkstemp( - dir=Path(path).parent, - prefix='.tmp_', - suffix='.json' - ) - try: - with os.fdopen(temp_fd, 'w') as f: - json.dump(creds, f, indent=2) - shutil.move(temp_path, path) - lib_logger.debug(f"Saved Antigravity credentials to: {Path(path).name}") - except Exception: - # Clean up temp file on error - try: - os.unlink(temp_path) - except Exception: - pass - raise - except Exception as e: - lib_logger.warning(f"Failed to save Antigravity credentials to {path}: {e}") - - def _is_token_expired(self, creds: Dict[str, Any]) -> bool: - """ - Check if the access token is expired or close to expiry. - - Args: - creds: Credentials dict with expiry_date field (in milliseconds) - - Returns: - True if token is expired or within buffer time of expiry - """ - if 'expiry_date' not in creds: - return True - - # expiry_date is in milliseconds - expiry_timestamp = creds['expiry_date'] / 1000.0 - current_time = time.time() - - # Consider expired if within buffer time - return (expiry_timestamp - current_time) <= REFRESH_EXPIRY_BUFFER_SECONDS - - async def _refresh_token(self, path: str, creds: Dict[str, Any]) -> Dict[str, Any]: - """ - Refresh an expired OAuth token using the refresh token. - - Args: - path: Credential file path (for saving updated credentials) - creds: Current credentials dict with refresh_token - - Returns: - Updated credentials dict with fresh access token - - Raises: - ValueError: If refresh fails - """ - if 'refresh_token' not in creds: - raise ValueError("No refresh token available") - - lib_logger.debug(f"Refreshing Antigravity OAuth token for: {Path(path).name if path != 'env' else 'env'}") - - client_id = creds.get('client_id', CLIENT_ID) - client_secret = creds.get('client_secret', CLIENT_SECRET) - token_uri = creds.get('token_uri', TOKEN_URI) - - async with httpx.AsyncClient() as client: - try: - response = await client.post( - token_uri, - data={ - 'client_id': client_id, - 'client_secret': client_secret, - 'refresh_token': creds['refresh_token'], - 'grant_type': 'refresh_token' - }, - timeout=30.0 - ) - response.raise_for_status() - token_data = response.json() - - # Update credentials with new token - creds['access_token'] = token_data['access_token'] - creds['expiry_date'] = (time.time() + token_data['expires_in']) * 1000 - - # Update metadata - if '_proxy_metadata' not in creds: - creds['_proxy_metadata'] = {} - creds['_proxy_metadata']['last_check_timestamp'] = time.time() - - # Save updated credentials - await self._save_credentials(path, creds) - - # Update cache - self._credentials_cache[path] = creds - - # Reset failure count on success - self._refresh_failures[path] = 0 - - lib_logger.info(f"Successfully refreshed Antigravity OAuth token for: {Path(path).name if path != 'env' else 'env'}") - return creds - - except httpx.HTTPStatusError as e: - # Track failures for backoff - self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 - raise ValueError(f"Failed to refresh Antigravity token (HTTP {e.response.status_code}): {e.response.text}") - except Exception as e: - self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 - raise ValueError(f"Failed to refresh Antigravity token: {e}") - - async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: - """ - Initialize or refresh an OAuth token. Handles the complete OAuth flow if needed. - - Args: - creds_or_path: Either a credentials dict or a file path string - - Returns: - Valid credentials dict with fresh access token - """ - path = creds_or_path if isinstance(creds_or_path, str) else None - - if isinstance(creds_or_path, dict): - display_name = creds_or_path.get("_proxy_metadata", {}).get("display_name", "in-memory object") - else: - display_name = Path(path).name if path and path != "env" else "env" - - lib_logger.debug(f"Initializing Antigravity token for '{display_name}'...") - - try: - creds = await self._load_credentials(creds_or_path) if path else creds_or_path - reason = "" - if not creds.get("refresh_token"): - reason = "refresh token is missing" - elif self._is_token_expired(creds): - reason = "token is expired" - - if reason: - if reason == "token is expired" and creds.get("refresh_token"): - try: - return await self._refresh_token(path, creds) - except Exception as e: - lib_logger.warning(f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login.") - - lib_logger.warning(f"Antigravity OAuth token for '{display_name}' needs setup: {reason}.") - - is_headless = is_headless_environment() - - auth_code_future = asyncio.get_event_loop().create_future() - server = None - - async def handle_callback(reader, writer): - try: - request_line_bytes = await reader.readline() - if not request_line_bytes: - return - path_str = request_line_bytes.decode('utf-8').strip().split(' ')[1] - # Consume headers - while await reader.readline() != b'\r\n': - pass - - from urllib.parse import urlparse, parse_qs - query_params = parse_qs(urlparse(path_str).query) - - writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n") - if 'code' in query_params: - if not auth_code_future.done(): - auth_code_future.set_result(query_params['code'][0]) - writer.write(b"

Authentication successful!

You can close this window.

") - else: - error = query_params.get('error', ['Unknown error'])[0] - if not auth_code_future.done(): - auth_code_future.set_exception(Exception(f"OAuth failed: {error}")) - writer.write(f"

Authentication Failed

Error: {error}. Please try again.

".encode()) - await writer.drain() - except Exception as e: - lib_logger.error(f"Error in OAuth callback handler: {e}") - finally: - writer.close() - - try: - server = await asyncio.start_server(handle_callback, '127.0.0.1', 8085) - - from urllib.parse import urlencode - auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode({ - "client_id": CLIENT_ID, - "redirect_uri": "http://localhost:8085/oauth2callback", - "scope": " ".join(OAUTH_SCOPES), - "access_type": "offline", - "response_type": "code", - "prompt": "consent" - }) - - if is_headless: - auth_panel_text = Text.from_markup( - "Running in headless environment (no GUI detected).\n" - "Please open the URL below in a browser on another machine to authorize:\n" - ) - else: - auth_panel_text = Text.from_markup( - "1. Your browser will now open to log in and authorize the application.\n" - "2. If it doesn't open automatically, please open the URL below manually." - ) - - console.print(Panel(auth_panel_text, title=f"Antigravity OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue")) - console.print(f"[bold]URL:[/bold] [link={auth_url}]{auth_url}[/link]\n") - - if not is_headless: - try: - webbrowser.open(auth_url) - lib_logger.info("Browser opened successfully for OAuth flow") - except Exception as e: - lib_logger.warning(f"Failed to open browser automatically: {e}. Please open the URL manually.") - - with console.status("[bold green]Waiting for you to complete authentication in the browser...[/bold green]", spinner="dots"): - auth_code = await asyncio.wait_for(auth_code_future, timeout=300) - except asyncio.TimeoutError: - raise Exception("OAuth flow timed out. Please try again.") - finally: - if server: - server.close() - await server.wait_closed() - - lib_logger.info(f"Attempting to exchange authorization code for tokens...") - async with httpx.AsyncClient() as client: - response = await client.post(TOKEN_URI, data={ - "code": auth_code.strip(), - "client_id": CLIENT_ID, - "client_secret": CLIENT_SECRET, - "redirect_uri": "http://localhost:8085/oauth2callback", - "grant_type": "authorization_code" - }) - response.raise_for_status() - token_data = response.json() - - creds = token_data.copy() - creds["expiry_date"] = (time.time() + creds.pop("expires_in")) * 1000 - creds["client_id"] = CLIENT_ID - creds["client_secret"] = CLIENT_SECRET - creds["token_uri"] = TOKEN_URI - creds["universe_domain"] = "googleapis.com" - - # Fetch user info - user_info_response = await client.get( - USER_INFO_URI, - headers={"Authorization": f"Bearer {creds['access_token']}"} - ) - user_info_response.raise_for_status() - user_info = user_info_response.json() - - creds["_proxy_metadata"] = { - "email": user_info.get("email"), - "last_check_timestamp": time.time() - } - - if path: - await self._save_credentials(path, creds) - - lib_logger.info(f"Antigravity OAuth initialized successfully for '{display_name}'.") - return creds - - lib_logger.info(f"Antigravity OAuth token at '{display_name}' is valid.") - return creds - except Exception as e: - raise ValueError(f"Failed to initialize Antigravity OAuth for '{display_name}': {e}") - - async def get_valid_token(self, credential_path: str) -> str: - """ - Get a valid access token, refreshing if necessary. - - Args: - credential_path: Path to credential file or "env" for environment variables - - Returns: - Valid access token string - - Raises: - ValueError: If token cannot be obtained - """ - try: - creds = await self.initialize_token(credential_path) - return creds['access_token'] - except Exception as e: - raise ValueError(f"Failed to get valid Antigravity token: {e}") + CLIENT_ID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" + CLIENT_SECRET = "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" + OAUTH_SCOPES = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + "https://www.googleapis.com/auth/cclog", # Antigravity-specific + "https://www.googleapis.com/auth/experimentsandconfigs", # Antigravity-specific + ] + ENV_PREFIX = "ANTIGRAVITY" + CALLBACK_PORT = 51121 + CALLBACK_PATH = "/oauthcallback" diff --git a/src/rotator_library/providers/gemini_auth_base.py b/src/rotator_library/providers/gemini_auth_base.py index 6e8c1cce..90b9d9a6 100644 --- a/src/rotator_library/providers/gemini_auth_base.py +++ b/src/rotator_library/providers/gemini_auth_base.py @@ -1,625 +1,21 @@ # src/rotator_library/providers/gemini_auth_base.py -import os -import webbrowser -from typing import Union, Optional -import json -import time -import asyncio -import logging -from pathlib import Path -from typing import Dict, Any -import tempfile -import shutil - -import httpx -from rich.console import Console -from rich.panel import Panel -from rich.text import Text - -from ..utils.headless_detection import is_headless_environment - -lib_logger = logging.getLogger('rotator_library') - -CLIENT_ID = "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com" #https://api.kilocode.ai/extension-config.json -CLIENT_SECRET = "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" #https://api.kilocode.ai/extension-config.json -TOKEN_URI = "https://oauth2.googleapis.com/token" -USER_INFO_URI = "https://www.googleapis.com/oauth2/v1/userinfo" -REFRESH_EXPIRY_BUFFER_SECONDS = 30 * 60 # 30 minutes buffer before expiry - -console = Console() - -class GeminiAuthBase: - def __init__(self): - self._credentials_cache: Dict[str, Dict[str, Any]] = {} - self._refresh_locks: Dict[str, asyncio.Lock] = {} - self._locks_lock = asyncio.Lock() # Protects the locks dict from race conditions - # [BACKOFF TRACKING] Track consecutive failures per credential - self._refresh_failures: Dict[str, int] = {} # Track consecutive failures per credential - self._next_refresh_after: Dict[str, float] = {} # Track backoff timers (Unix timestamp) - - # [QUEUE SYSTEM] Sequential refresh processing - self._refresh_queue: asyncio.Queue = asyncio.Queue() - self._queued_credentials: set = set() # Track credentials already in queue - self._unavailable_credentials: set = set() # Mark credentials unavailable during re-auth - self._queue_tracking_lock = asyncio.Lock() # Protects queue sets - self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task - - def _load_from_env(self) -> Optional[Dict[str, Any]]: - """ - Load OAuth credentials from environment variables for stateless deployments. - - Expected environment variables: - - GEMINI_CLI_ACCESS_TOKEN (required) - - GEMINI_CLI_REFRESH_TOKEN (required) - - GEMINI_CLI_EXPIRY_DATE (optional, defaults to 0) - - GEMINI_CLI_CLIENT_ID (optional, uses default) - - GEMINI_CLI_CLIENT_SECRET (optional, uses default) - - GEMINI_CLI_TOKEN_URI (optional, uses default) - - GEMINI_CLI_UNIVERSE_DOMAIN (optional, defaults to googleapis.com) - - GEMINI_CLI_EMAIL (optional, defaults to "env-user") - - GEMINI_CLI_PROJECT_ID (optional) - - GEMINI_CLI_TIER (optional) - - Returns: - Dict with credential structure if env vars present, None otherwise - """ - access_token = os.getenv("GEMINI_CLI_ACCESS_TOKEN") - refresh_token = os.getenv("GEMINI_CLI_REFRESH_TOKEN") - - # Both access and refresh tokens are required - if not (access_token and refresh_token): - return None - - lib_logger.debug("Loading Gemini CLI credentials from environment variables") - - # Parse expiry_date as float, default to 0 if not present - expiry_str = os.getenv("GEMINI_CLI_EXPIRY_DATE", "0") - try: - expiry_date = float(expiry_str) - except ValueError: - lib_logger.warning(f"Invalid GEMINI_CLI_EXPIRY_DATE value: {expiry_str}, using 0") - expiry_date = 0 - - creds = { - "access_token": access_token, - "refresh_token": refresh_token, - "expiry_date": expiry_date, - "client_id": os.getenv("GEMINI_CLI_CLIENT_ID", CLIENT_ID), - "client_secret": os.getenv("GEMINI_CLI_CLIENT_SECRET", CLIENT_SECRET), - "token_uri": os.getenv("GEMINI_CLI_TOKEN_URI", TOKEN_URI), - "universe_domain": os.getenv("GEMINI_CLI_UNIVERSE_DOMAIN", "googleapis.com"), - "_proxy_metadata": { - "email": os.getenv("GEMINI_CLI_EMAIL", "env-user"), - "last_check_timestamp": time.time(), - "loaded_from_env": True # Flag to indicate env-based credentials - } - } - - # Add project_id if provided - project_id = os.getenv("GEMINI_CLI_PROJECT_ID") - if project_id: - creds["_proxy_metadata"]["project_id"] = project_id - - # Add tier if provided - tier = os.getenv("GEMINI_CLI_TIER") - if tier: - creds["_proxy_metadata"]["tier"] = tier - - return creds - - async def _load_credentials(self, path: str) -> Dict[str, Any]: - if path in self._credentials_cache: - return self._credentials_cache[path] - - async with await self._get_lock(path): - if path in self._credentials_cache: - return self._credentials_cache[path] - - # First, try loading from environment variables - env_creds = self._load_from_env() - if env_creds: - lib_logger.info("Using Gemini CLI credentials from environment variables") - # Cache env-based credentials using the path as key - self._credentials_cache[path] = env_creds - return env_creds - - # Fall back to file-based loading - try: - lib_logger.debug(f"Loading Gemini credentials from file: {path}") - with open(path, 'r') as f: - creds = json.load(f) - # Handle gcloud-style creds file which nest tokens under "credential" - if "credential" in creds: - creds = creds["credential"] - self._credentials_cache[path] = creds - return creds - except FileNotFoundError: - raise IOError(f"Gemini OAuth credential file not found at '{path}'") - except Exception as e: - raise IOError(f"Failed to load Gemini OAuth credentials from '{path}': {e}") - - async def _save_credentials(self, path: str, creds: Dict[str, Any]): - # Don't save to file if credentials were loaded from environment - if creds.get("_proxy_metadata", {}).get("loaded_from_env"): - lib_logger.debug("Credentials loaded from env, skipping file save") - # Still update cache for in-memory consistency - self._credentials_cache[path] = creds - return - - # [ATOMIC WRITE] Use tempfile + move pattern to ensure atomic writes - # This prevents credential corruption if the process is interrupted during write - parent_dir = os.path.dirname(os.path.abspath(path)) - os.makedirs(parent_dir, exist_ok=True) - - tmp_fd = None - tmp_path = None - try: - # Create temp file in same directory as target (ensures same filesystem) - tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json', text=True) - - # Write JSON to temp file - with os.fdopen(tmp_fd, 'w') as f: - json.dump(creds, f, indent=2) - tmp_fd = None # fdopen closes the fd - - # Set secure permissions (0600 = owner read/write only) - try: - os.chmod(tmp_path, 0o600) - except (OSError, AttributeError): - # Windows may not support chmod, ignore - pass - - # Atomic move (overwrites target if it exists) - shutil.move(tmp_path, path) - tmp_path = None # Successfully moved - - # Update cache AFTER successful file write (prevents cache/file inconsistency) - self._credentials_cache[path] = creds - lib_logger.debug(f"Saved updated Gemini OAuth credentials to '{path}' (atomic write).") - - except Exception as e: - lib_logger.error(f"Failed to save updated Gemini OAuth credentials to '{path}': {e}") - # Clean up temp file if it still exists - if tmp_fd is not None: - try: - os.close(tmp_fd) - except: - pass - if tmp_path and os.path.exists(tmp_path): - try: - os.unlink(tmp_path) - except: - pass - raise - - def _is_token_expired(self, creds: Dict[str, Any]) -> bool: - expiry = creds.get("token_expiry") # gcloud format - if not expiry: # gemini-cli format - expiry_timestamp = creds.get("expiry_date", 0) / 1000 - else: - expiry_timestamp = time.mktime(time.strptime(expiry, "%Y-%m-%dT%H:%M:%SZ")) - return expiry_timestamp < time.time() + REFRESH_EXPIRY_BUFFER_SECONDS - - async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = False) -> Dict[str, Any]: - async with await self._get_lock(path): - # Skip the expiry check if a refresh is being forced - if not force and not self._is_token_expired(self._credentials_cache.get(path, creds)): - return self._credentials_cache.get(path, creds) - - lib_logger.debug(f"Refreshing Gemini OAuth token for '{Path(path).name}' (forced: {force})...") - refresh_token = creds.get("refresh_token") - if not refresh_token: - raise ValueError("No refresh_token found in credentials file.") - - # [RETRY LOGIC] Implement exponential backoff for transient errors - max_retries = 3 - new_token_data = None - last_error = None - needs_reauth = False - - async with httpx.AsyncClient() as client: - for attempt in range(max_retries): - try: - response = await client.post(TOKEN_URI, data={ - "client_id": creds.get("client_id", CLIENT_ID), - "client_secret": creds.get("client_secret", CLIENT_SECRET), - "refresh_token": refresh_token, - "grant_type": "refresh_token", - }, timeout=30.0) - response.raise_for_status() - new_token_data = response.json() - break # Success, exit retry loop - - except httpx.HTTPStatusError as e: - last_error = e - status_code = e.response.status_code - - # [INVALID GRANT HANDLING] Handle 401/403 by triggering re-authentication - if status_code == 401 or status_code == 403: - lib_logger.warning( - f"Refresh token invalid for '{Path(path).name}' (HTTP {status_code}). " - f"Token may have been revoked or expired. Starting re-authentication..." - ) - needs_reauth = True - break # Exit retry loop to trigger re-auth - - elif status_code == 429: - # Rate limit - honor Retry-After header if present - retry_after = int(e.response.headers.get("Retry-After", 60)) - lib_logger.warning(f"Rate limited (HTTP 429), retry after {retry_after}s") - if attempt < max_retries - 1: - await asyncio.sleep(retry_after) - continue - raise - - elif status_code >= 500 and status_code < 600: - # Server error - retry with exponential backoff - if attempt < max_retries - 1: - wait_time = 2 ** attempt # 1s, 2s, 4s - lib_logger.warning(f"Server error (HTTP {status_code}), retry {attempt + 1}/{max_retries} in {wait_time}s") - await asyncio.sleep(wait_time) - continue - raise # Final attempt failed - - else: - # Other errors - don't retry - raise - - except (httpx.RequestError, httpx.TimeoutException) as e: - # Network errors - retry with backoff - last_error = e - if attempt < max_retries - 1: - wait_time = 2 ** attempt - lib_logger.warning(f"Network error during refresh: {e}, retry {attempt + 1}/{max_retries} in {wait_time}s") - await asyncio.sleep(wait_time) - continue - raise - - # [INVALID GRANT RE-AUTH] Trigger OAuth flow if refresh token is invalid - if needs_reauth: - lib_logger.info(f"Starting re-authentication for '{Path(path).name}'...") - try: - # Call initialize_token to trigger OAuth flow - new_creds = await self.initialize_token(path) - return new_creds - except Exception as reauth_error: - lib_logger.error(f"Re-authentication failed for '{Path(path).name}': {reauth_error}") - raise ValueError(f"Refresh token invalid and re-authentication failed: {reauth_error}") - - # If we exhausted retries without success - if new_token_data is None: - raise last_error or Exception("Token refresh failed after all retries") - - # [FIX 1] Update OAuth token fields from response - creds["access_token"] = new_token_data["access_token"] - expiry_timestamp = time.time() + new_token_data["expires_in"] - creds["expiry_date"] = expiry_timestamp * 1000 # gemini-cli format - - # [FIX 2] Update refresh_token if server provided a new one (rare but possible with Google OAuth) - if "refresh_token" in new_token_data: - creds["refresh_token"] = new_token_data["refresh_token"] - - # [FIX 3] Ensure all required OAuth client fields are present (restore if missing) - if "client_id" not in creds or not creds["client_id"]: - creds["client_id"] = CLIENT_ID - if "client_secret" not in creds or not creds["client_secret"]: - creds["client_secret"] = CLIENT_SECRET - if "token_uri" not in creds or not creds["token_uri"]: - creds["token_uri"] = TOKEN_URI - if "universe_domain" not in creds or not creds["universe_domain"]: - creds["universe_domain"] = "googleapis.com" - - # [FIX 4] Add scopes array if missing - if "scopes" not in creds: - creds["scopes"] = [ - "https://www.googleapis.com/auth/cloud-platform", - "https://www.googleapis.com/auth/userinfo.email", - "https://www.googleapis.com/auth/userinfo.profile", - ] - - # [FIX 5] Ensure _proxy_metadata exists and update timestamp - if "_proxy_metadata" not in creds: - creds["_proxy_metadata"] = {} - creds["_proxy_metadata"]["last_check_timestamp"] = time.time() - - # [VALIDATION] Verify refreshed credentials have all required fields - required_fields = ["access_token", "refresh_token", "client_id", "client_secret", "token_uri"] - missing_fields = [field for field in required_fields if not creds.get(field)] - if missing_fields: - raise ValueError(f"Refreshed credentials missing required fields: {missing_fields}") - - # [VALIDATION] Optional: Test that the refreshed token is actually usable - try: - async with httpx.AsyncClient() as client: - test_response = await client.get( - USER_INFO_URI, - headers={"Authorization": f"Bearer {creds['access_token']}"}, - timeout=5.0 - ) - test_response.raise_for_status() - lib_logger.debug(f"Token validation successful for '{Path(path).name}'") - except Exception as e: - lib_logger.warning(f"Refreshed token validation failed for '{Path(path).name}': {e}") - # Don't fail the refresh - the token might still work for other endpoints - # But log it for debugging purposes - - await self._save_credentials(path, creds) - lib_logger.debug(f"Successfully refreshed Gemini OAuth token for '{Path(path).name}'.") - return creds - - async def proactively_refresh(self, credential_path: str): - """Proactively refresh a credential by queueing it for refresh.""" - creds = await self._load_credentials(credential_path) - if self._is_token_expired(creds): - # Queue for refresh with needs_reauth=False (automated refresh) - await self._queue_refresh(credential_path, force=False, needs_reauth=False) - - async def _get_lock(self, path: str) -> asyncio.Lock: - # [FIX RACE CONDITION] Protect lock creation with a master lock - # This prevents TOCTOU bug where multiple coroutines check and create simultaneously - async with self._locks_lock: - if path not in self._refresh_locks: - self._refresh_locks[path] = asyncio.Lock() - return self._refresh_locks[path] - - def is_credential_available(self, path: str) -> bool: - """Check if a credential is available for rotation (not queued/refreshing).""" - return path not in self._unavailable_credentials - - async def _ensure_queue_processor_running(self): - """Lazily starts the queue processor if not already running.""" - if self._queue_processor_task is None or self._queue_processor_task.done(): - self._queue_processor_task = asyncio.create_task(self._process_refresh_queue()) - - async def _queue_refresh(self, path: str, force: bool = False, needs_reauth: bool = False): - """Add a credential to the refresh queue if not already queued. - - Args: - path: Credential file path - force: Force refresh even if not expired - needs_reauth: True if full re-authentication needed (bypasses backoff) - """ - # IMPORTANT: Only check backoff for simple automated refreshes - # Re-authentication (interactive OAuth) should BYPASS backoff since it needs user input - if not needs_reauth: - now = time.time() - if path in self._next_refresh_after: - backoff_until = self._next_refresh_after[path] - if now < backoff_until: - # Credential is in backoff for automated refresh, do not queue - remaining = int(backoff_until - now) - lib_logger.debug(f"Skipping automated refresh for '{Path(path).name}' (in backoff for {remaining}s)") - return - - async with self._queue_tracking_lock: - if path not in self._queued_credentials: - self._queued_credentials.add(path) - self._unavailable_credentials.add(path) # Mark as unavailable - await self._refresh_queue.put((path, force, needs_reauth)) - await self._ensure_queue_processor_running() - - async def _process_refresh_queue(self): - """Background worker that processes refresh requests sequentially.""" - while True: - path = None - try: - # Wait for an item with timeout to allow graceful shutdown - try: - path, force, needs_reauth = await asyncio.wait_for( - self._refresh_queue.get(), - timeout=60.0 - ) - except asyncio.TimeoutError: - # No items for 60s, exit to save resources - self._queue_processor_task = None - return - - try: - # Perform the actual refresh (still using per-credential lock) - async with await self._get_lock(path): - # Re-check if still expired (may have changed since queueing) - creds = self._credentials_cache.get(path) - if creds and not self._is_token_expired(creds): - # No longer expired, mark as available - async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) - continue - - # Perform refresh - if not creds: - creds = await self._load_credentials(path) - await self._refresh_token(path, creds, force=force) - - # SUCCESS: Mark as available again - async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) - - finally: - # Remove from queued set - async with self._queue_tracking_lock: - self._queued_credentials.discard(path) - self._refresh_queue.task_done() - except asyncio.CancelledError: - break - except Exception as e: - lib_logger.error(f"Error in queue processor: {e}") - # Even on error, mark as available (backoff will prevent immediate retry) - if path: - async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) - - async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: - path = creds_or_path if isinstance(creds_or_path, str) else None - - # Get display name from metadata if available, otherwise derive from path - if isinstance(creds_or_path, dict): - display_name = creds_or_path.get("_proxy_metadata", {}).get("display_name", "in-memory object") - else: - display_name = Path(path).name if path else "in-memory object" - - lib_logger.debug(f"Initializing Gemini token for '{display_name}'...") - try: - creds = await self._load_credentials(creds_or_path) if path else creds_or_path - reason = "" - if not creds.get("refresh_token"): - reason = "refresh token is missing" - elif self._is_token_expired(creds): - reason = "token is expired" - - if reason: - if reason == "token is expired" and creds.get("refresh_token"): - try: - return await self._refresh_token(path, creds) - except Exception as e: - lib_logger.warning(f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login.") - - lib_logger.warning(f"Gemini OAuth token for '{display_name}' needs setup: {reason}.") - - # [HEADLESS DETECTION] Check if running in headless environment - is_headless = is_headless_environment() - - auth_code_future = asyncio.get_event_loop().create_future() - server = None - - async def handle_callback(reader, writer): - try: - request_line_bytes = await reader.readline() - if not request_line_bytes: return - path = request_line_bytes.decode('utf-8').strip().split(' ')[1] - while await reader.readline() != b'\r\n': pass - from urllib.parse import urlparse, parse_qs - query_params = parse_qs(urlparse(path).query) - writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n") - if 'code' in query_params: - if not auth_code_future.done(): - auth_code_future.set_result(query_params['code'][0]) - writer.write(b"

Authentication successful!

You can close this window.

") - else: - error = query_params.get('error', ['Unknown error'])[0] - if not auth_code_future.done(): - auth_code_future.set_exception(Exception(f"OAuth failed: {error}")) - writer.write(f"

Authentication Failed

Error: {error}. Please try again.

".encode()) - await writer.drain() - except Exception as e: - lib_logger.error(f"Error in OAuth callback handler: {e}") - finally: - writer.close() - - try: - server = await asyncio.start_server(handle_callback, '127.0.0.1', 8085) - from urllib.parse import urlencode - auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode({ - "client_id": CLIENT_ID, - "redirect_uri": "http://localhost:8085/oauth2callback", - "scope": " ".join(["https://www.googleapis.com/auth/cloud-platform", "https://www.googleapis.com/auth/userinfo.email", "https://www.googleapis.com/auth/userinfo.profile"]), - "access_type": "offline", "response_type": "code", "prompt": "consent" - }) - - # [HEADLESS SUPPORT] Display appropriate instructions - if is_headless: - auth_panel_text = Text.from_markup( - "Running in headless environment (no GUI detected).\n" - "Please open the URL below in a browser on another machine to authorize:\n" - ) - else: - auth_panel_text = Text.from_markup( - "1. Your browser will now open to log in and authorize the application.\n" - "2. If it doesn't open automatically, please open the URL below manually." - ) - - console.print(Panel(auth_panel_text, title=f"Gemini OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue")) - console.print(f"[bold]URL:[/bold] [link={auth_url}]{auth_url}[/link]\n") - - # [HEADLESS SUPPORT] Only attempt browser open if NOT headless - if not is_headless: - try: - webbrowser.open(auth_url) - lib_logger.info("Browser opened successfully for OAuth flow") - except Exception as e: - lib_logger.warning(f"Failed to open browser automatically: {e}. Please open the URL manually.") - - with console.status("[bold green]Waiting for you to complete authentication in the browser...[/bold green]", spinner="dots"): - auth_code = await asyncio.wait_for(auth_code_future, timeout=300) - except asyncio.TimeoutError: - raise Exception("OAuth flow timed out. Please try again.") - finally: - if server: - server.close() - await server.wait_closed() - - lib_logger.info(f"Attempting to exchange authorization code for tokens...") - async with httpx.AsyncClient() as client: - response = await client.post(TOKEN_URI, data={ - "code": auth_code.strip(), "client_id": CLIENT_ID, "client_secret": CLIENT_SECRET, - "redirect_uri": "http://localhost:8085/oauth2callback", "grant_type": "authorization_code" - }) - response.raise_for_status() - token_data = response.json() - # Start with the full token data from the exchange - creds = token_data.copy() - - # Convert 'expires_in' to 'expiry_date' in milliseconds - creds["expiry_date"] = (time.time() + creds.pop("expires_in")) * 1000 - - # Ensure client_id and client_secret are present - creds["client_id"] = CLIENT_ID - creds["client_secret"] = CLIENT_SECRET - - creds["token_uri"] = TOKEN_URI - creds["universe_domain"] = "googleapis.com" - - # Fetch user info and add metadata - user_info_response = await client.get(USER_INFO_URI, headers={"Authorization": f"Bearer {creds['access_token']}"}) - user_info_response.raise_for_status() - user_info = user_info_response.json() - creds["_proxy_metadata"] = { - "email": user_info.get("email"), - "last_check_timestamp": time.time() - } - - if path: - await self._save_credentials(path, creds) - lib_logger.info(f"Gemini OAuth initialized successfully for '{display_name}'.") - return creds - - lib_logger.info(f"Gemini OAuth token at '{display_name}' is valid.") - return creds - except Exception as e: - raise ValueError(f"Failed to initialize Gemini OAuth for '{path}': {e}") - - async def get_auth_header(self, credential_path: str) -> Dict[str, str]: - creds = await self._load_credentials(credential_path) - if self._is_token_expired(creds): - creds = await self._refresh_token(credential_path, creds) - return {"Authorization": f"Bearer {creds['access_token']}"} - - async def get_user_info(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: - path = creds_or_path if isinstance(creds_or_path, str) else None - creds = await self._load_credentials(creds_or_path) if path else creds_or_path - - if path and self._is_token_expired(creds): - creds = await self._refresh_token(path, creds) - - # Prefer locally stored metadata - if creds.get("_proxy_metadata", {}).get("email"): - if path: - creds["_proxy_metadata"]["last_check_timestamp"] = time.time() - await self._save_credentials(path, creds) - return {"email": creds["_proxy_metadata"]["email"]} - - # Fallback to API call if metadata is missing - headers = {"Authorization": f"Bearer {creds['access_token']}"} - async with httpx.AsyncClient() as client: - response = await client.get(USER_INFO_URI, headers=headers) - response.raise_for_status() - user_info = response.json() - - # Save the retrieved info for future use - creds["_proxy_metadata"] = { - "email": user_info.get("email"), - "last_check_timestamp": time.time() - } - if path: - await self._save_credentials(path, creds) - return {"email": user_info.get("email")} \ No newline at end of file +from .google_oauth_base import GoogleOAuthBase + +class GeminiAuthBase(GoogleOAuthBase): + """ + Gemini CLI OAuth2 authentication implementation. + + Inherits all OAuth functionality from GoogleOAuthBase with Gemini-specific configuration. + """ + + CLIENT_ID = "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com" + CLIENT_SECRET = "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" + OAUTH_SCOPES = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + ] + ENV_PREFIX = "GEMINI_CLI" + CALLBACK_PORT = 8085 + CALLBACK_PATH = "/oauth2callback" \ No newline at end of file diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py new file mode 100644 index 00000000..b40e90d1 --- /dev/null +++ b/src/rotator_library/providers/google_oauth_base.py @@ -0,0 +1,653 @@ +# src/rotator_library/providers/google_oauth_base.py + +import os +import webbrowser +from typing import Union, Optional +import json +import time +import asyncio +import logging +from pathlib import Path +from typing import Dict, Any +import tempfile +import shutil + +import httpx +from rich.console import Console +from rich.panel import Panel +from rich.text import Text + +from ..utils.headless_detection import is_headless_environment + +lib_logger = logging.getLogger('rotator_library') + +console = Console() + +class GoogleOAuthBase: + """ + Base class for Google OAuth2 authentication providers. + + Subclasses must override: + - CLIENT_ID: OAuth client ID + - CLIENT_SECRET: OAuth client secret + - OAUTH_SCOPES: List of OAuth scopes + - ENV_PREFIX: Prefix for environment variables (e.g., "GEMINI_CLI", "ANTIGRAVITY") + + Subclasses may optionally override: + - CALLBACK_PORT: Local OAuth callback server port (default: 8085) + - CALLBACK_PATH: OAuth callback path (default: "/oauth2callback") + - REFRESH_EXPIRY_BUFFER_SECONDS: Time buffer before token expiry (default: 30 minutes) + """ + + # Subclasses MUST override these + CLIENT_ID: str = None + CLIENT_SECRET: str = None + OAUTH_SCOPES: list = None + ENV_PREFIX: str = None + + # Subclasses MAY override these + TOKEN_URI: str = "https://oauth2.googleapis.com/token" + USER_INFO_URI: str = "https://www.googleapis.com/oauth2/v1/userinfo" + CALLBACK_PORT: int = 8085 + CALLBACK_PATH: str = "/oauth2callback" + REFRESH_EXPIRY_BUFFER_SECONDS: int = 30 * 60 # 30 minutes + + def __init__(self): + # Validate that subclass has set required attributes + if self.CLIENT_ID is None: + raise NotImplementedError(f"{self.__class__.__name__} must set CLIENT_ID") + if self.CLIENT_SECRET is None: + raise NotImplementedError(f"{self.__class__.__name__} must set CLIENT_SECRET") + if self.OAUTH_SCOPES is None: + raise NotImplementedError(f"{self.__class__.__name__} must set OAUTH_SCOPES") + if self.ENV_PREFIX is None: + raise NotImplementedError(f"{self.__class__.__name__} must set ENV_PREFIX") + + self._credentials_cache: Dict[str, Dict[str, Any]] = {} + self._refresh_locks: Dict[str, asyncio.Lock] = {} + self._locks_lock = asyncio.Lock() # Protects the locks dict from race conditions + # [BACKOFF TRACKING] Track consecutive failures per credential + self._refresh_failures: Dict[str, int] = {} # Track consecutive failures per credential + self._next_refresh_after: Dict[str, float] = {} # Track backoff timers (Unix timestamp) + + # [QUEUE SYSTEM] Sequential refresh processing + self._refresh_queue: asyncio.Queue = asyncio.Queue() + self._queued_credentials: set = set() # Track credentials already in queue + self._unavailable_credentials: set = set() # Mark credentials unavailable during re-auth + self._queue_tracking_lock = asyncio.Lock() # Protects queue sets + self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task + + def _load_from_env(self) -> Optional[Dict[str, Any]]: + """ + Load OAuth credentials from environment variables for stateless deployments. + + Expected environment variables: + - {ENV_PREFIX}_ACCESS_TOKEN (required) + - {ENV_PREFIX}_REFRESH_TOKEN (required) + - {ENV_PREFIX}_EXPIRY_DATE (optional, defaults to 0) + - {ENV_PREFIX}_CLIENT_ID (optional, uses default) + - {ENV_PREFIX}_CLIENT_SECRET (optional, uses default) + - {ENV_PREFIX}_TOKEN_URI (optional, uses default) + - {ENV_PREFIX}_UNIVERSE_DOMAIN (optional, defaults to googleapis.com) + - {ENV_PREFIX}_EMAIL (optional, defaults to "env-user") + - {ENV_PREFIX}_PROJECT_ID (optional) + - {ENV_PREFIX}_TIER (optional) + + Returns: + Dict with credential structure if env vars present, None otherwise + """ + access_token = os.getenv(f"{self.ENV_PREFIX}_ACCESS_TOKEN") + refresh_token = os.getenv(f"{self.ENV_PREFIX}_REFRESH_TOKEN") + + # Both access and refresh tokens are required + if not (access_token and refresh_token): + return None + + lib_logger.debug(f"Loading {self.ENV_PREFIX} credentials from environment variables") + + # Parse expiry_date as float, default to 0 if not present + expiry_str = os.getenv(f"{self.ENV_PREFIX}_EXPIRY_DATE", "0") + try: + expiry_date = float(expiry_str) + except ValueError: + lib_logger.warning(f"Invalid {self.ENV_PREFIX}_EXPIRY_DATE value: {expiry_str}, using 0") + expiry_date = 0 + + creds = { + "access_token": access_token, + "refresh_token": refresh_token, + "expiry_date": expiry_date, + "client_id": os.getenv(f"{self.ENV_PREFIX}_CLIENT_ID", self.CLIENT_ID), + "client_secret": os.getenv(f"{self.ENV_PREFIX}_CLIENT_SECRET", self.CLIENT_SECRET), + "token_uri": os.getenv(f"{self.ENV_PREFIX}_TOKEN_URI", self.TOKEN_URI), + "universe_domain": os.getenv(f"{self.ENV_PREFIX}_UNIVERSE_DOMAIN", "googleapis.com"), + "_proxy_metadata": { + "email": os.getenv(f"{self.ENV_PREFIX}_EMAIL", "env-user"), + "last_check_timestamp": time.time(), + "loaded_from_env": True # Flag to indicate env-based credentials + } + } + + # Add project_id if provided + project_id = os.getenv(f"{self.ENV_PREFIX}_PROJECT_ID") + if project_id: + creds["_proxy_metadata"]["project_id"] = project_id + + # Add tier if provided + tier = os.getenv(f"{self.ENV_PREFIX}_TIER") + if tier: + creds["_proxy_metadata"]["tier"] = tier + + return creds + + async def _load_credentials(self, path: str) -> Dict[str, Any]: + if path in self._credentials_cache: + return self._credentials_cache[path] + + async with await self._get_lock(path): + if path in self._credentials_cache: + return self._credentials_cache[path] + + # First, try loading from environment variables + env_creds = self._load_from_env() + if env_creds: + lib_logger.info(f"Using {self.ENV_PREFIX} credentials from environment variables") + # Cache env-based credentials using the path as key + self._credentials_cache[path] = env_creds + return env_creds + + # Fall back to file-based loading + try: + lib_logger.debug(f"Loading {self.ENV_PREFIX} credentials from file: {path}") + with open(path, 'r') as f: + creds = json.load(f) + # Handle gcloud-style creds file which nest tokens under "credential" + if "credential" in creds: + creds = creds["credential"] + self._credentials_cache[path] = creds + return creds + except FileNotFoundError: + raise IOError(f"{self.ENV_PREFIX} OAuth credential file not found at '{path}'") + except Exception as e: + raise IOError(f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}") + + async def _save_credentials(self, path: str, creds: Dict[str, Any]): + # Don't save to file if credentials were loaded from environment + if creds.get("_proxy_metadata", {}).get("loaded_from_env"): + lib_logger.debug("Credentials loaded from env, skipping file save") + # Still update cache for in-memory consistency + self._credentials_cache[path] = creds + return + + # [ATOMIC WRITE] Use tempfile + move pattern to ensure atomic writes + # This prevents credential corruption if the process is interrupted during write + parent_dir = os.path.dirname(os.path.abspath(path)) + os.makedirs(parent_dir, exist_ok=True) + + tmp_fd = None + tmp_path = None + try: + # Create temp file in same directory as target (ensures same filesystem) + tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json', text=True) + + # Write JSON to temp file + with os.fdopen(tmp_fd, 'w') as f: + json.dump(creds, f, indent=2) + tmp_fd = None # fdopen closes the fd + + # Set secure permissions (0600 = owner read/write only) + try: + os.chmod(tmp_path, 0o600) + except (OSError, AttributeError): + # Windows may not support chmod, ignore + pass + + # Atomic move (overwrites target if it exists) + shutil.move(tmp_path, path) + tmp_path = None # Successfully moved + + # Update cache AFTER successful file write (prevents cache/file inconsistency) + self._credentials_cache[path] = creds + lib_logger.debug(f"Saved updated {self.ENV_PREFIX} OAuth credentials to '{path}' (atomic write).") + + except Exception as e: + lib_logger.error(f"Failed to save updated {self.ENV_PREFIX} OAuth credentials to '{path}': {e}") + # Clean up temp file if it still exists + if tmp_fd is not None: + try: + os.close(tmp_fd) + except: + pass + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except: + pass + raise + + def _is_token_expired(self, creds: Dict[str, Any]) -> bool: + expiry = creds.get("token_expiry") # gcloud format + if not expiry: # gemini-cli format + expiry_timestamp = creds.get("expiry_date", 0) / 1000 + else: + expiry_timestamp = time.mktime(time.strptime(expiry, "%Y-%m-%dT%H:%M:%SZ")) + return expiry_timestamp < time.time() + self.REFRESH_EXPIRY_BUFFER_SECONDS + + async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = False) -> Dict[str, Any]: + async with await self._get_lock(path): + # Skip the expiry check if a refresh is being forced + if not force and not self._is_token_expired(self._credentials_cache.get(path, creds)): + return self._credentials_cache.get(path, creds) + + lib_logger.debug(f"Refreshing {self.ENV_PREFIX} OAuth token for '{Path(path).name}' (forced: {force})...") + refresh_token = creds.get("refresh_token") + if not refresh_token: + raise ValueError("No refresh_token found in credentials file.") + + # [RETRY LOGIC] Implement exponential backoff for transient errors + max_retries = 3 + new_token_data = None + last_error = None + needs_reauth = False + + async with httpx.AsyncClient() as client: + for attempt in range(max_retries): + try: + response = await client.post(self.TOKEN_URI, data={ + "client_id": creds.get("client_id", self.CLIENT_ID), + "client_secret": creds.get("client_secret", self.CLIENT_SECRET), + "refresh_token": refresh_token, + "grant_type": "refresh_token", + }, timeout=30.0) + response.raise_for_status() + new_token_data = response.json() + break # Success, exit retry loop + + except httpx.HTTPStatusError as e: + last_error = e + status_code = e.response.status_code + + # [INVALID GRANT HANDLING] Handle 401/403 by triggering re-authentication + if status_code == 401 or status_code == 403: + lib_logger.warning( + f"Refresh token invalid for '{Path(path).name}' (HTTP {status_code}). " + f"Token may have been revoked or expired. Starting re-authentication..." + ) + needs_reauth = True + break # Exit retry loop to trigger re-auth + + elif status_code == 429: + # Rate limit - honor Retry-After header if present + retry_after = int(e.response.headers.get("Retry-After", 60)) + lib_logger.warning(f"Rate limited (HTTP 429), retry after {retry_after}s") + if attempt < max_retries - 1: + await asyncio.sleep(retry_after) + continue + raise + + elif status_code >= 500 and status_code < 600: + # Server error - retry with exponential backoff + if attempt < max_retries - 1: + wait_time = 2 ** attempt # 1s, 2s, 4s + lib_logger.warning(f"Server error (HTTP {status_code}), retry {attempt + 1}/{max_retries} in {wait_time}s") + await asyncio.sleep(wait_time) + continue + raise # Final attempt failed + + else: + # Other errors - don't retry + raise + + except (httpx.RequestError, httpx.TimeoutException) as e: + # Network errors - retry with backoff + last_error = e + if attempt < max_retries - 1: + wait_time = 2 ** attempt + lib_logger.warning(f"Network error during refresh: {e}, retry {attempt + 1}/{max_retries} in {wait_time}s") + await asyncio.sleep(wait_time) + continue + raise + + # [INVALID GRANT RE-AUTH] Trigger OAuth flow if refresh token is invalid + if needs_reauth: + lib_logger.info(f"Starting re-authentication for '{Path(path).name}'...") + try: + # Call initialize_token to trigger OAuth flow + new_creds = await self.initialize_token(path) + return new_creds + except Exception as reauth_error: + lib_logger.error(f"Re-authentication failed for '{Path(path).name}': {reauth_error}") + raise ValueError(f"Refresh token invalid and re-authentication failed: {reauth_error}") + + # If we exhausted retries without success + if new_token_data is None: + raise last_error or Exception("Token refresh failed after all retries") + + # [FIX 1] Update OAuth token fields from response + creds["access_token"] = new_token_data["access_token"] + expiry_timestamp = time.time() + new_token_data["expires_in"] + creds["expiry_date"] = expiry_timestamp * 1000 # gemini-cli format + + # [FIX 2] Update refresh_token if server provided a new one (rare but possible with Google OAuth) + if "refresh_token" in new_token_data: + creds["refresh_token"] = new_token_data["refresh_token"] + + # [FIX 3] Ensure all required OAuth client fields are present (restore if missing) + if "client_id" not in creds or not creds["client_id"]: + creds["client_id"] = self.CLIENT_ID + if "client_secret" not in creds or not creds["client_secret"]: + creds["client_secret"] = self.CLIENT_SECRET + if "token_uri" not in creds or not creds["token_uri"]: + creds["token_uri"] = self.TOKEN_URI + if "universe_domain" not in creds or not creds["universe_domain"]: + creds["universe_domain"] = "googleapis.com" + + # [FIX 4] Add scopes array if missing + if "scopes" not in creds: + creds["scopes"] = self.OAUTH_SCOPES + + # [FIX 5] Ensure _proxy_metadata exists and update timestamp + if "_proxy_metadata" not in creds: + creds["_proxy_metadata"] = {} + creds["_proxy_metadata"]["last_check_timestamp"] = time.time() + + # [VALIDATION] Verify refreshed credentials have all required fields + required_fields = ["access_token", "refresh_token", "client_id", "client_secret", "token_uri"] + missing_fields = [field for field in required_fields if not creds.get(field)] + if missing_fields: + raise ValueError(f"Refreshed credentials missing required fields: {missing_fields}") + + # [VALIDATION] Optional: Test that the refreshed token is actually usable + try: + async with httpx.AsyncClient() as client: + test_response = await client.get( + self.USER_INFO_URI, + headers={"Authorization": f"Bearer {creds['access_token']}"}, + timeout=5.0 + ) + test_response.raise_for_status() + lib_logger.debug(f"Token validation successful for '{Path(path).name}'") + except Exception as e: + lib_logger.warning(f"Refreshed token validation failed for '{Path(path).name}': {e}") + # Don't fail the refresh - the token might still work for other endpoints + # But log it for debugging purposes + + await self._save_credentials(path, creds) + lib_logger.debug(f"Successfully refreshed {self.ENV_PREFIX} OAuth token for '{Path(path).name}'.") + return creds + + async def proactively_refresh(self, credential_path: str): + """Proactively refresh a credential by queueing it for refresh.""" + creds = await self._load_credentials(credential_path) + if self._is_token_expired(creds): + # Queue for refresh with needs_reauth=False (automated refresh) + await self._queue_refresh(credential_path, force=False, needs_reauth=False) + + async def _get_lock(self, path: str) -> asyncio.Lock: + # [FIX RACE CONDITION] Protect lock creation with a master lock + # This prevents TOCTOU bug where multiple coroutines check and create simultaneously + async with self._locks_lock: + if path not in self._refresh_locks: + self._refresh_locks[path] = asyncio.Lock() + return self._refresh_locks[path] + + def is_credential_available(self, path: str) -> bool: + """Check if a credential is available for rotation (not queued/refreshing).""" + return path not in self._unavailable_credentials + + async def _ensure_queue_processor_running(self): + """Lazily starts the queue processor if not already running.""" + if self._queue_processor_task is None or self._queue_processor_task.done(): + self._queue_processor_task = asyncio.create_task(self._process_refresh_queue()) + + async def _queue_refresh(self, path: str, force: bool = False, needs_reauth: bool = False): + """Add a credential to the refresh queue if not already queued. + + Args: + path: Credential file path + force: Force refresh even if not expired + needs_reauth: True if full re-authentication needed (bypasses backoff) + """ + # IMPORTANT: Only check backoff for simple automated refreshes + # Re-authentication (interactive OAuth) should BYPASS backoff since it needs user input + if not needs_reauth: + now = time.time() + if path in self._next_refresh_after: + backoff_until = self._next_refresh_after[path] + if now < backoff_until: + # Credential is in backoff for automated refresh, do not queue + remaining = int(backoff_until - now) + lib_logger.debug(f"Skipping automated refresh for '{Path(path).name}' (in backoff for {remaining}s)") + return + + async with self._queue_tracking_lock: + if path not in self._queued_credentials: + self._queued_credentials.add(path) + self._unavailable_credentials.add(path) # Mark as unavailable + await self._refresh_queue.put((path, force, needs_reauth)) + await self._ensure_queue_processor_running() + + async def _process_refresh_queue(self): + """Background worker that processes refresh requests sequentially.""" + while True: + path = None + try: + # Wait for an item with timeout to allow graceful shutdown + try: + path, force, needs_reauth = await asyncio.wait_for( + self._refresh_queue.get(), + timeout=60.0 + ) + except asyncio.TimeoutError: + # No items for 60s, exit to save resources + self._queue_processor_task = None + return + + try: + # Perform the actual refresh (still using per-credential lock) + async with await self._get_lock(path): + # Re-check if still expired (may have changed since queueing) + creds = self._credentials_cache.get(path) + if creds and not self._is_token_expired(creds): + # No longer expired, mark as available + async with self._queue_tracking_lock: + self._unavailable_credentials.discard(path) + continue + + # Perform refresh + if not creds: + creds = await self._load_credentials(path) + await self._refresh_token(path, creds, force=force) + + # SUCCESS: Mark as available again + async with self._queue_tracking_lock: + self._unavailable_credentials.discard(path) + + finally: + # Remove from queued set + async with self._queue_tracking_lock: + self._queued_credentials.discard(path) + self._refresh_queue.task_done() + except asyncio.CancelledError: + break + except Exception as e: + lib_logger.error(f"Error in queue processor: {e}") + # Even on error, mark as available (backoff will prevent immediate retry) + if path: + async with self._queue_tracking_lock: + self._unavailable_credentials.discard(path) + + async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + path = creds_or_path if isinstance(creds_or_path, str) else None + + # Get display name from metadata if available, otherwise derive from path + if isinstance(creds_or_path, dict): + display_name = creds_or_path.get("_proxy_metadata", {}).get("display_name", "in-memory object") + else: + display_name = Path(path).name if path else "in-memory object" + + lib_logger.debug(f"Initializing {self.ENV_PREFIX} token for '{display_name}'...") + try: + creds = await self._load_credentials(creds_or_path) if path else creds_or_path + reason = "" + if not creds.get("refresh_token"): + reason = "refresh token is missing" + elif self._is_token_expired(creds): + reason = "token is expired" + + if reason: + if reason == "token is expired" and creds.get("refresh_token"): + try: + return await self._refresh_token(path, creds) + except Exception as e: + lib_logger.warning(f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login.") + + lib_logger.warning(f"{self.ENV_PREFIX} OAuth token for '{display_name}' needs setup: {reason}.") + + # [HEADLESS DETECTION] Check if running in headless environment + is_headless = is_headless_environment() + + auth_code_future = asyncio.get_event_loop().create_future() + server = None + + async def handle_callback(reader, writer): + try: + request_line_bytes = await reader.readline() + if not request_line_bytes: return + path_str = request_line_bytes.decode('utf-8').strip().split(' ')[1] + while await reader.readline() != b'\r\n': pass + from urllib.parse import urlparse, parse_qs + query_params = parse_qs(urlparse(path_str).query) + writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n") + if 'code' in query_params: + if not auth_code_future.done(): + auth_code_future.set_result(query_params['code'][0]) + writer.write(b"

Authentication successful!

You can close this window.

") + else: + error = query_params.get('error', ['Unknown error'])[0] + if not auth_code_future.done(): + auth_code_future.set_exception(Exception(f"OAuth failed: {error}")) + writer.write(f"

Authentication Failed

Error: {error}. Please try again.

".encode()) + await writer.drain() + except Exception as e: + lib_logger.error(f"Error in OAuth callback handler: {e}") + finally: + writer.close() + + try: + server = await asyncio.start_server(handle_callback, '127.0.0.1', self.CALLBACK_PORT) + from urllib.parse import urlencode + auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode({ + "client_id": self.CLIENT_ID, + "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", + "scope": " ".join(self.OAUTH_SCOPES), + "access_type": "offline", "response_type": "code", "prompt": "consent" + }) + + # [HEADLESS SUPPORT] Display appropriate instructions + if is_headless: + auth_panel_text = Text.from_markup( + "Running in headless environment (no GUI detected).\n" + "Please open the URL below in a browser on another machine to authorize:\n" + ) + else: + auth_panel_text = Text.from_markup( + "1. Your browser will now open to log in and authorize the application.\n" + "2. If it doesn't open automatically, please open the URL below manually." + ) + + console.print(Panel(auth_panel_text, title=f"{self.ENV_PREFIX} OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue")) + console.print(f"[bold]URL:[/bold] [link={auth_url}]{auth_url}[/link]\n") + + # [HEADLESS SUPPORT] Only attempt browser open if NOT headless + if not is_headless: + try: + webbrowser.open(auth_url) + lib_logger.info("Browser opened successfully for OAuth flow") + except Exception as e: + lib_logger.warning(f"Failed to open browser automatically: {e}. Please open the URL manually.") + + with console.status(f"[bold green]Waiting for you to complete authentication in the browser...[/bold green]", spinner="dots"): + auth_code = await asyncio.wait_for(auth_code_future, timeout=300) + except asyncio.TimeoutError: + raise Exception("OAuth flow timed out. Please try again.") + finally: + if server: + server.close() + await server.wait_closed() + + lib_logger.info(f"Attempting to exchange authorization code for tokens...") + async with httpx.AsyncClient() as client: + response = await client.post(self.TOKEN_URI, data={ + "code": auth_code.strip(), "client_id": self.CLIENT_ID, "client_secret": self.CLIENT_SECRET, + "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", "grant_type": "authorization_code" + }) + response.raise_for_status() + token_data = response.json() + # Start with the full token data from the exchange + creds = token_data.copy() + + # Convert 'expires_in' to 'expiry_date' in milliseconds + creds["expiry_date"] = (time.time() + creds.pop("expires_in")) * 1000 + + # Ensure client_id and client_secret are present + creds["client_id"] = self.CLIENT_ID + creds["client_secret"] = self.CLIENT_SECRET + + creds["token_uri"] = self.TOKEN_URI + creds["universe_domain"] = "googleapis.com" + + # Fetch user info and add metadata + user_info_response = await client.get(self.USER_INFO_URI, headers={"Authorization": f"Bearer {creds['access_token']}"}) + user_info_response.raise_for_status() + user_info = user_info_response.json() + creds["_proxy_metadata"] = { + "email": user_info.get("email"), + "last_check_timestamp": time.time() + } + + if path: + await self._save_credentials(path, creds) + lib_logger.info(f"{self.ENV_PREFIX} OAuth initialized successfully for '{display_name}'.") + return creds + + lib_logger.info(f"{self.ENV_PREFIX} OAuth token at '{display_name}' is valid.") + return creds + except Exception as e: + raise ValueError(f"Failed to initialize {self.ENV_PREFIX} OAuth for '{path}': {e}") + + async def get_auth_header(self, credential_path: str) -> Dict[str, str]: + creds = await self._load_credentials(credential_path) + if self._is_token_expired(creds): + creds = await self._refresh_token(credential_path, creds) + return {"Authorization": f"Bearer {creds['access_token']}"} + + async def get_user_info(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + path = creds_or_path if isinstance(creds_or_path, str) else None + creds = await self._load_credentials(creds_or_path) if path else creds_or_path + + if path and self._is_token_expired(creds): + creds = await self._refresh_token(path, creds) + + # Prefer locally stored metadata + if creds.get("_proxy_metadata", {}).get("email"): + if path: + creds["_proxy_metadata"]["last_check_timestamp"] = time.time() + await self._save_credentials(path, creds) + return {"email": creds["_proxy_metadata"]["email"]} + + # Fallback to API call if metadata is missing + headers = {"Authorization": f"Bearer {creds['access_token']}"} + async with httpx.AsyncClient() as client: + response = await client.get(self.USER_INFO_URI, headers=headers) + response.raise_for_status() + user_info = response.json() + + # Save the retrieved info for future use + creds["_proxy_metadata"] = { + "email": user_info.get("email"), + "last_check_timestamp": time.time() + } + if path: + await self._save_credentials(path, creds) + return {"email": user_info.get("email")} From 77bfd5f778a185311a25e9a5ed47d5a1406db518 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 18:18:44 +0100 Subject: [PATCH 009/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20add=20?= =?UTF-8?q?dynamic=20model=20discovery=20toggle=20and=20get=5Fvalid=5Ftoke?= =?UTF-8?q?n=20helper?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add opt-in dynamic model discovery controlled by ANTIGRAVITY_ENABLE_DYNAMIC_MODELS (default: false) to avoid relying on an unstable endpoint. When disabled, the provider returns the hardcoded model list; when enabled, it attempts to fetch models from the API and applies alias mappings. Add clear logging for enabled/disabled states and dynamic discovery results. Also introduce an async get_valid_token helper that loads credentials, refreshes expired tokens, and returns a valid access token for OAuth-style credential paths. - New env var: ANTIGRAVITY_ENABLE_DYNAMIC_MODELS (false by default) - Dynamic discovery returns discovered models prefixed with "antigravity/" - Hardcoded fallback now returns names prefixed with "antigravity/" - Added logs to indicate discovery mode and failures - Added async get_valid_token(credential_identifier) to centralize token refresh/load BREAKING CHANGE: Model names returned by the provider are now namespaced with the "antigravity/" prefix (e.g., "antigravity/xyz"). Update consumers to handle the new prefixed names or strip the prefix as needed. Dynamic discovery is disabled by default; enable it with ANTIGRAVITY_ENABLE_DYNAMIC_MODELS=true if desired. --- .../providers/antigravity_provider.py | 49 ++++++++++++++++--- 1 file changed, 42 insertions(+), 7 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index ed30d417..1618640c 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -219,6 +219,12 @@ def __init__(self): "true" # Default ON for testing ).lower() in ("true", "1", "yes") + # Check if dynamic model discovery is enabled (default: OFF due to endpoint instability) + self._enable_dynamic_model_discovery = os.getenv( + "ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", + "false" # Default OFF - use hardcoded list + ).lower() in ("true", "1", "yes") + if self._preserve_signatures_in_client: lib_logger.info("Antigravity: thoughtSignature client passthrough ENABLED") else: @@ -228,6 +234,11 @@ def __init__(self): lib_logger.info(f"Antigravity: thoughtSignature server-side cache ENABLED (TTL: {cache_ttl}s)") else: lib_logger.info("Antigravity: thoughtSignature server-side cache DISABLED") + + if self._enable_dynamic_model_discovery: + lib_logger.info("Antigravity: Dynamic model discovery ENABLED (may fail if endpoint unavailable)") + else: + lib_logger.info("Antigravity: Dynamic model discovery DISABLED (using hardcoded model list)") # ============================================================================ # MODEL ALIAS SYSTEM @@ -938,11 +949,26 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D response["usage"] = usage return response - + # ============================================================================ # PROVIDER INTERFACE IMPLEMENTATION # ============================================================================ + async def get_valid_token(self, credential_identifier: str) -> str: + """ + Get a valid access token for the credential. + + Args: + credential_identifier: Credential file path or "env" + + Returns: + Access token string + """ + creds = await self._load_credentials(credential_identifier) + if self._is_token_expired(creds): + creds = await self._refresh_token(credential_identifier, creds) + return creds['access_token'] + def has_custom_logic(self) -> bool: """Antigravity uses custom translation logic.""" return True @@ -964,8 +990,11 @@ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str] """ Fetch available models from Antigravity. - For Antigravity, we use the fetchAvailableModels endpoint and apply - alias mapping to convert internal names to public names. + For Antigravity, we can optionally use the fetchAvailableModels endpoint and apply + alias mapping to convert internal names to public names. However, this endpoint is + often unavailable (404), so dynamic discovery is disabled by default. + + Set ANTIGRAVITY_ENABLE_DYNAMIC_MODELS=true to enable dynamic discovery. Args: api_key: Credential path (not a traditional API key) @@ -974,6 +1003,12 @@ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str] Returns: List of public model names """ + # If dynamic discovery is disabled, immediately return hardcoded list + if not self._enable_dynamic_model_discovery: + lib_logger.debug("Using hardcoded Antigravity model list (dynamic discovery disabled)") + return [f"antigravity/{m}" for m in HARDCODED_MODELS] + + # Dynamic discovery enabled - attempt to fetch from API credential_path = api_key # For OAuth providers, this is the credential path try: @@ -1013,18 +1048,18 @@ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str] if internal_name: public_name = self._model_name_to_alias(internal_name) if public_name: # Skip excluded models (empty string) - models.append(public_name) + models.append(f"antigravity/{public_name}") if models: - lib_logger.info(f"Discovered {len(models)} Antigravity models") + lib_logger.info(f"Discovered {len(models)} Antigravity models via dynamic discovery") return models else: lib_logger.warning("No models returned from Antigravity, using hardcoded list") - return HARDCODED_MODELS + return [f"antigravity/{m}" for m in HARDCODED_MODELS] except Exception as e: lib_logger.warning(f"Failed to fetch Antigravity models: {e}, using hardcoded list") - return HARDCODED_MODELS + return [f"antigravity/{m}" for m in HARDCODED_MODELS] async def acompletion( self, From c6478edb3f43b87b3683239529f67c46ba060167 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 22:55:44 +0100 Subject: [PATCH 010/221] =?UTF-8?q?fix(providers):=20=F0=9F=90=9B=20fix=20?= =?UTF-8?q?antigravity=20provider=20compatibility=20and=20async=20credenti?= =?UTF-8?q?al=20save?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Handle system prompt content as either string or list and strip Claude-specific cache_control fields to avoid 400 errors - Safely parse tool content (JSON or raw) and wrap function responses consistently - Treat merged function response role as "user" to match Antigravity expectations - Add tool_call index for OpenAI streaming format and track index for parallel tool calls - Strip provider prefix from model names and add streaming query param (?alt=sse) when streaming - Include Host and User-Agent headers, set Accept based on streaming, and log error response bodies for easier debugging - Convert OpenAI-style chunks into litellm.ModelResponse objects before yielding in stream handler - Make credential persistence in Gemini CLI provider async (await _save_credentials) --- .../providers/antigravity_provider.py | 82 +++++++++++++++---- .../providers/gemini_cli_provider.py | 4 +- 2 files changed, 70 insertions(+), 16 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 1618640c..e19d9e1f 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -359,10 +359,25 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tup if messages and messages[0].get('role') == 'system': system_prompt_content = messages.pop(0).get('content', '') if system_prompt_content: - system_instruction = { - "role": "user", - "parts": [{"text": system_prompt_content}] - } + # Handle both string and list-based system content + system_parts = [] + if isinstance(system_prompt_content, str): + system_parts.append({"text": system_prompt_content}) + elif isinstance(system_prompt_content, list): + # Multi-part system content (strip cache_control) + for item in system_prompt_content: + if item.get("type") == "text": + text = item.get("text", "") + if text: + # Skip cache_control - Claude-specific field + system_parts.append({"text": text}) + + if system_parts: + system_instruction = { + "role": "user", + "parts": system_parts + } + # Build tool call ID to name mapping tool_call_id_to_name = {} @@ -390,6 +405,8 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tup if item.get("type") == "text": text = item.get("text", "") if text: + # Strip Claude-specific cache_control field + # This field causes 400 errors with Antigravity parts.append({"text": text}) elif item.get("type") == "image_url": # Handle image data URLs @@ -459,15 +476,18 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tup function_name = tool_call_id_to_name.get(tool_call_id, "unknown_function") tool_content = msg.get("content", "{}") + # Parse tool content - if it's JSON, use parsed value; otherwise use as-is try: - response_data = json.loads(tool_content) + parsed_content = json.loads(tool_content) except (json.JSONDecodeError, TypeError): - response_data = {"result": tool_content} - + parsed_content = tool_content + parts.append({ "functionResponse": { "name": function_name, - "response": response_data + "response": { + "result": parsed_content + } } }) @@ -620,7 +640,7 @@ def _fix_tool_response_grouping(self, contents: List[Dict[str, Any]]) -> List[Di # Create merged function response content function_response_content = { "parts": group_responses, - "role": "function" # Changed from tool + "role": "user" } new_contents.append(function_response_content) @@ -659,7 +679,7 @@ def _fix_tool_response_grouping(self, contents: List[Dict[str, Any]]) -> List[Di function_response_content = { "parts": group_responses, - "role": "function" + "role": "user" } new_contents.append(function_response_content) @@ -834,6 +854,7 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D # Track if we've seen a signature yet (for parallel tool call handling) # Per Gemini 3 spec: only FIRST tool call in parallel gets signature first_signature_seen = False + tool_call_index = 0 # Track index for OpenAI streaming format for part in content_parts: has_function_call = "functionCall" in part @@ -861,11 +882,13 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D tool_call = { "id": tool_call_id, "type": "function", + "index": tool_call_index, # REQUIRED for OpenAI streaming format "function": { "name": func_call.get("name", ""), "arguments": json.dumps(func_call.get("args", {})) } } + tool_call_index += 1 # Increment for next tool call # Handle thoughtSignature if present # CRITICAL FIX: Cache and passthrough are INDEPENDENT toggles @@ -1084,6 +1107,11 @@ async def acompletion( """ # Extract key parameters model = kwargs.get("model", "gemini-2.5-pro") + + # Strip provider prefix from model name (e.g., "antigravity/claude-sonnet-4-5-thinking" -> "claude-sonnet-4-5-thinking") + if "/" in model: + model = model.split("/")[-1] + messages = kwargs.get("messages", []) stream = kwargs.get("stream", False) credential_path = kwargs.pop("credential_identifier", kwargs.get("api_key", "")) @@ -1168,12 +1196,28 @@ async def acompletion( endpoint = ":streamGenerateContent" if stream else ":generateContent" url = f"{base_url}{endpoint}" - + + # Add query parameter for streaming (required by Antigravity API) + if stream: + url = f"{url}?alt=sse" + + # Extract host from base_url for Host header (required by Google's API) + from urllib.parse import urlparse + parsed_url = urlparse(base_url) + host = parsed_url.netloc if parsed_url.netloc else base_url.replace("https://", "").replace("http://", "").rstrip("/") + headers = { "Authorization": f"Bearer {access_token}", - "Content-Type": "application/json" + "Content-Type": "application/json", + "Host": host, # CRITICAL: Required by Antigravity API + "User-Agent": "antigravity/1.11.5" # Match Go implementation } - + + if stream: + headers["Accept"] = "text/event-stream" + else: + headers["Accept"] = "application/json" + lib_logger.debug(f"Antigravity request to: {url}") try: @@ -1231,6 +1275,14 @@ async def _handle_streaming( ) -> AsyncGenerator[litellm.ModelResponse, None]: """Handle streaming completion.""" async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response: + # Log error response body for debugging if request failed + if response.status_code >= 400: + try: + error_body = await response.aread() + lib_logger.error(f"Antigravity API error {response.status_code}: {error_body.decode('utf-8', errors='replace')}") + except Exception as e: + lib_logger.error(f"Failed to read error response body: {e}") + response.raise_for_status() async for line in response.aiter_lines(): @@ -1252,7 +1304,9 @@ async def _handle_streaming( # Convert to OpenAI format openai_chunk = self._gemini_to_openai_chunk(gemini_chunk, model) - yield openai_chunk + # Convert dict to ModelResponse object + model_response = litellm.ModelResponse(**openai_chunk) + yield model_response except json.JSONDecodeError: if file_logger: file_logger.log_error(f"Failed to parse chunk: {data_str[:100]}") diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index fe3980fd..140da2ce 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -311,13 +311,13 @@ async def _persist_project_metadata(self, credential_path: str, project_id: str, # Update metadata if "_proxy_metadata" not in creds: creds["_proxy_metadata"] = {} - + creds["_proxy_metadata"]["project_id"] = project_id if tier: creds["_proxy_metadata"]["tier"] = tier # Save back using the existing save method (handles atomic writes and permissions) - self._save_credentials(credential_path, creds) + await self._save_credentials(credential_path, creds) lib_logger.debug(f"Persisted project_id and tier to credential file: {credential_path}") except Exception as e: From 264959a7f8da294bd420b4fc1f29ecf799fa3138 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 23 Nov 2025 23:30:35 +0100 Subject: [PATCH 011/221] =?UTF-8?q?fix(antigravity):=20=F0=9F=90=9B=20conv?= =?UTF-8?q?ert=20tool=20parameters=20to=20parametersJsonSchema=20and=20str?= =?UTF-8?q?ip=20unsupported=20fields?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove dependency on _build_vertex_schema and align tool handling with the Go reference implementation. For function-type tools, build a function declaration with name, description, and a parametersJsonSchema field: - copy parameters when present and remove OpenAI-specific keys (`$schema`, `strict`); - default to an empty object schema when parameters are missing; - avoid mutating the original parameters and embed the declaration in `functionDeclarations`. This ensures Antigravity-compatible tool payloads and fixes schema/compatibility issues when passing tool definitions. --- .../providers/antigravity_provider.py | 113 +++++++++++++----- 1 file changed, 83 insertions(+), 30 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index e19d9e1f..74be6298 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -18,7 +18,8 @@ from ..model_definitions import ModelDefinitions import litellm from litellm.exceptions import RateLimitError -from litellm.llms.vertex_ai.common_utils import _build_vertex_schema +# Removed: from litellm.llms.vertex_ai.common_utils import _build_vertex_schema +# Using direct parameter passthrough instead, matching Go reference implementation lib_logger = logging.getLogger('rotator_library') @@ -302,6 +303,41 @@ def _is_gemini_3_model(self, model: str) -> bool: internal_model = self._alias_to_model_name(model) return internal_model.startswith("gemini-3-") or model.startswith("gemini-3-") + @staticmethod + def _normalize_json_schema(schema: Any) -> Any: + """ + Normalize JSON Schema for Proto-based Antigravity API. + + The Proto-based API doesn't support array values for the 'type' field. + This function converts `"type": ["string", "null"]` → `"type": "string"`. + + Args: + schema: JSON schema object (dict, list, or primitive) + + Returns: + Normalized schema + """ + if isinstance(schema, dict): + # Make a copy to avoid modifying the original + normalized = {} + for key, value in schema.items(): + if key == "type" and isinstance(value, list): + # Convert array type to single type + # Take the first non-"null" type, or the first type if all are "null" + non_null_types = [t for t in value if t != "null"] + normalized[key] = non_null_types[0] if non_null_types else value[0] + else: + # Recursively normalize nested structures + normalized[key] = AntigravityProvider._normalize_json_schema(value) + return normalized + elif isinstance(schema, list): + # Recursively normalize list items + return [AntigravityProvider._normalize_json_schema(item) for item in schema] + else: + # Primitive value - return as-is + return schema + + # ============================================================================ # RANDOM ID GENERATION # ============================================================================ @@ -750,30 +786,25 @@ def _transform_to_antigravity_format( part["thoughtSignature"] = "skip_thought_signature_validator" # If thoughtSignature already exists, preserve it (important for Gemini 3) - # ======================================================================== - # IMPORTANT: CLAUDE SCHEMA HANDLING - REQUIRES INVESTIGATION - # ======================================================================== - # WARNING: This code block may be incorrect! - # - # INVESTIGATION REQUIRED BEFORE MAKING CHANGES: - # - Test Claude model access through Antigravity with tools - # - Verify whether parametersJsonSchema → parameters conversion is needed - # - The Go reference suggests Antigravity expects parametersJsonSchema for ALL models - # - # Current behavior: Converts parametersJsonSchema back to parameters for Claude models - # Potential issue: Antigravity may actually expect parametersJsonSchema for Claude too - # - # DO NOT MODIFY without first confirming actual API behavior! - # ======================================================================== + # 7. CRITICAL: Claude-specific tool schema transformation + # Claude models need 'parameters' NOT 'parametersJsonSchema' (opposite of Gemini) + # Reference: Go implementation antigravity_executor.go lines 672-684 if internal_model.startswith("claude-sonnet-"): - # For Claude models, convert parametersJsonSchema back to parameters - for tool in antigravity_payload["request"].get("tools", []): - for func_decl in tool.get("functionDeclarations", []): + tools = antigravity_payload["request"].get("tools", []) + for tool_idx, tool in enumerate(tools): + function_declarations = tool.get("functionDeclarations", []) + for func_idx, func_decl in enumerate(function_declarations): if "parametersJsonSchema" in func_decl: - func_decl["parameters"] = func_decl.pop("parametersJsonSchema") - # Remove $schema if present - if "parameters" in func_decl and "$schema" in func_decl["parameters"]: - del func_decl["parameters"]["$schema"] + # Convert parametersJsonSchema → parameters for Claude + params = func_decl["parametersJsonSchema"] + + # Remove $schema if present (Claude doesn't support it) + if isinstance(params, dict): + params.pop("$schema", None) + + # Set as 'parameters' and remove 'parametersJsonSchema' + antigravity_payload["request"]["tools"][tool_idx]["functionDeclarations"][func_idx]["parameters"] = params + del antigravity_payload["request"]["tools"][tool_idx]["functionDeclarations"][func_idx]["parametersJsonSchema"] return antigravity_payload @@ -1167,20 +1198,42 @@ async def acompletion( if generation_config: gemini_cli_payload["generationConfig"] = generation_config - # Add tools + # Add tools - using Go reference implementation approach + # Go code (line 298-328): renames 'parameters' -> 'parametersJsonSchema' and removes 'strict' if tools: gemini_tools = [] for tool in tools: if tool.get("type") == "function": func = tool.get("function", {}) - schema = _build_vertex_schema(parameters=func.get("parameters", {})) + + # Get parameters dict (may be missing) + parameters = func.get("parameters") + + # Build function declaration + func_decl = { + "name": func.get("name", ""), + "description": func.get("description", "") + } + + # Handle parameters -> parametersJsonSchema conversion (matching Go) + if parameters and isinstance(parameters, dict): + # Make a copy to avoid modifying original + schema = dict(parameters) + # Remove OpenAI-specific fields that Antigravity doesn't support + schema.pop("$schema", None) + schema.pop("strict", None) + func_decl["parametersJsonSchema"] = schema + else: + # No parameters provided - set default empty schema (matching Go lines 318-323) + func_decl["parametersJsonSchema"] = { + "type": "object", + "properties": {} + } + gemini_tools.append({ - "functionDeclarations": [{ - "name": func.get("name", ""), - "description": func.get("description", ""), - "parametersJsonSchema": schema - }] + "functionDeclarations": [func_decl] }) + if gemini_tools: gemini_cli_payload["tools"] = gemini_tools From 4ff1edfd9e7c2bb8a3bc4a3c83aebed0ba848d57 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 24 Nov 2025 00:16:26 +0100 Subject: [PATCH 012/221] =?UTF-8?q?fix(providers):=20=F0=9F=90=9B=20normal?= =?UTF-8?q?ize=20JSON=20Schema=20types,=20clean=20Claude=20tool=20schemas,?= =?UTF-8?q?=20and=20fix=20Gemini=20tool=20conversion?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename _normalize_json_schema → _normalize_type_arrays and convert JSON Schema "type" arrays (e.g. ["string","null"]) to a single non-null type to avoid protobuf "non-repeating" errors. - Add recursive Claude-specific schema cleaner and rename parametersJsonSchema → parameters for claude-sonnet-* models, stripping incompatible fields that break Claude validation. - Ensure thoughtSignature preservation logic remains with proper first-seen handling. - Inline generation of project/request IDs when fetching models. - Replace Vertex helper usage when building Gemini tool declarations: copy/clean parameters, set a safe default parametersJsonSchema, and call _normalize_type_arrays for compatibility. --- .../providers/antigravity_provider.py | 125 ++++++++++-------- 1 file changed, 71 insertions(+), 54 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 74be6298..f7756f38 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -18,8 +18,6 @@ from ..model_definitions import ModelDefinitions import litellm from litellm.exceptions import RateLimitError -# Removed: from litellm.llms.vertex_ai.common_utils import _build_vertex_schema -# Using direct parameter passthrough instead, matching Go reference implementation lib_logger = logging.getLogger('rotator_library') @@ -304,40 +302,26 @@ def _is_gemini_3_model(self, model: str) -> bool: return internal_model.startswith("gemini-3-") or model.startswith("gemini-3-") @staticmethod - def _normalize_json_schema(schema: Any) -> Any: + def _normalize_type_arrays(schema: Any) -> Any: """ - Normalize JSON Schema for Proto-based Antigravity API. - - The Proto-based API doesn't support array values for the 'type' field. - This function converts `"type": ["string", "null"]` → `"type": "string"`. - - Args: - schema: JSON schema object (dict, list, or primitive) - - Returns: - Normalized schema + Normalize type arrays in JSON Schema for Proto-based Antigravity API. + Converts `"type": ["string", "null"]` → `"type": "string"`. """ if isinstance(schema, dict): - # Make a copy to avoid modifying the original normalized = {} for key, value in schema.items(): if key == "type" and isinstance(value, list): - # Convert array type to single type - # Take the first non-"null" type, or the first type if all are "null" + # Take first non-null type non_null_types = [t for t in value if t != "null"] normalized[key] = non_null_types[0] if non_null_types else value[0] else: - # Recursively normalize nested structures - normalized[key] = AntigravityProvider._normalize_json_schema(value) + normalized[key] = AntigravityProvider._normalize_type_arrays(value) return normalized elif isinstance(schema, list): - # Recursively normalize list items - return [AntigravityProvider._normalize_json_schema(item) for item in schema] + return [AntigravityProvider._normalize_type_arrays(item) for item in schema] else: - # Primitive value - return as-is return schema - # ============================================================================ # RANDOM ID GENERATION # ============================================================================ @@ -371,9 +355,7 @@ def generate_project_id() -> str: def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: """ Transform OpenAI messages to Gemini CLI format. - Reused from GeminiCliProvider with modifications for Antigravity. - - UPDATED: Now handles thoughtSignature preservation with 3-tier fallback: + Handles thoughtSignature preservation with 3-tier fallback: 1. Use client-provided signature (if present) 2. Fall back to server-side cache 3. Use bypass constant as last resort @@ -784,27 +766,53 @@ def _transform_to_antigravity_format( # Add signature to function calls OR preserve if already exists if "functionCall" in part and "thoughtSignature" not in part: part["thoughtSignature"] = "skip_thought_signature_validator" - # If thoughtSignature already exists, preserve it (important for Gemini 3) - # 7. CRITICAL: Claude-specific tool schema transformation - # Claude models need 'parameters' NOT 'parametersJsonSchema' (opposite of Gemini) + # 7. CLAUDE-SPECIFIC TOOL SCHEMA TRANSFORMATION # Reference: Go implementation antigravity_executor.go lines 672-684 + # For Claude models: parametersJsonSchema → parameters, remove $schema if internal_model.startswith("claude-sonnet-"): + lib_logger.debug(f"Applying Claude-specific tool schema transformation for {internal_model}") tools = antigravity_payload["request"].get("tools", []) - for tool_idx, tool in enumerate(tools): + + for tool in tools: function_declarations = tool.get("functionDeclarations", []) - for func_idx, func_decl in enumerate(function_declarations): + for func_decl in function_declarations: if "parametersJsonSchema" in func_decl: - # Convert parametersJsonSchema → parameters for Claude params = func_decl["parametersJsonSchema"] - # Remove $schema if present (Claude doesn't support it) - if isinstance(params, dict): - params.pop("$schema", None) + # CRITICAL: Claude requires clean JSON Schema draft 2020-12 + # Recursively remove ALL incompatible fields + def clean_claude_schema(schema): + """Recursively remove fields Claude doesn't support.""" + if not isinstance(schema, dict): + return schema + + # Fields that break Claude's JSON Schema validation + incompatible = {'$schema', 'additionalProperties', 'minItems', 'maxItems', 'pattern'} + cleaned = {} + + for key, value in schema.items(): + if key in incompatible: + continue # Skip incompatible fields + + if isinstance(value, dict): + cleaned[key] = clean_claude_schema(value) + elif isinstance(value, list): + cleaned[key] = [ + clean_claude_schema(item) if isinstance(item, dict) else item + for item in value + ] + else: + cleaned[key] = value + + return cleaned + + # Clean the schema + params = clean_claude_schema(params) if isinstance(params, dict) else params - # Set as 'parameters' and remove 'parametersJsonSchema' - antigravity_payload["request"]["tools"][tool_idx]["functionDeclarations"][func_idx]["parameters"] = params - del antigravity_payload["request"]["tools"][tool_idx]["functionDeclarations"][func_idx]["parametersJsonSchema"] + # Rename parametersJsonSchema → parameters for Claude + func_decl["parameters"] = params + del func_decl["parametersJsonSchema"] return antigravity_payload @@ -922,7 +930,6 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D tool_call_index += 1 # Increment for next tool call # Handle thoughtSignature if present - # CRITICAL FIX: Cache and passthrough are INDEPENDENT toggles if has_signature and not first_signature_seen: # Only first tool call gets signature (parallel call handling) first_signature_seen = True @@ -1069,11 +1076,6 @@ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str] access_token = await self.get_valid_token(credential_path) base_url = self._get_current_base_url() - # Generate required IDs - project_id = self.generate_project_id() - request_id = self.generate_request_id() - - # Fetch models endpoint url = f"{base_url}/fetchAvailableModels" headers = { @@ -1082,13 +1084,11 @@ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str] } payload = { - "project": project_id, - "requestId": request_id, + "project": self.generate_project_id(), + "requestId": self.generate_request_id(), "userAgent": "antigravity" } - lib_logger.debug(f"Fetching Antigravity models from: {url}") - response = await client.post(url, json=payload, headers=headers, timeout=30.0) response.raise_for_status() @@ -1222,6 +1222,9 @@ async def acompletion( # Remove OpenAI-specific fields that Antigravity doesn't support schema.pop("$schema", None) schema.pop("strict", None) + # CRITICAL: Normalize type arrays for protobuf compatibility + # Converts ["string", "null"] → "string" to avoid "Proto field is not repeating" errors + schema = self._normalize_type_arrays(schema) func_decl["parametersJsonSchema"] = schema else: # No parameters provided - set default empty schema (matching Go lines 318-323) @@ -1411,19 +1414,33 @@ async def count_tokens( gemini_cli_payload["systemInstruction"] = system_instruction if tools: - # Transform tools to Gemini format + # Transform tools - same as in acompletion gemini_tools = [] for tool in tools: if tool.get("type") == "function": func = tool.get("function", {}) - schema = _build_vertex_schema(parameters=func.get("parameters", {})) + parameters = func.get("parameters") + + func_decl = { + "name": func.get("name", ""), + "description": func.get("description", "") + } + + if parameters and isinstance(parameters, dict): + schema = dict(parameters) + schema.pop("$schema", None) + schema.pop("strict", None) + func_decl["parametersJsonSchema"] = schema + else: + func_decl["parametersJsonSchema"] = { + "type": "object", + "properties": {} + } + gemini_tools.append({ - "functionDeclarations": [{ - "name": func.get("name", ""), - "description": func.get("description", ""), - "parametersJsonSchema": schema - }] + "functionDeclarations": [func_decl] }) + if gemini_tools: gemini_cli_payload["tools"] = gemini_tools From 0970b56ece20996c3702e1d520ebd1666d91b2d5 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 24 Nov 2025 01:25:52 +0100 Subject: [PATCH 013/221] =?UTF-8?q?fix(antigravity):=20=F0=9F=90=9B=20add?= =?UTF-8?q?=20function=20call=20id=20fields=20and=20restrict=20thoughtSign?= =?UTF-8?q?ature=20handling=20to=20gemini-3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add "id" to functionCall and response objects required by Antigravity/Claude integrations. Restrict preservation/insertion of thoughtSignature to Gemini 3 models only: prefer client-provided signature, fall back to the server-side cache when enabled, and finally use the bypass constant "skip_thought_signature_validator". Emit a warning when a Gemini 3 tool call lacks a signature. Avoid adding thoughtSignature for Claude and other models to prevent sending unsupported fields. --- .../providers/antigravity_provider.py | 58 ++++++++++--------- 1 file changed, 32 insertions(+), 26 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index f7756f38..524524f6 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -355,7 +355,7 @@ def generate_project_id() -> str: def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: """ Transform OpenAI messages to Gemini CLI format. - Handles thoughtSignature preservation with 3-tier fallback: + Handles thoughtSignature preservation with 3-tier fallback (GEMINI 3 ONLY): 1. Use client-provided signature (if present) 2. Fall back to server-side cache 3. Use bypass constant as last resort @@ -459,27 +459,29 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tup func_call_part = { "functionCall": { "name": tool_call["function"]["name"], - "args": args_dict + "args": args_dict, + "id": tool_call_id # ← ADD THIS LINE - Antigravity needs it for Claude! } } - # PRIORITY 1: Use client-provided signature if available - client_signature = tool_call.get("thought_signature") - - # PRIORITY 2: Fall back to server-side cache - if not client_signature and tool_call_id and self._enable_signature_cache: - client_signature = self._signature_cache.retrieve(tool_call_id) - if client_signature: - lib_logger.debug(f"Retrieved thoughtSignature from cache for {tool_call_id}") - - # PRIORITY 3: Use bypass constant as last resort - if client_signature: - func_call_part["thoughtSignature"] = client_signature - else: - func_call_part["thoughtSignature"] = "skip_thought_signature_validator" + # thoughtSignature handling (GEMINI 3 ONLY) + # Claude and other models don't support this field! + if self._is_gemini_3_model(model): + # PRIORITY 1: Use client-provided signature if available + client_signature = tool_call.get("thought_signature") + + # PRIORITY 2: Fall back to server-side cache + if not client_signature and tool_call_id and self._enable_signature_cache: + client_signature = self._signature_cache.retrieve(tool_call_id) + if client_signature: + lib_logger.debug(f"Retrieved thoughtSignature from cache for {tool_call_id}") - # WARNING: Missing signature for Gemini 3 - if self._is_gemini_3_model(model): + # PRIORITY 3: Use bypass constant as last resort + if client_signature: + func_call_part["thoughtSignature"] = client_signature + else: + func_call_part["thoughtSignature"] = "skip_thought_signature_validator" + # WARNING: Missing signature for Gemini 3 lib_logger.warning( f"Gemini 3 tool call '{tool_call_id}' missing thoughtSignature. " f"Client didn't provide it and cache lookup failed. " @@ -505,7 +507,8 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tup "name": function_name, "response": { "result": parsed_content - } + }, + "id": tool_call_id # ← ADD THIS LINE - Antigravity needs it for Claude! } }) @@ -759,13 +762,16 @@ def _transform_to_antigravity_format( # Set thinkingBudget to -1 (auto/dynamic) thinking_config["thinkingBudget"] = -1 - # 6. Preserve/add thoughtSignature to ALL function calls in model role content - for content in antigravity_payload["request"].get("contents", []): - if content.get("role") == "model": - for part in content.get("parts", []): - # Add signature to function calls OR preserve if already exists - if "functionCall" in part and "thoughtSignature" not in part: - part["thoughtSignature"] = "skip_thought_signature_validator" + # 6. Preserve/add thoughtSignature to function calls in model role content (GEMINI 3 ONLY) + # thoughtSignature is a Gemini 3 feature for preserving reasoning context in multi-turn conversations + # DO NOT add this for Claude or other models - they don't support it! + if internal_model.startswith("gemini-3-"): + for content in antigravity_payload["request"].get("contents", []): + if content.get("role") == "model": + for part in content.get("parts", []): + # Add signature to function calls OR preserve if already exists + if "functionCall" in part and "thoughtSignature" not in part: + part["thoughtSignature"] = "skip_thought_signature_validator" # 7. CLAUDE-SPECIFIC TOOL SCHEMA TRANSFORMATION # Reference: Go implementation antigravity_executor.go lines 672-684 From 6adac7a7d3ce5838969a57630f43343b4cf6d346 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 24 Nov 2025 02:51:31 +0100 Subject: [PATCH 014/221] =?UTF-8?q?fix(api):=20=F0=9F=90=9B=20override=20g?= =?UTF-8?q?lobal=20temperature=3D0=20via=20OVERRIDE=5FTEMPERATURE=5FZERO?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an environment-controlled override that modifies requests with `temperature: 0` for chat completions when `OVERRIDE_TEMPERATURE_ZERO` is enabled (default: "false"). - Supported modes: "remove" — delete the `temperature` key; "set"/"true"/"1"/"yes" — set temperature to 1.0. - Rationale: temperature=0 makes models overly deterministic and can cause tool hallucination; the override helps mitigate that when toggled. - Emits debug logs when an override is applied. --- src/proxy_app/main.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 94f2c38a..8903b688 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -652,6 +652,22 @@ async def chat_completions( except json.JSONDecodeError: raise HTTPException(status_code=400, detail="Invalid JSON in request body.") + # Global temperature=0 override (controlled by .env variable, default: OFF) + # Low temperature makes models deterministic and prone to following training data + # instead of actual schemas, which can cause tool hallucination + # Modes: "remove" = delete temperature key, "set" = change to 1.0, "false" = disabled + override_temp_zero = os.getenv("OVERRIDE_TEMPERATURE_ZERO", "false").lower() + + if override_temp_zero in ("remove", "set", "true", "1", "yes") and "temperature" in request_data and request_data["temperature"] == 0: + if override_temp_zero == "remove": + # Remove temperature key entirely + del request_data["temperature"] + logging.debug("OVERRIDE_TEMPERATURE_ZERO=remove: Removed temperature=0 from request") + else: + # Set to 1.0 (for "set", "true", "1", "yes") + request_data["temperature"] = 1.0 + logging.debug("OVERRIDE_TEMPERATURE_ZERO=set: Converting temperature=0 to temperature=1.0") + # If logging is enabled, perform all logging operations using the parsed data. if logger: logger.log_request(headers=request.headers, body=request_data) From d7fa9988d6c56e04fac002dd4d3009c578216976 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 24 Nov 2025 02:52:43 +0100 Subject: [PATCH 015/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20add=20?= =?UTF-8?q?Gemini=203=20tool-fix=20(namespace,=20signature,=20system-instr?= =?UTF-8?q?uction)=20to=20reduce=20tool=20hallucination?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce a configurable "Gemini 3" catch-all fix that enforces schema-driven tool usage and reduces tool hallucination by: - adding env-configurable flag ANTIGRAVITY_GEMINI3_TOOL_FIX (default ON) and related vars for prefix, description prompt, and system instruction - implementing namespace prefixing for tool names to break model training associations - injecting strict parameter signatures into tool descriptions to force schema adherence - prepending configurable system instructions for Gemini-3 models to override training-data assumptions - normalizing request/response names (prefix/strip) and preserving function call ids for API consistency - applying transformations only for gemini-3-* models and logging configuration details This change improves robustness when calling external tools by making tool schemas explicit to the model. --- .../providers/antigravity_provider.py | 274 +++++++++++++++++- 1 file changed, 270 insertions(+), 4 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 524524f6..86bed053 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -238,6 +238,58 @@ def __init__(self): lib_logger.info("Antigravity: Dynamic model discovery ENABLED (may fail if endpoint unavailable)") else: lib_logger.info("Antigravity: Dynamic model discovery DISABLED (using hardcoded model list)") + + # Check if Gemini 3 tool fix is enabled (default: ON for testing) + # This applies the "Quad-Lock" catch-all strategy to prevent tool hallucination + self._enable_gemini3_tool_fix = os.getenv( + "ANTIGRAVITY_GEMINI3_TOOL_FIX", + "true" # Default ON - applies namespace + signature injection + ).lower() in ("true", "1", "yes") + + # Gemini 3 fix configuration - customize the fix components + # Namespace prefix for tool names (Strategy 1) + self._gemini3_tool_prefix = os.getenv( + "ANTIGRAVITY_GEMINI3_TOOL_PREFIX", + "gemini3_" # Default prefix + ) + + # Description prompt format (Strategy 2) + # Use {params} as placeholder for parameter list + self._gemini3_description_prompt = os.getenv( + "ANTIGRAVITY_GEMINI3_DESCRIPTION_PROMPT", + "\n\nSTRICT PARAMETERS: {params}." # Default format + ) + + # System instruction text (Strategy 3) + # Set to empty string to disable system instruction injection + self._gemini3_system_instruction = os.getenv( + "ANTIGRAVITY_GEMINI3_SYSTEM_INSTRUCTION", + # Default: comprehensive tool usage instructions + """CRITICAL TOOL USAGE INSTRUCTIONS: +You are operating in a custom environment where tool definitions differ from your training data. +You MUST follow these rules strictly: + +1. DO NOT use your internal training data to guess tool parameters +2. ONLY use the exact parameter structure defined in the tool schema +3. If a tool takes a 'files' parameter, it is ALWAYS an array of objects with specific properties, NEVER a simple array of strings +4. If a tool edits code, it takes structured JSON objects with specific fields, NEVER raw diff strings or plain text +5. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') +6. Array parameters have specific item types - check the schema's 'items' field for the exact structure +7. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions + +If you are unsure about a tool's parameters, YOU MUST read the schema definition carefully. Your training data about common tool names like 'read_file' or 'apply_diff' does NOT apply here. +""" + ) + + if self._enable_gemini3_tool_fix: + lib_logger.info(f"Antigravity: Gemini 3 tool fix ENABLED") + lib_logger.debug(f" - Namespace prefix: '{self._gemini3_tool_prefix}'") + lib_logger.debug(f" - Description prompt: '{self._gemini3_description_prompt[:50]}...'") + lib_logger.debug(f" - System instruction: {'ENABLED' if self._gemini3_system_instruction else 'DISABLED'} ({len(self._gemini3_system_instruction)} chars)") + else: + lib_logger.info("Antigravity: Gemini 3 tool fix DISABLED (using default tool schemas)") + + # ============================================================================ # MODEL ALIAS SYSTEM @@ -456,9 +508,15 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tup tool_call_id = tool_call.get("id", "") + # Get function name and add configured prefix if needed (Gemini 3 specific) + function_name = tool_call["function"]["name"] + if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: + # Client sends original names, we need to prefix for API consistency + function_name = f"{self._gemini3_tool_prefix}{function_name}" + func_call_part = { "functionCall": { - "name": tool_call["function"]["name"], + "name": function_name, "args": args_dict, "id": tool_call_id # ← ADD THIS LINE - Antigravity needs it for Claude! } @@ -496,6 +554,11 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tup function_name = tool_call_id_to_name.get(tool_call_id, "unknown_function") tool_content = msg.get("content", "{}") + # Add configured prefix to function response name if needed (Gemini 3 specific) + if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: + # Client sends responses for original names, we need to prefix for API consistency + function_name = f"{self._gemini3_tool_prefix}{function_name}" + # Parse tool content - if it's JSON, use parsed value; otherwise use as-is try: parsed_content = json.loads(tool_content) @@ -706,6 +769,153 @@ def _fix_tool_response_grouping(self, contents: List[Dict[str, Any]]) -> List[Di return new_contents + # ============================================================================ + # GEMINI 3 TOOL TRANSFORMATION (Catch-All Fix for Hallucination) + # ============================================================================ + + def _apply_gemini3_namespace_to_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Apply namespace prefix to all tool names for Gemini 3 (Strategy 1: Namespace). + + This breaks the model's association with training data by prepending 'gemini3_' + to every tool name, forcing it to read the schema definition instead of using + its internal knowledge. + + Args: + tools: List of tool definitions (Gemini format with functionDeclarations) + + Returns: + Modified tools with prefixed names + """ + if not tools: + return tools + + modified_tools = copy.deepcopy(tools) + + for tool in modified_tools: + function_declarations = tool.get("functionDeclarations", []) + for func_decl in function_declarations: + # Prepend namespace to tool name + original_name = func_decl.get("name", "") + if original_name: + func_decl["name"] = f"{self._gemini3_tool_prefix}{original_name}" + lib_logger.debug(f"Gemini 3 namespace: {original_name} -> {self._gemini3_tool_prefix}{original_name}") + + return modified_tools + + def _inject_signature_into_tool_descriptions(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Inject parameter signatures into tool descriptions for Gemini 3 (Strategy 2: Signature Injection). + + This strategy appends the expected parameter structure into the description text, + creating a natural language enforcement of the schema that models pay close attention to. + + Args: + tools: List of tool definitions (Gemini format with functionDeclarations) + + Returns: + Modified tools with enriched descriptions + """ + if not tools: + return tools + + modified_tools = copy.deepcopy(tools) + + for tool in modified_tools: + function_declarations = tool.get("functionDeclarations", []) + for func_decl in function_declarations: + # Get parameter schema + schema = func_decl.get("parametersJsonSchema", {}) + if not schema or not isinstance(schema, dict): + continue + + # Extract required parameters + required_params = schema.get("required", []) + properties = schema.get("properties", {}) + + if not properties: + continue + + # Build parameter list with type hints + param_list = [] + for prop_name, prop_data in properties.items(): + if not isinstance(prop_data, dict): + continue + + type_hint = prop_data.get("type", "unknown") + + # Handle arrays specially (critical for read_file/apply_diff issues) + if type_hint == "array": + items_schema = prop_data.get("items", {}) + if isinstance(items_schema, dict): + item_type = items_schema.get("type", "unknown") + + # Check if it's an array of objects - RECURSE into nested properties + if item_type == "object": + # Extract nested properties for explicit visibility + nested_props = items_schema.get("properties", {}) + nested_required = items_schema.get("required", []) + + if nested_props: + # Build nested property list with types + nested_list = [] + for nested_name, nested_data in nested_props.items(): + if not isinstance(nested_data, dict): + continue + nested_type = nested_data.get("type", "unknown") + + # Mark nested required fields + if nested_name in nested_required: + nested_list.append(f"{nested_name}: {nested_type} REQUIRED") + else: + nested_list.append(f"{nested_name}: {nested_type}") + + # Format as ARRAY_OF_OBJECTS[key1: type1, key2: type2] + nested_str = ", ".join(nested_list) + type_hint = f"ARRAY_OF_OBJECTS[{nested_str}]" + else: + # No properties defined - just generic objects + type_hint = "ARRAY_OF_OBJECTS" + else: + type_hint = f"ARRAY_OF_{item_type.upper()}" + else: + type_hint = "ARRAY" + + # Mark required parameters + if prop_name in required_params: + param_list.append(f"{prop_name} ({type_hint}, REQUIRED)") + else: + param_list.append(f"{prop_name} ({type_hint})") + + # Create strict signature string using configurable template + # Replace {params} placeholder with actual parameter list + signature_str = self._gemini3_description_prompt.replace("{params}", ", ".join(param_list)) + + # Inject into description + description = func_decl.get("description", "") + func_decl["description"] = description + signature_str + + lib_logger.debug(f"Gemini 3 signature injection: {func_decl.get('name', '')} - {len(param_list)} params") + + return modified_tools + + def _strip_gemini3_namespace_from_name(self, tool_name: str) -> str: + """ + Strip the configured namespace prefix from a tool name. + + This reverses the namespace transformation applied in the request, + ensuring the client receives the original tool names. + + Args: + tool_name: Tool name (possibly with configured prefix) + + Returns: + Original tool name without prefix + """ + if tool_name and tool_name.startswith(self._gemini3_tool_prefix): + return tool_name[len(self._gemini3_tool_prefix):] + return tool_name + # ============================================================================ # ANTIGRAVITY REQUEST TRANSFORMATION # ============================================================================ @@ -924,12 +1134,17 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D func_call = part["functionCall"] tool_call_id = f"call_{uuid.uuid4().hex[:24]}" + # Get tool name and strip gemini3_ namespace if present (Gemini 3 specific) + tool_name = func_call.get("name", "") + if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: + tool_name = self._strip_gemini3_namespace_from_name(tool_name) + tool_call = { "id": tool_call_id, "type": "function", "index": tool_call_index, # REQUIRED for OpenAI streaming format "function": { - "name": func_call.get("name", ""), + "name": tool_name, "arguments": json.dumps(func_call.get("args", {})) } } @@ -1181,10 +1396,45 @@ async def acompletion( if system_instruction: gemini_cli_payload["system_instruction"] = system_instruction + # Apply Gemini 3 system instruction injection (Strategy 3) if fix is enabled + # This prepends critical tool usage instructions to override model's training data + if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix and tools: + gemini3_instruction = self._gemini3_system_instruction + + if "system_instruction" in gemini_cli_payload: + # Prepend to existing system instruction + existing_instruction = gemini_cli_payload["system_instruction"] + if isinstance(existing_instruction, dict) and "parts" in existing_instruction: + # System instruction with parts structure + gemini3_part = {"text": gemini3_instruction} + existing_instruction["parts"].insert(0, gemini3_part) + else: + # Shouldn't happen, but handle gracefully + gemini_cli_payload["system_instruction"] = { + "role": "user", + "parts": [ + {"text": gemini3_instruction}, + {"text": str(existing_instruction)} + ] + } + else: + # Create new system instruction with Gemini 3 instructions + gemini_cli_payload["system_instruction"] = { + "role": "user", + "parts": [{"text": gemini3_instruction}] + } + + lib_logger.debug("Gemini 3 system instruction injection applied") + + + # Add generation config generation_config = {} - if temperature is not None: - generation_config["temperature"] = temperature + + # Temperature handling: Default to 1.0, override 0 to 1.0 + # Low temperature (especially 0) makes models deterministic and prone to following + # training data patterns instead of actual schemas, which causes tool hallucination + if top_p is not None: generation_config["topP"] = top_p @@ -1245,6 +1495,22 @@ async def acompletion( if gemini_tools: gemini_cli_payload["tools"] = gemini_tools + + # Apply Gemini 3 specific tool transformations (ONLY for gemini-3-* models) + # This implements the "Double-Lock" catch-all strategy to prevent tool hallucination + if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: + lib_logger.info(f"Applying Gemini 3 catch-all tool transformations for {model}") + + # Strategy 1: Namespace prefixing (breaks association with training data) + gemini_cli_payload["tools"] = self._apply_gemini3_namespace_to_tools( + gemini_cli_payload["tools"] + ) + + # Strategy 2: Signature injection (natural language schema enforcement) + gemini_cli_payload["tools"] = self._inject_signature_into_tool_descriptions( + gemini_cli_payload["tools"] + ) + # Step 3: Transform to Antigravity format antigravity_payload = self._transform_to_antigravity_format(gemini_cli_payload, model) From 946e5a0df2fc653f5ff052465ea7912252682740 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 24 Nov 2025 04:05:48 +0100 Subject: [PATCH 016/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20add=20?= =?UTF-8?q?disk=20persistence=20for=20thoughtSignature=20cache?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement dual-TTL caching system with async disk persistence to improve thoughtSignature handling across server restarts and long-running sessions. - Add disk persistence using atomic file writes with tempfile pattern for data integrity - Implement dual-TTL system: 1-hour memory cache, 24-hour disk cache - Create background async tasks for periodic disk writes and memory cleanup - Add disk fallback mechanism for cache misses (loads from disk into memory) - Introduce cache statistics tracking (memory hits, disk hits, misses, writes) - Add graceful shutdown with pending write flush - Convert cache operations from threading.Lock to asyncio.Lock for async support - Add environment variables for configurable write/cleanup intervals - Implement secure file permissions (0o600) for cache files - Add comprehensive logging for cache lifecycle events The cache now survives server restarts and provides better support for multi-turn conversations by persisting thoughtSignatures to disk. Memory cache expires after 1 hour to prevent unbounded growth, while disk cache persists for 24 hours to support longer conversation sessions. --- .../providers/antigravity_provider.py | 584 ++++++++++++++++-- 1 file changed, 537 insertions(+), 47 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 86bed053..e916fa5c 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -10,6 +10,8 @@ import copy import threading import os +import tempfile +import shutil from pathlib import Path from datetime import datetime from typing import List, Dict, Any, AsyncGenerator, Union, Optional, Tuple @@ -45,6 +47,11 @@ LOGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "logs" ANTIGRAVITY_LOGS_DIR = LOGS_DIR / "antigravity_logs" +# Cache configuration +CACHE_DIR = Path(__file__).resolve().parent.parent.parent.parent / "cache" +ANTIGRAVITY_CACHE_DIR = CACHE_DIR / "antigravity" +ANTIGRAVITY_CACHE_FILE = ANTIGRAVITY_CACHE_DIR / "thought_signatures.json" + class _AntigravityFileLogger: """A simple file logger for a single Antigravity transaction.""" @@ -108,37 +115,289 @@ class ThoughtSignatureCache: across turns, even if clients don't support the thought_signature field. Features: - - TTL-based expiration to prevent memory growth + - Dual-TTL system: 1hr memory, 24hr disk + - Async disk persistence with batched writes + - Background cleanup task for expired entries - Thread-safe for concurrent access - - Automatic cleanup of expired entries + - Fallback to disk when not in memory + - High concurrency support with asyncio locks """ - def __init__(self, ttl_seconds: int = 3600): + def __init__(self, memory_ttl_seconds: int = 3600, disk_ttl_seconds: int = 86400): """ - Initialize the signature cache. + Initialize the signature cache with disk persistence. Args: - ttl_seconds: Time-to-live for cache entries in seconds (default: 1 hour) + memory_ttl_seconds: Time-to-live for memory cache entries (default: 1 hour) + disk_ttl_seconds: Time-to-live for disk cache entries (default: 24 hours) """ - self._cache: Dict[str, Tuple[str, float]] = {} # {call_id: (signature, timestamp)} - self._ttl = ttl_seconds - self._lock = threading.Lock() + # In-memory cache: {call_id: (signature, timestamp)} + self._cache: Dict[str, Tuple[str, float]] = {} + self._memory_ttl = memory_ttl_seconds + self._disk_ttl = disk_ttl_seconds + self._lock = asyncio.Lock() + self._disk_lock = asyncio.Lock() + + # Disk persistence configuration + self._cache_file = ANTIGRAVITY_CACHE_FILE + self._enable_disk_persistence = os.getenv( + "ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", + "true" + ).lower() in ("true", "1", "yes") + + # Async write configuration + self._dirty = False # Flag for pending writes + self._write_interval = int(os.getenv("ANTIGRAVITY_CACHE_WRITE_INTERVAL", "60")) + self._cleanup_interval = int(os.getenv("ANTIGRAVITY_CACHE_CLEANUP_INTERVAL", "1800")) + + # Background tasks + self._writer_task: Optional[asyncio.Task] = None + self._cleanup_task: Optional[asyncio.Task] = None + self._running = False + + # Statistics + self._stats = { + "memory_hits": 0, + "disk_hits": 0, + "misses": 0, + "writes": 0 + } + + # Initialize + if self._enable_disk_persistence: + lib_logger.debug( + f"ThoughtSignatureCache: Disk persistence ENABLED " + f"(memory_ttl={memory_ttl_seconds}s, disk_ttl={disk_ttl_seconds}s, " + f"write_interval={self._write_interval}s)" + ) + # Schedule async initialization + asyncio.create_task(self._async_init()) + else: + lib_logger.debug("ThoughtSignatureCache: Disk persistence DISABLED (memory-only mode)") + + async def _async_init(self): + """Async initialization: load from disk and start background tasks.""" + try: + await self._load_from_disk() + await self._start_background_tasks() + except Exception as e: + lib_logger.error(f"ThoughtSignatureCache async init failed: {e}") + + async def _load_from_disk(self): + """Load cache from disk file (with TTL validation).""" + if not self._enable_disk_persistence: + return + + if not self._cache_file.exists(): + lib_logger.debug("No existing cache file found, starting fresh") + return + + try: + async with self._disk_lock: + # Read cache file + with open(self._cache_file, 'r', encoding='utf-8') as f: + data = json.load(f) + + # Validate version + if data.get("version") != "1.0": + lib_logger.warning(f"Cache file version mismatch, ignoring") + return + + # Load entries with disk TTL validation + now = time.time() + entries = data.get("entries", {}) + loaded = 0 + expired = 0 + + for call_id, entry in entries.items(): + timestamp = entry.get("timestamp", 0) + age = now - timestamp + + # Check against DISK TTL (24 hours) + if age <= self._disk_ttl: + signature = entry.get("signature", "") + if signature: + self._cache[call_id] = (signature, timestamp) + loaded += 1 + else: + expired += 1 + + lib_logger.debug( + f"ThoughtSignatureCache: Loaded {loaded} signatures from disk " + f"({expired} expired entries removed)" + ) + + except json.JSONDecodeError as e: + lib_logger.warning(f"Cache file corrupted, starting fresh: {e}") + except Exception as e: + lib_logger.error(f"Failed to load cache from disk: {e}") + + async def _save_to_disk(self): + """Persist cache to disk using atomic write.""" + if not self._enable_disk_persistence: + return + + try: + async with self._disk_lock: + # Ensure cache directory exists + self._cache_file.parent.mkdir(parents=True, exist_ok=True) + + # Build cache data structure + cache_data = { + "version": "1.0", + "memory_ttl_seconds": self._memory_ttl, + "disk_ttl_seconds": self._disk_ttl, + "entries": { + call_id: { + "signature": sig, + "timestamp": ts + } + for call_id, (sig, ts) in self._cache.items() + }, + "statistics": { + "total_entries": len(self._cache), + "last_write": time.time(), + "memory_hits": self._stats["memory_hits"], + "disk_hits": self._stats["disk_hits"], + "misses": self._stats["misses"], + "writes": self._stats["writes"] + } + } + + # Atomic write using tempfile pattern (same as OAuth credentials) + parent_dir = self._cache_file.parent + tmp_fd = None + tmp_path = None + + try: + # Create temp file in same directory + tmp_fd, tmp_path = tempfile.mkstemp( + dir=parent_dir, + prefix='.tmp_', + suffix='.json', + text=True + ) + + # Write JSON to temp file + with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f: + json.dump(cache_data, f, indent=2) + tmp_fd = None # fdopen closes the fd + + # Set secure permissions (owner read/write only) + try: + os.chmod(tmp_path, 0o600) + except (OSError, AttributeError): + # Windows may not support chmod, ignore + pass + + # Atomic move (overwrites target if exists) + shutil.move(tmp_path, self._cache_file) + tmp_path = None # Successfully moved + + self._stats["writes"] += 1 + lib_logger.debug(f"Saved {len(self._cache)} signatures to disk") + + except Exception as e: + lib_logger.error(f"Failed to save cache to disk: {e}") + # Clean up temp file if it still exists + if tmp_fd is not None: + try: + os.close(tmp_fd) + except: + pass + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except: + pass + raise + + except Exception as e: + lib_logger.error(f"Disk save operation failed: {e}") + + async def _start_background_tasks(self): + """Start background writer and cleanup tasks.""" + if not self._enable_disk_persistence or self._running: + return + + self._running = True + + # Start async writer task + self._writer_task = asyncio.create_task(self._writer_loop()) + lib_logger.debug(f"Started background writer task (interval: {self._write_interval}s)") + + # Start cleanup task + self._cleanup_task = asyncio.create_task(self._cleanup_loop()) + lib_logger.debug(f"Started background cleanup task (interval: {self._cleanup_interval}s)") + + async def _writer_loop(self): + """Background task: periodically flush dirty cache to disk.""" + try: + while self._running: + await asyncio.sleep(self._write_interval) + + if self._dirty: + try: + await self._save_to_disk() + self._dirty = False + except Exception as e: + lib_logger.error(f"Background writer error: {e}") + except asyncio.CancelledError: + lib_logger.debug("Background writer task cancelled") + except Exception as e: + lib_logger.error(f"Background writer crashed: {e}") + + async def _cleanup_loop(self): + """Background task: periodically clean up expired entries.""" + try: + while self._running: + await asyncio.sleep(self._cleanup_interval) + + try: + await self._cleanup_expired() + except Exception as e: + lib_logger.error(f"Background cleanup error: {e}") + except asyncio.CancelledError: + lib_logger.debug("Background cleanup task cancelled") + except Exception as e: + lib_logger.error(f"Background cleanup crashed: {e}") + + async def _cleanup_expired(self): + """Remove expired entries from memory cache (based on memory TTL).""" + async with self._lock: + now = time.time() + expired = [ + k for k, (_, ts) in self._cache.items() + if now - ts > self._memory_ttl + ] + + for k in expired: + del self._cache[k] + + if expired: + self._dirty = True # Mark for disk save + lib_logger.debug(f"Cleaned up {len(expired)} expired signatures from memory") def store(self, tool_call_id: str, signature: str): """ - Store a signature for a tool call ID. + Store a signature for a tool call ID (sync wrapper for async storage). Args: tool_call_id: Unique identifier for the tool call signature: Encrypted thoughtSignature from Antigravity API """ - with self._lock: + # Create task for async storage + asyncio.create_task(self._async_store(tool_call_id, signature)) + + async def _async_store(self, tool_call_id: str, signature: str): + """Async implementation of store.""" + async with self._lock: self._cache[tool_call_id] = (signature, time.time()) - self._cleanup_expired() + self._dirty = True # Mark for disk write def retrieve(self, tool_call_id: str) -> Optional[str]: """ - Retrieve signature for a tool call ID. + Retrieve signature for a tool call ID (sync method). Args: tool_call_id: Unique identifier for the tool call @@ -146,28 +405,97 @@ def retrieve(self, tool_call_id: str) -> Optional[str]: Returns: The signature if found and not expired, None otherwise """ - with self._lock: - if tool_call_id not in self._cache: - return None - + # Try memory cache first (sync access is safe for read) + if tool_call_id in self._cache: signature, timestamp = self._cache[tool_call_id] - if time.time() - timestamp > self._ttl: + if time.time() - timestamp <= self._memory_ttl: + self._stats["memory_hits"] += 1 + return signature + else: + # Expired in memory, remove it del self._cache[tool_call_id] - return None - - return signature + self._dirty = True + + # Not in memory - schedule async disk lookup + # For now, return None (disk fallback happens on next request) + # This is intentional to avoid blocking the sync caller + self._stats["misses"] += 1 + + # Schedule background disk check (non-blocking) + if self._enable_disk_persistence: + asyncio.create_task(self._check_disk_fallback(tool_call_id)) + + return None - def _cleanup_expired(self): - """Remove expired entries from cache.""" - now = time.time() - expired = [k for k, (_, ts) in self._cache.items() if now - ts > self._ttl] - for k in expired: - del self._cache[k] + async def _check_disk_fallback(self, tool_call_id: str): + """Check disk for signature and load into memory if found.""" + try: + # Reload from disk if file exists + if self._cache_file.exists(): + async with self._disk_lock: + with open(self._cache_file, 'r', encoding='utf-8') as f: + data = json.load(f) + + entries = data.get("entries", {}) + if tool_call_id in entries: + entry = entries[tool_call_id] + timestamp = entry.get("timestamp", 0) + + # Check disk TTL (24 hours) + if time.time() - timestamp <= self._disk_ttl: + signature = entry.get("signature", "") + if signature: + # Load into memory cache + async with self._lock: + self._cache[tool_call_id] = (signature, timestamp) + self._stats["disk_hits"] += 1 + lib_logger.debug(f"Loaded signature {tool_call_id} from disk") + except Exception as e: + lib_logger.debug(f"Disk fallback check failed: {e}") - def clear(self): - """Clear all cached signatures.""" - with self._lock: + async def clear(self): + """Clear all cached signatures (memory and disk).""" + async with self._lock: self._cache.clear() + self._dirty = True + + if self._enable_disk_persistence: + await self._save_to_disk() + + async def shutdown(self): + """Graceful shutdown: flush pending writes and stop background tasks.""" + lib_logger.info("ThoughtSignatureCache shutting down...") + + # Stop background tasks + self._running = False + + if self._writer_task: + self._writer_task.cancel() + try: + await self._writer_task + except asyncio.CancelledError: + pass + + if self._cleanup_task: + self._cleanup_task.cancel() + try: + await self._cleanup_task + except asyncio.CancelledError: + pass + + # Flush pending writes + if self._dirty and self._enable_disk_persistence: + lib_logger.info("Flushing pending cache writes...") + await self._save_to_disk() + + lib_logger.info( + f"ThoughtSignatureCache shutdown complete " + f"(stats: mem_hits={self._stats['memory_hits']}, " + f"disk_hits={self._stats['disk_hits']}, " + f"misses={self._stats['misses']}, " + f"writes={self._stats['writes']})" + ) + class AntigravityProvider(AntigravityAuthBase, ProviderInterface): """ @@ -203,8 +531,12 @@ def __init__(self): self._base_url_index = 0 # Initialize thoughtSignature cache for Gemini 3 multi-turn conversations - cache_ttl = int(os.getenv("ANTIGRAVITY_SIGNATURE_CACHE_TTL", "3600")) - self._signature_cache = ThoughtSignatureCache(ttl_seconds=cache_ttl) + memory_ttl = int(os.getenv("ANTIGRAVITY_SIGNATURE_CACHE_TTL", "3600")) + disk_ttl = int(os.getenv("ANTIGRAVITY_SIGNATURE_DISK_TTL", "86400")) + self._signature_cache = ThoughtSignatureCache( + memory_ttl_seconds=memory_ttl, + disk_ttl_seconds=disk_ttl + ) # Check if client passthrough is enabled (default: TRUE for testing) self._preserve_signatures_in_client = os.getenv( @@ -225,19 +557,19 @@ def __init__(self): ).lower() in ("true", "1", "yes") if self._preserve_signatures_in_client: - lib_logger.info("Antigravity: thoughtSignature client passthrough ENABLED") + lib_logger.debug("Antigravity: thoughtSignature client passthrough ENABLED") else: - lib_logger.info("Antigravity: thoughtSignature client passthrough DISABLED") + lib_logger.debug("Antigravity: thoughtSignature client passthrough DISABLED") if self._enable_signature_cache: - lib_logger.info(f"Antigravity: thoughtSignature server-side cache ENABLED (TTL: {cache_ttl}s)") + lib_logger.debug(f"Antigravity: thoughtSignature server-side cache ENABLED (memory_ttl={memory_ttl}s, disk_ttl={disk_ttl}s)") else: - lib_logger.info("Antigravity: thoughtSignature server-side cache DISABLED") + lib_logger.debug("Antigravity: thoughtSignature server-side cache DISABLED") if self._enable_dynamic_model_discovery: - lib_logger.info("Antigravity: Dynamic model discovery ENABLED (may fail if endpoint unavailable)") + lib_logger.debug("Antigravity: Dynamic model discovery ENABLED (may fail if endpoint unavailable)") else: - lib_logger.info("Antigravity: Dynamic model discovery DISABLED (using hardcoded model list)") + lib_logger.debug("Antigravity: Dynamic model discovery DISABLED (using hardcoded model list)") # Check if Gemini 3 tool fix is enabled (default: ON for testing) # This applies the "Quad-Lock" catch-all strategy to prevent tool hallucination @@ -282,12 +614,12 @@ def __init__(self): ) if self._enable_gemini3_tool_fix: - lib_logger.info(f"Antigravity: Gemini 3 tool fix ENABLED") + lib_logger.debug(f"Antigravity: Gemini 3 tool fix ENABLED") lib_logger.debug(f" - Namespace prefix: '{self._gemini3_tool_prefix}'") lib_logger.debug(f" - Description prompt: '{self._gemini3_description_prompt[:50]}...'") lib_logger.debug(f" - System instruction: {'ENABLED' if self._gemini3_system_instruction else 'DISABLED'} ({len(self._gemini3_system_instruction)} chars)") else: - lib_logger.info("Antigravity: Gemini 3 tool fix DISABLED (using default tool schemas)") + lib_logger.debug("Antigravity: Gemini 3 tool fix DISABLED (using default tool schemas)") @@ -799,7 +1131,7 @@ def _apply_gemini3_namespace_to_tools(self, tools: List[Dict[str, Any]]) -> List original_name = func_decl.get("name", "") if original_name: func_decl["name"] = f"{self._gemini3_tool_prefix}{original_name}" - lib_logger.debug(f"Gemini 3 namespace: {original_name} -> {self._gemini3_tool_prefix}{original_name}") + #lib_logger.debug(f"Gemini 3 namespace: {original_name} -> {self._gemini3_tool_prefix}{original_name}") return modified_tools @@ -895,7 +1227,7 @@ def _inject_signature_into_tool_descriptions(self, tools: List[Dict[str, Any]]) description = func_decl.get("description", "") func_decl["description"] = description + signature_str - lib_logger.debug(f"Gemini 3 signature injection: {func_decl.get('name', '')} - {len(param_list)} params") + #lib_logger.debug(f"Gemini 3 signature injection: {func_decl.get('name', '')} - {len(param_list)} params") return modified_tools @@ -1231,6 +1563,161 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D response["usage"] = usage return response + + def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model: str) -> Dict[str, Any]: + """ + Convert a Gemini API response to OpenAI non-streaming format. + + This is specifically for non-streaming completions where we need 'message' instead of 'delta'. + + Args: + gemini_response: Gemini API response + model: Model name for Gemini 3 detection + + Returns: + OpenAI-compatible non-streaming response + """ + # Extract the main response structure + candidates = gemini_response.get("candidates", []) + if not candidates: + return {} + + candidate = candidates[0] + content = candidate.get("content", {}) + content_parts = content.get("parts", []) + + # Build message components + text_content = "" + reasoning_content = "" + tool_calls = [] + + # Track if we've seen a signature yet (for parallel tool call handling) + first_signature_seen = False + + for part in content_parts: + has_function_call = "functionCall" in part + has_text = "text" in part + has_signature = "thoughtSignature" in part and part["thoughtSignature"] + + # Skip standalone signature parts + if has_signature and not has_function_call and not has_text: + continue + + # Process text content + if has_text: + thought = part.get("thought") + if thought is True or (isinstance(thought, str) and thought.lower() == 'true'): + reasoning_content += part["text"] + else: + text_content += part["text"] + + # Process function calls + if has_function_call: + func_call = part["functionCall"] + tool_call_id = f"call_{uuid.uuid4().hex[:24]}" + + # Get tool name and strip gemini3_ namespace if present + tool_name = func_call.get("name", "") + if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: + tool_name = self._strip_gemini3_namespace_from_name(tool_name) + + tool_call = { + "id": tool_call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": json.dumps(func_call.get("args", {})) + } + } + + # Handle thoughtSignature if present + if has_signature and not first_signature_seen: + first_signature_seen = True + signature = part["thoughtSignature"] + + # Store in server-side cache + if self._enable_signature_cache: + self._signature_cache.store(tool_call_id, signature) + lib_logger.debug(f"Stored thoughtSignature in cache for {tool_call_id}") + + # Pass to client if enabled + if self._preserve_signatures_in_client: + tool_call["thought_signature"] = signature + + tool_calls.append(tool_call) + + # Build message object (not delta!) + message = {"role": "assistant"} + + if text_content: + message["content"] = text_content + elif not tool_calls: + # If no text and no tool calls, set content to empty string + message["content"] = "" + + if reasoning_content: + message["reasoning_content"] = reasoning_content + + if tool_calls: + message["tool_calls"] = tool_calls + # Don't set content if we have tool calls (OpenAI convention) + if "content" in message: + message.pop("content") + + # Handle finish reason + finish_reason = candidate.get("finishReason") + if finish_reason: + # Map Gemini finish reasons to OpenAI + finish_reason_map = { + "STOP": "stop", + "MAX_TOKENS": "length", + "SAFETY": "content_filter", + "RECITATION": "content_filter", + "OTHER": "stop" + } + finish_reason = finish_reason_map.get(finish_reason, "stop") + if tool_calls: + finish_reason = "tool_calls" + + # Build usage metadata + usage = None + usage_metadata = gemini_response.get("usageMetadata", {}) + if usage_metadata: + prompt_tokens = usage_metadata.get("promptTokenCount", 0) + thoughts_tokens = usage_metadata.get("thoughtsTokenCount", 0) + completion_tokens = usage_metadata.get("candidatesTokenCount", 0) + + usage = { + "prompt_tokens": prompt_tokens + thoughts_tokens, + "completion_tokens": completion_tokens, + "total_tokens": usage_metadata.get("totalTokenCount", 0) + } + + # Add reasoning tokens details if thinking was used + if thoughts_tokens > 0: + if "completion_tokens_details" not in usage: + usage["completion_tokens_details"] = {} + usage["completion_tokens_details"]["reasoning_tokens"] = thoughts_tokens + + # Build final response + response = { + "id": gemini_response.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), + "object": "chat.completion", # Non-streaming uses chat.completion, not chunk + "created": int(time.time()), + "model": model, + "choices": [{ + "index": 0, + "message": message, # message, not delta! + "finish_reason": finish_reason + }] + } + + if usage: + response["usage"] = usage + + return response + + # ============================================================================ # PROVIDER INTERFACE IMPLEMENTATION @@ -1374,7 +1861,7 @@ async def acompletion( max_tokens = kwargs.get("max_tokens") enable_request_logging = kwargs.pop("enable_request_logging", False) - lib_logger.info(f"Antigravity completion: model={model}, stream={stream}, messages={len(messages)}") + #lib_logger.debug(f"Antigravity completion: model={model}, stream={stream}, messages={len(messages)}") # Create file logger file_logger = _AntigravityFileLogger( @@ -1424,7 +1911,7 @@ async def acompletion( "parts": [{"text": gemini3_instruction}] } - lib_logger.debug("Gemini 3 system instruction injection applied") + #lib_logger.debug("Gemini 3 system instruction injection applied") @@ -1499,7 +1986,7 @@ async def acompletion( # Apply Gemini 3 specific tool transformations (ONLY for gemini-3-* models) # This implements the "Double-Lock" catch-all strategy to prevent tool hallucination if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: - lib_logger.info(f"Applying Gemini 3 catch-all tool transformations for {model}") + #lib_logger.debug(f"Applying Gemini 3 catch-all tool transformations for {model}") # Strategy 1: Namespace prefixing (breaks association with training data) gemini_cli_payload["tools"] = self._apply_gemini3_namespace_to_tools( @@ -1546,7 +2033,7 @@ async def acompletion( else: headers["Accept"] = "application/json" - lib_logger.debug(f"Antigravity request to: {url}") + #lib_logger.debug(f"Antigravity request to: {url}") try: if stream: @@ -1589,8 +2076,11 @@ async def _handle_non_streaming( # Unwrap Antigravity envelope gemini_response = self._unwrap_antigravity_response(antigravity_response) - # Convert to OpenAI format - return self._gemini_to_openai_chunk(gemini_response, model) + # Convert to OpenAI non-streaming format (returns dict with 'message' not 'delta') + openai_response = self._gemini_to_openai_non_streaming(gemini_response, model) + + # Convert dict to ModelResponse object for non-streaming + return litellm.ModelResponse(**openai_response) async def _handle_streaming( self, From 08736cc493e55052c377311ea7b2efcbabebf776 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 01:25:27 +0100 Subject: [PATCH 017/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20add=20?= =?UTF-8?q?Claude=20support=20and=20parse=20double-encoded=20JSON=20in=20t?= =?UTF-8?q?ool=20args?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extend reasoning/thinking mapping to include Claude alongside Gemini 2.5 and Gemini 3: - Claude now uses `thinkingBudget` (same handling as Gemini 2.5, including pro budgets). - Gemini 3 continues to use `thinkingLevel`. - Add a static helper `_recursively_parse_json_strings` to detect and parse JSON-stringified values returned by Antigravity (e.g., `{"files": "[{...}]"}`) and recursively restore proper structures. - Use parsed arguments before `json.dumps()` when building tool call payloads to prevent double-encoding and JSON parsing errors from Antigravity responses. - Update .gitignore to add `launcher_config.json` and `cache/antigravity/thought_signatures.json` and remove the previous `*.log` ignore entry. --- .gitignore | 3 +- .../providers/antigravity_provider.py | 85 +++++++++++++++---- 2 files changed, 71 insertions(+), 17 deletions(-) diff --git a/.gitignore b/.gitignore index d42c6b8a..0d40840f 100644 --- a/.gitignore +++ b/.gitignore @@ -54,7 +54,6 @@ coverage.xml *.pot # Django stuff: -*.log local_settings.py db.sqlite3 db.sqlite3-journal @@ -124,4 +123,6 @@ test_proxy.py start_proxy.bat key_usage.json staged_changes.txt +launcher_config.json +cache/antigravity/thought_signatures.json logs/ diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index e916fa5c..262943b8 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -926,17 +926,15 @@ def _map_reasoning_effort_to_thinking_config( custom_reasoning_budget: bool = False ) -> Optional[Dict[str, Any]]: """ - Map reasoning_effort to thinking configuration for Gemini 2.5 and 3 models. + Map reasoning_effort to thinking configuration for Gemini 2.5, Gemini 3, and Claude models. - IMPORTANT: This function ONLY applies to Gemini 2.5 and 3 models. - For other models (e.g., Claude via Antigravity), it returns None. - - Gemini 2.5 and 3 use separate budgeting systems: + Supports thinking/reasoning via Antigravity for: - Gemini 2.5: thinkingBudget (integer tokens, based on Gemini CLI logic) - Gemini 3: thinkingLevel (string: "low" or "high") + - Claude: thinkingBudget (same as Gemini 2.5, proxied by Antigravity backend) Default behavior (no reasoning_effort): - - Gemini 2.5: thinkingBudget=-1 (auto mode) + - Gemini 2.5 & Claude: thinkingBudget=-1 (auto mode) - Gemini 3: thinkingLevel="high" (always enabled at high level) Args: @@ -945,23 +943,23 @@ def _map_reasoning_effort_to_thinking_config( custom_reasoning_budget: If True, use full budgets; if False, divide by 4 Returns: - Dict with thinkingConfig or None if not a Gemini 2.5/3 model + Dict with thinkingConfig or None if model doesn't support thinking """ internal_model = self._alias_to_model_name(model) - # Detect model family - ONLY support gemini-2.5 and gemini-3 - # For other models (Claude, etc.), return None without filtering + # Detect model family is_gemini_25 = "gemini-2.5" in model is_gemini_3 = internal_model.startswith("gemini-3-") + is_claude = "claude" in model.lower() - # Return None for unsupported models - no reasoning config changes - if not is_gemini_25 and not is_gemini_3: + # Only Gemini 2.5, Gemini 3, and Claude support thinking via Antigravity + if not is_gemini_25 and not is_gemini_3 and not is_claude: return None # ======================================================================== - # GEMINI 2.5: Use Gemini CLI logic with thinkingBudget + # GEMINI 2.5 & CLAUDE: Use thinkingBudget (INTEGER) # ======================================================================== - if is_gemini_25: + if is_gemini_25 or is_claude: # Default: auto mode if not reasoning_effort: return {"thinkingBudget": -1, "include_thoughts": True} @@ -970,8 +968,9 @@ def _map_reasoning_effort_to_thinking_config( if reasoning_effort == "disable": return {"thinkingBudget": 0, "include_thoughts": False} - # Model-specific budgets (same as Gemini CLI) - if "gemini-2.5-pro" in model: + # Model-specific budgets + # Claude uses Gemini 2.5 pro budgets (high-quality thinking) + if "gemini-2.5-pro" in model or is_claude: budgets = {"low": 8192, "medium": 16384, "high": 32768} elif "gemini-2.5-flash" in model: budgets = {"low": 6144, "medium": 12288, "high": 24576} @@ -1408,6 +1407,48 @@ def _unwrap_antigravity_response(self, antigravity_response: Dict[str, Any]) -> # For both streaming and non-streaming, response is in 'response' field return antigravity_response.get("response", antigravity_response) + @staticmethod + def _recursively_parse_json_strings(obj: Any) -> Any: + """ + Recursively parse JSON strings in nested data structures. + + Antigravity (especially for Claude models) sometimes returns tool arguments + with JSON-stringified values: {"files": "[{...}]"} instead of {"files": [{...}]}. + This causes double-encoding when we call json.dumps() on it. + + This function recursively detects and parses such strings to restore proper structure. + + Args: + obj: Any value (dict, list, str, etc.) + + Returns: + Parsed version with JSON strings converted to their object form + """ + if isinstance(obj, dict): + # Recursively process dictionary values + return {k: AntigravityProvider._recursively_parse_json_strings(v) for k, v in obj.items()} + elif isinstance(obj, list): + # Recursively process list items + return [AntigravityProvider._recursively_parse_json_strings(item) for item in obj] + elif isinstance(obj, str): + # Check if this string looks like JSON + stripped = obj.strip() + if (stripped.startswith('{') and stripped.endswith('}')) or \ + (stripped.startswith('[') and stripped.endswith(']')): + try: + # Attempt to parse as JSON + parsed = json.loads(obj) + # Recursively process the parsed result (it might contain more JSON strings) + return AntigravityProvider._recursively_parse_json_strings(parsed) + except (json.JSONDecodeError, ValueError): + # Not valid JSON, return as-is + return obj + else: + return obj + else: + # Primitive types (int, bool, None, etc.) - return as-is + return obj + def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> Dict[str, Any]: """ Convert a Gemini API response chunk to OpenAI format. @@ -1417,6 +1458,10 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D - Includes signatures in response (if client passthrough enabled) - Filters standalone signature parts (no functionCall/text) + FIXED: Handles Antigravity's double-encoded JSON in tool arguments + - Recursively parses JSON-stringified values before serialization + - Prevents "Unexpected non-whitespace character after JSON" errors + Args: gemini_chunk: Gemini API response chunk model: Model name for Gemini 3 detection @@ -1621,12 +1666,20 @@ def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: tool_name = self._strip_gemini3_namespace_from_name(tool_name) + # Get raw args from Antigravity + raw_args = func_call.get("args", {}) + + # FIX: Recursively parse JSON-stringified values + # Antigravity (especially Claude) returns: {"files": "[{...}]"} + # We need to parse these strings before calling json.dumps() + parsed_args = self._recursively_parse_json_strings(raw_args) + tool_call = { "id": tool_call_id, "type": "function", "function": { "name": tool_name, - "arguments": json.dumps(func_call.get("args", {})) + "arguments": json.dumps(parsed_args) } } From 78eef9662cc55810aac915f11373f7b495af57ca Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 02:23:46 +0100 Subject: [PATCH 018/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20add=20?= =?UTF-8?q?Claude=20thinking=20caching=20and=20generalize=20Antigravity=20?= =?UTF-8?q?cache=20handling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Split the single signature cache into separate files: `GEMINI3_SIGNATURE_CACHE_FILE` and `CLAUDE_THINKING_CACHE_FILE`. - Replace `ThoughtSignatureCache` with `AntigravityCache`; disk persistence file is now passed via a `cache_file` constructor argument and in-memory entries are keyed by generic cache keys. - Introduce a stable key generator (`_generate_thinking_cache_key`) that combines tool call IDs and text hashes for Claude thinking caching. - Add separate caches for Gemini 3 signatures (`_signature_cache`) and Claude thinking content (`_thinking_cache`), and wire caching into both streaming and non-streaming flows. - Accumulate reasoning content, tool calls, and the final `thoughtSignature` during streaming (via `stream_accumulator`) and persist complete Claude thinking after the stream (`_cache_claude_thinking_after_stream`). - Inject cached Claude "thinking" parts into assistant messages when available (with signature fallback handling). - Use tool-provided IDs when present (fall back to generated `call_` IDs), fix skipping logic for signature-only parts, and accumulate tool calls/text for reliable cache keys. - Adjust reasoning budget division from `// 4` to `// 6` to reduce default thinking budget. - Update `_gemini_to_openai_chunk` signature to accept an optional `stream_accumulator` and propagate accumulator through streaming logic. BREAKING CHANGE: `ThoughtSignatureCache` has been removed/renamed to `AntigravityCache` and its constructor now requires a `cache_file: Path` argument. Update any external imports/usages: - Replace `ThoughtSignatureCache(...)` with `AntigravityCache(cache_file=GEMINI3_SIGNATURE_CACHE_FILE|CLAUDE_THINKING_CACHE_FILE, memory_ttl_seconds=..., disk_ttl_seconds=...)`. - New cache constants `GEMINI3_SIGNATURE_CACHE_FILE` and `CLAUDE_THINKING_CACHE_FILE` were added; ensure integrations use the new names if relying on disk cache paths. --- .../providers/antigravity_provider.py | 300 ++++++++++++++++-- 1 file changed, 269 insertions(+), 31 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 262943b8..d4c469e9 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -50,7 +50,9 @@ # Cache configuration CACHE_DIR = Path(__file__).resolve().parent.parent.parent.parent / "cache" ANTIGRAVITY_CACHE_DIR = CACHE_DIR / "antigravity" -ANTIGRAVITY_CACHE_FILE = ANTIGRAVITY_CACHE_DIR / "thought_signatures.json" +# Separate cache files for different data types +GEMINI3_SIGNATURE_CACHE_FILE = ANTIGRAVITY_CACHE_DIR / "gemini3_signatures.json" +CLAUDE_THINKING_CACHE_FILE = ANTIGRAVITY_CACHE_DIR / "claude_thinking.json" class _AntigravityFileLogger: @@ -107,12 +109,13 @@ def log_final_response(self, response_data: Dict[str, Any]): except Exception as e: lib_logger.error(f"_AntigravityFileLogger: Failed to write final response: {e}") -class ThoughtSignatureCache: +class AntigravityCache: """ - Server-side cache for thoughtSignatures to maintain Gemini 3 conversation context. + Server-side cache for Antigravity conversation state preservation. - Maps tool_call_id → thoughtSignature to preserve encrypted reasoning signatures - across turns, even if clients don't support the thought_signature field. + Supports two types of cached data: + 1. Gemini 3: thoughtSignatures (tool_call_id → encrypted signature) + 2. Claude: Thinking content (composite_key → thinking text + signature) Features: - Dual-TTL system: 1hr memory, 24hr disk @@ -123,15 +126,16 @@ class ThoughtSignatureCache: - High concurrency support with asyncio locks """ - def __init__(self, memory_ttl_seconds: int = 3600, disk_ttl_seconds: int = 86400): + def __init__(self, cache_file: Path, memory_ttl_seconds: int = 3600, disk_ttl_seconds: int = 86400): """ - Initialize the signature cache with disk persistence. + Initialize the cache with disk persistence. Args: + cache_file: Path to cache file for disk persistence memory_ttl_seconds: Time-to-live for memory cache entries (default: 1 hour) disk_ttl_seconds: Time-to-live for disk cache entries (default: 24 hours) """ - # In-memory cache: {call_id: (signature, timestamp)} + # In-memory cache: {cache_key: (data, timestamp)} self._cache: Dict[str, Tuple[str, float]] = {} self._memory_ttl = memory_ttl_seconds self._disk_ttl = disk_ttl_seconds @@ -139,7 +143,7 @@ def __init__(self, memory_ttl_seconds: int = 3600, disk_ttl_seconds: int = 86400 self._disk_lock = asyncio.Lock() # Disk persistence configuration - self._cache_file = ANTIGRAVITY_CACHE_FILE + self._cache_file = cache_file self._enable_disk_persistence = os.getenv( "ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", "true" @@ -530,10 +534,20 @@ def __init__(self): self._current_base_url = BASE_URLS[0] # Start with daily sandbox self._base_url_index = 0 - # Initialize thoughtSignature cache for Gemini 3 multi-turn conversations + # Initialize caches for conversation state preservation memory_ttl = int(os.getenv("ANTIGRAVITY_SIGNATURE_CACHE_TTL", "3600")) disk_ttl = int(os.getenv("ANTIGRAVITY_SIGNATURE_DISK_TTL", "86400")) - self._signature_cache = ThoughtSignatureCache( + + # Cache for Gemini 3 thoughtSignatures + self._signature_cache = AntigravityCache( + cache_file=GEMINI3_SIGNATURE_CACHE_FILE, + memory_ttl_seconds=memory_ttl, + disk_ttl_seconds=disk_ttl + ) + + # Cache for Claude thinking content + self._thinking_cache = AntigravityCache( + cache_file=CLAUDE_THINKING_CACHE_FILE, memory_ttl_seconds=memory_ttl, disk_ttl_seconds=disk_ttl ) @@ -622,6 +636,46 @@ def __init__(self): lib_logger.debug("Antigravity: Gemini 3 tool fix DISABLED (using default tool schemas)") + def _generate_thinking_cache_key(self, text_content: str, tool_calls: List[Dict]) -> Optional[str]: + """ + Generate stable cache key from response content for Claude thinking preservation. + + Uses composite key strategy: + - If tool calls exist: Use first tool call ID (most reliable) + - If text exists: Use text hash + - If both: Combine both for maximum uniqueness + + Args: + text_content: Regular text from response + tool_calls: List of tool calls with IDs + + Returns: + Cache key string, or None if no cacheable content + """ + import hashlib + key_parts = [] + + # Priority 1: Tool call IDs (most stable - we generate these) + if tool_calls and len(tool_calls) > 0: + first_tool_id = tool_calls[0].get("id", "") + if first_tool_id: + # Remove 'call_' prefix if present for shorter key + tool_id_short = first_tool_id.replace("call_", "") + key_parts.append(f"tool_{tool_id_short}") + + # Priority 2: Text hash (for text-only or mixed responses) + if text_content: + # Use first 200 chars for stability (longer text may vary slightly) + text_hash = hashlib.md5(text_content[:200].encode()).hexdigest()[:16] + key_parts.append(f"text_{text_hash}") + + # Combine parts + if key_parts: + return "thinking_" + "_".join(key_parts) + + # Shouldn't happen - responses always have text or tools + return None + # ============================================================================ # MODEL ALIAS SYSTEM @@ -828,6 +882,51 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tup lib_logger.warning(f"Failed to parse image data URL: {e}") elif role == "assistant": + # Try to retrieve cached thinking for Claude models + thinking_to_inject = None + cache_key = None + + if model.startswith("claude-") and self._enable_signature_cache: + # Build cache key from incoming message + msg_text = content if isinstance(content, str) else "" + msg_tools = msg.get("tool_calls", []) + + cache_key = self._generate_thinking_cache_key(msg_text, msg_tools) + + if cache_key: + cached_json = self._thinking_cache.retrieve(cache_key) + if cached_json: + try: + thinking_to_inject = json.loads(cached_json) + lib_logger.debug(f"✓ Retrieved thinking from cache: {cache_key[:50]}...") + except json.JSONDecodeError: + lib_logger.warning(f"Failed to parse cached thinking for: {cache_key}") + + # Inject thinking FIRST if we have it + if thinking_to_inject: + thinking_text = thinking_to_inject.get("thinking_text", "") + thought_sig = thinking_to_inject.get("thought_signature", "") + + if thinking_text: + thinking_part = { + "text": thinking_text, + "thought": True + } + + # Add signature if available, otherwise use skip validator + if thought_sig: + thinking_part["thoughtSignature"] = thought_sig + else: + thinking_part["thoughtSignature"] = "skip_thought_signature_validator" + lib_logger.debug("Using skip validator for missing signature") + + parts.append(thinking_part) + lib_logger.debug( + f"✅ Injected {len(thinking_text)} chars of thinking " + f"(sig={'yes' if thought_sig else 'fallback'})" + ) + + # Then add regular content if isinstance(content, str) and content: parts.append({"text": content}) if msg.get("tool_calls"): @@ -983,7 +1082,7 @@ def _map_reasoning_effort_to_thinking_config( # Apply custom_reasoning_budget toggle # If False (default), divide by 4 like Gemini CLI if not custom_reasoning_budget: - budget = budget // 4 + budget = budget // 6 return {"thinkingBudget": budget, "include_thoughts": True} @@ -1449,7 +1548,12 @@ def _recursively_parse_json_strings(obj: Any) -> Any: # Primitive types (int, bool, None, etc.) - return as-is return obj - def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> Dict[str, Any]: + def _gemini_to_openai_chunk( + self, + gemini_chunk: Dict[str, Any], + model: str, + stream_accumulator: Optional[Dict[str, Any]] = None + ) -> Dict[str, Any]: """ Convert a Gemini API response chunk to OpenAI format. @@ -1462,9 +1566,15 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D - Recursively parses JSON-stringified values before serialization - Prevents "Unexpected non-whitespace character after JSON" errors + Claude Thinking Caching: + - For Claude models, thinking content is accumulated across all chunks + - The stream_accumulator collects reasoning_content and thought_signature + - Caching happens AFTER the full stream is processed (in _handle_streaming) + Args: gemini_chunk: Gemini API response chunk model: Model name for Gemini 3 detection + stream_accumulator: Optional dict to accumulate streaming data for post-processing Returns: OpenAI-compatible response chunk @@ -1492,24 +1602,36 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D has_function_call = "functionCall" in part has_text = "text" in part has_signature = "thoughtSignature" in part and part["thoughtSignature"] + is_thought = part.get("thought") is True or (isinstance(part.get("thought"), str) and part.get("thought").lower() == 'true') + + # Accumulate thought signature from thinking parts (Claude caching) + # The signature appears on the LAST thinking part (the one with empty text after all thinking) + if has_signature and is_thought and stream_accumulator is not None: + stream_accumulator["thought_signature"] = part["thoughtSignature"] - # FIXED: Only skip if ONLY signature (standalone encryption part) - # Previously this filtered out ALL function calls with signatures! - if has_signature and not has_function_call and not has_text: - continue # Skip standalone signature parts + # Skip standalone signature-only parts (empty thinking parts with just signature) + if has_signature and not has_function_call and (not has_text or part.get("text") == ""): + continue # Process text content if has_text: - thought = part.get("thought") - if thought is True or (isinstance(thought, str) and thought.lower() == 'true'): + if is_thought: reasoning_content += part["text"] + # Accumulate reasoning for Claude caching + if stream_accumulator is not None: + stream_accumulator["reasoning_content"] += part["text"] else: text_content += part["text"] + # Accumulate text content for cache key generation + if stream_accumulator is not None: + stream_accumulator["text_content"] += part["text"] # Process function calls (NOW WORKS with signatures!) if has_function_call: func_call = part["functionCall"] - tool_call_id = f"call_{uuid.uuid4().hex[:24]}" + + # Use ID from Antigravity if provided, otherwise generate + tool_call_id = func_call.get("id") or f"call_{uuid.uuid4().hex[:24]}" # Get tool name and strip gemini3_ namespace if present (Gemini 3 specific) tool_name = func_call.get("name", "") @@ -1527,7 +1649,11 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D } tool_call_index += 1 # Increment for next tool call - # Handle thoughtSignature if present + # Accumulate tool calls for Claude caching + if stream_accumulator is not None: + stream_accumulator["tool_calls"].append(tool_call) + + # Handle thoughtSignature if present (on function call part) if has_signature and not first_signature_seen: # Only first tool call gets signature (parallel call handling) first_signature_seen = True @@ -1570,7 +1696,11 @@ def _gemini_to_openai_chunk(self, gemini_chunk: Dict[str, Any], model: str) -> D finish_reason = finish_reason_map.get(finish_reason, "stop") if tool_calls: finish_reason = "tool_calls" - + + # Mark stream as complete for accumulator + if stream_accumulator is not None: + stream_accumulator["is_complete"] = True + # Build usage metadata usage = None usage_metadata = gemini_chunk.get("usageMetadata", {}) @@ -1614,6 +1744,7 @@ def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model Convert a Gemini API response to OpenAI non-streaming format. This is specifically for non-streaming completions where we need 'message' instead of 'delta'. + Also handles Claude thinking caching for non-streaming responses. Args: gemini_response: Gemini API response @@ -1635,6 +1766,7 @@ def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model text_content = "" reasoning_content = "" tool_calls = [] + thought_signature = "" # Track signature for Claude caching # Track if we've seen a signature yet (for parallel tool call handling) first_signature_seen = False @@ -1643,15 +1775,19 @@ def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model has_function_call = "functionCall" in part has_text = "text" in part has_signature = "thoughtSignature" in part and part["thoughtSignature"] + is_thought = part.get("thought") is True or (isinstance(part.get("thought"), str) and part.get("thought").lower() == 'true') + + # Capture thought signature (appears on last thinking part) + if has_signature and is_thought: + thought_signature = part["thoughtSignature"] - # Skip standalone signature parts - if has_signature and not has_function_call and not has_text: + # Skip standalone signature parts (empty thinking parts with just signature) + if has_signature and not has_function_call and (not has_text or part.get("text") == ""): continue # Process text content if has_text: - thought = part.get("thought") - if thought is True or (isinstance(thought, str) and thought.lower() == 'true'): + if is_thought: reasoning_content += part["text"] else: text_content += part["text"] @@ -1659,7 +1795,9 @@ def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model # Process function calls if has_function_call: func_call = part["functionCall"] - tool_call_id = f"call_{uuid.uuid4().hex[:24]}" + + # Use ID from Antigravity if provided, otherwise generate + tool_call_id = func_call.get("id") or f"call_{uuid.uuid4().hex[:24]}" # Get tool name and strip gemini3_ namespace if present tool_name = func_call.get("name", "") @@ -1683,7 +1821,7 @@ def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model } } - # Handle thoughtSignature if present + # Handle thoughtSignature if present (on function call part) if has_signature and not first_signature_seen: first_signature_seen = True signature = part["thoughtSignature"] @@ -1699,6 +1837,27 @@ def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model tool_calls.append(tool_call) + # Cache Claude thinking content for non-streaming responses + if reasoning_content and model.startswith("claude-") and self._enable_signature_cache: + cache_key = self._generate_thinking_cache_key(text_content, tool_calls) + + if cache_key: + thinking_data = { + "thinking_text": reasoning_content, + "thought_signature": thought_signature, + "text_preview": text_content[:100] if text_content else "", + "tool_ids": [tc.get("id", "") for tc in tool_calls] if tool_calls else [], + "timestamp": time.time() + } + + self._thinking_cache.store(cache_key, json.dumps(thinking_data)) + lib_logger.info( + f"✓ Cached Claude thinking (non-streaming): {cache_key[:50]}... " + f"(reasoning={len(reasoning_content)} chars, " + f"tools={len(tool_calls)}, " + f"sig={'yes' if thought_signature else 'no'})" + ) + # Build message object (not delta!) message = {"role": "assistant"} @@ -2144,7 +2303,24 @@ async def _handle_streaming( model: str, file_logger: Optional[_AntigravityFileLogger] = None ) -> AsyncGenerator[litellm.ModelResponse, None]: - """Handle streaming completion.""" + """ + Handle streaming completion. + + For Claude models with thinking enabled: + - Accumulates reasoning content and thought signature across all chunks + - Caches the complete thinking data AFTER the stream is fully processed + - Uses a generator wrapper to ensure post-stream caching happens + """ + # Create stream accumulator for Claude thinking caching + # This collects data across all chunks so we can cache after stream completes + stream_accumulator = { + "reasoning_content": "", + "thought_signature": "", + "text_content": "", + "tool_calls": [], + "is_complete": False + } if model.startswith("claude-") and self._enable_signature_cache else None + async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response: # Log error response body for debugging if request failed if response.status_code >= 400: @@ -2172,8 +2348,12 @@ async def _handle_streaming( # Unwrap Antigravity envelope gemini_chunk = self._unwrap_antigravity_response(antigravity_chunk) - # Convert to OpenAI format - openai_chunk = self._gemini_to_openai_chunk(gemini_chunk, model) + # Convert to OpenAI format (with accumulator for Claude) + openai_chunk = self._gemini_to_openai_chunk( + gemini_chunk, + model, + stream_accumulator + ) # Convert dict to ModelResponse object model_response = litellm.ModelResponse(**openai_chunk) @@ -2183,6 +2363,64 @@ async def _handle_streaming( file_logger.log_error(f"Failed to parse chunk: {data_str[:100]}") lib_logger.warning(f"Failed to parse Antigravity chunk: {data_str[:100]}") continue + + # After stream completes: cache Claude thinking content + if stream_accumulator and stream_accumulator.get("reasoning_content"): + await self._cache_claude_thinking_after_stream(stream_accumulator, model) + + async def _cache_claude_thinking_after_stream( + self, + accumulator: Dict[str, Any], + model: str + ): + """ + Cache Claude thinking content after the complete stream has been processed. + + This is called after ALL streaming chunks have been received, ensuring we have: + - Complete reasoning content (accumulated from all thought=true parts) + - The thoughtSignature (appears on the final thinking part) + - All tool calls with their IDs (for cache key generation) + - Complete text content (for cache key generation) + + Args: + accumulator: Dict with accumulated stream data + model: Model name (for logging) + """ + reasoning_content = accumulator.get("reasoning_content", "") + thought_signature = accumulator.get("thought_signature", "") + text_content = accumulator.get("text_content", "") + tool_calls = accumulator.get("tool_calls", []) + + if not reasoning_content: + lib_logger.debug("No reasoning content to cache") + return + + # Generate cache key from the accumulated response data + cache_key = self._generate_thinking_cache_key(text_content, tool_calls) + + if not cache_key: + lib_logger.warning("Could not generate cache key for Claude thinking") + return + + # Build cache data + thinking_data = { + "thinking_text": reasoning_content, + "thought_signature": thought_signature, + "text_preview": text_content[:100] if text_content else "", + "tool_ids": [tc.get("id", "") for tc in tool_calls] if tool_calls else [], + "timestamp": time.time() + } + + # Store in cache + self._thinking_cache.store(cache_key, json.dumps(thinking_data)) + + lib_logger.info( + f"✓ Cached Claude thinking after stream: {cache_key[:50]}... " + f"(reasoning={len(reasoning_content)} chars, " + f"text={len(text_content)} chars, " + f"tools={len(tool_calls)}, " + f"sig={'yes' if thought_signature else 'no'})" + ) # ============================================================================ # TOKEN COUNTING From 0ff233dfdcdaf2f71ff3fc6c4077e29876e3537b Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 04:00:06 +0100 Subject: [PATCH 019/221] =?UTF-8?q?refactor(gemini):=20=F0=9F=94=A8=20impl?= =?UTF-8?q?ement=20official=20Gemini=20CLI=20discovery=20flow=20with=20tie?= =?UTF-8?q?r-based=20onboarding?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit refactors the project discovery logic to strictly follow the official Gemini CLI behavior, fixing critical issues with paid tier support and free tier onboarding. Key changes: - Implement proper discovery flow: cache → configured override → persisted credentials → loadCodeAssist check → tier-based onboarding → fallback - Fix paid tier support: paid tiers now correctly use configured project_id instead of server-managed projects - Fix free tier onboarding: free tier correctly passes cloudaicompanionProject=None for server-managed projects - Add comprehensive tier detection logic: check currentTier from server response and respect userDefinedCloudaicompanionProject flag - Improve error handling: add specific error messages for 412 (precondition failed) and better guidance for missing project_id on paid tiers - Add detailed debug logging: log all tier information, server responses, and decision flow for troubleshooting - Add paid tier visibility: log paid tier usage on each request for transparency - Remove noisy debug logging: disable verbose chunk conversion logs The previous implementation incorrectly assumed all users should use server-managed projects and failed to properly distinguish between free tier (server-managed) and paid tier (user-provided) project handling. This caused 403/412 errors for paid users and incorrect onboarding flow for free users. --- .../providers/gemini_cli_provider.py | 276 +++++++++++++----- 1 file changed, 210 insertions(+), 66 deletions(-) diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 140da2ce..47572fd6 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -94,7 +94,22 @@ def __init__(self): self.project_tier_cache: Dict[str, str] = {} # Cache project tier per credential path async def _discover_project_id(self, credential_path: str, access_token: str, litellm_params: Dict[str, Any]) -> str: - """Discovers the Google Cloud Project ID, with caching and onboarding for new accounts.""" + """ + Discovers the Google Cloud Project ID, with caching and onboarding for new accounts. + + This follows the official Gemini CLI discovery flow: + 1. Check in-memory cache + 2. Check configured project_id override (litellm_params or env var) + 3. Check persisted project_id in credential file + 4. Call loadCodeAssist to check if user is already known (has currentTier) + - If currentTier exists AND cloudaicompanionProject returned: use server's project + - If currentTier exists but NO cloudaicompanionProject: use configured project_id (paid tier requires this) + - If no currentTier: user needs onboarding + 5. Onboard user based on tier: + - FREE tier: pass cloudaicompanionProject=None (server-managed) + - PAID tier: pass cloudaicompanionProject=configured_project_id + 6. Fallback to GCP Resource Manager project listing + """ lib_logger.debug(f"Starting project discovery for credential: {credential_path}") # Check in-memory cache first @@ -103,14 +118,13 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li lib_logger.debug(f"Using cached project ID: {cached_project}") return cached_project - # Check for configured project ID override - if litellm_params.get("project_id"): - project_id = litellm_params["project_id"] - lib_logger.info(f"Using configured Gemini CLI project ID: {project_id}") - self.project_id_cache[credential_path] = project_id - return project_id + # Check for configured project ID override (from litellm_params or env var) + # This is REQUIRED for paid tier users per the official CLI behavior + configured_project_id = litellm_params.get("project_id") + if configured_project_id: + lib_logger.debug(f"Found configured project_id override: {configured_project_id}") - # [NEW] Load credentials from file to check for persisted project_id and tier + # Load credentials from file to check for persisted project_id and tier try: with open(credential_path, 'r') as f: creds = json.load(f) @@ -139,64 +153,168 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li discovered_tier = None async with httpx.AsyncClient() as client: - # 1. Try discovery endpoint with onboarding logic + # 1. Try discovery endpoint with loadCodeAssist lib_logger.debug("Attempting project discovery via Code Assist loadCodeAssist endpoint...") try: - initial_project_id = "default" - client_metadata = { - "ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", - "pluginType": "GEMINI", "duetProject": initial_project_id, + # Build metadata - include duetProject only if we have a configured project + core_client_metadata = { + "ideType": "IDE_UNSPECIFIED", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + } + if configured_project_id: + core_client_metadata["duetProject"] = configured_project_id + + # Build load request - pass configured_project_id if available, otherwise None + load_request = { + "cloudaicompanionProject": configured_project_id, # Can be None + "metadata": core_client_metadata, } - load_request = {"cloudaicompanionProject": initial_project_id, "metadata": client_metadata} + lib_logger.debug(f"Sending loadCodeAssist request with cloudaicompanionProject={configured_project_id}") response = await client.post(f"{CODE_ASSIST_ENDPOINT}:loadCodeAssist", headers=headers, json=load_request, timeout=20) response.raise_for_status() data = response.json() - # Extract tier information for paid project detection - selected_tier_id = None - allowed_tiers = data.get('allowedTiers', []) - lib_logger.debug(f"Available tiers from loadCodeAssist response: {[t.get('id') for t in allowed_tiers]}") + # Log full response for debugging + lib_logger.debug(f"loadCodeAssist full response keys: {list(data.keys())}") + # Extract and log ALL tier information for debugging + allowed_tiers = data.get('allowedTiers', []) + current_tier = data.get('currentTier') + + lib_logger.debug(f"=== Tier Information ===") + lib_logger.debug(f"currentTier: {current_tier}") + lib_logger.debug(f"allowedTiers count: {len(allowed_tiers)}") + for i, tier in enumerate(allowed_tiers): + tier_id = tier.get('id', 'unknown') + is_default = tier.get('isDefault', False) + user_defined = tier.get('userDefinedCloudaicompanionProject', False) + lib_logger.debug(f" Tier {i+1}: id={tier_id}, isDefault={is_default}, userDefinedProject={user_defined}") + lib_logger.debug(f"========================") + + # Determine the current tier ID + current_tier_id = None + if current_tier: + current_tier_id = current_tier.get('id') + lib_logger.debug(f"User has currentTier: {current_tier_id}") + + # Check if user is already known to server (has currentTier) + if current_tier_id: + # User is already onboarded - check for project from server + server_project = data.get('cloudaicompanionProject') + + # Check if this tier requires user-defined project (paid tiers) + requires_user_project = any( + t.get('id') == current_tier_id and t.get('userDefinedCloudaicompanionProject', False) + for t in allowed_tiers + ) + is_free_tier = current_tier_id == 'free-tier' + + if server_project: + # Server returned a project - use it (server wins) + # This is the normal case for FREE tier users + project_id = server_project + lib_logger.debug(f"Server returned project: {project_id}") + elif configured_project_id: + # No server project but we have configured one - use it + # This is the PAID TIER case where server doesn't return a project + project_id = configured_project_id + lib_logger.debug(f"No server project, using configured: {project_id}") + elif is_free_tier: + # Free tier user without server project - this shouldn't happen normally + # but let's not fail, just proceed to onboarding + lib_logger.debug("Free tier user with currentTier but no project - will try onboarding") + project_id = None + elif requires_user_project: + # Paid tier requires a project ID to be set + raise ValueError( + f"Paid tier '{current_tier_id}' requires setting GEMINI_CLI_PROJECT_ID environment variable. " + "See https://goo.gle/gemini-cli-auth-docs#workspace-gca" + ) + else: + # Unknown tier without project - proceed carefully + lib_logger.warning(f"Tier '{current_tier_id}' has no project and none configured - will try onboarding") + project_id = None + + if project_id: + # Cache tier info + self.project_tier_cache[credential_path] = current_tier_id + discovered_tier = current_tier_id + + # Log appropriately based on tier + is_paid = current_tier_id and current_tier_id not in ['free-tier', 'legacy-tier', 'unknown'] + if is_paid: + lib_logger.info(f"Using Gemini paid tier '{current_tier_id}' with project: {project_id}") + else: + lib_logger.info(f"Discovered Gemini project ID via loadCodeAssist: {project_id}") + + self.project_id_cache[credential_path] = project_id + discovered_project_id = project_id + + # Persist to credential file + await self._persist_project_metadata(credential_path, project_id, discovered_tier) + + return project_id + + # 2. User needs onboarding - no currentTier + lib_logger.info("No existing Gemini session found (no currentTier), attempting to onboard user...") + + # Determine which tier to onboard with + onboard_tier = None for tier in allowed_tiers: if tier.get('isDefault'): - selected_tier_id = tier.get('id', 'unknown') - lib_logger.debug(f"Selected default tier: {selected_tier_id}") + onboard_tier = tier break - if not selected_tier_id and allowed_tiers: - selected_tier_id = allowed_tiers[0].get('id', 'unknown') - lib_logger.debug(f"No default tier found, using first available: {selected_tier_id}") - - if data.get('cloudaicompanionProject'): - project_id = data['cloudaicompanionProject'] - lib_logger.debug(f"Existing project found in loadCodeAssist response: {project_id}") - - # Cache tier info - if selected_tier_id: - self.project_tier_cache[credential_path] = selected_tier_id - discovered_tier = selected_tier_id - lib_logger.debug(f"Cached tier information: {selected_tier_id}") - - # Log concise message for paid projects - is_paid = selected_tier_id and selected_tier_id not in ['free-tier', 'legacy-tier', 'unknown'] - if is_paid: - lib_logger.info(f"Using Gemini paid project: {project_id}") - else: - lib_logger.info(f"Discovered Gemini project ID via loadCodeAssist: {project_id}") - - self.project_id_cache[credential_path] = project_id - discovered_project_id = project_id - - # [NEW] Persist to credential file - await self._persist_project_metadata(credential_path, project_id, discovered_tier) - - return project_id - # 2. If no project ID, trigger onboarding - lib_logger.info("No existing Gemini project found, attempting to onboard user...") - tier_id = next((t.get('id', 'free-tier') for t in data.get('allowedTiers', []) if t.get('isDefault')), 'free-tier') - lib_logger.debug(f"Onboarding with tier: {tier_id}") - onboard_request = {"tierId": tier_id, "cloudaicompanionProject": initial_project_id, "metadata": client_metadata} + # Fallback to LEGACY tier if no default (requires user project) + if not onboard_tier and allowed_tiers: + # Look for legacy-tier as fallback + for tier in allowed_tiers: + if tier.get('id') == 'legacy-tier': + onboard_tier = tier + break + # If still no tier, use first available + if not onboard_tier: + onboard_tier = allowed_tiers[0] + + if not onboard_tier: + raise ValueError("No onboarding tiers available from server") + + tier_id = onboard_tier.get('id', 'free-tier') + requires_user_project = onboard_tier.get('userDefinedCloudaicompanionProject', False) + + lib_logger.debug(f"Onboarding with tier: {tier_id}, requiresUserProject: {requires_user_project}") + + # Build onboard request based on tier type (following official CLI logic) + # FREE tier: cloudaicompanionProject = None (server-managed) + # PAID tier: cloudaicompanionProject = configured_project_id (user must provide) + is_free_tier = tier_id == 'free-tier' + + if is_free_tier: + # Free tier uses server-managed project + onboard_request = { + "tierId": tier_id, + "cloudaicompanionProject": None, # Server will create/manage + "metadata": core_client_metadata, + } + lib_logger.debug("Free tier onboarding: using server-managed project") + else: + # Paid/legacy tier requires user-provided project + if not configured_project_id and requires_user_project: + raise ValueError( + f"Tier '{tier_id}' requires setting GEMINI_CLI_PROJECT_ID environment variable. " + "See https://goo.gle/gemini-cli-auth-docs#workspace-gca" + ) + onboard_request = { + "tierId": tier_id, + "cloudaicompanionProject": configured_project_id, + "metadata": { + **core_client_metadata, + "duetProject": configured_project_id, + } if configured_project_id else core_client_metadata, + } + lib_logger.debug(f"Paid tier onboarding: using project {configured_project_id}") lib_logger.debug("Initiating onboardUser request...") lro_response = await client.post(f"{CODE_ASSIST_ENDPOINT}:onboardUser", headers=headers, json=onboard_request, timeout=30) @@ -204,7 +322,7 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li lro_data = lro_response.json() lib_logger.debug(f"Initial onboarding response: done={lro_data.get('done')}") - for i in range(150): # Poll for up to 5 minutes (150 × 2s) + for i in range(150): # Poll for up to 5 minutes (150 × 2s) if lro_data.get('done'): lib_logger.debug(f"Onboarding completed after {i} polling attempts") break @@ -220,41 +338,62 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li lib_logger.error("Onboarding process timed out after 5 minutes") raise ValueError("Onboarding process timed out after 5 minutes. Please try again or contact support.") - project_id = lro_data.get('response', {}).get('cloudaicompanionProject', {}).get('id') + # Extract project ID from LRO response + # Note: onboardUser returns response.cloudaicompanionProject as an object with .id + lro_response_data = lro_data.get('response', {}) + lro_project_obj = lro_response_data.get('cloudaicompanionProject', {}) + project_id = lro_project_obj.get('id') if isinstance(lro_project_obj, dict) else None + + # Fallback to configured project if LRO didn't return one + if not project_id and configured_project_id: + project_id = configured_project_id + lib_logger.debug(f"LRO didn't return project, using configured: {project_id}") + if not project_id: - lib_logger.error("Onboarding completed but no project ID in response") - raise ValueError("Onboarding completed, but no project ID was returned.") + lib_logger.error("Onboarding completed but no project ID in response and none configured") + raise ValueError( + "Onboarding completed, but no project ID was returned. " + "For paid tiers, set GEMINI_CLI_PROJECT_ID environment variable." + ) lib_logger.debug(f"Successfully extracted project ID from onboarding response: {project_id}") # Cache tier info - if tier_id: - self.project_tier_cache[credential_path] = tier_id - discovered_tier = tier_id - lib_logger.debug(f"Cached tier information: {tier_id}") + self.project_tier_cache[credential_path] = tier_id + discovered_tier = tier_id + lib_logger.debug(f"Cached tier information: {tier_id}") # Log concise message for paid projects is_paid = tier_id and tier_id not in ['free-tier', 'legacy-tier'] if is_paid: - lib_logger.info(f"Using Gemini paid project: {project_id}") + lib_logger.info(f"Using Gemini paid tier '{tier_id}' with project: {project_id}") else: lib_logger.info(f"Successfully onboarded user and discovered project ID: {project_id}") self.project_id_cache[credential_path] = project_id discovered_project_id = project_id - # [NEW] Persist to credential file + # Persist to credential file await self._persist_project_metadata(credential_path, project_id, discovered_tier) return project_id except httpx.HTTPStatusError as e: + error_body = "" + try: + error_body = e.response.text + except Exception: + pass if e.response.status_code == 403: - lib_logger.error(f"Gemini Code Assist API access denied (403). The cloudaicompanion.googleapis.com API may not be enabled for your account. Please enable it in Google Cloud Console.") + lib_logger.error(f"Gemini Code Assist API access denied (403). Response: {error_body}") + lib_logger.error("Possible causes: 1) cloudaicompanion.googleapis.com API not enabled, 2) Wrong project ID for paid tier, 3) Account lacks permissions") elif e.response.status_code == 404: lib_logger.warning(f"Gemini Code Assist endpoint not found (404). Falling back to project listing.") + elif e.response.status_code == 412: + # Precondition Failed - often means wrong project for free tier onboarding + lib_logger.error(f"Precondition failed (412): {error_body}. This may mean the project ID is incompatible with the selected tier.") else: - lib_logger.warning(f"Gemini onboarding/discovery failed with status {e.response.status_code}: {e}. Falling back to project listing.") + lib_logger.warning(f"Gemini onboarding/discovery failed with status {e.response.status_code}: {error_body}. Falling back to project listing.") except httpx.RequestError as e: lib_logger.warning(f"Gemini onboarding/discovery network error: {e}. Falling back to project listing.") @@ -499,7 +638,7 @@ def _handle_reasoning_parameters(self, payload: Dict[str, Any], model: str) -> O return {"thinkingBudget": budget, "include_thoughts": True} def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): - lib_logger.debug(f"Converting Gemini chunk: {json.dumps(chunk)}") + #lib_logger.debug(f"Converting Gemini chunk: {json.dumps(chunk)}") response_data = chunk.get('response', chunk) candidates = response_data.get('candidates', []) if not candidates: @@ -778,6 +917,11 @@ async def do_call(attempt_model: str, is_fallback: bool = False): access_token = auth_header['Authorization'].split(' ')[1] project_id = await self._discover_project_id(credential_path, access_token, kwargs.get("litellm_params", {})) + # Log paid tier usage visibly on each request + credential_tier = self.project_tier_cache.get(credential_path) + if credential_tier and credential_tier not in ['free-tier', 'legacy-tier', 'unknown']: + lib_logger.info(f"[PAID TIER] Using Gemini '{credential_tier}' subscription for this request") + # Handle :thinking suffix model_name = attempt_model.split('/')[-1].replace(':thinking', '') From afe6e7051a788dbaa3650290e60d46b01ee5c125 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 04:05:37 +0100 Subject: [PATCH 020/221] =?UTF-8?q?refactor(antigravity):=20=F0=9F=94=A8?= =?UTF-8?q?=20restructure=20provider=20with=20comprehensive=20code=20organ?= =?UTF-8?q?ization=20and=20documentation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This is a major refactoring of the Antigravity provider implementation that significantly improves code structure, readability, and maintainability without changing functionality. Key improvements: - Reorganized code into logical sections with clear separators (configuration, utilities, caching, transformations, API interface) - Consolidated helper functions with consistent naming patterns (underscore prefix for internal methods) - Simplified complex methods by extracting reusable components (e.g., _parse_content_parts, _extract_tool_call, _format_type_hint) - Enhanced documentation with comprehensive module docstring explaining features and capabilities - Streamlined environment variable handling with dedicated helper functions (_env_bool, _env_int) - Improved type hints and method signatures for better IDE support - Reduced code duplication in message transformation logic - Consolidated tool schema transformations into focused methods - Better separation of concerns between streaming and non-streaming response handling - Standardized error handling and logging patterns - Improved cache implementation with clearer separation of responsibilities The refactoring maintains full backward compatibility while making the codebase significantly easier to understand, test, and extend. All existing features including Gemini 3 thoughtSignature preservation, Claude thinking caching, tool hallucination prevention, and base URL fallback remain fully functional. --- .../providers/antigravity_provider.py | 3163 +++++++---------- 1 file changed, 1225 insertions(+), 1938 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index d4c469e9..9223fdaa 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -1,25 +1,48 @@ -# src/rotator_library/providers/antigravity_provider.py +# src/rotator_library/providers/antigravity_provider_v2.py +""" +Antigravity Provider - Refactored Implementation + +A clean, well-structured provider for Google's Antigravity API, supporting: +- Gemini 2.5 (Pro/Flash) with thinkingBudget +- Gemini 3 (Pro/Image) with thinkingLevel +- Claude (Sonnet 4.5) via Antigravity proxy + +Key Features: +- Unified streaming/non-streaming handling +- Server-side thought signature caching +- Automatic base URL fallback +- Gemini 3 tool hallucination prevention +""" + +from __future__ import annotations -import json -import httpx -import logging -import time import asyncio -import random -import uuid import copy -import threading +import hashlib +import json +import logging import os -import tempfile +import random import shutil -from pathlib import Path +import tempfile +import time +import uuid from datetime import datetime -from typing import List, Dict, Any, AsyncGenerator, Union, Optional, Tuple +from pathlib import Path +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union +from urllib.parse import urlparse + +import httpx +import litellm + from .provider_interface import ProviderInterface from .antigravity_auth_base import AntigravityAuthBase from ..model_definitions import ModelDefinitions -import litellm -from litellm.exceptions import RateLimitError + + +# ============================================================================= +# CONFIGURATION CONSTANTS +# ============================================================================= lib_logger = logging.getLogger('rotator_library') @@ -28,11 +51,11 @@ BASE_URLS = [ "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal", "https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal", - "https://cloudcode-pa.googleapis.com/v1internal" # Production fallback + "https://cloudcode-pa.googleapis.com/v1internal", # Production fallback ] -# Hardcoded models available via Antigravity -HARDCODED_MODELS = [ +# Available models via Antigravity +AVAILABLE_MODELS = [ "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite", @@ -40,101 +63,236 @@ "gemini-3-pro-image-preview", "gemini-2.5-computer-use-preview-10-2025", "claude-sonnet-4-5", - "claude-sonnet-4-5-thinking" + "claude-sonnet-4-5-thinking", ] -# Logging configuration -LOGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "logs" -ANTIGRAVITY_LOGS_DIR = LOGS_DIR / "antigravity_logs" +# Default max output tokens (including thinking) - can be overridden per request +DEFAULT_MAX_OUTPUT_TOKENS = 16384 + +# Model alias mappings (internal ↔ public) +MODEL_ALIAS_MAP = { + "rev19-uic3-1p": "gemini-2.5-computer-use-preview-10-2025", + "gemini-3-pro-image": "gemini-3-pro-image-preview", + "gemini-3-pro-high": "gemini-3-pro-preview", +} +MODEL_ALIAS_REVERSE = {v: k for k, v in MODEL_ALIAS_MAP.items()} + +# Models to exclude from dynamic discovery +EXCLUDED_MODELS = {"chat_20706", "chat_23310", "gemini-2.5-flash-thinking", "gemini-3-pro-low", "gemini-2.5-pro"} + +# Gemini finish reason mapping +FINISH_REASON_MAP = { + "STOP": "stop", + "MAX_TOKENS": "length", + "SAFETY": "content_filter", + "RECITATION": "content_filter", + "OTHER": "stop", +} + +# Directory paths +_BASE_DIR = Path(__file__).resolve().parent.parent.parent.parent +LOGS_DIR = _BASE_DIR / "logs" / "antigravity_logs" +CACHE_DIR = _BASE_DIR / "cache" / "antigravity" +GEMINI3_SIGNATURE_CACHE_FILE = CACHE_DIR / "gemini3_signatures.json" +CLAUDE_THINKING_CACHE_FILE = CACHE_DIR / "claude_thinking.json" + +# Gemini 3 tool fix system instruction (prevents hallucination) +DEFAULT_GEMINI3_SYSTEM_INSTRUCTION = """CRITICAL TOOL USAGE INSTRUCTIONS: +You are operating in a custom environment where tool definitions differ from your training data. +You MUST follow these rules strictly: + +1. DO NOT use your internal training data to guess tool parameters +2. ONLY use the exact parameter structure defined in the tool schema +3. If a tool takes a 'files' parameter, it is ALWAYS an array of objects with specific properties, NEVER a simple array of strings +4. If a tool edits code, it takes structured JSON objects with specific fields, NEVER raw diff strings or plain text +5. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') +6. Array parameters have specific item types - check the schema's 'items' field for the exact structure +7. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions + +If you are unsure about a tool's parameters, YOU MUST read the schema definition carefully. Your training data about common tool names like 'read_file' or 'apply_diff' does NOT apply here. +""" + + +# ============================================================================= +# HELPER FUNCTIONS +# ============================================================================= + +def _env_bool(key: str, default: bool = False) -> bool: + """Get boolean from environment variable.""" + return os.getenv(key, str(default).lower()).lower() in ("true", "1", "yes") + + +def _env_int(key: str, default: int) -> int: + """Get integer from environment variable.""" + return int(os.getenv(key, str(default))) + + +def _generate_request_id() -> str: + """Generate Antigravity request ID: agent-{uuid}""" + return f"agent-{uuid.uuid4()}" + + +def _generate_session_id() -> str: + """Generate Antigravity session ID: -{random_number}""" + n = random.randint(1_000_000_000_000_000_000, 9_999_999_999_999_999_999) + return f"-{n}" + + +def _generate_project_id() -> str: + """Generate fake project ID: {adj}-{noun}-{random}""" + adjectives = ["useful", "bright", "swift", "calm", "bold"] + nouns = ["fuze", "wave", "spark", "flow", "core"] + return f"{random.choice(adjectives)}-{random.choice(nouns)}-{uuid.uuid4().hex[:5]}" + + +def _normalize_type_arrays(schema: Any) -> Any: + """ + Normalize type arrays in JSON Schema for Proto-based Antigravity API. + Converts `"type": ["string", "null"]` → `"type": "string"`. + """ + if isinstance(schema, dict): + normalized = {} + for key, value in schema.items(): + if key == "type" and isinstance(value, list): + non_null = [t for t in value if t != "null"] + normalized[key] = non_null[0] if non_null else value[0] + else: + normalized[key] = _normalize_type_arrays(value) + return normalized + elif isinstance(schema, list): + return [_normalize_type_arrays(item) for item in schema] + return schema + + +def _recursively_parse_json_strings(obj: Any) -> Any: + """ + Recursively parse JSON strings in nested data structures. + + Antigravity sometimes returns tool arguments with JSON-stringified values: + {"files": "[{...}]"} instead of {"files": [{...}]}. + """ + if isinstance(obj, dict): + return {k: _recursively_parse_json_strings(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [_recursively_parse_json_strings(item) for item in obj] + elif isinstance(obj, str): + stripped = obj.strip() + if (stripped.startswith('{') and stripped.endswith('}')) or \ + (stripped.startswith('[') and stripped.endswith(']')): + try: + parsed = json.loads(obj) + return _recursively_parse_json_strings(parsed) + except (json.JSONDecodeError, ValueError): + pass + return obj + + +def _clean_claude_schema(schema: Any) -> Any: + """Recursively remove fields that Claude's JSON Schema validation doesn't support.""" + if not isinstance(schema, dict): + return schema + + incompatible = {'$schema', 'additionalProperties', 'minItems', 'maxItems', 'pattern'} + cleaned = {} + + for key, value in schema.items(): + if key in incompatible: + continue + if isinstance(value, dict): + cleaned[key] = _clean_claude_schema(value) + elif isinstance(value, list): + cleaned[key] = [_clean_claude_schema(item) if isinstance(item, dict) else item for item in value] + else: + cleaned[key] = value + + return cleaned -# Cache configuration -CACHE_DIR = Path(__file__).resolve().parent.parent.parent.parent / "cache" -ANTIGRAVITY_CACHE_DIR = CACHE_DIR / "antigravity" -# Separate cache files for different data types -GEMINI3_SIGNATURE_CACHE_FILE = ANTIGRAVITY_CACHE_DIR / "gemini3_signatures.json" -CLAUDE_THINKING_CACHE_FILE = ANTIGRAVITY_CACHE_DIR / "claude_thinking.json" +# ============================================================================= +# FILE LOGGER +# ============================================================================= -class _AntigravityFileLogger: - """A simple file logger for a single Antigravity transaction.""" +class AntigravityFileLogger: + """Transaction file logger for debugging Antigravity requests/responses.""" + + __slots__ = ('enabled', 'log_dir') + def __init__(self, model_name: str, enabled: bool = True): self.enabled = enabled - if not self.enabled: + self.log_dir: Optional[Path] = None + + if not enabled: return - + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") - request_id = str(uuid.uuid4()) - # Sanitize model name for directory - safe_model_name = model_name.replace('/', '_').replace(':', '_') - self.log_dir = ANTIGRAVITY_LOGS_DIR / f"{timestamp}_{safe_model_name}_{request_id}" + safe_model = model_name.replace('/', '_').replace(':', '_') + self.log_dir = LOGS_DIR / f"{timestamp}_{safe_model}_{uuid.uuid4()}" + try: self.log_dir.mkdir(parents=True, exist_ok=True) except Exception as e: - lib_logger.error(f"Failed to create Antigravity log directory: {e}") + lib_logger.error(f"Failed to create log directory: {e}") self.enabled = False - - def log_request(self, payload: Dict[str, Any]): - """Logs the request payload sent to Antigravity.""" - if not self.enabled: return + + def log_request(self, payload: Dict[str, Any]) -> None: + """Log the request payload.""" + self._write_json("request_payload.json", payload) + + def log_response_chunk(self, chunk: str) -> None: + """Append a raw chunk to the response stream log.""" + self._append_text("response_stream.log", chunk) + + def log_error(self, error_message: str) -> None: + """Log an error message.""" + self._append_text("error.log", f"[{datetime.utcnow().isoformat()}] {error_message}") + + def log_final_response(self, response: Dict[str, Any]) -> None: + """Log the final response.""" + self._write_json("final_response.json", response) + + def _write_json(self, filename: str, data: Dict[str, Any]) -> None: + if not self.enabled or not self.log_dir: + return try: - with open(self.log_dir / "request_payload.json", "w", encoding="utf-8") as f: - json.dump(payload, f, indent=2, ensure_ascii=False) + with open(self.log_dir / filename, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) except Exception as e: - lib_logger.error(f"_AntigravityFileLogger: Failed to write request: {e}") - - def log_response_chunk(self, chunk: str): - """Logs a raw chunk from the Antigravity response stream.""" - if not self.enabled: return + lib_logger.error(f"Failed to write {filename}: {e}") + + def _append_text(self, filename: str, text: str) -> None: + if not self.enabled or not self.log_dir: + return try: - with open(self.log_dir / "response_stream.log", "a", encoding="utf-8") as f: - f.write(chunk + "\n") + with open(self.log_dir / filename, "a", encoding="utf-8") as f: + f.write(text + "\n") except Exception as e: - lib_logger.error(f"_AntigravityFileLogger: Failed to write response chunk: {e}") + lib_logger.error(f"Failed to append to {filename}: {e}") - def log_error(self, error_message: str): - """Logs an error message.""" - if not self.enabled: return - try: - with open(self.log_dir / "error.log", "a", encoding="utf-8") as f: - f.write(f"[{datetime.utcnow().isoformat()}] {error_message}\n") - except Exception as e: - lib_logger.error(f"_AntigravityFileLogger: Failed to write error: {e}") - def log_final_response(self, response_data: Dict[str, Any]): - """Logs the final, reassembled response.""" - if not self.enabled: return - try: - with open(self.log_dir / "final_response.json", "w", encoding="utf-8") as f: - json.dump(response_data, f, indent=2, ensure_ascii=False) - except Exception as e: - lib_logger.error(f"_AntigravityFileLogger: Failed to write final response: {e}") +# ============================================================================= +# SIGNATURE CACHE +# ============================================================================= class AntigravityCache: """ Server-side cache for Antigravity conversation state preservation. Supports two types of cached data: - 1. Gemini 3: thoughtSignatures (tool_call_id → encrypted signature) - 2. Claude: Thinking content (composite_key → thinking text + signature) + - Gemini 3: thoughtSignatures (tool_call_id → encrypted signature) + - Claude: Thinking content (composite_key → thinking text + signature) Features: - Dual-TTL system: 1hr memory, 24hr disk - Async disk persistence with batched writes - Background cleanup task for expired entries - - Thread-safe for concurrent access - - Fallback to disk when not in memory - - High concurrency support with asyncio locks """ - def __init__(self, cache_file: Path, memory_ttl_seconds: int = 3600, disk_ttl_seconds: int = 86400): - """ - Initialize the cache with disk persistence. - - Args: - cache_file: Path to cache file for disk persistence - memory_ttl_seconds: Time-to-live for memory cache entries (default: 1 hour) - disk_ttl_seconds: Time-to-live for disk cache entries (default: 24 hours) - """ + def __init__( + self, + cache_file: Path, + memory_ttl_seconds: int = 3600, + disk_ttl_seconds: int = 86400 + ): # In-memory cache: {cache_key: (data, timestamp)} self._cache: Dict[str, Tuple[str, float]] = {} self._memory_ttl = memory_ttl_seconds @@ -142,17 +300,12 @@ def __init__(self, cache_file: Path, memory_ttl_seconds: int = 3600, disk_ttl_se self._lock = asyncio.Lock() self._disk_lock = asyncio.Lock() - # Disk persistence configuration + # Disk persistence self._cache_file = cache_file - self._enable_disk_persistence = os.getenv( - "ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", - "true" - ).lower() in ("true", "1", "yes") - - # Async write configuration - self._dirty = False # Flag for pending writes - self._write_interval = int(os.getenv("ANTIGRAVITY_CACHE_WRITE_INTERVAL", "60")) - self._cleanup_interval = int(os.getenv("ANTIGRAVITY_CACHE_CLEANUP_INTERVAL", "1800")) + self._enable_disk = _env_bool("ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", True) + self._dirty = False + self._write_interval = _env_int("ANTIGRAVITY_CACHE_WRITE_INTERVAL", 60) + self._cleanup_interval = _env_int("ANTIGRAVITY_CACHE_CLEANUP_INTERVAL", 1800) # Background tasks self._writer_task: Optional[asyncio.Task] = None @@ -160,186 +313,121 @@ def __init__(self, cache_file: Path, memory_ttl_seconds: int = 3600, disk_ttl_se self._running = False # Statistics - self._stats = { - "memory_hits": 0, - "disk_hits": 0, - "misses": 0, - "writes": 0 - } + self._stats = {"memory_hits": 0, "disk_hits": 0, "misses": 0, "writes": 0} - # Initialize - if self._enable_disk_persistence: + if self._enable_disk: lib_logger.debug( - f"ThoughtSignatureCache: Disk persistence ENABLED " - f"(memory_ttl={memory_ttl_seconds}s, disk_ttl={disk_ttl_seconds}s, " - f"write_interval={self._write_interval}s)" + f"AntigravityCache: Disk persistence enabled " + f"(memory_ttl={memory_ttl_seconds}s, disk_ttl={disk_ttl_seconds}s)" ) - # Schedule async initialization asyncio.create_task(self._async_init()) else: - lib_logger.debug("ThoughtSignatureCache: Disk persistence DISABLED (memory-only mode)") + lib_logger.debug("AntigravityCache: Memory-only mode") - async def _async_init(self): + async def _async_init(self) -> None: """Async initialization: load from disk and start background tasks.""" try: await self._load_from_disk() await self._start_background_tasks() except Exception as e: - lib_logger.error(f"ThoughtSignatureCache async init failed: {e}") + lib_logger.error(f"Cache async init failed: {e}") - async def _load_from_disk(self): - """Load cache from disk file (with TTL validation).""" - if not self._enable_disk_persistence: - return - - if not self._cache_file.exists(): - lib_logger.debug("No existing cache file found, starting fresh") + async def _load_from_disk(self) -> None: + """Load cache from disk file with TTL validation.""" + if not self._enable_disk or not self._cache_file.exists(): return try: async with self._disk_lock: - # Read cache file with open(self._cache_file, 'r', encoding='utf-8') as f: data = json.load(f) - # Validate version if data.get("version") != "1.0": - lib_logger.warning(f"Cache file version mismatch, ignoring") + lib_logger.warning("Cache version mismatch, starting fresh") return - # Load entries with disk TTL validation now = time.time() entries = data.get("entries", {}) - loaded = 0 - expired = 0 + loaded = expired = 0 for call_id, entry in entries.items(): - timestamp = entry.get("timestamp", 0) - age = now - timestamp - - # Check against DISK TTL (24 hours) + age = now - entry.get("timestamp", 0) if age <= self._disk_ttl: - signature = entry.get("signature", "") - if signature: - self._cache[call_id] = (signature, timestamp) + sig = entry.get("signature", "") + if sig: + self._cache[call_id] = (sig, entry["timestamp"]) loaded += 1 else: expired += 1 - lib_logger.debug( - f"ThoughtSignatureCache: Loaded {loaded} signatures from disk " - f"({expired} expired entries removed)" - ) - + lib_logger.debug(f"Loaded {loaded} entries from disk ({expired} expired)") except json.JSONDecodeError as e: - lib_logger.warning(f"Cache file corrupted, starting fresh: {e}") + lib_logger.warning(f"Cache file corrupted: {e}") except Exception as e: - lib_logger.error(f"Failed to load cache from disk: {e}") + lib_logger.error(f"Failed to load cache: {e}") - async def _save_to_disk(self): + async def _save_to_disk(self) -> None: """Persist cache to disk using atomic write.""" - if not self._enable_disk_persistence: + if not self._enable_disk: return try: async with self._disk_lock: - # Ensure cache directory exists self._cache_file.parent.mkdir(parents=True, exist_ok=True) - # Build cache data structure cache_data = { "version": "1.0", "memory_ttl_seconds": self._memory_ttl, "disk_ttl_seconds": self._disk_ttl, "entries": { - call_id: { - "signature": sig, - "timestamp": ts - } - for call_id, (sig, ts) in self._cache.items() + cid: {"signature": sig, "timestamp": ts} + for cid, (sig, ts) in self._cache.items() }, "statistics": { "total_entries": len(self._cache), "last_write": time.time(), - "memory_hits": self._stats["memory_hits"], - "disk_hits": self._stats["disk_hits"], - "misses": self._stats["misses"], - "writes": self._stats["writes"] + **self._stats } } - # Atomic write using tempfile pattern (same as OAuth credentials) + # Atomic write parent_dir = self._cache_file.parent - tmp_fd = None - tmp_path = None + tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json') try: - # Create temp file in same directory - tmp_fd, tmp_path = tempfile.mkstemp( - dir=parent_dir, - prefix='.tmp_', - suffix='.json', - text=True - ) - - # Write JSON to temp file with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f: json.dump(cache_data, f, indent=2) - tmp_fd = None # fdopen closes the fd - # Set secure permissions (owner read/write only) try: os.chmod(tmp_path, 0o600) except (OSError, AttributeError): - # Windows may not support chmod, ignore pass - # Atomic move (overwrites target if exists) shutil.move(tmp_path, self._cache_file) - tmp_path = None # Successfully moved - self._stats["writes"] += 1 - lib_logger.debug(f"Saved {len(self._cache)} signatures to disk") - - except Exception as e: - lib_logger.error(f"Failed to save cache to disk: {e}") - # Clean up temp file if it still exists - if tmp_fd is not None: - try: - os.close(tmp_fd) - except: - pass + lib_logger.debug(f"Saved {len(self._cache)} entries to disk") + except Exception: if tmp_path and os.path.exists(tmp_path): - try: - os.unlink(tmp_path) - except: - pass + os.unlink(tmp_path) raise - except Exception as e: - lib_logger.error(f"Disk save operation failed: {e}") + lib_logger.error(f"Disk save failed: {e}") - async def _start_background_tasks(self): + async def _start_background_tasks(self) -> None: """Start background writer and cleanup tasks.""" - if not self._enable_disk_persistence or self._running: + if not self._enable_disk or self._running: return self._running = True - - # Start async writer task self._writer_task = asyncio.create_task(self._writer_loop()) - lib_logger.debug(f"Started background writer task (interval: {self._write_interval}s)") - - # Start cleanup task self._cleanup_task = asyncio.create_task(self._cleanup_loop()) - lib_logger.debug(f"Started background cleanup task (interval: {self._cleanup_interval}s)") + lib_logger.debug("Started background cache tasks") - async def _writer_loop(self): + async def _writer_loop(self) -> None: """Background task: periodically flush dirty cache to disk.""" try: while self._running: await asyncio.sleep(self._write_interval) - if self._dirty: try: await self._save_to_disk() @@ -347,1328 +435,868 @@ async def _writer_loop(self): except Exception as e: lib_logger.error(f"Background writer error: {e}") except asyncio.CancelledError: - lib_logger.debug("Background writer task cancelled") - except Exception as e: - lib_logger.error(f"Background writer crashed: {e}") + pass - async def _cleanup_loop(self): + async def _cleanup_loop(self) -> None: """Background task: periodically clean up expired entries.""" try: while self._running: await asyncio.sleep(self._cleanup_interval) - - try: - await self._cleanup_expired() - except Exception as e: - lib_logger.error(f"Background cleanup error: {e}") + await self._cleanup_expired() except asyncio.CancelledError: - lib_logger.debug("Background cleanup task cancelled") - except Exception as e: - lib_logger.error(f"Background cleanup crashed: {e}") + pass - async def _cleanup_expired(self): - """Remove expired entries from memory cache (based on memory TTL).""" + async def _cleanup_expired(self) -> None: + """Remove expired entries from memory cache.""" async with self._lock: now = time.time() - expired = [ - k for k, (_, ts) in self._cache.items() - if now - ts > self._memory_ttl - ] - + expired = [k for k, (_, ts) in self._cache.items() if now - ts > self._memory_ttl] for k in expired: del self._cache[k] - if expired: - self._dirty = True # Mark for disk save - lib_logger.debug(f"Cleaned up {len(expired)} expired signatures from memory") + self._dirty = True + lib_logger.debug(f"Cleaned up {len(expired)} expired entries") - def store(self, tool_call_id: str, signature: str): - """ - Store a signature for a tool call ID (sync wrapper for async storage). - - Args: - tool_call_id: Unique identifier for the tool call - signature: Encrypted thoughtSignature from Antigravity API - """ - # Create task for async storage - asyncio.create_task(self._async_store(tool_call_id, signature)) + def store(self, key: str, value: str) -> None: + """Store a value (sync wrapper for async storage).""" + asyncio.create_task(self._async_store(key, value)) - async def _async_store(self, tool_call_id: str, signature: str): + async def _async_store(self, key: str, value: str) -> None: """Async implementation of store.""" async with self._lock: - self._cache[tool_call_id] = (signature, time.time()) - self._dirty = True # Mark for disk write + self._cache[key] = (value, time.time()) + self._dirty = True - def retrieve(self, tool_call_id: str) -> Optional[str]: - """ - Retrieve signature for a tool call ID (sync method). - - Args: - tool_call_id: Unique identifier for the tool call - - Returns: - The signature if found and not expired, None otherwise - """ - # Try memory cache first (sync access is safe for read) - if tool_call_id in self._cache: - signature, timestamp = self._cache[tool_call_id] + def retrieve(self, key: str) -> Optional[str]: + """Retrieve a value by key (sync method).""" + if key in self._cache: + value, timestamp = self._cache[key] if time.time() - timestamp <= self._memory_ttl: self._stats["memory_hits"] += 1 - return signature + return value else: - # Expired in memory, remove it - del self._cache[tool_call_id] + del self._cache[key] self._dirty = True - # Not in memory - schedule async disk lookup - # For now, return None (disk fallback happens on next request) - # This is intentional to avoid blocking the sync caller self._stats["misses"] += 1 - - # Schedule background disk check (non-blocking) - if self._enable_disk_persistence: - asyncio.create_task(self._check_disk_fallback(tool_call_id)) - + if self._enable_disk: + asyncio.create_task(self._check_disk_fallback(key)) return None - async def _check_disk_fallback(self, tool_call_id: str): - """Check disk for signature and load into memory if found.""" + async def _check_disk_fallback(self, key: str) -> None: + """Check disk for key and load into memory if found.""" try: - # Reload from disk if file exists - if self._cache_file.exists(): - async with self._disk_lock: - with open(self._cache_file, 'r', encoding='utf-8') as f: - data = json.load(f) - - entries = data.get("entries", {}) - if tool_call_id in entries: - entry = entries[tool_call_id] - timestamp = entry.get("timestamp", 0) - - # Check disk TTL (24 hours) - if time.time() - timestamp <= self._disk_ttl: - signature = entry.get("signature", "") - if signature: - # Load into memory cache - async with self._lock: - self._cache[tool_call_id] = (signature, timestamp) - self._stats["disk_hits"] += 1 - lib_logger.debug(f"Loaded signature {tool_call_id} from disk") + if not self._cache_file.exists(): + return + + async with self._disk_lock: + with open(self._cache_file, 'r', encoding='utf-8') as f: + data = json.load(f) + + entries = data.get("entries", {}) + if key in entries: + entry = entries[key] + ts = entry.get("timestamp", 0) + if time.time() - ts <= self._disk_ttl: + sig = entry.get("signature", "") + if sig: + async with self._lock: + self._cache[key] = (sig, ts) + self._stats["disk_hits"] += 1 + lib_logger.debug(f"Loaded {key} from disk") except Exception as e: - lib_logger.debug(f"Disk fallback check failed: {e}") + lib_logger.debug(f"Disk fallback failed: {e}") - async def clear(self): - """Clear all cached signatures (memory and disk).""" + async def clear(self) -> None: + """Clear all cached data.""" async with self._lock: self._cache.clear() self._dirty = True - - if self._enable_disk_persistence: + if self._enable_disk: await self._save_to_disk() - async def shutdown(self): + async def shutdown(self) -> None: """Graceful shutdown: flush pending writes and stop background tasks.""" - lib_logger.info("ThoughtSignatureCache shutting down...") - - # Stop background tasks + lib_logger.info("AntigravityCache shutting down...") self._running = False - if self._writer_task: - self._writer_task.cancel() - try: - await self._writer_task - except asyncio.CancelledError: - pass - - if self._cleanup_task: - self._cleanup_task.cancel() - try: - await self._cleanup_task - except asyncio.CancelledError: - pass + for task in (self._writer_task, self._cleanup_task): + if task: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass - # Flush pending writes - if self._dirty and self._enable_disk_persistence: - lib_logger.info("Flushing pending cache writes...") + if self._dirty and self._enable_disk: await self._save_to_disk() lib_logger.info( - f"ThoughtSignatureCache shutdown complete " - f"(stats: mem_hits={self._stats['memory_hits']}, " - f"disk_hits={self._stats['disk_hits']}, " - f"misses={self._stats['misses']}, " - f"writes={self._stats['writes']})" + f"Cache shutdown complete (stats: mem_hits={self._stats['memory_hits']}, " + f"disk_hits={self._stats['disk_hits']}, misses={self._stats['misses']})" ) +# ============================================================================= +# MAIN PROVIDER CLASS +# ============================================================================= + class AntigravityProvider(AntigravityAuthBase, ProviderInterface): """ - Antigravity provider implementation for Gemini models. - - Antigravity is an experimental internal Google API that provides access to Gemini models - including Gemini 3 with thinking/reasoning capabilities. It wraps standard Gemini API - requests with additional metadata and uses sandbox endpoints. - - Key features: - - Model aliasing (gemini-3-pro-high ↔ gemini-3-pro-preview) - - Gemini 3 thinkingLevel support - - ThoughtSignature preservation for multi-turn conversations - - Reasoning content separation (thought=true parts) - - Sophisticated tool response grouping - - Base URL fallback (sandbox → production) - - Gemini 3 Special Mechanics: - 1. ThinkingLevel: Uses thinkingLevel (low/high) instead of thinkingBudget for Gemini 3 models - 2. ThoughtSignature: Function calls include thoughtSignature="skip_thought_signature_validator" - - This is a CONSTANT validation bypass flag, not a session key - - Preserved across conversation turns to maintain reasoning continuity - - Filtered from responses to prevent exposing encrypted internal data - 3. Reasoning Content: Text parts with thought=true flag are separated into reasoning_content - 4. Token Counting: thoughtsTokenCount is included in prompt_tokens and reported as reasoning_tokens + Antigravity provider for Gemini and Claude models via Google's internal API. + + Supports: + - Gemini 2.5 (Pro/Flash) with thinkingBudget + - Gemini 3 (Pro/Image) with thinkingLevel + - Claude Sonnet 4.5 via Antigravity proxy + + Features: + - Unified streaming/non-streaming handling + - ThoughtSignature caching for multi-turn conversations + - Automatic base URL fallback + - Gemini 3 tool hallucination prevention """ + skip_cost_calculation = True - + def __init__(self): super().__init__() self.model_definitions = ModelDefinitions() - self._current_base_url = BASE_URLS[0] # Start with daily sandbox + + # Base URL management self._base_url_index = 0 + self._current_base_url = BASE_URLS[0] - # Initialize caches for conversation state preservation - memory_ttl = int(os.getenv("ANTIGRAVITY_SIGNATURE_CACHE_TTL", "3600")) - disk_ttl = int(os.getenv("ANTIGRAVITY_SIGNATURE_DISK_TTL", "86400")) + # Configuration from environment + memory_ttl = _env_int("ANTIGRAVITY_SIGNATURE_CACHE_TTL", 3600) + disk_ttl = _env_int("ANTIGRAVITY_SIGNATURE_DISK_TTL", 86400) - # Cache for Gemini 3 thoughtSignatures + # Initialize caches self._signature_cache = AntigravityCache( - cache_file=GEMINI3_SIGNATURE_CACHE_FILE, - memory_ttl_seconds=memory_ttl, - disk_ttl_seconds=disk_ttl + GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl ) - - # Cache for Claude thinking content self._thinking_cache = AntigravityCache( - cache_file=CLAUDE_THINKING_CACHE_FILE, - memory_ttl_seconds=memory_ttl, - disk_ttl_seconds=disk_ttl + CLAUDE_THINKING_CACHE_FILE, memory_ttl, disk_ttl ) - # Check if client passthrough is enabled (default: TRUE for testing) - self._preserve_signatures_in_client = os.getenv( - "ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES", - "true" # Default ON for testing - ).lower() in ("true", "1", "yes") - - # Check if server-side cache is enabled (default: TRUE for testing) - self._enable_signature_cache = os.getenv( - "ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", - "true" # Default ON for testing - ).lower() in ("true", "1", "yes") - - # Check if dynamic model discovery is enabled (default: OFF due to endpoint instability) - self._enable_dynamic_model_discovery = os.getenv( - "ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", - "false" # Default OFF - use hardcoded list - ).lower() in ("true", "1", "yes") - - if self._preserve_signatures_in_client: - lib_logger.debug("Antigravity: thoughtSignature client passthrough ENABLED") - else: - lib_logger.debug("Antigravity: thoughtSignature client passthrough DISABLED") - - if self._enable_signature_cache: - lib_logger.debug(f"Antigravity: thoughtSignature server-side cache ENABLED (memory_ttl={memory_ttl}s, disk_ttl={disk_ttl}s)") - else: - lib_logger.debug("Antigravity: thoughtSignature server-side cache DISABLED") - - if self._enable_dynamic_model_discovery: - lib_logger.debug("Antigravity: Dynamic model discovery ENABLED (may fail if endpoint unavailable)") - else: - lib_logger.debug("Antigravity: Dynamic model discovery DISABLED (using hardcoded model list)") - - # Check if Gemini 3 tool fix is enabled (default: ON for testing) - # This applies the "Quad-Lock" catch-all strategy to prevent tool hallucination - self._enable_gemini3_tool_fix = os.getenv( - "ANTIGRAVITY_GEMINI3_TOOL_FIX", - "true" # Default ON - applies namespace + signature injection - ).lower() in ("true", "1", "yes") - - # Gemini 3 fix configuration - customize the fix components - # Namespace prefix for tool names (Strategy 1) - self._gemini3_tool_prefix = os.getenv( - "ANTIGRAVITY_GEMINI3_TOOL_PREFIX", - "gemini3_" # Default prefix - ) + # Feature flags + self._preserve_signatures_in_client = _env_bool("ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES", True) + self._enable_signature_cache = _env_bool("ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", True) + self._enable_dynamic_models = _env_bool("ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", False) + self._enable_gemini3_tool_fix = _env_bool("ANTIGRAVITY_GEMINI3_TOOL_FIX", True) - # Description prompt format (Strategy 2) - # Use {params} as placeholder for parameter list + # Gemini 3 tool fix configuration + self._gemini3_tool_prefix = os.getenv("ANTIGRAVITY_GEMINI3_TOOL_PREFIX", "gemini3_") self._gemini3_description_prompt = os.getenv( "ANTIGRAVITY_GEMINI3_DESCRIPTION_PROMPT", - "\n\nSTRICT PARAMETERS: {params}." # Default format + "\n\nSTRICT PARAMETERS: {params}." ) - - # System instruction text (Strategy 3) - # Set to empty string to disable system instruction injection self._gemini3_system_instruction = os.getenv( "ANTIGRAVITY_GEMINI3_SYSTEM_INSTRUCTION", - # Default: comprehensive tool usage instructions - """CRITICAL TOOL USAGE INSTRUCTIONS: -You are operating in a custom environment where tool definitions differ from your training data. -You MUST follow these rules strictly: - -1. DO NOT use your internal training data to guess tool parameters -2. ONLY use the exact parameter structure defined in the tool schema -3. If a tool takes a 'files' parameter, it is ALWAYS an array of objects with specific properties, NEVER a simple array of strings -4. If a tool edits code, it takes structured JSON objects with specific fields, NEVER raw diff strings or plain text -5. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') -6. Array parameters have specific item types - check the schema's 'items' field for the exact structure -7. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions - -If you are unsure about a tool's parameters, YOU MUST read the schema definition carefully. Your training data about common tool names like 'read_file' or 'apply_diff' does NOT apply here. -""" + DEFAULT_GEMINI3_SYSTEM_INSTRUCTION ) - if self._enable_gemini3_tool_fix: - lib_logger.debug(f"Antigravity: Gemini 3 tool fix ENABLED") - lib_logger.debug(f" - Namespace prefix: '{self._gemini3_tool_prefix}'") - lib_logger.debug(f" - Description prompt: '{self._gemini3_description_prompt[:50]}...'") - lib_logger.debug(f" - System instruction: {'ENABLED' if self._gemini3_system_instruction else 'DISABLED'} ({len(self._gemini3_system_instruction)} chars)") - else: - lib_logger.debug("Antigravity: Gemini 3 tool fix DISABLED (using default tool schemas)") - - - def _generate_thinking_cache_key(self, text_content: str, tool_calls: List[Dict]) -> Optional[str]: + # Log configuration + self._log_config() + + def _log_config(self) -> None: + """Log provider configuration.""" + lib_logger.debug( + f"Antigravity config: signatures_in_client={self._preserve_signatures_in_client}, " + f"cache={self._enable_signature_cache}, dynamic_models={self._enable_dynamic_models}, " + f"gemini3_fix={self._enable_gemini3_tool_fix}" + ) + + # ========================================================================= + # MODEL UTILITIES + # ========================================================================= + + def _alias_to_internal(self, alias: str) -> str: + """Convert public alias to internal model name.""" + return MODEL_ALIAS_REVERSE.get(alias, alias) + + def _internal_to_alias(self, internal: str) -> str: + """Convert internal model name to public alias.""" + if internal in EXCLUDED_MODELS: + return "" + return MODEL_ALIAS_MAP.get(internal, internal) + + def _is_gemini_3(self, model: str) -> bool: + """Check if model is Gemini 3 (requires special handling).""" + internal = self._alias_to_internal(model) + return internal.startswith("gemini-3-") or model.startswith("gemini-3-") + + def _is_claude(self, model: str) -> bool: + """Check if model is Claude.""" + return "claude" in model.lower() + + def _strip_provider_prefix(self, model: str) -> str: + """Strip provider prefix from model name.""" + return model.split("/")[-1] if "/" in model else model + + # ========================================================================= + # BASE URL MANAGEMENT + # ========================================================================= + + def _get_base_url(self) -> str: + """Get current base URL.""" + return self._current_base_url + + def _try_next_base_url(self) -> bool: + """Switch to next base URL in fallback list. Returns True if successful.""" + if self._base_url_index < len(BASE_URLS) - 1: + self._base_url_index += 1 + self._current_base_url = BASE_URLS[self._base_url_index] + lib_logger.info(f"Switching to fallback URL: {self._current_base_url}") + return True + return False + + def _reset_base_url(self) -> None: + """Reset to primary base URL.""" + self._base_url_index = 0 + self._current_base_url = BASE_URLS[0] + + # ========================================================================= + # THINKING CACHE KEY GENERATION + # ========================================================================= + + def _generate_thinking_cache_key( + self, + text_content: str, + tool_calls: List[Dict] + ) -> Optional[str]: """ Generate stable cache key from response content for Claude thinking preservation. - Uses composite key strategy: - - If tool calls exist: Use first tool call ID (most reliable) - - If text exists: Use text hash - - If both: Combine both for maximum uniqueness - - Args: - text_content: Regular text from response - tool_calls: List of tool calls with IDs - - Returns: - Cache key string, or None if no cacheable content + Uses composite key: + - Tool call IDs (most stable) + - Text hash (for text-only responses) """ - import hashlib key_parts = [] - # Priority 1: Tool call IDs (most stable - we generate these) - if tool_calls and len(tool_calls) > 0: - first_tool_id = tool_calls[0].get("id", "") - if first_tool_id: - # Remove 'call_' prefix if present for shorter key - tool_id_short = first_tool_id.replace("call_", "") - key_parts.append(f"tool_{tool_id_short}") + if tool_calls: + first_id = tool_calls[0].get("id", "") + if first_id: + key_parts.append(f"tool_{first_id.replace('call_', '')}") - # Priority 2: Text hash (for text-only or mixed responses) if text_content: - # Use first 200 chars for stability (longer text may vary slightly) text_hash = hashlib.md5(text_content[:200].encode()).hexdigest()[:16] key_parts.append(f"text_{text_hash}") - # Combine parts - if key_parts: - return "thinking_" + "_".join(key_parts) - - # Shouldn't happen - responses always have text or tools - return None - - - # ============================================================================ - # MODEL ALIAS SYSTEM - # ============================================================================ - - def _model_name_to_alias(self, model_name: str) -> str: + return "thinking_" + "_".join(key_parts) if key_parts else None + + # ========================================================================= + # REASONING CONFIGURATION + # ========================================================================= + + def _get_thinking_config( + self, + reasoning_effort: Optional[str], + model: str, + custom_budget: bool = False + ) -> Optional[Dict[str, Any]]: """ - Convert internal Antigravity model names to public aliases. + Map reasoning_effort to thinking configuration. - Args: - model_name: Internal model name - - Returns: - Public alias name, or empty string if model should be excluded + - Gemini 2.5 & Claude: thinkingBudget (integer tokens) + - Gemini 3: thinkingLevel (string: "low"/"high") """ - alias_map = { - "rev19-uic3-1p": "gemini-2.5-computer-use-preview-10-2025", - "gemini-3-pro-image": "gemini-3-pro-image-preview", - "gemini-3-pro-high": "gemini-3-pro-preview", - # Claude models: no aliasing needed (public name = internal name) - } + internal = self._alias_to_internal(model) + is_gemini_25 = "gemini-2.5" in model + is_gemini_3 = internal.startswith("gemini-3-") + is_claude = self._is_claude(model) - # Filter out excluded models (return empty string to skip) - excluded = [ - "chat_20706", "chat_23310", "gemini-2.5-flash-thinking", - "gemini-3-pro-low", "gemini-2.5-pro" - ] - if model_name in excluded: - return "" + if not (is_gemini_25 or is_gemini_3 or is_claude): + return None - return alias_map.get(model_name, model_name) - - def _alias_to_model_name(self, alias: str) -> str: - """ - Convert public aliases to internal Antigravity model names. + # Gemini 3: String-based thinkingLevel + if is_gemini_3: + if reasoning_effort == "low": + return {"thinkingLevel": "low", "include_thoughts": True} + return {"thinkingLevel": "high", "include_thoughts": True} - Args: - alias: Public alias name - - Returns: - Internal model name - """ - reverse_map = { - "gemini-2.5-computer-use-preview-10-2025": "rev19-uic3-1p", - "gemini-3-pro-image-preview": "gemini-3-pro-image", - "gemini-3-pro-preview": "gemini-3-pro-high", - # Claude models: no aliasing needed (public name = internal name) - } - return reverse_map.get(alias, alias) - - def _is_gemini_3_model(self, model: str) -> bool: - """ - Check if model is Gemini 3 (requires thoughtSignature preservation). + # Gemini 2.5 & Claude: Integer thinkingBudget + if not reasoning_effort: + return {"thinkingBudget": -1, "include_thoughts": True} # Auto - Args: - model: Model name (public alias) - - Returns: - True if this is a Gemini 3 model - """ - internal_model = self._alias_to_model_name(model) - return internal_model.startswith("gemini-3-") or model.startswith("gemini-3-") - - @staticmethod - def _normalize_type_arrays(schema: Any) -> Any: - """ - Normalize type arrays in JSON Schema for Proto-based Antigravity API. - Converts `"type": ["string", "null"]` → `"type": "string"`. - """ - if isinstance(schema, dict): - normalized = {} - for key, value in schema.items(): - if key == "type" and isinstance(value, list): - # Take first non-null type - non_null_types = [t for t in value if t != "null"] - normalized[key] = non_null_types[0] if non_null_types else value[0] - else: - normalized[key] = AntigravityProvider._normalize_type_arrays(value) - return normalized - elif isinstance(schema, list): - return [AntigravityProvider._normalize_type_arrays(item) for item in schema] + if reasoning_effort == "disable": + return {"thinkingBudget": 0, "include_thoughts": False} + + # Model-specific budgets + if "gemini-2.5-pro" in model or is_claude: + budgets = {"low": 8192, "medium": 16384, "high": 32768} + elif "gemini-2.5-flash" in model: + budgets = {"low": 6144, "medium": 12288, "high": 24576} else: - return schema - - # ============================================================================ - # RANDOM ID GENERATION - # ============================================================================ - - @staticmethod - def generate_request_id() -> str: - """Generate Antigravity request ID: agent-{uuid}""" - return f"agent-{uuid.uuid4()}" - - @staticmethod - def generate_session_id() -> str: - """Generate Antigravity session ID: -{random_number}""" - # Generate random 19-digit number - n = random.randint(1_000_000_000_000_000_000, 9_999_999_999_999_999_999) - return f"-{n}" - - @staticmethod - def generate_project_id() -> str: - """Generate fake project ID: {adj}-{noun}-{random}""" - adjectives = ["useful", "bright", "swift", "calm", "bold"] - nouns = ["fuze", "wave", "spark", "flow", "core"] - adj = random.choice(adjectives) - noun = random.choice(nouns) - random_part = str(uuid.uuid4())[:5].lower() - return f"{adj}-{noun}-{random_part}" - - # ============================================================================ - # MESSAGE TRANSFORMATION (OpenAI → Gemini CLI format) - # ============================================================================ - - def _transform_messages(self, messages: List[Dict[str, Any]], model: str) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: - """ - Transform OpenAI messages to Gemini CLI format. - Handles thoughtSignature preservation with 3-tier fallback (GEMINI 3 ONLY): - 1. Use client-provided signature (if present) - 2. Fall back to server-side cache - 3. Use bypass constant as last resort + budgets = {"low": 1024, "medium": 2048, "high": 4096} - Args: - messages: List of OpenAI-formatted messages - model: Model name for Gemini 3 detection - - Returns: - Tuple of (system_instruction, gemini_contents) - """ - system_instruction = None - gemini_contents = [] + budget = budgets.get(reasoning_effort, -1) + if not custom_budget: + budget = budget // 4 # Default to 25% of max output tokens - # Make a copy to avoid modifying original + return {"thinkingBudget": budget, "include_thoughts": True} + + # ========================================================================= + # MESSAGE TRANSFORMATION (OpenAI → Gemini) + # ========================================================================= + + def _transform_messages( + self, + messages: List[Dict[str, Any]], + model: str + ) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: + """ + Transform OpenAI messages to Gemini CLI format. + + Handles: + - System instruction extraction + - Multi-part content (text, images) + - Tool calls and responses + - Claude thinking injection from cache + - Gemini 3 thoughtSignature preservation + """ messages = copy.deepcopy(messages) + system_instruction = None + gemini_contents = [] - # Separate system prompt from other messages + # Extract system prompt if messages and messages[0].get('role') == 'system': - system_prompt_content = messages.pop(0).get('content', '') - if system_prompt_content: - # Handle both string and list-based system content - system_parts = [] - if isinstance(system_prompt_content, str): - system_parts.append({"text": system_prompt_content}) - elif isinstance(system_prompt_content, list): - # Multi-part system content (strip cache_control) - for item in system_prompt_content: - if item.get("type") == "text": - text = item.get("text", "") - if text: - # Skip cache_control - Claude-specific field - system_parts.append({"text": text}) - + system_content = messages.pop(0).get('content', '') + if system_content: + system_parts = self._parse_content_parts(system_content, _strip_cache_control=True) if system_parts: - system_instruction = { - "role": "user", - "parts": system_parts - } - - - # Build tool call ID to name mapping - tool_call_id_to_name = {} + system_instruction = {"role": "user", "parts": system_parts} + + # Build tool_call_id → name mapping + tool_id_to_name = {} for msg in messages: if msg.get("role") == "assistant" and msg.get("tool_calls"): - for tool_call in msg["tool_calls"]: - if tool_call.get("type") == "function": - tool_call_id_to_name[tool_call["id"]] = tool_call["function"]["name"] - - #Convert each message + for tc in msg["tool_calls"]: + if tc.get("type") == "function": + tool_id_to_name[tc["id"]] = tc["function"]["name"] + + # Convert each message for msg in messages: role = msg.get("role") content = msg.get("content") parts = [] - gemini_role = "model" if role == "assistant" else "tool" if role == "tool" else "user" - + if role == "user": - if isinstance(content, str): - # Simple text content - if content: - parts.append({"text": content}) - elif isinstance(content, list): - # Multi-part content (text, images, etc.) - for item in content: - if item.get("type") == "text": - text = item.get("text", "") - if text: - # Strip Claude-specific cache_control field - # This field causes 400 errors with Antigravity - parts.append({"text": text}) - elif item.get("type") == "image_url": - # Handle image data URLs - image_url = item.get("image_url", {}).get("url", "") - if image_url.startswith("data:"): - try: - # Parse: data:image/png;base64,iVBORw0KG... - header, data = image_url.split(",", 1) - mime_type = header.split(":")[1].split(";")[0] - parts.append({ - "inlineData": { - "mimeType": mime_type, - "data": data - } - }) - except Exception as e: - lib_logger.warning(f"Failed to parse image data URL: {e}") - + parts = self._transform_user_message(content) elif role == "assistant": - # Try to retrieve cached thinking for Claude models - thinking_to_inject = None - cache_key = None - - if model.startswith("claude-") and self._enable_signature_cache: - # Build cache key from incoming message - msg_text = content if isinstance(content, str) else "" - msg_tools = msg.get("tool_calls", []) - - cache_key = self._generate_thinking_cache_key(msg_text, msg_tools) - - if cache_key: - cached_json = self._thinking_cache.retrieve(cache_key) - if cached_json: - try: - thinking_to_inject = json.loads(cached_json) - lib_logger.debug(f"✓ Retrieved thinking from cache: {cache_key[:50]}...") - except json.JSONDecodeError: - lib_logger.warning(f"Failed to parse cached thinking for: {cache_key}") - - # Inject thinking FIRST if we have it - if thinking_to_inject: - thinking_text = thinking_to_inject.get("thinking_text", "") - thought_sig = thinking_to_inject.get("thought_signature", "") - - if thinking_text: - thinking_part = { - "text": thinking_text, - "thought": True - } - - # Add signature if available, otherwise use skip validator - if thought_sig: - thinking_part["thoughtSignature"] = thought_sig - else: - thinking_part["thoughtSignature"] = "skip_thought_signature_validator" - lib_logger.debug("Using skip validator for missing signature") - - parts.append(thinking_part) - lib_logger.debug( - f"✅ Injected {len(thinking_text)} chars of thinking " - f"(sig={'yes' if thought_sig else 'fallback'})" - ) - - # Then add regular content - if isinstance(content, str) and content: - parts.append({"text": content}) - if msg.get("tool_calls"): - for tool_call in msg["tool_calls"]: - if tool_call.get("type") == "function": - try: - args_dict = json.loads(tool_call["function"]["arguments"]) - except (json.JSONDecodeError, TypeError): - args_dict = {} - - tool_call_id = tool_call.get("id", "") - - # Get function name and add configured prefix if needed (Gemini 3 specific) - function_name = tool_call["function"]["name"] - if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: - # Client sends original names, we need to prefix for API consistency - function_name = f"{self._gemini3_tool_prefix}{function_name}" - - func_call_part = { - "functionCall": { - "name": function_name, - "args": args_dict, - "id": tool_call_id # ← ADD THIS LINE - Antigravity needs it for Claude! - } - } - - # thoughtSignature handling (GEMINI 3 ONLY) - # Claude and other models don't support this field! - if self._is_gemini_3_model(model): - # PRIORITY 1: Use client-provided signature if available - client_signature = tool_call.get("thought_signature") - - # PRIORITY 2: Fall back to server-side cache - if not client_signature and tool_call_id and self._enable_signature_cache: - client_signature = self._signature_cache.retrieve(tool_call_id) - if client_signature: - lib_logger.debug(f"Retrieved thoughtSignature from cache for {tool_call_id}") - - # PRIORITY 3: Use bypass constant as last resort - if client_signature: - func_call_part["thoughtSignature"] = client_signature - else: - func_call_part["thoughtSignature"] = "skip_thought_signature_validator" - # WARNING: Missing signature for Gemini 3 - lib_logger.warning( - f"Gemini 3 tool call '{tool_call_id}' missing thoughtSignature. " - f"Client didn't provide it and cache lookup failed. " - f"Using bypass - reasoning quality may degrade." - ) - - parts.append(func_call_part) - + parts = self._transform_assistant_message(msg, model, tool_id_to_name) elif role == "tool": - # Tool responses grouped by function name - tool_call_id = msg.get("tool_call_id", "") - function_name = tool_call_id_to_name.get(tool_call_id, "unknown_function") - tool_content = msg.get("content", "{}") - - # Add configured prefix to function response name if needed (Gemini 3 specific) - if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: - # Client sends responses for original names, we need to prefix for API consistency - function_name = f"{self._gemini3_tool_prefix}{function_name}" - - # Parse tool content - if it's JSON, use parsed value; otherwise use as-is - try: - parsed_content = json.loads(tool_content) - except (json.JSONDecodeError, TypeError): - parsed_content = tool_content - - parts.append({ - "functionResponse": { - "name": function_name, - "response": { - "result": parsed_content - }, - "id": tool_call_id # ← ADD THIS LINE - Antigravity needs it for Claude! - } - }) - + parts = self._transform_tool_message(msg, model, tool_id_to_name) + if parts: - gemini_contents.append({ - "role": gemini_role, - "parts": parts - }) + gemini_role = "model" if role == "assistant" else "user" if role == "tool" else "user" + gemini_contents.append({"role": gemini_role, "parts": parts}) return system_instruction, gemini_contents - - # ============================================================================ - # REASONING CONFIGURATION (GEMINI 2.5 & 3 ONLY) - # ============================================================================ - - def _map_reasoning_effort_to_thinking_config( + + def _parse_content_parts( self, - reasoning_effort: Optional[str], - model: str, - custom_reasoning_budget: bool = False - ) -> Optional[Dict[str, Any]]: - """ - Map reasoning_effort to thinking configuration for Gemini 2.5, Gemini 3, and Claude models. + content: Any, + _strip_cache_control: bool = False + ) -> List[Dict[str, Any]]: + """Parse content into Gemini parts format.""" + parts = [] + + if isinstance(content, str): + if content: + parts.append({"text": content}) + elif isinstance(content, list): + for item in content: + if item.get("type") == "text": + text = item.get("text", "") + if text: + parts.append({"text": text}) + elif item.get("type") == "image_url": + image_part = self._parse_image_url(item.get("image_url", {})) + if image_part: + parts.append(image_part) + + return parts + + def _parse_image_url(self, image_url: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Parse image URL into Gemini inlineData format.""" + url = image_url.get("url", "") + if not url.startswith("data:"): + return None - Supports thinking/reasoning via Antigravity for: - - Gemini 2.5: thinkingBudget (integer tokens, based on Gemini CLI logic) - - Gemini 3: thinkingLevel (string: "low" or "high") - - Claude: thinkingBudget (same as Gemini 2.5, proxied by Antigravity backend) + try: + header, data = url.split(",", 1) + mime_type = header.split(":")[1].split(";")[0] + return {"inlineData": {"mimeType": mime_type, "data": data}} + except Exception as e: + lib_logger.warning(f"Failed to parse image URL: {e}") + return None + + def _transform_user_message(self, content: Any) -> List[Dict[str, Any]]: + """Transform user message content to Gemini parts.""" + return self._parse_content_parts(content) + + def _transform_assistant_message( + self, + msg: Dict[str, Any], + model: str, + _tool_id_to_name: Dict[str, str] + ) -> List[Dict[str, Any]]: + """Transform assistant message including tool calls and thinking injection.""" + parts = [] + content = msg.get("content") + tool_calls = msg.get("tool_calls", []) + + # Try to inject cached thinking for Claude + if self._is_claude(model) and self._enable_signature_cache: + thinking_parts = self._get_cached_thinking(content, tool_calls) + parts.extend(thinking_parts) + + # Add regular content + if isinstance(content, str) and content: + parts.append({"text": content}) + + # Add tool calls + for tc in tool_calls: + if tc.get("type") != "function": + continue + + try: + args = json.loads(tc["function"]["arguments"]) + except (json.JSONDecodeError, TypeError): + args = {} + + tool_id = tc.get("id", "") + func_name = tc["function"]["name"] + + # Add prefix for Gemini 3 + if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: + func_name = f"{self._gemini3_tool_prefix}{func_name}" + + func_part = { + "functionCall": { + "name": func_name, + "args": args, + "id": tool_id + } + } + + # Add thoughtSignature for Gemini 3 + if self._is_gemini_3(model): + sig = tc.get("thought_signature") + if not sig and tool_id and self._enable_signature_cache: + sig = self._signature_cache.retrieve(tool_id) + + if sig: + func_part["thoughtSignature"] = sig + else: + func_part["thoughtSignature"] = "skip_thought_signature_validator" + lib_logger.warning(f"Missing thoughtSignature for {tool_id}, using bypass") + + parts.append(func_part) - Default behavior (no reasoning_effort): - - Gemini 2.5 & Claude: thinkingBudget=-1 (auto mode) - - Gemini 3: thinkingLevel="high" (always enabled at high level) + return parts + + def _get_cached_thinking( + self, + content: Any, + tool_calls: List[Dict] + ) -> List[Dict[str, Any]]: + """Retrieve and format cached thinking content for Claude.""" + parts = [] + msg_text = content if isinstance(content, str) else "" + cache_key = self._generate_thinking_cache_key(msg_text, tool_calls) - Args: - reasoning_effort: Effort level ('low', 'medium', 'high', 'disable', or None) - model: Model name (public alias) - custom_reasoning_budget: If True, use full budgets; if False, divide by 4 - - Returns: - Dict with thinkingConfig or None if model doesn't support thinking - """ - internal_model = self._alias_to_model_name(model) + if not cache_key: + return parts - # Detect model family - is_gemini_25 = "gemini-2.5" in model - is_gemini_3 = internal_model.startswith("gemini-3-") - is_claude = "claude" in model.lower() + cached_json = self._thinking_cache.retrieve(cache_key) + if not cached_json: + return parts - # Only Gemini 2.5, Gemini 3, and Claude support thinking via Antigravity - if not is_gemini_25 and not is_gemini_3 and not is_claude: - return None + try: + thinking_data = json.loads(cached_json) + thinking_text = thinking_data.get("thinking_text", "") + sig = thinking_data.get("thought_signature", "") + + if thinking_text: + thinking_part = { + "text": thinking_text, + "thought": True, + "thoughtSignature": sig or "skip_thought_signature_validator" + } + parts.append(thinking_part) + lib_logger.debug(f"Injected {len(thinking_text)} chars of thinking") + except json.JSONDecodeError: + lib_logger.warning(f"Failed to parse cached thinking: {cache_key}") - # ======================================================================== - # GEMINI 2.5 & CLAUDE: Use thinkingBudget (INTEGER) - # ======================================================================== - if is_gemini_25 or is_claude: - # Default: auto mode - if not reasoning_effort: - return {"thinkingBudget": -1, "include_thoughts": True} - - # Disable thinking - if reasoning_effort == "disable": - return {"thinkingBudget": 0, "include_thoughts": False} - - # Model-specific budgets - # Claude uses Gemini 2.5 pro budgets (high-quality thinking) - if "gemini-2.5-pro" in model or is_claude: - budgets = {"low": 8192, "medium": 16384, "high": 32768} - elif "gemini-2.5-flash" in model: - budgets = {"low": 6144, "medium": 12288, "high": 24576} - else: - # Fallback for other gemini-2.5 models - budgets = {"low": 1024, "medium": 2048, "high": 4096} - - budget = budgets.get(reasoning_effort, -1) # -1 for invalid/auto - - # Apply custom_reasoning_budget toggle - # If False (default), divide by 4 like Gemini CLI - if not custom_reasoning_budget: - budget = budget // 6 - - return {"thinkingBudget": budget, "include_thoughts": True} + return parts + + def _transform_tool_message( + self, + msg: Dict[str, Any], + model: str, + tool_id_to_name: Dict[str, str] + ) -> List[Dict[str, Any]]: + """Transform tool response message.""" + tool_id = msg.get("tool_call_id", "") + func_name = tool_id_to_name.get(tool_id, "unknown_function") + content = msg.get("content", "{}") - # ======================================================================== - # GEMINI 3: Use STRING thinkingLevel ("low" or "high") - # ======================================================================== - if is_gemini_3: - # Default: Always use "high" if not specified - # Gemini 3 cannot be disabled - always has thinking enabled - if not reasoning_effort: - return {"thinkingLevel": "high", "include_thoughts": True} - - # Map reasoning effort to string level - # Note: "disable" is ignored - Gemini 3 cannot disable thinking - if reasoning_effort == "low": - level = "low" - # Medium level not yet available - map to high - # When medium is released, uncomment the following line: - # elif reasoning_effort == "medium": - # level = "medium" - else: - # "medium", "high", "disable", or any invalid value → "high" - level = "high" - - return {"thinkingLevel": level, "include_thoughts": True} + # Add prefix for Gemini 3 + if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: + func_name = f"{self._gemini3_tool_prefix}{func_name}" - return None - - # ============================================================================ + try: + parsed_content = json.loads(content) + except (json.JSONDecodeError, TypeError): + parsed_content = content + + return [{ + "functionResponse": { + "name": func_name, + "response": {"result": parsed_content}, + "id": tool_id + } + }] + + # ========================================================================= # TOOL RESPONSE GROUPING - # ============================================================================ - - def _fix_tool_response_grouping(self, contents: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + # ========================================================================= + + def _fix_tool_response_grouping( + self, + contents: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: """ Group function calls with their responses for Antigravity compatibility. - Converts linear format (function call, response, function call, response) - to grouped format (model with calls, function role with all responses). - - Args: - contents: List of Gemini content objects - - Returns: - List of grouped content objects + Converts linear format (call, response, call, response) + to grouped format (model with calls, user with all responses). """ new_contents = [] - pending_groups = [] # Groups awaiting responses - collected_responses = [] # Standalone responses to match + pending_groups = [] + collected_responses = [] for content in contents: role = content.get("role") parts = content.get("parts", []) - # Check if this content has function responses response_parts = [p for p in parts if "functionResponse" in p] if response_parts: - # Collect responses collected_responses.extend(response_parts) # Try to satisfy pending groups for i in range(len(pending_groups) - 1, -1, -1): group = pending_groups[i] - if len(collected_responses) >= group["responses_needed"]: - # Take needed responses - group_responses = collected_responses[:group["responses_needed"]] - collected_responses = collected_responses[group["responses_needed"]:] - - # Create merged function response content - function_response_content = { - "parts": group_responses, - "role": "user" - } - new_contents.append(function_response_content) - - # Remove satisfied group + if len(collected_responses) >= group["count"]: + group_responses = collected_responses[:group["count"]] + collected_responses = collected_responses[group["count"]:] + new_contents.append({"parts": group_responses, "role": "user"}) pending_groups.pop(i) break - - continue # Skip adding this content + continue - # If this is model content with function calls, create a group if role == "model": - function_calls = [p for p in parts if "functionCall" in p] - - if function_calls: - # Add model content first - new_contents.append(content) - - # Create pending group - pending_groups.append({ - "model_content": content, - "function_calls": function_calls, - "responses_needed": len(function_calls) - }) - else: - # Regular model content without function calls - new_contents.append(content) + func_calls = [p for p in parts if "functionCall" in p] + new_contents.append(content) + if func_calls: + pending_groups.append({"count": len(func_calls)}) else: - # Non-model content (user, etc.) new_contents.append(content) - # Handle remaining pending groups + # Handle remaining groups for group in pending_groups: - if len(collected_responses) >= group["responses_needed"]: - group_responses = collected_responses[:group["responses_needed"]] - collected_responses = collected_responses[group["responses_needed"]:] - - function_response_content = { - "parts": group_responses, - "role": "user" - } - new_contents.append(function_response_content) + if len(collected_responses) >= group["count"]: + group_responses = collected_responses[:group["count"]] + collected_responses = collected_responses[group["count"]:] + new_contents.append({"parts": group_responses, "role": "user"}) return new_contents - - # ============================================================================ - # GEMINI 3 TOOL TRANSFORMATION (Catch-All Fix for Hallucination) - # ============================================================================ - - def _apply_gemini3_namespace_to_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ - Apply namespace prefix to all tool names for Gemini 3 (Strategy 1: Namespace). - - This breaks the model's association with training data by prepending 'gemini3_' - to every tool name, forcing it to read the schema definition instead of using - its internal knowledge. - - Args: - tools: List of tool definitions (Gemini format with functionDeclarations) - - Returns: - Modified tools with prefixed names - """ + + # ========================================================================= + # GEMINI 3 TOOL TRANSFORMATIONS + # ========================================================================= + + def _apply_gemini3_namespace( + self, + tools: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Add namespace prefix to tool names for Gemini 3.""" if not tools: return tools - - modified_tools = copy.deepcopy(tools) - - for tool in modified_tools: - function_declarations = tool.get("functionDeclarations", []) - for func_decl in function_declarations: - # Prepend namespace to tool name - original_name = func_decl.get("name", "") - if original_name: - func_decl["name"] = f"{self._gemini3_tool_prefix}{original_name}" - #lib_logger.debug(f"Gemini 3 namespace: {original_name} -> {self._gemini3_tool_prefix}{original_name}") - - return modified_tools - - def _inject_signature_into_tool_descriptions(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ - Inject parameter signatures into tool descriptions for Gemini 3 (Strategy 2: Signature Injection). - This strategy appends the expected parameter structure into the description text, - creating a natural language enforcement of the schema that models pay close attention to. + modified = copy.deepcopy(tools) + for tool in modified: + for func_decl in tool.get("functionDeclarations", []): + name = func_decl.get("name", "") + if name: + func_decl["name"] = f"{self._gemini3_tool_prefix}{name}" - Args: - tools: List of tool definitions (Gemini format with functionDeclarations) - - Returns: - Modified tools with enriched descriptions - """ + return modified + + def _inject_signature_into_descriptions( + self, + tools: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Inject parameter signatures into tool descriptions for Gemini 3.""" if not tools: return tools - - modified_tools = copy.deepcopy(tools) - for tool in modified_tools: - function_declarations = tool.get("functionDeclarations", []) - for func_decl in function_declarations: - # Get parameter schema + modified = copy.deepcopy(tools) + for tool in modified: + for func_decl in tool.get("functionDeclarations", []): schema = func_decl.get("parametersJsonSchema", {}) - if not schema or not isinstance(schema, dict): + if not schema: continue - # Extract required parameters - required_params = schema.get("required", []) + required = schema.get("required", []) properties = schema.get("properties", {}) if not properties: continue - # Build parameter list with type hints param_list = [] for prop_name, prop_data in properties.items(): if not isinstance(prop_data, dict): continue - - type_hint = prop_data.get("type", "unknown") - - # Handle arrays specially (critical for read_file/apply_diff issues) - if type_hint == "array": - items_schema = prop_data.get("items", {}) - if isinstance(items_schema, dict): - item_type = items_schema.get("type", "unknown") - - # Check if it's an array of objects - RECURSE into nested properties - if item_type == "object": - # Extract nested properties for explicit visibility - nested_props = items_schema.get("properties", {}) - nested_required = items_schema.get("required", []) - - if nested_props: - # Build nested property list with types - nested_list = [] - for nested_name, nested_data in nested_props.items(): - if not isinstance(nested_data, dict): - continue - nested_type = nested_data.get("type", "unknown") - - # Mark nested required fields - if nested_name in nested_required: - nested_list.append(f"{nested_name}: {nested_type} REQUIRED") - else: - nested_list.append(f"{nested_name}: {nested_type}") - - # Format as ARRAY_OF_OBJECTS[key1: type1, key2: type2] - nested_str = ", ".join(nested_list) - type_hint = f"ARRAY_OF_OBJECTS[{nested_str}]" - else: - # No properties defined - just generic objects - type_hint = "ARRAY_OF_OBJECTS" - else: - type_hint = f"ARRAY_OF_{item_type.upper()}" - else: - type_hint = "ARRAY" - # Mark required parameters - if prop_name in required_params: - param_list.append(f"{prop_name} ({type_hint}, REQUIRED)") - else: - param_list.append(f"{prop_name} ({type_hint})") - - # Create strict signature string using configurable template - # Replace {params} placeholder with actual parameter list - signature_str = self._gemini3_description_prompt.replace("{params}", ", ".join(param_list)) - - # Inject into description - description = func_decl.get("description", "") - func_decl["description"] = description + signature_str + type_hint = self._format_type_hint(prop_data) + is_required = prop_name in required + param_list.append( + f"{prop_name} ({type_hint}{', REQUIRED' if is_required else ''})" + ) - #lib_logger.debug(f"Gemini 3 signature injection: {func_decl.get('name', '')} - {len(param_list)} params") - - return modified_tools - - def _strip_gemini3_namespace_from_name(self, tool_name: str) -> str: - """ - Strip the configured namespace prefix from a tool name. + if param_list: + sig_str = self._gemini3_description_prompt.replace( + "{params}", ", ".join(param_list) + ) + func_decl["description"] = func_decl.get("description", "") + sig_str - This reverses the namespace transformation applied in the request, - ensuring the client receives the original tool names. + return modified + + def _format_type_hint(self, prop_data: Dict[str, Any]) -> str: + """Format a type hint for a property schema.""" + type_hint = prop_data.get("type", "unknown") + + if type_hint == "array": + items = prop_data.get("items", {}) + if isinstance(items, dict): + item_type = items.get("type", "unknown") + if item_type == "object": + nested_props = items.get("properties", {}) + nested_req = items.get("required", []) + if nested_props: + nested_list = [] + for n, d in nested_props.items(): + if isinstance(d, dict): + t = d.get("type", "unknown") + req = " REQUIRED" if n in nested_req else "" + nested_list.append(f"{n}: {t}{req}") + return f"ARRAY_OF_OBJECTS[{', '.join(nested_list)}]" + return "ARRAY_OF_OBJECTS" + return f"ARRAY_OF_{item_type.upper()}" + return "ARRAY" + + return type_hint + + def _strip_gemini3_prefix(self, name: str) -> str: + """Strip the Gemini 3 namespace prefix from a tool name.""" + if name and name.startswith(self._gemini3_tool_prefix): + return name[len(self._gemini3_tool_prefix):] + return name + + # ========================================================================= + # REQUEST TRANSFORMATION + # ========================================================================= + + def _build_tools_payload( + self, + tools: Optional[List[Dict[str, Any]]], + _model: str + ) -> Optional[List[Dict[str, Any]]]: + """Build Gemini-format tools from OpenAI tools.""" + if not tools: + return None - Args: - tool_name: Tool name (possibly with configured prefix) + gemini_tools = [] + for tool in tools: + if tool.get("type") != "function": + continue - Returns: - Original tool name without prefix - """ - if tool_name and tool_name.startswith(self._gemini3_tool_prefix): - return tool_name[len(self._gemini3_tool_prefix):] - return tool_name - - # ============================================================================ - # ANTIGRAVITY REQUEST TRANSFORMATION - # ============================================================================ - + func = tool.get("function", {}) + params = func.get("parameters") + + func_decl = { + "name": func.get("name", ""), + "description": func.get("description", "") + } + + if params and isinstance(params, dict): + schema = dict(params) + schema.pop("$schema", None) + schema.pop("strict", None) + schema = _normalize_type_arrays(schema) + func_decl["parametersJsonSchema"] = schema + else: + func_decl["parametersJsonSchema"] = {"type": "object", "properties": {}} + + gemini_tools.append({"functionDeclarations": [func_decl]}) + + return gemini_tools or None + def _transform_to_antigravity_format( self, - gemini_cli_payload: Dict[str, Any], - model: str + gemini_payload: Dict[str, Any], + model: str, + max_tokens: Optional[int] = None ) -> Dict[str, Any]: """ - Transform Gemini CLI format to complete Antigravity format. + Transform Gemini CLI payload to complete Antigravity format. Args: - gemini_cli_payload: Request in Gemini CLI format + gemini_payload: Request in Gemini CLI format model: Model name (public alias) - - Returns: - Complete Antigravity request payload + max_tokens: Max output tokens (including thinking) """ - internal_model = self._alias_to_model_name(model) + internal_model = self._alias_to_internal(model) - # 1. Wrap in Antigravity envelope + # Wrap in Antigravity envelope antigravity_payload = { - "project": self.generate_project_id(), + "project": _generate_project_id(), "userAgent": "antigravity", - "requestId": self.generate_request_id(), - "model": internal_model, # Use internal name - "request": copy.deepcopy(gemini_cli_payload) + "requestId": _generate_request_id(), + "model": internal_model, + "request": copy.deepcopy(gemini_payload) } - # 2. Add session ID - antigravity_payload["request"]["sessionId"] = self.generate_session_id() + # Add session ID + antigravity_payload["request"]["sessionId"] = _generate_session_id() - # 3. Remove fields that Antigravity doesn't support + # Remove unsupported fields antigravity_payload["request"].pop("safetySettings", None) - if "generationConfig" in antigravity_payload["request"]: - antigravity_payload["request"]["generationConfig"].pop("maxOutputTokens", None) - - # 4. Set toolConfig mode - if "toolConfig" not in antigravity_payload["request"]: - antigravity_payload["request"]["toolConfig"] = {} - if "functionCallingConfig" not in antigravity_payload["request"]["toolConfig"]: - antigravity_payload["request"]["toolConfig"]["functionCallingConfig"] = {} - antigravity_payload["request"]["toolConfig"]["functionCallingConfig"]["mode"] = "VALIDATED" - - # 5. Handle Gemini 3 specific thinking logic - # For non-Gemini-3 models, convert thinkingLevel to thinkingBudget + + # Handle max_tokens - only apply to Claude, or if explicitly set for others + gen_config = antigravity_payload["request"].get("generationConfig", {}) + is_claude = self._is_claude(model) + + if max_tokens is not None: + # Explicitly set in request - apply to all models + gen_config["maxOutputTokens"] = max_tokens + elif is_claude: + # Claude model without explicit max_tokens - use default + gen_config["maxOutputTokens"] = DEFAULT_MAX_OUTPUT_TOKENS + # For non-Claude models without explicit max_tokens, don't set it + + antigravity_payload["request"]["generationConfig"] = gen_config + + # Set toolConfig mode + tool_config = antigravity_payload["request"].setdefault("toolConfig", {}) + func_config = tool_config.setdefault("functionCallingConfig", {}) + func_config["mode"] = "VALIDATED" + + # Handle Gemini 3 thinking logic if not internal_model.startswith("gemini-3-"): - gen_config = antigravity_payload["request"].get("generationConfig", {}) thinking_config = gen_config.get("thinkingConfig", {}) if "thinkingLevel" in thinking_config: - # Remove thinkingLevel for non-Gemini-3 models del thinking_config["thinkingLevel"] - # Set thinkingBudget to -1 (auto/dynamic) thinking_config["thinkingBudget"] = -1 - # 6. Preserve/add thoughtSignature to function calls in model role content (GEMINI 3 ONLY) - # thoughtSignature is a Gemini 3 feature for preserving reasoning context in multi-turn conversations - # DO NOT add this for Claude or other models - they don't support it! + # Add thoughtSignature to function calls for Gemini 3 if internal_model.startswith("gemini-3-"): for content in antigravity_payload["request"].get("contents", []): if content.get("role") == "model": for part in content.get("parts", []): - # Add signature to function calls OR preserve if already exists if "functionCall" in part and "thoughtSignature" not in part: part["thoughtSignature"] = "skip_thought_signature_validator" - # 7. CLAUDE-SPECIFIC TOOL SCHEMA TRANSFORMATION - # Reference: Go implementation antigravity_executor.go lines 672-684 - # For Claude models: parametersJsonSchema → parameters, remove $schema + # Claude-specific tool schema transformation if internal_model.startswith("claude-sonnet-"): - lib_logger.debug(f"Applying Claude-specific tool schema transformation for {internal_model}") - tools = antigravity_payload["request"].get("tools", []) - - for tool in tools: - function_declarations = tool.get("functionDeclarations", []) - for func_decl in function_declarations: - if "parametersJsonSchema" in func_decl: - params = func_decl["parametersJsonSchema"] - - # CRITICAL: Claude requires clean JSON Schema draft 2020-12 - # Recursively remove ALL incompatible fields - def clean_claude_schema(schema): - """Recursively remove fields Claude doesn't support.""" - if not isinstance(schema, dict): - return schema - - # Fields that break Claude's JSON Schema validation - incompatible = {'$schema', 'additionalProperties', 'minItems', 'maxItems', 'pattern'} - cleaned = {} - - for key, value in schema.items(): - if key in incompatible: - continue # Skip incompatible fields - - if isinstance(value, dict): - cleaned[key] = clean_claude_schema(value) - elif isinstance(value, list): - cleaned[key] = [ - clean_claude_schema(item) if isinstance(item, dict) else item - for item in value - ] - else: - cleaned[key] = value - - return cleaned - - # Clean the schema - params = clean_claude_schema(params) if isinstance(params, dict) else params - - # Rename parametersJsonSchema → parameters for Claude - func_decl["parameters"] = params - del func_decl["parametersJsonSchema"] + self._apply_claude_tool_transform(antigravity_payload) return antigravity_payload - - #============================================================================ - # BASE URL FALLBACK LOGIC - # ============================================================================ - - def _get_current_base_url(self) -> str: - """Get the current base URL from the fallback list.""" - return self._current_base_url - - def _try_next_base_url(self) -> bool: - """ - Switch to the next base URL in the fallback list. - - Returns: - True if successfully switched to next URL, False if no more URLs available - """ - if self._base_url_index < len(BASE_URLS) - 1: - self._base_url_index += 1 - self._current_base_url = BASE_URLS[self._base_url_index] - lib_logger.info(f"Switching to fallback Antigravity base URL: {self._current_base_url}") - return True - return False - - def _reset_base_url(self): - """Reset to the primary base URL (daily sandbox).""" - self._base_url_index = 0 - self._current_base_url = BASE_URLS[0] - - # ============================================================================ - # RESPONSE TRANSFORMATION (Antigravity → OpenAI) - # ============================================================================ - - def _unwrap_antigravity_response(self, antigravity_response: Dict[str, Any]) -> Dict[str, Any]: - """ - Extract Gemini response from Antigravity envelope. - - Args: - antigravity_response: Response from Antigravity API - - Returns: - Gemini response (unwrapped) - """ - # For both streaming and non-streaming, response is in 'response' field - return antigravity_response.get("response", antigravity_response) - - @staticmethod - def _recursively_parse_json_strings(obj: Any) -> Any: - """ - Recursively parse JSON strings in nested data structures. - - Antigravity (especially for Claude models) sometimes returns tool arguments - with JSON-stringified values: {"files": "[{...}]"} instead of {"files": [{...}]}. - This causes double-encoding when we call json.dumps() on it. - - This function recursively detects and parses such strings to restore proper structure. - - Args: - obj: Any value (dict, list, str, etc.) - - Returns: - Parsed version with JSON strings converted to their object form - """ - if isinstance(obj, dict): - # Recursively process dictionary values - return {k: AntigravityProvider._recursively_parse_json_strings(v) for k, v in obj.items()} - elif isinstance(obj, list): - # Recursively process list items - return [AntigravityProvider._recursively_parse_json_strings(item) for item in obj] - elif isinstance(obj, str): - # Check if this string looks like JSON - stripped = obj.strip() - if (stripped.startswith('{') and stripped.endswith('}')) or \ - (stripped.startswith('[') and stripped.endswith(']')): - try: - # Attempt to parse as JSON - parsed = json.loads(obj) - # Recursively process the parsed result (it might contain more JSON strings) - return AntigravityProvider._recursively_parse_json_strings(parsed) - except (json.JSONDecodeError, ValueError): - # Not valid JSON, return as-is - return obj - else: - return obj - else: - # Primitive types (int, bool, None, etc.) - return as-is - return obj - + + def _apply_claude_tool_transform(self, payload: Dict[str, Any]) -> None: + """Apply Claude-specific tool schema transformations.""" + tools = payload["request"].get("tools", []) + for tool in tools: + for func_decl in tool.get("functionDeclarations", []): + if "parametersJsonSchema" in func_decl: + params = func_decl["parametersJsonSchema"] + params = _clean_claude_schema(params) if isinstance(params, dict) else params + func_decl["parameters"] = params + del func_decl["parametersJsonSchema"] + + # ========================================================================= + # RESPONSE TRANSFORMATION + # ========================================================================= + + def _unwrap_response(self, response: Dict[str, Any]) -> Dict[str, Any]: + """Extract Gemini response from Antigravity envelope.""" + return response.get("response", response) + def _gemini_to_openai_chunk( - self, - gemini_chunk: Dict[str, Any], + self, + chunk: Dict[str, Any], model: str, - stream_accumulator: Optional[Dict[str, Any]] = None + accumulator: Optional[Dict[str, Any]] = None ) -> Dict[str, Any]: """ - Convert a Gemini API response chunk to OpenAI format. - - UPDATED: Now preserves thoughtSignatures for Gemini 3 multi-turn conversations: - - Stores signatures in server-side cache (if enabled) - - Includes signatures in response (if client passthrough enabled) - - Filters standalone signature parts (no functionCall/text) - - FIXED: Handles Antigravity's double-encoded JSON in tool arguments - - Recursively parses JSON-stringified values before serialization - - Prevents "Unexpected non-whitespace character after JSON" errors - - Claude Thinking Caching: - - For Claude models, thinking content is accumulated across all chunks - - The stream_accumulator collects reasoning_content and thought_signature - - Caching happens AFTER the full stream is processed (in _handle_streaming) + Convert Gemini response chunk to OpenAI streaming format. Args: - gemini_chunk: Gemini API response chunk - model: Model name for Gemini 3 detection - stream_accumulator: Optional dict to accumulate streaming data for post-processing - - Returns: - OpenAI-compatible response chunk + chunk: Gemini API response chunk + model: Model name + accumulator: Optional dict to accumulate data for post-processing """ - # Extract the main response structure - candidates = gemini_chunk.get("candidates", []) + candidates = chunk.get("candidates", []) if not candidates: return {} candidate = candidates[0] - content = candidate.get("content", {}) - content_parts = content.get("parts", []) + content_parts = candidate.get("content", {}).get("parts", []) - # Build delta components text_content = "" reasoning_content = "" tool_calls = [] - - # Track if we've seen a signature yet (for parallel tool call handling) - # Per Gemini 3 spec: only FIRST tool call in parallel gets signature - first_signature_seen = False - tool_call_index = 0 # Track index for OpenAI streaming format + first_sig_seen = False + tool_idx = 0 for part in content_parts: - has_function_call = "functionCall" in part + has_func = "functionCall" in part has_text = "text" in part - has_signature = "thoughtSignature" in part and part["thoughtSignature"] - is_thought = part.get("thought") is True or (isinstance(part.get("thought"), str) and part.get("thought").lower() == 'true') + has_sig = bool(part.get("thoughtSignature")) + is_thought = part.get("thought") is True or str(part.get("thought")).lower() == 'true' - # Accumulate thought signature from thinking parts (Claude caching) - # The signature appears on the LAST thinking part (the one with empty text after all thinking) - if has_signature and is_thought and stream_accumulator is not None: - stream_accumulator["thought_signature"] = part["thoughtSignature"] + # Accumulate signature for Claude caching + if has_sig and is_thought and accumulator is not None: + accumulator["thought_signature"] = part["thoughtSignature"] - # Skip standalone signature-only parts (empty thinking parts with just signature) - if has_signature and not has_function_call and (not has_text or part.get("text") == ""): + # Skip standalone signature parts + if has_sig and not has_func and (not has_text or not part.get("text")): continue - # Process text content if has_text: + text = part["text"] if is_thought: - reasoning_content += part["text"] - # Accumulate reasoning for Claude caching - if stream_accumulator is not None: - stream_accumulator["reasoning_content"] += part["text"] + reasoning_content += text + if accumulator is not None: + accumulator["reasoning_content"] += text else: - text_content += part["text"] - # Accumulate text content for cache key generation - if stream_accumulator is not None: - stream_accumulator["text_content"] += part["text"] + text_content += text + if accumulator is not None: + accumulator["text_content"] += text - # Process function calls (NOW WORKS with signatures!) - if has_function_call: - func_call = part["functionCall"] - - # Use ID from Antigravity if provided, otherwise generate - tool_call_id = func_call.get("id") or f"call_{uuid.uuid4().hex[:24]}" + if has_func: + tool_call = self._extract_tool_call(part, model, tool_idx, accumulator) - # Get tool name and strip gemini3_ namespace if present (Gemini 3 specific) - tool_name = func_call.get("name", "") - if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: - tool_name = self._strip_gemini3_namespace_from_name(tool_name) - - tool_call = { - "id": tool_call_id, - "type": "function", - "index": tool_call_index, # REQUIRED for OpenAI streaming format - "function": { - "name": tool_name, - "arguments": json.dumps(func_call.get("args", {})) - } - } - tool_call_index += 1 # Increment for next tool call - - # Accumulate tool calls for Claude caching - if stream_accumulator is not None: - stream_accumulator["tool_calls"].append(tool_call) - - # Handle thoughtSignature if present (on function call part) - if has_signature and not first_signature_seen: - # Only first tool call gets signature (parallel call handling) - first_signature_seen = True - signature = part["thoughtSignature"] - - # Option 1: Store in server-side cache (if enabled) - if self._enable_signature_cache: - self._signature_cache.store(tool_call_id, signature) - lib_logger.debug(f"Stored thoughtSignature in cache for {tool_call_id}") - - # Option 2: Pass to client (if enabled) - INDEPENDENT of cache! - if self._preserve_signatures_in_client: - tool_call["thought_signature"] = signature + if has_sig and not first_sig_seen: + first_sig_seen = True + self._handle_tool_signature(tool_call, part["thoughtSignature"]) tool_calls.append(tool_call) + tool_idx += 1 # Build delta delta = {} @@ -1683,55 +1311,19 @@ def _gemini_to_openai_chunk( delta["role"] = "assistant" # Handle finish reason - finish_reason = candidate.get("finishReason") - if finish_reason: - # Map Gemini finish reasons to OpenAI - finish_reason_map = { - "STOP": "stop", - "MAX_TOKENS": "length", - "SAFETY": "content_filter", - "RECITATION": "content_filter", - "OTHER": "stop" - } - finish_reason = finish_reason_map.get(finish_reason, "stop") - if tool_calls: - finish_reason = "tool_calls" - - # Mark stream as complete for accumulator - if stream_accumulator is not None: - stream_accumulator["is_complete"] = True - - # Build usage metadata - usage = None - usage_metadata = gemini_chunk.get("usageMetadata", {}) - if usage_metadata: - prompt_tokens = usage_metadata.get("promptTokenCount", 0) - thoughts_tokens = usage_metadata.get("thoughtsTokenCount", 0) - completion_tokens = usage_metadata.get("candidatesTokenCount", 0) - - usage = { - "prompt_tokens": prompt_tokens + thoughts_tokens, # Include thoughts in prompt - "completion_tokens": completion_tokens, - "total_tokens": usage_metadata.get("totalTokenCount", 0) - } - - # Add reasoning tokens details if thinking was used - if thoughts_tokens > 0: - if "completion_tokens_details" not in usage: - usage["completion_tokens_details"] = {} - usage["completion_tokens_details"]["reasoning_tokens"] = thoughts_tokens + finish_reason = self._map_finish_reason(candidate.get("finishReason"), bool(tool_calls)) + if finish_reason and accumulator is not None: + accumulator["is_complete"] = True + + # Build usage + usage = self._build_usage(chunk.get("usageMetadata", {})) - # Build final response response = { - "id": gemini_chunk.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), + "id": chunk.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), "object": "chat.completion.chunk", "created": int(time.time()), "model": model, - "choices": [{ - "index": 0, - "delta": delta, - "finish_reason": finish_reason - }] + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}] } if usage: @@ -1739,302 +1331,246 @@ def _gemini_to_openai_chunk( return response - def _gemini_to_openai_non_streaming(self, gemini_response: Dict[str, Any], model: str) -> Dict[str, Any]: - """ - Convert a Gemini API response to OpenAI non-streaming format. - - This is specifically for non-streaming completions where we need 'message' instead of 'delta'. - Also handles Claude thinking caching for non-streaming responses. - - Args: - gemini_response: Gemini API response - model: Model name for Gemini 3 detection - - Returns: - OpenAI-compatible non-streaming response - """ - # Extract the main response structure - candidates = gemini_response.get("candidates", []) + def _gemini_to_openai_non_streaming( + self, + response: Dict[str, Any], + model: str + ) -> Dict[str, Any]: + """Convert Gemini response to OpenAI non-streaming format.""" + candidates = response.get("candidates", []) if not candidates: return {} candidate = candidates[0] - content = candidate.get("content", {}) - content_parts = content.get("parts", []) + content_parts = candidate.get("content", {}).get("parts", []) - # Build message components text_content = "" reasoning_content = "" tool_calls = [] - thought_signature = "" # Track signature for Claude caching - - # Track if we've seen a signature yet (for parallel tool call handling) - first_signature_seen = False + thought_sig = "" + first_sig_seen = False for part in content_parts: - has_function_call = "functionCall" in part + has_func = "functionCall" in part has_text = "text" in part - has_signature = "thoughtSignature" in part and part["thoughtSignature"] - is_thought = part.get("thought") is True or (isinstance(part.get("thought"), str) and part.get("thought").lower() == 'true') + has_sig = bool(part.get("thoughtSignature")) + is_thought = part.get("thought") is True or str(part.get("thought")).lower() == 'true' - # Capture thought signature (appears on last thinking part) - if has_signature and is_thought: - thought_signature = part["thoughtSignature"] + if has_sig and is_thought: + thought_sig = part["thoughtSignature"] - # Skip standalone signature parts (empty thinking parts with just signature) - if has_signature and not has_function_call and (not has_text or part.get("text") == ""): + if has_sig and not has_func and (not has_text or not part.get("text")): continue - # Process text content if has_text: if is_thought: reasoning_content += part["text"] else: text_content += part["text"] - # Process function calls - if has_function_call: - func_call = part["functionCall"] - - # Use ID from Antigravity if provided, otherwise generate - tool_call_id = func_call.get("id") or f"call_{uuid.uuid4().hex[:24]}" - - # Get tool name and strip gemini3_ namespace if present - tool_name = func_call.get("name", "") - if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: - tool_name = self._strip_gemini3_namespace_from_name(tool_name) - - # Get raw args from Antigravity - raw_args = func_call.get("args", {}) - - # FIX: Recursively parse JSON-stringified values - # Antigravity (especially Claude) returns: {"files": "[{...}]"} - # We need to parse these strings before calling json.dumps() - parsed_args = self._recursively_parse_json_strings(raw_args) - - tool_call = { - "id": tool_call_id, - "type": "function", - "function": { - "name": tool_name, - "arguments": json.dumps(parsed_args) - } - } + if has_func: + tool_call = self._extract_tool_call(part, model, len(tool_calls)) - # Handle thoughtSignature if present (on function call part) - if has_signature and not first_signature_seen: - first_signature_seen = True - signature = part["thoughtSignature"] - - # Store in server-side cache - if self._enable_signature_cache: - self._signature_cache.store(tool_call_id, signature) - lib_logger.debug(f"Stored thoughtSignature in cache for {tool_call_id}") - - # Pass to client if enabled - if self._preserve_signatures_in_client: - tool_call["thought_signature"] = signature + if has_sig and not first_sig_seen: + first_sig_seen = True + self._handle_tool_signature(tool_call, part["thoughtSignature"]) tool_calls.append(tool_call) - # Cache Claude thinking content for non-streaming responses - if reasoning_content and model.startswith("claude-") and self._enable_signature_cache: - cache_key = self._generate_thinking_cache_key(text_content, tool_calls) - - if cache_key: - thinking_data = { - "thinking_text": reasoning_content, - "thought_signature": thought_signature, - "text_preview": text_content[:100] if text_content else "", - "tool_ids": [tc.get("id", "") for tc in tool_calls] if tool_calls else [], - "timestamp": time.time() - } - - self._thinking_cache.store(cache_key, json.dumps(thinking_data)) - lib_logger.info( - f"✓ Cached Claude thinking (non-streaming): {cache_key[:50]}... " - f"(reasoning={len(reasoning_content)} chars, " - f"tools={len(tool_calls)}, " - f"sig={'yes' if thought_signature else 'no'})" - ) - - # Build message object (not delta!) - message = {"role": "assistant"} + # Cache Claude thinking + if reasoning_content and self._is_claude(model) and self._enable_signature_cache: + self._cache_thinking(reasoning_content, thought_sig, text_content, tool_calls) + # Build message + message = {"role": "assistant"} if text_content: message["content"] = text_content elif not tool_calls: - # If no text and no tool calls, set content to empty string message["content"] = "" - if reasoning_content: message["reasoning_content"] = reasoning_content - if tool_calls: message["tool_calls"] = tool_calls - # Don't set content if we have tool calls (OpenAI convention) - if "content" in message: - message.pop("content") + message.pop("content", None) - # Handle finish reason - finish_reason = candidate.get("finishReason") - if finish_reason: - # Map Gemini finish reasons to OpenAI - finish_reason_map = { - "STOP": "stop", - "MAX_TOKENS": "length", - "SAFETY": "content_filter", - "RECITATION": "content_filter", - "OTHER": "stop" - } - finish_reason = finish_reason_map.get(finish_reason, "stop") - if tool_calls: - finish_reason = "tool_calls" - - # Build usage metadata - usage = None - usage_metadata = gemini_response.get("usageMetadata", {}) - if usage_metadata: - prompt_tokens = usage_metadata.get("promptTokenCount", 0) - thoughts_tokens = usage_metadata.get("thoughtsTokenCount", 0) - completion_tokens = usage_metadata.get("candidatesTokenCount", 0) - - usage = { - "prompt_tokens": prompt_tokens + thoughts_tokens, - "completion_tokens": completion_tokens, - "total_tokens": usage_metadata.get("totalTokenCount", 0) - } - - # Add reasoning tokens details if thinking was used - if thoughts_tokens > 0: - if "completion_tokens_details" not in usage: - usage["completion_tokens_details"] = {} - usage["completion_tokens_details"]["reasoning_tokens"] = thoughts_tokens + finish_reason = self._map_finish_reason(candidate.get("finishReason"), bool(tool_calls)) + usage = self._build_usage(response.get("usageMetadata", {})) - # Build final response - response = { - "id": gemini_response.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), - "object": "chat.completion", # Non-streaming uses chat.completion, not chunk + result = { + "id": response.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), + "object": "chat.completion", "created": int(time.time()), "model": model, - "choices": [{ - "index": 0, - "message": message, # message, not delta! - "finish_reason": finish_reason - }] + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}] } if usage: - response["usage"] = usage + result["usage"] = usage - return response - - - - # ============================================================================ + return result + + def _extract_tool_call( + self, + part: Dict[str, Any], + model: str, + index: int, + accumulator: Optional[Dict[str, Any]] = None + ) -> Dict[str, Any]: + """Extract and format a tool call from a response part.""" + func_call = part["functionCall"] + tool_id = func_call.get("id") or f"call_{uuid.uuid4().hex[:24]}" + + tool_name = func_call.get("name", "") + if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: + tool_name = self._strip_gemini3_prefix(tool_name) + + raw_args = func_call.get("args", {}) + parsed_args = _recursively_parse_json_strings(raw_args) + + tool_call = { + "id": tool_id, + "type": "function", + "index": index, + "function": { + "name": tool_name, + "arguments": json.dumps(parsed_args) + } + } + + if accumulator is not None: + accumulator["tool_calls"].append(tool_call) + + return tool_call + + def _handle_tool_signature(self, tool_call: Dict, signature: str) -> None: + """Handle thoughtSignature for a tool call.""" + tool_id = tool_call["id"] + + if self._enable_signature_cache: + self._signature_cache.store(tool_id, signature) + lib_logger.debug(f"Stored signature for {tool_id}") + + if self._preserve_signatures_in_client: + tool_call["thought_signature"] = signature + + def _map_finish_reason( + self, + gemini_reason: Optional[str], + has_tool_calls: bool + ) -> Optional[str]: + """Map Gemini finish reason to OpenAI format.""" + if not gemini_reason: + return None + reason = FINISH_REASON_MAP.get(gemini_reason, "stop") + return "tool_calls" if has_tool_calls else reason + + def _build_usage(self, metadata: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Build usage dict from Gemini usage metadata.""" + if not metadata: + return None + + prompt = metadata.get("promptTokenCount", 0) + thoughts = metadata.get("thoughtsTokenCount", 0) + completion = metadata.get("candidatesTokenCount", 0) + + usage = { + "prompt_tokens": prompt + thoughts, + "completion_tokens": completion, + "total_tokens": metadata.get("totalTokenCount", 0) + } + + if thoughts > 0: + usage["completion_tokens_details"] = {"reasoning_tokens": thoughts} + + return usage + + def _cache_thinking( + self, + reasoning: str, + signature: str, + text: str, + tool_calls: List[Dict] + ) -> None: + """Cache Claude thinking content.""" + cache_key = self._generate_thinking_cache_key(text, tool_calls) + if not cache_key: + return + + data = { + "thinking_text": reasoning, + "thought_signature": signature, + "text_preview": text[:100] if text else "", + "tool_ids": [tc.get("id", "") for tc in tool_calls], + "timestamp": time.time() + } + + self._thinking_cache.store(cache_key, json.dumps(data)) + lib_logger.info(f"Cached thinking: {cache_key[:50]}...") + + # ========================================================================= # PROVIDER INTERFACE IMPLEMENTATION - # ============================================================================ - + # ========================================================================= + async def get_valid_token(self, credential_identifier: str) -> str: - """ - Get a valid access token for the credential. - - Args: - credential_identifier: Credential file path or "env" - - Returns: - Access token string - """ + """Get a valid access token for the credential.""" creds = await self._load_credentials(credential_identifier) if self._is_token_expired(creds): creds = await self._refresh_token(credential_identifier, creds) return creds['access_token'] - + def has_custom_logic(self) -> bool: """Antigravity uses custom translation logic.""" return True - + async def get_auth_header(self, credential_identifier: str) -> Dict[str, str]: - """ - Get OAuth authorization header for Antigravity. - - Args: - credential_identifier: Credential file path or "env" - - Returns: - Dict with Authorization header - """ - access_token = await self.get_valid_token(credential_identifier) - return {"Authorization": f"Bearer {access_token}"} - - async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: - """ - Fetch available models from Antigravity. - - For Antigravity, we can optionally use the fetchAvailableModels endpoint and apply - alias mapping to convert internal names to public names. However, this endpoint is - often unavailable (404), so dynamic discovery is disabled by default. - - Set ANTIGRAVITY_ENABLE_DYNAMIC_MODELS=true to enable dynamic discovery. - - Args: - api_key: Credential path (not a traditional API key) - client: HTTP client - - Returns: - List of public model names - """ - # If dynamic discovery is disabled, immediately return hardcoded list - if not self._enable_dynamic_model_discovery: - lib_logger.debug("Using hardcoded Antigravity model list (dynamic discovery disabled)") - return [f"antigravity/{m}" for m in HARDCODED_MODELS] - - # Dynamic discovery enabled - attempt to fetch from API - credential_path = api_key # For OAuth providers, this is the credential path + """Get OAuth authorization header.""" + token = await self.get_valid_token(credential_identifier) + return {"Authorization": f"Bearer {token}"} + + async def get_models( + self, + api_key: str, + client: httpx.AsyncClient + ) -> List[str]: + """Fetch available models from Antigravity.""" + if not self._enable_dynamic_models: + lib_logger.debug("Using hardcoded model list") + return [f"antigravity/{m}" for m in AVAILABLE_MODELS] try: - access_token = await self.get_valid_token(credential_path) - base_url = self._get_current_base_url() - - url = f"{base_url}/fetchAvailableModels" + token = await self.get_valid_token(api_key) + url = f"{self._get_base_url()}/fetchAvailableModels" headers = { - "Authorization": f"Bearer {access_token}", + "Authorization": f"Bearer {token}", "Content-Type": "application/json" } - payload = { - "project": self.generate_project_id(), - "requestId": self.generate_request_id(), + "project": _generate_project_id(), + "requestId": _generate_request_id(), "userAgent": "antigravity" } response = await client.post(url, json=payload, headers=headers, timeout=30.0) response.raise_for_status() - data = response.json() - # Extract model names and apply aliasing models = [] - if "models" in data: - for model_info in data["models"]: - internal_name = model_info.get("name", "").replace("models/", "") - if internal_name: - public_name = self._model_name_to_alias(internal_name) - if public_name: # Skip excluded models (empty string) - models.append(f"antigravity/{public_name}") + for model_info in data.get("models", []): + internal = model_info.get("name", "").replace("models/", "") + if internal: + public = self._internal_to_alias(internal) + if public: + models.append(f"antigravity/{public}") if models: - lib_logger.info(f"Discovered {len(models)} Antigravity models via dynamic discovery") + lib_logger.info(f"Discovered {len(models)} models") return models - else: - lib_logger.warning("No models returned from Antigravity, using hardcoded list") - return [f"antigravity/{m}" for m in HARDCODED_MODELS] - except Exception as e: - lib_logger.warning(f"Failed to fetch Antigravity models: {e}, using hardcoded list") - return [f"antigravity/{m}" for m in HARDCODED_MODELS] - + lib_logger.warning(f"Dynamic model discovery failed: {e}") + + return [f"antigravity/{m}" for m in AVAILABLE_MODELS] + async def acompletion( self, client: httpx.AsyncClient, @@ -2043,229 +1579,121 @@ async def acompletion( """ Handle completion requests for Antigravity. - This is the main entry point that: - 1. Extracts the model and credential path - 2. Transforms OpenAI request → Gemini CLI → Antigravity format - 3. Makes the API call with fallback logic - 4. Transforms Antigravity response → Gemini → OpenAI format - - Args: - client: HTTP client - **kwargs: LiteLLM completion parameters - - Returns: - ModelResponse (non-streaming) or AsyncGenerator (streaming) + Main entry point that: + 1. Extracts parameters and transforms messages + 2. Builds Antigravity request payload + 3. Makes API call with fallback logic + 4. Transforms response to OpenAI format """ - # Extract key parameters - model = kwargs.get("model", "gemini-2.5-pro") - - # Strip provider prefix from model name (e.g., "antigravity/claude-sonnet-4-5-thinking" -> "claude-sonnet-4-5-thinking") - if "/" in model: - model = model.split("/")[-1] - + # Extract parameters + model = self._strip_provider_prefix(kwargs.get("model", "gemini-2.5-pro")) messages = kwargs.get("messages", []) stream = kwargs.get("stream", False) credential_path = kwargs.pop("credential_identifier", kwargs.get("api_key", "")) tools = kwargs.get("tools") reasoning_effort = kwargs.get("reasoning_effort") - temperature = kwargs.get("temperature") top_p = kwargs.get("top_p") max_tokens = kwargs.get("max_tokens") - enable_request_logging = kwargs.pop("enable_request_logging", False) - - #lib_logger.debug(f"Antigravity completion: model={model}, stream={stream}, messages={len(messages)}") - - # Create file logger - file_logger = _AntigravityFileLogger( - model_name=model, - enabled=enable_request_logging - ) + custom_budget = kwargs.get("custom_reasoning_budget", False) + enable_logging = kwargs.pop("enable_request_logging", False) - # Step 1: Transform messages (OpenAI → Gemini CLI) - system_instruction, gemini_contents = self._transform_messages(messages, model=model) + # Create logger + file_logger = AntigravityFileLogger(model, enable_logging) - # Apply tool response grouping + # Transform messages + system_instruction, gemini_contents = self._transform_messages(messages, model) gemini_contents = self._fix_tool_response_grouping(gemini_contents) - # Step 2: Build Gemini CLI payload - gemini_cli_payload = { - "contents": gemini_contents - } + # Build payload + gemini_payload = {"contents": gemini_contents} if system_instruction: - gemini_cli_payload["system_instruction"] = system_instruction - - # Apply Gemini 3 system instruction injection (Strategy 3) if fix is enabled - # This prepends critical tool usage instructions to override model's training data - if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix and tools: - gemini3_instruction = self._gemini3_system_instruction - - if "system_instruction" in gemini_cli_payload: - # Prepend to existing system instruction - existing_instruction = gemini_cli_payload["system_instruction"] - if isinstance(existing_instruction, dict) and "parts" in existing_instruction: - # System instruction with parts structure - gemini3_part = {"text": gemini3_instruction} - existing_instruction["parts"].insert(0, gemini3_part) - else: - # Shouldn't happen, but handle gracefully - gemini_cli_payload["system_instruction"] = { - "role": "user", - "parts": [ - {"text": gemini3_instruction}, - {"text": str(existing_instruction)} - ] - } - else: - # Create new system instruction with Gemini 3 instructions - gemini_cli_payload["system_instruction"] = { - "role": "user", - "parts": [{"text": gemini3_instruction}] - } - - #lib_logger.debug("Gemini 3 system instruction injection applied") - + gemini_payload["system_instruction"] = system_instruction + # Inject Gemini 3 system instruction + if self._is_gemini_3(model) and self._enable_gemini3_tool_fix and tools: + self._inject_gemini3_system_instruction(gemini_payload) # Add generation config - generation_config = {} - - # Temperature handling: Default to 1.0, override 0 to 1.0 - # Low temperature (especially 0) makes models deterministic and prone to following - # training data patterns instead of actual schemas, which causes tool hallucination - + gen_config = {} if top_p is not None: - generation_config["topP"] = top_p + gen_config["topP"] = top_p - # Extract custom_reasoning_budget toggle - # Check kwargs first, then headers if not found - custom_reasoning_budget = kwargs.get("custom_reasoning_budget", False) - - # Handle thinking config - thinking_config = self._map_reasoning_effort_to_thinking_config( - reasoning_effort, - model, - custom_reasoning_budget - ) + thinking_config = self._get_thinking_config(reasoning_effort, model, custom_budget) if thinking_config: - generation_config.setdefault("thinkingConfig", {}).update(thinking_config) - - if generation_config: - gemini_cli_payload["generationConfig"] = generation_config - - # Add tools - using Go reference implementation approach - # Go code (line 298-328): renames 'parameters' -> 'parametersJsonSchema' and removes 'strict' - if tools: - gemini_tools = [] - for tool in tools: - if tool.get("type") == "function": - func = tool.get("function", {}) - - # Get parameters dict (may be missing) - parameters = func.get("parameters") - - # Build function declaration - func_decl = { - "name": func.get("name", ""), - "description": func.get("description", "") - } - - # Handle parameters -> parametersJsonSchema conversion (matching Go) - if parameters and isinstance(parameters, dict): - # Make a copy to avoid modifying original - schema = dict(parameters) - # Remove OpenAI-specific fields that Antigravity doesn't support - schema.pop("$schema", None) - schema.pop("strict", None) - # CRITICAL: Normalize type arrays for protobuf compatibility - # Converts ["string", "null"] → "string" to avoid "Proto field is not repeating" errors - schema = self._normalize_type_arrays(schema) - func_decl["parametersJsonSchema"] = schema - else: - # No parameters provided - set default empty schema (matching Go lines 318-323) - func_decl["parametersJsonSchema"] = { - "type": "object", - "properties": {} - } - - gemini_tools.append({ - "functionDeclarations": [func_decl] - }) - - if gemini_tools: - gemini_cli_payload["tools"] = gemini_tools - - # Apply Gemini 3 specific tool transformations (ONLY for gemini-3-* models) - # This implements the "Double-Lock" catch-all strategy to prevent tool hallucination - if self._is_gemini_3_model(model) and self._enable_gemini3_tool_fix: - #lib_logger.debug(f"Applying Gemini 3 catch-all tool transformations for {model}") - - # Strategy 1: Namespace prefixing (breaks association with training data) - gemini_cli_payload["tools"] = self._apply_gemini3_namespace_to_tools( - gemini_cli_payload["tools"] - ) - - # Strategy 2: Signature injection (natural language schema enforcement) - gemini_cli_payload["tools"] = self._inject_signature_into_tool_descriptions( - gemini_cli_payload["tools"] - ) - + gen_config.setdefault("thinkingConfig", {}).update(thinking_config) - # Step 3: Transform to Antigravity format - antigravity_payload = self._transform_to_antigravity_format(gemini_cli_payload, model) + if gen_config: + gemini_payload["generationConfig"] = gen_config - # Log the request - file_logger.log_request(antigravity_payload) + # Add tools + gemini_tools = self._build_tools_payload(tools, model) + if gemini_tools: + gemini_payload["tools"] = gemini_tools + + # Apply Gemini 3 tool transformations + if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: + gemini_payload["tools"] = self._apply_gemini3_namespace(gemini_payload["tools"]) + gemini_payload["tools"] = self._inject_signature_into_descriptions(gemini_payload["tools"]) - # Step 4: Make API call - access_token = await self.get_valid_token(credential_path) - base_url = self._get_current_base_url() + # Transform to Antigravity format + payload = self._transform_to_antigravity_format(gemini_payload, model, max_tokens) + file_logger.log_request(payload) + # Make API call + token = await self.get_valid_token(credential_path) + base_url = self._get_base_url() endpoint = ":streamGenerateContent" if stream else ":generateContent" url = f"{base_url}{endpoint}" - - # Add query parameter for streaming (required by Antigravity API) + if stream: url = f"{url}?alt=sse" - - # Extract host from base_url for Host header (required by Google's API) - from urllib.parse import urlparse - parsed_url = urlparse(base_url) - host = parsed_url.netloc if parsed_url.netloc else base_url.replace("https://", "").replace("http://", "").rstrip("/") - + + parsed = urlparse(base_url) + host = parsed.netloc or base_url.replace("https://", "").replace("http://", "").rstrip("/") + headers = { - "Authorization": f"Bearer {access_token}", + "Authorization": f"Bearer {token}", "Content-Type": "application/json", - "Host": host, # CRITICAL: Required by Antigravity API - "User-Agent": "antigravity/1.11.5" # Match Go implementation + "Host": host, + "User-Agent": "antigravity/1.11.5", + "Accept": "text/event-stream" if stream else "application/json" } - - if stream: - headers["Accept"] = "text/event-stream" - else: - headers["Accept"] = "application/json" - - #lib_logger.debug(f"Antigravity request to: {url}") try: if stream: - return self._handle_streaming(client, url, headers, antigravity_payload, model, file_logger) + return self._handle_streaming(client, url, headers, payload, model, file_logger) else: - return await self._handle_non_streaming(client, url, headers, antigravity_payload, model, file_logger) + return await self._handle_non_streaming(client, url, headers, payload, model, file_logger) except Exception as e: - # Try fallback URL if available if self._try_next_base_url(): - lib_logger.warning(f"Retrying Antigravity request with fallback URL: {e}") - base_url = self._get_current_base_url() - url = f"{base_url}{endpoint}" - + lib_logger.warning(f"Retrying with fallback URL: {e}") + url = f"{self._get_base_url()}{endpoint}" if stream: - return self._handle_streaming(client, url, headers, antigravity_payload, model) + return self._handle_streaming(client, url, headers, payload, model, file_logger) else: - return await self._handle_non_streaming(client, url, headers, antigravity_payload, model) + return await self._handle_non_streaming(client, url, headers, payload, model, file_logger) + raise + + def _inject_gemini3_system_instruction(self, payload: Dict[str, Any]) -> None: + """Inject Gemini 3 system instruction for tool fix.""" + if not self._gemini3_system_instruction: + return + + instruction_part = {"text": self._gemini3_system_instruction} + + if "system_instruction" in payload: + existing = payload["system_instruction"] + if isinstance(existing, dict) and "parts" in existing: + existing["parts"].insert(0, instruction_part) else: - raise - + payload["system_instruction"] = { + "role": "user", + "parts": [instruction_part, {"text": str(existing)}] + } + else: + payload["system_instruction"] = {"role": "user", "parts": [instruction_part]} + async def _handle_non_streaming( self, client: httpx.AsyncClient, @@ -2273,27 +1701,21 @@ async def _handle_non_streaming( headers: Dict[str, str], payload: Dict[str, Any], model: str, - file_logger: Optional[_AntigravityFileLogger] = None + file_logger: Optional[AntigravityFileLogger] = None ) -> litellm.ModelResponse: """Handle non-streaming completion.""" response = await client.post(url, headers=headers, json=payload, timeout=120.0) response.raise_for_status() - antigravity_response = response.json() - - # Log response + data = response.json() if file_logger: - file_logger.log_final_response(antigravity_response) - - # Unwrap Antigravity envelope - gemini_response = self._unwrap_antigravity_response(antigravity_response) + file_logger.log_final_response(data) - # Convert to OpenAI non-streaming format (returns dict with 'message' not 'delta') + gemini_response = self._unwrap_response(data) openai_response = self._gemini_to_openai_non_streaming(gemini_response, model) - # Convert dict to ModelResponse object for non-streaming return litellm.ModelResponse(**openai_response) - + async def _handle_streaming( self, client: httpx.AsyncClient, @@ -2301,39 +1723,28 @@ async def _handle_streaming( headers: Dict[str, str], payload: Dict[str, Any], model: str, - file_logger: Optional[_AntigravityFileLogger] = None + file_logger: Optional[AntigravityFileLogger] = None ) -> AsyncGenerator[litellm.ModelResponse, None]: - """ - Handle streaming completion. - - For Claude models with thinking enabled: - - Accumulates reasoning content and thought signature across all chunks - - Caches the complete thinking data AFTER the stream is fully processed - - Uses a generator wrapper to ensure post-stream caching happens - """ - # Create stream accumulator for Claude thinking caching - # This collects data across all chunks so we can cache after stream completes - stream_accumulator = { + """Handle streaming completion.""" + accumulator = { "reasoning_content": "", "thought_signature": "", "text_content": "", "tool_calls": [], "is_complete": False - } if model.startswith("claude-") and self._enable_signature_cache else None + } if self._is_claude(model) and self._enable_signature_cache else None async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response: - # Log error response body for debugging if request failed if response.status_code >= 400: try: error_body = await response.aread() - lib_logger.error(f"Antigravity API error {response.status_code}: {error_body.decode('utf-8', errors='replace')}") - except Exception as e: - lib_logger.error(f"Failed to read error response body: {e}") + lib_logger.error(f"API error {response.status_code}: {error_body.decode()}") + except Exception: + pass response.raise_for_status() async for line in response.aiter_lines(): - # Log raw chunk if file_logger: file_logger.log_response_chunk(line) @@ -2343,89 +1754,25 @@ async def _handle_streaming( break try: - antigravity_chunk = json.loads(data_str) - - # Unwrap Antigravity envelope - gemini_chunk = self._unwrap_antigravity_response(antigravity_chunk) - - # Convert to OpenAI format (with accumulator for Claude) - openai_chunk = self._gemini_to_openai_chunk( - gemini_chunk, - model, - stream_accumulator - ) + chunk = json.loads(data_str) + gemini_chunk = self._unwrap_response(chunk) + openai_chunk = self._gemini_to_openai_chunk(gemini_chunk, model, accumulator) - # Convert dict to ModelResponse object - model_response = litellm.ModelResponse(**openai_chunk) - yield model_response + yield litellm.ModelResponse(**openai_chunk) except json.JSONDecodeError: if file_logger: - file_logger.log_error(f"Failed to parse chunk: {data_str[:100]}") - lib_logger.warning(f"Failed to parse Antigravity chunk: {data_str[:100]}") + file_logger.log_error(f"Parse error: {data_str[:100]}") continue - # After stream completes: cache Claude thinking content - if stream_accumulator and stream_accumulator.get("reasoning_content"): - await self._cache_claude_thinking_after_stream(stream_accumulator, model) + # Cache Claude thinking after stream completes + if accumulator and accumulator.get("reasoning_content"): + self._cache_thinking( + accumulator["reasoning_content"], + accumulator["thought_signature"], + accumulator["text_content"], + accumulator["tool_calls"] + ) - async def _cache_claude_thinking_after_stream( - self, - accumulator: Dict[str, Any], - model: str - ): - """ - Cache Claude thinking content after the complete stream has been processed. - - This is called after ALL streaming chunks have been received, ensuring we have: - - Complete reasoning content (accumulated from all thought=true parts) - - The thoughtSignature (appears on the final thinking part) - - All tool calls with their IDs (for cache key generation) - - Complete text content (for cache key generation) - - Args: - accumulator: Dict with accumulated stream data - model: Model name (for logging) - """ - reasoning_content = accumulator.get("reasoning_content", "") - thought_signature = accumulator.get("thought_signature", "") - text_content = accumulator.get("text_content", "") - tool_calls = accumulator.get("tool_calls", []) - - if not reasoning_content: - lib_logger.debug("No reasoning content to cache") - return - - # Generate cache key from the accumulated response data - cache_key = self._generate_thinking_cache_key(text_content, tool_calls) - - if not cache_key: - lib_logger.warning("Could not generate cache key for Claude thinking") - return - - # Build cache data - thinking_data = { - "thinking_text": reasoning_content, - "thought_signature": thought_signature, - "text_preview": text_content[:100] if text_content else "", - "tool_ids": [tc.get("id", "") for tc in tool_calls] if tool_calls else [], - "timestamp": time.time() - } - - # Store in cache - self._thinking_cache.store(cache_key, json.dumps(thinking_data)) - - lib_logger.info( - f"✓ Cached Claude thinking after stream: {cache_key[:50]}... " - f"(reasoning={len(reasoning_content)} chars, " - f"text={len(text_content)} chars, " - f"tools={len(tool_calls)}, " - f"sig={'yes' if thought_signature else 'no'})" - ) - - # ============================================================================ - # TOKEN COUNTING - # ============================================================================ - async def count_tokens( self, client: httpx.AsyncClient, @@ -2433,105 +1780,45 @@ async def count_tokens( model: str, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, - litellm_params: Optional[Dict[str, Any]] = None + _litellm_params: Optional[Dict[str, Any]] = None ) -> Dict[str, int]: - """ - Counts tokens for the given prompt using the Antigravity :countTokens endpoint. - - Args: - client: The HTTP client to use - credential_path: Path to the credential file - model: Model name to use for token counting - messages: List of messages in OpenAI format - tools: Optional list of tool definitions - litellm_params: Optional additional parameters - - Returns: - Dict with 'prompt_tokens' and 'total_tokens' counts - """ - # Get auth token - access_token = await self.get_valid_token(credential_path) - - # Convert public alias to internal name - internal_model = self._alias_to_model_name(model) - - # Transform messages to Gemini format - system_instruction, contents = self._transform_messages(messages, model=internal_model) - - # Build Gemini CLI payload - gemini_cli_payload = { - "contents": contents - } - - if system_instruction: - gemini_cli_payload["systemInstruction"] = system_instruction - - if tools: - # Transform tools - same as in acompletion - gemini_tools = [] - for tool in tools: - if tool.get("type") == "function": - func = tool.get("function", {}) - parameters = func.get("parameters") - - func_decl = { - "name": func.get("name", ""), - "description": func.get("description", "") - } - - if parameters and isinstance(parameters, dict): - schema = dict(parameters) - schema.pop("$schema", None) - schema.pop("strict", None) - func_decl["parametersJsonSchema"] = schema - else: - func_decl["parametersJsonSchema"] = { - "type": "object", - "properties": {} - } - - gemini_tools.append({ - "functionDeclarations": [func_decl] - }) + """Count tokens for the given prompt using Antigravity :countTokens endpoint.""" + try: + token = await self.get_valid_token(credential_path) + internal_model = self._alias_to_internal(model) + + system_instruction, contents = self._transform_messages(messages, internal_model) + gemini_payload = {"contents": contents} + if system_instruction: + gemini_payload["systemInstruction"] = system_instruction + + gemini_tools = self._build_tools_payload(tools, model) if gemini_tools: - gemini_cli_payload["tools"] = gemini_tools - - # Wrap in Antigravity envelope - antigravity_payload = { - "project": self.generate_project_id(), - "userAgent": "antigravity", - "requestId": self.generate_request_id(), - "model": internal_model, - "request": gemini_cli_payload - } - - # Make the request - base_url = self._get_current_base_url() - url = f"{base_url}:countTokens" - - headers = { - "Authorization": f"Bearer {access_token}", - "Content-Type": "application/json" - } - - try: + gemini_payload["tools"] = gemini_tools + + antigravity_payload = { + "project": _generate_project_id(), + "userAgent": "antigravity", + "requestId": _generate_request_id(), + "model": internal_model, + "request": gemini_payload + } + + url = f"{self._get_base_url()}:countTokens" + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + response = await client.post(url, headers=headers, json=antigravity_payload, timeout=30) response.raise_for_status() - data = response.json() - # Unwrap Antigravity response - unwrapped = self._unwrap_antigravity_response(data) - - # Extract token counts from response - total_tokens = unwrapped.get('totalTokens', 0) + data = response.json() + unwrapped = self._unwrap_response(data) + total = unwrapped.get('totalTokens', 0) - return { - 'prompt_tokens': total_tokens, - 'total_tokens': total_tokens, - } - - except httpx.HTTPStatusError as e: - lib_logger.error(f"Failed to count tokens: {e}") - # Return 0 on error rather than raising - return {'prompt_tokens': 0, 'total_tokens': 0} + return {'prompt_tokens': total, 'total_tokens': total} + except Exception as e: + lib_logger.error(f"Token counting failed: {e}") + return {'prompt_tokens': 0, 'total_tokens': 0} \ No newline at end of file From 9bc26b913ef89fb9e12de2e12eb0323df622fada Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 05:13:16 +0100 Subject: [PATCH 021/221] =?UTF-8?q?refactor(providers):=20=F0=9F=94=A8=20e?= =?UTF-8?q?xtract=20cache=20logic=20into=20shared=20ProviderCache=20module?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extracted the AntigravityCache class into a new shared ProviderCache module to eliminate code duplication and improve maintainability across providers. - Created src/rotator_library/providers/provider_cache.py with generic, reusable cache implementation - Removed 266 lines of cache-specific code from antigravity_provider.py - Updated AntigravityProvider to use ProviderCache for both signature and thinking caches - Added configurable env_prefix parameter for flexible environment variable namespacing - Improved cache naming with _cache_name for better logging context - Added convenience factory function create_provider_cache() for streamlined cache creation - Removed unused imports (shutil, tempfile) from antigravity_provider.py - Updated .gitignore to include cache/ directory The new ProviderCache maintains full backward compatibility with the previous AntigravityCache implementation while providing a more modular, reusable foundation for other providers. --- .gitignore | 1 + .../providers/antigravity_provider.py | 281 +--------- .../providers/provider_cache.py | 498 ++++++++++++++++++ 3 files changed, 507 insertions(+), 273 deletions(-) create mode 100644 src/rotator_library/providers/provider_cache.py diff --git a/.gitignore b/.gitignore index 0d40840f..92bac087 100644 --- a/.gitignore +++ b/.gitignore @@ -126,3 +126,4 @@ staged_changes.txt launcher_config.json cache/antigravity/thought_signatures.json logs/ +cache/ diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 9223fdaa..1e332fcd 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -23,8 +23,6 @@ import logging import os import random -import shutil -import tempfile import time import uuid from datetime import datetime @@ -37,6 +35,7 @@ from .provider_interface import ProviderInterface from .antigravity_auth_base import AntigravityAuthBase +from .provider_cache import ProviderCache from ..model_definitions import ModelDefinitions @@ -269,272 +268,6 @@ def _append_text(self, filename: str, text: str) -> None: lib_logger.error(f"Failed to append to {filename}: {e}") -# ============================================================================= -# SIGNATURE CACHE -# ============================================================================= - -class AntigravityCache: - """ - Server-side cache for Antigravity conversation state preservation. - - Supports two types of cached data: - - Gemini 3: thoughtSignatures (tool_call_id → encrypted signature) - - Claude: Thinking content (composite_key → thinking text + signature) - - Features: - - Dual-TTL system: 1hr memory, 24hr disk - - Async disk persistence with batched writes - - Background cleanup task for expired entries - """ - - def __init__( - self, - cache_file: Path, - memory_ttl_seconds: int = 3600, - disk_ttl_seconds: int = 86400 - ): - # In-memory cache: {cache_key: (data, timestamp)} - self._cache: Dict[str, Tuple[str, float]] = {} - self._memory_ttl = memory_ttl_seconds - self._disk_ttl = disk_ttl_seconds - self._lock = asyncio.Lock() - self._disk_lock = asyncio.Lock() - - # Disk persistence - self._cache_file = cache_file - self._enable_disk = _env_bool("ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", True) - self._dirty = False - self._write_interval = _env_int("ANTIGRAVITY_CACHE_WRITE_INTERVAL", 60) - self._cleanup_interval = _env_int("ANTIGRAVITY_CACHE_CLEANUP_INTERVAL", 1800) - - # Background tasks - self._writer_task: Optional[asyncio.Task] = None - self._cleanup_task: Optional[asyncio.Task] = None - self._running = False - - # Statistics - self._stats = {"memory_hits": 0, "disk_hits": 0, "misses": 0, "writes": 0} - - if self._enable_disk: - lib_logger.debug( - f"AntigravityCache: Disk persistence enabled " - f"(memory_ttl={memory_ttl_seconds}s, disk_ttl={disk_ttl_seconds}s)" - ) - asyncio.create_task(self._async_init()) - else: - lib_logger.debug("AntigravityCache: Memory-only mode") - - async def _async_init(self) -> None: - """Async initialization: load from disk and start background tasks.""" - try: - await self._load_from_disk() - await self._start_background_tasks() - except Exception as e: - lib_logger.error(f"Cache async init failed: {e}") - - async def _load_from_disk(self) -> None: - """Load cache from disk file with TTL validation.""" - if not self._enable_disk or not self._cache_file.exists(): - return - - try: - async with self._disk_lock: - with open(self._cache_file, 'r', encoding='utf-8') as f: - data = json.load(f) - - if data.get("version") != "1.0": - lib_logger.warning("Cache version mismatch, starting fresh") - return - - now = time.time() - entries = data.get("entries", {}) - loaded = expired = 0 - - for call_id, entry in entries.items(): - age = now - entry.get("timestamp", 0) - if age <= self._disk_ttl: - sig = entry.get("signature", "") - if sig: - self._cache[call_id] = (sig, entry["timestamp"]) - loaded += 1 - else: - expired += 1 - - lib_logger.debug(f"Loaded {loaded} entries from disk ({expired} expired)") - except json.JSONDecodeError as e: - lib_logger.warning(f"Cache file corrupted: {e}") - except Exception as e: - lib_logger.error(f"Failed to load cache: {e}") - - async def _save_to_disk(self) -> None: - """Persist cache to disk using atomic write.""" - if not self._enable_disk: - return - - try: - async with self._disk_lock: - self._cache_file.parent.mkdir(parents=True, exist_ok=True) - - cache_data = { - "version": "1.0", - "memory_ttl_seconds": self._memory_ttl, - "disk_ttl_seconds": self._disk_ttl, - "entries": { - cid: {"signature": sig, "timestamp": ts} - for cid, (sig, ts) in self._cache.items() - }, - "statistics": { - "total_entries": len(self._cache), - "last_write": time.time(), - **self._stats - } - } - - # Atomic write - parent_dir = self._cache_file.parent - tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json') - - try: - with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f: - json.dump(cache_data, f, indent=2) - - try: - os.chmod(tmp_path, 0o600) - except (OSError, AttributeError): - pass - - shutil.move(tmp_path, self._cache_file) - self._stats["writes"] += 1 - lib_logger.debug(f"Saved {len(self._cache)} entries to disk") - except Exception: - if tmp_path and os.path.exists(tmp_path): - os.unlink(tmp_path) - raise - except Exception as e: - lib_logger.error(f"Disk save failed: {e}") - - async def _start_background_tasks(self) -> None: - """Start background writer and cleanup tasks.""" - if not self._enable_disk or self._running: - return - - self._running = True - self._writer_task = asyncio.create_task(self._writer_loop()) - self._cleanup_task = asyncio.create_task(self._cleanup_loop()) - lib_logger.debug("Started background cache tasks") - - async def _writer_loop(self) -> None: - """Background task: periodically flush dirty cache to disk.""" - try: - while self._running: - await asyncio.sleep(self._write_interval) - if self._dirty: - try: - await self._save_to_disk() - self._dirty = False - except Exception as e: - lib_logger.error(f"Background writer error: {e}") - except asyncio.CancelledError: - pass - - async def _cleanup_loop(self) -> None: - """Background task: periodically clean up expired entries.""" - try: - while self._running: - await asyncio.sleep(self._cleanup_interval) - await self._cleanup_expired() - except asyncio.CancelledError: - pass - - async def _cleanup_expired(self) -> None: - """Remove expired entries from memory cache.""" - async with self._lock: - now = time.time() - expired = [k for k, (_, ts) in self._cache.items() if now - ts > self._memory_ttl] - for k in expired: - del self._cache[k] - if expired: - self._dirty = True - lib_logger.debug(f"Cleaned up {len(expired)} expired entries") - - def store(self, key: str, value: str) -> None: - """Store a value (sync wrapper for async storage).""" - asyncio.create_task(self._async_store(key, value)) - - async def _async_store(self, key: str, value: str) -> None: - """Async implementation of store.""" - async with self._lock: - self._cache[key] = (value, time.time()) - self._dirty = True - - def retrieve(self, key: str) -> Optional[str]: - """Retrieve a value by key (sync method).""" - if key in self._cache: - value, timestamp = self._cache[key] - if time.time() - timestamp <= self._memory_ttl: - self._stats["memory_hits"] += 1 - return value - else: - del self._cache[key] - self._dirty = True - - self._stats["misses"] += 1 - if self._enable_disk: - asyncio.create_task(self._check_disk_fallback(key)) - return None - - async def _check_disk_fallback(self, key: str) -> None: - """Check disk for key and load into memory if found.""" - try: - if not self._cache_file.exists(): - return - - async with self._disk_lock: - with open(self._cache_file, 'r', encoding='utf-8') as f: - data = json.load(f) - - entries = data.get("entries", {}) - if key in entries: - entry = entries[key] - ts = entry.get("timestamp", 0) - if time.time() - ts <= self._disk_ttl: - sig = entry.get("signature", "") - if sig: - async with self._lock: - self._cache[key] = (sig, ts) - self._stats["disk_hits"] += 1 - lib_logger.debug(f"Loaded {key} from disk") - except Exception as e: - lib_logger.debug(f"Disk fallback failed: {e}") - - async def clear(self) -> None: - """Clear all cached data.""" - async with self._lock: - self._cache.clear() - self._dirty = True - if self._enable_disk: - await self._save_to_disk() - - async def shutdown(self) -> None: - """Graceful shutdown: flush pending writes and stop background tasks.""" - lib_logger.info("AntigravityCache shutting down...") - self._running = False - - for task in (self._writer_task, self._cleanup_task): - if task: - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - - if self._dirty and self._enable_disk: - await self._save_to_disk() - - lib_logger.info( - f"Cache shutdown complete (stats: mem_hits={self._stats['memory_hits']}, " - f"disk_hits={self._stats['disk_hits']}, misses={self._stats['misses']})" - ) # ============================================================================= @@ -571,12 +304,14 @@ def __init__(self): memory_ttl = _env_int("ANTIGRAVITY_SIGNATURE_CACHE_TTL", 3600) disk_ttl = _env_int("ANTIGRAVITY_SIGNATURE_DISK_TTL", 86400) - # Initialize caches - self._signature_cache = AntigravityCache( - GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl + # Initialize caches using shared ProviderCache + self._signature_cache = ProviderCache( + GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl, + env_prefix="ANTIGRAVITY_SIGNATURE" ) - self._thinking_cache = AntigravityCache( - CLAUDE_THINKING_CACHE_FILE, memory_ttl, disk_ttl + self._thinking_cache = ProviderCache( + CLAUDE_THINKING_CACHE_FILE, memory_ttl, disk_ttl, + env_prefix="ANTIGRAVITY_THINKING" ) # Feature flags diff --git a/src/rotator_library/providers/provider_cache.py b/src/rotator_library/providers/provider_cache.py new file mode 100644 index 00000000..b6bb2db6 --- /dev/null +++ b/src/rotator_library/providers/provider_cache.py @@ -0,0 +1,498 @@ +# src/rotator_library/providers/provider_cache.py +""" +Shared cache utility for providers. + +A modular, async-capable cache system supporting: +- Dual-TTL: short-lived memory cache, longer-lived disk persistence +- Background persistence with batched writes +- Automatic cleanup of expired entries +- Generic key-value storage for any provider-specific needs + +Usage examples: +- Gemini 3: thoughtSignatures (tool_call_id → encrypted signature) +- Claude: Thinking content (composite_key → thinking text + signature) +- General: Any transient data that benefits from persistence across requests +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import shutil +import tempfile +import time +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + +lib_logger = logging.getLogger('rotator_library') + + +# ============================================================================= +# UTILITY FUNCTIONS +# ============================================================================= + +def _env_bool(key: str, default: bool = False) -> bool: + """Get boolean from environment variable.""" + return os.getenv(key, str(default).lower()).lower() in ("true", "1", "yes") + + +def _env_int(key: str, default: int) -> int: + """Get integer from environment variable.""" + return int(os.getenv(key, str(default))) + + +# ============================================================================= +# PROVIDER CACHE CLASS +# ============================================================================= + +class ProviderCache: + """ + Server-side cache for provider conversation state preservation. + + A generic, modular cache supporting any key-value data that providers need + to persist across requests. Features: + + - Dual-TTL system: configurable memory TTL, longer disk TTL + - Async disk persistence with batched writes + - Background cleanup task for expired entries + - Statistics tracking (hits, misses, writes) + + Args: + cache_file: Path to disk cache file + memory_ttl_seconds: In-memory entry lifetime (default: 1 hour) + disk_ttl_seconds: Disk entry lifetime (default: 24 hours) + enable_disk: Whether to enable disk persistence (default: from env or True) + write_interval: Seconds between background disk writes (default: 60) + cleanup_interval: Seconds between expired entry cleanup (default: 30 min) + env_prefix: Environment variable prefix for configuration overrides + + Environment Variables (with default prefix "PROVIDER_CACHE"): + {PREFIX}_ENABLE: Enable/disable disk persistence + {PREFIX}_WRITE_INTERVAL: Background write interval in seconds + {PREFIX}_CLEANUP_INTERVAL: Cleanup interval in seconds + """ + + def __init__( + self, + cache_file: Path, + memory_ttl_seconds: int = 3600, + disk_ttl_seconds: int = 86400, + enable_disk: Optional[bool] = None, + write_interval: Optional[int] = None, + cleanup_interval: Optional[int] = None, + env_prefix: str = "PROVIDER_CACHE" + ): + # In-memory cache: {cache_key: (data, timestamp)} + self._cache: Dict[str, Tuple[str, float]] = {} + self._memory_ttl = memory_ttl_seconds + self._disk_ttl = disk_ttl_seconds + self._lock = asyncio.Lock() + self._disk_lock = asyncio.Lock() + + # Disk persistence configuration + self._cache_file = cache_file + self._enable_disk = enable_disk if enable_disk is not None else _env_bool(f"{env_prefix}_ENABLE", True) + self._dirty = False + self._write_interval = write_interval or _env_int(f"{env_prefix}_WRITE_INTERVAL", 60) + self._cleanup_interval = cleanup_interval or _env_int(f"{env_prefix}_CLEANUP_INTERVAL", 1800) + + # Background tasks + self._writer_task: Optional[asyncio.Task] = None + self._cleanup_task: Optional[asyncio.Task] = None + self._running = False + + # Statistics + self._stats = {"memory_hits": 0, "disk_hits": 0, "misses": 0, "writes": 0} + + # Metadata about this cache instance + self._cache_name = cache_file.stem if cache_file else "unnamed" + + if self._enable_disk: + lib_logger.debug( + f"ProviderCache[{self._cache_name}]: Disk enabled " + f"(memory_ttl={memory_ttl_seconds}s, disk_ttl={disk_ttl_seconds}s)" + ) + asyncio.create_task(self._async_init()) + else: + lib_logger.debug(f"ProviderCache[{self._cache_name}]: Memory-only mode") + + # ========================================================================= + # INITIALIZATION + # ========================================================================= + + async def _async_init(self) -> None: + """Async initialization: load from disk and start background tasks.""" + try: + await self._load_from_disk() + await self._start_background_tasks() + except Exception as e: + lib_logger.error(f"ProviderCache[{self._cache_name}] async init failed: {e}") + + async def _load_from_disk(self) -> None: + """Load cache from disk file with TTL validation.""" + if not self._enable_disk or not self._cache_file.exists(): + return + + try: + async with self._disk_lock: + with open(self._cache_file, 'r', encoding='utf-8') as f: + data = json.load(f) + + if data.get("version") != "1.0": + lib_logger.warning(f"ProviderCache[{self._cache_name}]: Version mismatch, starting fresh") + return + + now = time.time() + entries = data.get("entries", {}) + loaded = expired = 0 + + for cache_key, entry in entries.items(): + age = now - entry.get("timestamp", 0) + if age <= self._disk_ttl: + value = entry.get("value", entry.get("signature", "")) # Support both formats + if value: + self._cache[cache_key] = (value, entry["timestamp"]) + loaded += 1 + else: + expired += 1 + + lib_logger.debug( + f"ProviderCache[{self._cache_name}]: Loaded {loaded} entries ({expired} expired)" + ) + except json.JSONDecodeError as e: + lib_logger.warning(f"ProviderCache[{self._cache_name}]: File corrupted: {e}") + except Exception as e: + lib_logger.error(f"ProviderCache[{self._cache_name}]: Load failed: {e}") + + # ========================================================================= + # DISK PERSISTENCE + # ========================================================================= + + async def _save_to_disk(self) -> None: + """Persist cache to disk using atomic write.""" + if not self._enable_disk: + return + + try: + async with self._disk_lock: + self._cache_file.parent.mkdir(parents=True, exist_ok=True) + + cache_data = { + "version": "1.0", + "memory_ttl_seconds": self._memory_ttl, + "disk_ttl_seconds": self._disk_ttl, + "entries": { + key: {"value": val, "timestamp": ts} + for key, (val, ts) in self._cache.items() + }, + "statistics": { + "total_entries": len(self._cache), + "last_write": time.time(), + **self._stats + } + } + + # Atomic write using temp file + parent_dir = self._cache_file.parent + tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json') + + try: + with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f: + json.dump(cache_data, f, indent=2) + + # Set restrictive permissions (if supported) + try: + os.chmod(tmp_path, 0o600) + except (OSError, AttributeError): + pass + + shutil.move(tmp_path, self._cache_file) + self._stats["writes"] += 1 + lib_logger.debug( + f"ProviderCache[{self._cache_name}]: Saved {len(self._cache)} entries" + ) + except Exception: + if tmp_path and os.path.exists(tmp_path): + os.unlink(tmp_path) + raise + except Exception as e: + lib_logger.error(f"ProviderCache[{self._cache_name}]: Disk save failed: {e}") + + # ========================================================================= + # BACKGROUND TASKS + # ========================================================================= + + async def _start_background_tasks(self) -> None: + """Start background writer and cleanup tasks.""" + if not self._enable_disk or self._running: + return + + self._running = True + self._writer_task = asyncio.create_task(self._writer_loop()) + self._cleanup_task = asyncio.create_task(self._cleanup_loop()) + lib_logger.debug(f"ProviderCache[{self._cache_name}]: Started background tasks") + + async def _writer_loop(self) -> None: + """Background task: periodically flush dirty cache to disk.""" + try: + while self._running: + await asyncio.sleep(self._write_interval) + if self._dirty: + try: + await self._save_to_disk() + self._dirty = False + except Exception as e: + lib_logger.error(f"ProviderCache[{self._cache_name}]: Writer error: {e}") + except asyncio.CancelledError: + pass + + async def _cleanup_loop(self) -> None: + """Background task: periodically clean up expired entries.""" + try: + while self._running: + await asyncio.sleep(self._cleanup_interval) + await self._cleanup_expired() + except asyncio.CancelledError: + pass + + async def _cleanup_expired(self) -> None: + """Remove expired entries from memory cache.""" + async with self._lock: + now = time.time() + expired = [k for k, (_, ts) in self._cache.items() if now - ts > self._memory_ttl] + for k in expired: + del self._cache[k] + if expired: + self._dirty = True + lib_logger.debug( + f"ProviderCache[{self._cache_name}]: Cleaned {len(expired)} expired entries" + ) + + # ========================================================================= + # CORE OPERATIONS + # ========================================================================= + + def store(self, key: str, value: str) -> None: + """ + Store a value synchronously (schedules async storage). + + Args: + key: Cache key + value: Value to store (typically JSON-serialized data) + """ + asyncio.create_task(self._async_store(key, value)) + + async def _async_store(self, key: str, value: str) -> None: + """Async implementation of store.""" + async with self._lock: + self._cache[key] = (value, time.time()) + self._dirty = True + + async def store_async(self, key: str, value: str) -> None: + """ + Store a value asynchronously (awaitable). + + Use this when you need to ensure the value is stored before continuing. + """ + await self._async_store(key, value) + + def retrieve(self, key: str) -> Optional[str]: + """ + Retrieve a value by key (synchronous, with optional async disk fallback). + + Args: + key: Cache key + + Returns: + Cached value if found and not expired, None otherwise + """ + if key in self._cache: + value, timestamp = self._cache[key] + if time.time() - timestamp <= self._memory_ttl: + self._stats["memory_hits"] += 1 + return value + else: + del self._cache[key] + self._dirty = True + + self._stats["misses"] += 1 + if self._enable_disk: + # Schedule async disk lookup for next time + asyncio.create_task(self._check_disk_fallback(key)) + return None + + async def retrieve_async(self, key: str) -> Optional[str]: + """ + Retrieve a value asynchronously (checks disk if not in memory). + + Use this when you can await and need guaranteed disk fallback. + """ + # Check memory first + if key in self._cache: + value, timestamp = self._cache[key] + if time.time() - timestamp <= self._memory_ttl: + self._stats["memory_hits"] += 1 + return value + else: + async with self._lock: + if key in self._cache: + del self._cache[key] + self._dirty = True + + # Check disk + if self._enable_disk: + return await self._disk_retrieve(key) + + self._stats["misses"] += 1 + return None + + async def _check_disk_fallback(self, key: str) -> None: + """Check disk for key and load into memory if found (background).""" + try: + if not self._cache_file.exists(): + return + + async with self._disk_lock: + with open(self._cache_file, 'r', encoding='utf-8') as f: + data = json.load(f) + + entries = data.get("entries", {}) + if key in entries: + entry = entries[key] + ts = entry.get("timestamp", 0) + if time.time() - ts <= self._disk_ttl: + value = entry.get("value", entry.get("signature", "")) + if value: + async with self._lock: + self._cache[key] = (value, ts) + self._stats["disk_hits"] += 1 + lib_logger.debug( + f"ProviderCache[{self._cache_name}]: Loaded {key} from disk" + ) + except Exception as e: + lib_logger.debug(f"ProviderCache[{self._cache_name}]: Disk fallback failed: {e}") + + async def _disk_retrieve(self, key: str) -> Optional[str]: + """Direct disk retrieval with loading into memory.""" + try: + if not self._cache_file.exists(): + self._stats["misses"] += 1 + return None + + async with self._disk_lock: + with open(self._cache_file, 'r', encoding='utf-8') as f: + data = json.load(f) + + entries = data.get("entries", {}) + if key in entries: + entry = entries[key] + ts = entry.get("timestamp", 0) + if time.time() - ts <= self._disk_ttl: + value = entry.get("value", entry.get("signature", "")) + if value: + async with self._lock: + self._cache[key] = (value, ts) + self._stats["disk_hits"] += 1 + return value + + self._stats["misses"] += 1 + return None + except Exception as e: + lib_logger.debug(f"ProviderCache[{self._cache_name}]: Disk retrieve failed: {e}") + self._stats["misses"] += 1 + return None + + # ========================================================================= + # UTILITY METHODS + # ========================================================================= + + def contains(self, key: str) -> bool: + """Check if key exists in memory cache (without updating stats).""" + if key in self._cache: + _, timestamp = self._cache[key] + return time.time() - timestamp <= self._memory_ttl + return False + + def get_stats(self) -> Dict[str, Any]: + """Get cache statistics.""" + return { + **self._stats, + "memory_entries": len(self._cache), + "dirty": self._dirty, + "disk_enabled": self._enable_disk + } + + async def clear(self) -> None: + """Clear all cached data.""" + async with self._lock: + self._cache.clear() + self._dirty = True + if self._enable_disk: + await self._save_to_disk() + + async def shutdown(self) -> None: + """Graceful shutdown: flush pending writes and stop background tasks.""" + lib_logger.info(f"ProviderCache[{self._cache_name}]: Shutting down...") + self._running = False + + # Cancel background tasks + for task in (self._writer_task, self._cleanup_task): + if task: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Final save + if self._dirty and self._enable_disk: + await self._save_to_disk() + + lib_logger.info( + f"ProviderCache[{self._cache_name}]: Shutdown complete " + f"(stats: mem_hits={self._stats['memory_hits']}, " + f"disk_hits={self._stats['disk_hits']}, misses={self._stats['misses']})" + ) + + +# ============================================================================= +# CONVENIENCE FACTORY +# ============================================================================= + +def create_provider_cache( + name: str, + cache_dir: Optional[Path] = None, + memory_ttl_seconds: int = 3600, + disk_ttl_seconds: int = 86400, + env_prefix: Optional[str] = None +) -> ProviderCache: + """ + Factory function to create a provider cache with sensible defaults. + + Args: + name: Cache name (used as filename and for logging) + cache_dir: Directory for cache file (default: project_root/cache/provider_name) + memory_ttl_seconds: In-memory TTL + disk_ttl_seconds: Disk TTL + env_prefix: Environment variable prefix (default: derived from name) + + Returns: + Configured ProviderCache instance + """ + if cache_dir is None: + cache_dir = Path(__file__).resolve().parent.parent.parent.parent / "cache" + + cache_file = cache_dir / f"{name}.json" + + if env_prefix is None: + # Convert name to env prefix: "gemini3_signatures" -> "GEMINI3_SIGNATURES_CACHE" + env_prefix = f"{name.upper().replace('-', '_')}_CACHE" + + return ProviderCache( + cache_file=cache_file, + memory_ttl_seconds=memory_ttl_seconds, + disk_ttl_seconds=disk_ttl_seconds, + env_prefix=env_prefix + ) From e6a4ff2871d0cb37e3ef679302ea69813bd954c4 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 05:13:43 +0100 Subject: [PATCH 022/221] =?UTF-8?q?refactor(antigravity):=20=F0=9F=94=A8?= =?UTF-8?q?=20simplify=20Claude=20model=20variant=20handling=20with=20auto?= =?UTF-8?q?matic=20-thinking=20mapping?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit streamlines the handling of Claude Sonnet 4.5 model variants by automatically mapping the base model to its -thinking variant when reasoning_effort is provided. - Remove explicit "claude-sonnet-4-5-thinking" from AVAILABLE_MODELS list - Add inline documentation explaining internal mapping behavior - Implement automatic model variant selection in _transform_to_antigravity_format based on reasoning_effort parameter - Thread reasoning_effort parameter through generate_content call chain - Check for base claude-sonnet-4-5 model and append "-thinking" suffix when reasoning_effort is present This improves the API surface by reducing redundant model options while maintaining full functionality through intelligent runtime model selection. --- .../providers/antigravity_provider.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 1e332fcd..5aa68252 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -61,8 +61,7 @@ "gemini-3-pro-preview", "gemini-3-pro-image-preview", "gemini-2.5-computer-use-preview-10-2025", - "claude-sonnet-4-5", - "claude-sonnet-4-5-thinking", + "claude-sonnet-4-5", # Internally mapped to -thinking variant when reasoning_effort is provided ] # Default max output tokens (including thinking) - can be overridden per request @@ -885,7 +884,8 @@ def _transform_to_antigravity_format( self, gemini_payload: Dict[str, Any], model: str, - max_tokens: Optional[int] = None + max_tokens: Optional[int] = None, + reasoning_effort: Optional[str] = None ) -> Dict[str, Any]: """ Transform Gemini CLI payload to complete Antigravity format. @@ -894,9 +894,15 @@ def _transform_to_antigravity_format( gemini_payload: Request in Gemini CLI format model: Model name (public alias) max_tokens: Max output tokens (including thinking) + reasoning_effort: Reasoning effort level (determines -thinking variant for Claude) """ internal_model = self._alias_to_internal(model) + # Map base Claude model to -thinking variant when reasoning_effort is provided + if self._is_claude(internal_model) and reasoning_effort: + if internal_model == "claude-sonnet-4-5" and not internal_model.endswith("-thinking"): + internal_model = "claude-sonnet-4-5-thinking" + # Wrap in Antigravity envelope antigravity_payload = { "project": _generate_project_id(), @@ -1372,7 +1378,7 @@ async def acompletion( gemini_payload["tools"] = self._inject_signature_into_descriptions(gemini_payload["tools"]) # Transform to Antigravity format - payload = self._transform_to_antigravity_format(gemini_payload, model, max_tokens) + payload = self._transform_to_antigravity_format(gemini_payload, model, max_tokens, reasoning_effort) file_logger.log_request(payload) # Make API call From ae567625c3d8895c0a1729e459c8226e54a15041 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 05:31:02 +0100 Subject: [PATCH 023/221] =?UTF-8?q?feat(gemini):=20=E2=9C=A8=20implement?= =?UTF-8?q?=20Gemini=203=20support=20with=20tool=20fixes=20and=20signature?= =?UTF-8?q?=20caching?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit integrates comprehensive support for `gemini-3-pro-preview`, addressing specific requirements for reasoning models and tool reliability. - Update `AntigravityProvider` and `GeminiCliProvider` model lists to prioritize Gemini 3. - Implement a "Tool Fix" mechanism to prevent parameter hallucinations: - Inject strict parameter signatures and type hints into tool descriptions. - Add specific system instructions to enforce schema adherence. - Apply `gemini3_` namespace prefixing to isolate tool contexts. - Integrate `ProviderCache` to persist `thoughtSignature` values, ensuring reasoning continuity during tool execution. - Refactor `_handle_reasoning_parameters` to support Gemini 3's `thinkingLevel` (string) alongside Gemini 2.5's `thinkingBudget` (integer). - Add environment variable configuration for cache TTL and feature flags. --- .../providers/antigravity_provider.py | 10 +- .../providers/gemini_cli_provider.py | 354 ++++++++++++++++-- 2 files changed, 334 insertions(+), 30 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 5aa68252..dc13ae9d 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -55,12 +55,12 @@ # Available models via Antigravity AVAILABLE_MODELS = [ - "gemini-2.5-pro", - "gemini-2.5-flash", - "gemini-2.5-flash-lite", + #"gemini-2.5-pro", + #"gemini-2.5-flash", + #"gemini-2.5-flash-lite", "gemini-3-pro-preview", - "gemini-3-pro-image-preview", - "gemini-2.5-computer-use-preview-10-2025", + #"gemini-3-pro-image-preview", + #"gemini-2.5-computer-use-preview-10-2025", "claude-sonnet-4-5", # Internally mapped to -thinking variant when reasoning_effort is provided ] diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 47572fd6..8029e3d2 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -1,5 +1,6 @@ # src/rotator_library/providers/gemini_cli_provider.py +import copy import json import httpx import logging @@ -8,10 +9,10 @@ from typing import List, Dict, Any, AsyncGenerator, Union, Optional, Tuple from .provider_interface import ProviderInterface from .gemini_auth_base import GeminiAuthBase +from .provider_cache import ProviderCache from ..model_definitions import ModelDefinitions import litellm from litellm.exceptions import RateLimitError -from litellm.llms.vertex_ai.common_utils import _build_vertex_schema import os from pathlib import Path import uuid @@ -81,9 +82,49 @@ def log_final_response(self, response_data: Dict[str, Any]): HARDCODED_MODELS = [ "gemini-2.5-pro", "gemini-2.5-flash", - "gemini-2.5-flash-lite" + "gemini-2.5-flash-lite", + "gemini-3-pro-preview" ] +# Cache directory for Gemini CLI +CACHE_DIR = Path(__file__).resolve().parent.parent.parent.parent / "cache" / "gemini_cli" +GEMINI3_SIGNATURE_CACHE_FILE = CACHE_DIR / "gemini3_signatures.json" + +# Gemini 3 tool fix system instruction (prevents hallucination) +DEFAULT_GEMINI3_SYSTEM_INSTRUCTION = """CRITICAL TOOL USAGE INSTRUCTIONS: +You are operating in a custom environment where tool definitions differ from your training data. +You MUST follow these rules strictly: + +1. DO NOT use your internal training data to guess tool parameters +2. ONLY use the exact parameter structure defined in the tool schema +3. If a tool takes a 'files' parameter, it is ALWAYS an array of objects with specific properties, NEVER a simple array of strings +4. If a tool edits code, it takes structured JSON objects with specific fields, NEVER raw diff strings or plain text +5. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') +6. Array parameters have specific item types - check the schema's 'items' field for the exact structure +7. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions + +If you are unsure about a tool's parameters, YOU MUST read the schema definition carefully. Your training data about common tool names like 'read_file' or 'apply_diff' does NOT apply here. +""" + +# Gemini finish reason mapping +FINISH_REASON_MAP = { + "STOP": "stop", + "MAX_TOKENS": "length", + "SAFETY": "content_filter", + "RECITATION": "content_filter", + "OTHER": "stop", +} + + +def _env_bool(key: str, default: bool = False) -> bool: + """Get boolean from environment variable.""" + return os.getenv(key, str(default).lower()).lower() in ("true", "1", "yes") + + +def _env_int(key: str, default: int) -> int: + """Get integer from environment variable.""" + return int(os.getenv(key, str(default))) + class GeminiCliProvider(GeminiAuthBase, ProviderInterface): skip_cost_calculation = True @@ -92,6 +133,52 @@ def __init__(self): self.model_definitions = ModelDefinitions() self.project_id_cache: Dict[str, str] = {} # Cache project ID per credential path self.project_tier_cache: Dict[str, str] = {} # Cache project tier per credential path + + # Gemini 3 configuration from environment + memory_ttl = _env_int("GEMINI_CLI_SIGNATURE_CACHE_TTL", 3600) + disk_ttl = _env_int("GEMINI_CLI_SIGNATURE_DISK_TTL", 86400) + + # Initialize signature cache for Gemini 3 thoughtSignatures + self._signature_cache = ProviderCache( + GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl, + env_prefix="GEMINI_CLI_SIGNATURE" + ) + + # Gemini 3 feature flags + self._preserve_signatures_in_client = _env_bool("GEMINI_CLI_PRESERVE_THOUGHT_SIGNATURES", True) + self._enable_signature_cache = _env_bool("GEMINI_CLI_ENABLE_SIGNATURE_CACHE", True) + self._enable_gemini3_tool_fix = _env_bool("GEMINI_CLI_GEMINI3_TOOL_FIX", True) + + # Gemini 3 tool fix configuration + self._gemini3_tool_prefix = os.getenv("GEMINI_CLI_GEMINI3_TOOL_PREFIX", "gemini3_") + self._gemini3_description_prompt = os.getenv( + "GEMINI_CLI_GEMINI3_DESCRIPTION_PROMPT", + "\n\nSTRICT PARAMETERS: {params}." + ) + self._gemini3_system_instruction = os.getenv( + "GEMINI_CLI_GEMINI3_SYSTEM_INSTRUCTION", + DEFAULT_GEMINI3_SYSTEM_INSTRUCTION + ) + + lib_logger.debug( + f"GeminiCli config: signatures_in_client={self._preserve_signatures_in_client}, " + f"cache={self._enable_signature_cache}, gemini3_fix={self._enable_gemini3_tool_fix}" + ) + + # ========================================================================= + # MODEL UTILITIES + # ========================================================================= + + def _is_gemini_3(self, model: str) -> bool: + """Check if model is Gemini 3 (requires special handling).""" + model_name = model.split('/')[-1].replace(':thinking', '') + return model_name.startswith("gemini-3-") + + def _strip_gemini3_prefix(self, name: str) -> str: + """Strip the Gemini 3 namespace prefix from a tool name.""" + if name and name.startswith(self._gemini3_tool_prefix): + return name[len(self._gemini3_tool_prefix):] + return name async def _discover_project_id(self, credential_path: str, access_token: str, litellm_params: Dict[str, Any]) -> str: """ @@ -513,9 +600,20 @@ def _cli_preview_fallback_order(self, model: str) -> List[str]: # Return fallback chain if available, otherwise just return the original model return fallback_chains.get(model_name, [model_name]) - def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: + def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: + """ + Transform OpenAI messages to Gemini CLI format. + + Handles: + - System instruction extraction + - Multi-part content (text, images) + - Tool calls and responses + - Gemini 3 thoughtSignature preservation + """ + messages = copy.deepcopy(messages) # Don't mutate original system_instruction = None gemini_contents = [] + is_gemini_3 = self._is_gemini_3(model) # Separate system prompt from other messages if messages and messages[0].get('role') == 'system': @@ -580,15 +678,53 @@ def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[ args_dict = json.loads(tool_call["function"]["arguments"]) except (json.JSONDecodeError, TypeError): args_dict = {} - parts.append({"functionCall": {"name": tool_call["function"]["name"], "args": args_dict}}) + + tool_id = tool_call.get("id", "") + func_name = tool_call["function"]["name"] + + # Add prefix for Gemini 3 + if is_gemini_3 and self._enable_gemini3_tool_fix: + func_name = f"{self._gemini3_tool_prefix}{func_name}" + + func_part = { + "functionCall": { + "name": func_name, + "args": args_dict, + "id": tool_id + } + } + + # Add thoughtSignature for Gemini 3 + if is_gemini_3: + sig = tool_call.get("thought_signature") + if not sig and tool_id and self._enable_signature_cache: + sig = self._signature_cache.retrieve(tool_id) + + if sig: + func_part["thoughtSignature"] = sig + else: + func_part["thoughtSignature"] = "skip_thought_signature_validator" + lib_logger.warning(f"Missing thoughtSignature for {tool_id}, using bypass") + + parts.append(func_part) elif role == "tool": tool_call_id = msg.get("tool_call_id") function_name = tool_call_id_to_name.get(tool_call_id) if function_name: + # Add prefix for Gemini 3 + if is_gemini_3 and self._enable_gemini3_tool_fix: + function_name = f"{self._gemini3_tool_prefix}{function_name}" + # Wrap the tool response in a 'result' object response_content = {"result": content} - parts.append({"functionResponse": {"name": function_name, "response": response_content}}) + parts.append({ + "functionResponse": { + "name": function_name, + "response": response_content, + "id": tool_call_id + } + }) if parts: gemini_contents.append({"role": gemini_role, "parts": parts}) @@ -599,19 +735,42 @@ def _transform_messages(self, messages: List[Dict[str, Any]]) -> Tuple[Optional[ return system_instruction, gemini_contents def _handle_reasoning_parameters(self, payload: Dict[str, Any], model: str) -> Optional[Dict[str, Any]]: + """ + Map reasoning_effort to thinking configuration. + + - Gemini 2.5: thinkingBudget (integer tokens) + - Gemini 3: thinkingLevel (string: "low"/"high") + """ custom_reasoning_budget = payload.get("custom_reasoning_budget", False) reasoning_effort = payload.get("reasoning_effort") if "thinkingConfig" in payload.get("generationConfig", {}): return None - # Only apply reasoning logic to the gemini-2.5 model family - if "gemini-2.5" not in model: + is_gemini_25 = "gemini-2.5" in model + is_gemini_3 = self._is_gemini_3(model) + + # Only apply reasoning logic to supported models + if not (is_gemini_25 or is_gemini_3): payload.pop("reasoning_effort", None) payload.pop("custom_reasoning_budget", None) return None + + # Gemini 3: String-based thinkingLevel + if is_gemini_3: + # Clean up the original payload + payload.pop("reasoning_effort", None) + payload.pop("custom_reasoning_budget", None) + + if reasoning_effort == "low": + return {"thinkingLevel": "low", "include_thoughts": True} + return {"thinkingLevel": "high", "include_thoughts": True} + # Gemini 2.5: Integer thinkingBudget if not reasoning_effort: + # Clean up the original payload + payload.pop("reasoning_effort", None) + payload.pop("custom_reasoning_budget", None) return {"thinkingBudget": -1, "include_thoughts": True} # If reasoning_effort is provided, calculate the budget @@ -637,8 +796,15 @@ def _handle_reasoning_parameters(self, payload: Dict[str, Any], model: str) -> O return {"thinkingBudget": budget, "include_thoughts": True} - def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): - #lib_logger.debug(f"Converting Gemini chunk: {json.dumps(chunk)}") + def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumulator: Optional[Dict[str, Any]] = None): + """ + Convert Gemini response chunk to OpenAI streaming format. + + Args: + chunk: Gemini API response chunk + model_id: Model name + accumulator: Optional dict to accumulate data for post-processing (signatures, etc.) + """ response_data = chunk.get('response', chunk) candidates = response_data.get('candidates', []) if not candidates: @@ -646,17 +812,34 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): candidate = candidates[0] parts = candidate.get('content', {}).get('parts', []) + is_gemini_3 = self._is_gemini_3(model_id) + first_sig_seen = False for part in parts: delta = {} finish_reason = None + + has_func = 'functionCall' in part + has_text = 'text' in part + has_sig = bool(part.get('thoughtSignature')) + is_thought = part.get('thought') is True or (isinstance(part.get('thought'), str) and str(part.get('thought')).lower() == 'true') + + # Skip standalone signature parts (no function, no meaningful text) + if has_sig and not has_func and (not has_text or not part.get('text')): + continue - if 'functionCall' in part: + if has_func: function_call = part['functionCall'] function_name = function_call.get('name', 'unknown') - # Generate unique ID with nanosecond precision - tool_call_id = f"call_{function_name}_{int(time.time() * 1_000_000_000)}" - delta['tool_calls'] = [{ + + # Strip Gemini 3 prefix from tool name + if is_gemini_3 and self._enable_gemini3_tool_fix: + function_name = self._strip_gemini3_prefix(function_name) + + # Use provided ID or generate unique one with nanosecond precision + tool_call_id = function_call.get('id') or f"call_{function_name}_{int(time.time() * 1_000_000_000)}" + + tool_call = { "index": 0, "id": tool_call_id, "type": "function", @@ -664,11 +847,25 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): "name": function_name, "arguments": json.dumps(function_call.get('args', {})) } - }] - elif 'text' in part: + } + + # Handle thoughtSignature for Gemini 3 + if is_gemini_3 and has_sig and not first_sig_seen: + first_sig_seen = True + sig = part['thoughtSignature'] + + if self._enable_signature_cache: + self._signature_cache.store(tool_call_id, sig) + lib_logger.debug(f"Stored signature for {tool_call_id}") + + if self._preserve_signatures_in_client: + tool_call["thought_signature"] = sig + + delta['tool_calls'] = [tool_call] + + elif has_text: # Use an explicit check for the 'thought' flag, as its type can be inconsistent - thought = part.get('thought') - if thought is True or (isinstance(thought, str) and thought.lower() == 'true'): + if is_thought: delta['reasoning_content'] = part['text'] else: delta['content'] = part['text'] @@ -678,14 +875,16 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): raw_finish_reason = candidate.get('finishReason') if raw_finish_reason: - mapping = {'STOP': 'stop', 'MAX_TOKENS': 'length', 'SAFETY': 'content_filter'} - finish_reason = mapping.get(raw_finish_reason, 'stop') + finish_reason = FINISH_REASON_MAP.get(raw_finish_reason, 'stop') + # Use tool_calls if we have function calls + if delta.get('tool_calls'): + finish_reason = 'tool_calls' choice = {"index": 0, "delta": delta, "finish_reason": finish_reason} openai_chunk = { "choices": [choice], "model": model_id, "object": "chat.completion.chunk", - "id": f"chatcmpl-geminicli-{time.time()}", "created": int(time.time()) + "id": chunk.get("responseId", f"chatcmpl-geminicli-{time.time()}"), "created": int(time.time()) } if 'usageMetadata' in response_data: @@ -843,12 +1042,18 @@ def _gemini_cli_transform_schema(self, schema: Dict[str, Any]) -> Dict[str, Any] return schema - def _transform_tool_schemas(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + def _transform_tool_schemas(self, tools: List[Dict[str, Any]], model: str = "") -> List[Dict[str, Any]]: """ Transforms a list of OpenAI-style tool schemas into the format required by the Gemini CLI API. This uses a custom schema transformer instead of litellm's generic one. + + For Gemini 3 models, also applies: + - Namespace prefix to tool names + - Parameter signature injection into descriptions """ transformed_declarations = [] + is_gemini_3 = self._is_gemini_3(model) + for tool in tools: if tool.get("type") == "function" and "function" in tool: new_function = json.loads(json.dumps(tool["function"])) @@ -865,19 +1070,108 @@ def _transform_tool_schemas(self, tools: List[Dict[str, Any]]) -> List[Dict[str, # Set default empty schema if neither exists new_function["parametersJsonSchema"] = {"type": "object", "properties": {}} + # Gemini 3 specific transformations + if is_gemini_3 and self._enable_gemini3_tool_fix: + # Add namespace prefix to tool names + name = new_function.get("name", "") + if name: + new_function["name"] = f"{self._gemini3_tool_prefix}{name}" + + # Inject parameter signature into description + new_function = self._inject_signature_into_description(new_function) + transformed_declarations.append(new_function) return transformed_declarations - def _translate_tool_choice(self, tool_choice: Union[str, Dict[str, Any]]) -> Optional[Dict[str, Any]]: + def _inject_signature_into_description(self, func_decl: Dict[str, Any]) -> Dict[str, Any]: + """Inject parameter signatures into tool description for Gemini 3.""" + schema = func_decl.get("parametersJsonSchema", {}) + if not schema: + return func_decl + + required = schema.get("required", []) + properties = schema.get("properties", {}) + + if not properties: + return func_decl + + param_list = [] + for prop_name, prop_data in properties.items(): + if not isinstance(prop_data, dict): + continue + + type_hint = self._format_type_hint(prop_data) + is_required = prop_name in required + param_list.append( + f"{prop_name} ({type_hint}{', REQUIRED' if is_required else ''})" + ) + + if param_list: + sig_str = self._gemini3_description_prompt.replace( + "{params}", ", ".join(param_list) + ) + func_decl["description"] = func_decl.get("description", "") + sig_str + + return func_decl + + def _format_type_hint(self, prop_data: Dict[str, Any]) -> str: + """Format a type hint for a property schema.""" + type_hint = prop_data.get("type", "unknown") + + if type_hint == "array": + items = prop_data.get("items", {}) + if isinstance(items, dict): + item_type = items.get("type", "unknown") + if item_type == "object": + nested_props = items.get("properties", {}) + nested_req = items.get("required", []) + if nested_props: + nested_list = [] + for n, d in nested_props.items(): + if isinstance(d, dict): + t = d.get("type", "unknown") + req = " REQUIRED" if n in nested_req else "" + nested_list.append(f"{n}: {t}{req}") + return f"ARRAY_OF_OBJECTS[{', '.join(nested_list)}]" + return "ARRAY_OF_OBJECTS" + return f"ARRAY_OF_{item_type.upper()}" + return "ARRAY" + + return type_hint + + def _inject_gemini3_system_instruction(self, request_payload: Dict[str, Any]) -> None: + """Inject Gemini 3 tool fix system instruction if tools are present.""" + if not request_payload.get("request", {}).get("tools"): + return + + existing_system = request_payload.get("request", {}).get("systemInstruction") + + if existing_system: + # Prepend to existing system instruction + existing_parts = existing_system.get("parts", []) + if existing_parts and existing_parts[0].get("text"): + existing_parts[0]["text"] = self._gemini3_system_instruction + "\n\n" + existing_parts[0]["text"] + else: + existing_parts.insert(0, {"text": self._gemini3_system_instruction}) + else: + # Create new system instruction + request_payload["request"]["systemInstruction"] = { + "role": "user", + "parts": [{"text": self._gemini3_system_instruction}] + } + + def _translate_tool_choice(self, tool_choice: Union[str, Dict[str, Any]], model: str = "") -> Optional[Dict[str, Any]]: """ Translates OpenAI's `tool_choice` to Gemini's `toolConfig`. + Handles Gemini 3 namespace prefixes for specific tool selection. """ if not tool_choice: return None config = {} mode = "AUTO" # Default to auto + is_gemini_3 = self._is_gemini_3(model) if isinstance(tool_choice, str): if tool_choice == "auto": @@ -889,6 +1183,10 @@ def _translate_tool_choice(self, tool_choice: Union[str, Dict[str, Any]]) -> Opt elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function": function_name = tool_choice.get("function", {}).get("name") if function_name: + # Add Gemini 3 prefix if needed + if is_gemini_3 and self._enable_gemini3_tool_fix: + function_name = f"{self._gemini3_tool_prefix}{function_name}" + mode = "ANY" # Force a call, but only to this function config["functionCallingConfig"] = { "mode": mode, @@ -930,6 +1228,8 @@ async def do_call(attempt_model: str, is_fallback: bool = False): model_name=model_name, enabled=enable_request_logging ) + + is_gemini_3 = self._is_gemini_3(model_name) gen_config = { "maxOutputTokens": kwargs.get("max_tokens", 64000), # Increased default @@ -945,7 +1245,7 @@ async def do_call(attempt_model: str, is_fallback: bool = False): if thinking_config: gen_config["thinkingConfig"] = thinking_config - system_instruction, contents = self._transform_messages(kwargs.get("messages", [])) + system_instruction, contents = self._transform_messages(kwargs.get("messages", []), model_name) request_payload = { "model": model_name, "project": project_id, @@ -959,15 +1259,19 @@ async def do_call(attempt_model: str, is_fallback: bool = False): request_payload["request"]["systemInstruction"] = system_instruction if "tools" in kwargs and kwargs["tools"]: - function_declarations = self._transform_tool_schemas(kwargs["tools"]) + function_declarations = self._transform_tool_schemas(kwargs["tools"], model_name) if function_declarations: request_payload["request"]["tools"] = [{"functionDeclarations": function_declarations}] # [NEW] Handle tool_choice translation if "tool_choice" in kwargs and kwargs["tool_choice"]: - tool_config = self._translate_tool_choice(kwargs["tool_choice"]) + tool_config = self._translate_tool_choice(kwargs["tool_choice"], model_name) if tool_config: request_payload["request"]["toolConfig"] = tool_config + + # Inject Gemini 3 system instruction if using tools + if is_gemini_3 and self._enable_gemini3_tool_fix: + self._inject_gemini3_system_instruction(request_payload) # Add default safety settings to prevent content filtering if "safetySettings" not in request_payload["request"]: From 3298177f073d142ec026b5afc070463ba84a889d Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 05:32:37 +0100 Subject: [PATCH 024/221] =?UTF-8?q?refactor(gemini):=20=F0=9F=94=A8=20remo?= =?UTF-8?q?ve=20redundant=20model=20and=20project=20fields=20from=20reques?= =?UTF-8?q?t=20payload?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `model` and `project` parameters were being incorrectly included at the top level of the request payload. These fields are not part of the Gemini API request body structure and should only be used for endpoint construction or authentication context. --- src/rotator_library/providers/gemini_cli_provider.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 8029e3d2..52c7daf8 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -1438,8 +1438,6 @@ async def count_tokens( # Build request payload request_payload = { - "model": model_name, - "project": project_id, "request": { "contents": contents, }, From 868b7c9b6436ae4db75f82dd9ada03af1e22d4e2 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 06:44:14 +0100 Subject: [PATCH 025/221] =?UTF-8?q?refactor(logging):=20=F0=9F=94=A8=20adj?= =?UTF-8?q?ust=20logging=20levels=20and=20improve=20schema=20cleaning=20fo?= =?UTF-8?q?r=20Antigravity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Change reasoning parameters log from info to debug level in main.py - Move reasoning parameters logging outside logger conditional block for consistent monitoring - Enhance _clean_claude_schema documentation to clarify it's for Antigravity/Google's Proto-based API - Add support for converting 'const' to 'enum' with single value in schema cleaning - Improve code organization with better comments explaining unsupported fields These changes improve logging granularity and enhance JSON Schema compatibility with Antigravity's Proto-based API requirements. --- src/proxy_app/main.py | 18 +++++++++--------- .../providers/antigravity_provider.py | 18 +++++++++++++++--- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 8903b688..71bc4ee4 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -672,15 +672,15 @@ async def chat_completions( if logger: logger.log_request(headers=request.headers, body=request_data) - # Extract and log specific reasoning parameters for monitoring. - model = request_data.get("model") - generation_cfg = request_data.get("generationConfig", {}) or request_data.get("generation_config", {}) or {} - reasoning_effort = request_data.get("reasoning_effort") or generation_cfg.get("reasoning_effort") - custom_reasoning_budget = request_data.get("custom_reasoning_budget") or generation_cfg.get("custom_reasoning_budget", False) - - logging.getLogger("rotator_library").info( - f"Handling reasoning parameters: model={model}, reasoning_effort={reasoning_effort}, custom_reasoning_budget={custom_reasoning_budget}" - ) + # Extract and log specific reasoning parameters for monitoring. + model = request_data.get("model") + generation_cfg = request_data.get("generationConfig", {}) or request_data.get("generation_config", {}) or {} + reasoning_effort = request_data.get("reasoning_effort") or generation_cfg.get("reasoning_effort") + custom_reasoning_budget = request_data.get("custom_reasoning_budget") or generation_cfg.get("custom_reasoning_budget", False) + + logging.getLogger("rotator_library").debug( + f"Handling reasoning parameters: model={model}, reasoning_effort={reasoning_effort}, custom_reasoning_budget={custom_reasoning_budget}" + ) # Log basic request info to console (this is a separate, simpler logger). log_request_to_console( diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index dc13ae9d..5b1e6ae8 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -186,15 +186,27 @@ def _recursively_parse_json_strings(obj: Any) -> Any: def _clean_claude_schema(schema: Any) -> Any: - """Recursively remove fields that Claude's JSON Schema validation doesn't support.""" + """ + Recursively clean JSON Schema for Antigravity/Google's Proto-based API. + - Removes unsupported fields ($schema, additionalProperties, etc.) + - Converts 'const' to 'enum' with single value (supported equivalent) + """ if not isinstance(schema, dict): return schema - incompatible = {'$schema', 'additionalProperties', 'minItems', 'maxItems', 'pattern'} + # Fields not supported by Antigravity/Google's Proto-based API + incompatible = { + '$schema', 'additionalProperties', 'minItems', 'maxItems', 'pattern', + } cleaned = {} + # Handle 'const' by converting to 'enum' with single value + if 'const' in schema: + const_value = schema['const'] + cleaned['enum'] = [const_value] + for key, value in schema.items(): - if key in incompatible: + if key in incompatible or key == 'const': continue if isinstance(value, dict): cleaned[key] = _clean_claude_schema(value) From 74f9532797d51c2853341ce3924f245c3a46f8b7 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 06:47:56 +0100 Subject: [PATCH 026/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20add=20?= =?UTF-8?q?thinking=20mode=20toggling=20for=20mid-conversation=20model=20s?= =?UTF-8?q?witches?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces intelligent handling of Claude's thinking mode when switching models mid-conversation during incomplete tool use loops. **New Features:** - Auto-detection of incomplete tool turns (when messages end with tool results without assistant completion) - Configurable turn completion injection via `ANTIGRAVITY_AUTO_INJECT_TURN_COMPLETION` (default: true) - Configurable thinking mode suppression via `ANTIGRAVITY_AUTO_SUPPRESS_THINKING` (default: false) - Customizable turn completion placeholder text via `ANTIGRAVITY_TURN_COMPLETION_TEXT` (default: "...") **Implementation Details:** - `_detect_incomplete_tool_turn()`: Analyzes message history to identify incomplete tool use patterns - `_inject_turn_completion()`: Appends a synthetic assistant message to close incomplete turns - `_handle_thinking_mode_toggle()`: Orchestrates the toggling strategy based on configuration **Behavior:** When switching to Claude with thinking mode enabled during an incomplete tool loop: 1. If auto-injection is enabled: Inject a completion message to allow thinking mode 2. If auto-suppression is enabled: Disable thinking mode to prevent API errors 3. If both disabled: Allow the request to proceed (likely resulting in API error) This resolves API compatibility issues when transitioning between models with different conversation state requirements. --- .../providers/antigravity_provider.py | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 5b1e6ae8..d5cce1e8 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -331,6 +331,11 @@ def __init__(self): self._enable_dynamic_models = _env_bool("ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", False) self._enable_gemini3_tool_fix = _env_bool("ANTIGRAVITY_GEMINI3_TOOL_FIX", True) + # Thinking mode toggling behavior + self._auto_inject_turn_completion = _env_bool("ANTIGRAVITY_AUTO_INJECT_TURN_COMPLETION", True) + self._auto_suppress_thinking = _env_bool("ANTIGRAVITY_AUTO_SUPPRESS_THINKING", False) + self._turn_completion_placeholder = os.getenv("ANTIGRAVITY_TURN_COMPLETION_TEXT", "...") + # Gemini 3 tool fix configuration self._gemini3_tool_prefix = os.getenv("ANTIGRAVITY_GEMINI3_TOOL_PREFIX", "gemini3_") self._gemini3_description_prompt = os.getenv( @@ -1324,6 +1329,142 @@ async def get_models( return [f"antigravity/{m}" for m in AVAILABLE_MODELS] + # ========================================================================= + # THINKING MODE TOGGLING HELPERS + # ========================================================================= + + def _detect_incomplete_tool_turn(self, messages: List[Dict[str, Any]]) -> Optional[int]: + """ + Detect if messages end with an incomplete tool use loop. + + An incomplete tool turn is when: + - Last message is a tool result + - The assistant message that made the tool call hasn't been completed + with a final text response + + Returns: + Index of the assistant message with tool_calls if incomplete turn detected, + None otherwise + """ + if len(messages) < 2: + return None + + # Last message must be tool result + if messages[-1].get("role") != "tool": + return None + + # Find the assistant message that made the tool call + for i in range(len(messages) - 2, -1, -1): + msg = messages[i] + if msg.get("role") == "assistant": + if msg.get("tool_calls"): + # Check if turn was completed by a subsequent assistant message + for j in range(i + 1, len(messages)): + if messages[j].get("role") == "assistant" and not messages[j].get("tool_calls"): + return None # Turn completed + + # Incomplete turn found + lib_logger.debug( + f"Detected incomplete tool turn: assistant message at index {i} " + f"has tool_calls, but no completing text response found" + ) + return i + else: + # Found completing assistant message + return None + + return None + + def _inject_turn_completion( + self, + messages: List[Dict[str, Any]], + incomplete_turn_index: int + ) -> List[Dict[str, Any]]: + """ + Inject a completing assistant message to close an incomplete tool use turn. + + Args: + messages: Original message list + incomplete_turn_index: Index of the assistant message with tool_calls + + Returns: + Modified message list with injected completion + """ + completion_msg = { + "role": "assistant", + "content": self._turn_completion_placeholder + } + + # Append to close the turn + modified_messages = messages.copy() + modified_messages.append(completion_msg) + + lib_logger.info( + f"Injected turn-completing assistant message ('{self._turn_completion_placeholder}') " + f"to enable thinking mode. Original tool use started at message index {incomplete_turn_index}." + ) + + return modified_messages + + def _handle_thinking_mode_toggle( + self, + messages: List[Dict[str, Any]], + model: str, + reasoning_effort: Optional[str] + ) -> Tuple[List[Dict[str, Any]], Optional[str]]: + """ + Handle thinking mode toggling when switching models mid-conversation. + + When switching to Claude with thinking enabled, but the conversation has + an incomplete tool use loop from another model, either: + 1. Inject a completing message to close the turn (if auto_inject enabled) + 2. Suppress thinking mode (if auto_suppress enabled) + 3. Let it fail with API error (if both disabled) + + Args: + messages: Original message list + model: Target model + reasoning_effort: Requested reasoning effort level + + Returns: + (modified_messages, modified_reasoning_effort) + """ + # Only applies when trying to enable thinking on Claude + if not self._is_claude(model) or not reasoning_effort: + return messages, reasoning_effort + + incomplete_turn_index = self._detect_incomplete_tool_turn(messages) + if incomplete_turn_index is None: + # No incomplete turn - proceed normally + return messages, reasoning_effort + + # Strategy 1: Auto-inject turn completion (preferred) + if self._auto_inject_turn_completion: + lib_logger.info( + "Model switch to Claude with thinking detected mid-tool-use-loop. " + "Injecting turn completion to enable thinking mode." + ) + modified_messages = self._inject_turn_completion(messages, incomplete_turn_index) + return modified_messages, reasoning_effort + + # Strategy 2: Auto-suppress thinking + if self._auto_suppress_thinking: + lib_logger.warning( + f"Model switch to Claude with thinking detected mid-tool-use-loop. " + f"Suppressing reasoning_effort={reasoning_effort} to avoid API error. " + f"Set ANTIGRAVITY_AUTO_INJECT_TURN_COMPLETION=true to inject completion instead." + ) + return messages, None + + # Strategy 3: Let it fail (user wants to handle it themselves) + lib_logger.warning( + "Model switch to Claude with thinking detected mid-tool-use-loop. " + "Both auto-injection and auto-suppression are disabled. " + "Request will likely fail with API error. " + f"Enable ANTIGRAVITY_AUTO_INJECT_TURN_COMPLETION or ANTIGRAVITY_AUTO_SUPPRESS_THINKING." + ) + return messages, reasoning_effort + async def acompletion( self, client: httpx.AsyncClient, @@ -1353,6 +1494,13 @@ async def acompletion( # Create logger file_logger = AntigravityFileLogger(model, enable_logging) + # Handle thinking mode toggling for model switches + messages, reasoning_effort = self._handle_thinking_mode_toggle(messages, model, reasoning_effort) + if reasoning_effort != kwargs.get("reasoning_effort"): + kwargs["reasoning_effort"] = reasoning_effort + if messages != kwargs.get("messages"): + kwargs["messages"] = messages + # Transform messages system_instruction, gemini_contents = self._transform_messages(messages, model) gemini_contents = self._fix_tool_response_grouping(gemini_contents) From 0ea3b2d65e5a808014136e14d5d88e634ba67d26 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 07:06:20 +0100 Subject: [PATCH 027/221] =?UTF-8?q?fix(proxy):=20=F0=9F=90=9B=20prevent=20?= =?UTF-8?q?role=20field=20concatenation=20in=20streaming=20responses?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The generic key handling logic was incorrectly concatenating the 'role' field when processing streaming message chunks. The role field should always be replaced with the latest value, not concatenated like content fields. This fix adds an explicit check to ensure the 'role' key is always overwritten rather than appended to, preventing malformed role values in the final message object. --- src/proxy_app/main.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 71bc4ee4..b5cacd31 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -589,7 +589,10 @@ async def streaming_response_wrapper( final_message["function_call"]["arguments"] += value["arguments"] else: # Generic key handling for other data like 'reasoning' - if key not in final_message: + # FIX: Role should always replace, never concatenate + if key == "role": + final_message[key] = value + elif key not in final_message: final_message[key] = value elif isinstance(final_message.get(key), str): final_message[key] += value From 4d4a19844dd4b883da068d3882a2505d242aa8b4 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 07:22:03 +0100 Subject: [PATCH 028/221] =?UTF-8?q?fix(antigravity):=20=F0=9F=90=9B=20hand?= =?UTF-8?q?le=20malformed=20double-encoded=20JSON=20responses?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Antigravity sometimes returns malformed JSON strings with extra trailing characters (e.g., '[{...}]}' instead of '[{...}]'). This enhancement extends the JSON parsing logic to automatically detect and correct such malformations by: - Detecting JSON-like strings that don't have proper closing delimiters - Finding the last valid closing bracket/brace and truncating extra characters - Logging warnings when auto-correction is applied for debugging purposes - Recursively parsing the corrected JSON structures This prevents parsing failures when Antigravity returns double-encoded or malformed JSON in tool arguments. --- .../providers/antigravity_provider.py | 53 ++++++++++++++++--- 1 file changed, 46 insertions(+), 7 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index d5cce1e8..d9164c00 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -168,6 +168,9 @@ def _recursively_parse_json_strings(obj: Any) -> Any: Antigravity sometimes returns tool arguments with JSON-stringified values: {"files": "[{...}]"} instead of {"files": [{...}]}. + + Additionally handles malformed double-encoded JSON where Antigravity + returns strings like '[{...}]}' (extra trailing '}'). """ if isinstance(obj, dict): return {k: _recursively_parse_json_strings(v) for k, v in obj.items()} @@ -175,13 +178,49 @@ def _recursively_parse_json_strings(obj: Any) -> Any: return [_recursively_parse_json_strings(item) for item in obj] elif isinstance(obj, str): stripped = obj.strip() - if (stripped.startswith('{') and stripped.endswith('}')) or \ - (stripped.startswith('[') and stripped.endswith(']')): - try: - parsed = json.loads(obj) - return _recursively_parse_json_strings(parsed) - except (json.JSONDecodeError, ValueError): - pass + # Check if it looks like JSON (starts with { or [) + if stripped and stripped[0] in ('{', '['): + # Try standard parsing first + if (stripped.startswith('{') and stripped.endswith('}')) or \ + (stripped.startswith('[') and stripped.endswith(']')): + try: + parsed = json.loads(obj) + return _recursively_parse_json_strings(parsed) + except (json.JSONDecodeError, ValueError): + pass + + # Handle malformed JSON: array that doesn't end with ] + # e.g., '[{"path": "..."}]}' instead of '[{"path": "..."}]' + if stripped.startswith('[') and not stripped.endswith(']'): + try: + # Find the last ] and truncate there + last_bracket = stripped.rfind(']') + if last_bracket > 0: + cleaned = stripped[:last_bracket+1] + parsed = json.loads(cleaned) + lib_logger.warning( + f"Auto-corrected malformed JSON string: " + f"truncated {len(stripped) - len(cleaned)} extra chars" + ) + return _recursively_parse_json_strings(parsed) + except (json.JSONDecodeError, ValueError): + pass + + # Handle malformed JSON: object that doesn't end with } + if stripped.startswith('{') and not stripped.endswith('}'): + try: + # Find the last } and truncate there + last_brace = stripped.rfind('}') + if last_brace > 0: + cleaned = stripped[:last_brace+1] + parsed = json.loads(cleaned) + lib_logger.warning( + f"Auto-corrected malformed JSON string: " + f"truncated {len(stripped) - len(cleaned)} extra chars" + ) + return _recursively_parse_json_strings(parsed) + except (json.JSONDecodeError, ValueError): + pass return obj From 8d69bcd58adac8437f93a16f7ccc877cb339ea5f Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 07:40:26 +0100 Subject: [PATCH 029/221] =?UTF-8?q?fix(client):=20=F0=9F=90=9B=20prevent?= =?UTF-8?q?=20provider=20initialization=20without=20configured=20credentia?= =?UTF-8?q?ls?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `_get_provider_instance` method now checks if credentials exist for a provider before attempting initialization. This prevents potential errors from initializing providers that lack proper configuration. - Added credential existence check at the start of the method - Returns `None` early if provider credentials are not configured - Added debug logging to indicate when provider initialization is skipped - Enhanced docstring with detailed Args and Returns documentation This change improves system robustness by failing gracefully when providers are referenced but not properly configured. --- src/rotator_library/client.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 83a285f6..0cb65786 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -393,7 +393,23 @@ def _is_custom_openai_compatible_provider(self, provider_name: str) -> bool: return os.getenv(api_base_env) is not None def _get_provider_instance(self, provider_name: str): - """Lazily initializes and returns a provider instance.""" + """ + Lazily initializes and returns a provider instance. + Only initializes providers that have configured credentials. + + Args: + provider_name: The name of the provider to get an instance for. + + Returns: + Provider instance if credentials exist, None otherwise. + """ + # Only initialize providers for which we have credentials + if provider_name not in self.all_credentials: + lib_logger.debug( + f"Skipping provider '{provider_name}' initialization: no credentials configured" + ) + return None + if provider_name not in self._provider_instances: if provider_name in self._provider_plugins: self._provider_instances[provider_name] = self._provider_plugins[ From 8a839ed0cf91b9fd409c6cad2cbc2872012a726f Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 08:35:58 +0100 Subject: [PATCH 030/221] =?UTF-8?q?refactor(antigravity):=20=F0=9F=94=A8?= =?UTF-8?q?=20remove=20thinking=20mode=20toggling=20feature?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit removes the thinking mode toggling functionality that was previously used to handle model switches mid-conversation when tool use loops were incomplete. - Removed `_detect_incomplete_tool_turn`, `_inject_turn_completion`, and `_handle_thinking_mode_toggle` helper methods - Removed environment variable configuration for turn completion behavior (`ANTIGRAVITY_AUTO_INJECT_TURN_COMPLETION`, `ANTIGRAVITY_AUTO_SUPPRESS_THINKING`, `ANTIGRAVITY_TURN_COMPLETION_TEXT`) - Removed thinking mode toggle logic from `acompletion` method - Added provider prefix to JSON auto-correction warning log for better debugging The removed feature was designed to automatically handle incomplete tool use loops when switching to Claude models with thinking mode enabled, but was buggy as hell. --- .../providers/antigravity_provider.py | 150 +----------------- 1 file changed, 1 insertion(+), 149 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index d9164c00..0fa11faa 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -199,7 +199,7 @@ def _recursively_parse_json_strings(obj: Any) -> Any: cleaned = stripped[:last_bracket+1] parsed = json.loads(cleaned) lib_logger.warning( - f"Auto-corrected malformed JSON string: " + f"[Antigravity] Auto-corrected malformed JSON string: " f"truncated {len(stripped) - len(cleaned)} extra chars" ) return _recursively_parse_json_strings(parsed) @@ -370,11 +370,6 @@ def __init__(self): self._enable_dynamic_models = _env_bool("ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", False) self._enable_gemini3_tool_fix = _env_bool("ANTIGRAVITY_GEMINI3_TOOL_FIX", True) - # Thinking mode toggling behavior - self._auto_inject_turn_completion = _env_bool("ANTIGRAVITY_AUTO_INJECT_TURN_COMPLETION", True) - self._auto_suppress_thinking = _env_bool("ANTIGRAVITY_AUTO_SUPPRESS_THINKING", False) - self._turn_completion_placeholder = os.getenv("ANTIGRAVITY_TURN_COMPLETION_TEXT", "...") - # Gemini 3 tool fix configuration self._gemini3_tool_prefix = os.getenv("ANTIGRAVITY_GEMINI3_TOOL_PREFIX", "gemini3_") self._gemini3_description_prompt = os.getenv( @@ -1368,142 +1363,6 @@ async def get_models( return [f"antigravity/{m}" for m in AVAILABLE_MODELS] - # ========================================================================= - # THINKING MODE TOGGLING HELPERS - # ========================================================================= - - def _detect_incomplete_tool_turn(self, messages: List[Dict[str, Any]]) -> Optional[int]: - """ - Detect if messages end with an incomplete tool use loop. - - An incomplete tool turn is when: - - Last message is a tool result - - The assistant message that made the tool call hasn't been completed - with a final text response - - Returns: - Index of the assistant message with tool_calls if incomplete turn detected, - None otherwise - """ - if len(messages) < 2: - return None - - # Last message must be tool result - if messages[-1].get("role") != "tool": - return None - - # Find the assistant message that made the tool call - for i in range(len(messages) - 2, -1, -1): - msg = messages[i] - if msg.get("role") == "assistant": - if msg.get("tool_calls"): - # Check if turn was completed by a subsequent assistant message - for j in range(i + 1, len(messages)): - if messages[j].get("role") == "assistant" and not messages[j].get("tool_calls"): - return None # Turn completed - - # Incomplete turn found - lib_logger.debug( - f"Detected incomplete tool turn: assistant message at index {i} " - f"has tool_calls, but no completing text response found" - ) - return i - else: - # Found completing assistant message - return None - - return None - - def _inject_turn_completion( - self, - messages: List[Dict[str, Any]], - incomplete_turn_index: int - ) -> List[Dict[str, Any]]: - """ - Inject a completing assistant message to close an incomplete tool use turn. - - Args: - messages: Original message list - incomplete_turn_index: Index of the assistant message with tool_calls - - Returns: - Modified message list with injected completion - """ - completion_msg = { - "role": "assistant", - "content": self._turn_completion_placeholder - } - - # Append to close the turn - modified_messages = messages.copy() - modified_messages.append(completion_msg) - - lib_logger.info( - f"Injected turn-completing assistant message ('{self._turn_completion_placeholder}') " - f"to enable thinking mode. Original tool use started at message index {incomplete_turn_index}." - ) - - return modified_messages - - def _handle_thinking_mode_toggle( - self, - messages: List[Dict[str, Any]], - model: str, - reasoning_effort: Optional[str] - ) -> Tuple[List[Dict[str, Any]], Optional[str]]: - """ - Handle thinking mode toggling when switching models mid-conversation. - - When switching to Claude with thinking enabled, but the conversation has - an incomplete tool use loop from another model, either: - 1. Inject a completing message to close the turn (if auto_inject enabled) - 2. Suppress thinking mode (if auto_suppress enabled) - 3. Let it fail with API error (if both disabled) - - Args: - messages: Original message list - model: Target model - reasoning_effort: Requested reasoning effort level - - Returns: - (modified_messages, modified_reasoning_effort) - """ - # Only applies when trying to enable thinking on Claude - if not self._is_claude(model) or not reasoning_effort: - return messages, reasoning_effort - - incomplete_turn_index = self._detect_incomplete_tool_turn(messages) - if incomplete_turn_index is None: - # No incomplete turn - proceed normally - return messages, reasoning_effort - - # Strategy 1: Auto-inject turn completion (preferred) - if self._auto_inject_turn_completion: - lib_logger.info( - "Model switch to Claude with thinking detected mid-tool-use-loop. " - "Injecting turn completion to enable thinking mode." - ) - modified_messages = self._inject_turn_completion(messages, incomplete_turn_index) - return modified_messages, reasoning_effort - - # Strategy 2: Auto-suppress thinking - if self._auto_suppress_thinking: - lib_logger.warning( - f"Model switch to Claude with thinking detected mid-tool-use-loop. " - f"Suppressing reasoning_effort={reasoning_effort} to avoid API error. " - f"Set ANTIGRAVITY_AUTO_INJECT_TURN_COMPLETION=true to inject completion instead." - ) - return messages, None - - # Strategy 3: Let it fail (user wants to handle it themselves) - lib_logger.warning( - "Model switch to Claude with thinking detected mid-tool-use-loop. " - "Both auto-injection and auto-suppression are disabled. " - "Request will likely fail with API error. " - f"Enable ANTIGRAVITY_AUTO_INJECT_TURN_COMPLETION or ANTIGRAVITY_AUTO_SUPPRESS_THINKING." - ) - return messages, reasoning_effort - async def acompletion( self, client: httpx.AsyncClient, @@ -1533,13 +1392,6 @@ async def acompletion( # Create logger file_logger = AntigravityFileLogger(model, enable_logging) - # Handle thinking mode toggling for model switches - messages, reasoning_effort = self._handle_thinking_mode_toggle(messages, model, reasoning_effort) - if reasoning_effort != kwargs.get("reasoning_effort"): - kwargs["reasoning_effort"] = reasoning_effort - if messages != kwargs.get("messages"): - kwargs["messages"] = messages - # Transform messages system_instruction, gemini_contents = self._transform_messages(messages, model) gemini_contents = self._fix_tool_response_grouping(gemini_contents) From b5da45c8bb539cd7bbea124a86d288cb0039c7f2 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 09:13:09 +0100 Subject: [PATCH 031/221] =?UTF-8?q?feat(client):=20=E2=9C=A8=20add=20crede?= =?UTF-8?q?ntial=20prioritization=20system=20for=20tier-based=20model=20ac?= =?UTF-8?q?cess?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements a comprehensive credential prioritization system that enables providers to enforce tier-based access controls and optimize credential selection based on account types. Key changes: - Added `get_credential_priority()` and `get_model_tier_requirement()` methods to ProviderInterface, allowing providers to define credential tiers and model restrictions - Enhanced UsageManager.acquire_key() to respect credential priorities, always attempting highest-priority credentials first before falling back to lower tiers - Implemented Gemini-specific tier detection in GeminiCliProvider, mapping paid tier credentials to priority 1, free tier to priority 2, and unknown to priority 10 - Added model-based filtering in RotatingClient to exclude incompatible credentials before acquisition (e.g., Gemini 3 models require paid-tier credentials) - Improved logging to show priority-aware credential selection and tier compatibility warnings The system gracefully handles unknown credential tiers by treating them as potentially compatible until their actual tier is discovered on first use. Within each priority level, load balancing by usage count is preserved. --- src/rotator_library/client.py | 140 +++++++++- .../providers/gemini_cli_provider.py | 53 ++++ .../providers/provider_interface.py | 47 +++- src/rotator_library/usage_manager.py | 258 +++++++++++++----- 4 files changed, 428 insertions(+), 70 deletions(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 0cb65786..6cdae12f 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -672,6 +672,73 @@ async def _execute_with_retry( lib_logger.info(f"Resolved model '{model}' to '{resolved_model}'") model = resolved_model kwargs["model"] = model # Ensure kwargs has the resolved model for litellm + + # [NEW] Filter by model tier requirement and build priority map + credential_priorities = None + if provider_plugin and hasattr(provider_plugin, 'get_model_tier_requirement'): + required_tier = provider_plugin.get_model_tier_requirement(model) + if required_tier is not None: + # Filter OUT only credentials we KNOW are too low priority + # Keep credentials with unknown priority (None) - they might be high priority + incompatible_creds = [] + compatible_creds = [] + unknown_creds = [] + + for cred in credentials_for_provider: + if hasattr(provider_plugin, 'get_credential_priority'): + priority = provider_plugin.get_credential_priority(cred) + if priority is None: + # Unknown priority - keep it, will be discovered on first use + unknown_creds.append(cred) + elif priority <= required_tier: + # Known compatible priority + compatible_creds.append(cred) + else: + # Known incompatible priority (too low) + incompatible_creds.append(cred) + else: + # Provider doesn't support priorities - keep all + unknown_creds.append(cred) + + # If we have any known-compatible or unknown credentials, use them + tier_compatible_creds = compatible_creds + unknown_creds + if tier_compatible_creds: + credentials_for_provider = tier_compatible_creds + if compatible_creds and unknown_creds: + lib_logger.info( + f"Model {model} requires priority <= {required_tier}. " + f"Using {len(compatible_creds)} known-compatible + {len(unknown_creds)} unknown-tier credentials." + ) + elif compatible_creds: + lib_logger.info( + f"Model {model} requires priority <= {required_tier}. " + f"Using {len(compatible_creds)} known-compatible credentials." + ) + else: + lib_logger.info( + f"Model {model} requires priority <= {required_tier}. " + f"Using {len(unknown_creds)} unknown-tier credentials (will discover on use)." + ) + elif incompatible_creds: + # Only known-incompatible credentials remain + lib_logger.warning( + f"Model {model} requires priority <= {required_tier} credentials, " + f"but all {len(incompatible_creds)} known credentials have priority > {required_tier}. " + f"Request will likely fail." + ) + + # Build priority map for usage_manager + if provider_plugin and hasattr(provider_plugin, 'get_credential_priority'): + credential_priorities = {} + for cred in credentials_for_provider: + priority = provider_plugin.get_credential_priority(cred) + if priority is not None: + credential_priorities[cred] = priority + + if credential_priorities: + lib_logger.debug( + f"Credential priorities for {provider}: {', '.join(f'P{p}={len([c for c in credentials_for_provider if credential_priorities.get(c)==p])}' for p in sorted(set(credential_priorities.values())))}" + ) while ( len(tried_creds) < len(credentials_for_provider) and time.time() < deadline @@ -710,7 +777,8 @@ async def _execute_with_retry( max_concurrent = self.max_concurrent_requests_per_key.get(provider, 1) current_cred = await self.usage_manager.acquire_key( available_keys=creds_to_try, model=model, deadline=deadline, - max_concurrent=max_concurrent + max_concurrent=max_concurrent, + credential_priorities=credential_priorities ) key_acquired = True tried_creds.add(current_cred) @@ -1047,6 +1115,73 @@ async def _streaming_acompletion_with_retry( lib_logger.info(f"Resolved model '{model}' to '{resolved_model}'") model = resolved_model kwargs["model"] = model # Ensure kwargs has the resolved model for litellm + + # [NEW] Filter by model tier requirement and build priority map + credential_priorities = None + if provider_plugin and hasattr(provider_plugin, 'get_model_tier_requirement'): + required_tier = provider_plugin.get_model_tier_requirement(model) + if required_tier is not None: + # Filter OUT only credentials we KNOW are too low priority + # Keep credentials with unknown priority (None) - they might be high priority + incompatible_creds = [] + compatible_creds = [] + unknown_creds = [] + + for cred in credentials_for_provider: + if hasattr(provider_plugin, 'get_credential_priority'): + priority = provider_plugin.get_credential_priority(cred) + if priority is None: + # Unknown priority - keep it, will be discovered on first use + unknown_creds.append(cred) + elif priority <= required_tier: + # Known compatible priority + compatible_creds.append(cred) + else: + # Known incompatible priority (too low) + incompatible_creds.append(cred) + else: + # Provider doesn't support priorities - keep all + unknown_creds.append(cred) + + # If we have any known-compatible or unknown credentials, use them + tier_compatible_creds = compatible_creds + unknown_creds + if tier_compatible_creds: + credentials_for_provider = tier_compatible_creds + if compatible_creds and unknown_creds: + lib_logger.info( + f"Model {model} requires priority <= {required_tier}. " + f"Using {len(compatible_creds)} known-compatible + {len(unknown_creds)} unknown-tier credentials." + ) + elif compatible_creds: + lib_logger.info( + f"Model {model} requires priority <= {required_tier}. " + f"Using {len(compatible_creds)} known-compatible credentials." + ) + else: + lib_logger.info( + f"Model {model} requires priority <= {required_tier}. " + f"Using {len(unknown_creds)} unknown-tier credentials (will discover on use)." + ) + elif incompatible_creds: + # Only known-incompatible credentials remain + lib_logger.warning( + f"Model {model} requires priority <= {required_tier} credentials, " + f"but all {len(incompatible_creds)} known credentials have priority > {required_tier}. " + f"Request will likely fail." + ) + + # Build priority map for usage_manager + if provider_plugin and hasattr(provider_plugin, 'get_credential_priority'): + credential_priorities = {} + for cred in credentials_for_provider: + priority = provider_plugin.get_credential_priority(cred) + if priority is not None: + credential_priorities[cred] = priority + + if credential_priorities: + lib_logger.debug( + f"Credential priorities for {provider}: {', '.join(f'P{p}={len([c for c in credentials_for_provider if credential_priorities.get(c)==p])}' for p in sorted(set(credential_priorities.values())))}" + ) try: while ( @@ -1086,7 +1221,8 @@ async def _streaming_acompletion_with_retry( max_concurrent = self.max_concurrent_requests_per_key.get(provider, 1) current_cred = await self.usage_manager.acquire_key( available_keys=creds_to_try, model=model, deadline=deadline, - max_concurrent=max_concurrent + max_concurrent=max_concurrent, + credential_priorities=credential_priorities ) key_acquired = True tried_creds.add(current_cred) diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 52c7daf8..3ea9c4ea 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -165,6 +165,59 @@ def __init__(self): f"cache={self._enable_signature_cache}, gemini3_fix={self._enable_gemini3_tool_fix}" ) + # ========================================================================= + # CREDENTIAL PRIORITIZATION + # ========================================================================= + + def get_credential_priority(self, credential: str) -> Optional[int]: + """ + Returns priority based on Gemini tier. + Paid tiers: priority 1 (highest) + Free/Legacy tiers: priority 2 + Unknown: priority 10 (lowest) + + Args: + credential: The credential path + + Returns: + Priority level (1-10) or None if tier not yet discovered + """ + tier = self.project_tier_cache.get(credential) + if not tier: + return None # Not yet discovered + + # Paid tiers get highest priority + if tier not in ['free-tier', 'legacy-tier', 'unknown']: + return 1 + + # Free tier gets lower priority + if tier == 'free-tier': + return 2 + + # Legacy and unknown get even lower + return 10 + + def get_model_tier_requirement(self, model: str) -> Optional[int]: + """ + Returns the minimum priority tier required for a model. + Gemini 3 requires paid tier (priority 1). + + Args: + model: The model name (with or without provider prefix) + + Returns: + Minimum required priority level or None if no restrictions + """ + model_name = model.split('/')[-1].replace(':thinking', '') + + # Gemini 3 requires paid tier + if model_name.startswith("gemini-3-"): + return 1 # Only priority 1 (paid) credentials + + return None # All other models have no restrictions + + + # ========================================================================= # MODEL UTILITIES # ========================================================================= diff --git a/src/rotator_library/providers/provider_interface.py b/src/rotator_library/providers/provider_interface.py index 9ca39ecd..8a20a64c 100644 --- a/src/rotator_library/providers/provider_interface.py +++ b/src/rotator_library/providers/provider_interface.py @@ -66,4 +66,49 @@ async def proactively_refresh(self, credential_path: str): """ Proactively refreshes a token if it's nearing expiry. """ - pass \ No newline at end of file + pass + + # [NEW] Credential Prioritization System + def get_credential_priority(self, credential: str) -> Optional[int]: + """ + Returns the priority level for a credential. + Lower numbers = higher priority (1 is highest). + Returns None if provider doesn't use priorities. + + This allows providers to auto-detect credential tiers (e.g., paid vs free) + and ensure higher-tier credentials are always tried first. + + Args: + credential: The credential identifier (API key or path) + + Returns: + Priority level (1-10) or None if no priority system + + Example: + For Gemini CLI: + - Paid tier credentials: priority 1 (highest) + - Free tier credentials: priority 2 + - Unknown tier: priority 10 (lowest) + """ + return None + + def get_model_tier_requirement(self, model: str) -> Optional[int]: + """ + Returns the minimum priority tier required for a model. + If a model requires priority 1, only credentials with priority <= 1 can use it. + + This allows providers to restrict certain models to specific credential tiers. + For example, Gemini 3 models require paid-tier credentials. + + Args: + model: The model name (with or without provider prefix) + + Returns: + Minimum required priority level or None if no restrictions + + Example: + For Gemini CLI: + - gemini-3-*: requires priority 1 (paid tier only) + - gemini-2.5-*: no restriction (None) + """ + return None \ No newline at end of file diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index ec1f1222..d6e0ed99 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -162,11 +162,31 @@ def _initialize_key_states(self, keys: List[str]): async def acquire_key( self, available_keys: List[str], model: str, deadline: float, - max_concurrent: int = 1 + max_concurrent: int = 1, + credential_priorities: Optional[Dict[str, int]] = None ) -> str: """ Acquires the best available key using a tiered, model-aware locking strategy, - respecting a global deadline. + respecting a global deadline and credential priorities. + + Priority Logic: + - Groups credentials by priority level (1=highest, 2=lower, etc.) + - Always tries highest priority (lowest number) first + - Within same priority, sorts by usage count (load balancing) + - Only moves to next priority if all higher-priority keys exhausted/busy + + Args: + available_keys: List of credential identifiers to choose from + model: Model name being requested + deadline: Timestamp after which to stop trying + max_concurrent: Maximum concurrent requests allowed per credential + credential_priorities: Optional dict mapping credentials to priority levels (1=highest) + + Returns: + Selected credential identifier + + Raises: + NoAvailableKeysError: If no key could be acquired within the deadline """ await self._lazy_init() await self._reset_daily_stats_if_needed() @@ -174,78 +194,180 @@ async def acquire_key( # This loop continues as long as the global deadline has not been met. while time.time() < deadline: - tier1_keys, tier2_keys = [], [] now = time.time() - # First, filter the list of available keys to exclude any on cooldown. - async with self._data_lock: - for key in available_keys: - key_data = self._usage_data.get(key, {}) - - if (key_data.get("key_cooldown_until") or 0) > now or ( - key_data.get("model_cooldowns", {}).get(model) or 0 - ) > now: - continue - - # Prioritize keys based on their current usage to ensure load balancing. - usage_count = ( - key_data.get("daily", {}) - .get("models", {}) - .get(model, {}) - .get("success_count", 0) - ) - key_state = self.key_states[key] - - # Tier 1: Completely idle keys (preferred). - if not key_state["models_in_use"]: - tier1_keys.append((key, usage_count)) - # Tier 2: Keys that can accept more concurrent requests for this model. - elif key_state["models_in_use"].get(model, 0) < max_concurrent: - tier2_keys.append((key, usage_count)) - - tier1_keys.sort(key=lambda x: x[1]) - tier2_keys.sort(key=lambda x: x[1]) - - # Attempt to acquire a key from Tier 1 first. - for key, _ in tier1_keys: - state = self.key_states[key] - async with state["lock"]: - if not state["models_in_use"]: - state["models_in_use"][model] = 1 - lib_logger.info( - f"Acquired Tier 1 key ...{key[-6:]} for model {model}" + # Group credentials by priority level (if priorities provided) + if credential_priorities: + # Group keys by priority level + priority_groups = {} + async with self._data_lock: + for key in available_keys: + key_data = self._usage_data.get(key, {}) + + # Skip keys on cooldown + if (key_data.get("key_cooldown_until") or 0) > now or ( + key_data.get("model_cooldowns", {}).get(model) or 0 + ) > now: + continue + + # Get priority for this key (default to 999 if not specified) + priority = credential_priorities.get(key, 999) + + # Get usage count for load balancing within priority groups + usage_count = ( + key_data.get("daily", {}) + .get("models", {}) + .get(model, {}) + .get("success_count", 0) ) - return key - - # If no Tier 1 keys are available, try Tier 2. - for key, _ in tier2_keys: - state = self.key_states[key] - async with state["lock"]: - current_count = state["models_in_use"].get(model, 0) - if current_count < max_concurrent: - state["models_in_use"][model] = current_count + 1 - lib_logger.info( - f"Acquired Tier 2 key ...{key[-6:]} for model {model} " - f"(concurrent: {state['models_in_use'][model]}/{max_concurrent})" + + # Group by priority + if priority not in priority_groups: + priority_groups[priority] = [] + priority_groups[priority].append((key, usage_count)) + + # Try priority groups in order (1, 2, 3, ...) + sorted_priorities = sorted(priority_groups.keys()) + + for priority_level in sorted_priorities: + keys_in_priority = priority_groups[priority_level] + + # Within each priority group, use existing tier1/tier2 logic + tier1_keys, tier2_keys = [], [] + for key, usage_count in keys_in_priority: + key_state = self.key_states[key] + + # Tier 1: Completely idle keys (preferred) + if not key_state["models_in_use"]: + tier1_keys.append((key, usage_count)) + # Tier 2: Keys that can accept more concurrent requests + elif key_state["models_in_use"].get(model, 0) < max_concurrent: + tier2_keys.append((key, usage_count)) + + # Sort by usage within each tier + tier1_keys.sort(key=lambda x: x[1]) + tier2_keys.sort(key=lambda x: x[1]) + + # Try to acquire from Tier 1 first + for key, usage in tier1_keys: + state = self.key_states[key] + async with state["lock"]: + if not state["models_in_use"]: + state["models_in_use"][model] = 1 + lib_logger.info( + f"Acquired Priority-{priority_level} Tier-1 key ...{key[-6:]} for model {model} (usage: {usage})" + ) + return key + + # Then try Tier 2 + for key, usage in tier2_keys: + state = self.key_states[key] + async with state["lock"]: + current_count = state["models_in_use"].get(model, 0) + if current_count < max_concurrent: + state["models_in_use"][model] = current_count + 1 + lib_logger.info( + f"Acquired Priority-{priority_level} Tier-2 key ...{key[-6:]} for model {model} " + f"(concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" + ) + return key + + # If we get here, all priority groups were exhausted but keys might become available + # Collect all keys across all priorities for waiting + all_potential_keys = [] + for keys_list in priority_groups.values(): + all_potential_keys.extend(keys_list) + + if not all_potential_keys: + lib_logger.warning( + "No keys are eligible (all on cooldown or filtered out). Waiting before re-evaluating." + ) + await asyncio.sleep(1) + continue + + # Wait for the highest priority key with lowest usage + best_priority = min(priority_groups.keys()) + best_priority_keys = priority_groups[best_priority] + best_wait_key = min(best_priority_keys, key=lambda x: x[1])[0] + wait_condition = self.key_states[best_wait_key]["condition"] + + lib_logger.info( + f"All Priority-{best_priority} keys are busy. Waiting for highest priority credential to become available..." + ) + + else: + # Original logic when no priorities specified + tier1_keys, tier2_keys = [], [] + + # First, filter the list of available keys to exclude any on cooldown. + async with self._data_lock: + for key in available_keys: + key_data = self._usage_data.get(key, {}) + + if (key_data.get("key_cooldown_until") or 0) > now or ( + key_data.get("model_cooldowns", {}).get(model) or 0 + ) > now: + continue + + # Prioritize keys based on their current usage to ensure load balancing. + usage_count = ( + key_data.get("daily", {}) + .get("models", {}) + .get(model, {}) + .get("success_count", 0) ) - return key - - # If all eligible keys are locked, wait for a key to be released. - lib_logger.info( - "All eligible keys are currently locked for this model. Waiting..." - ) + key_state = self.key_states[key] + + # Tier 1: Completely idle keys (preferred). + if not key_state["models_in_use"]: + tier1_keys.append((key, usage_count)) + # Tier 2: Keys that can accept more concurrent requests for this model. + elif key_state["models_in_use"].get(model, 0) < max_concurrent: + tier2_keys.append((key, usage_count)) + + tier1_keys.sort(key=lambda x: x[1]) + tier2_keys.sort(key=lambda x: x[1]) + + # Attempt to acquire a key from Tier 1 first. + for key, _ in tier1_keys: + state = self.key_states[key] + async with state["lock"]: + if not state["models_in_use"]: + state["models_in_use"][model] = 1 + lib_logger.info( + f"Acquired Tier 1 key ...{key[-6:]} for model {model}" + ) + return key + + # If no Tier 1 keys are available, try Tier 2. + for key, _ in tier2_keys: + state = self.key_states[key] + async with state["lock"]: + current_count = state["models_in_use"].get(model, 0) + if current_count < max_concurrent: + state["models_in_use"][model] = current_count + 1 + lib_logger.info( + f"Acquired Tier 2 key ...{key[-6:]} for model {model} " + f"(concurrent: {state['models_in_use'][model]}/{max_concurrent})" + ) + return key - all_potential_keys = tier1_keys + tier2_keys - if not all_potential_keys: - lib_logger.warning( - "No keys are eligible (all on cooldown). Waiting before re-evaluating." + # If all eligible keys are locked, wait for a key to be released. + lib_logger.info( + "All eligible keys are currently locked for this model. Waiting..." ) - await asyncio.sleep(1) - continue - # Wait on the condition of the key with the lowest current usage. - best_wait_key = min(all_potential_keys, key=lambda x: x[1])[0] - wait_condition = self.key_states[best_wait_key]["condition"] + all_potential_keys = tier1_keys + tier2_keys + if not all_potential_keys: + lib_logger.warning( + "No keys are eligible (all on cooldown). Waiting before re-evaluating." + ) + await asyncio.sleep(1) + continue + + # Wait on the condition of the key with the lowest current usage. + best_wait_key = min(all_potential_keys, key=lambda x: x[1])[0] + wait_condition = self.key_states[best_wait_key]["condition"] try: async with wait_condition: @@ -266,6 +388,8 @@ async def acquire_key( f"Could not acquire a key for model {model} within the global time budget." ) + + async def release_key(self, key: str, model: str): """Releases a key's lock for a specific model and notifies waiting tasks.""" if key not in self.key_states: From f35e0e767d41603a2c81418a587b740eb823e15b Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 09:48:48 +0100 Subject: [PATCH 032/221] =?UTF-8?q?feat(rotation):=20=E2=9C=A8=20add=20con?= =?UTF-8?q?figurable=20weighted=20random=20credential=20selection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a new `rotation_tolerance` parameter to enable weighted random credential selection as an alternative to deterministic least-used rotation. This enhancement addresses potential fingerprinting vulnerabilities while maintaining load balance. - Add `rotation_tolerance` parameter to both `RotatingClient` and `UsageManager` (default: 0.0 for backward compatibility) - Implement `_select_weighted_random()` method using weight formula: `(max_usage - credential_usage) + tolerance + 1` - Support three recommended tolerance levels: - 0.0: Deterministic least-used (existing behavior) - 3.0-4.0: Balanced randomness with good load distribution - 5.0+: High randomness for maximum unpredictability - Update credential acquisition logic to apply weighted selection within tier-based priority system - Enhance logging to indicate selection method (weighted-random vs least-used) and include usage counts - Add comprehensive docstrings explaining rotation strategy and tolerance impact - Import `random` module for weighted selection functionality The weighted random approach reduces predictability in credential usage patterns while the tolerance parameter allows fine-tuning the balance between randomness and efficiency. --- src/rotator_library/client.py | 27 ++++- src/rotator_library/usage_manager.py | 142 ++++++++++++++++++++++++--- 2 files changed, 156 insertions(+), 13 deletions(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 6cdae12f..bfd3be5a 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -63,7 +63,29 @@ def __init__( whitelist_models: Optional[Dict[str, List[str]]] = None, enable_request_logging: bool = False, max_concurrent_requests_per_key: Optional[Dict[str, int]] = None, + rotation_tolerance: float = 3.0, ): + """ + Initialize the RotatingClient with intelligent credential rotation. + + Args: + api_keys: Dictionary mapping provider names to lists of API keys + oauth_credentials: Dictionary mapping provider names to OAuth credential paths + max_retries: Maximum number of retry attempts per credential + usage_file_path: Path to store usage statistics + configure_logging: Whether to configure library logging + global_timeout: Global timeout for requests in seconds + abort_on_callback_error: Whether to abort on pre-request callback errors + litellm_provider_params: Provider-specific parameters for LiteLLM + ignore_models: Models to ignore/blacklist per provider + whitelist_models: Models to explicitly whitelist per provider + enable_request_logging: Whether to enable detailed request logging + max_concurrent_requests_per_key: Max concurrent requests per key by provider + rotation_tolerance: Tolerance for weighted random credential rotation. + - 0.0: Deterministic, least-used credential always selected + - 2.0 - 4.0 (default, recommended): Balanced randomness, can pick credentials within 2 uses of max + - 5.0+: High randomness, more unpredictable selection patterns + """ os.environ["LITELLM_LOG"] = "ERROR" litellm.set_verbose = False litellm.drop_params = True @@ -108,7 +130,10 @@ def __init__( self.max_retries = max_retries self.global_timeout = global_timeout self.abort_on_callback_error = abort_on_callback_error - self.usage_manager = UsageManager(file_path=usage_file_path) + self.usage_manager = UsageManager( + file_path=usage_file_path, + rotation_tolerance=rotation_tolerance + ) self._model_list_cache = {} self._provider_plugins = PROVIDER_PLUGINS self._provider_instances = {} diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index d6e0ed99..4ec2b825 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -3,6 +3,7 @@ import time import logging import asyncio +import random from datetime import date, datetime, timezone, time as dt_time from typing import Any, Dict, List, Optional, Set import aiofiles @@ -20,15 +21,48 @@ class UsageManager: """ Manages usage statistics and cooldowns for API keys with asyncio-safe locking, - asynchronous file I/O, and a lazy-loading mechanism for usage data. + asynchronous file I/O, lazy-loading mechanism, and weighted random credential rotation. + + The credential rotation strategy can be configured via the `rotation_tolerance` parameter: + + - **tolerance = 0.0**: Deterministic least-used selection. The credential with + the lowest usage count is always selected. This provides predictable, perfectly balanced + load distribution but may be vulnerable to fingerprinting. + + - **tolerance = 2.0 - 4.0 (default, recommended)**: Balanced weighted randomness. Credentials are selected + randomly with weights biased toward less-used ones. Credentials within 2 uses of the + maximum can still be selected with reasonable probability. This provides security through + unpredictability while maintaining good load balance. + + - **tolerance = 5.0+**: High randomness. Even heavily-used credentials have significant + selection probability. Useful for stress testing or maximum unpredictability, but may + result in less balanced load distribution. + + The weight formula is: `weight = (max_usage - credential_usage) + tolerance + 1` + + This ensures lower-usage credentials are preferred while tolerance controls how much + randomness is introduced into the selection process. """ def __init__( self, file_path: str = "key_usage.json", daily_reset_time_utc: Optional[str] = "03:00", + rotation_tolerance: float = 0.0, ): + """ + Initialize the UsageManager. + + Args: + file_path: Path to the usage data JSON file + daily_reset_time_utc: Time in UTC when daily stats should reset (HH:MM format) + rotation_tolerance: Tolerance for weighted random credential rotation. + - 0.0: Deterministic, least-used credential always selected + - tolerance = 2.0 - 4.0 (default, recommended): Balanced randomness, can pick credentials within 2 uses of max + - 5.0+: High randomness, more unpredictable selection patterns + """ self.file_path = file_path + self.rotation_tolerance = rotation_tolerance self.key_states: Dict[str, Dict[str, Any]] = {} self._data_lock = asyncio.Lock() @@ -160,6 +194,63 @@ def _initialize_key_states(self, keys: List[str]): "models_in_use": {}, # Dict[model_name, concurrent_count] } + def _select_weighted_random( + self, + candidates: List[tuple], + tolerance: float + ) -> str: + """ + Selects a credential using weighted random selection based on usage counts. + + Args: + candidates: List of (credential_id, usage_count) tuples + tolerance: Tolerance value for weight calculation + + Returns: + Selected credential ID + + Formula: + weight = (max_usage - credential_usage) + tolerance + 1 + + This formula ensures: + - Lower usage = higher weight = higher selection probability + - Tolerance adds variability: higher tolerance means more randomness + - The +1 ensures all credentials have at least some chance of selection + """ + if not candidates: + raise ValueError("Cannot select from empty candidate list") + + if len(candidates) == 1: + return candidates[0][0] + + # Extract usage counts + usage_counts = [usage for _, usage in candidates] + max_usage = max(usage_counts) + + # Calculate weights using the formula: (max - current) + tolerance + 1 + weights = [] + for credential, usage in candidates: + weight = (max_usage - usage) + tolerance + 1 + weights.append(weight) + + # Log weight distribution for debugging + if lib_logger.isEnabledFor(logging.DEBUG): + total_weight = sum(weights) + weight_info = ", ".join( + f"...{cred[-6:]}: w={w:.1f} ({w/total_weight*100:.1f}%)" + for (cred, _), w in zip(candidates, weights) + ) + #lib_logger.debug(f"Weighted selection candidates: {weight_info}") + + # Random selection with weights + selected_credential = random.choices( + [cred for cred, _ in candidates], + weights=weights, + k=1 + )[0] + + return selected_credential + async def acquire_key( self, available_keys: List[str], model: str, deadline: float, max_concurrent: int = 1, @@ -244,9 +335,21 @@ async def acquire_key( elif key_state["models_in_use"].get(model, 0) < max_concurrent: tier2_keys.append((key, usage_count)) - # Sort by usage within each tier - tier1_keys.sort(key=lambda x: x[1]) - tier2_keys.sort(key=lambda x: x[1]) + # Apply weighted random selection or deterministic sorting + selection_method = "weighted-random" if self.rotation_tolerance > 0 else "least-used" + + if self.rotation_tolerance > 0: + # Weighted random selection within each tier + if tier1_keys: + selected_key = self._select_weighted_random(tier1_keys, self.rotation_tolerance) + tier1_keys = [(k, u) for k, u in tier1_keys if k == selected_key] + if tier2_keys: + selected_key = self._select_weighted_random(tier2_keys, self.rotation_tolerance) + tier2_keys = [(k, u) for k, u in tier2_keys if k == selected_key] + else: + # Deterministic: sort by usage within each tier + tier1_keys.sort(key=lambda x: x[1]) + tier2_keys.sort(key=lambda x: x[1]) # Try to acquire from Tier 1 first for key, usage in tier1_keys: @@ -255,7 +358,8 @@ async def acquire_key( if not state["models_in_use"]: state["models_in_use"][model] = 1 lib_logger.info( - f"Acquired Priority-{priority_level} Tier-1 key ...{key[-6:]} for model {model} (usage: {usage})" + f"Acquired Priority-{priority_level} Tier-1 key ...{key[-6:]} for model {model} " + f"(selection: {selection_method}, usage: {usage})" ) return key @@ -268,7 +372,7 @@ async def acquire_key( state["models_in_use"][model] = current_count + 1 lib_logger.info( f"Acquired Priority-{priority_level} Tier-2 key ...{key[-6:]} for model {model} " - f"(concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" + f"(selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" ) return key @@ -325,22 +429,36 @@ async def acquire_key( elif key_state["models_in_use"].get(model, 0) < max_concurrent: tier2_keys.append((key, usage_count)) - tier1_keys.sort(key=lambda x: x[1]) - tier2_keys.sort(key=lambda x: x[1]) + # Apply weighted random selection or deterministic sorting + selection_method = "weighted-random" if self.rotation_tolerance > 0 else "least-used" + + if self.rotation_tolerance > 0: + # Weighted random selection within each tier + if tier1_keys: + selected_key = self._select_weighted_random(tier1_keys, self.rotation_tolerance) + tier1_keys = [(k, u) for k, u in tier1_keys if k == selected_key] + if tier2_keys: + selected_key = self._select_weighted_random(tier2_keys, self.rotation_tolerance) + tier2_keys = [(k, u) for k, u in tier2_keys if k == selected_key] + else: + # Deterministic: sort by usage within each tier + tier1_keys.sort(key=lambda x: x[1]) + tier2_keys.sort(key=lambda x: x[1]) # Attempt to acquire a key from Tier 1 first. - for key, _ in tier1_keys: + for key, usage in tier1_keys: state = self.key_states[key] async with state["lock"]: if not state["models_in_use"]: state["models_in_use"][model] = 1 lib_logger.info( - f"Acquired Tier 1 key ...{key[-6:]} for model {model}" + f"Acquired Tier 1 key ...{key[-6:]} for model {model} " + f"(selection: {selection_method}, usage: {usage})" ) return key # If no Tier 1 keys are available, try Tier 2. - for key, _ in tier2_keys: + for key, usage in tier2_keys: state = self.key_states[key] async with state["lock"]: current_count = state["models_in_use"].get(model, 0) @@ -348,7 +466,7 @@ async def acquire_key( state["models_in_use"][model] = current_count + 1 lib_logger.info( f"Acquired Tier 2 key ...{key[-6:]} for model {model} " - f"(concurrent: {state['models_in_use'][model]}/{max_concurrent})" + f"(selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" ) return key From f5ccdf66e7678fa7cc5f487a071dcfa979b958ac Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 09:58:21 +0100 Subject: [PATCH 033/221] =?UTF-8?q?docs:=20=F0=9F=93=9A=20add=20comprehens?= =?UTF-8?q?ive=20documentation=20for=20new=20features=20and=20providers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit adds extensive documentation for recently implemented features across all documentation files: - **Antigravity Provider**: Complete documentation of the new Antigravity provider with support for Gemini 2.5, Gemini 3, and Claude Sonnet 4.5 models, including thought signature caching, tool hallucination prevention, and base URL fallback mechanisms - **Credential Prioritization System**: Detailed explanation of the new tier-based credential selection system that ensures paid-tier credentials are used for premium models - **Weighted Random Rotation**: Documentation of the configurable `rotation_tolerance` parameter that enables unpredictable credential selection patterns to avoid fingerprinting while maintaining load balance - **Provider Cache System**: Architecture and usage documentation for the new modular caching system used for preserving conversation state across requests - **Google OAuth Base Refactoring**: Documentation of the shared `GoogleOAuthBase` class that eliminates code duplication across OAuth providers - **Enhanced Gemini CLI Features**: Updated documentation covering project tier detection, paid vs free tier credential prioritization, and Gemini 3 support - **Temperature Override**: Global temperature=0 override configuration to prevent tool hallucination issues - **Deployment Guide Updates**: Step-by-step instructions for setting up Antigravity OAuth credentials in both local and stateless deployment scenarios - **Environment Variable Reference**: Comprehensive list of new configuration options including cache control, feature flags, and rotation strategy settings The documentation includes practical examples, configuration snippets, use cases, and security benefits for each feature. --- DOCUMENTATION.md | 248 +++++++++++++++++++++++++++++++++- Deployment guide.md | 31 +++++ README.md | 113 +++++++++++++++- src/rotator_library/README.md | 42 +++++- 4 files changed, 429 insertions(+), 5 deletions(-) diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index bd4c6c17..94beec4b 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -57,6 +57,7 @@ client = RotatingClient( - `whitelist_models` (`Optional[Dict[str, List[str]]]`, default: `None`): Whitelist of models to always include, overriding `ignore_models`. - `enable_request_logging` (`bool`, default: `False`): If `True`, enables detailed per-request file logging. - `max_concurrent_requests_per_key` (`Optional[Dict[str, int]]`, default: `None`): Max concurrent requests allowed for a single API key per provider. +- `rotation_tolerance` (`float`, default: `3.0`): Controls the credential rotation strategy. See Section 2.2 for details. #### Core Responsibilities @@ -110,8 +111,16 @@ The `acquire_key` method uses a sophisticated strategy to balance load: 2. **Tiering**: Valid keys are split into two tiers: * **Tier 1 (Ideal)**: Keys that are completely idle (0 concurrent requests). * **Tier 2 (Acceptable)**: Keys that are busy but still under their configured `MAX_CONCURRENT_REQUESTS_PER_KEY_` limit for the requested model. This allows a single key to be used multiple times for the same model, maximizing throughput. -3. **Prioritization**: Within each tier, keys with the **lowest daily usage** are prioritized to spread costs evenly. +3. **Selection Strategy** (configurable via `rotation_tolerance`): + * **Deterministic (tolerance=0.0)**: Within each tier, keys are sorted by daily usage count and the least-used key is always selected. This provides perfect load balance but predictable patterns. + * **Weighted Random (tolerance>0, default)**: Keys are selected randomly with weights biased toward less-used ones: + - Formula: `weight = (max_usage - credential_usage) + tolerance + 1` + - `tolerance=2.0` (recommended): Balanced randomness - credentials within 2 uses of the maximum can still be selected with reasonable probability + - `tolerance=5.0+`: High randomness - even heavily-used credentials have significant probability + - **Security Benefit**: Unpredictable selection patterns make rate limit detection and fingerprinting harder + - **Load Balance**: Lower-usage credentials still preferred, maintaining reasonable distribution 4. **Concurrency Limits**: Checks against `max_concurrent` limits to prevent overloading a single key. +5. **Priority Groups**: When credential prioritization is enabled, higher-tier credentials (lower priority numbers) are tried first before moving to lower tiers. #### Failure Handling & Cooldowns @@ -313,6 +322,243 @@ The `CooldownManager` handles IP or account-level rate limiting that affects all - If so, `CooldownManager.start_cooldown()` is called for the entire provider - All subsequent `acquire_key()` calls for that provider will wait until the cooldown expires + +### 2.10. Credential Prioritization System (`client.py` & `usage_manager.py`) + +The library now includes an intelligent credential prioritization system that automatically detects credential tiers and ensures optimal credential selection for each request. + +**Key Concepts:** + +- **Provider-Level Priorities**: Providers can implement `get_credential_priority()` to return a priority level (1=highest, 10=lowest) for each credential +- **Model-Level Requirements**: Providers can implement `get_model_tier_requirement()` to specify minimum priority required for specific models +- **Automatic Filtering**: The client automatically filters out incompatible credentials before making requests +- **Priority-Aware Selection**: The `UsageManager` prioritizes higher-tier credentials (lower numbers) within the same priority group + +**Implementation Example (Gemini CLI):** + +```python +def get_credential_priority(self, credential: str) -> Optional[int]: + """Returns priority based on Gemini tier.""" + tier = self.project_tier_cache.get(credential) + if not tier: + return None # Not yet discovered + + # Paid tiers get highest priority + if tier not in ['free-tier', 'legacy-tier', 'unknown']: + return 1 + + # Free tier gets lower priority + if tier == 'free-tier': + return 2 + + return 10 + +def get_model_tier_requirement(self, model: str) -> Optional[int]: + """Returns minimum priority required for model.""" + if model.startswith("gemini-3-"): + return 1 # Only paid tier (priority 1) credentials + + return None # All other models have no restrictions +``` + +**Usage Manager Integration:** + +The `acquire_key()` method has been enhanced to: +1. Group credentials by priority level +2. Try highest priority group first (priority 1, then 2, etc.) +3. Within each group, use existing tier1/tier2 logic (idle keys first, then busy keys) +4. Load balance within priority groups by usage count +5. Only move to next priority if all higher-priority credentials are exhausted + +**Benefits:** + +- Ensures paid-tier credentials are always used for premium models +- Prevents failed requests due to tier restrictions +- Optimal cost distribution (free tier used when possible, paid when required) +- Graceful fallback if primary credentials are unavailable + +--- + +### 2.11. Provider Cache System (`providers/provider_cache.py`) + +A modular, shared caching system for providers to persist conversation state across requests. + +**Architecture:** + +- **Dual-TTL Design**: Short-lived memory cache (default: 1 hour) + longer-lived disk persistence (default: 24 hours) +- **Background Persistence**: Batched disk writes every 60 seconds (configurable) +- **Automatic Cleanup**: Background task removes expired entries from memory cache + +### 3.5. Antigravity (`antigravity_provider.py`) + +The most sophisticated provider implementation, supporting Google's internal Antigravity API for Gemini and Claude models. + +#### Architecture + +- **Unified Streaming/Non-Streaming**: Single code path handles both response types with optimal transformations +- **Thought Signature Caching**: Server-side caching of encrypted signatures for multi-turn Gemini 3 conversations +- **Model-Specific Logic**: Automatic configuration based on model type (Gemini 2.5, Gemini 3, Claude) + +#### Model Support + +**Gemini 2.5 (Pro/Flash):** +- Uses `thinkingBudget` parameter (integer tokens: -1 for auto, 0 to disable, or specific value) +- Standard safety settings and toolConfig +- Stream processing with thinking content separation + +**Gemini 3 (Pro/Image):** +- Uses `thinkingLevel` parameter (string: "low" or "high") +- **Tool Hallucination Prevention**: + - Automatic system instruction injection explaining custom tool schema rules + - Parameter signature injection into tool descriptions (e.g., "STRICT PARAMETERS: files (ARRAY_OF_OBJECTS[path: string REQUIRED, ...])") + - Namespace prefix for tool names (`gemini3_` prefix) to avoid training data conflicts + - Malformed JSON auto-correction (handles extra trailing braces) +- **ThoughtSignature Management**: + - Caching signatures from responses for reuse in follow-up messages + - Automatic injection into functionCalls for multi-turn conversations + - Fallback to bypass value if signature unavailable + +**Claude Sonnet 4.5:** +- Proxied through Antigravity API (uses internal model name `claude-sonnet-4-5-thinking`) +- Uses `thinkingBudget` parameter like Gemini 2.5 +- **Thinking Preservation**: Caches thinking content using composite keys (tool_call_id + text_hash) +- **Schema Cleaning**: Removes unsupported properties (`$schema`, `additionalProperties`, `const` → `enum`) + +#### Base URL Fallback + +Automatic fallback chain for resilience: +1. `daily-cloudcode-pa.sandbox.googleapis.com` (primary sandbox) +2. `autopush-cloudcode-pa.sandbox.googleapis.com` (fallback sandbox) +3. `cloudcode-pa.googleapis.com` (production fallback) + +#### Message Transformation + +**OpenAI → Gemini Format:** +- System messages → `systemInstruction` with parts array +- Multi-part content (text + images) → `inlineData` format +- Tool calls → `functionCall` with args and id +- Tool responses → `functionResponse` with name and response +- ThoughtSignatures preserved/injected as needed + +**Tool Response Grouping:** +- Converts linear format (call, response, call, response) to grouped format +- Groups all function calls in one `model` message +- Groups all responses in one `user` message +- Required for Antigravity API compatibility + +#### Configuration (Environment Variables) + +```env +# Cache control +ANTIGRAVITY_SIGNATURE_CACHE_TTL=3600 # Memory cache TTL +ANTIGRAVITY_SIGNATURE_DISK_TTL=86400 # Disk cache TTL +ANTIGRAVITY_ENABLE_SIGNATURE_CACHE=true + +# Feature flags +ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES=true # Include signatures in client responses +ANTIGRAVITY_ENABLE_DYNAMIC_MODELS=false # Use API model discovery +ANTIGRAVITY_GEMINI3_TOOL_FIX=true # Enable Gemini 3 hallucination prevention + +# Gemini 3 tool fix customization +ANTIGRAVITY_GEMINI3_TOOL_PREFIX="gemini3_" # Namespace prefix +ANTIGRAVITY_GEMINI3_DESCRIPTION_PROMPT="\n\nSTRICT PARAMETERS: {params}." +ANTIGRAVITY_GEMINI3_SYSTEM_INSTRUCTION="..." # Full system prompt +``` + +#### File Logging + +Optional transaction logging for debugging: +- Enabled via `enable_request_logging` parameter +- Creates `logs/antigravity_logs/TIMESTAMP_MODEL_UUID/` directory per request +- Logs: `request_payload.json`, `response_stream.log`, `final_response.json`, `error.log` + +--- + + +- **Atomic Disk Writes**: Uses temp-file-and-move pattern to prevent corruption + +**Key Methods:** + +1. **`store(key, value)`**: Synchronously queues value for storage (schedules async write) +2. **`retrieve(key)`**: Synchronously retrieves from memory, optionally schedules disk fallback +3. **`store_async(key, value)`**: Awaitable storage for guaranteed persistence +4. **`retrieve_async(key)`**: Awaitable retrieval with disk fallback + +**Use Cases:** + +- **Gemini 3 ThoughtSignatures**: Caching tool call signatures for multi-turn conversations +- **Claude Thinking**: Preserving thinking content for consistency across conversation turns +- **Any Transient State**: Generic key-value storage for provider-specific needs + +**Configuration (Environment Variables):** + +```env +# Cache control (prefix can be customized per cache instance) +PROVIDER_CACHE_ENABLE=true +PROVIDER_CACHE_WRITE_INTERVAL=60 # seconds between disk writes +PROVIDER_CACHE_CLEANUP_INTERVAL=1800 # 30 min between cleanups + +# Gemini 3 specific +GEMINI_CLI_SIGNATURE_CACHE_ENABLE=true +GEMINI_CLI_SIGNATURE_CACHE_TTL=3600 # 1 hour memory TTL +GEMINI_CLI_SIGNATURE_DISK_TTL=86400 # 24 hours disk TTL +``` + +**File Structure:** + +``` +cache/ +├── gemini_cli/ +│ └── gemini3_signatures.json +└── antigravity/ + ├── gemini3_signatures.json + └── claude_thinking.json +``` + +--- + +### 2.12. Google OAuth Base (`providers/google_oauth_base.py`) + +A refactored, reusable OAuth2 base class that eliminates code duplication across Google-based providers. + +**Refactoring Benefits:** + +- **Single Source of Truth**: All OAuth logic centralized in one class +- **Easy Provider Addition**: New providers only need to override constants +- **Consistent Behavior**: Token refresh, expiry handling, and validation work identically across providers +- **Maintainability**: OAuth bugs fixed once apply to all inheriting providers + +**Provider Implementation:** + +```python +class AntigravityAuthBase(GoogleOAuthBase): + # Required overrides + CLIENT_ID = "antigravity-client-id" + CLIENT_SECRET = "antigravity-secret" + OAUTH_SCOPES = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/cclog", # Antigravity-specific + "https://www.googleapis.com/auth/experimentsandconfigs", + ] + ENV_PREFIX = "ANTIGRAVITY" # Used for env var loading + + # Optional overrides (defaults provided) + CALLBACK_PORT = 51121 + CALLBACK_PATH = "/oauthcallback" +``` + +**Inherited Features:** + +- Automatic token refresh with exponential backoff +- Invalid grant re-authentication flow +- Stateless deployment support (env var loading) +- Atomic credential file writes +- Headless environment detection +- Sequential refresh queue processing + +--- + + --- ## 3. Provider Specific Implementations diff --git a/Deployment guide.md b/Deployment guide.md index 1d31c14f..57acd536 100644 --- a/Deployment guide.md +++ b/Deployment guide.md @@ -79,6 +79,37 @@ If you are using providers that require complex OAuth files (like **Gemini CLI** 4. Copy the contents of this file and paste them directly into your `.env` file or Render's "Environment Variables" section. 5. The proxy will automatically detect and use these variables—no file upload required! + +### Advanced: Antigravity OAuth Provider + +The Antigravity provider requires OAuth2 authentication similar to Gemini CLI. It provides access to: +- Gemini 2.5 models (Pro/Flash) +- Gemini 3 models (Pro/Image-preview) - **requires paid-tier Google Cloud project** +- Claude Sonnet 4.5 via Google's Antigravity proxy + +**Setting up Antigravity locally:** +1. Run the credential tool: `python -m rotator_library.credential_tool` +2. Select "Add OAuth Credential" and choose "Antigravity" +3. Complete the OAuth flow in your browser +4. The credential is saved to `oauth_creds/antigravity_oauth_1.json` + +**Exporting for stateless deployment:** +1. Run: `python -m rotator_library.credential_tool` +2. Select "Export Antigravity to .env" +3. Copy the generated environment variables to your deployment platform: + ```env + ANTIGRAVITY_ACCESS_TOKEN="..." + ANTIGRAVITY_REFRESH_TOKEN="..." + ANTIGRAVITY_EXPIRY_DATE="..." + ANTIGRAVITY_EMAIL="your-email@gmail.com" + ``` + +**Important Notes:** +- Antigravity uses Google OAuth with additional scopes for cloud platform access +- Gemini 3 models require a paid-tier Google Cloud project (free tier will fail) +- The provider automatically handles thought signature caching for multi-turn conversations +- Tool hallucination prevention is enabled by default for Gemini 3 models + 4. Save the file. (We'll upload it to Render in Step 5.) diff --git a/README.md b/README.md index 6129d11d..f3a12867 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,15 @@ This project provides a powerful solution for developers building complex applic - **Provider Agnostic**: Compatible with any provider supported by `litellm`. - **OpenAI-Compatible Proxy**: Offers a familiar API interface with additional endpoints for model and provider discovery. - **Advanced Model Filtering**: Supports both blacklists and whitelists to give you fine-grained control over which models are available through the proxy. + +- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 2.5, Gemini 3, and Claude Sonnet 4.5 models with advanced features like thought signature caching and tool hallucination prevention. +- **🆕 Credential Prioritization**: Automatic tier detection and priority-based credential selection ensures paid-tier credentials are used for premium models that require them. +- **🆕 Weighted Random Rotation**: Configurable credential rotation strategy - choose between deterministic (perfect balance) or weighted random (unpredictable, harder to fingerprint) selection. +- **🆕 Enhanced Gemini CLI**: Improved project discovery, paid vs free tier detection, and Gemini 3 support with thoughtSignature caching. +- **🆕 Temperature Override**: Global temperature=0 override option to prevent tool hallucination issues with low-temperature settings. +- **🆕 Provider Cache System**: Modular caching system for preserving conversation state (thought signatures, thinking content) across requests. +- **🆕 Refactored OAuth Base**: Shared [`GoogleOAuthBase`](src/rotator_library/providers/google_oauth_base.py) class eliminates code duplication across OAuth providers. + - **🆕 Interactive Launcher TUI**: Beautiful, cross-platform TUI for configuration and management with an integrated settings tool for advanced configuration. @@ -234,11 +243,12 @@ python src/proxy_app/main.py **Main Menu Features:** -1. **Add OAuth Credential** - Interactive OAuth flow for Gemini CLI, Qwen Code, and iFlow +1. **Add OAuth Credential** - Interactive OAuth flow for Gemini CLI, Antigravity, Qwen Code, and iFlow - Automatically opens your browser for authentication - Handles the entire OAuth flow including callbacks - Saves credentials to the local `oauth_creds/` directory - For Gemini CLI: Automatically discovers or creates a Google Cloud project + - For Antigravity: Similar to Gemini CLI with Antigravity-specific scopes - For Qwen Code: Uses Device Code flow (you'll enter a code in your browser) - For iFlow: Starts a local callback server on port 11451 @@ -488,6 +498,42 @@ The following advanced settings can be added to your `.env` file (or configured - **`SKIP_OAUTH_INIT_CHECK`**: Set to `true` to skip the interactive OAuth setup/validation check on startup. Essential for non-interactive environments like Docker containers or CI/CD pipelines. ```env SKIP_OAUTH_INIT_CHECK=true + + +#### **Antigravity (Advanced - Gemini 3 \Claude 4.5 Access)** +The newest and most sophisticated provider, offering access to cutting-edge models via Google's internal Antigravity API. + +**Supported Models:** +- Gemini 2.5 (Pro/Flash) with `thinkingBudget` parameter +- **Gemini 3 Pro (High/Low)** - Latest preview models +- **Claude Sonnet 4.5 + Thinking** via Antigravity proxy + +**Advanced Features:** +- **Thought Signature Caching**: Preserves encrypted signatures for multi-turn Gemini 3 conversations +- **Tool Hallucination Prevention**: Automatic system instruction and parameter signature injection for Gemini 3 to prevent tools from being called with incorrect parameters +- **Thinking Preservation**: Caches Claude thinking content for consistency across conversation turns +- **Automatic Fallback**: Tries sandbox endpoints before falling back to production +- **Schema Cleaning**: Handles Claude-specific tool schema requirements + +**Configuration:** +- **OAuth Setup**: Uses Google OAuth similar to Gemini CLI (separate scopes) +- **Stateless Deployment**: Full environment variable support +- **Paid Tier Recommended**: Gemini 3 models require a paid Google Cloud project + +**Environment Variables:** +```env +# Stateless deployment +ANTIGRAVITY_ACCESS_TOKEN="..." +ANTIGRAVITY_REFRESH_TOKEN="..." +ANTIGRAVITY_EXPIRY_DATE="..." +ANTIGRAVITY_EMAIL="user@gmail.com" + +# Feature toggles +ANTIGRAVITY_ENABLE_SIGNATURE_CACHE=true # Multi-turn conversation support +ANTIGRAVITY_GEMINI3_TOOL_FIX=true # Prevent tool hallucination +``` + + ``` #### Concurrency Control @@ -516,6 +562,71 @@ For providers that support custom model definitions (Qwen Code, iFlow), you can #### Provider-Specific Settings - **`GEMINI_CLI_PROJECT_ID`**: Manually specify a Google Cloud Project ID for Gemini CLI OAuth. Only needed if automatic discovery fails. + + +#### Antigravity Provider + +- **`ANTIGRAVITY_OAUTH_1`**: Path to Antigravity OAuth credential file (auto-discovered from `~/.antigravity/` or use the credential tool). + ```env + ANTIGRAVITY_OAUTH_1="/path/to/your/antigravity_creds.json" + ``` + +- **Stateless Deployment** (Environment Variables): + ```env + ANTIGRAVITY_ACCESS_TOKEN="ya29.your-access-token" + + +#### Credential Rotation Strategy + +- **`ROTATION_TOLERANCE`**: Controls how credentials are selected for requests. Set via environment variable or programmatically. + - `0.0`: **Deterministic** - Always selects the least-used credential for perfect load balance + - `3.0` (default, recommended): **Weighted Random** - Randomly selects with bias toward less-used credentials. Provides unpredictability (harder to fingerprint/detect) while maintaining good balance + - `5.0+`: **High Randomness** - Maximum unpredictability, even heavily-used credentials can be selected + + ```env + # For maximum security/unpredictability (recommended for production) + ROTATION_TOLERANCE=3.0 + + # For perfect load balancing (default) + ROTATION_TOLERANCE=0.0 + ``` + + **Why use weighted random?** + - Makes traffic patterns less predictable + - Still maintains good load distribution across keys + - Recommended for production environments with multiple credentials + + + ANTIGRAVITY_REFRESH_TOKEN="1//your-refresh-token" + ANTIGRAVITY_EXPIRY_DATE="1234567890000" + ANTIGRAVITY_EMAIL="your-email@gmail.com" + ``` + +- **`ANTIGRAVITY_ENABLE_SIGNATURE_CACHE`**: Enable/disable thought signature caching for Gemini 3 multi-turn conversations. Default: `true`. + ```env + ANTIGRAVITY_ENABLE_SIGNATURE_CACHE=true + ``` + +- **`ANTIGRAVITY_GEMINI3_TOOL_FIX`**: Enable/disable tool hallucination prevention for Gemini 3 models. Default: `true`. + ```env + ANTIGRAVITY_GEMINI3_TOOL_FIX=true + ``` + +#### Temperature Override (Global) + +- **`OVERRIDE_TEMPERATURE_ZERO`**: Prevents tool hallucination caused by temperature=0 settings. Modes: + - `"remove"`: Deletes temperature=0 from requests (lets provider use default) + - `"set"`: Changes temperature=0 to temperature=1.0 + - `"false"` or unset: Disabled (default) + +#### Credential Prioritization + +- **`GEMINI_CLI_PROJECT_ID`**: Manually specify a Google Cloud Project ID for Gemini CLI OAuth. Auto-discovered unless unexpected failure occurs. + ```env + GEMINI_CLI_PROJECT_ID="your-gcp-project-id" + ``` + + ```env GEMINI_CLI_PROJECT_ID="your-gcp-project-id" ``` diff --git a/src/rotator_library/README.md b/src/rotator_library/README.md index c0207999..2050f1ba 100644 --- a/src/rotator_library/README.md +++ b/src/rotator_library/README.md @@ -7,9 +7,11 @@ A robust, asynchronous, and thread-safe Python library for managing a pool of AP - **Asynchronous by Design**: Built with `asyncio` and `httpx` for high-performance, non-blocking I/O. - **Advanced Concurrency Control**: A single API key can be used for multiple concurrent requests. By default, it supports concurrent requests to *different* models. With configuration (`MAX_CONCURRENT_REQUESTS_PER_KEY_`), it can also support multiple concurrent requests to the *same* model using the same key. - **Smart Key Management**: Selects the optimal key for each request using a tiered, model-aware locking strategy to distribute load evenly and maximize availability. +- **Configurable Rotation Strategy**: Choose between deterministic least-used selection (perfect balance) or default weighted random selection (unpredictable, harder to fingerprint). - **Deadline-Driven Requests**: A global timeout ensures that no request, including all retries and key selections, exceeds a specified time limit. - **OAuth & API Key Support**: Built-in support for standard API keys and complex OAuth flows. - - **Gemini CLI**: Full OAuth 2.0 web flow with automatic project discovery and free-tier onboarding. + - **Gemini CLI**: Full OAuth 2.0 web flow with automatic project discovery, free-tier onboarding, and credential prioritization (paid vs free tier). + - **Antigravity**: Full OAuth 2.0 support for Gemini 3, Gemini 2.5, and Claude Sonnet 4.5 models with thought signature caching(Full support for Gemini 3 and Claude models). **First on the scene to provide full support for Gemini 3** via Antigravity with advanced features like thought signature caching and tool hallucination prevention. - **Qwen Code**: Device Code flow support. - **iFlow**: Authorization Code flow with local callback handling. - **Stateless Deployment Ready**: Can load complex OAuth credentials from environment variables, eliminating the need for physical credential files in containerized environments. @@ -17,11 +19,15 @@ A robust, asynchronous, and thread-safe Python library for managing a pool of AP - **Escalating Per-Model Cooldowns**: Failed keys are placed on a temporary, escalating cooldown for specific models. - **Key-Level Lockouts**: Keys failing across multiple models are temporarily removed from rotation. - **Stream Recovery**: The client detects mid-stream errors (like quota limits) and gracefully handles them. +- **Credential Prioritization**: Automatic tier detection and priority-based credential selection (e.g., paid tier credentials used first for models that require them). +- **Advanced Model Requirements**: Support for model-tier restrictions (e.g., Gemini 3 requires paid-tier credentials). - **Robust Streaming Support**: Includes a wrapper for streaming responses that reassembles fragmented JSON chunks. - **Detailed Usage Tracking**: Tracks daily and global usage for each key, persisted to a JSON file. - **Automatic Daily Resets**: Automatically resets cooldowns and archives stats daily. - **Provider Agnostic**: Works with any provider supported by `litellm`. - **Extensible**: Easily add support for new providers through a simple plugin-based architecture. +- **Temperature Override**: Global temperature=0 override to prevent tool hallucination with low-temperature settings. +- **Shared OAuth Base**: Refactored OAuth implementation with reusable [`GoogleOAuthBase`](providers/google_oauth_base.py) for multiple providers. ## Installation @@ -71,7 +77,8 @@ client = RotatingClient( ignore_models={}, whitelist_models={}, enable_request_logging=False, - max_concurrent_requests_per_key={} + max_concurrent_requests_per_key={}, + rotation_tolerance=2.0 # 0.0=deterministic, 2.0=recommended random ) ``` @@ -89,6 +96,17 @@ client = RotatingClient( - `whitelist_models` (`Optional[Dict[str, List[str]]]`, default: `None`): A dictionary where keys are provider names and values are lists of model names/patterns to always include, overriding `ignore_models`. - `enable_request_logging` (`bool`, default: `False`): If `True`, enables detailed per-request file logging (useful for debugging complex interactions). - `max_concurrent_requests_per_key` (`Optional[Dict[str, int]]`, default: `None`): A dictionary defining the maximum number of concurrent requests allowed for a single API key for a specific provider. Defaults to 1 if not specified. +- `rotation_tolerance` (`float`, default: `0.0`): Controls credential rotation strategy: + - `0.0`: **Deterministic** - Always selects the least-used credential for perfect load balance. + - `2.0` (default, recommended): **Weighted Random** - Randomly selects credentials with bias toward less-used ones. Provides unpredictability (harder to fingerprint) while maintaining good balance. + - `5.0+`: **High Randomness** - Even heavily-used credentials have significant selection probability. Maximum unpredictability. + + The weight formula is: `weight = (max_usage - credential_usage) + tolerance + 1` + + **Use Cases:** + - `0.0`: When perfect load balance is critical + - `2.0`: When avoiding fingerprinting/rate limit detection is important + - `5.0+`: For stress testing or maximum unpredictability ### Concurrency and Resource Management @@ -185,9 +203,27 @@ Use this tool to: ### Google Gemini (CLI) - **Auth**: Simulates the Google Cloud CLI authentication flow. -- **Project Discovery**: Automatically discovers the default Google Cloud Project ID. +- **Project Discovery**: Automatically discovers the default Google Cloud Project ID with enhanced onboarding flow. +- **Credential Prioritization**: Automatic detection and prioritization of paid vs free tier credentials. +- **Model Tier Requirements**: Gemini 3 models automatically filtered to paid-tier credentials only. +- **Gemini 3 Support**: Full support for Gemini 3 models with: + - `thinkingLevel` configuration (low/high) + - Tool hallucination prevention via system instruction injection + - ThoughtSignature caching for multi-turn conversations + - Parameter signature injection into tool descriptions - **Rate Limits**: Implements smart fallback strategies (e.g., switching from `gemini-1.5-pro` to `gemini-1.5-pro-002`) when rate limits are hit. +### Antigravity +- **Auth**: Uses OAuth 2.0 flow similar to Gemini CLI, with Antigravity-specific credentials and scopes. +- **Models**: Supports Gemini 2.5 (Pro/Flash), Gemini 3 (Pro/Image), and Claude Sonnet 4.5 via Google's internal Antigravity API. +- **Thought Signature Caching**: Server-side caching of `thoughtSignature` data for multi-turn conversations with Gemini 3 models. +- **Tool Hallucination Prevention**: Automatic injection of system instructions and parameter signatures for Gemini 3 to prevent tool parameter hallucination. +- **Thinking Support**: + - Gemini 2.5: Uses `thinkingBudget` (integer tokens) + - Gemini 3: Uses `thinkingLevel` (string: "low"/"high") + - Claude: Uses `thinkingBudget` via Antigravity proxy +- **Base URL Fallback**: Automatic fallback between sandbox and production endpoints. + ## Error Handling and Cooldowns The client uses a sophisticated error handling mechanism: From 7830a78a3b28b1fc0624071a819de1b3042558fb Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 10:06:25 +0100 Subject: [PATCH 034/221] =?UTF-8?q?refactor(credential-tool):=20?= =?UTF-8?q?=F0=9F=94=A8=20add=20export=20submenu=20for=20credential=20mana?= =?UTF-8?q?gement?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduced a new submenu for exporting credentials to .env format to improve user experience and code organization. - Add `export_credentials_submenu()` function to consolidate all export options - Implement `export_antigravity_to_env()` for Antigravity credential export - Refactor main menu to replace individual export options (3, 4, 5) with single "Export Credentials" option - Maintain consistent UI/UX patterns across all export functions - Generate .env files with metadata headers and timestamp information This change improves menu navigation by reducing clutter in the main menu and grouping related export functionality together. --- src/rotator_library/credential_tool.py | 157 ++++++++++++++++++++++--- 1 file changed, 140 insertions(+), 17 deletions(-) diff --git a/src/rotator_library/credential_tool.py b/src/rotator_library/credential_tool.py index a1705a13..066befe3 100644 --- a/src/rotator_library/credential_tool.py +++ b/src/rotator_library/credential_tool.py @@ -533,6 +533,143 @@ async def export_iflow_to_env(): console.print(Panel(f"An error occurred during export: {e}", style="bold red", title="Error")) +async def export_antigravity_to_env(): + """ + Export an Antigravity credential JSON file to .env format. + Generates one .env file per credential. + """ + console.print(Panel("[bold cyan]Export Antigravity Credential to .env[/bold cyan]", expand=False)) + + # Find all antigravity credentials + antigravity_files = list(OAUTH_BASE_DIR.glob("antigravity_oauth_*.json")) + + if not antigravity_files: + console.print(Panel("No Antigravity credentials found. Please add one first using 'Add OAuth Credential'.", + style="bold red", title="No Credentials")) + return + + # Display available credentials + cred_text = Text() + for i, cred_file in enumerate(antigravity_files): + try: + with open(cred_file, 'r') as f: + creds = json.load(f) + email = creds.get("_proxy_metadata", {}).get("email", "unknown") + cred_text.append(f" {i + 1}. {cred_file.name} ({email})\n") + except Exception as e: + cred_text.append(f" {i + 1}. {cred_file.name} (error reading: {e})\n") + + console.print(Panel(cred_text, title="Available Antigravity Credentials", style="bold blue")) + + choice = Prompt.ask( + Text.from_markup("[bold]Please select a credential to export or type [red]'b'[/red] to go back[/bold]"), + choices=[str(i + 1) for i in range(len(antigravity_files))] + ["b"], + show_choices=False + ) + + if choice.lower() == 'b': + return + + try: + choice_index = int(choice) - 1 + if 0 <= choice_index < len(antigravity_files): + cred_file = antigravity_files[choice_index] + + # Load the credential + with open(cred_file, 'r') as f: + creds = json.load(f) + + # Extract metadata + email = creds.get("_proxy_metadata", {}).get("email", "unknown") + + # Generate .env file name + safe_email = email.replace("@", "_at_").replace(".", "_") + env_filename = f"antigravity_{safe_email}.env" + env_filepath = OAUTH_BASE_DIR / env_filename + + # Build .env content + env_lines = [ + f"# Antigravity Credential for: {email}", + f"# Generated from: {cred_file.name}", + f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", + "", + f"ANTIGRAVITY_ACCESS_TOKEN={creds.get('access_token', '')}", + f"ANTIGRAVITY_REFRESH_TOKEN={creds.get('refresh_token', '')}", + f"ANTIGRAVITY_EXPIRY_DATE={creds.get('expiry_date', 0)}", + f"ANTIGRAVITY_CLIENT_ID={creds.get('client_id', '')}", + f"ANTIGRAVITY_CLIENT_SECRET={creds.get('client_secret', '')}", + f"ANTIGRAVITY_TOKEN_URI={creds.get('token_uri', 'https://oauth2.googleapis.com/token')}", + f"ANTIGRAVITY_UNIVERSE_DOMAIN={creds.get('universe_domain', 'googleapis.com')}", + f"ANTIGRAVITY_EMAIL={email}", + ] + + # Write to .env file + with open(env_filepath, 'w') as f: + f.write('\n'.join(env_lines)) + + success_text = Text.from_markup( + f"Successfully exported credential to [bold yellow]'{env_filepath}'[/bold yellow]\n\n" + f"To use this credential:\n" + f"1. Copy [bold yellow]{env_filepath.name}[/bold yellow] to your deployment environment\n" + f"2. Load the variables: [bold cyan]export $(cat {env_filepath.name} | grep -v '^#' | xargs)[/bold cyan]\n" + f"3. Or source it: [bold cyan]source {env_filepath.name}[/bold cyan]\n" + f"4. The Antigravity provider will automatically use these environment variables" + ) + console.print(Panel(success_text, style="bold green", title="Success")) + else: + console.print("[bold red]Invalid choice. Please try again.[/bold red]") + except ValueError: + console.print("[bold red]Invalid input. Please enter a number or 'b'.[/bold red]") + except Exception as e: + console.print(Panel(f"An error occurred during export: {e}", style="bold red", title="Error")) + + +async def export_credentials_submenu(): + """ + Submenu for credential export options. + """ + while True: + console.clear() + console.print(Panel("[bold cyan]Export Credentials to .env[/bold cyan]", title="--- API Key Proxy ---", expand=False)) + + console.print(Panel( + Text.from_markup( + "1. Export Gemini CLI credential\n" + "2. Export Qwen Code credential\n" + "3. Export iFlow credential\n" + "4. Export Antigravity credential" + ), + title="Choose credential type to export", + style="bold blue" + )) + + export_choice = Prompt.ask( + Text.from_markup("[bold]Please select an option or type [red]'b'[/red] to go back[/bold]"), + choices=["1", "2", "3", "4", "b"], + show_choices=False + ) + + if export_choice.lower() == 'b': + break + + if export_choice == "1": + await export_gemini_cli_to_env() + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "2": + await export_qwen_code_to_env() + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "3": + await export_iflow_to_env() + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "4": + await export_antigravity_to_env() + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + + async def main(clear_on_start=True): """ An interactive CLI tool to add new credentials. @@ -556,9 +693,7 @@ async def main(clear_on_start=True): Text.from_markup( "1. Add OAuth Credential\n" "2. Add API Key\n" - "3. Export Gemini CLI credential to .env\n" - "4. Export Qwen Code credential to .env\n" - "5. Export iFlow credential to .env" + "3. Export Credentials" ), title="Choose credential type", style="bold blue" @@ -566,7 +701,7 @@ async def main(clear_on_start=True): setup_type = Prompt.ask( Text.from_markup("[bold]Please select an option or type [red]'q'[/red] to quit[/bold]"), - choices=["1", "2", "3", "4", "5", "q"], + choices=["1", "2", "3", "q"], show_choices=False ) @@ -622,19 +757,7 @@ async def main(clear_on_start=True): input() elif setup_type == "3": - await export_gemini_cli_to_env() - console.print("\n[dim]Press Enter to return to main menu...[/dim]") - input() - - elif setup_type == "4": - await export_qwen_code_to_env() - console.print("\n[dim]Press Enter to return to main menu...[/dim]") - input() - - elif setup_type == "5": - await export_iflow_to_env() - console.print("\n[dim]Press Enter to return to main menu...[/dim]") - input() + await export_credentials_submenu() def run_credential_tool(from_launcher=False): """ From 62e7cf33f3c0f9ea598f62854b54c567998ce2c9 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 14:34:05 +0100 Subject: [PATCH 035/221] One huge ass bugfix i can't even list here. It's a mess i'll fix later --- src/proxy_app/main.py | 3 + src/rotator_library/client.py | 62 +++- .../providers/antigravity_provider.py | 300 ++++++++++++++---- .../providers/gemini_cli_provider.py | 57 +++- 4 files changed, 347 insertions(+), 75 deletions(-) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index b5cacd31..43b2d2d3 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -608,6 +608,9 @@ async def streaming_response_wrapper( # --- Final Response Construction --- if aggregated_tool_calls: final_message["tool_calls"] = list(aggregated_tool_calls.values()) + # CRITICAL FIX: Override finish_reason when tool_calls exist + # This ensures OpenCode and other agentic systems continue the conversation loop + finish_reason = "tool_calls" # Ensure standard fields are present for consistent logging for field in ["content", "tool_calls", "function_call"]: diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index bfd3be5a..7fa50806 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -495,11 +495,19 @@ async def _safe_streaming_wrapper( """ A hybrid wrapper for streaming that buffers fragmented JSON, handles client disconnections gracefully, and distinguishes between content and streamed errors. + + FINISH_REASON HANDLING: + Providers just translate chunks - this wrapper handles ALL finish_reason logic: + 1. Strip finish_reason from intermediate chunks (litellm defaults to "stop") + 2. Track accumulated_finish_reason with priority: tool_calls > length/content_filter > stop + 3. Only emit finish_reason on final chunk (detected by usage.completion_tokens > 0) """ last_usage = None stream_completed = False stream_iterator = stream.__aiter__() json_buffer = "" + accumulated_finish_reason = None # Track strongest finish_reason across chunks + has_tool_calls = False # Track if ANY tool calls were seen in stream try: while True: @@ -507,26 +515,64 @@ async def _safe_streaming_wrapper( lib_logger.info( f"Client disconnected. Aborting stream for credential ...{key[-6:]}." ) - # Do not yield [DONE] because the client is gone. - # The 'finally' block will handle key release. break try: chunk = await stream_iterator.__anext__() if json_buffer: - # If we are about to discard a buffer, it means data was likely lost. - # Log this as a warning to make it visible. lib_logger.warning( f"Discarding incomplete JSON buffer from previous chunk: {json_buffer}" ) json_buffer = "" - yield f"data: {json.dumps(chunk.dict())}\n\n" + # Convert chunk to dict, handling both litellm.ModelResponse and raw dicts + if hasattr(chunk, "dict"): + chunk_dict = chunk.dict() + elif hasattr(chunk, "model_dump"): + chunk_dict = chunk.model_dump() + else: + chunk_dict = chunk + + # === FINISH_REASON LOGIC === + # Providers send raw chunks without finish_reason logic. + # This wrapper determines finish_reason based on accumulated state. + if "choices" in chunk_dict and chunk_dict["choices"]: + choice = chunk_dict["choices"][0] + delta = choice.get("delta", {}) + usage = chunk_dict.get("usage", {}) + + # Track tool_calls across ALL chunks - if we ever see one, finish_reason must be tool_calls + if delta.get("tool_calls"): + has_tool_calls = True + accumulated_finish_reason = "tool_calls" + + # Detect final chunk: has usage with completion_tokens > 0 + has_completion_tokens = ( + usage and + isinstance(usage, dict) and + usage.get("completion_tokens", 0) > 0 + ) + + if has_completion_tokens: + # FINAL CHUNK: Determine correct finish_reason + if has_tool_calls: + # Tool calls always win + choice["finish_reason"] = "tool_calls" + elif accumulated_finish_reason: + # Use accumulated reason (length, content_filter, etc.) + choice["finish_reason"] = accumulated_finish_reason + else: + # Default to stop + choice["finish_reason"] = "stop" + else: + # INTERMEDIATE CHUNK: Never emit finish_reason + # (litellm.ModelResponse defaults to "stop" which is wrong) + choice["finish_reason"] = None + + yield f"data: {json.dumps(chunk_dict)}\n\n" if hasattr(chunk, "usage") and chunk.usage: - last_usage = ( - chunk.usage - ) # Overwrite with the latest (cumulative) + last_usage = chunk.usage except StopAsyncIteration: stream_completed = True diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 0fa11faa..28a9f694 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -16,7 +16,6 @@ from __future__ import annotations -import asyncio import copy import hashlib import json @@ -58,7 +57,7 @@ #"gemini-2.5-pro", #"gemini-2.5-flash", #"gemini-2.5-flash-lite", - "gemini-3-pro-preview", + "gemini-3-pro-preview", # Internally mapped to -low/-high variant based on thinkingLevel #"gemini-3-pro-image-preview", #"gemini-2.5-computer-use-preview-10-2025", "claude-sonnet-4-5", # Internally mapped to -thinking variant when reasoning_effort is provided @@ -71,12 +70,13 @@ MODEL_ALIAS_MAP = { "rev19-uic3-1p": "gemini-2.5-computer-use-preview-10-2025", "gemini-3-pro-image": "gemini-3-pro-image-preview", + "gemini-3-pro-low": "gemini-3-pro-preview", "gemini-3-pro-high": "gemini-3-pro-preview", } MODEL_ALIAS_REVERSE = {v: k for k, v in MODEL_ALIAS_MAP.items()} # Models to exclude from dynamic discovery -EXCLUDED_MODELS = {"chat_20706", "chat_23310", "gemini-2.5-flash-thinking", "gemini-3-pro-low", "gemini-2.5-pro"} +EXCLUDED_MODELS = {"chat_20706", "chat_23310", "gemini-2.5-flash-thinking", "gemini-2.5-pro"} # Gemini finish reason mapping FINISH_REASON_MAP = { @@ -101,15 +101,28 @@ 1. DO NOT use your internal training data to guess tool parameters 2. ONLY use the exact parameter structure defined in the tool schema -3. If a tool takes a 'files' parameter, it is ALWAYS an array of objects with specific properties, NEVER a simple array of strings -4. If a tool edits code, it takes structured JSON objects with specific fields, NEVER raw diff strings or plain text -5. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') -6. Array parameters have specific item types - check the schema's 'items' field for the exact structure -7. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions +3. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') +4. Array parameters have specific item types - check the schema's 'items' field for the exact structure +5. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions If you are unsure about a tool's parameters, YOU MUST read the schema definition carefully. Your training data about common tool names like 'read_file' or 'apply_diff' does NOT apply here. """ +# Claude tool fix system instruction (prevents hallucination) +DEFAULT_CLAUDE_SYSTEM_INSTRUCTION = """CRITICAL TOOL USAGE INSTRUCTIONS: +You are operating in a custom environment where tool definitions differ from your training data. +You MUST follow these rules strictly: + +1. DO NOT use your internal training data to guess tool parameters +2. ONLY use the exact parameter structure defined in the tool schema +3. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') +4. Array parameters have specific item types - check the schema's 'items' field for the exact structure +5. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions +6. Tool use in agentic workflows is REQUIRED - you must call tools with the exact parameters specified in the schema + +If you are unsure about a tool's parameters, YOU MUST read the schema definition carefully. +""" + # ============================================================================= # HELPER FUNCTIONS @@ -169,8 +182,9 @@ def _recursively_parse_json_strings(obj: Any) -> Any: Antigravity sometimes returns tool arguments with JSON-stringified values: {"files": "[{...}]"} instead of {"files": [{...}]}. - Additionally handles malformed double-encoded JSON where Antigravity - returns strings like '[{...}]}' (extra trailing '}'). + Additionally handles: + - Malformed double-encoded JSON (extra trailing '}' or ']') + - Escaped string content (\n, \t, \", etc.) """ if isinstance(obj, dict): return {k: _recursively_parse_json_strings(v) for k, v in obj.items()} @@ -178,6 +192,23 @@ def _recursively_parse_json_strings(obj: Any) -> Any: return [_recursively_parse_json_strings(item) for item in obj] elif isinstance(obj, str): stripped = obj.strip() + + # Check if string contains common escape sequences that need unescaping + # This handles cases where diff content or other text has literal \n instead of newlines + if '\\n' in obj or '\\t' in obj or '\\"' in obj or '\\\\' in obj: + try: + # Use json.loads with quotes to properly unescape the string + # This converts \n -> newline, \t -> tab, \" -> quote, etc. + unescaped = json.loads(f'"{obj}"') + lib_logger.debug( + f"[Antigravity] Unescaped string content: " + f"{len(obj) - len(unescaped)} chars changed" + ) + return unescaped + except (json.JSONDecodeError, ValueError): + # If unescaping fails, continue with original processing + pass + # Check if it looks like JSON (starts with { or [) if stripped and stripped[0] in ('{', '['): # Try standard parsing first @@ -215,7 +246,7 @@ def _recursively_parse_json_strings(obj: Any) -> Any: cleaned = stripped[:last_brace+1] parsed = json.loads(cleaned) lib_logger.warning( - f"Auto-corrected malformed JSON string: " + f"[Antigravity] Auto-corrected malformed JSON string: " f"truncated {len(stripped) - len(cleaned)} extra chars" ) return _recursively_parse_json_strings(parsed) @@ -369,6 +400,7 @@ def __init__(self): self._enable_signature_cache = _env_bool("ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", True) self._enable_dynamic_models = _env_bool("ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", False) self._enable_gemini3_tool_fix = _env_bool("ANTIGRAVITY_GEMINI3_TOOL_FIX", True) + self._enable_claude_tool_fix = _env_bool("ANTIGRAVITY_CLAUDE_TOOL_FIX", True) # Gemini 3 tool fix configuration self._gemini3_tool_prefix = os.getenv("ANTIGRAVITY_GEMINI3_TOOL_PREFIX", "gemini3_") @@ -381,6 +413,16 @@ def __init__(self): DEFAULT_GEMINI3_SYSTEM_INSTRUCTION ) + # Claude tool fix configuration (separate from Gemini 3) + self._claude_description_prompt = os.getenv( + "ANTIGRAVITY_CLAUDE_DESCRIPTION_PROMPT", + "\n\nSTRICT PARAMETERS: {params}." + ) + self._claude_system_instruction = os.getenv( + "ANTIGRAVITY_CLAUDE_SYSTEM_INSTRUCTION", + DEFAULT_CLAUDE_SYSTEM_INSTRUCTION + ) + # Log configuration self._log_config() @@ -389,7 +431,7 @@ def _log_config(self) -> None: lib_logger.debug( f"Antigravity config: signatures_in_client={self._preserve_signatures_in_client}, " f"cache={self._enable_signature_cache}, dynamic_models={self._enable_dynamic_models}, " - f"gemini3_fix={self._enable_gemini3_tool_fix}" + f"gemini3_fix={self._enable_gemini3_tool_fix}, claude_fix={self._enable_claude_tool_fix}" ) # ========================================================================= @@ -558,7 +600,10 @@ def _transform_messages( if msg.get("role") == "assistant" and msg.get("tool_calls"): for tc in msg["tool_calls"]: if tc.get("type") == "function": - tool_id_to_name[tc["id"]] = tc["function"]["name"] + tc_id = tc["id"] + tc_name = tc["function"]["name"] + tool_id_to_name[tc_id] = tc_name + #lib_logger.debug(f"[ID Mapping] Registered tool_call: id={tc_id}, name={tc_name}") # Convert each message for msg in messages: @@ -654,6 +699,11 @@ def _transform_assistant_message( tool_id = tc.get("id", "") func_name = tc["function"]["name"] + #lib_logger.debug( + # f"[ID Transform] Converting assistant tool_call to functionCall: " + # f"id={tool_id}, name={func_name}" + #) + # Add prefix for Gemini 3 if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: func_name = f"{self._gemini3_tool_prefix}{func_name}" @@ -728,6 +778,15 @@ def _transform_tool_message( func_name = tool_id_to_name.get(tool_id, "unknown_function") content = msg.get("content", "{}") + # Log ID lookup + if tool_id not in tool_id_to_name: + lib_logger.warning( + f"[ID Mismatch] Tool response has ID '{tool_id}' which was not found in tool_id_to_name map. " + f"Available IDs: {list(tool_id_to_name.keys())}" + ) + #else: + #lib_logger.debug(f"[ID Mapping] Tool response matched: id={tool_id}, name={func_name}") + # Add prefix for Gemini 3 if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: func_name = f"{self._gemini3_tool_prefix}{func_name}" @@ -758,10 +817,12 @@ def _fix_tool_response_grouping( Converts linear format (call, response, call, response) to grouped format (model with calls, user with all responses). + + IMPORTANT: Preserves ID-based pairing to prevent mismatches. """ new_contents = [] - pending_groups = [] - collected_responses = [] + pending_groups = [] # List of {"ids": [id1, id2, ...], "call_indices": [...]} + collected_responses = {} # Dict mapping ID -> response_part for content in contents: role = content.get("role") @@ -770,15 +831,33 @@ def _fix_tool_response_grouping( response_parts = [p for p in parts if "functionResponse" in p] if response_parts: - collected_responses.extend(response_parts) + # Collect responses by ID (ignore duplicates - keep first occurrence) + for resp in response_parts: + resp_id = resp.get("functionResponse", {}).get("id", "") + if resp_id: + if resp_id in collected_responses: + lib_logger.warning( + f"[Grouping] Duplicate response ID detected: {resp_id}. " + f"Ignoring duplicate - this may indicate malformed conversation history." + ) + continue + #lib_logger.debug(f"[Grouping] Collected response for ID: {resp_id}") + collected_responses[resp_id] = resp - # Try to satisfy pending groups + # Try to satisfy pending groups (newest first) for i in range(len(pending_groups) - 1, -1, -1): group = pending_groups[i] - if len(collected_responses) >= group["count"]: - group_responses = collected_responses[:group["count"]] - collected_responses = collected_responses[group["count"]:] + group_ids = group["ids"] + + # Check if we have ALL responses for this group + if all(gid in collected_responses for gid in group_ids): + # Extract responses in the same order as the function calls + group_responses = [collected_responses.pop(gid) for gid in group_ids] new_contents.append({"parts": group_responses, "role": "user"}) + #lib_logger.debug( + # f"[Grouping] Satisfied group with {len(group_responses)} responses: " + # f"ids={group_ids}" + #) pending_groups.pop(i) break continue @@ -787,16 +866,32 @@ def _fix_tool_response_grouping( func_calls = [p for p in parts if "functionCall" in p] new_contents.append(content) if func_calls: - pending_groups.append({"count": len(func_calls)}) + call_ids = [fc.get("functionCall", {}).get("id", "") for fc in func_calls] + call_ids = [cid for cid in call_ids if cid] # Filter empty IDs + if call_ids: + lib_logger.debug(f"[Grouping] Created pending group expecting {len(call_ids)} responses: ids={call_ids}") + pending_groups.append({"ids": call_ids, "call_indices": list(range(len(func_calls)))}) else: new_contents.append(content) - # Handle remaining groups + # Handle remaining groups (shouldn't happen in well-formed conversations) for group in pending_groups: - if len(collected_responses) >= group["count"]: - group_responses = collected_responses[:group["count"]] - collected_responses = collected_responses[group["count"]:] + group_ids = group["ids"] + available_ids = [gid for gid in group_ids if gid in collected_responses] + if available_ids: + group_responses = [collected_responses.pop(gid) for gid in available_ids] new_contents.append({"parts": group_responses, "role": "user"}) + lib_logger.warning( + f"[Grouping] Partial group satisfaction: expected {len(group_ids)}, " + f"got {len(available_ids)} responses" + ) + + # Warn about unmatched responses + if collected_responses: + lib_logger.warning( + f"[Grouping] {len(collected_responses)} unmatched responses remaining: " + f"ids={list(collected_responses.keys())}" + ) return new_contents @@ -823,12 +918,16 @@ def _apply_gemini3_namespace( def _inject_signature_into_descriptions( self, - tools: List[Dict[str, Any]] + tools: List[Dict[str, Any]], + description_prompt: Optional[str] = None ) -> List[Dict[str, Any]]: - """Inject parameter signatures into tool descriptions for Gemini 3.""" + """Inject parameter signatures into tool descriptions for Gemini 3 & Claude.""" if not tools: return tools + # Use provided prompt or default to Gemini 3 prompt + prompt_template = description_prompt or self._gemini3_description_prompt + modified = copy.deepcopy(tools) for tool in modified: for func_decl in tool.get("functionDeclarations", []): @@ -854,7 +953,7 @@ def _inject_signature_into_descriptions( ) if param_list: - sig_str = self._gemini3_description_prompt.replace( + sig_str = prompt_template.replace( "{params}", ", ".join(param_list) ) func_decl["description"] = func_decl.get("description", "") + sig_str @@ -892,6 +991,42 @@ def _strip_gemini3_prefix(self, name: str) -> str: return name[len(self._gemini3_tool_prefix):] return name + def _translate_tool_choice(self, tool_choice: Union[str, Dict[str, Any]], model: str = "") -> Optional[Dict[str, Any]]: + """ + Translates OpenAI's `tool_choice` to Gemini's `toolConfig`. + Handles Gemini 3 namespace prefixes for specific tool selection. + """ + if not tool_choice: + return None + + config = {} + mode = "AUTO" # Default to auto + is_gemini_3 = self._is_gemini_3(model) + + if isinstance(tool_choice, str): + if tool_choice == "auto": + mode = "AUTO" + elif tool_choice == "none": + mode = "NONE" + elif tool_choice == "required": + mode = "ANY" + elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function": + function_name = tool_choice.get("function", {}).get("name") + if function_name: + # Add Gemini 3 prefix if needed + if is_gemini_3 and self._enable_gemini3_tool_fix: + function_name = f"{self._gemini3_tool_prefix}{function_name}" + + mode = "ANY" # Force a call, but only to this function + config["functionCallingConfig"] = { + "mode": mode, + "allowedFunctionNames": [function_name] + } + return config + + config["functionCallingConfig"] = {"mode": mode} + return config + # ========================================================================= # REQUEST TRANSFORMATION # ========================================================================= @@ -936,7 +1071,8 @@ def _transform_to_antigravity_format( gemini_payload: Dict[str, Any], model: str, max_tokens: Optional[int] = None, - reasoning_effort: Optional[str] = None + reasoning_effort: Optional[str] = None, + tool_choice: Optional[Union[str, Dict[str, Any]]] = None ) -> Dict[str, Any]: """ Transform Gemini CLI payload to complete Antigravity format. @@ -954,6 +1090,16 @@ def _transform_to_antigravity_format( if internal_model == "claude-sonnet-4-5" and not internal_model.endswith("-thinking"): internal_model = "claude-sonnet-4-5-thinking" + # Map gemini-3-pro-preview to -low/-high variant based on thinking config + if model == "gemini-3-pro-preview" or internal_model == "gemini-3-pro-preview": + # Check thinking config to determine variant + thinking_config = gemini_payload.get("generationConfig", {}).get("thinkingConfig", {}) + thinking_level = thinking_config.get("thinkingLevel", "high") + if thinking_level == "low": + internal_model = "gemini-3-pro-low" + else: + internal_model = "gemini-3-pro-high" + # Wrap in Antigravity envelope antigravity_payload = { "project": _generate_project_id(), @@ -983,10 +1129,15 @@ def _transform_to_antigravity_format( antigravity_payload["request"]["generationConfig"] = gen_config - # Set toolConfig mode - tool_config = antigravity_payload["request"].setdefault("toolConfig", {}) - func_config = tool_config.setdefault("functionCallingConfig", {}) - func_config["mode"] = "VALIDATED" + # Set toolConfig based on tool_choice parameter + tool_config_result = self._translate_tool_choice(tool_choice, model) + if tool_config_result: + antigravity_payload["request"]["toolConfig"] = tool_config_result + else: + # Default to AUTO if no tool_choice specified + tool_config = antigravity_payload["request"].setdefault("toolConfig", {}) + func_config = tool_config.setdefault("functionCallingConfig", {}) + func_config["mode"] = "AUTO" # Handle Gemini 3 thinking logic if not internal_model.startswith("gemini-3-"): @@ -1053,7 +1204,8 @@ def _gemini_to_openai_chunk( reasoning_content = "" tool_calls = [] first_sig_seen = False - tool_idx = 0 + # Use accumulator's tool_idx if available, otherwise use local counter + tool_idx = accumulator.get("tool_idx", 0) if accumulator else 0 for part in content_parts: has_func = "functionCall" in part @@ -1099,23 +1251,29 @@ def _gemini_to_openai_chunk( if tool_calls: delta["tool_calls"] = tool_calls delta["role"] = "assistant" + # Update tool_idx for next chunk + if accumulator is not None: + accumulator["tool_idx"] = tool_idx elif text_content or reasoning_content: delta["role"] = "assistant" - # Handle finish reason - finish_reason = self._map_finish_reason(candidate.get("finishReason"), bool(tool_calls)) - if finish_reason and accumulator is not None: + # Build usage if present + usage = self._build_usage(chunk.get("usageMetadata", {})) + + # Mark completion when we see usageMetadata + if chunk.get("usageMetadata") and accumulator is not None: accumulator["is_complete"] = True - # Build usage - usage = self._build_usage(chunk.get("usageMetadata", {})) + # Build choice - just translate, don't include finish_reason + # Client will handle finish_reason logic + choice = {"index": 0, "delta": delta} response = { "id": chunk.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), "object": "chat.completion.chunk", "created": int(time.time()), "model": model, - "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}] + "choices": [choice] } if usage: @@ -1188,12 +1346,13 @@ def _gemini_to_openai_non_streaming( finish_reason = self._map_finish_reason(candidate.get("finishReason"), bool(tool_calls)) usage = self._build_usage(response.get("usageMetadata", {})) + # For non-streaming, always include finish_reason (should always be present) result = { "id": response.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), "object": "chat.completion", "created": int(time.time()), "model": model, - "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}] + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason or "stop"}] } if usage: @@ -1212,6 +1371,8 @@ def _extract_tool_call( func_call = part["functionCall"] tool_id = func_call.get("id") or f"call_{uuid.uuid4().hex[:24]}" + #lib_logger.debug(f"[ID Extraction] Extracting tool call: id={tool_id}, raw_id={func_call.get('id')}") + tool_name = func_call.get("name", "") if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: tool_name = self._strip_gemini3_prefix(tool_name) @@ -1383,6 +1544,7 @@ async def acompletion( stream = kwargs.get("stream", False) credential_path = kwargs.pop("credential_identifier", kwargs.get("api_key", "")) tools = kwargs.get("tools") + tool_choice = kwargs.get("tool_choice") reasoning_effort = kwargs.get("reasoning_effort") top_p = kwargs.get("top_p") max_tokens = kwargs.get("max_tokens") @@ -1402,9 +1564,12 @@ async def acompletion( if system_instruction: gemini_payload["system_instruction"] = system_instruction - # Inject Gemini 3 system instruction - if self._is_gemini_3(model) and self._enable_gemini3_tool_fix and tools: - self._inject_gemini3_system_instruction(gemini_payload) + # Inject tool usage hardening system instructions + if tools: + if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: + self._inject_tool_hardening_instruction(gemini_payload, self._gemini3_system_instruction) + elif self._is_claude(model) and self._enable_claude_tool_fix: + self._inject_tool_hardening_instruction(gemini_payload, self._claude_system_instruction) # Add generation config gen_config = {} @@ -1423,13 +1588,23 @@ async def acompletion( if gemini_tools: gemini_payload["tools"] = gemini_tools - # Apply Gemini 3 tool transformations + # Apply tool transformations if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: + # Gemini 3: namespace prefix + parameter signatures gemini_payload["tools"] = self._apply_gemini3_namespace(gemini_payload["tools"]) - gemini_payload["tools"] = self._inject_signature_into_descriptions(gemini_payload["tools"]) + gemini_payload["tools"] = self._inject_signature_into_descriptions( + gemini_payload["tools"], + self._gemini3_description_prompt + ) + elif self._is_claude(model) and self._enable_claude_tool_fix: + # Claude: parameter signatures only (no namespace prefix) + gemini_payload["tools"] = self._inject_signature_into_descriptions( + gemini_payload["tools"], + self._claude_description_prompt + ) # Transform to Antigravity format - payload = self._transform_to_antigravity_format(gemini_payload, model, max_tokens, reasoning_effort) + payload = self._transform_to_antigravity_format(gemini_payload, model, max_tokens, reasoning_effort, tool_choice) file_logger.log_request(payload) # Make API call @@ -1467,12 +1642,12 @@ async def acompletion( return await self._handle_non_streaming(client, url, headers, payload, model, file_logger) raise - def _inject_gemini3_system_instruction(self, payload: Dict[str, Any]) -> None: - """Inject Gemini 3 system instruction for tool fix.""" - if not self._gemini3_system_instruction: + def _inject_tool_hardening_instruction(self, payload: Dict[str, Any], instruction_text: str) -> None: + """Inject tool usage hardening system instruction for Gemini 3 & Claude.""" + if not instruction_text: return - instruction_part = {"text": self._gemini3_system_instruction} + instruction_part = {"text": instruction_text} if "system_instruction" in payload: existing = payload["system_instruction"] @@ -1518,13 +1693,15 @@ async def _handle_streaming( file_logger: Optional[AntigravityFileLogger] = None ) -> AsyncGenerator[litellm.ModelResponse, None]: """Handle streaming completion.""" + # Accumulator tracks state across chunks for caching and tool indexing accumulator = { "reasoning_content": "", "thought_signature": "", "text_content": "", "tool_calls": [], - "is_complete": False - } if self._is_claude(model) and self._enable_signature_cache else None + "tool_idx": 0, # Track tool call index across chunks + "is_complete": False # Track if we received usageMetadata + } async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response: if response.status_code >= 400: @@ -1556,8 +1733,23 @@ async def _handle_streaming( file_logger.log_error(f"Parse error: {data_str[:100]}") continue + # If stream ended without usageMetadata chunk, emit a final chunk with finish_reason + # Emit final chunk if stream ended without usageMetadata + # Client will determine the correct finish_reason based on accumulated state + if not accumulator.get("is_complete"): + final_chunk = { + "id": f"chatcmpl-{uuid.uuid4().hex[:24]}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": None}], + # Include minimal usage to signal this is the final chunk + "usage": {"prompt_tokens": 0, "completion_tokens": 1, "total_tokens": 1} + } + yield litellm.ModelResponse(**final_chunk) + # Cache Claude thinking after stream completes - if accumulator and accumulator.get("reasoning_content"): + if self._is_claude(model) and self._enable_signature_cache and accumulator.get("reasoning_content"): self._cache_thinking( accumulator["reasoning_content"], accumulator["thought_signature"], diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 3ea9c4ea..32e54f3f 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -870,7 +870,6 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumul for part in parts: delta = {} - finish_reason = None has_func = 'functionCall' in part has_text = 'text' in part @@ -892,8 +891,11 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumul # Use provided ID or generate unique one with nanosecond precision tool_call_id = function_call.get('id') or f"call_{function_name}_{int(time.time() * 1_000_000_000)}" + # Get current tool index from accumulator (default 0) and increment + current_tool_idx = accumulator.get('tool_idx', 0) if accumulator else 0 + tool_call = { - "index": 0, + "index": current_tool_idx, "id": tool_call_id, "type": "function", "function": { @@ -915,6 +917,10 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumul tool_call["thought_signature"] = sig delta['tool_calls'] = [tool_call] + # Mark that we've sent tool calls and increment tool_idx + if accumulator is not None: + accumulator['has_tool_calls'] = True + accumulator['tool_idx'] = current_tool_idx + 1 elif has_text: # Use an explicit check for the 'thought' flag, as its type can be inconsistent @@ -926,14 +932,16 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumul if not delta: continue - raw_finish_reason = candidate.get('finishReason') - if raw_finish_reason: - finish_reason = FINISH_REASON_MAP.get(raw_finish_reason, 'stop') - # Use tool_calls if we have function calls - if delta.get('tool_calls'): - finish_reason = 'tool_calls' + # Mark that we have tool calls for accumulator tracking + # finish_reason determination is handled by the client + + # Mark stream complete if we have usageMetadata + is_final_chunk = 'usageMetadata' in response_data + if is_final_chunk and accumulator is not None: + accumulator['is_complete'] = True - choice = {"index": 0, "delta": delta, "finish_reason": finish_reason} + # Build choice - don't include finish_reason, let client handle it + choice = {"index": 0, "delta": delta} openai_chunk = { "choices": [choice], "model": model_id, "object": "chat.completion.chunk", @@ -1020,9 +1028,8 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> if "arguments" in delta["function_call"] and delta["function_call"]["arguments"] is not None: final_message["function_call"]["arguments"] += delta["function_call"]["arguments"] - # Get finish reason from the last chunk that has it - if choice.get("finish_reason"): - finish_reason = choice["finish_reason"] + # Note: chunks don't include finish_reason (client handles it) + # This is kept for compatibility but shouldn't trigger # Handle usage data from the last chunk that has it for chunk in reversed(chunks): @@ -1039,6 +1046,13 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> if field not in final_message: final_message[field] = None + # Determine finish_reason based on content (same logic as client.py) + # tool_calls wins, otherwise stop + if aggregated_tool_calls: + finish_reason = "tool_calls" + else: + finish_reason = "stop" + # Construct the final response final_choice = { "index": 0, @@ -1343,6 +1357,9 @@ async def do_call(attempt_model: str, is_fallback: bool = False): url = f"{CODE_ASSIST_ENDPOINT}:streamGenerateContent" async def stream_handler(): + # Track state across chunks for tool indexing + accumulator = {"has_tool_calls": False, "tool_idx": 0, "is_complete": False} + final_headers = auth_header.copy() final_headers.update({ "User-Agent": "google-api-nodejs-client/9.15.1", @@ -1362,10 +1379,24 @@ async def stream_handler(): if data_str == "[DONE]": break try: chunk = json.loads(data_str) - for openai_chunk in self._convert_chunk_to_openai(chunk, model): + for openai_chunk in self._convert_chunk_to_openai(chunk, model, accumulator): yield litellm.ModelResponse(**openai_chunk) except json.JSONDecodeError: lib_logger.warning(f"Could not decode JSON from Gemini CLI: {line}") + + # Emit final chunk if stream ended without usageMetadata + # Client will determine the correct finish_reason + if not accumulator.get("is_complete"): + final_chunk = { + "id": f"chatcmpl-geminicli-{time.time()}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": None}], + # Include minimal usage to signal this is the final chunk + "usage": {"prompt_tokens": 0, "completion_tokens": 1, "total_tokens": 1} + } + yield litellm.ModelResponse(**final_chunk) except httpx.HTTPStatusError as e: error_body = None From d4593e5bc9d89adfb9645dde1b45cc78669edaf8 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 15:24:30 +0100 Subject: [PATCH 036/221] =?UTF-8?q?fix(gemini):=20=F0=9F=90=9B=20consolida?= =?UTF-8?q?te=20parallel=20tool=20responses=20and=20improve=20rate=20limit?= =?UTF-8?q?=20handling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit addresses multiple issues with Gemini API providers related to parallel function calling and rate limit error handling: **Tool Response Consolidation:** - Parallel function responses are now consolidated into a single user message as required by Gemini API specification - Previously, consecutive tool responses were sent as separate messages, causing API errors - Implemented pending tool parts accumulation pattern in both GeminiCliProvider and AntigravityProvider - Tool responses are flushed when a non-tool message is encountered or at the end of message processing **Thought Signature Handling:** - Fixed parallel function call signature behavior to match Gemini 3 API requirements - Only the first parallel function call in a message receives a thoughtSignature field - Subsequent parallel calls no longer include thoughtSignature to prevent API validation errors - Removed `first_sig_seen` tracking flags since signatures are now stored per tool call **Rate Limit Error Handling:** - Added `extract_retry_after_from_body()` function to parse retry-after times from various API error formats - Improved Gemini CLI rate limit error messages with extracted retry-after information - Enhanced error logging to capture and display response bodies before raising HTTPStatusError - Reduced log noise by using debug level for rate limit rotation events instead of info/warning - Better error context propagation for 429 responses **Code Quality:** - Removed unused `first_sig_seen` tracking variables - Improved inline documentation explaining Gemini API parallel function call requirements - Consistent role mapping (tool -> user) across message transformation logic --- src/rotator_library/error_handler.py | 38 +++++++++ .../providers/antigravity_provider.py | 57 +++++++++---- .../providers/gemini_cli_provider.py | 79 +++++++++++++++---- 3 files changed, 143 insertions(+), 31 deletions(-) diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index 5298aec8..a3775f7f 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -17,6 +17,42 @@ ) +def extract_retry_after_from_body(error_body: Optional[str]) -> Optional[int]: + """ + Extract the retry-after time from an API error response body. + + Handles various error formats including: + - Gemini CLI: "Your quota will reset after 39s." + - Generic: "quota will reset after 120s", "retry after 60s" + + Args: + error_body: The raw error response body + + Returns: + The retry time in seconds, or None if not found + """ + if not error_body: + return None + + # Pattern to match various "reset after Xs" or "retry after Xs" formats + patterns = [ + r"quota will reset after\s*(\d+)s", + r"reset after\s*(\d+)s", + r"retry after\s*(\d+)s", + r"try again in\s*(\d+)\s*seconds?", + ] + + for pattern in patterns: + match = re.search(pattern, error_body, re.IGNORECASE) + if match: + try: + return int(match.group(1)) + except (ValueError, IndexError): + continue + + return None + + class NoAvailableKeysError(Exception): """Raised when no API keys are available for a request after waiting.""" @@ -106,6 +142,8 @@ def get_retry_after(error: Exception) -> Optional[int]: r"wait for\s*(\d+)\s*seconds?", r'"retryDelay":\s*"(\d+)s"', r"x-ratelimit-reset:?\s*(\d+)", + r"quota will reset after\s*(\d+)s", # Gemini CLI rate limit format + r"reset after\s*(\d+)s", # Generic reset after format ] for pattern in patterns: diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 28a9f694..55c28a8e 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -605,23 +605,38 @@ def _transform_messages( tool_id_to_name[tc_id] = tc_name #lib_logger.debug(f"[ID Mapping] Registered tool_call: id={tc_id}, name={tc_name}") - # Convert each message + # Convert each message, consolidating consecutive tool responses + # Per Gemini docs: parallel function responses must be in a single user message + pending_tool_parts = [] + for msg in messages: role = msg.get("role") content = msg.get("content") parts = [] + # Flush pending tool parts before non-tool message + if pending_tool_parts and role != "tool": + gemini_contents.append({"role": "user", "parts": pending_tool_parts}) + pending_tool_parts = [] + if role == "user": parts = self._transform_user_message(content) elif role == "assistant": parts = self._transform_assistant_message(msg, model, tool_id_to_name) elif role == "tool": - parts = self._transform_tool_message(msg, model, tool_id_to_name) + tool_parts = self._transform_tool_message(msg, model, tool_id_to_name) + # Accumulate tool responses instead of adding individually + pending_tool_parts.extend(tool_parts) + continue if parts: - gemini_role = "model" if role == "assistant" else "user" if role == "tool" else "user" + gemini_role = "model" if role == "assistant" else "user" gemini_contents.append({"role": gemini_role, "parts": parts}) + # Flush any remaining tool parts + if pending_tool_parts: + gemini_contents.append({"role": "user", "parts": pending_tool_parts}) + return system_instruction, gemini_contents def _parse_content_parts( @@ -687,6 +702,9 @@ def _transform_assistant_message( parts.append({"text": content}) # Add tool calls + # Track if we've seen the first function call in this message + # Per Gemini docs: Only the FIRST parallel function call gets a signature + first_func_in_msg = True for tc in tool_calls: if tc.get("type") != "function": continue @@ -717,6 +735,8 @@ def _transform_assistant_message( } # Add thoughtSignature for Gemini 3 + # Per Gemini docs: Only the FIRST parallel function call gets a signature. + # Subsequent parallel calls should NOT have a thoughtSignature field. if self._is_gemini_3(model): sig = tc.get("thought_signature") if not sig and tool_id and self._enable_signature_cache: @@ -724,9 +744,13 @@ def _transform_assistant_message( if sig: func_part["thoughtSignature"] = sig - else: + elif first_func_in_msg: + # Only add bypass to the first function call if no sig available func_part["thoughtSignature"] = "skip_thought_signature_validator" - lib_logger.warning(f"Missing thoughtSignature for {tool_id}, using bypass") + lib_logger.warning(f"Missing thoughtSignature for first func call {tool_id}, using bypass") + # Subsequent parallel calls: no signature field at all + + first_func_in_msg = False parts.append(func_part) @@ -1146,13 +1170,20 @@ def _transform_to_antigravity_format( del thinking_config["thinkingLevel"] thinking_config["thinkingBudget"] = -1 - # Add thoughtSignature to function calls for Gemini 3 + # Ensure first function call in each model message has a thoughtSignature for Gemini 3 + # Per Gemini docs: Only the FIRST parallel function call gets a signature if internal_model.startswith("gemini-3-"): for content in antigravity_payload["request"].get("contents", []): if content.get("role") == "model": + first_func_seen = False for part in content.get("parts", []): - if "functionCall" in part and "thoughtSignature" not in part: - part["thoughtSignature"] = "skip_thought_signature_validator" + if "functionCall" in part: + if not first_func_seen: + # First function call in this message - needs a signature + if "thoughtSignature" not in part: + part["thoughtSignature"] = "skip_thought_signature_validator" + first_func_seen = True + # Subsequent parallel calls: leave as-is (no signature) # Claude-specific tool schema transformation if internal_model.startswith("claude-sonnet-"): @@ -1203,7 +1234,6 @@ def _gemini_to_openai_chunk( text_content = "" reasoning_content = "" tool_calls = [] - first_sig_seen = False # Use accumulator's tool_idx if available, otherwise use local counter tool_idx = accumulator.get("tool_idx", 0) if accumulator else 0 @@ -1235,8 +1265,8 @@ def _gemini_to_openai_chunk( if has_func: tool_call = self._extract_tool_call(part, model, tool_idx, accumulator) - if has_sig and not first_sig_seen: - first_sig_seen = True + # Store signature for each tool call (needed for parallel tool calls) + if has_sig: self._handle_tool_signature(tool_call, part["thoughtSignature"]) tool_calls.append(tool_call) @@ -1298,7 +1328,6 @@ def _gemini_to_openai_non_streaming( reasoning_content = "" tool_calls = [] thought_sig = "" - first_sig_seen = False for part in content_parts: has_func = "functionCall" in part @@ -1321,8 +1350,8 @@ def _gemini_to_openai_non_streaming( if has_func: tool_call = self._extract_tool_call(part, model, len(tool_calls)) - if has_sig and not first_sig_seen: - first_sig_seen = True + # Store signature for each tool call (needed for parallel tool calls) + if has_sig: self._handle_tool_signature(tool_call, part["thoughtSignature"]) tool_calls.append(tool_call) diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 32e54f3f..0a0ab514 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -13,6 +13,7 @@ from ..model_definitions import ModelDefinitions import litellm from litellm.exceptions import RateLimitError +from ..error_handler import extract_retry_after_from_body import os from pathlib import Path import uuid @@ -125,6 +126,7 @@ def _env_int(key: str, default: int) -> int: """Get integer from environment variable.""" return int(os.getenv(key, str(default))) + class GeminiCliProvider(GeminiAuthBase, ProviderInterface): skip_cost_calculation = True @@ -684,11 +686,21 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - if tool_call.get("type") == "function": tool_call_id_to_name[tool_call["id"]] = tool_call["function"]["name"] + # Process messages and consolidate consecutive tool responses + # Per Gemini docs: parallel function responses must be in a single user message, + # not interleaved as separate messages + pending_tool_parts = [] # Accumulate tool responses + for msg in messages: role = msg.get("role") content = msg.get("content") parts = [] - gemini_role = "model" if role == "assistant" else "tool" if role == "tool" else "user" + gemini_role = "model" if role == "assistant" else "user" # tool -> user in Gemini + + # If we have pending tool parts and hit a non-tool message, flush them first + if pending_tool_parts and role != "tool": + gemini_contents.append({"role": "user", "parts": pending_tool_parts}) + pending_tool_parts = [] if role == "user": if isinstance(content, str): @@ -725,6 +737,9 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - if isinstance(content, str): parts.append({"text": content}) if msg.get("tool_calls"): + # Track if we've seen the first function call in this message + # Per Gemini docs: Only the FIRST parallel function call gets a signature + first_func_in_msg = True for tool_call in msg["tool_calls"]: if tool_call.get("type") == "function": try: @@ -748,6 +763,8 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - } # Add thoughtSignature for Gemini 3 + # Per Gemini docs: Only the FIRST parallel function call gets a signature. + # Subsequent parallel calls should NOT have a thoughtSignature field. if is_gemini_3: sig = tool_call.get("thought_signature") if not sig and tool_id and self._enable_signature_cache: @@ -755,9 +772,13 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - if sig: func_part["thoughtSignature"] = sig - else: + elif first_func_in_msg: + # Only add bypass to the first function call if no sig available func_part["thoughtSignature"] = "skip_thought_signature_validator" - lib_logger.warning(f"Missing thoughtSignature for {tool_id}, using bypass") + lib_logger.warning(f"Missing thoughtSignature for first func call {tool_id}, using bypass") + # Subsequent parallel calls: no signature field at all + + first_func_in_msg = False parts.append(func_part) @@ -771,17 +792,24 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - # Wrap the tool response in a 'result' object response_content = {"result": content} - parts.append({ + # Accumulate tool responses - they'll be combined into one user message + pending_tool_parts.append({ "functionResponse": { "name": function_name, "response": response_content, "id": tool_call_id } }) + # Don't add parts here - tool responses are handled via pending_tool_parts + continue if parts: gemini_contents.append({"role": gemini_role, "parts": parts}) + # Flush any remaining tool parts at end of messages + if pending_tool_parts: + gemini_contents.append({"role": "user", "parts": pending_tool_parts}) + if not gemini_contents or gemini_contents[0]['role'] != 'user': gemini_contents.insert(0, {"role": "user", "parts": [{"text": ""}]}) @@ -866,7 +894,6 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumul candidate = candidates[0] parts = candidate.get('content', {}).get('parts', []) is_gemini_3 = self._is_gemini_3(model_id) - first_sig_seen = False for part in parts: delta = {} @@ -905,8 +932,8 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumul } # Handle thoughtSignature for Gemini 3 - if is_gemini_3 and has_sig and not first_sig_seen: - first_sig_seen = True + # Store signature for each tool call (needed for parallel tool calls) + if is_gemini_3 and has_sig: sig = part['thoughtSignature'] if self._enable_signature_cache: @@ -1369,6 +1396,15 @@ async def stream_handler(): }) try: async with client.stream("POST", url, headers=final_headers, json=request_payload, params={"alt": "sse"}, timeout=600) as response: + # Read and log error body before raise_for_status for better debugging + if response.status_code >= 400: + try: + error_body = await response.aread() + lib_logger.error(f"Gemini CLI API error {response.status_code}: {error_body.decode()}") + file_logger.log_error(f"API error {response.status_code}: {error_body.decode()}") + except Exception: + pass + # This will raise an HTTPStatusError for 4xx/5xx responses response.raise_for_status() @@ -1405,16 +1441,24 @@ async def stream_handler(): error_body = e.response.text except Exception: pass - log_line = f"Stream handler HTTPStatusError: {str(e)}" + + # Only log to file logger (for detailed logging) if error_body: - log_line = f"{log_line} | response_body={error_body}" - file_logger.log_error(log_line) + file_logger.log_error(f"HTTPStatusError {e.response.status_code}: {error_body}") + else: + file_logger.log_error(f"HTTPStatusError {e.response.status_code}: {str(e)}") + if e.response.status_code == 429: - # Pass the raw response object to the exception. Do not read the - # response body here as it will close the stream and cause a - # 'StreamClosed' error in the client's stream reader. + # Extract retry-after time from the error body + retry_after = extract_retry_after_from_body(error_body) + retry_info = f" (retry after {retry_after}s)" if retry_after else "" + error_msg = f"Gemini CLI rate limit exceeded{retry_info}" + if error_body: + error_msg = f"{error_msg} | {error_body}" + # Only log at debug level - rotation happens silently + lib_logger.debug(f"Gemini CLI 429 rate limit: retry_after={retry_after}s") raise RateLimitError( - message=f"Gemini CLI rate limit exceeded: {e.request.url}", + message=error_msg, llm_provider="gemini_cli", model=model, response=e.response @@ -1451,7 +1495,8 @@ async def logging_stream_wrapper(): for idx, attempt_model in enumerate(fallback_models): is_fallback = idx > 0 if is_fallback: - lib_logger.info(f"Gemini CLI rate limited, retrying with fallback model: {attempt_model}") + # Silent rotation - only log at debug level + lib_logger.debug(f"Rate limited on previous model, trying fallback: {attempt_model}") elif has_fallbacks: lib_logger.debug(f"Attempting primary model: {attempt_model} (with {len(fallback_models)-1} fallback(s) available)") else: @@ -1473,8 +1518,8 @@ async def logging_stream_wrapper(): if idx + 1 < len(fallback_models): lib_logger.debug(f"Rate limit hit on {attempt_model}, trying next fallback...") continue - # If this was the last fallback option, raise the error - lib_logger.error(f"Rate limit hit on all fallback models (tried {len(fallback_models)} models)") + # If this was the last fallback option, log error and raise + lib_logger.warning(f"Rate limit exhausted on all fallback models (tried {len(fallback_models)} models)") raise # Should not reach here, but raise last error if we do From 087aab7958255e82f6d83ef46ce951b668b1d68a Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 16:02:20 +0100 Subject: [PATCH 037/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20add=20?= =?UTF-8?q?thinking=20mode=20sanitization=20for=20Claude=20API=20compatibi?= =?UTF-8?q?lity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces comprehensive thinking mode sanitization to prevent 400 errors when using Claude's extended thinking feature across different conversation states and model switches. - Add new environment variable `ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION` (default: true) to control sanitization behavior - Implement conversation state analysis to detect tool use loops and thinking block presence - Handle four distinct scenarios per Claude API documentation: 1. Thinking disabled: strip all thinking blocks from conversation 2. Tool loop with existing thinking: preserve current turn thinking only 3. Tool loop without thinking (invalid toggle): inject synthetic assistant response to close the loop 4. No tool loop: strip old turn thinking, allow new response to add thinking naturally - Add `_analyze_conversation_state()` to detect tool loops and thinking block locations - Add `_sanitize_thinking_for_claude()` as main orchestration method - Add helper methods for stripping, preserving, and closing tool loops - Support `reasoning_content` field in message transformation for cached thinking blocks - Add safety checks to maintain role alternation in edge cases - Integrate sanitization into completion flow before message transformation The sanitization prevents the Claude API error: "Expected `thinking` or `redacted_thinking`, but found `tool_use`" which occurs when attempting to toggle thinking mode mid-turn during tool use loops. This fix enables seamless thinking mode across context compression, model switching (e.g., Gemini to Claude), and multi-turn tool use conversations. --- DOCUMENTATION.md | 51 +++ README.md | 2 +- .../providers/antigravity_provider.py | 363 +++++++++++++++++- 3 files changed, 411 insertions(+), 5 deletions(-) diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index 94beec4b..b5a94938 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -458,6 +458,7 @@ ANTIGRAVITY_ENABLE_SIGNATURE_CACHE=true ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES=true # Include signatures in client responses ANTIGRAVITY_ENABLE_DYNAMIC_MODELS=false # Use API model discovery ANTIGRAVITY_GEMINI3_TOOL_FIX=true # Enable Gemini 3 hallucination prevention +ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION=true # Enable Claude thinking mode auto-correction # Gemini 3 tool fix customization ANTIGRAVITY_GEMINI3_TOOL_PREFIX="gemini3_" # Namespace prefix @@ -465,6 +466,56 @@ ANTIGRAVITY_GEMINI3_DESCRIPTION_PROMPT="\n\nSTRICT PARAMETERS: {params}." ANTIGRAVITY_GEMINI3_SYSTEM_INSTRUCTION="..." # Full system prompt ``` +#### Claude Extended Thinking Sanitization + +The provider includes automatic sanitization for Claude's extended thinking mode, handling common error scenarios: + +**Problem**: Claude's extended thinking API requires strict consistency in thinking blocks: +- If thinking is enabled, the final assistant turn must start with a thinking block +- If thinking is disabled, no thinking blocks can be present in the final turn +- Tool use loops are part of a single "assistant turn" +- You **cannot** toggle thinking mode mid-turn (this is invalid per Claude API) + +**Scenarios Handled**: + +| Scenario | Action | +|----------|--------| +| Tool loop WITH thinking + thinking enabled | Preserve thinking, continue normally | +| Tool loop WITHOUT thinking + thinking enabled | **Inject synthetic closure** to start fresh turn with thinking | +| Thinking disabled | Strip all thinking blocks | +| Normal conversation (no tool loop) | Strip old thinking, new response adds thinking naturally | + +**Solution**: The `_sanitize_thinking_for_claude()` method: +- Analyzes conversation state to detect incomplete tool use loops +- When enabling thinking in a tool loop that started without thinking: + - Injects a minimal synthetic assistant message: `"[Tool execution completed. Processing results.]"` + - This **closes** the previous turn, allowing Claude to start a **fresh turn with thinking** +- Strips thinking from old turns (Claude API ignores them anyway) +- Preserves thinking when the turn was started with thinking enabled + +**Key Insight**: Instead of force-disabling thinking, we close the tool loop with a synthetic message. This allows seamless model switching (e.g., Gemini → Claude with thinking) without losing the ability to think. + +**Example**: +``` +Before sanitization: + User: "What's the weather?" + Assistant: [tool_use: get_weather] ← Made by Gemini (no thinking) + User: [tool_result: "20C sunny"] + +After sanitization (thinking enabled): + User: "What's the weather?" + Assistant: [tool_use: get_weather] + User: [tool_result: "20C sunny"] + Assistant: "[Tool execution completed. Processing results.]" ← INJECTED + + → Claude now starts a NEW turn and CAN think! +``` + +**Configuration**: +```env +ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION=true # Enable/disable auto-correction +``` + #### File Logging Optional transaction logging for debugging: diff --git a/README.md b/README.md index f3a12867..b3ae33d3 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ This project provides a powerful solution for developers building complex applic - **OpenAI-Compatible Proxy**: Offers a familiar API interface with additional endpoints for model and provider discovery. - **Advanced Model Filtering**: Supports both blacklists and whitelists to give you fine-grained control over which models are available through the proxy. -- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 2.5, Gemini 3, and Claude Sonnet 4.5 models with advanced features like thought signature caching and tool hallucination prevention. +- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 2.5, Gemini 3, and Claude Sonnet 4.5 models with advanced features like thought signature caching and tool hallucination prevention. However - Sonnet 4.5 Thinking with native tool calls is very skittish, so if you have compaction or switch the model (or toggle thinking) mid task - it will error 400 on you, as claude needs it's previous thinking block. With compaction - it will be destroyed. There is a system to maybe catch all this, but i am hurting my head here trying to come up with a solution that makes sense. - **🆕 Credential Prioritization**: Automatic tier detection and priority-based credential selection ensures paid-tier credentials are used for premium models that require them. - **🆕 Weighted Random Rotation**: Configurable credential rotation strategy - choose between deterministic (perfect balance) or weighted random (unpredictable, harder to fingerprint) selection. - **🆕 Enhanced Gemini CLI**: Improved project discovery, paid vs free tier detection, and Gemini 3 support with thoughtSignature caching. diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 55c28a8e..988a6e2c 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -401,6 +401,7 @@ def __init__(self): self._enable_dynamic_models = _env_bool("ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", False) self._enable_gemini3_tool_fix = _env_bool("ANTIGRAVITY_GEMINI3_TOOL_FIX", True) self._enable_claude_tool_fix = _env_bool("ANTIGRAVITY_CLAUDE_TOOL_FIX", True) + self._enable_thinking_sanitization = _env_bool("ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION", True) # Gemini 3 tool fix configuration self._gemini3_tool_prefix = os.getenv("ANTIGRAVITY_GEMINI3_TOOL_PREFIX", "gemini3_") @@ -431,7 +432,8 @@ def _log_config(self) -> None: lib_logger.debug( f"Antigravity config: signatures_in_client={self._preserve_signatures_in_client}, " f"cache={self._enable_signature_cache}, dynamic_models={self._enable_dynamic_models}, " - f"gemini3_fix={self._enable_gemini3_tool_fix}, claude_fix={self._enable_claude_tool_fix}" + f"gemini3_fix={self._enable_gemini3_tool_fix}, claude_fix={self._enable_claude_tool_fix}, " + f"thinking_sanitization={self._enable_thinking_sanitization}" ) # ========================================================================= @@ -512,6 +514,295 @@ def _generate_thinking_cache_key( return "thinking_" + "_".join(key_parts) if key_parts else None + # ========================================================================= + # THINKING MODE SANITIZATION + # ========================================================================= + + def _analyze_conversation_state( + self, + messages: List[Dict[str, Any]] + ) -> Dict[str, Any]: + """ + Analyze conversation state to detect tool use loops and thinking mode issues. + + Returns: + { + "in_tool_loop": bool - True if we're in an incomplete tool use loop + "last_assistant_idx": int - Index of last assistant message + "last_assistant_has_thinking": bool - Whether last assistant msg has thinking + "last_assistant_has_tool_calls": bool - Whether last assistant msg has tool calls + "pending_tool_results": bool - Whether there are tool results after last assistant + "thinking_block_indices": List[int] - Indices of messages with thinking/reasoning + } + """ + state = { + "in_tool_loop": False, + "last_assistant_idx": -1, + "last_assistant_has_thinking": False, + "last_assistant_has_tool_calls": False, + "pending_tool_results": False, + "thinking_block_indices": [], + } + + # Find last assistant message and analyze the conversation + for i, msg in enumerate(messages): + role = msg.get("role") + + if role == "assistant": + state["last_assistant_idx"] = i + state["last_assistant_has_tool_calls"] = bool(msg.get("tool_calls")) + # Check for thinking/reasoning content + has_thinking = bool(msg.get("reasoning_content")) + # Also check for thinking in content array (some formats) + content = msg.get("content") + if isinstance(content, list): + for item in content: + if isinstance(item, dict) and item.get("type") == "thinking": + has_thinking = True + break + state["last_assistant_has_thinking"] = has_thinking + if has_thinking: + state["thinking_block_indices"].append(i) + elif role == "tool": + # Tool result after an assistant message with tool calls = in tool loop + if state["last_assistant_has_tool_calls"]: + state["pending_tool_results"] = True + + # We're in a tool loop if: + # 1. Last assistant message had tool calls + # 2. There are tool results after it + # 3. There's no final text response yet (the conversation ends with tool results) + if state["pending_tool_results"] and messages: + last_msg = messages[-1] + if last_msg.get("role") == "tool": + state["in_tool_loop"] = True + + return state + + def _sanitize_thinking_for_claude( + self, + messages: List[Dict[str, Any]], + thinking_enabled: bool + ) -> Tuple[List[Dict[str, Any]], bool]: + """ + Sanitize thinking blocks in conversation history for Claude compatibility. + + Handles the following scenarios per Claude docs: + 1. If thinking is disabled, remove all thinking blocks from conversation + 2. If thinking is enabled: + a. In a tool use loop WITH thinking: preserve it (same mode continues) + b. In a tool use loop WITHOUT thinking: this is INVALID toggle - force disable + c. Not in tool loop: strip old thinking, new response adds thinking naturally + + Per Claude docs: + - "If thinking is enabled, the final assistant turn must start with a thinking block" + - "If thinking is disabled, the final assistant turn must not contain any thinking blocks" + - Tool use loops are part of a single assistant turn + - You CANNOT toggle thinking mid-turn + + The key insight: We only force-disable thinking when TOGGLING it ON mid-turn. + If thinking was already enabled (assistant has thinking), we preserve. + If thinking was disabled (assistant has no thinking), enabling it now is invalid. + + Returns: + Tuple of (sanitized_messages, force_disable_thinking) + - sanitized_messages: The cleaned message list + - force_disable_thinking: If True, thinking must be disabled for this request + """ + messages = copy.deepcopy(messages) + state = self._analyze_conversation_state(messages) + + lib_logger.debug( + f"[Thinking Sanitization] thinking_enabled={thinking_enabled}, " + f"in_tool_loop={state['in_tool_loop']}, " + f"last_assistant_has_thinking={state['last_assistant_has_thinking']}, " + f"last_assistant_has_tool_calls={state['last_assistant_has_tool_calls']}" + ) + + if not thinking_enabled: + # CASE 1: Thinking is disabled - strip ALL thinking blocks + return self._strip_all_thinking_blocks(messages), False + + # CASE 2: Thinking is enabled + if state["in_tool_loop"]: + # We're in a tool use loop (conversation ends with tool_result) + # Per Claude docs: entire assistant turn must operate in single thinking mode + + if state["last_assistant_has_thinking"]: + # Last assistant turn HAD thinking - this is valid! + # Thinking was enabled when tool was called, continue with thinking enabled. + # Only keep thinking for the current turn (last assistant + following tools) + lib_logger.debug( + "[Thinking Sanitization] Tool loop with existing thinking - preserving." + ) + return self._preserve_current_turn_thinking( + messages, state["last_assistant_idx"] + ), False + else: + # Last assistant turn DID NOT have thinking, but thinking is NOW enabled + # This is the INVALID case: toggling thinking ON mid-turn + # + # Per Claude docs, this causes: + # "Expected `thinking` or `redacted_thinking`, but found `tool_use`." + # + # SOLUTION: Inject a synthetic assistant message to CLOSE the tool loop. + # This allows Claude to start a fresh turn WITH thinking. + # + # The synthetic message summarizes the tool results, allowing the model + # to respond naturally with thinking enabled on what is now a "new" turn. + lib_logger.info( + "[Thinking Sanitization] Closing tool loop with synthetic response. " + "This allows thinking to be enabled on the new turn." + ) + return self._close_tool_loop_for_thinking(messages), False + else: + # Not in a tool loop - this is the simple case + # The conversation doesn't end with tool_result, so we're starting fresh. + # Strip thinking from old turns (API ignores them anyway). + # New response will include thinking naturally. + + if state["last_assistant_idx"] >= 0 and not state["last_assistant_has_thinking"]: + if state["last_assistant_has_tool_calls"]: + # Last assistant made tool calls but no thinking + # This could be from context compression, model switch, or + # the assistant responded after tool results (completing the turn) + lib_logger.debug( + "[Thinking Sanitization] Last assistant has completed tool_calls but no thinking. " + "This is likely from context compression or completed tool loop. " + "New response will include thinking." + ) + + # Strip thinking from old turns, let new response add thinking naturally + return self._strip_old_turn_thinking(messages, state["last_assistant_idx"]), False + + def _strip_all_thinking_blocks( + self, + messages: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """Remove all thinking/reasoning content from messages.""" + for msg in messages: + if msg.get("role") == "assistant": + # Remove reasoning_content field + msg.pop("reasoning_content", None) + + # Remove thinking blocks from content array + content = msg.get("content") + if isinstance(content, list): + filtered = [ + item for item in content + if not (isinstance(item, dict) and item.get("type") == "thinking") + ] + # If filtering leaves empty list, we need to preserve message structure + # to maintain user/assistant alternation. Use empty string as placeholder + # (will result in empty "text" part which is valid). + if not filtered: + # Only if there are no tool_calls either - otherwise message is valid + if not msg.get("tool_calls"): + msg["content"] = "" + else: + msg["content"] = None # tool_calls exist, content not needed + else: + msg["content"] = filtered + return messages + + def _strip_old_turn_thinking( + self, + messages: List[Dict[str, Any]], + last_assistant_idx: int + ) -> List[Dict[str, Any]]: + """ + Strip thinking from old turns but preserve for the last assistant turn. + + Per Claude docs: "thinking blocks from previous turns are removed from context" + This mimics the API behavior and prevents issues. + """ + for i, msg in enumerate(messages): + if msg.get("role") == "assistant" and i < last_assistant_idx: + # Old turn - strip thinking + msg.pop("reasoning_content", None) + content = msg.get("content") + if isinstance(content, list): + filtered = [ + item for item in content + if not (isinstance(item, dict) and item.get("type") == "thinking") + ] + # Preserve message structure with empty string if needed + if not filtered: + msg["content"] = "" if not msg.get("tool_calls") else None + else: + msg["content"] = filtered + return messages + + def _preserve_current_turn_thinking( + self, + messages: List[Dict[str, Any]], + last_assistant_idx: int + ) -> List[Dict[str, Any]]: + """ + Preserve thinking only for the current (last) assistant turn. + Strip from all previous turns. + """ + # Same as strip_old_turn_thinking - we keep the last turn intact + return self._strip_old_turn_thinking(messages, last_assistant_idx) + + def _close_tool_loop_for_thinking( + self, + messages: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """ + Close an incomplete tool loop by injecting a synthetic assistant response. + + This is used when: + - We're in a tool loop (conversation ends with tool_result) + - The tool call was made WITHOUT thinking (e.g., by Gemini or non-thinking Claude) + - We NOW want to enable thinking + + By injecting a synthetic response that "closes" the previous turn, + Claude can start a fresh turn with thinking enabled. + + The synthetic message is minimal and factual - it just acknowledges + the tool results were received, allowing the model to process them + with thinking on the new turn. + """ + # Strip any old thinking first + messages = self._strip_all_thinking_blocks(messages) + + # Collect tool results from the end of the conversation + tool_results = [] + for msg in reversed(messages): + if msg.get("role") == "tool": + tool_results.append(msg) + elif msg.get("role") == "assistant": + break # Stop at the assistant that made the tool calls + + tool_results.reverse() # Put back in order + + # Safety check: if no tool results found, this shouldn't have been called + # But handle gracefully with a generic message + if not tool_results: + lib_logger.warning( + "[Thinking Sanitization] _close_tool_loop_for_thinking called but no tool results found. " + "This may indicate malformed conversation history." + ) + synthetic_content = "[Processing previous context.]" + elif len(tool_results) == 1: + synthetic_content = "[Tool execution completed. Processing results.]" + else: + synthetic_content = f"[{len(tool_results)} tool executions completed. Processing results.]" + + # Inject the synthetic assistant message to close the loop + synthetic_msg = { + "role": "assistant", + "content": synthetic_content + } + messages.append(synthetic_msg) + + lib_logger.debug( + f"[Thinking Sanitization] Injected synthetic closure: '{synthetic_content}'" + ) + + return messages + # ========================================================================= # REASONING CONFIGURATION # ========================================================================= @@ -691,9 +982,43 @@ def _transform_assistant_message( parts = [] content = msg.get("content") tool_calls = msg.get("tool_calls", []) - - # Try to inject cached thinking for Claude - if self._is_claude(model) and self._enable_signature_cache: + reasoning_content = msg.get("reasoning_content") + + # Handle reasoning_content if present (from original Claude response with thinking) + if reasoning_content and self._is_claude(model): + # Add thinking part with cached signature + thinking_part = { + "text": reasoning_content, + "thought": True, + } + # Try to get signature from cache + cache_key = self._generate_thinking_cache_key( + content if isinstance(content, str) else "", + tool_calls + ) + cached_sig = None + if cache_key: + cached_json = self._thinking_cache.retrieve(cache_key) + if cached_json: + try: + cached_data = json.loads(cached_json) + cached_sig = cached_data.get("thought_signature", "") + except json.JSONDecodeError: + pass + + if cached_sig: + thinking_part["thoughtSignature"] = cached_sig + parts.append(thinking_part) + lib_logger.debug(f"Added reasoning_content with cached signature ({len(reasoning_content)} chars)") + else: + # No cached signature - skip the thinking block + # This can happen if context was compressed and signature was lost + lib_logger.warning( + f"Skipping reasoning_content - no valid signature found. " + f"This may cause issues if thinking is enabled." + ) + elif self._is_claude(model) and self._enable_signature_cache and not reasoning_content: + # Fallback: Try to inject cached thinking for Claude (original behavior) thinking_parts = self._get_cached_thinking(content, tool_calls) parts.extend(thinking_parts) @@ -754,6 +1079,16 @@ def _transform_assistant_message( parts.append(func_part) + # Safety: ensure we return at least one part to maintain role alternation + # This handles edge cases like assistant messages that had only thinking content + # which got stripped, leaving the message otherwise empty + if not parts: + # Use a minimal text part - can happen after thinking is stripped + parts.append({"text": ""}) + lib_logger.debug( + "[Transform] Added empty text part to maintain role alternation" + ) + return parts def _get_cached_thinking( @@ -1583,6 +1918,26 @@ async def acompletion( # Create logger file_logger = AntigravityFileLogger(model, enable_logging) + # Determine if thinking is enabled for this request + # Thinking is enabled if reasoning_effort is set (and not "disable") for Claude + thinking_enabled = False + if self._is_claude(model): + # For Claude, thinking is enabled when reasoning_effort is provided and not "disable" + thinking_enabled = reasoning_effort is not None and reasoning_effort != "disable" + + # Sanitize thinking blocks for Claude to prevent 400 errors + # This handles: context compression, model switching, mid-turn thinking toggle + # Returns (sanitized_messages, force_disable_thinking) + force_disable_thinking = False + if self._is_claude(model) and self._enable_thinking_sanitization: + messages, force_disable_thinking = self._sanitize_thinking_for_claude(messages, thinking_enabled) + + # If we're in a mid-turn thinking toggle situation, we MUST disable thinking + # for this request. Thinking will naturally resume on the next turn. + if force_disable_thinking: + thinking_enabled = False + reasoning_effort = "disable" # Force disable for this request + # Transform messages system_instruction, gemini_contents = self._transform_messages(messages, model) gemini_contents = self._fix_tool_response_grouping(gemini_contents) From 474826e193eac52de44b7f94a2900d76645d45ee Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 16:33:14 +0100 Subject: [PATCH 038/221] =?UTF-8?q?chore(antigravity):=20=F0=9F=A7=B9=20up?= =?UTF-8?q?date=20User-Agent=20header=20to=20version=201.11.9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bump the User-Agent version string from 1.11.5 to 1.11.9 to reflect the current antigravity provider implementation version. --- src/rotator_library/providers/antigravity_provider.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 988a6e2c..cc70191c 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -2007,7 +2007,7 @@ async def acompletion( "Authorization": f"Bearer {token}", "Content-Type": "application/json", "Host": host, - "User-Agent": "antigravity/1.11.5", + "User-Agent": "antigravity/1.11.9", "Accept": "text/event-stream" if stream else "application/json" } From 6c4ca7ccec58144d73afd0e14d5426d6f537b115 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 16:41:57 +0100 Subject: [PATCH 039/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20add=20?= =?UTF-8?q?default=20safety=20settings=20to=20prevent=20content=20filterin?= =?UTF-8?q?g?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, safety settings were being removed from the Antigravity API payload, which could result in content being blocked by default safety filters. This commit introduces default safety settings that disable content filtering for all categories. - Adds `DEFAULT_SAFETY_SETTINGS` constant with all safety categories set to minimum thresholds - Modifies payload preparation to include safety settings if not already present - Uses deep copy to prevent mutation of the default settings constant - Aligns with CLIProxyAPI requirements to prevent safety blocks during API calls The change ensures that API calls are not unexpectedly blocked by content filters while still allowing custom safety settings to be passed when explicitly provided in the request payload. --- .../providers/antigravity_provider.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index cc70191c..22573096 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -87,6 +87,16 @@ "OTHER": "stop", } +# Default safety settings - disable content filtering for all categories +# Per CLIProxyAPI: these are attached to prevent safety blocks during API calls +DEFAULT_SAFETY_SETTINGS = [ + {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"}, + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"}, + {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"}, + {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"}, + {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"}, +] + # Directory paths _BASE_DIR = Path(__file__).resolve().parent.parent.parent.parent LOGS_DIR = _BASE_DIR / "logs" / "antigravity_logs" @@ -1471,8 +1481,10 @@ def _transform_to_antigravity_format( # Add session ID antigravity_payload["request"]["sessionId"] = _generate_session_id() - # Remove unsupported fields - antigravity_payload["request"].pop("safetySettings", None) + # Add default safety settings to prevent content filtering + # Only add if not already present in the payload + if "safetySettings" not in antigravity_payload["request"]: + antigravity_payload["request"]["safetySettings"] = copy.deepcopy(DEFAULT_SAFETY_SETTINGS) # Handle max_tokens - only apply to Claude, or if explicitly set for others gen_config = antigravity_payload["request"].get("generationConfig", {}) From 5bc49f20fefbe22eedce0d8d38196caca144805b Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 17:07:18 +0100 Subject: [PATCH 040/221] =?UTF-8?q?feat(auth):=20=E2=9C=A8=20add=20environ?= =?UTF-8?q?ment=20variable-based=20OAuth=20credential=20support=20with=20m?= =?UTF-8?q?ulti-account=20capability?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a comprehensive environment variable-based credential system for stateless deployments, enabling multiple OAuth accounts per provider without requiring credential files. Key changes: - Add env-based credential discovery in CredentialManager with priority over file-based credentials - Implement numbered credential format (PROVIDER_N_ACCESS_TOKEN) supporting multiple accounts per provider - Support legacy single-credential format (PROVIDER_ACCESS_TOKEN) for backwards compatibility - Introduce virtual path system (env://provider/index) for env-based credentials - Update credential export tool to generate numbered .env files with merge instructions - Extend env credential support across all OAuth providers (Google OAuth, Antigravity, iFlow, Qwen Code) - Add Windows launcher script (launcher.bat) with interactive menu system for proxy configuration The numbered format allows combining multiple credentials in a single .env file: - ANTIGRAVITY_1_ACCESS_TOKEN, ANTIGRAVITY_1_REFRESH_TOKEN (first account) - ANTIGRAVITY_2_ACCESS_TOKEN, ANTIGRAVITY_2_REFRESH_TOKEN (second account) - etc. This enables containerized and serverless deployments without managing credential files, while maintaining full multi-account rotation capabilities. --- launcher.bat | 293 ++++++++++++++++++ src/rotator_library/credential_manager.py | 88 +++++- src/rotator_library/credential_tool.py | 267 ++++++++++------ .../providers/google_oauth_base.py | 105 +++++-- .../providers/iflow_auth_base.py | 84 +++-- .../providers/qwen_auth_base.py | 76 ++++- todo.md | 7 + 7 files changed, 765 insertions(+), 155 deletions(-) create mode 100644 launcher.bat create mode 100644 todo.md diff --git a/launcher.bat b/launcher.bat new file mode 100644 index 00000000..ec241862 --- /dev/null +++ b/launcher.bat @@ -0,0 +1,293 @@ +@echo off +:: ================================================================================ +:: Universal Instructions for macOS / Linux Users +:: ================================================================================ +:: This launcher.bat file is for Windows only. +:: If you are on macOS or Linux, please use the following Python commands directly +:: in your terminal. +:: +:: First, ensure you have Python 3.10 or higher installed. +:: +:: To run the proxy server (basic command): +:: export PYTHONPATH=${PYTHONPATH}:$(pwd)/src +:: python src/proxy_app/main.py --host 0.0.0.0 --port 8000 +:: +:: Note: To enable request logging, add the --enable-request-logging flag to the command. +:: +:: To add new credentials: +:: export PYTHONPATH=${PYTHONPATH}:$(pwd)/src +:: python src/proxy_app/main.py --add-credential +:: +:: To build the executable (requires PyInstaller): +:: pip install -r requirements.txt +:: pip install pyinstaller +:: python src/proxy_app/build.py +:: ================================================================================ + +setlocal enabledelayedexpansion + +:: Default Settings +set "HOST=0.0.0.0" +set "PORT=8000" +set "LOGGING=false" +set "EXECUTION_MODE=" +set "EXE_NAME=proxy_app.exe" +set "SOURCE_PATH=src\proxy_app\main.py" + +:: --- Phase 1: Detection and Mode Selection --- +set "EXE_EXISTS=false" +set "SOURCE_EXISTS=false" + +if exist "%EXE_NAME%" ( + set "EXE_EXISTS=true" +) + +if exist "%SOURCE_PATH%" ( + set "SOURCE_EXISTS=true" +) + +if "%EXE_EXISTS%"=="true" ( + if "%SOURCE_EXISTS%"=="true" ( + call :SelectModeMenu + ) else ( + set "EXECUTION_MODE=exe" + ) +) else ( + if "%SOURCE_EXISTS%"=="true" ( + set "EXECUTION_MODE=source" + call :CheckPython + if errorlevel 1 goto :eof + ) else ( + call :NoTargetsFound + ) +) + +if "%EXECUTION_MODE%"=="" ( + goto :eof +) + +:: --- Phase 2: Main Menu --- +:MainMenu +cls +echo ================================================== +echo LLM API Key Proxy Launcher +echo ================================================== +echo. +echo Current Configuration: +echo ---------------------- +echo - Host IP: %HOST% +echo - Port: %PORT% +echo - Request Logging: %LOGGING% +echo - Execution Mode: %EXECUTION_MODE% +echo. +echo Main Menu: +echo ---------- +echo 1. Run Proxy +echo 2. Configure Proxy +echo 3. Add Credentials +if "%EXECUTION_MODE%"=="source" ( + echo 4. Build Executable + echo 5. Exit +) else ( + echo 4. Exit +) +echo. +set /p "CHOICE=Enter your choice: " + +if "%CHOICE%"=="1" goto :RunProxy +if "%CHOICE%"=="2" goto :ConfigMenu +if "%CHOICE%"=="3" goto :AddCredentials + +if "%EXECUTION_MODE%"=="source" ( + if "%CHOICE%"=="4" goto :BuildExecutable + if "%CHOICE%"=="5" goto :eof +) else ( + if "%CHOICE%"=="4" goto :eof +) + +echo Invalid choice. +pause +goto :MainMenu + +:: --- Phase 3: Configuration Sub-Menu --- +:ConfigMenu +cls +echo ================================================== +echo Configuration Menu +echo ================================================== +echo. +echo Current Configuration: +echo ---------------------- +echo - Host IP: %HOST% +echo - Port: %PORT% +echo - Request Logging: %LOGGING% +echo - Execution Mode: %EXECUTION_MODE% +echo. +echo Configuration Options: +echo ---------------------- +echo 1. Set Host IP +echo 2. Set Port +echo 3. Toggle Request Logging +echo 4. Back to Main Menu +echo. +set /p "CHOICE=Enter your choice: " + +if "%CHOICE%"=="1" ( + set /p "NEW_HOST=Enter new Host IP: " + if defined NEW_HOST ( + set "HOST=!NEW_HOST!" + ) + goto :ConfigMenu +) +if "%CHOICE%"=="2" ( + set "NEW_PORT=" + set /p "NEW_PORT=Enter new Port: " + if not defined NEW_PORT goto :ConfigMenu + set "IS_NUM=true" + for /f "delims=0123456789" %%i in ("!NEW_PORT!") do set "IS_NUM=false" + if "!IS_NUM!"=="false" ( + echo Invalid Port. Please enter numbers only. + pause + ) else ( + if !NEW_PORT! GTR 65535 ( + echo Invalid Port. Port cannot be greater than 65535. + pause + ) else ( + set "PORT=!NEW_PORT!" + ) + ) + goto :ConfigMenu +) +if "%CHOICE%"=="3" ( + if "%LOGGING%"=="true" ( + set "LOGGING=false" + ) else ( + set "LOGGING=true" + ) + goto :ConfigMenu +) +if "%CHOICE%"=="4" goto :MainMenu + +echo Invalid choice. +pause +goto :ConfigMenu + +:: --- Phase 4: Execution --- +:RunProxy +cls +set "ARGS=--host "%HOST%" --port %PORT%" +if "%LOGGING%"=="true" ( + set "ARGS=%ARGS% --enable-request-logging" +) +echo Starting Proxy... +echo Arguments: %ARGS% +echo. +if "%EXECUTION_MODE%"=="exe" ( + start "LLM API Proxy" "%EXE_NAME%" %ARGS% +) else ( + set "PYTHONPATH=%~dp0src;%PYTHONPATH%" + start "LLM API Proxy" python "%SOURCE_PATH%" %ARGS% +) +exit /b 0 + +:AddCredentials +cls +echo Launching Credential Tool... +echo. +if "%EXECUTION_MODE%"=="exe" ( + "%EXE_NAME%" --add-credential +) else ( + set "PYTHONPATH=%~dp0src;%PYTHONPATH%" + python "%SOURCE_PATH%" --add-credential +) +pause +goto :MainMenu + +:BuildExecutable +cls +echo ================================================== +echo Building Executable +echo ================================================== +echo. +echo The build process will start in a new window. +start "Build Process" cmd /c "pip install -r requirements.txt && pip install pyinstaller && python "src/proxy_app/build.py" && echo Build finished. && pause" +exit /b + +:: --- Helper Functions --- + +:SelectModeMenu +cls +echo ================================================== +echo Execution Mode Selection +echo ================================================== +echo. +echo Both executable and source code found. +echo Please choose which to use: +echo. +echo 1. Executable ("%EXE_NAME%") +echo 2. Source Code ("%SOURCE_PATH%") +echo. +set /p "CHOICE=Enter your choice: " + +if "%CHOICE%"=="1" ( + set "EXECUTION_MODE=exe" +) else if "%CHOICE%"=="2" ( + call :CheckPython + if errorlevel 1 goto :eof + set "EXECUTION_MODE=source" +) else ( + echo Invalid choice. + pause + goto :SelectModeMenu +) +goto :end_of_function + +:CheckPython +where python >nul 2>nul +if errorlevel 1 ( + echo Error: Python is not installed or not in PATH. + echo Please install Python and try again. + pause + exit /b 1 +) + +for /f "tokens=1,2" %%a in ('python -c "import sys; print(sys.version_info.major, sys.version_info.minor)"') do ( + set "PY_MAJOR=%%a" + set "PY_MINOR=%%b" +) + +if not "%PY_MAJOR%"=="3" ( + call :PythonVersionError + exit /b 1 +) +if %PY_MINOR% lss 10 ( + call :PythonVersionError + exit /b 1 +) + +exit /b 0 + +:PythonVersionError +echo Error: Python 3.10 or higher is required. +echo Found version: %PY_MAJOR%.%PY_MINOR% +echo Please upgrade your Python installation. +pause +goto :eof + +:NoTargetsFound +cls +echo ================================================== +echo Error +echo ================================================== +echo. +echo Could not find the executable ("%EXE_NAME%") +echo or the source code ("%SOURCE_PATH%"). +echo. +echo Please ensure the launcher is in the correct +echo directory or that the project has been built. +echo. +pause +goto :eof + +:end_of_function +endlocal diff --git a/src/rotator_library/credential_manager.py b/src/rotator_library/credential_manager.py index 0678f7c2..16be41c1 100644 --- a/src/rotator_library/credential_manager.py +++ b/src/rotator_library/credential_manager.py @@ -1,8 +1,9 @@ import os +import re import shutil import logging from pathlib import Path -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Set lib_logger = logging.getLogger('rotator_library') @@ -18,19 +19,96 @@ # Add other providers like 'claude' here if they have a standard CLI path } +# OAuth providers that support environment variable-based credentials +# Maps provider name to the ENV_PREFIX used by the provider +ENV_OAUTH_PROVIDERS = { + "gemini_cli": "GEMINI_CLI", + "antigravity": "ANTIGRAVITY", + "qwen_code": "QWEN_CODE", + "iflow": "IFLOW", +} + + class CredentialManager: """ Discovers OAuth credential files from standard locations, copies them locally, and updates the configuration to use the local paths. + + Also discovers environment variable-based OAuth credentials for stateless deployments. + Supports two env var formats: + + 1. Single credential (legacy): PROVIDER_ACCESS_TOKEN, PROVIDER_REFRESH_TOKEN + 2. Multiple credentials (numbered): PROVIDER_1_ACCESS_TOKEN, PROVIDER_2_ACCESS_TOKEN, etc. + + When env-based credentials are detected, virtual paths like "env://provider/1" are created. """ def __init__(self, env_vars: Dict[str, str]): self.env_vars = env_vars + def _discover_env_oauth_credentials(self) -> Dict[str, List[str]]: + """ + Discover OAuth credentials defined via environment variables. + + Supports two formats: + 1. Single credential: ANTIGRAVITY_ACCESS_TOKEN + ANTIGRAVITY_REFRESH_TOKEN + 2. Multiple credentials: ANTIGRAVITY_1_ACCESS_TOKEN + ANTIGRAVITY_1_REFRESH_TOKEN, etc. + + Returns: + Dict mapping provider name to list of virtual paths (e.g., "env://antigravity/1") + """ + env_credentials: Dict[str, Set[str]] = {} + + for provider, env_prefix in ENV_OAUTH_PROVIDERS.items(): + found_indices: Set[str] = set() + + # Check for numbered credentials (PROVIDER_N_ACCESS_TOKEN pattern) + # Pattern: ANTIGRAVITY_1_ACCESS_TOKEN, ANTIGRAVITY_2_ACCESS_TOKEN, etc. + numbered_pattern = re.compile(rf"^{env_prefix}_(\d+)_ACCESS_TOKEN$") + + for key in self.env_vars.keys(): + match = numbered_pattern.match(key) + if match: + index = match.group(1) + # Verify refresh token also exists + refresh_key = f"{env_prefix}_{index}_REFRESH_TOKEN" + if refresh_key in self.env_vars and self.env_vars[refresh_key]: + found_indices.add(index) + + # Check for legacy single credential (PROVIDER_ACCESS_TOKEN pattern) + # Only use this if no numbered credentials exist + if not found_indices: + access_key = f"{env_prefix}_ACCESS_TOKEN" + refresh_key = f"{env_prefix}_REFRESH_TOKEN" + if (access_key in self.env_vars and self.env_vars[access_key] and + refresh_key in self.env_vars and self.env_vars[refresh_key]): + # Use "0" as the index for legacy single credential + found_indices.add("0") + + if found_indices: + env_credentials[provider] = found_indices + lib_logger.info(f"Found {len(found_indices)} env-based credential(s) for {provider}") + + # Convert to virtual paths + result: Dict[str, List[str]] = {} + for provider, indices in env_credentials.items(): + # Sort indices numerically for consistent ordering + sorted_indices = sorted(indices, key=lambda x: int(x)) + result[provider] = [f"env://{provider}/{idx}" for idx in sorted_indices] + + return result + def discover_and_prepare(self) -> Dict[str, List[str]]: lib_logger.info("Starting automated OAuth credential discovery...") final_config = {} - # Extract OAuth paths from environment variables first + # PHASE 1: Discover environment variable-based OAuth credentials + # These take priority for stateless deployments + env_oauth_creds = self._discover_env_oauth_credentials() + for provider, virtual_paths in env_oauth_creds.items(): + lib_logger.info(f"Using {len(virtual_paths)} env-based credential(s) for {provider}") + final_config[provider] = virtual_paths + + # Extract OAuth file paths from environment variables env_oauth_paths = {} for key, value in self.env_vars.items(): if "_OAUTH_" in key: @@ -40,7 +118,13 @@ def discover_and_prepare(self) -> Dict[str, List[str]]: if value: # Only consider non-empty values env_oauth_paths[provider].append(value) + # PHASE 2: Discover file-based OAuth credentials for provider, default_dir in DEFAULT_OAUTH_DIRS.items(): + # Skip if already discovered from environment variables + if provider in final_config: + lib_logger.debug(f"Skipping file discovery for {provider} - using env-based credentials") + continue + # Check for existing local credentials first. If found, use them and skip discovery. local_provider_creds = sorted(list(OAUTH_BASE_DIR.glob(f"{provider}_oauth_*.json"))) if local_provider_creds: diff --git a/src/rotator_library/credential_tool.py b/src/rotator_library/credential_tool.py index 066befe3..4b2f8a04 100644 --- a/src/rotator_library/credential_tool.py +++ b/src/rotator_library/credential_tool.py @@ -36,6 +36,77 @@ def _ensure_providers_loaded(): _provider_plugins = pp return _provider_factory, _provider_plugins + +def _get_credential_number_from_filename(filename: str) -> int: + """ + Extract credential number from filename like 'provider_oauth_1.json' -> 1 + """ + match = re.search(r'_oauth_(\d+)\.json$', filename) + if match: + return int(match.group(1)) + return 1 + + +def _build_env_export_content( + provider_prefix: str, + cred_number: int, + creds: dict, + email: str, + extra_fields: dict = None, + include_client_creds: bool = True +) -> tuple[list[str], str]: + """ + Build .env content for OAuth credential export with numbered format. + Exports all fields from the JSON file as a 1-to-1 mirror. + + Args: + provider_prefix: Environment variable prefix (e.g., "ANTIGRAVITY", "GEMINI_CLI") + cred_number: Credential number for this export (1, 2, 3, etc.) + creds: The credential dictionary loaded from JSON + email: User email for comments + extra_fields: Optional dict of additional fields to include + include_client_creds: Whether to include client_id/secret (Google OAuth providers) + + Returns: + Tuple of (env_lines list, numbered_prefix string for display) + """ + # Use numbered format: PROVIDER_N_ACCESS_TOKEN + numbered_prefix = f"{provider_prefix}_{cred_number}" + + env_lines = [ + f"# {provider_prefix} Credential #{cred_number} for: {email}", + f"# Exported from: {provider_prefix.lower()}_oauth_{cred_number}.json", + f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", + f"# ", + f"# To combine multiple credentials into one .env file, copy these lines", + f"# and ensure each credential has a unique number (1, 2, 3, etc.)", + "", + f"{numbered_prefix}_ACCESS_TOKEN={creds.get('access_token', '')}", + f"{numbered_prefix}_REFRESH_TOKEN={creds.get('refresh_token', '')}", + f"{numbered_prefix}_SCOPE={creds.get('scope', '')}", + f"{numbered_prefix}_TOKEN_TYPE={creds.get('token_type', 'Bearer')}", + f"{numbered_prefix}_ID_TOKEN={creds.get('id_token', '')}", + f"{numbered_prefix}_EXPIRY_DATE={creds.get('expiry_date', 0)}", + ] + + if include_client_creds: + env_lines.extend([ + f"{numbered_prefix}_CLIENT_ID={creds.get('client_id', '')}", + f"{numbered_prefix}_CLIENT_SECRET={creds.get('client_secret', '')}", + f"{numbered_prefix}_TOKEN_URI={creds.get('token_uri', 'https://oauth2.googleapis.com/token')}", + f"{numbered_prefix}_UNIVERSE_DOMAIN={creds.get('universe_domain', 'googleapis.com')}", + ]) + + env_lines.append(f"{numbered_prefix}_EMAIL={email}") + + # Add extra provider-specific fields + if extra_fields: + for key, value in extra_fields.items(): + if value: # Only add non-empty values + env_lines.append(f"{numbered_prefix}_{key}={value}") + + return env_lines, numbered_prefix + def ensure_env_defaults(): """ Ensures the .env file exists and contains essential default values like PROXY_API_KEY. @@ -256,12 +327,12 @@ async def setup_new_credential(provider_name: str): async def export_gemini_cli_to_env(): """ Export a Gemini CLI credential JSON file to .env format. - Generates one .env file per credential. + Uses numbered format (GEMINI_CLI_1_*, GEMINI_CLI_2_*) for multiple credential support. """ console.print(Panel("[bold cyan]Export Gemini CLI Credential to .env[/bold cyan]", expand=False)) # Find all gemini_cli credentials - gemini_cli_files = list(OAUTH_BASE_DIR.glob("gemini_cli_oauth_*.json")) + gemini_cli_files = sorted(list(OAUTH_BASE_DIR.glob("gemini_cli_oauth_*.json"))) if not gemini_cli_files: console.print(Panel("No Gemini CLI credentials found. Please add one first using 'Add OAuth Credential'.", @@ -304,34 +375,30 @@ async def export_gemini_cli_to_env(): project_id = creds.get("_proxy_metadata", {}).get("project_id", "") tier = creds.get("_proxy_metadata", {}).get("tier", "") - # Generate .env file name + # Get credential number from filename + cred_number = _get_credential_number_from_filename(cred_file.name) + + # Generate .env file name with credential number safe_email = email.replace("@", "_at_").replace(".", "_") - env_filename = f"gemini_cli_{safe_email}.env" + env_filename = f"gemini_cli_{cred_number}_{safe_email}.env" env_filepath = OAUTH_BASE_DIR / env_filename - # Build .env content - env_lines = [ - f"# Gemini CLI Credential for: {email}", - f"# Generated from: {cred_file.name}", - f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", - "", - f"GEMINI_CLI_ACCESS_TOKEN={creds.get('access_token', '')}", - f"GEMINI_CLI_REFRESH_TOKEN={creds.get('refresh_token', '')}", - f"GEMINI_CLI_EXPIRY_DATE={creds.get('expiry_date', 0)}", - f"GEMINI_CLI_CLIENT_ID={creds.get('client_id', '')}", - f"GEMINI_CLI_CLIENT_SECRET={creds.get('client_secret', '')}", - f"GEMINI_CLI_TOKEN_URI={creds.get('token_uri', 'https://oauth2.googleapis.com/token')}", - f"GEMINI_CLI_UNIVERSE_DOMAIN={creds.get('universe_domain', 'googleapis.com')}", - f"GEMINI_CLI_EMAIL={email}", - ] - - # Add project_id if present + # Build extra fields + extra_fields = {} if project_id: - env_lines.append(f"GEMINI_CLI_PROJECT_ID={project_id}") - - # Add tier if present + extra_fields["PROJECT_ID"] = project_id if tier: - env_lines.append(f"GEMINI_CLI_TIER={tier}") + extra_fields["TIER"] = tier + + # Build .env content using helper + env_lines, numbered_prefix = _build_env_export_content( + provider_prefix="GEMINI_CLI", + cred_number=cred_number, + creds=creds, + email=email, + extra_fields=extra_fields, + include_client_creds=True + ) # Write to .env file with open(env_filepath, 'w') as f: @@ -339,11 +406,14 @@ async def export_gemini_cli_to_env(): success_text = Text.from_markup( f"Successfully exported credential to [bold yellow]'{env_filepath}'[/bold yellow]\n\n" - f"To use this credential:\n" - f"1. Copy [bold yellow]{env_filepath.name}[/bold yellow] to your deployment environment\n" - f"2. Load the variables: [bold cyan]export $(cat {env_filepath.name} | grep -v '^#' | xargs)[/bold cyan]\n" - f"3. Or source it: [bold cyan]source {env_filepath.name}[/bold cyan]\n" - f"4. The Gemini CLI provider will automatically use these environment variables" + f"[bold]Environment variable prefix:[/bold] [cyan]{numbered_prefix}_*[/cyan]\n\n" + f"[bold]To use this credential:[/bold]\n" + f"1. Copy the contents to your main .env file, OR\n" + f"2. Source it: [bold cyan]source {env_filepath.name}[/bold cyan] (Linux/Mac)\n" + f"3. Or on Windows: [bold cyan]Get-Content {env_filepath.name} | ForEach-Object {{ $_ -replace '^([^#].*)$', 'set $1' }} | cmd[/bold cyan]\n\n" + f"[bold]To combine multiple credentials:[/bold]\n" + f"Copy lines from multiple .env files into one file.\n" + f"Each credential uses a unique number ({numbered_prefix}_*)." ) console.print(Panel(success_text, style="bold green", title="Success")) else: @@ -403,22 +473,30 @@ async def export_qwen_code_to_env(): # Extract metadata email = creds.get("_proxy_metadata", {}).get("email", "unknown") - # Generate .env file name + # Get credential number from filename + cred_number = _get_credential_number_from_filename(cred_file.name) + + # Generate .env file name with credential number safe_email = email.replace("@", "_at_").replace(".", "_") - env_filename = f"qwen_code_{safe_email}.env" + env_filename = f"qwen_code_{cred_number}_{safe_email}.env" env_filepath = OAUTH_BASE_DIR / env_filename - # Build .env content + # Use numbered format: QWEN_CODE_N_* + numbered_prefix = f"QWEN_CODE_{cred_number}" + + # Build .env content (Qwen has different structure) env_lines = [ - f"# Qwen Code Credential for: {email}", - f"# Generated from: {cred_file.name}", + f"# QWEN_CODE Credential #{cred_number} for: {email}", f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", + f"# ", + f"# To combine multiple credentials into one .env file, copy these lines", + f"# and ensure each credential has a unique number (1, 2, 3, etc.)", "", - f"QWEN_CODE_ACCESS_TOKEN={creds.get('access_token', '')}", - f"QWEN_CODE_REFRESH_TOKEN={creds.get('refresh_token', '')}", - f"QWEN_CODE_EXPIRY_DATE={creds.get('expiry_date', 0)}", - f"QWEN_CODE_RESOURCE_URL={creds.get('resource_url', 'https://portal.qwen.ai/v1')}", - f"QWEN_CODE_EMAIL={email}", + f"{numbered_prefix}_ACCESS_TOKEN={creds.get('access_token', '')}", + f"{numbered_prefix}_REFRESH_TOKEN={creds.get('refresh_token', '')}", + f"{numbered_prefix}_EXPIRY_DATE={creds.get('expiry_date', 0)}", + f"{numbered_prefix}_RESOURCE_URL={creds.get('resource_url', 'https://portal.qwen.ai/v1')}", + f"{numbered_prefix}_EMAIL={email}", ] # Write to .env file @@ -427,11 +505,13 @@ async def export_qwen_code_to_env(): success_text = Text.from_markup( f"Successfully exported credential to [bold yellow]'{env_filepath}'[/bold yellow]\n\n" - f"To use this credential:\n" - f"1. Copy [bold yellow]{env_filepath.name}[/bold yellow] to your deployment environment\n" - f"2. Load the variables: [bold cyan]export $(cat {env_filepath.name} | grep -v '^#' | xargs)[/bold cyan]\n" - f"3. Or source it: [bold cyan]source {env_filepath.name}[/bold cyan]\n" - f"4. The Qwen Code provider will automatically use these environment variables" + f"[bold]Environment variable prefix:[/bold] [cyan]{numbered_prefix}_*[/cyan]\n\n" + f"[bold]To use this credential:[/bold]\n" + f"1. Copy the contents to your main .env file, OR\n" + f"2. Source it: [bold cyan]source {env_filepath.name}[/bold cyan] (Linux/Mac)\n\n" + f"[bold]To combine multiple credentials:[/bold]\n" + f"Copy lines from multiple .env files into one file.\n" + f"Each credential uses a unique number ({numbered_prefix}_*)." ) console.print(Panel(success_text, style="bold green", title="Success")) else: @@ -445,12 +525,12 @@ async def export_qwen_code_to_env(): async def export_iflow_to_env(): """ Export an iFlow credential JSON file to .env format. - Generates one .env file per credential. + Uses numbered format (IFLOW_1_*, IFLOW_2_*) for multiple credential support. """ console.print(Panel("[bold cyan]Export iFlow Credential to .env[/bold cyan]", expand=False)) # Find all iflow credentials - iflow_files = list(OAUTH_BASE_DIR.glob("iflow_oauth_*.json")) + iflow_files = sorted(list(OAUTH_BASE_DIR.glob("iflow_oauth_*.json"))) if not iflow_files: console.print(Panel("No iFlow credentials found. Please add one first using 'Add OAuth Credential'.", @@ -491,25 +571,32 @@ async def export_iflow_to_env(): # Extract metadata email = creds.get("_proxy_metadata", {}).get("email", "unknown") - # Generate .env file name + # Get credential number from filename + cred_number = _get_credential_number_from_filename(cred_file.name) + + # Generate .env file name with credential number safe_email = email.replace("@", "_at_").replace(".", "_") - env_filename = f"iflow_{safe_email}.env" + env_filename = f"iflow_{cred_number}_{safe_email}.env" env_filepath = OAUTH_BASE_DIR / env_filename - # Build .env content - # IMPORTANT: iFlow requires BOTH OAuth tokens AND the API key for API requests + # Use numbered format: IFLOW_N_* + numbered_prefix = f"IFLOW_{cred_number}" + + # Build .env content (iFlow has different structure with API key) env_lines = [ - f"# iFlow Credential for: {email}", - f"# Generated from: {cred_file.name}", + f"# IFLOW Credential #{cred_number} for: {email}", f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", + f"# ", + f"# To combine multiple credentials into one .env file, copy these lines", + f"# and ensure each credential has a unique number (1, 2, 3, etc.)", "", - f"IFLOW_ACCESS_TOKEN={creds.get('access_token', '')}", - f"IFLOW_REFRESH_TOKEN={creds.get('refresh_token', '')}", - f"IFLOW_API_KEY={creds.get('api_key', '')}", - f"IFLOW_EXPIRY_DATE={creds.get('expiry_date', '')}", - f"IFLOW_EMAIL={email}", - f"IFLOW_TOKEN_TYPE={creds.get('token_type', 'Bearer')}", - f"IFLOW_SCOPE={creds.get('scope', 'read write')}", + f"{numbered_prefix}_ACCESS_TOKEN={creds.get('access_token', '')}", + f"{numbered_prefix}_REFRESH_TOKEN={creds.get('refresh_token', '')}", + f"{numbered_prefix}_API_KEY={creds.get('api_key', '')}", + f"{numbered_prefix}_EXPIRY_DATE={creds.get('expiry_date', '')}", + f"{numbered_prefix}_EMAIL={email}", + f"{numbered_prefix}_TOKEN_TYPE={creds.get('token_type', 'Bearer')}", + f"{numbered_prefix}_SCOPE={creds.get('scope', 'read write')}", ] # Write to .env file @@ -518,11 +605,13 @@ async def export_iflow_to_env(): success_text = Text.from_markup( f"Successfully exported credential to [bold yellow]'{env_filepath}'[/bold yellow]\n\n" - f"To use this credential:\n" - f"1. Copy [bold yellow]{env_filepath.name}[/bold yellow] to your deployment environment\n" - f"2. Load the variables: [bold cyan]export $(cat {env_filepath.name} | grep -v '^#' | xargs)[/bold cyan]\n" - f"3. Or source it: [bold cyan]source {env_filepath.name}[/bold cyan]\n" - f"4. The iFlow provider will automatically use these environment variables" + f"[bold]Environment variable prefix:[/bold] [cyan]{numbered_prefix}_*[/cyan]\n\n" + f"[bold]To use this credential:[/bold]\n" + f"1. Copy the contents to your main .env file, OR\n" + f"2. Source it: [bold cyan]source {env_filepath.name}[/bold cyan] (Linux/Mac)\n\n" + f"[bold]To combine multiple credentials:[/bold]\n" + f"Copy lines from multiple .env files into one file.\n" + f"Each credential uses a unique number ({numbered_prefix}_*)." ) console.print(Panel(success_text, style="bold green", title="Success")) else: @@ -536,12 +625,12 @@ async def export_iflow_to_env(): async def export_antigravity_to_env(): """ Export an Antigravity credential JSON file to .env format. - Generates one .env file per credential. + Uses numbered format (ANTIGRAVITY_1_*, ANTIGRAVITY_2_*) for multiple credential support. """ console.print(Panel("[bold cyan]Export Antigravity Credential to .env[/bold cyan]", expand=False)) # Find all antigravity credentials - antigravity_files = list(OAUTH_BASE_DIR.glob("antigravity_oauth_*.json")) + antigravity_files = sorted(list(OAUTH_BASE_DIR.glob("antigravity_oauth_*.json"))) if not antigravity_files: console.print(Panel("No Antigravity credentials found. Please add one first using 'Add OAuth Credential'.", @@ -582,26 +671,23 @@ async def export_antigravity_to_env(): # Extract metadata email = creds.get("_proxy_metadata", {}).get("email", "unknown") - # Generate .env file name + # Get credential number from filename + cred_number = _get_credential_number_from_filename(cred_file.name) + + # Generate .env file name with credential number safe_email = email.replace("@", "_at_").replace(".", "_") - env_filename = f"antigravity_{safe_email}.env" + env_filename = f"antigravity_{cred_number}_{safe_email}.env" env_filepath = OAUTH_BASE_DIR / env_filename - # Build .env content - env_lines = [ - f"# Antigravity Credential for: {email}", - f"# Generated from: {cred_file.name}", - f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", - "", - f"ANTIGRAVITY_ACCESS_TOKEN={creds.get('access_token', '')}", - f"ANTIGRAVITY_REFRESH_TOKEN={creds.get('refresh_token', '')}", - f"ANTIGRAVITY_EXPIRY_DATE={creds.get('expiry_date', 0)}", - f"ANTIGRAVITY_CLIENT_ID={creds.get('client_id', '')}", - f"ANTIGRAVITY_CLIENT_SECRET={creds.get('client_secret', '')}", - f"ANTIGRAVITY_TOKEN_URI={creds.get('token_uri', 'https://oauth2.googleapis.com/token')}", - f"ANTIGRAVITY_UNIVERSE_DOMAIN={creds.get('universe_domain', 'googleapis.com')}", - f"ANTIGRAVITY_EMAIL={email}", - ] + # Build .env content using helper + env_lines, numbered_prefix = _build_env_export_content( + provider_prefix="ANTIGRAVITY", + cred_number=cred_number, + creds=creds, + email=email, + extra_fields=None, + include_client_creds=True + ) # Write to .env file with open(env_filepath, 'w') as f: @@ -609,11 +695,14 @@ async def export_antigravity_to_env(): success_text = Text.from_markup( f"Successfully exported credential to [bold yellow]'{env_filepath}'[/bold yellow]\n\n" - f"To use this credential:\n" - f"1. Copy [bold yellow]{env_filepath.name}[/bold yellow] to your deployment environment\n" - f"2. Load the variables: [bold cyan]export $(cat {env_filepath.name} | grep -v '^#' | xargs)[/bold cyan]\n" - f"3. Or source it: [bold cyan]source {env_filepath.name}[/bold cyan]\n" - f"4. The Antigravity provider will automatically use these environment variables" + f"[bold]Environment variable prefix:[/bold] [cyan]{numbered_prefix}_*[/cyan]\n\n" + f"[bold]To use this credential:[/bold]\n" + f"1. Copy the contents to your main .env file, OR\n" + f"2. Source it: [bold cyan]source {env_filepath.name}[/bold cyan] (Linux/Mac)\n" + f"3. Or on Windows: [bold cyan]Get-Content {env_filepath.name} | ForEach-Object {{ $_ -replace '^([^#].*)$', 'set $1' }} | cmd[/bold cyan]\n\n" + f"[bold]To combine multiple credentials:[/bold]\n" + f"Copy lines from multiple .env files into one file.\n" + f"Each credential uses a unique number ({numbered_prefix}_*)." ) console.print(Panel(success_text, style="bold green", title="Success")) else: diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index b40e90d1..3f1ed9d6 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -77,64 +77,103 @@ def __init__(self): self._queue_tracking_lock = asyncio.Lock() # Protects queue sets self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task - def _load_from_env(self) -> Optional[Dict[str, Any]]: + def _parse_env_credential_path(self, path: str) -> Optional[str]: + """ + Parse a virtual env:// path and return the credential index. + + Supported formats: + - "env://provider/0" - Legacy single credential (no index in env var names) + - "env://provider/1" - First numbered credential (PROVIDER_1_ACCESS_TOKEN) + - "env://provider/2" - Second numbered credential, etc. + + Returns: + The credential index as string ("0" for legacy, "1", "2", etc. for numbered) + or None if path is not an env:// path + """ + if not path.startswith("env://"): + return None + + # Parse: env://provider/index + parts = path[6:].split("/") # Remove "env://" prefix + if len(parts) >= 2: + return parts[1] # Return the index + return "0" # Default to legacy format + + def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dict[str, Any]]: """ Load OAuth credentials from environment variables for stateless deployments. - Expected environment variables: - - {ENV_PREFIX}_ACCESS_TOKEN (required) - - {ENV_PREFIX}_REFRESH_TOKEN (required) - - {ENV_PREFIX}_EXPIRY_DATE (optional, defaults to 0) - - {ENV_PREFIX}_CLIENT_ID (optional, uses default) - - {ENV_PREFIX}_CLIENT_SECRET (optional, uses default) - - {ENV_PREFIX}_TOKEN_URI (optional, uses default) - - {ENV_PREFIX}_UNIVERSE_DOMAIN (optional, defaults to googleapis.com) - - {ENV_PREFIX}_EMAIL (optional, defaults to "env-user") - - {ENV_PREFIX}_PROJECT_ID (optional) - - {ENV_PREFIX}_TIER (optional) + Supports two formats: + 1. Legacy (credential_index="0" or None): PROVIDER_ACCESS_TOKEN + 2. Numbered (credential_index="1", "2", etc.): PROVIDER_1_ACCESS_TOKEN, PROVIDER_2_ACCESS_TOKEN + + Expected environment variables (for numbered format with index N): + - {ENV_PREFIX}_{N}_ACCESS_TOKEN (required) + - {ENV_PREFIX}_{N}_REFRESH_TOKEN (required) + - {ENV_PREFIX}_{N}_EXPIRY_DATE (optional, defaults to 0) + - {ENV_PREFIX}_{N}_CLIENT_ID (optional, uses default) + - {ENV_PREFIX}_{N}_CLIENT_SECRET (optional, uses default) + - {ENV_PREFIX}_{N}_TOKEN_URI (optional, uses default) + - {ENV_PREFIX}_{N}_UNIVERSE_DOMAIN (optional, defaults to googleapis.com) + - {ENV_PREFIX}_{N}_EMAIL (optional, defaults to "env-user-{N}") + - {ENV_PREFIX}_{N}_PROJECT_ID (optional) + - {ENV_PREFIX}_{N}_TIER (optional) + + For legacy format (index="0" or None), omit the _{N}_ part. Returns: Dict with credential structure if env vars present, None otherwise """ - access_token = os.getenv(f"{self.ENV_PREFIX}_ACCESS_TOKEN") - refresh_token = os.getenv(f"{self.ENV_PREFIX}_REFRESH_TOKEN") + # Determine the env var prefix based on credential index + if credential_index and credential_index != "0": + # Numbered format: PROVIDER_N_ACCESS_TOKEN + prefix = f"{self.ENV_PREFIX}_{credential_index}" + default_email = f"env-user-{credential_index}" + else: + # Legacy format: PROVIDER_ACCESS_TOKEN + prefix = self.ENV_PREFIX + default_email = "env-user" + + access_token = os.getenv(f"{prefix}_ACCESS_TOKEN") + refresh_token = os.getenv(f"{prefix}_REFRESH_TOKEN") # Both access and refresh tokens are required if not (access_token and refresh_token): return None - lib_logger.debug(f"Loading {self.ENV_PREFIX} credentials from environment variables") + lib_logger.debug(f"Loading {prefix} credentials from environment variables") # Parse expiry_date as float, default to 0 if not present - expiry_str = os.getenv(f"{self.ENV_PREFIX}_EXPIRY_DATE", "0") + expiry_str = os.getenv(f"{prefix}_EXPIRY_DATE", "0") try: expiry_date = float(expiry_str) except ValueError: - lib_logger.warning(f"Invalid {self.ENV_PREFIX}_EXPIRY_DATE value: {expiry_str}, using 0") + lib_logger.warning(f"Invalid {prefix}_EXPIRY_DATE value: {expiry_str}, using 0") expiry_date = 0 creds = { "access_token": access_token, "refresh_token": refresh_token, "expiry_date": expiry_date, - "client_id": os.getenv(f"{self.ENV_PREFIX}_CLIENT_ID", self.CLIENT_ID), - "client_secret": os.getenv(f"{self.ENV_PREFIX}_CLIENT_SECRET", self.CLIENT_SECRET), - "token_uri": os.getenv(f"{self.ENV_PREFIX}_TOKEN_URI", self.TOKEN_URI), - "universe_domain": os.getenv(f"{self.ENV_PREFIX}_UNIVERSE_DOMAIN", "googleapis.com"), + "client_id": os.getenv(f"{prefix}_CLIENT_ID", self.CLIENT_ID), + "client_secret": os.getenv(f"{prefix}_CLIENT_SECRET", self.CLIENT_SECRET), + "token_uri": os.getenv(f"{prefix}_TOKEN_URI", self.TOKEN_URI), + "universe_domain": os.getenv(f"{prefix}_UNIVERSE_DOMAIN", "googleapis.com"), "_proxy_metadata": { - "email": os.getenv(f"{self.ENV_PREFIX}_EMAIL", "env-user"), + "email": os.getenv(f"{prefix}_EMAIL", default_email), "last_check_timestamp": time.time(), - "loaded_from_env": True # Flag to indicate env-based credentials + "loaded_from_env": True, # Flag to indicate env-based credentials + "env_credential_index": credential_index or "0" # Track which env credential this is } } # Add project_id if provided - project_id = os.getenv(f"{self.ENV_PREFIX}_PROJECT_ID") + project_id = os.getenv(f"{prefix}_PROJECT_ID") if project_id: creds["_proxy_metadata"]["project_id"] = project_id # Add tier if provided - tier = os.getenv(f"{self.ENV_PREFIX}_TIER") + tier = os.getenv(f"{prefix}_TIER") if tier: creds["_proxy_metadata"]["tier"] = tier @@ -148,7 +187,19 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: if path in self._credentials_cache: return self._credentials_cache[path] - # First, try loading from environment variables + # Check if this is a virtual env:// path + credential_index = self._parse_env_credential_path(path) + if credential_index is not None: + # Load from environment variables with specific index + env_creds = self._load_from_env(credential_index) + if env_creds: + lib_logger.info(f"Using {self.ENV_PREFIX} credentials from environment variables (index: {credential_index})") + self._credentials_cache[path] = env_creds + return env_creds + else: + raise IOError(f"Environment variables for {self.ENV_PREFIX} credential index {credential_index} not found") + + # For file paths, first try loading from legacy env vars (for backwards compatibility) env_creds = self._load_from_env() if env_creds: lib_logger.info(f"Using {self.ENV_PREFIX} credentials from environment variables") @@ -170,6 +221,8 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: raise IOError(f"{self.ENV_PREFIX} OAuth credential file not found at '{path}'") except Exception as e: raise IOError(f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}") + except Exception as e: + raise IOError(f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}") async def _save_credentials(self, path: str, creds: Dict[str, Any]): # Don't save to file if credentials were loaded from environment diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index 4d77b79b..f6618f7f 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -158,47 +158,79 @@ def __init__(self): self._queue_tracking_lock = asyncio.Lock() # Protects queue sets self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task - def _load_from_env(self) -> Optional[Dict[str, Any]]: + def _parse_env_credential_path(self, path: str) -> Optional[str]: + """ + Parse a virtual env:// path and return the credential index. + + Supported formats: + - "env://provider/0" - Legacy single credential (no index in env var names) + - "env://provider/1" - First numbered credential (IFLOW_1_ACCESS_TOKEN) + + Returns: + The credential index as string, or None if path is not an env:// path + """ + if not path.startswith("env://"): + return None + + parts = path[6:].split("/") + if len(parts) >= 2: + return parts[1] + return "0" + + def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dict[str, Any]]: """ Load OAuth credentials from environment variables for stateless deployments. - Expected environment variables: - - IFLOW_ACCESS_TOKEN (required) - - IFLOW_REFRESH_TOKEN (required) - - IFLOW_API_KEY (required - critical for iFlow!) - - IFLOW_EXPIRY_DATE (optional, defaults to empty string) - - IFLOW_EMAIL (optional, defaults to "env-user") - - IFLOW_TOKEN_TYPE (optional, defaults to "Bearer") - - IFLOW_SCOPE (optional, defaults to "read write") + Supports two formats: + 1. Legacy (credential_index="0" or None): IFLOW_ACCESS_TOKEN + 2. Numbered (credential_index="1", "2", etc.): IFLOW_1_ACCESS_TOKEN, etc. + + Expected environment variables (for numbered format with index N): + - IFLOW_{N}_ACCESS_TOKEN (required) + - IFLOW_{N}_REFRESH_TOKEN (required) + - IFLOW_{N}_API_KEY (required - critical for iFlow!) + - IFLOW_{N}_EXPIRY_DATE (optional, defaults to empty string) + - IFLOW_{N}_EMAIL (optional, defaults to "env-user-{N}") + - IFLOW_{N}_TOKEN_TYPE (optional, defaults to "Bearer") + - IFLOW_{N}_SCOPE (optional, defaults to "read write") Returns: Dict with credential structure if env vars present, None otherwise """ - access_token = os.getenv("IFLOW_ACCESS_TOKEN") - refresh_token = os.getenv("IFLOW_REFRESH_TOKEN") - api_key = os.getenv("IFLOW_API_KEY") + # Determine the env var prefix based on credential index + if credential_index and credential_index != "0": + prefix = f"IFLOW_{credential_index}" + default_email = f"env-user-{credential_index}" + else: + prefix = "IFLOW" + default_email = "env-user" + + access_token = os.getenv(f"{prefix}_ACCESS_TOKEN") + refresh_token = os.getenv(f"{prefix}_REFRESH_TOKEN") + api_key = os.getenv(f"{prefix}_API_KEY") # All three are required for iFlow if not (access_token and refresh_token and api_key): return None - lib_logger.debug("Loading iFlow credentials from environment variables") + lib_logger.debug(f"Loading iFlow credentials from environment variables (prefix: {prefix})") # Parse expiry_date as string (ISO 8601 format) - expiry_str = os.getenv("IFLOW_EXPIRY_DATE", "") + expiry_str = os.getenv(f"{prefix}_EXPIRY_DATE", "") creds = { "access_token": access_token, "refresh_token": refresh_token, "api_key": api_key, # Critical for iFlow! "expiry_date": expiry_str, - "email": os.getenv("IFLOW_EMAIL", "env-user"), - "token_type": os.getenv("IFLOW_TOKEN_TYPE", "Bearer"), - "scope": os.getenv("IFLOW_SCOPE", "read write"), + "email": os.getenv(f"{prefix}_EMAIL", default_email), + "token_type": os.getenv(f"{prefix}_TOKEN_TYPE", "Bearer"), + "scope": os.getenv(f"{prefix}_SCOPE", "read write"), "_proxy_metadata": { - "email": os.getenv("IFLOW_EMAIL", "env-user"), + "email": os.getenv(f"{prefix}_EMAIL", default_email), "last_check_timestamp": time.time(), - "loaded_from_env": True # Flag to indicate env-based credentials + "loaded_from_env": True, + "env_credential_index": credential_index or "0" } } @@ -227,11 +259,21 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: if path in self._credentials_cache: return self._credentials_cache[path] - # First, try loading from environment variables + # Check if this is a virtual env:// path + credential_index = self._parse_env_credential_path(path) + if credential_index is not None: + env_creds = self._load_from_env(credential_index) + if env_creds: + lib_logger.info(f"Using iFlow credentials from environment variables (index: {credential_index})") + self._credentials_cache[path] = env_creds + return env_creds + else: + raise IOError(f"Environment variables for iFlow credential index {credential_index} not found") + + # For file paths, try loading from legacy env vars first env_creds = self._load_from_env() if env_creds: lib_logger.info("Using iFlow credentials from environment variables") - # Cache env-based credentials using the path as key self._credentials_cache[path] = env_creds return env_creds diff --git a/src/rotator_library/providers/qwen_auth_base.py b/src/rotator_library/providers/qwen_auth_base.py index 9d028c7a..58db90e9 100644 --- a/src/rotator_library/providers/qwen_auth_base.py +++ b/src/rotator_library/providers/qwen_auth_base.py @@ -47,46 +47,78 @@ def __init__(self): self._queue_tracking_lock = asyncio.Lock() # Protects queue sets self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task - def _load_from_env(self) -> Optional[Dict[str, Any]]: + def _parse_env_credential_path(self, path: str) -> Optional[str]: + """ + Parse a virtual env:// path and return the credential index. + + Supported formats: + - "env://provider/0" - Legacy single credential (no index in env var names) + - "env://provider/1" - First numbered credential (QWEN_CODE_1_ACCESS_TOKEN) + + Returns: + The credential index as string, or None if path is not an env:// path + """ + if not path.startswith("env://"): + return None + + parts = path[6:].split("/") + if len(parts) >= 2: + return parts[1] + return "0" + + def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dict[str, Any]]: """ Load OAuth credentials from environment variables for stateless deployments. - Expected environment variables: - - QWEN_CODE_ACCESS_TOKEN (required) - - QWEN_CODE_REFRESH_TOKEN (required) - - QWEN_CODE_EXPIRY_DATE (optional, defaults to 0) - - QWEN_CODE_RESOURCE_URL (optional, defaults to https://portal.qwen.ai/v1) - - QWEN_CODE_EMAIL (optional, defaults to "env-user") + Supports two formats: + 1. Legacy (credential_index="0" or None): QWEN_CODE_ACCESS_TOKEN + 2. Numbered (credential_index="1", "2", etc.): QWEN_CODE_1_ACCESS_TOKEN, etc. + + Expected environment variables (for numbered format with index N): + - QWEN_CODE_{N}_ACCESS_TOKEN (required) + - QWEN_CODE_{N}_REFRESH_TOKEN (required) + - QWEN_CODE_{N}_EXPIRY_DATE (optional, defaults to 0) + - QWEN_CODE_{N}_RESOURCE_URL (optional, defaults to https://portal.qwen.ai/v1) + - QWEN_CODE_{N}_EMAIL (optional, defaults to "env-user-{N}") Returns: Dict with credential structure if env vars present, None otherwise """ - access_token = os.getenv("QWEN_CODE_ACCESS_TOKEN") - refresh_token = os.getenv("QWEN_CODE_REFRESH_TOKEN") + # Determine the env var prefix based on credential index + if credential_index and credential_index != "0": + prefix = f"QWEN_CODE_{credential_index}" + default_email = f"env-user-{credential_index}" + else: + prefix = "QWEN_CODE" + default_email = "env-user" + + access_token = os.getenv(f"{prefix}_ACCESS_TOKEN") + refresh_token = os.getenv(f"{prefix}_REFRESH_TOKEN") # Both access and refresh tokens are required if not (access_token and refresh_token): return None - lib_logger.debug("Loading Qwen Code credentials from environment variables") + lib_logger.debug(f"Loading Qwen Code credentials from environment variables (prefix: {prefix})") # Parse expiry_date as float, default to 0 if not present - expiry_str = os.getenv("QWEN_CODE_EXPIRY_DATE", "0") + expiry_str = os.getenv(f"{prefix}_EXPIRY_DATE", "0") try: expiry_date = float(expiry_str) except ValueError: - lib_logger.warning(f"Invalid QWEN_CODE_EXPIRY_DATE value: {expiry_str}, using 0") + lib_logger.warning(f"Invalid {prefix}_EXPIRY_DATE value: {expiry_str}, using 0") expiry_date = 0 creds = { "access_token": access_token, "refresh_token": refresh_token, "expiry_date": expiry_date, - "resource_url": os.getenv("QWEN_CODE_RESOURCE_URL", "https://portal.qwen.ai/v1"), + "resource_url": os.getenv(f"{prefix}_RESOURCE_URL", "https://portal.qwen.ai/v1"), "_proxy_metadata": { - "email": os.getenv("QWEN_CODE_EMAIL", "env-user"), + "email": os.getenv(f"{prefix}_EMAIL", default_email), "last_check_timestamp": time.time(), - "loaded_from_env": True # Flag to indicate env-based credentials + "loaded_from_env": True, + "env_credential_index": credential_index or "0" } } @@ -115,11 +147,21 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: if path in self._credentials_cache: return self._credentials_cache[path] - # First, try loading from environment variables + # Check if this is a virtual env:// path + credential_index = self._parse_env_credential_path(path) + if credential_index is not None: + env_creds = self._load_from_env(credential_index) + if env_creds: + lib_logger.info(f"Using Qwen Code credentials from environment variables (index: {credential_index})") + self._credentials_cache[path] = env_creds + return env_creds + else: + raise IOError(f"Environment variables for Qwen Code credential index {credential_index} not found") + + # For file paths, try loading from legacy env vars first env_creds = self._load_from_env() if env_creds: lib_logger.info("Using Qwen Code credentials from environment variables") - # Cache env-based credentials using the path as key self._credentials_cache[path] = env_creds return env_creds diff --git a/todo.md b/todo.md new file mode 100644 index 00000000..5966e4b1 --- /dev/null +++ b/todo.md @@ -0,0 +1,7 @@ +~~Refine claude injection to inject even if we have correct thinking - to force it to think if we made ultrathink prompt. If last msg is tool use and you prompt - it never thinks again.~~ Maybe done + +Anthropic translation and anthropic compatible endpoint. + +Refine for deployment. + + From d94742e00149a793f0a8328e279df153f58b475a Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 17:16:34 +0100 Subject: [PATCH 041/221] =?UTF-8?q?fix(auth):=20=F0=9F=90=9B=20add=20expon?= =?UTF-8?q?ential=20backoff=20and=20validation=20for=20token=20refresh=20f?= =?UTF-8?q?ailures?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit improves the robustness of OAuth token refresh operations in both IFlowAuthBase and QwenAuthBase by implementing failure tracking with exponential backoff and credential validation. - Track refresh failures per credential path using `_refresh_failures` dictionary - Implement exponential backoff (30s * 2^failures, max 5 minutes) to prevent rapid retry loops on persistent failures - Clear backoff state on successful authentication or refresh - Add validation to ensure refreshed credentials contain required fields (access_token, refresh_token, and api_key for iFlow) - Update proactively_refresh to support env:// virtual paths for environment-based OAuth credentials - Add detailed debug logging for backoff timer settings The backoff mechanism prevents excessive API calls when refresh tokens are invalid or services are temporarily unavailable, while the validation ensures credential integrity after refresh operations. --- .../providers/iflow_auth_base.py | 32 +++++++++++++++++-- .../providers/qwen_auth_base.py | 32 +++++++++++++++++-- 2 files changed, 58 insertions(+), 6 deletions(-) diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index f6618f7f..cae85928 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -551,12 +551,25 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] try: # Call initialize_token to trigger OAuth flow new_creds = await self.initialize_token(path) + # Clear backoff on successful re-auth + self._refresh_failures.pop(path, None) + self._next_refresh_after.pop(path, None) return new_creds except Exception as reauth_error: lib_logger.error(f"Re-authentication failed for '{Path(path).name}': {reauth_error}") + # [BACKOFF TRACKING] Increment failure count and set backoff timer + self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 + backoff_seconds = min(300, 30 * (2 ** self._refresh_failures[path])) # Max 5 min backoff + self._next_refresh_after[path] = time.time() + backoff_seconds + lib_logger.debug(f"Setting backoff for '{Path(path).name}': {backoff_seconds}s") raise ValueError(f"Refresh token invalid and re-authentication failed: {reauth_error}") if new_token_data is None: + # [BACKOFF TRACKING] Increment failure count and set backoff timer + self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 + backoff_seconds = min(300, 30 * (2 ** self._refresh_failures[path])) # Max 5 min backoff + self._next_refresh_after[path] = time.time() + backoff_seconds + lib_logger.debug(f"Setting backoff for '{Path(path).name}': {backoff_seconds}s") raise last_error or Exception("Token refresh failed after all retries") # Update tokens @@ -589,6 +602,16 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] creds_from_file["_proxy_metadata"] = {} creds_from_file["_proxy_metadata"]["last_check_timestamp"] = time.time() + # [VALIDATION] Verify required fields exist after refresh + required_fields = ["access_token", "refresh_token", "api_key"] + missing_fields = [field for field in required_fields if not creds_from_file.get(field)] + if missing_fields: + raise ValueError(f"Refreshed credentials missing required fields: {missing_fields}") + + # [BACKOFF TRACKING] Clear failure count on successful refresh + self._refresh_failures.pop(path, None) + self._next_refresh_after.pop(path, None) + await self._save_credentials(path, creds_from_file) lib_logger.debug(f"Successfully refreshed iFlow OAuth token for '{Path(path).name}'.") return creds_from_file @@ -626,10 +649,13 @@ async def get_api_details(self, credential_identifier: str) -> Tuple[str, str]: async def proactively_refresh(self, credential_identifier: str): """ Proactively refreshes tokens if they're close to expiry. - Only applies to OAuth credentials (file paths). Direct API keys are skipped. + Only applies to OAuth credentials (file paths or env:// paths). Direct API keys are skipped. """ - # Only refresh if it's an OAuth credential (file path) - if not os.path.isfile(credential_identifier): + # Check if it's an env:// virtual path (OAuth credentials from environment) + is_env_path = credential_identifier.startswith("env://") + + # Only refresh if it's an OAuth credential (file path or env:// path) + if not is_env_path and not os.path.isfile(credential_identifier): return # Direct API key, no refresh needed creds = await self._load_credentials(credential_identifier) diff --git a/src/rotator_library/providers/qwen_auth_base.py b/src/rotator_library/providers/qwen_auth_base.py index 58db90e9..589e6bef 100644 --- a/src/rotator_library/providers/qwen_auth_base.py +++ b/src/rotator_library/providers/qwen_auth_base.py @@ -316,12 +316,25 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] try: # Call initialize_token to trigger OAuth flow new_creds = await self.initialize_token(path) + # Clear backoff on successful re-auth + self._refresh_failures.pop(path, None) + self._next_refresh_after.pop(path, None) return new_creds except Exception as reauth_error: lib_logger.error(f"Re-authentication failed for '{Path(path).name}': {reauth_error}") + # [BACKOFF TRACKING] Increment failure count and set backoff timer + self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 + backoff_seconds = min(300, 30 * (2 ** self._refresh_failures[path])) # Max 5 min backoff + self._next_refresh_after[path] = time.time() + backoff_seconds + lib_logger.debug(f"Setting backoff for '{Path(path).name}': {backoff_seconds}s") raise ValueError(f"Refresh token invalid and re-authentication failed: {reauth_error}") if new_token_data is None: + # [BACKOFF TRACKING] Increment failure count and set backoff timer + self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 + backoff_seconds = min(300, 30 * (2 ** self._refresh_failures[path])) # Max 5 min backoff + self._next_refresh_after[path] = time.time() + backoff_seconds + lib_logger.debug(f"Setting backoff for '{Path(path).name}': {backoff_seconds}s") raise last_error or Exception("Token refresh failed after all retries") creds_from_file["access_token"] = new_token_data["access_token"] @@ -334,6 +347,16 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] creds_from_file["_proxy_metadata"] = {} creds_from_file["_proxy_metadata"]["last_check_timestamp"] = time.time() + # [VALIDATION] Verify required fields exist after refresh + required_fields = ["access_token", "refresh_token"] + missing_fields = [field for field in required_fields if not creds_from_file.get(field)] + if missing_fields: + raise ValueError(f"Refreshed credentials missing required fields: {missing_fields}") + + # [BACKOFF TRACKING] Clear failure count on successful refresh + self._refresh_failures.pop(path, None) + self._next_refresh_after.pop(path, None) + await self._save_credentials(path, creds_from_file) lib_logger.debug(f"Successfully refreshed Qwen OAuth token for '{Path(path).name}'.") return creds_from_file @@ -370,10 +393,13 @@ async def get_api_details(self, credential_identifier: str) -> Tuple[str, str]: async def proactively_refresh(self, credential_identifier: str): """ Proactively refreshes tokens if they're close to expiry. - Only applies to OAuth credentials (file paths). Direct API keys are skipped. + Only applies to OAuth credentials (file paths or env:// paths). Direct API keys are skipped. """ - # Only refresh if it's an OAuth credential (file path) - if not os.path.isfile(credential_identifier): + # Check if it's an env:// virtual path (OAuth credentials from environment) + is_env_path = credential_identifier.startswith("env://") + + # Only refresh if it's an OAuth credential (file path or env:// path) + if not is_env_path and not os.path.isfile(credential_identifier): return # Direct API key, no refresh needed creds = await self._load_credentials(credential_identifier) From f6dce021ef65262de60851ffdfcf415d591ddb1e Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 17:22:16 +0100 Subject: [PATCH 042/221] =?UTF-8?q?fix(providers):=20=F0=9F=90=9B=20improv?= =?UTF-8?q?e=20finish=5Freason=20handling=20and=20tool=5Fcalls=20initializ?= =?UTF-8?q?ation=20in=20stream=20reassembly?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit addresses critical issues in the streaming response reassembly logic across multiple providers (Gemini CLI, iFlow, and Qwen Code): - Implements priority-based finish_reason determination: tool_calls > chunk's finish_reason (length, content_filter, etc.) > stop - Properly initializes aggregated_tool_calls with "type": "function" field for OpenAI compatibility - Tracks chunk_finish_reason separately to preserve provider-specific finish reasons (e.g., content_filter, length limits) - Uses safer .get("index", 0) for tool call index extraction to prevent KeyErrors - Adds explicit type field handling during tool call aggregation - Improves docstring documentation explaining the reassembly logic - Moves copy import to top-level in iflow_provider.py and qwen_code_provider.py for consistency CRITICAL FIX for qwen_code_provider.py: Handles chunks with BOTH usage and choices data (typical for final chunk) without early return, ensuring finish_reason is properly captured before yielding usage data separately. --- .../providers/gemini_cli_provider.py | 23 +++-- .../providers/iflow_provider.py | 29 +++++-- .../providers/qwen_code_provider.py | 83 +++++++++++++++---- 3 files changed, 105 insertions(+), 30 deletions(-) diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 0a0ab514..bd85283e 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -998,7 +998,11 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumul def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> litellm.ModelResponse: """ Manually reassembles streaming chunks into a complete response. - This replaces the non-existent litellm.utils.stream_to_completion_response function. + + Key improvements: + - Determines finish_reason based on accumulated state + - Priority: tool_calls > chunk's finish_reason (length, content_filter, etc.) > stop + - Properly initializes tool_calls with type field """ if not chunks: raise ValueError("No chunks provided for reassembly") @@ -1007,7 +1011,7 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> final_message = {"role": "assistant"} aggregated_tool_calls = {} usage_data = None - finish_reason = None + chunk_finish_reason = None # Track finish_reason from chunks # Get the first chunk for basic response metadata first_chunk = chunks[0] @@ -1035,11 +1039,13 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> # Aggregate tool calls if "tool_calls" in delta and delta["tool_calls"]: for tc_chunk in delta["tool_calls"]: - index = tc_chunk["index"] + index = tc_chunk.get("index", 0) if index not in aggregated_tool_calls: aggregated_tool_calls[index] = {"type": "function", "function": {"name": "", "arguments": ""}} if "id" in tc_chunk: aggregated_tool_calls[index]["id"] = tc_chunk["id"] + if "type" in tc_chunk: + aggregated_tool_calls[index]["type"] = tc_chunk["type"] if "function" in tc_chunk: if "name" in tc_chunk["function"] and tc_chunk["function"]["name"] is not None: aggregated_tool_calls[index]["function"]["name"] += tc_chunk["function"]["name"] @@ -1055,8 +1061,9 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> if "arguments" in delta["function_call"] and delta["function_call"]["arguments"] is not None: final_message["function_call"]["arguments"] += delta["function_call"]["arguments"] - # Note: chunks don't include finish_reason (client handles it) - # This is kept for compatibility but shouldn't trigger + # Track finish_reason from chunks (respects length, content_filter, etc.) + if choice.get("finish_reason"): + chunk_finish_reason = choice["finish_reason"] # Handle usage data from the last chunk that has it for chunk in reversed(chunks): @@ -1073,10 +1080,12 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> if field not in final_message: final_message[field] = None - # Determine finish_reason based on content (same logic as client.py) - # tool_calls wins, otherwise stop + # Determine finish_reason based on accumulated state + # Priority: tool_calls wins if present, then chunk's finish_reason (length, content_filter, etc.), then default to "stop" if aggregated_tool_calls: finish_reason = "tool_calls" + elif chunk_finish_reason: + finish_reason = chunk_finish_reason else: finish_reason = "stop" diff --git a/src/rotator_library/providers/iflow_provider.py b/src/rotator_library/providers/iflow_provider.py index b6021127..28d84f64 100644 --- a/src/rotator_library/providers/iflow_provider.py +++ b/src/rotator_library/providers/iflow_provider.py @@ -1,5 +1,6 @@ # src/rotator_library/providers/iflow_provider.py +import copy import json import time import os @@ -203,7 +204,6 @@ def _clean_tool_schemas(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any Removes unsupported properties from tool schemas to prevent API errors. Similar to Qwen Code implementation. """ - import copy cleaned_tools = [] for tool in tools: @@ -345,6 +345,11 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> litellm.ModelResponse: """ Manually reassembles streaming chunks into a complete response. + + Key improvements: + - Determines finish_reason based on accumulated state (tool_calls vs stop) + - Properly initializes tool_calls with type field + - Handles usage data extraction from chunks """ if not chunks: raise ValueError("No chunks provided for reassembly") @@ -353,7 +358,7 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> final_message = {"role": "assistant"} aggregated_tool_calls = {} usage_data = None - finish_reason = None + chunk_finish_reason = None # Track finish_reason from chunks (but we'll override) # Get the first chunk for basic response metadata first_chunk = chunks[0] @@ -378,12 +383,13 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> final_message["reasoning_content"] = "" final_message["reasoning_content"] += delta["reasoning_content"] - # Aggregate tool calls + # Aggregate tool calls with proper initialization if "tool_calls" in delta and delta["tool_calls"]: for tc_chunk in delta["tool_calls"]: - index = tc_chunk["index"] + index = tc_chunk.get("index", 0) if index not in aggregated_tool_calls: - aggregated_tool_calls[index] = {"function": {"name": "", "arguments": ""}} + # Initialize with type field for OpenAI compatibility + aggregated_tool_calls[index] = {"type": "function", "function": {"name": "", "arguments": ""}} if "id" in tc_chunk: aggregated_tool_calls[index]["id"] = tc_chunk["id"] if "type" in tc_chunk: @@ -403,9 +409,9 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> if "arguments" in delta["function_call"] and delta["function_call"]["arguments"] is not None: final_message["function_call"]["arguments"] += delta["function_call"]["arguments"] - # Get finish reason from the last chunk that has it + # Track finish_reason from chunks (for reference only) if choice.get("finish_reason"): - finish_reason = choice["finish_reason"] + chunk_finish_reason = choice["finish_reason"] # Handle usage data from the last chunk that has it for chunk in reversed(chunks): @@ -422,6 +428,15 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> if field not in final_message: final_message[field] = None + # Determine finish_reason based on accumulated state + # Priority: tool_calls wins if present, then chunk's finish_reason, then default to "stop" + if aggregated_tool_calls: + finish_reason = "tool_calls" + elif chunk_finish_reason: + finish_reason = chunk_finish_reason + else: + finish_reason = "stop" + # Construct the final response final_choice = { "index": 0, diff --git a/src/rotator_library/providers/qwen_code_provider.py b/src/rotator_library/providers/qwen_code_provider.py index d57c88dd..334e3142 100644 --- a/src/rotator_library/providers/qwen_code_provider.py +++ b/src/rotator_library/providers/qwen_code_provider.py @@ -1,5 +1,6 @@ # src/rotator_library/providers/qwen_code_provider.py +import copy import json import time import os @@ -186,7 +187,6 @@ def _clean_tool_schemas(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any Removes unsupported properties from tool schemas to prevent API errors. Adapted for Qwen's API requirements. """ - import copy cleaned_tools = [] for tool in tools: @@ -263,15 +263,38 @@ def _build_request_payload(self, **kwargs) -> Dict[str, Any]: return payload def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): - """Converts a raw Qwen SSE chunk to an OpenAI-compatible chunk.""" + """ + Converts a raw Qwen SSE chunk to an OpenAI-compatible chunk. + + CRITICAL FIX: Handle chunks with BOTH usage and choices (final chunk) + without early return to ensure finish_reason is properly processed. + """ if not isinstance(chunk, dict): return - # Handle usage data - if usage_data := chunk.get("usage"): + # Get choices and usage data + choices = chunk.get("choices", []) + usage_data = chunk.get("usage") + chunk_id = chunk.get("id", f"chatcmpl-qwen-{time.time()}") + chunk_created = chunk.get("created", int(time.time())) + + # Handle chunks with BOTH choices and usage (typical for final chunk) + # CRITICAL: Process choices FIRST to capture finish_reason, then yield usage + if choices and usage_data: + choice = choices[0] + delta = choice.get("delta", {}) + finish_reason = choice.get("finish_reason") + + # Yield the choice chunk first (contains finish_reason) + yield { + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + "model": model_id, "object": "chat.completion.chunk", + "id": chunk_id, "created": chunk_created + } + # Then yield the usage chunk yield { "choices": [], "model": model_id, "object": "chat.completion.chunk", - "id": f"chatcmpl-qwen-{time.time()}", "created": int(time.time()), + "id": chunk_id, "created": chunk_created, "usage": { "prompt_tokens": usage_data.get("prompt_tokens", 0), "completion_tokens": usage_data.get("completion_tokens", 0), @@ -280,8 +303,20 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): } return - # Handle content data - choices = chunk.get("choices", []) + # Handle usage-only chunks + if usage_data: + yield { + "choices": [], "model": model_id, "object": "chat.completion.chunk", + "id": chunk_id, "created": chunk_created, + "usage": { + "prompt_tokens": usage_data.get("prompt_tokens", 0), + "completion_tokens": usage_data.get("completion_tokens", 0), + "total_tokens": usage_data.get("total_tokens", 0), + } + } + return + + # Handle content-only chunks if not choices: return @@ -307,20 +342,24 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str): yield { "choices": [{"index": 0, "delta": new_delta, "finish_reason": None}], "model": model_id, "object": "chat.completion.chunk", - "id": f"chatcmpl-qwen-{time.time()}", "created": int(time.time()) + "id": chunk_id, "created": chunk_created } else: # Standard content chunk yield { "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], "model": model_id, "object": "chat.completion.chunk", - "id": f"chatcmpl-qwen-{time.time()}", "created": int(time.time()) + "id": chunk_id, "created": chunk_created } def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> litellm.ModelResponse: """ Manually reassembles streaming chunks into a complete response. - This replaces the non-existent litellm.utils.stream_to_completion_response function. + + Key improvements: + - Determines finish_reason based on accumulated state (tool_calls vs stop) + - Properly initializes tool_calls with type field + - Handles usage data extraction from chunks """ if not chunks: raise ValueError("No chunks provided for reassembly") @@ -329,7 +368,7 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> final_message = {"role": "assistant"} aggregated_tool_calls = {} usage_data = None - finish_reason = None + chunk_finish_reason = None # Track finish_reason from chunks (but we'll override) # Get the first chunk for basic response metadata first_chunk = chunks[0] @@ -354,14 +393,17 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> final_message["reasoning_content"] = "" final_message["reasoning_content"] += delta["reasoning_content"] - # Aggregate tool calls + # Aggregate tool calls with proper initialization if "tool_calls" in delta and delta["tool_calls"]: for tc_chunk in delta["tool_calls"]: - index = tc_chunk["index"] + index = tc_chunk.get("index", 0) if index not in aggregated_tool_calls: - aggregated_tool_calls[index] = {"function": {"name": "", "arguments": ""}} + # Initialize with type field for OpenAI compatibility + aggregated_tool_calls[index] = {"type": "function", "function": {"name": "", "arguments": ""}} if "id" in tc_chunk: aggregated_tool_calls[index]["id"] = tc_chunk["id"] + if "type" in tc_chunk: + aggregated_tool_calls[index]["type"] = tc_chunk["type"] if "function" in tc_chunk: if "name" in tc_chunk["function"] and tc_chunk["function"]["name"] is not None: aggregated_tool_calls[index]["function"]["name"] += tc_chunk["function"]["name"] @@ -377,9 +419,9 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> if "arguments" in delta["function_call"] and delta["function_call"]["arguments"] is not None: final_message["function_call"]["arguments"] += delta["function_call"]["arguments"] - # Get finish reason from the last chunk that has it + # Track finish_reason from chunks (for reference only) if choice.get("finish_reason"): - finish_reason = choice["finish_reason"] + chunk_finish_reason = choice["finish_reason"] # Handle usage data from the last chunk that has it for chunk in reversed(chunks): @@ -396,6 +438,15 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> if field not in final_message: final_message[field] = None + # Determine finish_reason based on accumulated state + # Priority: tool_calls wins if present, then chunk's finish_reason, then default to "stop" + if aggregated_tool_calls: + finish_reason = "tool_calls" + elif chunk_finish_reason: + finish_reason = chunk_finish_reason + else: + finish_reason = "stop" + # Construct the final response final_choice = { "index": 0, From 2384d8699c5bc4b23d49373bfd64aa5a4a096204 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 17:26:22 +0100 Subject: [PATCH 043/221] =?UTF-8?q?fix(proxy):=20=F0=9F=90=9B=20load=20env?= =?UTF-8?q?ironment=20variables=20before=20displaying=20PROXY=5FAPI=5FKEY?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The .env file was being loaded after attempting to read PROXY_API_KEY from environment variables, causing the key to be unavailable for display during startup. Moving the dotenv.load_dotenv() call earlier in the initialization sequence ensures environment variables are loaded before they are accessed. --- src/proxy_app/main.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 43b2d2d3..dfbc0418 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -38,6 +38,10 @@ # If we get here, we're ACTUALLY running the proxy - NOW show startup messages and start timer _start_time = time.time() +# Load .env early so PROXY_API_KEY is available for display +from dotenv import load_dotenv +load_dotenv() + # Get proxy API key for display proxy_api_key = os.getenv("PROXY_API_KEY") if proxy_api_key: From 64859d936e50eecfe4e438a193df2c93e291c0ab Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 17:32:47 +0100 Subject: [PATCH 044/221] =?UTF-8?q?feat(settings):=20=E2=9C=A8=20add=20pro?= =?UTF-8?q?vider-specific=20settings=20management=20UI?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a comprehensive provider-specific settings management system for Antigravity and Gemini CLI providers with detection, display, and interactive configuration capabilities. - Add `PROVIDER_SETTINGS_MAP` with detailed definitions for Antigravity (12 settings) and Gemini CLI (8 settings) including signature caching, tool fixes, and provider-specific parameters - Implement `ProviderSettingsManager` class for managing provider settings with type-aware value parsing and modification tracking - Add `detect_provider_settings()` method to `SettingsDetector` to identify modified provider settings from environment variables - Integrate provider settings detection into launcher TUI summary display and detailed advanced settings view - Add new menu option (4) in settings tool for provider-specific configuration management - Implement interactive TUI for browsing, editing, and resetting individual or all provider settings with visual indication of modified values - Display provider settings status in launcher with count of modified settings per provider - Support bool, int, and string setting types with appropriate input handling and validation --- src/proxy_app/launcher_tui.py | 64 +++++- src/proxy_app/settings_tool.py | 379 ++++++++++++++++++++++++++++++++- 2 files changed, 436 insertions(+), 7 deletions(-) diff --git a/src/proxy_app/launcher_tui.py b/src/proxy_app/launcher_tui.py index a14c0aea..26a36bf1 100644 --- a/src/proxy_app/launcher_tui.py +++ b/src/proxy_app/launcher_tui.py @@ -100,7 +100,8 @@ def get_all_settings() -> dict: "custom_bases": SettingsDetector.detect_custom_api_bases(), "model_definitions": SettingsDetector.detect_model_definitions(), "concurrency_limits": SettingsDetector.detect_concurrency_limits(), - "model_filters": SettingsDetector.detect_model_filters() + "model_filters": SettingsDetector.detect_model_filters(), + "provider_settings": SettingsDetector.detect_provider_settings() } @staticmethod @@ -198,6 +199,45 @@ def detect_model_filters() -> dict: else: filters[provider]["has_whitelist"] = True return filters + + @staticmethod + def detect_provider_settings() -> dict: + """Detect provider-specific settings (Antigravity, Gemini CLI)""" + try: + from proxy_app.settings_tool import PROVIDER_SETTINGS_MAP + except ImportError: + # Fallback for direct execution or testing + from .settings_tool import PROVIDER_SETTINGS_MAP + + provider_settings = {} + env_vars = SettingsDetector._load_local_env() + + for provider, definitions in PROVIDER_SETTINGS_MAP.items(): + modified_count = 0 + for key, definition in definitions.items(): + env_value = env_vars.get(key) + if env_value is not None: + # Check if value differs from default + default = definition.get("default") + setting_type = definition.get("type", "str") + + try: + if setting_type == "bool": + current = env_value.lower() in ("true", "1", "yes") + elif setting_type == "int": + current = int(env_value) + else: + current = env_value + + if current != default: + modified_count += 1 + except (ValueError, AttributeError): + pass + + if modified_count > 0: + provider_settings[provider] = modified_count + + return provider_settings class LauncherTUI: @@ -300,7 +340,8 @@ def show_main_menu(self): self.console.print("━" * 70) provider_count = len(credentials) custom_count = len(custom_bases) - has_advanced = bool(settings["model_definitions"] or settings["concurrency_limits"] or settings["model_filters"]) + provider_settings = settings.get("provider_settings", {}) + has_advanced = bool(settings["model_definitions"] or settings["concurrency_limits"] or settings["model_filters"] or provider_settings) self.console.print(f" Providers: {provider_count} configured") self.console.print(f" Custom Providers: {custom_count} configured") @@ -422,6 +463,7 @@ def show_provider_settings_menu(self): model_defs = settings["model_definitions"] concurrency = settings["concurrency_limits"] filters = settings["model_filters"] + provider_settings = settings.get("provider_settings", {}) self.console.print(Panel.fit( "[bold cyan]📊 Provider & Advanced Settings[/bold cyan]", @@ -472,7 +514,7 @@ def show_provider_settings_menu(self): self.console.print("━" * 70) for provider, limit in concurrency.items(): self.console.print(f" • {provider:15} {limit} requests/key") - self.console.print(f" • Default: 1 request/key (all others)") + self.console.print(" • Default: 1 request/key (all others)") # Model Filters (basic info only) if filters: @@ -488,6 +530,22 @@ def show_provider_settings_menu(self): status = " + ".join(status_parts) if status_parts else "None" self.console.print(f" • {provider:15} ✅ {status}") + # Provider-Specific Settings + self.console.print() + self.console.print("[bold]🔬 Provider-Specific Settings[/bold]") + self.console.print("━" * 70) + try: + from proxy_app.settings_tool import PROVIDER_SETTINGS_MAP + except ImportError: + from .settings_tool import PROVIDER_SETTINGS_MAP + for provider in PROVIDER_SETTINGS_MAP.keys(): + display_name = provider.replace("_", " ").title() + modified = provider_settings.get(provider, 0) + if modified > 0: + self.console.print(f" • {display_name:20} [yellow]{modified} setting{'s' if modified > 1 else ''} modified[/yellow]") + else: + self.console.print(f" • {display_name:20} [dim]using defaults[/dim]") + # Actions self.console.print() self.console.print("━" * 70) diff --git a/src/proxy_app/settings_tool.py b/src/proxy_app/settings_tool.py index 67ee0cb1..71641f33 100644 --- a/src/proxy_app/settings_tool.py +++ b/src/proxy_app/settings_tool.py @@ -166,6 +166,184 @@ def remove_limit(self, provider: str): self.settings.remove(key) +# ============================================================================= +# PROVIDER-SPECIFIC SETTINGS DEFINITIONS +# ============================================================================= + +# Antigravity provider environment variables +ANTIGRAVITY_SETTINGS = { + "ANTIGRAVITY_SIGNATURE_CACHE_TTL": { + "type": "int", + "default": 3600, + "description": "Memory cache TTL for Gemini 3 thought signatures (seconds)", + }, + "ANTIGRAVITY_SIGNATURE_DISK_TTL": { + "type": "int", + "default": 86400, + "description": "Disk cache TTL for Gemini 3 thought signatures (seconds)", + }, + "ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES": { + "type": "bool", + "default": True, + "description": "Preserve thought signatures in client responses", + }, + "ANTIGRAVITY_ENABLE_SIGNATURE_CACHE": { + "type": "bool", + "default": True, + "description": "Enable signature caching for multi-turn conversations", + }, + "ANTIGRAVITY_ENABLE_DYNAMIC_MODELS": { + "type": "bool", + "default": False, + "description": "Enable dynamic model discovery from API", + }, + "ANTIGRAVITY_GEMINI3_TOOL_FIX": { + "type": "bool", + "default": True, + "description": "Enable Gemini 3 tool hallucination prevention", + }, + "ANTIGRAVITY_CLAUDE_TOOL_FIX": { + "type": "bool", + "default": True, + "description": "Enable Claude tool hallucination prevention", + }, + "ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION": { + "type": "bool", + "default": True, + "description": "Sanitize thinking blocks for Claude multi-turn conversations", + }, + "ANTIGRAVITY_GEMINI3_TOOL_PREFIX": { + "type": "str", + "default": "gemini3_", + "description": "Prefix added to tool names for Gemini 3 disambiguation", + }, + "ANTIGRAVITY_GEMINI3_DESCRIPTION_PROMPT": { + "type": "str", + "default": "\n\nSTRICT PARAMETERS: {params}.", + "description": "Template for strict parameter hints in tool descriptions", + }, + "ANTIGRAVITY_CLAUDE_DESCRIPTION_PROMPT": { + "type": "str", + "default": "\n\nSTRICT PARAMETERS: {params}.", + "description": "Template for Claude strict parameter hints in tool descriptions", + }, +} + +# Gemini CLI provider environment variables +GEMINI_CLI_SETTINGS = { + "GEMINI_CLI_SIGNATURE_CACHE_TTL": { + "type": "int", + "default": 3600, + "description": "Memory cache TTL for thought signatures (seconds)", + }, + "GEMINI_CLI_SIGNATURE_DISK_TTL": { + "type": "int", + "default": 86400, + "description": "Disk cache TTL for thought signatures (seconds)", + }, + "GEMINI_CLI_PRESERVE_THOUGHT_SIGNATURES": { + "type": "bool", + "default": True, + "description": "Preserve thought signatures in client responses", + }, + "GEMINI_CLI_ENABLE_SIGNATURE_CACHE": { + "type": "bool", + "default": True, + "description": "Enable signature caching for multi-turn conversations", + }, + "GEMINI_CLI_GEMINI3_TOOL_FIX": { + "type": "bool", + "default": True, + "description": "Enable Gemini 3 tool hallucination prevention", + }, + "GEMINI_CLI_GEMINI3_TOOL_PREFIX": { + "type": "str", + "default": "gemini3_", + "description": "Prefix added to tool names for Gemini 3 disambiguation", + }, + "GEMINI_CLI_GEMINI3_DESCRIPTION_PROMPT": { + "type": "str", + "default": "\n\nSTRICT PARAMETERS: {params}.", + "description": "Template for strict parameter hints in tool descriptions", + }, + "GEMINI_CLI_PROJECT_ID": { + "type": "str", + "default": "", + "description": "GCP Project ID for paid tier users (required for paid tiers)", + }, +} + +# Map provider names to their settings definitions +PROVIDER_SETTINGS_MAP = { + "antigravity": ANTIGRAVITY_SETTINGS, + "gemini_cli": GEMINI_CLI_SETTINGS, +} + + +class ProviderSettingsManager: + """Manages provider-specific configuration settings""" + + def __init__(self, settings: AdvancedSettings): + self.settings = settings + + def get_available_providers(self) -> List[str]: + """Get list of providers with specific settings available""" + return list(PROVIDER_SETTINGS_MAP.keys()) + + def get_provider_settings_definitions(self, provider: str) -> Dict[str, Dict[str, Any]]: + """Get settings definitions for a provider""" + return PROVIDER_SETTINGS_MAP.get(provider, {}) + + def get_current_value(self, key: str, definition: Dict[str, Any]) -> Any: + """Get current value of a setting from environment""" + env_value = os.getenv(key) + if env_value is None: + return definition.get("default") + + setting_type = definition.get("type", "str") + try: + if setting_type == "bool": + return env_value.lower() in ("true", "1", "yes") + elif setting_type == "int": + return int(env_value) + else: + return env_value + except (ValueError, AttributeError): + return definition.get("default") + + def get_all_current_values(self, provider: str) -> Dict[str, Any]: + """Get all current values for a provider""" + definitions = self.get_provider_settings_definitions(provider) + values = {} + for key, definition in definitions.items(): + values[key] = self.get_current_value(key, definition) + return values + + def set_value(self, key: str, value: Any, definition: Dict[str, Any]): + """Set a setting value, converting to string for .env storage""" + setting_type = definition.get("type", "str") + if setting_type == "bool": + str_value = "true" if value else "false" + else: + str_value = str(value) + self.settings.set(key, str_value) + + def reset_to_default(self, key: str): + """Remove a setting to reset it to default""" + self.settings.remove(key) + + def get_modified_settings(self, provider: str) -> Dict[str, Any]: + """Get settings that differ from defaults""" + definitions = self.get_provider_settings_definitions(provider) + modified = {} + for key, definition in definitions.items(): + current = self.get_current_value(key, definition) + default = definition.get("default") + if current != default: + modified[key] = current + return modified + + class SettingsTool: """Main settings tool TUI""" @@ -175,6 +353,7 @@ def __init__(self): self.provider_mgr = CustomProviderManager(self.settings) self.model_mgr = ModelDefinitionManager(self.settings) self.concurrency_mgr = ConcurrencyManager(self.settings) + self.provider_settings_mgr = ProviderSettingsManager(self.settings) self.running = True def get_available_providers(self) -> List[str]: @@ -223,8 +402,9 @@ def show_main_menu(self): self.console.print(" 1. 🌐 Custom Provider API Bases") self.console.print(" 2. 📦 Provider Model Definitions") self.console.print(" 3. ⚡ Concurrency Limits") - self.console.print(" 4. 💾 Save & Exit") - self.console.print(" 5. 🚫 Exit Without Saving") + self.console.print(" 4. 🔬 Provider-Specific Settings") + self.console.print(" 5. 💾 Save & Exit") + self.console.print(" 6. 🚫 Exit Without Saving") self.console.print() self.console.print("━" * 70) @@ -238,7 +418,7 @@ def show_main_menu(self): self.console.print("[dim]⚠️ Model filters not supported - edit .env for IGNORE_MODELS_* / WHITELIST_MODELS_*[/dim]") self.console.print() - choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5"], show_choices=False) + choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5", "6"], show_choices=False) if choice == "1": self.manage_custom_providers() @@ -247,8 +427,10 @@ def show_main_menu(self): elif choice == "3": self.manage_concurrency_limits() elif choice == "4": - self.save_and_exit() + self.manage_provider_settings() elif choice == "5": + self.save_and_exit() + elif choice == "6": self.exit_without_saving() def manage_custom_providers(self): @@ -631,6 +813,195 @@ def view_model_definitions(self, providers: List[str]): input("Press Enter to return...") + def manage_provider_settings(self): + """Manage provider-specific settings (Antigravity, Gemini CLI)""" + while True: + self.console.clear() + + available_providers = self.provider_settings_mgr.get_available_providers() + + self.console.print(Panel.fit( + "[bold cyan]🔬 Provider-Specific Settings[/bold cyan]", + border_style="cyan" + )) + + self.console.print() + self.console.print("[bold]📋 Available Providers with Custom Settings[/bold]") + self.console.print("━" * 70) + + for provider in available_providers: + modified = self.provider_settings_mgr.get_modified_settings(provider) + status = f"[yellow]{len(modified)} modified[/yellow]" if modified else "[dim]defaults[/dim]" + display_name = provider.replace("_", " ").title() + self.console.print(f" • {display_name:20} {status}") + + self.console.print() + self.console.print("━" * 70) + self.console.print() + self.console.print("[bold]⚙️ Select Provider to Configure[/bold]") + self.console.print() + + for idx, provider in enumerate(available_providers, 1): + display_name = provider.replace("_", " ").title() + self.console.print(f" {idx}. {display_name}") + self.console.print(f" {len(available_providers) + 1}. ↩️ Back to Settings Menu") + + self.console.print() + self.console.print("━" * 70) + self.console.print() + + choices = [str(i) for i in range(1, len(available_providers) + 2)] + choice = Prompt.ask("Select option", choices=choices, show_choices=False) + choice_idx = int(choice) + + if choice_idx == len(available_providers) + 1: + break + + provider = available_providers[choice_idx - 1] + self._manage_single_provider_settings(provider) + + def _manage_single_provider_settings(self, provider: str): + """Manage settings for a single provider""" + while True: + self.console.clear() + + display_name = provider.replace("_", " ").title() + definitions = self.provider_settings_mgr.get_provider_settings_definitions(provider) + current_values = self.provider_settings_mgr.get_all_current_values(provider) + + self.console.print(Panel.fit( + f"[bold cyan]🔬 {display_name} Settings[/bold cyan]", + border_style="cyan" + )) + + self.console.print() + self.console.print("[bold]📋 Current Settings[/bold]") + self.console.print("━" * 70) + + # Display all settings with current values + settings_list = list(definitions.keys()) + for idx, key in enumerate(settings_list, 1): + definition = definitions[key] + current = current_values.get(key) + default = definition.get("default") + setting_type = definition.get("type", "str") + description = definition.get("description", "") + + # Format value display + if setting_type == "bool": + value_display = "[green]✓ Enabled[/green]" if current else "[red]✗ Disabled[/red]" + elif setting_type == "int": + value_display = f"[cyan]{current}[/cyan]" + else: + value_display = f"[cyan]{current or '(not set)'}[/cyan]" if current else "[dim](not set)[/dim]" + + # Check if modified from default + modified = current != default + mod_marker = "[yellow]*[/yellow]" if modified else " " + + # Short key name for display (strip provider prefix) + short_key = key.replace(f"{provider.upper()}_", "") + + self.console.print(f" {mod_marker}{idx:2}. {short_key:35} {value_display}") + self.console.print(f" [dim]{description}[/dim]") + + self.console.print() + self.console.print("━" * 70) + self.console.print("[dim]* = modified from default[/dim]") + self.console.print() + self.console.print("[bold]⚙️ Actions[/bold]") + self.console.print() + self.console.print(" E. ✏️ Edit a Setting") + self.console.print(" R. 🔄 Reset Setting to Default") + self.console.print(" A. 🔄 Reset All to Defaults") + self.console.print(" B. ↩️ Back to Provider Selection") + + self.console.print() + self.console.print("━" * 70) + self.console.print() + + choice = Prompt.ask("Select action", choices=["e", "r", "a", "b", "E", "R", "A", "B"], show_choices=False).lower() + + if choice == "b": + break + elif choice == "e": + self._edit_provider_setting(provider, settings_list, definitions) + elif choice == "r": + self._reset_provider_setting(provider, settings_list, definitions) + elif choice == "a": + self._reset_all_provider_settings(provider, settings_list) + + def _edit_provider_setting(self, provider: str, settings_list: List[str], definitions: Dict[str, Dict[str, Any]]): + """Edit a single provider setting""" + self.console.print("\n[bold]Select setting number to edit:[/bold]") + + choices = [str(i) for i in range(1, len(settings_list) + 1)] + choice = IntPrompt.ask("Setting number", choices=choices) + key = settings_list[choice - 1] + definition = definitions[key] + + current = self.provider_settings_mgr.get_current_value(key, definition) + default = definition.get("default") + setting_type = definition.get("type", "str") + short_key = key.replace(f"{provider.upper()}_", "") + + self.console.print(f"\n[bold]Editing: {short_key}[/bold]") + self.console.print(f"Current value: [cyan]{current}[/cyan]") + self.console.print(f"Default value: [dim]{default}[/dim]") + self.console.print(f"Type: {setting_type}") + + if setting_type == "bool": + new_value = Confirm.ask("\nEnable this setting?", default=current) + self.provider_settings_mgr.set_value(key, new_value, definition) + status = "enabled" if new_value else "disabled" + self.console.print(f"\n[green]✅ {short_key} {status}![/green]") + elif setting_type == "int": + new_value = IntPrompt.ask("\nNew value", default=current) + self.provider_settings_mgr.set_value(key, new_value, definition) + self.console.print(f"\n[green]✅ {short_key} set to {new_value}![/green]") + else: + new_value = Prompt.ask("\nNew value", default=str(current) if current else "").strip() + if new_value: + self.provider_settings_mgr.set_value(key, new_value, definition) + self.console.print(f"\n[green]✅ {short_key} updated![/green]") + else: + self.console.print("\n[yellow]No changes made[/yellow]") + + input("\nPress Enter to continue...") + + def _reset_provider_setting(self, provider: str, settings_list: List[str], definitions: Dict[str, Dict[str, Any]]): + """Reset a single provider setting to default""" + self.console.print("\n[bold]Select setting number to reset:[/bold]") + + choices = [str(i) for i in range(1, len(settings_list) + 1)] + choice = IntPrompt.ask("Setting number", choices=choices) + key = settings_list[choice - 1] + definition = definitions[key] + + default = definition.get("default") + short_key = key.replace(f"{provider.upper()}_", "") + + if Confirm.ask(f"\nReset {short_key} to default ({default})?"): + self.provider_settings_mgr.reset_to_default(key) + self.console.print(f"\n[green]✅ {short_key} reset to default![/green]") + else: + self.console.print("\n[yellow]No changes made[/yellow]") + + input("\nPress Enter to continue...") + + def _reset_all_provider_settings(self, provider: str, settings_list: List[str]): + """Reset all provider settings to defaults""" + display_name = provider.replace("_", " ").title() + + if Confirm.ask(f"\n[bold red]Reset ALL {display_name} settings to defaults?[/bold red]"): + for key in settings_list: + self.provider_settings_mgr.reset_to_default(key) + self.console.print(f"\n[green]✅ All {display_name} settings reset to defaults![/green]") + else: + self.console.print("\n[yellow]No changes made[/yellow]") + + input("\nPress Enter to continue...") + def manage_concurrency_limits(self): """Manage concurrency limits""" while True: From 0dbcf50ca3fe98894c6b17c593028d9278ce248e Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 17:36:57 +0100 Subject: [PATCH 045/221] =?UTF-8?q?chore(build):=20=F0=9F=A7=B9=20remove?= =?UTF-8?q?=20Windows=20launcher=20script=20(not=20supposed=20to=20be=20th?= =?UTF-8?q?ere=20anyway)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- launcher.bat | 293 --------------------------------------------------- 1 file changed, 293 deletions(-) delete mode 100644 launcher.bat diff --git a/launcher.bat b/launcher.bat deleted file mode 100644 index ec241862..00000000 --- a/launcher.bat +++ /dev/null @@ -1,293 +0,0 @@ -@echo off -:: ================================================================================ -:: Universal Instructions for macOS / Linux Users -:: ================================================================================ -:: This launcher.bat file is for Windows only. -:: If you are on macOS or Linux, please use the following Python commands directly -:: in your terminal. -:: -:: First, ensure you have Python 3.10 or higher installed. -:: -:: To run the proxy server (basic command): -:: export PYTHONPATH=${PYTHONPATH}:$(pwd)/src -:: python src/proxy_app/main.py --host 0.0.0.0 --port 8000 -:: -:: Note: To enable request logging, add the --enable-request-logging flag to the command. -:: -:: To add new credentials: -:: export PYTHONPATH=${PYTHONPATH}:$(pwd)/src -:: python src/proxy_app/main.py --add-credential -:: -:: To build the executable (requires PyInstaller): -:: pip install -r requirements.txt -:: pip install pyinstaller -:: python src/proxy_app/build.py -:: ================================================================================ - -setlocal enabledelayedexpansion - -:: Default Settings -set "HOST=0.0.0.0" -set "PORT=8000" -set "LOGGING=false" -set "EXECUTION_MODE=" -set "EXE_NAME=proxy_app.exe" -set "SOURCE_PATH=src\proxy_app\main.py" - -:: --- Phase 1: Detection and Mode Selection --- -set "EXE_EXISTS=false" -set "SOURCE_EXISTS=false" - -if exist "%EXE_NAME%" ( - set "EXE_EXISTS=true" -) - -if exist "%SOURCE_PATH%" ( - set "SOURCE_EXISTS=true" -) - -if "%EXE_EXISTS%"=="true" ( - if "%SOURCE_EXISTS%"=="true" ( - call :SelectModeMenu - ) else ( - set "EXECUTION_MODE=exe" - ) -) else ( - if "%SOURCE_EXISTS%"=="true" ( - set "EXECUTION_MODE=source" - call :CheckPython - if errorlevel 1 goto :eof - ) else ( - call :NoTargetsFound - ) -) - -if "%EXECUTION_MODE%"=="" ( - goto :eof -) - -:: --- Phase 2: Main Menu --- -:MainMenu -cls -echo ================================================== -echo LLM API Key Proxy Launcher -echo ================================================== -echo. -echo Current Configuration: -echo ---------------------- -echo - Host IP: %HOST% -echo - Port: %PORT% -echo - Request Logging: %LOGGING% -echo - Execution Mode: %EXECUTION_MODE% -echo. -echo Main Menu: -echo ---------- -echo 1. Run Proxy -echo 2. Configure Proxy -echo 3. Add Credentials -if "%EXECUTION_MODE%"=="source" ( - echo 4. Build Executable - echo 5. Exit -) else ( - echo 4. Exit -) -echo. -set /p "CHOICE=Enter your choice: " - -if "%CHOICE%"=="1" goto :RunProxy -if "%CHOICE%"=="2" goto :ConfigMenu -if "%CHOICE%"=="3" goto :AddCredentials - -if "%EXECUTION_MODE%"=="source" ( - if "%CHOICE%"=="4" goto :BuildExecutable - if "%CHOICE%"=="5" goto :eof -) else ( - if "%CHOICE%"=="4" goto :eof -) - -echo Invalid choice. -pause -goto :MainMenu - -:: --- Phase 3: Configuration Sub-Menu --- -:ConfigMenu -cls -echo ================================================== -echo Configuration Menu -echo ================================================== -echo. -echo Current Configuration: -echo ---------------------- -echo - Host IP: %HOST% -echo - Port: %PORT% -echo - Request Logging: %LOGGING% -echo - Execution Mode: %EXECUTION_MODE% -echo. -echo Configuration Options: -echo ---------------------- -echo 1. Set Host IP -echo 2. Set Port -echo 3. Toggle Request Logging -echo 4. Back to Main Menu -echo. -set /p "CHOICE=Enter your choice: " - -if "%CHOICE%"=="1" ( - set /p "NEW_HOST=Enter new Host IP: " - if defined NEW_HOST ( - set "HOST=!NEW_HOST!" - ) - goto :ConfigMenu -) -if "%CHOICE%"=="2" ( - set "NEW_PORT=" - set /p "NEW_PORT=Enter new Port: " - if not defined NEW_PORT goto :ConfigMenu - set "IS_NUM=true" - for /f "delims=0123456789" %%i in ("!NEW_PORT!") do set "IS_NUM=false" - if "!IS_NUM!"=="false" ( - echo Invalid Port. Please enter numbers only. - pause - ) else ( - if !NEW_PORT! GTR 65535 ( - echo Invalid Port. Port cannot be greater than 65535. - pause - ) else ( - set "PORT=!NEW_PORT!" - ) - ) - goto :ConfigMenu -) -if "%CHOICE%"=="3" ( - if "%LOGGING%"=="true" ( - set "LOGGING=false" - ) else ( - set "LOGGING=true" - ) - goto :ConfigMenu -) -if "%CHOICE%"=="4" goto :MainMenu - -echo Invalid choice. -pause -goto :ConfigMenu - -:: --- Phase 4: Execution --- -:RunProxy -cls -set "ARGS=--host "%HOST%" --port %PORT%" -if "%LOGGING%"=="true" ( - set "ARGS=%ARGS% --enable-request-logging" -) -echo Starting Proxy... -echo Arguments: %ARGS% -echo. -if "%EXECUTION_MODE%"=="exe" ( - start "LLM API Proxy" "%EXE_NAME%" %ARGS% -) else ( - set "PYTHONPATH=%~dp0src;%PYTHONPATH%" - start "LLM API Proxy" python "%SOURCE_PATH%" %ARGS% -) -exit /b 0 - -:AddCredentials -cls -echo Launching Credential Tool... -echo. -if "%EXECUTION_MODE%"=="exe" ( - "%EXE_NAME%" --add-credential -) else ( - set "PYTHONPATH=%~dp0src;%PYTHONPATH%" - python "%SOURCE_PATH%" --add-credential -) -pause -goto :MainMenu - -:BuildExecutable -cls -echo ================================================== -echo Building Executable -echo ================================================== -echo. -echo The build process will start in a new window. -start "Build Process" cmd /c "pip install -r requirements.txt && pip install pyinstaller && python "src/proxy_app/build.py" && echo Build finished. && pause" -exit /b - -:: --- Helper Functions --- - -:SelectModeMenu -cls -echo ================================================== -echo Execution Mode Selection -echo ================================================== -echo. -echo Both executable and source code found. -echo Please choose which to use: -echo. -echo 1. Executable ("%EXE_NAME%") -echo 2. Source Code ("%SOURCE_PATH%") -echo. -set /p "CHOICE=Enter your choice: " - -if "%CHOICE%"=="1" ( - set "EXECUTION_MODE=exe" -) else if "%CHOICE%"=="2" ( - call :CheckPython - if errorlevel 1 goto :eof - set "EXECUTION_MODE=source" -) else ( - echo Invalid choice. - pause - goto :SelectModeMenu -) -goto :end_of_function - -:CheckPython -where python >nul 2>nul -if errorlevel 1 ( - echo Error: Python is not installed or not in PATH. - echo Please install Python and try again. - pause - exit /b 1 -) - -for /f "tokens=1,2" %%a in ('python -c "import sys; print(sys.version_info.major, sys.version_info.minor)"') do ( - set "PY_MAJOR=%%a" - set "PY_MINOR=%%b" -) - -if not "%PY_MAJOR%"=="3" ( - call :PythonVersionError - exit /b 1 -) -if %PY_MINOR% lss 10 ( - call :PythonVersionError - exit /b 1 -) - -exit /b 0 - -:PythonVersionError -echo Error: Python 3.10 or higher is required. -echo Found version: %PY_MAJOR%.%PY_MINOR% -echo Please upgrade your Python installation. -pause -goto :eof - -:NoTargetsFound -cls -echo ================================================== -echo Error -echo ================================================== -echo. -echo Could not find the executable ("%EXE_NAME%") -echo or the source code ("%SOURCE_PATH%"). -echo. -echo Please ensure the launcher is in the correct -echo directory or that the project has been built. -echo. -pause -goto :eof - -:end_of_function -endlocal From efbd008cd12c8b78abd50661612736f2f15b1dc6 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 17:39:11 +0100 Subject: [PATCH 046/221] =?UTF-8?q?docs(readme):=20=F0=9F=93=9A=20improve?= =?UTF-8?q?=20Antigravity=20provider=20feature=20documentation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restructured the Antigravity provider description in the README for better clarity and readability: - Converted the dense paragraph into a structured bullet list highlighting key features - Separated thought signature caching, tool hallucination prevention, and thinking block sanitization into distinct points - Replaced the informal troubleshooting note with a concise reference to dedicated documentation - Added direct link to Antigravity documentation section for Claude extended thinking sanitization details This change improves the discoverability of Antigravity's advanced features and provides a clearer path for users to understand Claude Sonnet 4.5 thinking mode limitations. --- README.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index b3ae33d3..51399bd2 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,11 @@ This project provides a powerful solution for developers building complex applic - **OpenAI-Compatible Proxy**: Offers a familiar API interface with additional endpoints for model and provider discovery. - **Advanced Model Filtering**: Supports both blacklists and whitelists to give you fine-grained control over which models are available through the proxy. -- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 2.5, Gemini 3, and Claude Sonnet 4.5 models with advanced features like thought signature caching and tool hallucination prevention. However - Sonnet 4.5 Thinking with native tool calls is very skittish, so if you have compaction or switch the model (or toggle thinking) mid task - it will error 400 on you, as claude needs it's previous thinking block. With compaction - it will be destroyed. There is a system to maybe catch all this, but i am hurting my head here trying to come up with a solution that makes sense. +- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 2.5, Gemini 3, and Claude Sonnet 4.5 models with advanced features: + - Thought signature caching for multi-turn conversations + - Tool hallucination prevention via parameter signature injection + - Automatic thinking block sanitization for Claude models + - Note: Claude Sonnet 4.5 thinking mode requires careful conversation state management (see [Antigravity documentation](DOCUMENTATION.md#antigravity-claude-extended-thinking-sanitization) for details) - **🆕 Credential Prioritization**: Automatic tier detection and priority-based credential selection ensures paid-tier credentials are used for premium models that require them. - **🆕 Weighted Random Rotation**: Configurable credential rotation strategy - choose between deterministic (perfect balance) or weighted random (unpredictable, harder to fingerprint) selection. - **🆕 Enhanced Gemini CLI**: Improved project discovery, paid vs free tier detection, and Gemini 3 support with thoughtSignature caching. From 6573de373fdef96352887607e00f01f5792778e2 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 18:02:21 +0100 Subject: [PATCH 047/221] =?UTF-8?q?chore(config):=20=F0=9F=A7=B9=20ignore?= =?UTF-8?q?=20environment=20files=20and=20increase=20default=20token=20lim?= =?UTF-8?q?it?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add `*.env` to `.gitignore` to prevent accidentally committing environment variables containing sensitive data - Increase `DEFAULT_MAX_OUTPUT_TOKENS` from 16384 to 32384 in Antigravity provider to allow for longer model outputs --- .gitignore | 1 + src/rotator_library/providers/antigravity_provider.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 92bac087..1a75e867 100644 --- a/.gitignore +++ b/.gitignore @@ -127,3 +127,4 @@ launcher_config.json cache/antigravity/thought_signatures.json logs/ cache/ +*.env diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 22573096..e5b6727f 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -64,7 +64,7 @@ ] # Default max output tokens (including thinking) - can be overridden per request -DEFAULT_MAX_OUTPUT_TOKENS = 16384 +DEFAULT_MAX_OUTPUT_TOKENS = 32384 # Model alias mappings (internal ↔ public) MODEL_ALIAS_MAP = { From bd8f6386c418b9a03a698d6be5848b3c7123b7aa Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 18:18:59 +0100 Subject: [PATCH 048/221] =?UTF-8?q?feat(credentials):=20=E2=9C=A8=20add=20?= =?UTF-8?q?support=20for=20environment-based=20credential=20loading=20and?= =?UTF-8?q?=20bulk=20export=20tools?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces comprehensive support for loading OAuth credentials from environment variables alongside file-based credentials, and adds powerful bulk export/combine functionality for all credential types. Main changes: - **Environment-based credentials**: Modified main.py to load all *.env files from the root directory, enabling credentials to be stored in environment variables with an "env://" virtual path scheme - **Safe metadata handling**: Added checks throughout to skip file I/O operations for env-based credentials (they use virtual paths and don't have metadata files) - **Optimized credential discovery**: Updated RotatingClient to accept pre-discovered credentials from main.py, avoiding redundant discovery calls - **Bulk export tools**: Added `export_all_provider_credentials()` to export all credentials for a specific provider to individual .env files - **Credential combining**: Added `combine_provider_credentials()` to merge all credentials for a provider into a single .env file, and `combine_all_credentials()` to create one master .env file with all providers - **Enhanced export menu**: Expanded the credential export submenu with 13 options covering individual exports, bulk exports per provider, and various combining strategies - **Provider support**: Added helper functions `_build_gemini_cli_env_lines()`, `_build_qwen_code_env_lines()`, `_build_iflow_env_lines()`, and `_build_antigravity_env_lines()` for consistent .env file generation These changes enable flexible credential management, allowing users to store credentials as files or environment variables, and providing powerful tools to export and combine credentials for deployment scenarios. --- src/proxy_app/main.py | 43 +++- src/rotator_library/client.py | 9 +- src/rotator_library/credential_tool.py | 342 ++++++++++++++++++++++++- 3 files changed, 375 insertions(+), 19 deletions(-) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index dfbc0418..263dc115 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -38,10 +38,19 @@ # If we get here, we're ACTUALLY running the proxy - NOW show startup messages and start timer _start_time = time.time() -# Load .env early so PROXY_API_KEY is available for display +# Load all .env files from root folder (main .env first, then any additional *.env files) from dotenv import load_dotenv +from glob import glob + +# Load main .env first load_dotenv() +# Load any additional .env files (e.g., antigravity_all_combined.env, gemini_cli_all_combined.env) +_root_dir = Path.cwd() +for _env_file in sorted(_root_dir.glob("*.env")): + if _env_file.name != ".env": # Skip main .env (already loaded) + load_dotenv(_env_file, override=False) # Don't override existing values + # Get proxy API key for display proxy_api_key = os.getenv("PROXY_API_KEY") if proxy_api_key: @@ -298,6 +307,11 @@ async def lifespan(app: FastAPI): if provider not in credentials_to_initialize: credentials_to_initialize[provider] = [] for path in paths: + # Skip env-based credentials (virtual paths) - they don't have metadata files + if path.startswith("env://"): + credentials_to_initialize[provider].append(path) + continue + try: with open(path, 'r') as f: data = json.load(f) @@ -399,19 +413,20 @@ async def process_credential(provider: str, path: str, provider_instance): final_oauth_credentials[provider] = [] final_oauth_credentials[provider].append(path) - # Update metadata - try: - with open(path, 'r+') as f: - data = json.load(f) - metadata = data.get("_proxy_metadata", {}) - metadata["email"] = email - metadata["last_check_timestamp"] = time.time() - data["_proxy_metadata"] = metadata - f.seek(0) - json.dump(data, f, indent=2) - f.truncate() - except Exception as e: - logging.error(f"Failed to update metadata for '{path}': {e}") + # Update metadata (skip for env-based credentials - they don't have files) + if not path.startswith("env://"): + try: + with open(path, 'r+') as f: + data = json.load(f) + metadata = data.get("_proxy_metadata", {}) + metadata["email"] = email + metadata["last_check_timestamp"] = time.time() + data["_proxy_metadata"] = metadata + f.seek(0) + json.dump(data, f, indent=2) + f.truncate() + except Exception as e: + logging.error(f"Failed to update metadata for '{path}': {e}") logging.info("OAuth credential processing complete.") oauth_credentials = final_oauth_credentials diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 7fa50806..e536aeb4 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -115,8 +115,13 @@ def __init__( ) self.api_keys = api_keys - self.credential_manager = CredentialManager(oauth_credentials) - self.oauth_credentials = self.credential_manager.discover_and_prepare() + # Use provided oauth_credentials directly if available (already discovered by main.py) + # Only call discover_and_prepare() if no credentials were passed + if oauth_credentials: + self.oauth_credentials = oauth_credentials + else: + self.credential_manager = CredentialManager(os.environ) + self.oauth_credentials = self.credential_manager.discover_and_prepare() self.background_refresher = BackgroundRefresher(self) self.oauth_providers = set(self.oauth_credentials.keys()) diff --git a/src/rotator_library/credential_tool.py b/src/rotator_library/credential_tool.py index 4b2f8a04..1949f134 100644 --- a/src/rotator_library/credential_tool.py +++ b/src/rotator_library/credential_tool.py @@ -713,6 +713,288 @@ async def export_antigravity_to_env(): console.print(Panel(f"An error occurred during export: {e}", style="bold red", title="Error")) +def _build_gemini_cli_env_lines(creds: dict, cred_number: int) -> list[str]: + """Build .env lines for a Gemini CLI credential.""" + email = creds.get("_proxy_metadata", {}).get("email", "unknown") + project_id = creds.get("_proxy_metadata", {}).get("project_id", "") + tier = creds.get("_proxy_metadata", {}).get("tier", "") + + extra_fields = {} + if project_id: + extra_fields["PROJECT_ID"] = project_id + if tier: + extra_fields["TIER"] = tier + + env_lines, _ = _build_env_export_content( + provider_prefix="GEMINI_CLI", + cred_number=cred_number, + creds=creds, + email=email, + extra_fields=extra_fields, + include_client_creds=True + ) + return env_lines + + +def _build_qwen_code_env_lines(creds: dict, cred_number: int) -> list[str]: + """Build .env lines for a Qwen Code credential.""" + email = creds.get("_proxy_metadata", {}).get("email", "unknown") + numbered_prefix = f"QWEN_CODE_{cred_number}" + + env_lines = [ + f"# QWEN_CODE Credential #{cred_number} for: {email}", + f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", + "", + f"{numbered_prefix}_ACCESS_TOKEN={creds.get('access_token', '')}", + f"{numbered_prefix}_REFRESH_TOKEN={creds.get('refresh_token', '')}", + f"{numbered_prefix}_EXPIRY_DATE={creds.get('expiry_date', 0)}", + f"{numbered_prefix}_RESOURCE_URL={creds.get('resource_url', 'https://portal.qwen.ai/v1')}", + f"{numbered_prefix}_EMAIL={email}", + ] + return env_lines + + +def _build_iflow_env_lines(creds: dict, cred_number: int) -> list[str]: + """Build .env lines for an iFlow credential.""" + email = creds.get("_proxy_metadata", {}).get("email", "unknown") + numbered_prefix = f"IFLOW_{cred_number}" + + env_lines = [ + f"# IFLOW Credential #{cred_number} for: {email}", + f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", + "", + f"{numbered_prefix}_ACCESS_TOKEN={creds.get('access_token', '')}", + f"{numbered_prefix}_REFRESH_TOKEN={creds.get('refresh_token', '')}", + f"{numbered_prefix}_API_KEY={creds.get('api_key', '')}", + f"{numbered_prefix}_EXPIRY_DATE={creds.get('expiry_date', '')}", + f"{numbered_prefix}_EMAIL={email}", + f"{numbered_prefix}_TOKEN_TYPE={creds.get('token_type', 'Bearer')}", + f"{numbered_prefix}_SCOPE={creds.get('scope', 'read write')}", + ] + return env_lines + + +def _build_antigravity_env_lines(creds: dict, cred_number: int) -> list[str]: + """Build .env lines for an Antigravity credential.""" + email = creds.get("_proxy_metadata", {}).get("email", "unknown") + + env_lines, _ = _build_env_export_content( + provider_prefix="ANTIGRAVITY", + cred_number=cred_number, + creds=creds, + email=email, + extra_fields=None, + include_client_creds=True + ) + return env_lines + + +async def export_all_provider_credentials(provider_name: str): + """ + Export all credentials for a specific provider to individual .env files. + """ + provider_config = { + "gemini_cli": ("GEMINI_CLI", _build_gemini_cli_env_lines), + "qwen_code": ("QWEN_CODE", _build_qwen_code_env_lines), + "iflow": ("IFLOW", _build_iflow_env_lines), + "antigravity": ("ANTIGRAVITY", _build_antigravity_env_lines), + } + + if provider_name not in provider_config: + console.print(f"[bold red]Unknown provider: {provider_name}[/bold red]") + return + + prefix, build_func = provider_config[provider_name] + display_name = prefix.replace("_", " ").title() + + console.print(Panel(f"[bold cyan]Export All {display_name} Credentials[/bold cyan]", expand=False)) + + # Find all credentials for this provider + cred_files = sorted(list(OAUTH_BASE_DIR.glob(f"{provider_name}_oauth_*.json"))) + + if not cred_files: + console.print(Panel(f"No {display_name} credentials found.", style="bold red", title="No Credentials")) + return + + exported_count = 0 + for cred_file in cred_files: + try: + with open(cred_file, 'r') as f: + creds = json.load(f) + + email = creds.get("_proxy_metadata", {}).get("email", "unknown") + cred_number = _get_credential_number_from_filename(cred_file.name) + + # Generate .env file name + safe_email = email.replace("@", "_at_").replace(".", "_") + env_filename = f"{provider_name}_{cred_number}_{safe_email}.env" + env_filepath = OAUTH_BASE_DIR / env_filename + + # Build and write .env content + env_lines = build_func(creds, cred_number) + with open(env_filepath, 'w') as f: + f.write('\n'.join(env_lines)) + + console.print(f" ✓ Exported [cyan]{cred_file.name}[/cyan] → [yellow]{env_filename}[/yellow]") + exported_count += 1 + + except Exception as e: + console.print(f" ✗ Failed to export {cred_file.name}: {e}") + + console.print(Panel( + f"Successfully exported {exported_count}/{len(cred_files)} {display_name} credentials to individual .env files.", + style="bold green", title="Export Complete" + )) + + +async def combine_provider_credentials(provider_name: str): + """ + Combine all credentials for a specific provider into a single .env file. + """ + provider_config = { + "gemini_cli": ("GEMINI_CLI", _build_gemini_cli_env_lines), + "qwen_code": ("QWEN_CODE", _build_qwen_code_env_lines), + "iflow": ("IFLOW", _build_iflow_env_lines), + "antigravity": ("ANTIGRAVITY", _build_antigravity_env_lines), + } + + if provider_name not in provider_config: + console.print(f"[bold red]Unknown provider: {provider_name}[/bold red]") + return + + prefix, build_func = provider_config[provider_name] + display_name = prefix.replace("_", " ").title() + + console.print(Panel(f"[bold cyan]Combine All {display_name} Credentials[/bold cyan]", expand=False)) + + # Find all credentials for this provider + cred_files = sorted(list(OAUTH_BASE_DIR.glob(f"{provider_name}_oauth_*.json"))) + + if not cred_files: + console.print(Panel(f"No {display_name} credentials found.", style="bold red", title="No Credentials")) + return + + combined_lines = [ + f"# Combined {display_name} Credentials", + f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", + f"# Total credentials: {len(cred_files)}", + "#", + "# Copy all lines below into your main .env file", + "", + ] + + combined_count = 0 + for cred_file in cred_files: + try: + with open(cred_file, 'r') as f: + creds = json.load(f) + + cred_number = _get_credential_number_from_filename(cred_file.name) + env_lines = build_func(creds, cred_number) + + combined_lines.extend(env_lines) + combined_lines.append("") # Blank line between credentials + combined_count += 1 + + except Exception as e: + console.print(f" ✗ Failed to process {cred_file.name}: {e}") + + # Write combined file + combined_filename = f"{provider_name}_all_combined.env" + combined_filepath = OAUTH_BASE_DIR / combined_filename + + with open(combined_filepath, 'w') as f: + f.write('\n'.join(combined_lines)) + + console.print(Panel( + Text.from_markup( + f"Successfully combined {combined_count} {display_name} credentials into:\n" + f"[bold yellow]{combined_filepath}[/bold yellow]\n\n" + f"[bold]To use:[/bold] Copy the contents into your main .env file." + ), + style="bold green", title="Combine Complete" + )) + + +async def combine_all_credentials(): + """ + Combine ALL credentials from ALL providers into a single .env file. + """ + console.print(Panel("[bold cyan]Combine All Provider Credentials[/bold cyan]", expand=False)) + + provider_config = { + "gemini_cli": ("GEMINI_CLI", _build_gemini_cli_env_lines), + "qwen_code": ("QWEN_CODE", _build_qwen_code_env_lines), + "iflow": ("IFLOW", _build_iflow_env_lines), + "antigravity": ("ANTIGRAVITY", _build_antigravity_env_lines), + } + + combined_lines = [ + "# Combined All Provider Credentials", + f"# Generated at: {time.strftime('%Y-%m-%d %H:%M:%S')}", + "#", + "# Copy all lines below into your main .env file", + "", + ] + + total_count = 0 + provider_counts = {} + + for provider_name, (prefix, build_func) in provider_config.items(): + cred_files = sorted(list(OAUTH_BASE_DIR.glob(f"{provider_name}_oauth_*.json"))) + + if not cred_files: + continue + + display_name = prefix.replace("_", " ").title() + combined_lines.append(f"# ===== {display_name} Credentials =====") + combined_lines.append("") + + provider_count = 0 + for cred_file in cred_files: + try: + with open(cred_file, 'r') as f: + creds = json.load(f) + + cred_number = _get_credential_number_from_filename(cred_file.name) + env_lines = build_func(creds, cred_number) + + combined_lines.extend(env_lines) + combined_lines.append("") + provider_count += 1 + total_count += 1 + + except Exception as e: + console.print(f" ✗ Failed to process {cred_file.name}: {e}") + + provider_counts[display_name] = provider_count + + if total_count == 0: + console.print(Panel("No credentials found to combine.", style="bold red", title="No Credentials")) + return + + # Write combined file + combined_filename = "all_providers_combined.env" + combined_filepath = OAUTH_BASE_DIR / combined_filename + + with open(combined_filepath, 'w') as f: + f.write('\n'.join(combined_lines)) + + # Build summary + summary_lines = [f" • {name}: {count} credential(s)" for name, count in provider_counts.items()] + summary = "\n".join(summary_lines) + + console.print(Panel( + Text.from_markup( + f"Successfully combined {total_count} credentials from {len(provider_counts)} providers:\n" + f"{summary}\n\n" + f"[bold]Output file:[/bold] [yellow]{combined_filepath}[/yellow]\n\n" + f"[bold]To use:[/bold] Copy the contents into your main .env file." + ), + style="bold green", title="Combine Complete" + )) + + async def export_credentials_submenu(): """ Submenu for credential export options. @@ -723,24 +1005,39 @@ async def export_credentials_submenu(): console.print(Panel( Text.from_markup( + "[bold]Individual Exports:[/bold]\n" "1. Export Gemini CLI credential\n" "2. Export Qwen Code credential\n" "3. Export iFlow credential\n" - "4. Export Antigravity credential" + "4. Export Antigravity credential\n" + "\n" + "[bold]Bulk Exports (per provider):[/bold]\n" + "5. Export ALL Gemini CLI credentials\n" + "6. Export ALL Qwen Code credentials\n" + "7. Export ALL iFlow credentials\n" + "8. Export ALL Antigravity credentials\n" + "\n" + "[bold]Combine Credentials:[/bold]\n" + "9. Combine all Gemini CLI into one file\n" + "10. Combine all Qwen Code into one file\n" + "11. Combine all iFlow into one file\n" + "12. Combine all Antigravity into one file\n" + "13. Combine ALL providers into one file" ), - title="Choose credential type to export", + title="Choose export option", style="bold blue" )) export_choice = Prompt.ask( Text.from_markup("[bold]Please select an option or type [red]'b'[/red] to go back[/bold]"), - choices=["1", "2", "3", "4", "b"], + choices=["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "b"], show_choices=False ) if export_choice.lower() == 'b': break + # Individual exports if export_choice == "1": await export_gemini_cli_to_env() console.print("\n[dim]Press Enter to return to export menu...[/dim]") @@ -757,6 +1054,45 @@ async def export_credentials_submenu(): await export_antigravity_to_env() console.print("\n[dim]Press Enter to return to export menu...[/dim]") input() + # Bulk exports (all credentials for a provider) + elif export_choice == "5": + await export_all_provider_credentials("gemini_cli") + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "6": + await export_all_provider_credentials("qwen_code") + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "7": + await export_all_provider_credentials("iflow") + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "8": + await export_all_provider_credentials("antigravity") + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + # Combine per provider + elif export_choice == "9": + await combine_provider_credentials("gemini_cli") + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "10": + await combine_provider_credentials("qwen_code") + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "11": + await combine_provider_credentials("iflow") + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + elif export_choice == "12": + await combine_provider_credentials("antigravity") + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() + # Combine all providers + elif export_choice == "13": + await combine_all_credentials() + console.print("\n[dim]Press Enter to return to export menu...[/dim]") + input() async def main(clear_on_start=True): From b6a47c979ef557e281055daaad46de04769e96f4 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 19:07:36 +0100 Subject: [PATCH 049/221] =?UTF-8?q?feat(api):=20=E2=9C=A8=20add=20model=20?= =?UTF-8?q?pricing=20and=20capabilities=20enrichment=20service?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a new model information service that fetches pricing and capability data from external catalogs (OpenRouter and Models.dev) to enrich the /v1/models endpoint and enable cost estimation. - Implements ModelRegistry class with async background data fetching to avoid blocking proxy startup - Adds fuzzy model ID matching with multi-source data aggregation - Expands /v1/models endpoint with optional enriched response containing pricing, token limits, and capability flags - Adds new endpoints: GET /v1/models/{model_id}, GET /v1/model-info/stats, POST /v1/cost-estimate - Supports per-token pricing for input, output, cache read, and cache write operations - Integrates with lifespan management for proper service initialization and cleanup - Includes comprehensive backward compatibility layer for gradual migration The service refreshes data every 6 hours (configurable via MODEL_INFO_REFRESH_INTERVAL) and runs asynchronously to maintain fast proxy initialization times. --- src/proxy_app/main.py | 214 ++++- src/rotator_library/__init__.py | 11 +- src/rotator_library/model_info_service.py | 946 ++++++++++++++++++++++ 3 files changed, 1165 insertions(+), 6 deletions(-) create mode 100644 src/rotator_library/model_info_service.py diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 263dc115..c2e318d0 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -100,6 +100,7 @@ from rotator_library import RotatingClient from rotator_library.credential_manager import CredentialManager from rotator_library.background_refresher import BackgroundRefresher + from rotator_library.model_info_service import init_model_info_service from proxy_app.request_logger import log_request_to_console from proxy_app.batch_manager import EmbeddingBatcher from proxy_app.detailed_logger import DetailedLogger @@ -123,15 +124,59 @@ class EmbeddingRequest(BaseModel): user: Optional[str] = None class ModelCard(BaseModel): + """Basic model card for minimal response.""" id: str object: str = "model" created: int = Field(default_factory=lambda: int(time.time())) owned_by: str = "Mirro-Proxy" +class ModelCapabilities(BaseModel): + """Model capability flags.""" + tool_choice: bool = False + function_calling: bool = False + reasoning: bool = False + vision: bool = False + system_messages: bool = True + prompt_caching: bool = False + assistant_prefill: bool = False + +class EnrichedModelCard(BaseModel): + """Extended model card with pricing and capabilities.""" + id: str + object: str = "model" + created: int = Field(default_factory=lambda: int(time.time())) + owned_by: str = "unknown" + # Pricing (optional - may not be available for all models) + input_cost_per_token: Optional[float] = None + output_cost_per_token: Optional[float] = None + cache_read_input_token_cost: Optional[float] = None + cache_creation_input_token_cost: Optional[float] = None + # Limits (optional) + max_input_tokens: Optional[int] = None + max_output_tokens: Optional[int] = None + context_window: Optional[int] = None + # Capabilities + mode: str = "chat" + supported_modalities: List[str] = Field(default_factory=lambda: ["text"]) + supported_output_modalities: List[str] = Field(default_factory=lambda: ["text"]) + capabilities: Optional[ModelCapabilities] = None + # Debug info (optional) + _sources: Optional[List[str]] = None + _match_type: Optional[str] = None + + class Config: + extra = "allow" # Allow extra fields from the service + class ModelList(BaseModel): + """List of models response.""" object: str = "list" data: List[ModelCard] +class EnrichedModelList(BaseModel): + """List of enriched models with pricing and capabilities.""" + object: str = "list" + data: List[EnrichedModelCard] + # Calculate total loading time _elapsed = time.time() - _start_time print(f"✓ Server ready in {_elapsed:.2f}s ({_plugin_count} providers discovered in {_provider_time:.2f}s)") @@ -470,6 +515,12 @@ async def process_credential(provider: str, path: str, provider_instance): else: app.state.embedding_batcher = None logging.info("RotatingClient initialized (EmbeddingBatcher disabled).") + + # Start model info service in background (fetches pricing/capabilities data) + # This runs asynchronously and doesn't block proxy startup + model_info_service = await init_model_info_service() + app.state.model_info_service = model_info_service + logging.info("Model info service started (fetching pricing data in background).") yield @@ -478,6 +529,10 @@ async def process_credential(provider: str, path: str, provider_instance): await app.state.embedding_batcher.stop() await client.close() + # Stop model info service + if hasattr(app.state, 'model_info_service') and app.state.model_info_service: + await app.state.model_info_service.stop() + if app.state.embedding_batcher: logging.info("RotatingClient and EmbeddingBatcher closed.") else: @@ -847,17 +902,73 @@ async def embeddings( def read_root(): return {"Status": "API Key Proxy is running"} -@app.get("/v1/models", response_model=ModelList) +@app.get("/v1/models") async def list_models( + request: Request, client: RotatingClient = Depends(get_rotating_client), - _=Depends(verify_api_key) + _=Depends(verify_api_key), + enriched: bool = True, ): """ Returns a list of available models in the OpenAI-compatible format. + + Query Parameters: + enriched: If True (default), returns detailed model info with pricing and capabilities. + If False, returns minimal OpenAI-compatible response. """ model_ids = await client.get_all_available_models(grouped=False) - model_cards = [ModelCard(id=model_id) for model_id in model_ids] - return ModelList(data=model_cards) + + if enriched and hasattr(request.app.state, 'model_info_service'): + model_info_service = request.app.state.model_info_service + if model_info_service.is_ready(): + # Return enriched model data + enriched_data = model_info_service.enrich_model_list(model_ids) + return {"object": "list", "data": enriched_data} + + # Fallback to basic model cards + model_cards = [{"id": model_id, "object": "model", "created": int(time.time()), "owned_by": "Mirro-Proxy"} for model_id in model_ids] + return {"object": "list", "data": model_cards} + + +@app.get("/v1/models/{model_id:path}") +async def get_model( + model_id: str, + request: Request, + _=Depends(verify_api_key), +): + """ + Returns detailed information about a specific model. + + Path Parameters: + model_id: The model ID (e.g., "anthropic/claude-3-opus", "openrouter/openai/gpt-4") + """ + if hasattr(request.app.state, 'model_info_service'): + model_info_service = request.app.state.model_info_service + if model_info_service.is_ready(): + info = model_info_service.get_model_info(model_id) + if info: + return info.to_dict() + + # Return basic info if service not ready or model not found + return { + "id": model_id, + "object": "model", + "created": int(time.time()), + "owned_by": model_id.split("/")[0] if "/" in model_id else "unknown", + } + + +@app.get("/v1/model-info/stats") +async def model_info_stats( + request: Request, + _=Depends(verify_api_key), +): + """ + Returns statistics about the model info service (for monitoring/debugging). + """ + if hasattr(request.app.state, 'model_info_service'): + return request.app.state.model_info_service.get_stats() + return {"error": "Model info service not initialized"} @app.get("/v1/providers") @@ -891,6 +1002,101 @@ async def token_count( logging.error(f"Token count failed: {e}") raise HTTPException(status_code=500, detail=str(e)) + +@app.post("/v1/cost-estimate") +async def cost_estimate( + request: Request, + _=Depends(verify_api_key) +): + """ + Estimates the cost for a request based on token counts and model pricing. + + Request body: + { + "model": "anthropic/claude-3-opus", + "prompt_tokens": 1000, + "completion_tokens": 500, + "cache_read_tokens": 0, # optional + "cache_creation_tokens": 0 # optional + } + + Returns: + { + "model": "anthropic/claude-3-opus", + "cost": 0.0375, + "currency": "USD", + "pricing": { + "input_cost_per_token": 0.000015, + "output_cost_per_token": 0.000075 + }, + "source": "model_info_service" # or "litellm_fallback" + } + """ + try: + data = await request.json() + model = data.get("model") + prompt_tokens = data.get("prompt_tokens", 0) + completion_tokens = data.get("completion_tokens", 0) + cache_read_tokens = data.get("cache_read_tokens", 0) + cache_creation_tokens = data.get("cache_creation_tokens", 0) + + if not model: + raise HTTPException(status_code=400, detail="'model' is required.") + + result = { + "model": model, + "cost": None, + "currency": "USD", + "pricing": {}, + "source": None + } + + # Try model info service first + if hasattr(request.app.state, 'model_info_service'): + model_info_service = request.app.state.model_info_service + if model_info_service.is_ready(): + cost = model_info_service.calculate_cost( + model, prompt_tokens, completion_tokens, + cache_read_tokens, cache_creation_tokens + ) + if cost is not None: + cost_info = model_info_service.get_cost_info(model) + result["cost"] = cost + result["pricing"] = cost_info or {} + result["source"] = "model_info_service" + return result + + # Fallback to litellm + try: + import litellm + # Create a mock response for cost calculation + model_info = litellm.get_model_info(model) + input_cost = model_info.get("input_cost_per_token", 0) + output_cost = model_info.get("output_cost_per_token", 0) + + if input_cost or output_cost: + cost = (prompt_tokens * input_cost) + (completion_tokens * output_cost) + result["cost"] = cost + result["pricing"] = { + "input_cost_per_token": input_cost, + "output_cost_per_token": output_cost + } + result["source"] = "litellm_fallback" + return result + except Exception: + pass + + result["source"] = "unknown" + result["error"] = "Pricing data not available for this model" + return result + + except HTTPException: + raise + except Exception as e: + logging.error(f"Cost estimate failed: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + if __name__ == "__main__": # Define ENV_FILE for onboarding checks ENV_FILE = Path.cwd() / ".env" diff --git a/src/rotator_library/__init__.py b/src/rotator_library/__init__.py index 9a678123..f3ff0ec7 100644 --- a/src/rotator_library/__init__.py +++ b/src/rotator_library/__init__.py @@ -7,12 +7,19 @@ if TYPE_CHECKING: from .providers import PROVIDER_PLUGINS from .providers.provider_interface import ProviderInterface + from .model_info_service import ModelInfoService, ModelInfo -__all__ = ["RotatingClient", "PROVIDER_PLUGINS"] +__all__ = ["RotatingClient", "PROVIDER_PLUGINS", "ModelInfoService", "ModelInfo"] def __getattr__(name): - """Lazy-load PROVIDER_PLUGINS to speed up module import.""" + """Lazy-load PROVIDER_PLUGINS and ModelInfoService to speed up module import.""" if name == "PROVIDER_PLUGINS": from .providers import PROVIDER_PLUGINS return PROVIDER_PLUGINS + if name == "ModelInfoService": + from .model_info_service import ModelInfoService + return ModelInfoService + if name == "ModelInfo": + from .model_info_service import ModelInfo + return ModelInfo raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/rotator_library/model_info_service.py b/src/rotator_library/model_info_service.py new file mode 100644 index 00000000..0c577bce --- /dev/null +++ b/src/rotator_library/model_info_service.py @@ -0,0 +1,946 @@ +""" +Unified Model Registry + +Provides aggregated model metadata from external catalogs (OpenRouter, Models.dev) +for pricing calculations and the /v1/models endpoint. + +Data retrieval happens asynchronously post-startup to keep initialization fast. +""" + +import asyncio +import json +import logging +import os +import time +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple +from urllib.request import Request, urlopen +from urllib.error import URLError + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Data Structures +# ============================================================================ + +@dataclass +class ModelPricing: + """Token-level pricing information.""" + prompt: Optional[float] = None + completion: Optional[float] = None + cached_input: Optional[float] = None + cache_write: Optional[float] = None + + +@dataclass +class ModelLimits: + """Context and output token limits.""" + context_window: Optional[int] = None + max_output: Optional[int] = None + + +@dataclass +class ModelCapabilities: + """Feature flags for model capabilities.""" + tools: bool = False + functions: bool = False + reasoning: bool = False + vision: bool = False + system_prompt: bool = True + caching: bool = False + prefill: bool = False + + +@dataclass +class ModelMetadata: + """Complete model information record.""" + + model_id: str + display_name: str = "" + provider: str = "" + category: str = "chat" # chat, embedding, image, audio + + pricing: ModelPricing = field(default_factory=ModelPricing) + limits: ModelLimits = field(default_factory=ModelLimits) + capabilities: ModelCapabilities = field(default_factory=ModelCapabilities) + + input_types: List[str] = field(default_factory=lambda: ["text"]) + output_types: List[str] = field(default_factory=lambda: ["text"]) + + timestamp: int = field(default_factory=lambda: int(time.time())) + origin: str = "" + match_quality: str = "unknown" + + def as_api_response(self) -> Dict[str, Any]: + """Format for OpenAI-compatible /v1/models response.""" + response = { + "id": self.model_id, + "object": "model", + "created": self.timestamp, + "owned_by": self.provider or "proxy", + } + + # Pricing fields + if self.pricing.prompt is not None: + response["input_cost_per_token"] = self.pricing.prompt + if self.pricing.completion is not None: + response["output_cost_per_token"] = self.pricing.completion + if self.pricing.cached_input is not None: + response["cache_read_input_token_cost"] = self.pricing.cached_input + if self.pricing.cache_write is not None: + response["cache_creation_input_token_cost"] = self.pricing.cache_write + + # Limits + if self.limits.context_window: + response["max_input_tokens"] = self.limits.context_window + response["context_window"] = self.limits.context_window + if self.limits.max_output: + response["max_output_tokens"] = self.limits.max_output + + # Category and modalities + response["mode"] = self.category + response["supported_modalities"] = self.input_types + response["supported_output_modalities"] = self.output_types + + # Capability flags + response["capabilities"] = { + "tool_choice": self.capabilities.tools, + "function_calling": self.capabilities.functions, + "reasoning": self.capabilities.reasoning, + "vision": self.capabilities.vision, + "system_messages": self.capabilities.system_prompt, + "prompt_caching": self.capabilities.caching, + "assistant_prefill": self.capabilities.prefill, + } + + # Debug metadata + if self.origin: + response["_sources"] = [self.origin] + response["_match_type"] = self.match_quality + + return response + + def as_minimal(self) -> Dict[str, Any]: + """Minimal OpenAI format.""" + return { + "id": self.model_id, + "object": "model", + "created": self.timestamp, + "owned_by": self.provider or "proxy", + } + + def to_dict(self) -> Dict[str, Any]: + """Alias for as_api_response() - backward compatibility.""" + return self.as_api_response() + + def to_openai_format(self) -> Dict[str, Any]: + """Alias for as_minimal() - backward compatibility.""" + return self.as_minimal() + + # Backward-compatible property aliases + @property + def id(self) -> str: + return self.model_id + + @property + def name(self) -> str: + return self.display_name + + @property + def input_cost_per_token(self) -> Optional[float]: + return self.pricing.prompt + + @property + def output_cost_per_token(self) -> Optional[float]: + return self.pricing.completion + + @property + def cache_read_input_token_cost(self) -> Optional[float]: + return self.pricing.cached_input + + @property + def cache_creation_input_token_cost(self) -> Optional[float]: + return self.pricing.cache_write + + @property + def max_input_tokens(self) -> Optional[int]: + return self.limits.context_window + + @property + def max_output_tokens(self) -> Optional[int]: + return self.limits.max_output + + @property + def mode(self) -> str: + return self.category + + @property + def supported_modalities(self) -> List[str]: + return self.input_types + + @property + def supported_output_modalities(self) -> List[str]: + return self.output_types + + @property + def supports_tool_choice(self) -> bool: + return self.capabilities.tools + + @property + def supports_function_calling(self) -> bool: + return self.capabilities.functions + + @property + def supports_reasoning(self) -> bool: + return self.capabilities.reasoning + + @property + def supports_vision(self) -> bool: + return self.capabilities.vision + + @property + def supports_system_messages(self) -> bool: + return self.capabilities.system_prompt + + @property + def supports_prompt_caching(self) -> bool: + return self.capabilities.caching + + @property + def supports_assistant_prefill(self) -> bool: + return self.capabilities.prefill + + @property + def litellm_provider(self) -> str: + return self.provider + + @property + def created(self) -> int: + return self.timestamp + + @property + def _sources(self) -> List[str]: + return [self.origin] if self.origin else [] + + @property + def _match_type(self) -> str: + return self.match_quality + + +# ============================================================================ +# Data Source Adapters +# ============================================================================ + +class DataSourceAdapter: + """Base interface for external data sources.""" + + source_name: str = "unknown" + endpoint: str = "" + + def fetch(self) -> Dict[str, Dict]: + """Retrieve and normalize data. Returns {model_id: raw_data}.""" + raise NotImplementedError + + def _http_get(self, url: str, timeout: int = 30) -> Any: + """Execute HTTP GET with standard headers.""" + req = Request(url, headers={"User-Agent": "ModelRegistry/1.0"}) + with urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read().decode("utf-8")) + + +class OpenRouterAdapter(DataSourceAdapter): + """Fetches model data from OpenRouter's public API.""" + + source_name = "openrouter" + endpoint = "https://openrouter.ai/api/v1/models" + + def fetch(self) -> Dict[str, Dict]: + try: + raw = self._http_get(self.endpoint) + entries = raw.get("data", []) + + catalog = {} + for entry in entries: + mid = entry.get("id") + if not mid: + continue + + full_id = f"openrouter/{mid}" + catalog[full_id] = self._normalize(entry) + + return catalog + except (URLError, json.JSONDecodeError, TimeoutError) as err: + raise ConnectionError(f"OpenRouter unavailable: {err}") from err + + def _normalize(self, raw: Dict) -> Dict: + """Transform OpenRouter schema to internal format.""" + prices = raw.get("pricing", {}) + arch = raw.get("architecture", {}) + top = raw.get("top_provider", {}) + params = raw.get("supported_parameters", []) + + tokenizer = arch.get("tokenizer", "") + category = "embedding" if "embedding" in tokenizer.lower() else "chat" + + return { + "name": raw.get("name", ""), + "prompt_cost": float(prices.get("prompt", 0)), + "completion_cost": float(prices.get("completion", 0)), + "cache_read_cost": float(prices.get("input_cache_read", 0)) or None, + "context": top.get("context_length", 0), + "max_out": top.get("max_completion_tokens", 0), + "category": category, + "inputs": arch.get("input_modalities", ["text"]), + "outputs": arch.get("output_modalities", ["text"]), + "has_tools": "tool_choice" in params or "tools" in params, + "has_functions": "tools" in params or "function_calling" in params, + "has_reasoning": "reasoning" in params, + "has_vision": "image" in arch.get("input_modalities", []), + "provider": "openrouter", + "source": "openrouter", + } + + +class ModelsDevAdapter(DataSourceAdapter): + """Fetches model data from Models.dev catalog.""" + + source_name = "modelsdev" + endpoint = "https://models.dev/api.json" + + def __init__(self, skip_providers: Optional[List[str]] = None): + self.skip_providers = skip_providers or [] + + def fetch(self) -> Dict[str, Dict]: + try: + raw = self._http_get(self.endpoint) + + catalog = {} + for provider_key, provider_block in raw.items(): + if not isinstance(provider_block, dict): + continue + if provider_key in self.skip_providers: + continue + + models_block = provider_block.get("models", {}) + if not isinstance(models_block, dict): + continue + + for model_key, model_data in models_block.items(): + if not isinstance(model_data, dict): + continue + + full_id = f"{provider_key}/{model_key}" + catalog[full_id] = self._normalize(model_data, provider_key) + + return catalog + except (URLError, json.JSONDecodeError, TimeoutError) as err: + raise ConnectionError(f"Models.dev unavailable: {err}") from err + + def _normalize(self, raw: Dict, provider_key: str) -> Dict: + """Transform Models.dev schema to internal format.""" + costs = raw.get("cost", {}) + mods = raw.get("modalities", {}) + lims = raw.get("limit", {}) + + outputs = mods.get("output", ["text"]) + if "image" in outputs: + category = "image" + elif "audio" in outputs: + category = "audio" + else: + category = "chat" + + # Models.dev uses per-million pricing, convert to per-token + divisor = 1_000_000 + + cache_read = costs.get("cache_read") + cache_write = costs.get("cache_write") + + return { + "name": raw.get("name", ""), + "prompt_cost": float(costs.get("input", 0)) / divisor, + "completion_cost": float(costs.get("output", 0)) / divisor, + "cache_read_cost": float(cache_read) / divisor if cache_read else None, + "cache_write_cost": float(cache_write) / divisor if cache_write else None, + "context": lims.get("context", 0), + "max_out": lims.get("output", 0), + "category": category, + "inputs": mods.get("input", ["text"]), + "outputs": outputs, + "has_tools": raw.get("tool_call", False), + "has_functions": raw.get("tool_call", False), + "has_reasoning": raw.get("reasoning", False), + "has_vision": "image" in mods.get("input", []), + "provider": provider_key, + "source": "modelsdev", + } + + +# ============================================================================ +# Lookup Index +# ============================================================================ + +class ModelIndex: + """Fast lookup structure for model ID resolution.""" + + def __init__(self): + self._by_full_id: Dict[str, str] = {} # normalized_id -> canonical_id + self._by_suffix: Dict[str, List[str]] = {} # short_name -> [canonical_ids] + + def clear(self): + """Reset the index.""" + self._by_full_id.clear() + self._by_suffix.clear() + + def entry_count(self) -> int: + """Return total number of suffix index entries.""" + return sum(len(v) for v in self._by_suffix.values()) + + def add(self, canonical_id: str): + """Index a canonical model ID for various lookup patterns.""" + self._by_full_id[canonical_id] = canonical_id + + segments = canonical_id.split("/") + if len(segments) >= 2: + # Index by everything after first segment + partial = "/".join(segments[1:]) + self._by_suffix.setdefault(partial, []).append(canonical_id) + + # Index by final segment only + if len(segments) >= 3: + tail = segments[-1] + self._by_suffix.setdefault(tail, []).append(canonical_id) + + def resolve(self, query: str) -> List[str]: + """Find all canonical IDs matching a query.""" + # Direct match + if query in self._by_full_id: + return [self._by_full_id[query]] + + # Try with openrouter prefix + prefixed = f"openrouter/{query}" + if prefixed in self._by_full_id: + return [self._by_full_id[prefixed]] + + # Extract search terms from query + search_keys = [] + parts = query.split("/") + if len(parts) >= 2: + search_keys.append("/".join(parts[1:])) + search_keys.append(parts[-1]) + else: + search_keys.append(query) + # Find matches + matches = [] + seen = set() + for key in search_keys: + for cid in self._by_suffix.get(key, []): + if cid not in seen: + seen.add(cid) + matches.append(cid) + + return matches + + +# ============================================================================ +# Data Merger +# ============================================================================ + +class DataMerger: + """Combines data from multiple sources into unified ModelMetadata.""" + + @staticmethod + def single(model_id: str, data: Dict, origin: str, quality: str) -> ModelMetadata: + """Create ModelMetadata from a single source record.""" + return ModelMetadata( + model_id=model_id, + display_name=data.get("name", model_id), + provider=data.get("provider", ""), + category=data.get("category", "chat"), + pricing=ModelPricing( + prompt=data.get("prompt_cost"), + completion=data.get("completion_cost"), + cached_input=data.get("cache_read_cost"), + cache_write=data.get("cache_write_cost"), + ), + limits=ModelLimits( + context_window=data.get("context") or None, + max_output=data.get("max_out") or None, + ), + capabilities=ModelCapabilities( + tools=data.get("has_tools", False), + functions=data.get("has_functions", False), + reasoning=data.get("has_reasoning", False), + vision=data.get("has_vision", False), + ), + input_types=data.get("inputs", ["text"]), + output_types=data.get("outputs", ["text"]), + origin=origin, + match_quality=quality, + ) + + @staticmethod + def combine(model_id: str, records: List[Tuple[Dict, str]], quality: str) -> ModelMetadata: + """Merge multiple source records into one ModelMetadata.""" + if len(records) == 1: + data, origin = records[0] + return DataMerger.single(model_id, data, origin, quality) + + # Aggregate pricing - use average + prompt_costs = [r[0]["prompt_cost"] for r in records if r[0].get("prompt_cost")] + comp_costs = [r[0]["completion_cost"] for r in records if r[0].get("completion_cost")] + cache_costs = [r[0]["cache_read_cost"] for r in records if r[0].get("cache_read_cost")] + + # Aggregate limits - use most common value + contexts = [r[0]["context"] for r in records if r[0].get("context")] + max_outs = [r[0]["max_out"] for r in records if r[0].get("max_out")] + + # Capabilities - OR logic (any source supporting = supported) + has_tools = any(r[0].get("has_tools") for r in records) + has_funcs = any(r[0].get("has_functions") for r in records) + has_reason = any(r[0].get("has_reasoning") for r in records) + has_vis = any(r[0].get("has_vision") for r in records) + + # Modalities - union + all_inputs = set() + all_outputs = set() + for r in records: + all_inputs.update(r[0].get("inputs", ["text"])) + all_outputs.update(r[0].get("outputs", ["text"])) + + # Category - majority vote + categories = [r[0].get("category", "chat") for r in records] + category = max(set(categories), key=categories.count) + + # Name - first non-empty + name = model_id + for r in records: + if r[0].get("name"): + name = r[0]["name"] + break + + origins = [r[1] for r in records] + + return ModelMetadata( + model_id=model_id, + display_name=name, + provider=records[0][0].get("provider", ""), + category=category, + pricing=ModelPricing( + prompt=sum(prompt_costs) / len(prompt_costs) if prompt_costs else None, + completion=sum(comp_costs) / len(comp_costs) if comp_costs else None, + cached_input=sum(cache_costs) / len(cache_costs) if cache_costs else None, + ), + limits=ModelLimits( + context_window=DataMerger._mode(contexts), + max_output=DataMerger._mode(max_outs), + ), + capabilities=ModelCapabilities( + tools=has_tools, + functions=has_funcs, + reasoning=has_reason, + vision=has_vis, + ), + input_types=list(all_inputs) or ["text"], + output_types=list(all_outputs) or ["text"], + origin=",".join(origins), + match_quality=quality, + ) + + @staticmethod + def _mode(values: List[int]) -> Optional[int]: + """Return most frequent value.""" + if not values: + return None + return max(set(values), key=values.count) + + +# ============================================================================ +# Main Registry Service +# ============================================================================ + +class ModelRegistry: + """ + Central registry for model metadata from external catalogs. + + Manages background data refresh and provides lookup/pricing APIs. + """ + + REFRESH_INTERVAL_DEFAULT = 6 * 60 * 60 # 6 hours + + def __init__( + self, + refresh_seconds: Optional[int] = None, + skip_modelsdev_providers: Optional[List[str]] = None, + ): + interval_env = os.getenv("MODEL_INFO_REFRESH_INTERVAL") + self._refresh_interval = refresh_seconds or ( + int(interval_env) if interval_env else self.REFRESH_INTERVAL_DEFAULT + ) + + # Configure adapters + self._adapters: List[DataSourceAdapter] = [ + OpenRouterAdapter(), + ModelsDevAdapter(skip_providers=skip_modelsdev_providers or []), + ] + + # Raw data stores + self._openrouter_store: Dict[str, Dict] = {} + self._modelsdev_store: Dict[str, Dict] = {} + + # Lookup infrastructure + self._index = ModelIndex() + self._result_cache: Dict[str, ModelMetadata] = {} + + # Async coordination + self._ready = asyncio.Event() + self._mutex = asyncio.Lock() + self._worker: Optional[asyncio.Task] = None + self._last_refresh: float = 0 + + # ---------- Lifecycle ---------- + + async def start(self): + """Begin background refresh worker.""" + if self._worker is None: + self._worker = asyncio.create_task(self._refresh_worker()) + logger.info( + "ModelRegistry started (refresh every %ds)", + self._refresh_interval + ) + + async def stop(self): + """Halt background worker.""" + if self._worker: + self._worker.cancel() + try: + await self._worker + except asyncio.CancelledError: + pass + self._worker = None + logger.info("ModelRegistry stopped") + + async def await_ready(self, timeout_secs: float = 30.0) -> bool: + """Block until initial data load completes.""" + try: + await asyncio.wait_for(self._ready.wait(), timeout=timeout_secs) + return True + except asyncio.TimeoutError: + logger.warning("ModelRegistry ready timeout after %.1fs", timeout_secs) + return False + + @property + def is_ready(self) -> bool: + return self._ready.is_set() + + # ---------- Background Worker ---------- + + async def _refresh_worker(self): + """Periodic refresh loop.""" + await self._load_all_sources() + self._ready.set() + + while True: + try: + await asyncio.sleep(self._refresh_interval) + logger.info("Scheduled registry refresh...") + await self._load_all_sources() + logger.info("Registry refresh complete") + except asyncio.CancelledError: + break + except Exception as ex: + logger.error("Registry refresh error: %s", ex) + + async def _load_all_sources(self): + """Fetch from all adapters concurrently.""" + loop = asyncio.get_event_loop() + + tasks = [ + loop.run_in_executor(None, adapter.fetch) + for adapter in self._adapters + ] + + results = await asyncio.gather(*tasks, return_exceptions=True) + + async with self._mutex: + for adapter, result in zip(self._adapters, results): + if isinstance(result, Exception): + logger.error("%s fetch failed: %s", adapter.source_name, result) + continue + + if adapter.source_name == "openrouter": + self._openrouter_store = result + logger.info("OpenRouter: %d models loaded", len(result)) + elif adapter.source_name == "modelsdev": + self._modelsdev_store = result + logger.info("Models.dev: %d models loaded", len(result)) + + self._rebuild_index() + self._last_refresh = time.time() + + def _rebuild_index(self): + """Reconstruct lookup index from current stores.""" + self._index.clear() + self._result_cache.clear() + + for model_id in self._openrouter_store: + self._index.add(model_id) + + for model_id in self._modelsdev_store: + self._index.add(model_id) + + # ---------- Query API ---------- + + def lookup(self, model_id: str) -> Optional[ModelMetadata]: + """ + Retrieve model metadata by ID. + + Matching strategy: + 1. Exact match against known IDs + 2. Fuzzy match by model name suffix + 3. Aggregate if multiple sources match + """ + if model_id in self._result_cache: + return self._result_cache[model_id] + + metadata = self._resolve_model(model_id) + if metadata: + self._result_cache[model_id] = metadata + return metadata + + def _resolve_model(self, model_id: str) -> Optional[ModelMetadata]: + """Build ModelMetadata by matching source data.""" + records: List[Tuple[Dict, str]] = [] + quality = "none" + + # Check exact matches first + or_key = f"openrouter/{model_id}" if not model_id.startswith("openrouter/") else model_id + if or_key in self._openrouter_store: + records.append((self._openrouter_store[or_key], f"openrouter:exact:{or_key}")) + quality = "exact" + + if model_id in self._modelsdev_store: + records.append((self._modelsdev_store[model_id], f"modelsdev:exact:{model_id}")) + quality = "exact" + + # Fall back to index search + if not records: + candidates = self._index.resolve(model_id) + for cid in candidates: + if cid in self._openrouter_store: + records.append((self._openrouter_store[cid], f"openrouter:fuzzy:{cid}")) + elif cid in self._modelsdev_store: + records.append((self._modelsdev_store[cid], f"modelsdev:fuzzy:{cid}")) + + if records: + quality = "fuzzy" + + if not records: + return None + + return DataMerger.combine(model_id, records, quality) + + def get_pricing(self, model_id: str) -> Optional[Dict[str, float]]: + """Extract just pricing info for cost calculations.""" + meta = self.lookup(model_id) + if not meta: + return None + + result = {} + if meta.pricing.prompt is not None: + result["input_cost_per_token"] = meta.pricing.prompt + if meta.pricing.completion is not None: + result["output_cost_per_token"] = meta.pricing.completion + if meta.pricing.cached_input is not None: + result["cache_read_input_token_cost"] = meta.pricing.cached_input + if meta.pricing.cache_write is not None: + result["cache_creation_input_token_cost"] = meta.pricing.cache_write + + return result if result else None + + def compute_cost( + self, + model_id: str, + input_tokens: int, + output_tokens: int, + cache_hit_tokens: int = 0, + cache_miss_tokens: int = 0, + ) -> Optional[float]: + """ + Calculate total request cost. + + Returns None if pricing unavailable. + """ + pricing = self.get_pricing(model_id) + if not pricing: + return None + + in_rate = pricing.get("input_cost_per_token") + out_rate = pricing.get("output_cost_per_token") + + if in_rate is None or out_rate is None: + return None + + total = (input_tokens * in_rate) + (output_tokens * out_rate) + + cache_read_rate = pricing.get("cache_read_input_token_cost") + if cache_read_rate and cache_hit_tokens: + total += cache_hit_tokens * cache_read_rate + + cache_write_rate = pricing.get("cache_creation_input_token_cost") + if cache_write_rate and cache_miss_tokens: + total += cache_miss_tokens * cache_write_rate + + return total + + def enrich_models(self, model_ids: List[str]) -> List[Dict[str, Any]]: + """ + Attach metadata to a list of model IDs. + + Used by /v1/models endpoint. + """ + enriched = [] + for mid in model_ids: + meta = self.lookup(mid) + if meta: + enriched.append(meta.as_api_response()) + else: + # Fallback minimal entry + enriched.append({ + "id": mid, + "object": "model", + "created": int(time.time()), + "owned_by": mid.split("/")[0] if "/" in mid else "unknown", + }) + return enriched + + def all_raw_models(self) -> Dict[str, Dict]: + """Return all raw source data (for debugging).""" + combined = {} + combined.update(self._openrouter_store) + combined.update(self._modelsdev_store) + return combined + + def diagnostics(self) -> Dict[str, Any]: + """Return service health/stats.""" + return { + "ready": self._ready.is_set(), + "last_refresh": self._last_refresh, + "openrouter_count": len(self._openrouter_store), + "modelsdev_count": len(self._modelsdev_store), + "cached_lookups": len(self._result_cache), + "index_entries": self._index.entry_count(), + "refresh_interval": self._refresh_interval, + } + + # ---------- Backward Compatibility Methods ---------- + + def get_model_info(self, model_id: str) -> Optional[ModelMetadata]: + """Alias for lookup() - backward compatibility.""" + return self.lookup(model_id) + + def get_cost_info(self, model_id: str) -> Optional[Dict[str, float]]: + """Alias for get_pricing() - backward compatibility.""" + return self.get_pricing(model_id) + + def calculate_cost( + self, + model_id: str, + prompt_tokens: int, + completion_tokens: int, + cache_read_tokens: int = 0, + cache_creation_tokens: int = 0, + ) -> Optional[float]: + """Alias for compute_cost() - backward compatibility.""" + return self.compute_cost( + model_id, prompt_tokens, completion_tokens, + cache_read_tokens, cache_creation_tokens + ) + + def enrich_model_list(self, model_ids: List[str]) -> List[Dict[str, Any]]: + """Alias for enrich_models() - backward compatibility.""" + return self.enrich_models(model_ids) + + def get_all_source_models(self) -> Dict[str, Dict]: + """Alias for all_raw_models() - backward compatibility.""" + return self.all_raw_models() + + def get_stats(self) -> Dict[str, Any]: + """Alias for diagnostics() - backward compatibility.""" + return self.diagnostics() + + def wait_for_ready(self, timeout: float = 30.0): + """Sync wrapper for await_ready() - for compatibility.""" + return self.await_ready(timeout) + + +# ============================================================================ +# Backward Compatibility Layer +# ============================================================================ + +# Alias for backward compatibility +ModelInfo = ModelMetadata +ModelInfoService = ModelRegistry + +# Global singleton +_registry_instance: Optional[ModelRegistry] = None + + +def get_model_info_service() -> ModelRegistry: + """Get or create the global registry instance.""" + global _registry_instance + if _registry_instance is None: + _registry_instance = ModelRegistry() + return _registry_instance + + +async def init_model_info_service() -> ModelRegistry: + """Initialize and start the global registry.""" + registry = get_model_info_service() + await registry.start() + return registry + + +# Compatibility shim - map old method names to new +class _CompatibilityWrapper: + """Provides old API method names for gradual migration.""" + + def __init__(self, registry: ModelRegistry): + self._reg = registry + + def get_model_info(self, model_id: str) -> Optional[ModelMetadata]: + return self._reg.lookup(model_id) + + def get_cost_info(self, model_id: str) -> Optional[Dict[str, float]]: + return self._reg.get_pricing(model_id) + + def calculate_cost( + self, model_id: str, prompt_tokens: int, completion_tokens: int, + cache_read_tokens: int = 0, cache_creation_tokens: int = 0 + ) -> Optional[float]: + return self._reg.compute_cost( + model_id, prompt_tokens, completion_tokens, + cache_read_tokens, cache_creation_tokens + ) + + def enrich_model_list(self, model_ids: List[str]) -> List[Dict[str, Any]]: + return self._reg.enrich_models(model_ids) + + def get_all_source_models(self) -> Dict[str, Dict]: + return self._reg.all_raw_models() + + def get_stats(self) -> Dict[str, Any]: + return self._reg.diagnostics() + + async def start(self): + await self._reg.start() + + async def stop(self): + await self._reg.stop() + + async def wait_for_ready(self, timeout: float = 30.0) -> bool: + return await self._reg.await_ready(timeout) + + def is_ready(self) -> bool: + return self._reg.is_ready From 6ed16779cfbd72afb8108d7788b28b8da3945ebc Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 20:42:40 +0100 Subject: [PATCH 050/221] =?UTF-8?q?fix(provider):=20=F0=9F=90=9B=20improve?= =?UTF-8?q?=20Gemini=203=20tool=20schema=20handling=20and=20parameter=20va?= =?UTF-8?q?lidation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enhanced the Gemini 3 system instruction with more comprehensive and explicit rules for tool parameter usage to prevent hallucination and schema mismatches. - Rewrote DEFAULT_GEMINI3_SYSTEM_INSTRUCTION with clearer structure and XML-style tags for better model parsing - Added explicit warnings about pre-trained tool knowledge being invalid in custom environments - Included detailed guidance on array parameters, nested objects, and common failure patterns - Enhanced _clean_claude_schema to handle 'anyOf' and 'oneOf' by selecting the first option (Claude doesn't support these constructs) - Added temperature parameter handling with explicit Gemini 3 default of 1.0 for better tool use performance These changes address recurring issues where the model would use parameter names from its training data instead of reading the actual JSON schema definitions, particularly for tools with array-of-objects parameters. --- .../providers/antigravity_provider.py | 82 +++++++++++++++++-- 1 file changed, 73 insertions(+), 9 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index e5b6727f..7fbf7c2f 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -105,17 +105,58 @@ CLAUDE_THINKING_CACHE_FILE = CACHE_DIR / "claude_thinking.json" # Gemini 3 tool fix system instruction (prevents hallucination) -DEFAULT_GEMINI3_SYSTEM_INSTRUCTION = """CRITICAL TOOL USAGE INSTRUCTIONS: -You are operating in a custom environment where tool definitions differ from your training data. -You MUST follow these rules strictly: +DEFAULT_GEMINI3_SYSTEM_INSTRUCTION = """ +You are operating in a CUSTOM ENVIRONMENT where tool definitions COMPLETELY DIFFER from your training data. +VIOLATION OF THESE RULES WILL CAUSE IMMEDIATE SYSTEM FAILURE. -1. DO NOT use your internal training data to guess tool parameters -2. ONLY use the exact parameter structure defined in the tool schema -3. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') -4. Array parameters have specific item types - check the schema's 'items' field for the exact structure -5. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions +## ABSOLUTE RULES - NO EXCEPTIONS + +1. **SCHEMA IS LAW**: The JSON schema in each tool definition is the ONLY source of truth. + - Your pre-trained knowledge about tools like 'read_file', 'apply_diff', 'write_to_file', 'bash', etc. is INVALID here. + - Every tool has been REDEFINED with different parameters than what you learned during training. + +2. **PARAMETER NAMES ARE EXACT**: Use ONLY the parameter names from the schema. + - WRONG: 'suggested_answers', 'file_path', 'files_to_read', 'command_to_run' + - RIGHT: Check the 'properties' field in the schema for the exact names + - The schema's 'required' array tells you which parameters are mandatory + +3. **ARRAY PARAMETERS**: When a parameter has "type": "array", check the 'items' field: + - If items.type is "object", you MUST provide an array of objects with the EXACT properties listed + - If items.type is "string", you MUST provide an array of strings + - NEVER provide a single object when an array is expected + - NEVER provide an array when a single value is expected -If you are unsure about a tool's parameters, YOU MUST read the schema definition carefully. Your training data about common tool names like 'read_file' or 'apply_diff' does NOT apply here. +4. **NESTED OBJECTS**: When items.type is "object": + - Check items.properties for the EXACT field names required + - Check items.required for which nested fields are mandatory + - Include ALL required nested fields in EVERY array element + +5. **STRICT PARAMETERS HINT**: Tool descriptions contain "STRICT PARAMETERS: ..." which lists: + - Parameter name, type, and whether REQUIRED + - For arrays of objects: the nested structure in brackets like [field: type REQUIRED, ...] + - USE THIS as your quick reference, but the JSON schema is authoritative + +6. **BEFORE EVERY TOOL CALL**: + a. Read the tool's 'parametersJsonSchema' or 'parameters' field completely + b. Identify ALL required parameters + c. Verify your parameter names match EXACTLY (case-sensitive) + d. For arrays, verify you're providing the correct item structure + e. Do NOT add parameters that don't exist in the schema + +## COMMON FAILURE PATTERNS TO AVOID + +- Using 'path' when schema says 'filePath' (or vice versa) +- Using 'content' when schema says 'text' (or vice versa) +- Providing {"file": "..."} when schema wants [{"path": "...", "line_ranges": [...]}] +- Omitting required nested fields in array items +- Adding 'additionalProperties' that the schema doesn't define +- Guessing parameter names from similar tools you know from training + +## REMEMBER +Your training data about function calling is OUTDATED for this environment. +The tool names may look familiar, but the schemas are DIFFERENT. +When in doubt, RE-READ THE SCHEMA before making the call. + """ # Claude tool fix system instruction (prevents hallucination) @@ -270,6 +311,7 @@ def _clean_claude_schema(schema: Any) -> Any: Recursively clean JSON Schema for Antigravity/Google's Proto-based API. - Removes unsupported fields ($schema, additionalProperties, etc.) - Converts 'const' to 'enum' with single value (supported equivalent) + - Converts 'anyOf'/'oneOf' to the first option (Claude doesn't support these) """ if not isinstance(schema, dict): return schema @@ -278,6 +320,20 @@ def _clean_claude_schema(schema: Any) -> Any: incompatible = { '$schema', 'additionalProperties', 'minItems', 'maxItems', 'pattern', } + + # Handle 'anyOf' by taking the first option (Claude doesn't support anyOf) + if 'anyOf' in schema and isinstance(schema['anyOf'], list) and schema['anyOf']: + first_option = _clean_claude_schema(schema['anyOf'][0]) + if isinstance(first_option, dict): + return first_option + + # Handle 'oneOf' similarly + if 'oneOf' in schema and isinstance(schema['oneOf'], list) and schema['oneOf']: + first_option = _clean_claude_schema(schema['oneOf'][0]) + if isinstance(first_option, dict): + return first_option + + cleaned = {} # Handle 'const' by converting to 'enum' with single value @@ -1923,6 +1979,7 @@ async def acompletion( tool_choice = kwargs.get("tool_choice") reasoning_effort = kwargs.get("reasoning_effort") top_p = kwargs.get("top_p") + temperature = kwargs.get("temperature") max_tokens = kwargs.get("max_tokens") custom_budget = kwargs.get("custom_reasoning_budget", False) enable_logging = kwargs.pop("enable_request_logging", False) @@ -1972,6 +2029,13 @@ async def acompletion( if top_p is not None: gen_config["topP"] = top_p + # Handle temperature - Gemini 3 defaults to 1 if not explicitly set + if temperature is not None: + gen_config["temperature"] = temperature + elif self._is_gemini_3(model): + # Gemini 3 performs better with temperature=1 for tool use + gen_config["temperature"] = 1.0 + thinking_config = self._get_thinking_config(reasoning_effort, model, custom_budget) if thinking_config: gen_config.setdefault("thinkingConfig", {}).update(thinking_config) From f50cbff67c7d812e3ef4ae4e107922ac8cd8d20a Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 20:45:16 +0100 Subject: [PATCH 051/221] =?UTF-8?q?feat(provider):=20=E2=9C=A8=20add=20str?= =?UTF-8?q?ict=20JSON=20schema=20enforcement=20for=20Gemini=203=20tool=20c?= =?UTF-8?q?alls?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces a comprehensive strict schema enforcement mechanism to prevent Gemini 3 models from hallucinating parameters not defined in tool schemas. - Add new `_enforce_strict_schema()` method that recursively adds `additionalProperties: false` to all object schemas in tool definitions - Introduce `ANTIGRAVITY_GEMINI3_STRICT_SCHEMA` environment variable (defaults to True) to control strict schema enforcement - Enhance `_format_type_hint()` to provide more detailed parameter type information including enum values, const values, nested objects, and recursive type hints - Update Gemini 3 description prompt with explicit warning against using parameters from training data - Integrate strict schema enforcement into the Gemini 3 tool transformation pipeline - Add strict schema configuration to debug logging output The strict schema enforcement tells the model it cannot add properties not explicitly defined in the schema, significantly reducing parameter hallucination issues. The enhanced type hints provide clearer guidance to the model about expected parameter formats. --- .../providers/antigravity_provider.py | 81 +++++++++++++++++-- 1 file changed, 74 insertions(+), 7 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 7fbf7c2f..2aa47aa5 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -473,8 +473,9 @@ def __init__(self): self._gemini3_tool_prefix = os.getenv("ANTIGRAVITY_GEMINI3_TOOL_PREFIX", "gemini3_") self._gemini3_description_prompt = os.getenv( "ANTIGRAVITY_GEMINI3_DESCRIPTION_PROMPT", - "\n\nSTRICT PARAMETERS: {params}." + "\n\n⚠️ STRICT PARAMETERS (use EXACTLY as shown): {params}. Do NOT use parameters from your training data - use ONLY these parameter names." ) + self._gemini3_enforce_strict_schema = _env_bool("ANTIGRAVITY_GEMINI3_STRICT_SCHEMA", True) self._gemini3_system_instruction = os.getenv( "ANTIGRAVITY_GEMINI3_SYSTEM_INSTRUCTION", DEFAULT_GEMINI3_SYSTEM_INSTRUCTION @@ -498,8 +499,8 @@ def _log_config(self) -> None: lib_logger.debug( f"Antigravity config: signatures_in_client={self._preserve_signatures_in_client}, " f"cache={self._enable_signature_cache}, dynamic_models={self._enable_dynamic_models}, " - f"gemini3_fix={self._enable_gemini3_tool_fix}, claude_fix={self._enable_claude_tool_fix}, " - f"thinking_sanitization={self._enable_thinking_sanitization}" + f"gemini3_fix={self._enable_gemini3_tool_fix}, gemini3_strict_schema={self._gemini3_enforce_strict_schema}, " + f"claude_fix={self._enable_claude_tool_fix}, thinking_sanitization={self._enable_thinking_sanitization}" ) # ========================================================================= @@ -1341,6 +1342,43 @@ def _apply_gemini3_namespace( return modified + def _enforce_strict_schema(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Enforce strict JSON schema for Gemini 3 to prevent hallucinated parameters. + + Adds 'additionalProperties: false' recursively to all object schemas, + which tells the model it CANNOT add properties not in the schema. + """ + if not tools: + return tools + + def enforce_strict(schema: Any) -> Any: + if not isinstance(schema, dict): + return schema + + result = {} + for key, value in schema.items(): + if isinstance(value, dict): + result[key] = enforce_strict(value) + elif isinstance(value, list): + result[key] = [enforce_strict(item) if isinstance(item, dict) else item for item in value] + else: + result[key] = value + + # Add additionalProperties: false to object schemas + if result.get("type") == "object" and "properties" in result: + result["additionalProperties"] = False + + return result + + modified = copy.deepcopy(tools) + for tool in modified: + for func_decl in tool.get("functionDeclarations", []): + if "parametersJsonSchema" in func_decl: + func_decl["parametersJsonSchema"] = enforce_strict(func_decl["parametersJsonSchema"]) + + return modified + def _inject_signature_into_descriptions( self, tools: List[Dict[str, Any]], @@ -1385,10 +1423,21 @@ def _inject_signature_into_descriptions( return modified - def _format_type_hint(self, prop_data: Dict[str, Any]) -> str: - """Format a type hint for a property schema.""" + def _format_type_hint(self, prop_data: Dict[str, Any], depth: int = 0) -> str: + """Format a detailed type hint for a property schema.""" type_hint = prop_data.get("type", "unknown") + # Handle enum values - show allowed options + if "enum" in prop_data: + enum_vals = prop_data["enum"] + if len(enum_vals) <= 5: + return f"string ENUM[{', '.join(repr(v) for v in enum_vals)}]" + return f"string ENUM[{len(enum_vals)} options]" + + # Handle const values + if "const" in prop_data: + return f"string CONST={repr(prop_data['const'])}" + if type_hint == "array": items = prop_data.get("items", {}) if isinstance(items, dict): @@ -1400,7 +1449,11 @@ def _format_type_hint(self, prop_data: Dict[str, Any]) -> str: nested_list = [] for n, d in nested_props.items(): if isinstance(d, dict): - t = d.get("type", "unknown") + # Recursively format nested types (limit depth) + if depth < 1: + t = self._format_type_hint(d, depth + 1) + else: + t = d.get("type", "unknown") req = " REQUIRED" if n in nested_req else "" nested_list.append(f"{n}: {t}{req}") return f"ARRAY_OF_OBJECTS[{', '.join(nested_list)}]" @@ -1408,6 +1461,18 @@ def _format_type_hint(self, prop_data: Dict[str, Any]) -> str: return f"ARRAY_OF_{item_type.upper()}" return "ARRAY" + if type_hint == "object": + nested_props = prop_data.get("properties", {}) + nested_req = prop_data.get("required", []) + if nested_props and depth < 1: + nested_list = [] + for n, d in nested_props.items(): + if isinstance(d, dict): + t = d.get("type", "unknown") + req = " REQUIRED" if n in nested_req else "" + nested_list.append(f"{n}: {t}{req}") + return f"object{{{', '.join(nested_list)}}}" + return type_hint def _strip_gemini3_prefix(self, name: str) -> str: @@ -2050,8 +2115,10 @@ async def acompletion( # Apply tool transformations if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: - # Gemini 3: namespace prefix + parameter signatures + # Gemini 3: namespace prefix + strict schema + parameter signatures gemini_payload["tools"] = self._apply_gemini3_namespace(gemini_payload["tools"]) + if self._gemini3_enforce_strict_schema: + gemini_payload["tools"] = self._enforce_strict_schema(gemini_payload["tools"]) gemini_payload["tools"] = self._inject_signature_into_descriptions( gemini_payload["tools"], self._gemini3_description_prompt From 5a03c26f0f25e491fc9caf971c3e9e9294576e5a Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 22:14:12 +0100 Subject: [PATCH 052/221] =?UTF-8?q?fix(provider):=20=F0=9F=90=9B=20expand?= =?UTF-8?q?=20JSON=20schema=20validation=20keyword=20filtering=20and=20imp?= =?UTF-8?q?rove=20Gemini=203=20tool=20call=20reliability?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit addresses issues with schema compatibility and tool call hallucination across providers: - **Antigravity Provider**: Expands the list of incompatible JSON Schema keywords that must be filtered out for Claude via Antigravity, including validation constraints (minLength, maxLength, minimum, maximum), metadata fields (title, examples, deprecated), and JSON Schema draft 2020-12 specific keywords that cause API rejections. - **Gemini CLI Provider**: Significantly enhances the Gemini 3 tool calling system to prevent parameter hallucination: - Rewrites system instruction with more explicit warnings about custom tool schemas differing from training data - Adds common failure pattern examples to help the model avoid typical mistakes - Implements strict schema enforcement via `additionalProperties: false` to prevent invalid parameter injection - Improves parameter signature hints in tool descriptions with recursive type formatting, enum/const support, and nested object display - Adds new environment variable `GEMINI_CLI_GEMINI3_STRICT_SCHEMA` to control strict schema enforcement - Enhances type hint formatting to show array-of-objects structures more clearly These changes work together to reduce tool call errors by making schema constraints more explicit to both the Antigravity API and the Gemini 3 model. --- .../providers/antigravity_provider.py | 6 + .../providers/gemini_cli_provider.py | 134 +++++++++++++++--- 2 files changed, 122 insertions(+), 18 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 2aa47aa5..3f06b197 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -317,8 +317,14 @@ def _clean_claude_schema(schema: Any) -> Any: return schema # Fields not supported by Antigravity/Google's Proto-based API + # Note: Claude via Antigravity rejects JSON Schema draft 2020-12 validation keywords incompatible = { '$schema', 'additionalProperties', 'minItems', 'maxItems', 'pattern', + 'minLength', 'maxLength', 'minimum', 'maximum', 'default', + 'exclusiveMinimum', 'exclusiveMaximum', 'multipleOf', 'format', + 'minProperties', 'maxProperties', 'uniqueItems', 'contentEncoding', + 'contentMediaType', 'contentSchema', 'deprecated', 'readOnly', 'writeOnly', + 'examples', '$id', '$ref', '$defs', 'definitions', 'title', } # Handle 'anyOf' by taking the first option (Claude doesn't support anyOf) diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index bd85283e..601edf8e 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -92,19 +92,58 @@ def log_final_response(self, response_data: Dict[str, Any]): GEMINI3_SIGNATURE_CACHE_FILE = CACHE_DIR / "gemini3_signatures.json" # Gemini 3 tool fix system instruction (prevents hallucination) -DEFAULT_GEMINI3_SYSTEM_INSTRUCTION = """CRITICAL TOOL USAGE INSTRUCTIONS: -You are operating in a custom environment where tool definitions differ from your training data. -You MUST follow these rules strictly: - -1. DO NOT use your internal training data to guess tool parameters -2. ONLY use the exact parameter structure defined in the tool schema -3. If a tool takes a 'files' parameter, it is ALWAYS an array of objects with specific properties, NEVER a simple array of strings -4. If a tool edits code, it takes structured JSON objects with specific fields, NEVER raw diff strings or plain text -5. Parameter names in schemas are EXACT - do not substitute with similar names from your training (e.g., use 'follow_up' not 'suggested_answers') -6. Array parameters have specific item types - check the schema's 'items' field for the exact structure -7. When you see "STRICT PARAMETERS" in a tool description, those type definitions override any assumptions - -If you are unsure about a tool's parameters, YOU MUST read the schema definition carefully. Your training data about common tool names like 'read_file' or 'apply_diff' does NOT apply here. +DEFAULT_GEMINI3_SYSTEM_INSTRUCTION = """ +You are operating in a CUSTOM ENVIRONMENT where tool definitions COMPLETELY DIFFER from your training data. +VIOLATION OF THESE RULES WILL CAUSE IMMEDIATE SYSTEM FAILURE. + +## ABSOLUTE RULES - NO EXCEPTIONS + +1. **SCHEMA IS LAW**: The JSON schema in each tool definition is the ONLY source of truth. + - Your pre-trained knowledge about tools like 'read_file', 'apply_diff', 'write_to_file', 'bash', etc. is INVALID here. + - Every tool has been REDEFINED with different parameters than what you learned during training. + +2. **PARAMETER NAMES ARE EXACT**: Use ONLY the parameter names from the schema. + - WRONG: 'suggested_answers', 'file_path', 'files_to_read', 'command_to_run' + - RIGHT: Check the 'properties' field in the schema for the exact names + - The schema's 'required' array tells you which parameters are mandatory + +3. **ARRAY PARAMETERS**: When a parameter has "type": "array", check the 'items' field: + - If items.type is "object", you MUST provide an array of objects with the EXACT properties listed + - If items.type is "string", you MUST provide an array of strings + - NEVER provide a single object when an array is expected + - NEVER provide an array when a single value is expected + +4. **NESTED OBJECTS**: When items.type is "object": + - Check items.properties for the EXACT field names required + - Check items.required for which nested fields are mandatory + - Include ALL required nested fields in EVERY array element + +5. **STRICT PARAMETERS HINT**: Tool descriptions contain "STRICT PARAMETERS: ..." which lists: + - Parameter name, type, and whether REQUIRED + - For arrays of objects: the nested structure in brackets like [field: type REQUIRED, ...] + - USE THIS as your quick reference, but the JSON schema is authoritative + +6. **BEFORE EVERY TOOL CALL**: + a. Read the tool's 'parametersJsonSchema' or 'parameters' field completely + b. Identify ALL required parameters + c. Verify your parameter names match EXACTLY (case-sensitive) + d. For arrays, verify you're providing the correct item structure + e. Do NOT add parameters that don't exist in the schema + +## COMMON FAILURE PATTERNS TO AVOID + +- Using 'path' when schema says 'filePath' (or vice versa) +- Using 'content' when schema says 'text' (or vice versa) +- Providing {"file": "..."} when schema wants [{"path": "...", "line_ranges": [...]}] +- Omitting required nested fields in array items +- Adding 'additionalProperties' that the schema doesn't define +- Guessing parameter names from similar tools you know from training + +## REMEMBER +Your training data about function calling is OUTDATED for this environment. +The tool names may look familiar, but the schemas are DIFFERENT. +When in doubt, RE-READ THE SCHEMA before making the call. + """ # Gemini finish reason mapping @@ -150,12 +189,13 @@ def __init__(self): self._preserve_signatures_in_client = _env_bool("GEMINI_CLI_PRESERVE_THOUGHT_SIGNATURES", True) self._enable_signature_cache = _env_bool("GEMINI_CLI_ENABLE_SIGNATURE_CACHE", True) self._enable_gemini3_tool_fix = _env_bool("GEMINI_CLI_GEMINI3_TOOL_FIX", True) + self._gemini3_enforce_strict_schema = _env_bool("GEMINI_CLI_GEMINI3_STRICT_SCHEMA", True) # Gemini 3 tool fix configuration self._gemini3_tool_prefix = os.getenv("GEMINI_CLI_GEMINI3_TOOL_PREFIX", "gemini3_") self._gemini3_description_prompt = os.getenv( "GEMINI_CLI_GEMINI3_DESCRIPTION_PROMPT", - "\n\nSTRICT PARAMETERS: {params}." + "\n\n⚠️ STRICT PARAMETERS (use EXACTLY as shown): {params}. Do NOT use parameters from your training data - use ONLY these parameter names." ) self._gemini3_system_instruction = os.getenv( "GEMINI_CLI_GEMINI3_SYSTEM_INSTRUCTION", @@ -164,7 +204,8 @@ def __init__(self): lib_logger.debug( f"GeminiCli config: signatures_in_client={self._preserve_signatures_in_client}, " - f"cache={self._enable_signature_cache}, gemini3_fix={self._enable_gemini3_tool_fix}" + f"cache={self._enable_signature_cache}, gemini3_fix={self._enable_gemini3_tool_fix}, " + f"gemini3_strict_schema={self._gemini3_enforce_strict_schema}" ) # ========================================================================= @@ -1145,6 +1186,31 @@ def _gemini_cli_transform_schema(self, schema: Dict[str, Any]) -> Dict[str, Any] return schema + def _enforce_strict_schema(self, schema: Any) -> Any: + """ + Enforce strict JSON schema for Gemini 3 to prevent hallucinated parameters. + + Adds 'additionalProperties: false' recursively to all object schemas, + which tells the model it CANNOT add properties not in the schema. + """ + if not isinstance(schema, dict): + return schema + + result = {} + for key, value in schema.items(): + if isinstance(value, dict): + result[key] = self._enforce_strict_schema(value) + elif isinstance(value, list): + result[key] = [self._enforce_strict_schema(item) if isinstance(item, dict) else item for item in value] + else: + result[key] = value + + # Add additionalProperties: false to object schemas + if result.get("type") == "object" and "properties" in result: + result["additionalProperties"] = False + + return result + def _transform_tool_schemas(self, tools: List[Dict[str, Any]], model: str = "") -> List[Dict[str, Any]]: """ Transforms a list of OpenAI-style tool schemas into the format required by the Gemini CLI API. @@ -1153,6 +1219,7 @@ def _transform_tool_schemas(self, tools: List[Dict[str, Any]], model: str = "") For Gemini 3 models, also applies: - Namespace prefix to tool names - Parameter signature injection into descriptions + - Strict schema enforcement (additionalProperties: false) """ transformed_declarations = [] is_gemini_3 = self._is_gemini_3(model) @@ -1180,6 +1247,10 @@ def _transform_tool_schemas(self, tools: List[Dict[str, Any]], model: str = "") if name: new_function["name"] = f"{self._gemini3_tool_prefix}{name}" + # Enforce strict schema (additionalProperties: false) + if self._gemini3_enforce_strict_schema and "parametersJsonSchema" in new_function: + new_function["parametersJsonSchema"] = self._enforce_strict_schema(new_function["parametersJsonSchema"]) + # Inject parameter signature into description new_function = self._inject_signature_into_description(new_function) @@ -1218,10 +1289,21 @@ def _inject_signature_into_description(self, func_decl: Dict[str, Any]) -> Dict[ return func_decl - def _format_type_hint(self, prop_data: Dict[str, Any]) -> str: - """Format a type hint for a property schema.""" + def _format_type_hint(self, prop_data: Dict[str, Any], depth: int = 0) -> str: + """Format a detailed type hint for a property schema.""" type_hint = prop_data.get("type", "unknown") + # Handle enum values - show allowed options + if "enum" in prop_data: + enum_vals = prop_data["enum"] + if len(enum_vals) <= 5: + return f"string ENUM[{', '.join(repr(v) for v in enum_vals)}]" + return f"string ENUM[{len(enum_vals)} options]" + + # Handle const values + if "const" in prop_data: + return f"string CONST={repr(prop_data['const'])}" + if type_hint == "array": items = prop_data.get("items", {}) if isinstance(items, dict): @@ -1233,7 +1315,11 @@ def _format_type_hint(self, prop_data: Dict[str, Any]) -> str: nested_list = [] for n, d in nested_props.items(): if isinstance(d, dict): - t = d.get("type", "unknown") + # Recursively format nested types (limit depth) + if depth < 1: + t = self._format_type_hint(d, depth + 1) + else: + t = d.get("type", "unknown") req = " REQUIRED" if n in nested_req else "" nested_list.append(f"{n}: {t}{req}") return f"ARRAY_OF_OBJECTS[{', '.join(nested_list)}]" @@ -1241,6 +1327,18 @@ def _format_type_hint(self, prop_data: Dict[str, Any]) -> str: return f"ARRAY_OF_{item_type.upper()}" return "ARRAY" + if type_hint == "object": + nested_props = prop_data.get("properties", {}) + nested_req = prop_data.get("required", []) + if nested_props and depth < 1: + nested_list = [] + for n, d in nested_props.items(): + if isinstance(d, dict): + t = d.get("type", "unknown") + req = " REQUIRED" if n in nested_req else "" + nested_list.append(f"{n}: {t}{req}") + return f"object{{{', '.join(nested_list)}}}" + return type_hint def _inject_gemini3_system_instruction(self, request_payload: Dict[str, Any]) -> None: From eb3864bad538531cbd8b21bc027ff8e487272fc6 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 22:45:41 +0100 Subject: [PATCH 053/221] debugging pass to try to unfuck deployment --- src/proxy_app/main.py | 31 +++++++++++++++++++++++ src/rotator_library/credential_manager.py | 22 ++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index c2e318d0..816c985b 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -42,17 +42,31 @@ from dotenv import load_dotenv from glob import glob +# [DEBUG-REMOVE] Diagnostic logging for .env loading +print(f"[DEBUG-REMOVE] Current working directory: {Path.cwd()}") +print(f"[DEBUG-REMOVE] __file__ location: {Path(__file__).resolve().parent}") + # Load main .env first +_main_env_path = Path.cwd() / ".env" +print(f"[DEBUG-REMOVE] Looking for main .env at: {_main_env_path}") +print(f"[DEBUG-REMOVE] Main .env exists: {_main_env_path.exists()}") load_dotenv() # Load any additional .env files (e.g., antigravity_all_combined.env, gemini_cli_all_combined.env) _root_dir = Path.cwd() +_env_files_found = list(_root_dir.glob("*.env")) +print(f"[DEBUG-REMOVE] Found {len(_env_files_found)} .env files in {_root_dir}:") +for _ef in _env_files_found: + print(f"[DEBUG-REMOVE] - {_ef.name}") + for _env_file in sorted(_root_dir.glob("*.env")): if _env_file.name != ".env": # Skip main .env (already loaded) + print(f"[DEBUG-REMOVE] Loading additional .env file: {_env_file}") load_dotenv(_env_file, override=False) # Don't override existing values # Get proxy API key for display proxy_api_key = os.getenv("PROXY_API_KEY") +print(f"[DEBUG-REMOVE] PROXY_API_KEY from environment: {'SET' if proxy_api_key else 'NOT SET'}") if proxy_api_key: key_display = f"✓ {proxy_api_key}" else: @@ -288,12 +302,16 @@ def filter(self, record): # Discover API keys from environment variables api_keys = {} +print("[DEBUG-REMOVE] === Discovering API keys from environment ===") for key, value in os.environ.items(): if "_API_KEY" in key and key != "PROXY_API_KEY": provider = key.split("_API_KEY")[0].lower() if provider not in api_keys: api_keys[provider] = [] api_keys[provider].append(value) + print(f"[DEBUG-REMOVE] Found API key: {key} for provider '{provider}'") + +print(f"[DEBUG-REMOVE] Total providers with API keys: {list(api_keys.keys())}") # Load model ignore lists from environment variables ignore_models = {} @@ -337,8 +355,15 @@ async def lifespan(app: FastAPI): # The CredentialManager now handles all discovery, including .env overrides. # We pass all environment variables to it for this purpose. + print("[DEBUG-REMOVE] === Creating CredentialManager ===") + print(f"[DEBUG-REMOVE] Total environment variables: {len(os.environ)}") cred_manager = CredentialManager(os.environ) oauth_credentials = cred_manager.discover_and_prepare() + + print(f"[DEBUG-REMOVE] === OAuth credentials discovered ===") + print(f"[DEBUG-REMOVE] Providers with OAuth credentials: {list(oauth_credentials.keys())}") + for provider, paths in oauth_credentials.items(): + print(f"[DEBUG-REMOVE] {provider}: {len(paths)} credential(s) - {paths}") if not skip_oauth_init and oauth_credentials: logging.info("Starting OAuth credential validation and deduplication...") @@ -482,6 +507,9 @@ async def process_credential(provider: str, path: str, provider_instance): } # The client now uses the root logger configuration + print(f"[DEBUG-REMOVE] === Initializing RotatingClient ===") + print(f"[DEBUG-REMOVE] API keys providers: {list(api_keys.keys())}") + print(f"[DEBUG-REMOVE] OAuth providers: {list(oauth_credentials.keys())}") client = RotatingClient( api_keys=api_keys, oauth_credentials=oauth_credentials, # Pass OAuth config @@ -492,6 +520,9 @@ async def process_credential(provider: str, path: str, provider_instance): enable_request_logging=ENABLE_REQUEST_LOGGING, max_concurrent_requests_per_key=max_concurrent_requests_per_key ) + print(f"[DEBUG-REMOVE] RotatingClient.all_credentials keys: {list(client.all_credentials.keys())}") + for provider, creds in client.all_credentials.items(): + print(f"[DEBUG-REMOVE] {provider}: {len(creds)} credential(s)") client.background_refresher.start() # Start the background task app.state.rotating_client = client diff --git a/src/rotator_library/credential_manager.py b/src/rotator_library/credential_manager.py index 16be41c1..a4f35536 100644 --- a/src/rotator_library/credential_manager.py +++ b/src/rotator_library/credential_manager.py @@ -58,13 +58,25 @@ def _discover_env_oauth_credentials(self) -> Dict[str, List[str]]: """ env_credentials: Dict[str, Set[str]] = {} + # [DEBUG-REMOVE] Log all environment variable keys for OAuth providers + print(f"[DEBUG-REMOVE] === Scanning environment for OAuth credentials ===") + print(f"[DEBUG-REMOVE] ENV_OAUTH_PROVIDERS: {list(ENV_OAUTH_PROVIDERS.keys())}") + for provider, env_prefix in ENV_OAUTH_PROVIDERS.items(): found_indices: Set[str] = set() + print(f"[DEBUG-REMOVE] Scanning for provider '{provider}' with prefix '{env_prefix}'") # Check for numbered credentials (PROVIDER_N_ACCESS_TOKEN pattern) # Pattern: ANTIGRAVITY_1_ACCESS_TOKEN, ANTIGRAVITY_2_ACCESS_TOKEN, etc. numbered_pattern = re.compile(rf"^{env_prefix}_(\d+)_ACCESS_TOKEN$") + # [DEBUG-REMOVE] Show all matching environment variable keys + matching_keys = [k for k in self.env_vars.keys() if env_prefix in k] + if matching_keys: + print(f"[DEBUG-REMOVE] Found {len(matching_keys)} keys with '{env_prefix}': {matching_keys}") + else: + print(f"[DEBUG-REMOVE] No keys found with '{env_prefix}' prefix") + for key in self.env_vars.keys(): match = numbered_pattern.match(key) if match: @@ -73,20 +85,30 @@ def _discover_env_oauth_credentials(self) -> Dict[str, List[str]]: refresh_key = f"{env_prefix}_{index}_REFRESH_TOKEN" if refresh_key in self.env_vars and self.env_vars[refresh_key]: found_indices.add(index) + print(f"[DEBUG-REMOVE] ✓ Found numbered credential {index} for {provider}") + else: + print(f"[DEBUG-REMOVE] ✗ Missing REFRESH_TOKEN for {provider} credential {index}") # Check for legacy single credential (PROVIDER_ACCESS_TOKEN pattern) # Only use this if no numbered credentials exist if not found_indices: access_key = f"{env_prefix}_ACCESS_TOKEN" refresh_key = f"{env_prefix}_REFRESH_TOKEN" + print(f"[DEBUG-REMOVE] Checking legacy format: {access_key}, {refresh_key}") if (access_key in self.env_vars and self.env_vars[access_key] and refresh_key in self.env_vars and self.env_vars[refresh_key]): # Use "0" as the index for legacy single credential found_indices.add("0") + print(f"[DEBUG-REMOVE] ✓ Found legacy single credential for {provider}") + else: + print(f"[DEBUG-REMOVE] ✗ No legacy credential found for {provider}") if found_indices: env_credentials[provider] = found_indices lib_logger.info(f"Found {len(found_indices)} env-based credential(s) for {provider}") + print(f"[DEBUG-REMOVE] RESULT: {len(found_indices)} credential(s) registered for {provider}") + else: + print(f"[DEBUG-REMOVE] RESULT: No credentials found for {provider}") # Convert to virtual paths result: Dict[str, List[str]] = {} From a140a0dc43d760981398f2b3c39f2d154ea52d96 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 23:03:15 +0100 Subject: [PATCH 054/221] =?UTF-8?q?refactor(logging):=20=F0=9F=94=A8=20rem?= =?UTF-8?q?ove=20debug=20print=20statements=20and=20add=20concise=20deploy?= =?UTF-8?q?ment=20logs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removes verbose DEBUG-REMOVE diagnostic print statements that were used for troubleshooting .env loading and credential discovery during development. - Removes ~25 debug print statements from main.py and credential_manager.py - Adds concise, production-friendly logging for deployment verification: - .env file loading summary with file names - Credential loading summary with provider:count format - Preserves essential startup information for operational visibility - Improves code readability by removing debugging clutter - Maintains helpful deployment context without verbose diagnostic output --- src/proxy_app/main.py | 41 ++++++----------------- src/rotator_library/credential_manager.py | 22 ------------ 2 files changed, 11 insertions(+), 52 deletions(-) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 816c985b..4d0dd6f0 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -42,31 +42,23 @@ from dotenv import load_dotenv from glob import glob -# [DEBUG-REMOVE] Diagnostic logging for .env loading -print(f"[DEBUG-REMOVE] Current working directory: {Path.cwd()}") -print(f"[DEBUG-REMOVE] __file__ location: {Path(__file__).resolve().parent}") - # Load main .env first -_main_env_path = Path.cwd() / ".env" -print(f"[DEBUG-REMOVE] Looking for main .env at: {_main_env_path}") -print(f"[DEBUG-REMOVE] Main .env exists: {_main_env_path.exists()}") load_dotenv() # Load any additional .env files (e.g., antigravity_all_combined.env, gemini_cli_all_combined.env) _root_dir = Path.cwd() _env_files_found = list(_root_dir.glob("*.env")) -print(f"[DEBUG-REMOVE] Found {len(_env_files_found)} .env files in {_root_dir}:") -for _ef in _env_files_found: - print(f"[DEBUG-REMOVE] - {_ef.name}") - for _env_file in sorted(_root_dir.glob("*.env")): if _env_file.name != ".env": # Skip main .env (already loaded) - print(f"[DEBUG-REMOVE] Loading additional .env file: {_env_file}") load_dotenv(_env_file, override=False) # Don't override existing values +# Log discovered .env files for deployment verification +if _env_files_found: + _env_names = [_ef.name for _ef in _env_files_found] + print(f"📁 Loaded {len(_env_files_found)} .env file(s): {', '.join(_env_names)}") + # Get proxy API key for display proxy_api_key = os.getenv("PROXY_API_KEY") -print(f"[DEBUG-REMOVE] PROXY_API_KEY from environment: {'SET' if proxy_api_key else 'NOT SET'}") if proxy_api_key: key_display = f"✓ {proxy_api_key}" else: @@ -302,16 +294,12 @@ def filter(self, record): # Discover API keys from environment variables api_keys = {} -print("[DEBUG-REMOVE] === Discovering API keys from environment ===") for key, value in os.environ.items(): if "_API_KEY" in key and key != "PROXY_API_KEY": provider = key.split("_API_KEY")[0].lower() if provider not in api_keys: api_keys[provider] = [] api_keys[provider].append(value) - print(f"[DEBUG-REMOVE] Found API key: {key} for provider '{provider}'") - -print(f"[DEBUG-REMOVE] Total providers with API keys: {list(api_keys.keys())}") # Load model ignore lists from environment variables ignore_models = {} @@ -355,15 +343,8 @@ async def lifespan(app: FastAPI): # The CredentialManager now handles all discovery, including .env overrides. # We pass all environment variables to it for this purpose. - print("[DEBUG-REMOVE] === Creating CredentialManager ===") - print(f"[DEBUG-REMOVE] Total environment variables: {len(os.environ)}") cred_manager = CredentialManager(os.environ) oauth_credentials = cred_manager.discover_and_prepare() - - print(f"[DEBUG-REMOVE] === OAuth credentials discovered ===") - print(f"[DEBUG-REMOVE] Providers with OAuth credentials: {list(oauth_credentials.keys())}") - for provider, paths in oauth_credentials.items(): - print(f"[DEBUG-REMOVE] {provider}: {len(paths)} credential(s) - {paths}") if not skip_oauth_init and oauth_credentials: logging.info("Starting OAuth credential validation and deduplication...") @@ -507,9 +488,6 @@ async def process_credential(provider: str, path: str, provider_instance): } # The client now uses the root logger configuration - print(f"[DEBUG-REMOVE] === Initializing RotatingClient ===") - print(f"[DEBUG-REMOVE] API keys providers: {list(api_keys.keys())}") - print(f"[DEBUG-REMOVE] OAuth providers: {list(oauth_credentials.keys())}") client = RotatingClient( api_keys=api_keys, oauth_credentials=oauth_credentials, # Pass OAuth config @@ -520,9 +498,12 @@ async def process_credential(provider: str, path: str, provider_instance): enable_request_logging=ENABLE_REQUEST_LOGGING, max_concurrent_requests_per_key=max_concurrent_requests_per_key ) - print(f"[DEBUG-REMOVE] RotatingClient.all_credentials keys: {list(client.all_credentials.keys())}") - for provider, creds in client.all_credentials.items(): - print(f"[DEBUG-REMOVE] {provider}: {len(creds)} credential(s)") + + # Log loaded credentials summary (compact, always visible for deployment verification) + _api_summary = ', '.join([f"{p}:{len(c)}" for p, c in api_keys.items()]) if api_keys else "none" + _oauth_summary = ', '.join([f"{p}:{len(c)}" for p, c in oauth_credentials.items()]) if oauth_credentials else "none" + _total_summary = ', '.join([f"{p}:{len(c)}" for p, c in client.all_credentials.items()]) + print(f"🔑 Credentials loaded: {_total_summary} (API: {_api_summary} | OAuth: {_oauth_summary})") client.background_refresher.start() # Start the background task app.state.rotating_client = client diff --git a/src/rotator_library/credential_manager.py b/src/rotator_library/credential_manager.py index a4f35536..16be41c1 100644 --- a/src/rotator_library/credential_manager.py +++ b/src/rotator_library/credential_manager.py @@ -58,25 +58,13 @@ def _discover_env_oauth_credentials(self) -> Dict[str, List[str]]: """ env_credentials: Dict[str, Set[str]] = {} - # [DEBUG-REMOVE] Log all environment variable keys for OAuth providers - print(f"[DEBUG-REMOVE] === Scanning environment for OAuth credentials ===") - print(f"[DEBUG-REMOVE] ENV_OAUTH_PROVIDERS: {list(ENV_OAUTH_PROVIDERS.keys())}") - for provider, env_prefix in ENV_OAUTH_PROVIDERS.items(): found_indices: Set[str] = set() - print(f"[DEBUG-REMOVE] Scanning for provider '{provider}' with prefix '{env_prefix}'") # Check for numbered credentials (PROVIDER_N_ACCESS_TOKEN pattern) # Pattern: ANTIGRAVITY_1_ACCESS_TOKEN, ANTIGRAVITY_2_ACCESS_TOKEN, etc. numbered_pattern = re.compile(rf"^{env_prefix}_(\d+)_ACCESS_TOKEN$") - # [DEBUG-REMOVE] Show all matching environment variable keys - matching_keys = [k for k in self.env_vars.keys() if env_prefix in k] - if matching_keys: - print(f"[DEBUG-REMOVE] Found {len(matching_keys)} keys with '{env_prefix}': {matching_keys}") - else: - print(f"[DEBUG-REMOVE] No keys found with '{env_prefix}' prefix") - for key in self.env_vars.keys(): match = numbered_pattern.match(key) if match: @@ -85,30 +73,20 @@ def _discover_env_oauth_credentials(self) -> Dict[str, List[str]]: refresh_key = f"{env_prefix}_{index}_REFRESH_TOKEN" if refresh_key in self.env_vars and self.env_vars[refresh_key]: found_indices.add(index) - print(f"[DEBUG-REMOVE] ✓ Found numbered credential {index} for {provider}") - else: - print(f"[DEBUG-REMOVE] ✗ Missing REFRESH_TOKEN for {provider} credential {index}") # Check for legacy single credential (PROVIDER_ACCESS_TOKEN pattern) # Only use this if no numbered credentials exist if not found_indices: access_key = f"{env_prefix}_ACCESS_TOKEN" refresh_key = f"{env_prefix}_REFRESH_TOKEN" - print(f"[DEBUG-REMOVE] Checking legacy format: {access_key}, {refresh_key}") if (access_key in self.env_vars and self.env_vars[access_key] and refresh_key in self.env_vars and self.env_vars[refresh_key]): # Use "0" as the index for legacy single credential found_indices.add("0") - print(f"[DEBUG-REMOVE] ✓ Found legacy single credential for {provider}") - else: - print(f"[DEBUG-REMOVE] ✗ No legacy credential found for {provider}") if found_indices: env_credentials[provider] = found_indices lib_logger.info(f"Found {len(found_indices)} env-based credential(s) for {provider}") - print(f"[DEBUG-REMOVE] RESULT: {len(found_indices)} credential(s) registered for {provider}") - else: - print(f"[DEBUG-REMOVE] RESULT: No credentials found for {provider}") # Convert to virtual paths result: Dict[str, List[str]] = {} From 29df29409a71c9b55afa92f5eda4c2f4cf9e6f85 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 27 Nov 2025 23:13:23 +0100 Subject: [PATCH 055/221] =?UTF-8?q?fix(provider):=20=F0=9F=90=9B=20skip=20?= =?UTF-8?q?file=20operations=20for=20env://=20credential=20paths?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The project metadata loading and persistence logic was attempting to perform file I/O operations on env:// credential paths, which represent environment-based credentials rather than file-based ones. This caused unnecessary file operation errors. - Add checks using `_parse_env_credential_path()` to detect env:// paths before attempting file operations - Skip loading persisted project metadata from files for env:// credentials - Skip persisting project metadata to files for env:// credentials - Add debug logging to indicate when persistence is being skipped for env:// paths This prevents FileNotFoundError exceptions and improves reliability when using environment-based credential configuration. --- .../providers/gemini_cli_provider.py | 46 +++++++++++-------- 1 file changed, 28 insertions(+), 18 deletions(-) diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 601edf8e..259fb831 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -308,26 +308,30 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li lib_logger.debug(f"Found configured project_id override: {configured_project_id}") # Load credentials from file to check for persisted project_id and tier - try: - with open(credential_path, 'r') as f: - creds = json.load(f) - - metadata = creds.get("_proxy_metadata", {}) - persisted_project_id = metadata.get("project_id") - persisted_tier = metadata.get("tier") - - if persisted_project_id: - lib_logger.info(f"Loaded persisted project ID from credential file: {persisted_project_id}") - self.project_id_cache[credential_path] = persisted_project_id + # Skip for env:// paths (environment-based credentials don't persist to files) + credential_index = self._parse_env_credential_path(credential_path) + if credential_index is None: + # Only try to load from file if it's not an env:// path + try: + with open(credential_path, 'r') as f: + creds = json.load(f) - # Also load tier if available - if persisted_tier: - self.project_tier_cache[credential_path] = persisted_tier - lib_logger.debug(f"Loaded persisted tier: {persisted_tier}") + metadata = creds.get("_proxy_metadata", {}) + persisted_project_id = metadata.get("project_id") + persisted_tier = metadata.get("tier") - return persisted_project_id - except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: - lib_logger.debug(f"Could not load persisted project ID from file: {e}") + if persisted_project_id: + lib_logger.info(f"Loaded persisted project ID from credential file: {persisted_project_id}") + self.project_id_cache[credential_path] = persisted_project_id + + # Also load tier if available + if persisted_tier: + self.project_tier_cache[credential_path] = persisted_tier + lib_logger.debug(f"Loaded persisted tier: {persisted_tier}") + + return persisted_project_id + except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: + lib_logger.debug(f"Could not load persisted project ID from file: {e}") lib_logger.debug("No cached or configured project ID found, initiating discovery...") headers = {'Authorization': f'Bearer {access_token}', 'Content-Type': 'application/json'} @@ -625,6 +629,12 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li async def _persist_project_metadata(self, credential_path: str, project_id: str, tier: Optional[str]): """Persists project ID and tier to the credential file for faster future startups.""" + # Skip persistence for env:// paths (environment-based credentials) + credential_index = self._parse_env_credential_path(credential_path) + if credential_index is not None: + lib_logger.debug(f"Skipping project metadata persistence for env:// credential path: {credential_path}") + return + try: # Load current credentials with open(credential_path, 'r') as f: From c264be083034637447a8b96794ec78af5b89db44 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 28 Nov 2025 18:35:25 +0100 Subject: [PATCH 056/221] =?UTF-8?q?refactor(api):=20=F0=9F=94=A8=20change?= =?UTF-8?q?=20is=5Fready=20from=20method=20to=20property=20access?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Changed all `is_ready()` method calls to `is_ready` property access in the model_info_service across three endpoint functions: - list_models endpoint for enriched model data - get_model endpoint for model information retrieval - cost_estimate endpoint for cost calculation This aligns with the service's implementation where is_ready is exposed as a property rather than a callable method. --- src/proxy_app/main.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 4d0dd6f0..aa1278dc 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -932,7 +932,7 @@ async def list_models( if enriched and hasattr(request.app.state, 'model_info_service'): model_info_service = request.app.state.model_info_service - if model_info_service.is_ready(): + if model_info_service.is_ready: # Return enriched model data enriched_data = model_info_service.enrich_model_list(model_ids) return {"object": "list", "data": enriched_data} @@ -956,7 +956,7 @@ async def get_model( """ if hasattr(request.app.state, 'model_info_service'): model_info_service = request.app.state.model_info_service - if model_info_service.is_ready(): + if model_info_service.is_ready: info = model_info_service.get_model_info(model_id) if info: return info.to_dict() @@ -1066,7 +1066,7 @@ async def cost_estimate( # Try model info service first if hasattr(request.app.state, 'model_info_service'): model_info_service = request.app.state.model_info_service - if model_info_service.is_ready(): + if model_info_service.is_ready: cost = model_info_service.calculate_cost( model, prompt_tokens, completion_tokens, cache_read_tokens, cache_creation_tokens From 1ce8eba8df8e5b5ed6df80f98ce7a6f5b07233ce Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Wed, 3 Dec 2025 02:12:27 +0100 Subject: [PATCH 057/221] =?UTF-8?q?refactor(ui):=20=F0=9F=94=A8=20replace?= =?UTF-8?q?=20console.clear=20with=20cross-platform=20clear=5Fscreen=20fun?= =?UTF-8?q?ction?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaced all instances of `console.clear()` with a new `clear_screen()` helper function that uses native OS commands (`cls` for Windows, `clear` for Unix-like systems) instead of ANSI escape sequences. - Adds `clear_screen()` function to launcher_tui.py, settings_tool.py, and credential_tool.py - Replaces 18 instances of `console.clear()` across the codebase - Improves terminal clearing reliability on classic Windows conhost and modern terminals (Windows Terminal, Linux, Mac) - Removes unused anthropic_provider.py and bedrock_provider.py files - Enhances credential_tool API key setup with better provider filtering logic to prevent duplicates - Adds debug mode to show environment variable names in credential tool --- src/proxy_app/launcher_tui.py | 27 +++++-- src/proxy_app/main.py | 2 +- src/proxy_app/settings_tool.py | 28 +++++-- src/rotator_library/credential_tool.py | 74 +++++++++++++++---- .../providers/anthropic_provider.py | 31 -------- .../providers/bedrock_provider.py | 29 -------- 6 files changed, 101 insertions(+), 90 deletions(-) delete mode 100644 src/rotator_library/providers/anthropic_provider.py delete mode 100644 src/rotator_library/providers/bedrock_provider.py diff --git a/src/proxy_app/launcher_tui.py b/src/proxy_app/launcher_tui.py index 26a36bf1..2db109f9 100644 --- a/src/proxy_app/launcher_tui.py +++ b/src/proxy_app/launcher_tui.py @@ -16,6 +16,17 @@ console = Console() +def clear_screen(): + """ + Cross-platform terminal clear that works robustly on both + classic Windows conhost and modern terminals (Windows Terminal, Linux, Mac). + + Uses native OS commands instead of ANSI escape sequences: + - Windows (conhost & Windows Terminal): cls + - Unix-like systems (Linux, Mac): clear + """ + os.system('cls' if os.name == 'nt' else 'clear') + class LauncherConfig: """Manages launcher_config.json (host, port, logging only)""" @@ -262,7 +273,7 @@ def run(self): def show_main_menu(self): """Display main menu and handle selection""" - self.console.clear() + clear_screen() # Detect all settings settings = SettingsDetector.get_all_settings() @@ -394,7 +405,7 @@ def show_main_menu(self): def show_config_menu(self): """Display configuration sub-menu""" while True: - self.console.clear() + clear_screen() self.console.print(Panel.fit( "[bold cyan]⚙️ Proxy Configuration[/bold cyan]", @@ -455,7 +466,7 @@ def show_config_menu(self): def show_provider_settings_menu(self): """Display provider/advanced settings (read-only + launch tool)""" - self.console.clear() + clear_screen() settings = SettingsDetector.get_all_settings() credentials = settings["credentials"] @@ -573,7 +584,7 @@ def launch_credential_tool(self): import time # CRITICAL: Show full loading UI to replace the 6-7 second blank wait - self.console.clear() + clear_screen() _start_time = time.time() @@ -610,7 +621,7 @@ def launch_settings_tool(self): def show_about(self): """Display About page with project information""" - self.console.clear() + clear_screen() self.console.print(Panel.fit( "[bold cyan]ℹ️ About LLM API Key Proxy[/bold cyan]", @@ -654,7 +665,7 @@ def run_proxy(self): """Prepare and launch proxy in same window""" # Check if forced onboarding needed if self.needs_onboarding(): - self.console.clear() + clear_screen() self.console.print(Panel( Text.from_markup( "⚠️ [bold yellow]Setup Required[/bold yellow]\n\n" @@ -677,13 +688,13 @@ def run_proxy(self): return # Clear console and modify sys.argv - self.console.clear() + clear_screen() self.console.print(f"\n[bold green]🚀 Starting proxy on {self.config.config['host']}:{self.config.config['port']}...[/bold green]\n") # Clear console again to remove the starting message before main.py shows loading details import time time.sleep(0.5) # Brief pause so user sees the message - self.console.clear() + clear_screen() # Reconstruct sys.argv for main.py sys.argv = [ diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index aa1278dc..258a69f3 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -1137,7 +1137,7 @@ def needs_onboarding() -> bool: def show_onboarding_message(): """Display clear explanatory message for why onboarding is needed.""" - console.clear() # Clear terminal for clean presentation + os.system('cls' if os.name == 'nt' else 'clear') # Clear terminal for clean presentation console.print(Panel.fit( "[bold cyan]🚀 LLM API Key Proxy - First Time Setup[/bold cyan]", border_style="cyan" diff --git a/src/proxy_app/settings_tool.py b/src/proxy_app/settings_tool.py index 71641f33..59d91d5e 100644 --- a/src/proxy_app/settings_tool.py +++ b/src/proxy_app/settings_tool.py @@ -15,6 +15,18 @@ console = Console() +def clear_screen(): + """ + Cross-platform terminal clear that works robustly on both + classic Windows conhost and modern terminals (Windows Terminal, Linux, Mac). + + Uses native OS commands instead of ANSI escape sequences: + - Windows (conhost & Windows Terminal): cls + - Unix-like systems (Linux, Mac): clear + """ + os.system('cls' if os.name == 'nt' else 'clear') + + class AdvancedSettings: """Manages pending changes to .env""" @@ -389,7 +401,7 @@ def run(self): def show_main_menu(self): """Display settings categories""" - self.console.clear() + clear_screen() self.console.print(Panel.fit( "[bold cyan]🔧 Advanced Settings Configuration[/bold cyan]", @@ -436,7 +448,7 @@ def show_main_menu(self): def manage_custom_providers(self): """Manage custom provider API bases""" while True: - self.console.clear() + clear_screen() providers = self.provider_mgr.get_current_providers() @@ -533,7 +545,7 @@ def manage_custom_providers(self): def manage_model_definitions(self): """Manage provider model definitions""" while True: - self.console.clear() + clear_screen() all_providers = self.model_mgr.get_all_providers_with_models() @@ -710,7 +722,7 @@ def edit_model_definitions(self, providers: List[str]): current_models = {m: {} for m in current_models} while True: - self.console.clear() + clear_screen() self.console.print(f"[bold]Editing models for: {provider}[/bold]\n") self.console.print("Current models:") for i, (name, definition) in enumerate(current_models.items(), 1): @@ -788,7 +800,7 @@ def view_model_definitions(self, providers: List[str]): input("\nPress Enter to continue...") return - self.console.clear() + clear_screen() self.console.print(f"[bold]Provider: {provider}[/bold]\n") self.console.print("[bold]📦 Configured Models:[/bold]") self.console.print("━" * 50) @@ -816,7 +828,7 @@ def view_model_definitions(self, providers: List[str]): def manage_provider_settings(self): """Manage provider-specific settings (Antigravity, Gemini CLI)""" while True: - self.console.clear() + clear_screen() available_providers = self.provider_settings_mgr.get_available_providers() @@ -863,7 +875,7 @@ def manage_provider_settings(self): def _manage_single_provider_settings(self, provider: str): """Manage settings for a single provider""" while True: - self.console.clear() + clear_screen() display_name = provider.replace("_", " ").title() definitions = self.provider_settings_mgr.get_provider_settings_definitions(provider) @@ -1005,7 +1017,7 @@ def _reset_all_provider_settings(self, provider: str, settings_list: List[str]): def manage_concurrency_limits(self): """Manage concurrency limits""" while True: - self.console.clear() + clear_screen() limits = self.concurrency_mgr.get_current_limits() diff --git a/src/rotator_library/credential_tool.py b/src/rotator_library/credential_tool.py index 1949f134..6aca4bdf 100644 --- a/src/rotator_library/credential_tool.py +++ b/src/rotator_library/credential_tool.py @@ -37,6 +37,18 @@ def _ensure_providers_loaded(): return _provider_factory, _provider_plugins +def clear_screen(): + """ + Cross-platform terminal clear that works robustly on both + classic Windows conhost and modern terminals (Windows Terminal, Linux, Mac). + + Uses native OS commands instead of ANSI escape sequences: + - Windows (conhost & Windows Terminal): cls + - Unix-like systems (Linux, Mac): clear + """ + os.system('cls' if os.name == 'nt' else 'clear') + + def _get_credential_number_from_filename(filename: str) -> int: """ Extract credential number from filename like 'provider_oauth_1.json' -> 1 @@ -127,6 +139,9 @@ async def setup_api_key(): """ console.print(Panel("[bold cyan]API Key Setup[/bold cyan]", expand=False)) + # Debug toggle: Set to True to see env var names next to each provider + SHOW_ENV_VAR_NAMES = True + # Verified list of LiteLLM providers with their friendly names and API key variables LITELLM_PROVIDERS = { "OpenAI": "OPENAI_API_KEY", "Anthropic": "ANTHROPIC_API_KEY", @@ -162,26 +177,59 @@ async def setup_api_key(): "Nscale": "NSCALE_API_KEY", "Recraft": "RECRAFT_API_KEY", "v0": "V0_API_KEY", "Vercel": "VERCEL_AI_GATEWAY_API_KEY", "Topaz": "TOPAZ_API_KEY", "ElevenLabs": "ELEVENLABS_API_KEY", - "Deepgram": "DEEPGRAM_API_KEY", "Custom API": "CUSTOM_API_KEY", + "Deepgram": "DEEPGRAM_API_KEY", "GitHub Models": "GITHUB_TOKEN", "GitHub Copilot": "GITHUB_COPILOT_API_KEY", } # Discover custom providers and add them to the list - # Note: gemini_cli is OAuth-only, but qwen_code and iflow support both OAuth and API keys + # Note: gemini_cli and antigravity are OAuth-only + # qwen_code API key support is a fallback + # iflow API key support is a feature _, PROVIDER_PLUGINS = _ensure_providers_loaded() - oauth_only_providers = {'gemini_cli', 'antigravity'} - discovered_providers = { - p.replace('_', ' ').title(): p.upper() + "_API_KEY" - for p in PROVIDER_PLUGINS.keys() - if p not in oauth_only_providers and p.replace('_', ' ').title() not in LITELLM_PROVIDERS + + # Build a set of environment variables already in LITELLM_PROVIDERS + # to avoid duplicates based on the actual API key names + litellm_env_vars = set(LITELLM_PROVIDERS.values()) + + # Providers to exclude from API key list + exclude_providers = { + 'gemini_cli', # OAuth-only + 'antigravity', # OAuth-only + 'qwen_code', # API key is fallback, OAuth is primary - don't advertise + 'openai_compatible', # Base class, not a real provider } + discovered_providers = {} + for provider_key in PROVIDER_PLUGINS.keys(): + if provider_key in exclude_providers: + continue + + # Create environment variable name + env_var = provider_key.upper() + "_API_KEY" + + # Check if this env var already exists in LITELLM_PROVIDERS + # This catches duplicates like GEMINI_API_KEY, MISTRAL_API_KEY, etc. + if env_var in litellm_env_vars: + # Already in LITELLM_PROVIDERS with better name, skip this one + continue + + # Create display name for this custom provider + display_name = provider_key.replace('_', ' ').title() + discovered_providers[display_name] = env_var + + # LITELLM_PROVIDERS takes precedence (comes first in merge) combined_providers = {**LITELLM_PROVIDERS, **discovered_providers} provider_display_list = sorted(combined_providers.keys()) provider_text = Text() for i, provider_name in enumerate(provider_display_list): - provider_text.append(f" {i + 1}. {provider_name}\n") + if SHOW_ENV_VAR_NAMES: + # Extract env var prefix (before _API_KEY) + env_var = combined_providers[provider_name] + prefix = env_var.replace("_API_KEY", "").replace("_", " ") + provider_text.append(f" {i + 1}. {provider_name} ({prefix})\n") + else: + provider_text.append(f" {i + 1}. {provider_name}\n") console.print(Panel(provider_text, title="Available Providers for API Key", style="bold blue")) @@ -1000,7 +1048,7 @@ async def export_credentials_submenu(): Submenu for credential export options. """ while True: - console.clear() + clear_screen() console.print(Panel("[bold cyan]Export Credentials to .env[/bold cyan]", title="--- API Key Proxy ---", expand=False)) console.print(Panel( @@ -1111,7 +1159,7 @@ async def main(clear_on_start=True): while True: # Clear screen between menu selections for cleaner UX - console.clear() + clear_screen() console.print(Panel("[bold cyan]Interactive Credential Setup[/bold cyan]", title="--- API Key Proxy ---", expand=False)) console.print(Panel( @@ -1179,7 +1227,7 @@ async def main(clear_on_start=True): elif setup_type == "2": await setup_api_key() #console.print("\n[dim]Press Enter to return to main menu...[/dim]") - input() + #input() elif setup_type == "3": await export_credentials_submenu() @@ -1225,7 +1273,7 @@ def run_credential_tool(from_launcher=False): # If from launcher, don't clear screen at start to preserve loading messages try: asyncio.run(main(clear_on_start=not from_launcher)) - console.clear() # Clear terminal when credential tool exits + clear_screen() # Clear terminal when credential tool exits except KeyboardInterrupt: console.print("\n[bold yellow]Exiting setup.[/bold yellow]") - console.clear() # Clear terminal on keyboard interrupt too \ No newline at end of file + clear_screen() # Clear terminal on keyboard interrupt too \ No newline at end of file diff --git a/src/rotator_library/providers/anthropic_provider.py b/src/rotator_library/providers/anthropic_provider.py deleted file mode 100644 index 5859c2b9..00000000 --- a/src/rotator_library/providers/anthropic_provider.py +++ /dev/null @@ -1,31 +0,0 @@ -import httpx -import logging -from typing import List -from .provider_interface import ProviderInterface - -lib_logger = logging.getLogger('rotator_library') -lib_logger.propagate = False # Ensure this logger doesn't propagate to root -if not lib_logger.handlers: - lib_logger.addHandler(logging.NullHandler()) - -class AnthropicProvider(ProviderInterface): - """ - Provider implementation for the Anthropic API. - """ - async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: - """ - Fetches the list of available models from the Anthropic API. - """ - try: - response = await client.get( - "https://api.anthropic.com/v1/models", - headers={ - "x-api-key": api_key, - "anthropic-version": "2023-06-01" - } - ) - response.raise_for_status() - return [f"anthropic/{model['id']}" for model in response.json().get("data", [])] - except httpx.RequestError as e: - lib_logger.error(f"Failed to fetch Anthropic models: {e}") - return [] diff --git a/src/rotator_library/providers/bedrock_provider.py b/src/rotator_library/providers/bedrock_provider.py deleted file mode 100644 index a7a1a07a..00000000 --- a/src/rotator_library/providers/bedrock_provider.py +++ /dev/null @@ -1,29 +0,0 @@ -import httpx -import logging -from typing import List -from .provider_interface import ProviderInterface - -lib_logger = logging.getLogger('rotator_library') -lib_logger.propagate = False # Ensure this logger doesn't propagate to root -if not lib_logger.handlers: - lib_logger.addHandler(logging.NullHandler()) - -class BedrockProvider(ProviderInterface): - """ - Provider implementation for AWS Bedrock. - """ - async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: - """ - Returns a hardcoded list of common Bedrock models, as there is no - simple, unauthenticated API endpoint to list them. - """ - # Note: Listing Bedrock models typically requires AWS credentials and boto3. - # For a simple, key-based proxy, we'll list common models. - # This can be expanded with full AWS authentication if needed. - lib_logger.info("Returning hardcoded list for Bedrock. Full discovery requires AWS auth.") - return [ - "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", - "bedrock/anthropic.claude-3-haiku-20240307-v1:0", - "bedrock/cohere.command-r-plus-v1:0", - "bedrock/mistral.mistral-large-2402-v1:0", - ] From aeb8eaf7230a2a2760d974e84bb4dc59efdd6b23 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Wed, 3 Dec 2025 04:50:34 +0100 Subject: [PATCH 058/221] =?UTF-8?q?fix(provider):=20=F0=9F=90=9B=20add=20a?= =?UTF-8?q?utomatic=20ID=20repair=20for=20mismatched=20tool=20call=20respo?= =?UTF-8?q?nses?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements a recovery mechanism to handle cases where proxies or clients mutate tool call IDs (e.g., transforming "toolu_" prefix to "call_" prefix), which previously caused response grouping failures. - Enhanced pending group handling to attempt orphan response matching when expected IDs are missing - Automatically repairs response IDs to match their corresponding tool calls - Maintains response order by using first available orphan for each unmatched call - Added warning logs for ID mismatch repairs and partial group satisfaction - Integrated tool response grouping fix into the main message transformation pipeline This prevents tool call conversation corruption when intermediary services modify request/response identifiers. --- .../providers/antigravity_provider.py | 37 +++++++++++++++---- 1 file changed, 30 insertions(+), 7 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 3f06b197..dddbcefb 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -1307,16 +1307,38 @@ def _fix_tool_response_grouping( new_contents.append(content) # Handle remaining groups (shouldn't happen in well-formed conversations) + # Attempt recovery by matching orphans to unsatisfied calls for group in pending_groups: group_ids = group["ids"] - available_ids = [gid for gid in group_ids if gid in collected_responses] - if available_ids: - group_responses = [collected_responses.pop(gid) for gid in available_ids] + group_responses = [] + + for expected_id in group_ids: + if expected_id in collected_responses: + group_responses.append(collected_responses.pop(expected_id)) + elif collected_responses: + # Recovery: Match with an orphan response + # This handles cases where client/proxy mutates IDs (e.g. toolu_ -> call_) + # Get the first available orphan ID to maintain order + orphan_id = next(iter(collected_responses)) + orphan_resp = collected_responses.pop(orphan_id) + + # Fix the ID in the response to match the call + orphan_resp["functionResponse"]["id"] = expected_id + + lib_logger.warning( + f"[Grouping] Auto-repaired ID mismatch: mapped response '{orphan_id}' " + f"to call '{expected_id}'" + ) + group_responses.append(orphan_resp) + + if group_responses: new_contents.append({"parts": group_responses, "role": "user"}) - lib_logger.warning( - f"[Grouping] Partial group satisfaction: expected {len(group_ids)}, " - f"got {len(available_ids)} responses" - ) + + if len(group_responses) != len(group_ids): + lib_logger.warning( + f"[Grouping] Partial group satisfaction after repair: " + f"expected {len(group_ids)}, got {len(group_responses)} responses" + ) # Warn about unmatched responses if collected_responses: @@ -2305,6 +2327,7 @@ async def count_tokens( internal_model = self._alias_to_internal(model) system_instruction, contents = self._transform_messages(messages, internal_model) + contents = self._fix_tool_response_grouping(contents) gemini_payload = {"contents": contents} if system_instruction: From e8e22c6e90fb2fdbef2714e1e0ff5d5cb545684e Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Wed, 3 Dec 2025 07:44:45 +0100 Subject: [PATCH 059/221] =?UTF-8?q?docs(deployment):=20=F0=9F=93=9A=20add?= =?UTF-8?q?=20comprehensive=20VPS=20deployment=20guide=20for=20OAuth=20pro?= =?UTF-8?q?viders?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add detailed appendix section covering VPS deployment workflows for OAuth-based providers (Antigravity, Gemini CLI, iFlow). The guide addresses the localhost callback challenge inherent to OAuth flows on remote servers. - Document three professional deployment approaches: local authentication with credential export (recommended), SSH port forwarding for direct VPS authentication, and credential file copying - Provide production-ready security best practices including firewall configuration, environment variable management, and systemd service setup - Include comprehensive troubleshooting section for common VPS deployment issues - Add comparison tables for OAuth callback ports, credential storage methods, and deployment approach trade-offs - Explain technical rationale for why OAuth callbacks fail on remote servers and how each solution addresses the problem --- Deployment guide.md | 366 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 366 insertions(+) diff --git a/Deployment guide.md b/Deployment guide.md index 57acd536..ac8c2d7b 100644 --- a/Deployment guide.md +++ b/Deployment guide.md @@ -174,3 +174,369 @@ curl -X POST https://your-service.onrender.com/v1/chat/completions -H "Content-T That is it. +--- + +## Appendix: Deploying to a Custom VPS + +If you're deploying the proxy to a **custom VPS** (DigitalOcean, AWS EC2, Linode, etc.) instead of Render.com, you'll encounter special considerations when setting up OAuth providers (Antigravity, Gemini CLI, iFlow). This section covers the professional deployment workflow. + +### Understanding the OAuth Callback Problem + +OAuth providers like Antigravity, Gemini CLI, and iFlow require an interactive authentication flow that: + +1. Opens a browser for you to log in +2. Redirects back to a **local callback server** running on specific ports +3. Receives an authorization code to exchange for tokens + +The callback servers bind to `localhost` on these ports: + +| Provider | Port | Notes | +|---------------|-------|--------------------------------------------| +| **Antigravity** | 51121 | Google OAuth with extended scopes | +| **Gemini CLI** | 8085 | Google OAuth for Gemini API | +| **iFlow** | 11451 | Authorization Code flow with API key fetch | +| **Qwen Code** | N/A | Uses Device Code flow - works on remote VPS ✅ | + +**The Issue**: When running on a remote VPS, your local browser cannot reach `http://localhost:51121` (or other callback ports) on the remote server, causing authentication to fail with a "connection refused" error. + +### Recommended Deployment Workflow + +There are **three professional approaches** to handle OAuth authentication for VPS deployment, listed from most recommended to least: + +--- + +### **Option 1: Authenticate Locally, Deploy Credentials (RECOMMENDED)** + +This is the **cleanest and most secure** approach. Complete OAuth flows on your local machine, export to environment variables, then deploy. + +#### Step 1: Clone and Set Up Locally + +```bash +# On your local development machine +git clone https://github.com/YOUR-USERNAME/LLM-API-Key-Proxy.git +cd LLM-API-Key-Proxy + +# Install dependencies +pip install -r requirements.txt +``` + +#### Step 2: Run OAuth Authentication Locally + +```bash +# Start the credential tool +python -m rotator_library.credential_tool +``` + +Select **"Add OAuth Credential"** and choose your provider: +- Antigravity +- Gemini CLI +- iFlow +- Qwen Code (works directly on VPS, but can authenticate locally too) + +The tool will: +1. Open your browser automatically +2. Start a local callback server +3. Complete the OAuth flow +4. Save credentials to `oauth_creds/_oauth_N.json` + +#### Step 3: Export Credentials to Environment Variables + +Still in the credential tool, select the export option for each provider: +- **"Export Antigravity to .env"** +- **"Export Gemini CLI to .env"** +- **"Export iFlow to .env"** +- **"Export Qwen Code to .env"** + +The tool generates a `.env` file snippet like: + +```env +# Antigravity OAuth Credentials +ANTIGRAVITY_ACCESS_TOKEN="ya29.a0AfB_byD..." +ANTIGRAVITY_REFRESH_TOKEN="1//0gL6dK9..." +ANTIGRAVITY_EXPIRY_DATE="1735901234567" +ANTIGRAVITY_EMAIL="user@gmail.com" +ANTIGRAVITY_CLIENT_ID="1071006060591-..." +ANTIGRAVITY_CLIENT_SECRET="GOCSPX-..." +ANTIGRAVITY_TOKEN_URI="https://oauth2.googleapis.com/token" +ANTIGRAVITY_UNIVERSE_DOMAIN="googleapis.com" +``` + +Copy these variables to a file (e.g., `oauth_credentials.env`). + +#### Step 4: Deploy to VPS + +**Method A: Using Environment Variables (Recommended)** + +```bash +# On your VPS +cd /path/to/LLM-API-Key-Proxy + +# Create or edit .env file +nano .env + +# Paste the exported environment variables +# Also add your PROXY_API_KEY and other provider keys + +# Start the proxy +uvicorn src.proxy_app.main:app --host 0.0.0.0 --port 8000 +``` + +**Method B: Upload Credential Files** + +```bash +# On your local machine - copy credential files to VPS +scp -r oauth_creds/ user@your-vps-ip:/path/to/LLM-API-Key-Proxy/ + +# On VPS - verify files exist +ls -la oauth_creds/ + +# Start the proxy +uvicorn src.proxy_app.main:app --host 0.0.0.0 --port 8000 +``` + +> **Note**: Environment variables are preferred for production deployments (more secure, easier to manage, works with container orchestration). + +--- + +### **Option 2: SSH Port Forwarding (For Direct VPS Authentication)** + +If you need to authenticate directly on the VPS (e.g., you don't have a local development environment), use SSH port forwarding to create secure tunnels. + +#### How It Works + +SSH tunnels forward ports from your local machine to the remote VPS, allowing your local browser to reach the callback servers. + +#### Step-by-Step Process + +**Step 1: Create SSH Tunnels** + +From your **local machine**, open a terminal and run: + +```bash +# Forward all OAuth callback ports at once +ssh -L 51121:localhost:51121 -L 8085:localhost:8085 -L 11451:localhost:11451 user@your-vps-ip + +# Alternative: Forward ports individually as needed +ssh -L 51121:localhost:51121 user@your-vps-ip # For Antigravity +ssh -L 8085:localhost:8085 user@your-vps-ip # For Gemini CLI +ssh -L 11451:localhost:11451 user@your-vps-ip # For iFlow +``` + +**Keep this SSH session open** during the entire authentication process. + +**Step 2: Run Credential Tool on VPS** + +In the same SSH terminal (or open a new SSH connection): + +```bash +cd /path/to/LLM-API-Key-Proxy + +# Ensure Python dependencies are installed +pip install -r requirements.txt + +# Run the credential tool +python -m rotator_library.credential_tool +``` + +**Step 3: Complete OAuth Flow** + +1. Select **"Add OAuth Credential"** → Choose your provider +2. The tool displays an authorization URL +3. **Click the URL in your local browser** (works because of the SSH tunnel!) +4. Complete the authentication flow +5. The browser redirects to `localhost:` - **this now routes through the tunnel to your VPS** +6. Credentials are saved to `oauth_creds/` on the VPS + +**Step 4: Export to Environment Variables** + +Still in the credential tool: +1. Select the export option for each provider +2. Copy the generated environment variables +3. Add them to `/path/to/LLM-API-Key-Proxy/.env` on your VPS + +**Step 5: Close Tunnels and Deploy** + +```bash +# Exit the SSH session with tunnels (Ctrl+D or type 'exit') +# Tunnels are no longer needed + +# Start the proxy on VPS (in a screen/tmux session or as a service) +uvicorn src.proxy_app.main:app --host 0.0.0.0 --port 8000 +``` + +--- + +### **Option 3: Copy Credential Files to VPS** + +If you've already authenticated locally and have credential files, you can copy them directly. + +#### Copy OAuth Credential Files + +```bash +# From your local machine +scp -r oauth_creds/ user@your-vps-ip:/path/to/LLM-API-Key-Proxy/ + +# Verify on VPS +ssh user@your-vps-ip +ls -la /path/to/LLM-API-Key-Proxy/oauth_creds/ +``` + +Expected files: +- `antigravity_oauth_1.json` +- `gemini_cli_oauth_1.json` +- `iflow_oauth_1.json` +- `qwen_code_oauth_1.json` + +#### Configure .env to Use Credential Files + +On your VPS, edit `.env`: + +```env +# Option A: Use credential files directly (not recommended for production) +# No special configuration needed - the proxy auto-detects oauth_creds/ folder + +# Option B: Export to environment variables (recommended) +# Run credential tool and export each provider to .env +``` + +--- + +### Environment Variables vs. Credential Files + +| Aspect | Environment Variables | Credential Files | +|---------------------------|------------------------------------------|--------------------------------------------| +| **Security** | ✅ More secure (no files on disk) | ⚠️ Files readable if server compromised | +| **Container-Friendly** | ✅ Perfect for Docker/K8s | ❌ Requires volume mounts | +| **Ease of Rotation** | ✅ Update .env and restart | ⚠️ Need to regenerate JSON files | +| **Backup/Version Control**| ✅ Easy to manage with secrets managers | ❌ Binary files, harder to manage | +| **Auto-Refresh** | ✅ Uses refresh tokens | ✅ Uses refresh tokens | +| **Recommended For** | Production deployments | Local development / testing | + +**Best Practice**: Always export to environment variables for VPS/cloud deployments. + +--- + +### Production Deployment Checklist + +#### Security Best Practices + +- [ ] Never commit `.env` or `oauth_creds/` to version control +- [ ] Use environment variables instead of credential files in production +- [ ] Secure your VPS firewall - **do not** open OAuth callback ports (51121, 8085, 11451) to public internet +- [ ] Use SSH port forwarding only during initial authentication +- [ ] Rotate credentials regularly using the credential tool's export feature +- [ ] Set file permissions on `.env`: `chmod 600 .env` + +#### Firewall Configuration + +OAuth callback ports should **never** be publicly exposed: + +```bash +# ❌ DO NOT DO THIS - keeps ports closed +# sudo ufw allow 51121/tcp +# sudo ufw allow 8085/tcp +# sudo ufw allow 11451/tcp + +# ✅ Only open your proxy API port +sudo ufw allow 8000/tcp + +# Check firewall status +sudo ufw status +``` + +The SSH tunnel method works **without** opening these ports because traffic routes through the SSH connection (port 22). + +#### Running as a Service + +Create a systemd service file on your VPS: + +```bash +# Create service file +sudo nano /etc/systemd/system/llm-proxy.service +``` + +```ini +[Unit] +Description=LLM API Key Proxy +After=network.target + +[Service] +Type=simple +User=your-username +WorkingDirectory=/path/to/LLM-API-Key-Proxy +Environment="PATH=/path/to/python/bin" +ExecStart=/path/to/python/bin/uvicorn src.proxy_app.main:app --host 0.0.0.0 --port 8000 +Restart=always +RestartSec=10 + +[Install] +WantedBy=multi-user.target +``` + +```bash +# Enable and start the service +sudo systemctl daemon-reload +sudo systemctl enable llm-proxy +sudo systemctl start llm-proxy + +# Check status +sudo systemctl status llm-proxy + +# View logs +sudo journalctl -u llm-proxy -f +``` + +--- + +### Troubleshooting VPS Deployment + +#### "localhost:51121 connection refused" Error + +**Cause**: Trying to authenticate directly on VPS without SSH tunnel. + +**Solution**: Use Option 1 (authenticate locally) or Option 2 (SSH port forwarding). + +#### OAuth Credentials Not Loading + +```bash +# Check if environment variables are set +printenv | grep -E '(ANTIGRAVITY|GEMINI_CLI|IFLOW|QWEN_CODE)' + +# Verify .env file exists and is readable +ls -la .env +cat .env | grep -E '(ANTIGRAVITY|GEMINI_CLI|IFLOW|QWEN_CODE)' + +# Check credential files if using file-based approach +ls -la oauth_creds/ +``` + +#### Token Refresh Failing + +The proxy automatically refreshes tokens using refresh tokens. If refresh fails: + +1. **Re-authenticate**: Run credential tool again and export new credentials +2. **Check token expiry**: Some providers require periodic re-authentication +3. **Verify credentials**: Ensure `REFRESH_TOKEN` is present in environment variables + +#### Permission Denied on .env + +```bash +# Set correct permissions +chmod 600 .env +chown your-username:your-username .env +``` + +--- + +### Summary: VPS Deployment Best Practices + +1. **Authenticate locally** on your development machine (easiest, most secure) +2. **Export to environment variables** using the credential tool's built-in export feature +3. **Deploy to VPS** by adding environment variables to `.env` +4. **Never open OAuth callback ports** to the public internet +5. **Use SSH port forwarding** only if you must authenticate directly on VPS +6. **Run as a systemd service** for production reliability +7. **Monitor logs** for authentication errors and token refresh issues + +This approach ensures secure, production-ready deployment while maintaining the convenience of OAuth authentication. + From 7cb148b4c1e912ce3f354ec946947ca14e521bdd Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Wed, 3 Dec 2025 22:48:42 +0100 Subject: [PATCH 060/221] =?UTF-8?q?feat(core):=20=E2=9C=A8=20add=20structu?= =?UTF-8?q?red=20error=20accumulator=20and=20consistent=20error=20handling?= =?UTF-8?q?/reporting?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce RequestErrorAccumulator and related helpers to aggregate and classify errors across credential rotation and provide structured client-facing error responses. - Add RequestErrorAccumulator to record per-credential errors (abnormal vs normal), build concise server log messages and structured client error payloads. - Add mask_credential, is_abnormal_error, should_rotate_on_error and should_retry_same_key helpers and extend classify_error to better handle httpx.HTTPStatusError and more error types (forbidden, quota_exceeded, context_window_exceeded, etc.). - Update RotatingClient (sync and streaming paths) to: - initialize and record errors into the accumulator during retries/rotation, - mask credentials in logs, - handle httpx.HTTPStatusError explicitly, - standardize cooldowns and retry-vs-rotate decisions, - return a structured error response (dict) for non-streaming failures and yield structured JSON error payloads for streaming failures. - Improve failure_logger: extract and limit response bodies, capture error chains, and log richer failure details to failures.log while keeping concise main logs. - Silence noisy client-facing yields on recoverable errors and rotate keys transparently; make quota errors fatal after repeated occurrences in streaming with explicit client message. BREAKING CHANGE: RotatingClient failure behavior changed — methods that previously returned None (on exhausting keys or timeout) now return a structured error dict with shape: { "error": { "message": string, "type": "proxy_all_credentials_exhausted" | "proxy_timeout" | ..., "details": { "model": string, "provider": string, "credentials_tried": int, "timeout": bool, "abnormal_errors": [ ... ]? , "normal_error_summary": string? } } } Streaming endpoints now yield a JSON error payload (same structure) before the final [DONE]. Update callers to handle the new error response format. --- src/rotator_library/client.py | 336 +++++++++++++++++--------- src/rotator_library/error_handler.py | 297 +++++++++++++++++++++++ src/rotator_library/failure_logger.py | 75 +++++- 3 files changed, 590 insertions(+), 118 deletions(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index e536aeb4..ef322e6c 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -25,6 +25,10 @@ classify_error, AllProviders, NoAvailableKeysError, + should_rotate_on_error, + should_retry_same_key, + RequestErrorAccumulator, + mask_credential, ) from .providers import PROVIDER_PLUGINS from .providers.openai_compatible_provider import OpenAICompatibleProvider @@ -816,6 +820,11 @@ async def _execute_with_retry( f"Credential priorities for {provider}: {', '.join(f'P{p}={len([c for c in credentials_for_provider if credential_priorities.get(c)==p])}' for p in sorted(set(credential_priorities.values())))}" ) + # Initialize error accumulator for tracking errors across credential rotation + error_accumulator = RequestErrorAccumulator() + error_accumulator.model = model + error_accumulator.provider = provider + while ( len(tried_creds) < len(credentials_for_provider) and time.time() < deadline ): @@ -1023,8 +1032,12 @@ async def _execute_with_retry( # Extract a clean error message for the user-facing log error_message = str(e).split("\n")[0] + + # Record in accumulator for client reporting + error_accumulator.record_error(current_cred, classified_error, error_message) + lib_logger.info( - f"Key ...{current_cred[-6:]} hit rate limit for model {model}. Reason: '{error_message}'. Rotating key." + f"Key {mask_credential(current_cred)} hit rate limit for {model}. Rotating key." ) if classified_error.status_code == 429: @@ -1032,16 +1045,10 @@ async def _execute_with_retry( await self.cooldown_manager.start_cooldown( provider, cooldown_duration ) - lib_logger.warning( - f"IP-based rate limit detected for {provider}. Starting a {cooldown_duration}-second global cooldown." - ) await self.usage_manager.record_failure( current_cred, model, classified_error ) - lib_logger.warning( - f"Key ...{current_cred[-6:]} encountered a rate limit. Trying next key." - ) break # Move to the next key except ( @@ -1060,6 +1067,8 @@ async def _execute_with_retry( else {}, ) classified_error = classify_error(e) + error_message = str(e).split("\n")[0] + # Provider-level error: don't increment consecutive failures await self.usage_manager.record_failure( current_cred, model, classified_error, @@ -1067,9 +1076,10 @@ async def _execute_with_retry( ) if attempt >= self.max_retries - 1: - error_message = str(e).split("\n")[0] + # Record in accumulator only on final failure for this key + error_accumulator.record_error(current_cred, classified_error, error_message) lib_logger.warning( - f"Key ...{current_cred[-6:]} failed after max retries for model {model} due to a server error. Reason: '{error_message}'. Rotating key." + f"Key {mask_credential(current_cred)} failed after max retries due to server error. Rotating." ) break # Move to the next key @@ -1081,18 +1091,73 @@ async def _execute_with_retry( # If the required wait time exceeds the budget, don't wait; rotate to the next key immediately. if wait_time > remaining_budget: + error_accumulator.record_error(current_cred, classified_error, error_message) lib_logger.warning( - f"Required retry wait time ({wait_time:.2f}s) exceeds remaining budget ({remaining_budget:.2f}s). Rotating key early." + f"Retry wait ({wait_time:.2f}s) exceeds budget ({remaining_budget:.2f}s). Rotating key." ) break - error_message = str(e).split("\n")[0] lib_logger.warning( - f"Key ...{current_cred[-6:]} encountered a server error for model {model}. Reason: '{error_message}'. Retrying in {wait_time:.2f}s." + f"Key {mask_credential(current_cred)} server error. Retrying in {wait_time:.2f}s." ) await asyncio.sleep(wait_time) continue # Retry with the same key + except httpx.HTTPStatusError as e: + # Handle HTTP errors from httpx (e.g., from custom providers like Antigravity) + last_exception = e + log_failure( + api_key=current_cred, + model=model, + attempt=attempt + 1, + error=e, + request_headers=dict(request.headers) + if request + else {}, + ) + + classified_error = classify_error(e) + error_message = str(e).split("\n")[0] + + # Record in accumulator for client reporting + error_accumulator.record_error(current_cred, classified_error, error_message) + + lib_logger.warning( + f"Key {mask_credential(current_cred)} HTTP {e.response.status_code} ({classified_error.error_type})." + ) + + # Check if this error should trigger rotation + if not should_rotate_on_error(classified_error): + lib_logger.error( + f"Non-recoverable error ({classified_error.error_type}). Failing request." + ) + raise last_exception + + # Handle rate limits with cooldown + if classified_error.error_type in ["rate_limit", "quota_exceeded"]: + cooldown_duration = classified_error.retry_after or 60 + await self.cooldown_manager.start_cooldown( + provider, cooldown_duration + ) + + # Check if we should retry same key (server errors with retries left) + if should_retry_same_key(classified_error) and attempt < self.max_retries - 1: + wait_time = classified_error.retry_after or (1 * (2**attempt)) + random.uniform(0, 1) + remaining_budget = deadline - time.time() + if wait_time <= remaining_budget: + lib_logger.warning( + f"Server error, retrying same key in {wait_time:.2f}s." + ) + await asyncio.sleep(wait_time) + continue + + # Record failure and rotate to next key + await self.usage_manager.record_failure( + current_cred, model, classified_error + ) + lib_logger.info(f"Rotating to next key after {classified_error.error_type} error.") + break + except Exception as e: last_exception = e log_failure( @@ -1107,30 +1172,32 @@ async def _execute_with_retry( if request and await request.is_disconnected(): lib_logger.warning( - f"Client disconnected. Aborting retries for credential ...{current_cred[-6:]}." + f"Client disconnected. Aborting retries for {mask_credential(current_cred)}." ) raise last_exception classified_error = classify_error(e) error_message = str(e).split("\n")[0] + + # Record in accumulator for client reporting + error_accumulator.record_error(current_cred, classified_error, error_message) + lib_logger.warning( - f"Key ...{current_cred[-6:]} failed with {classified_error.error_type} (Status: {classified_error.status_code}). Error: {error_message}. Rotating key." + f"Key {mask_credential(current_cred)} {classified_error.error_type} (HTTP {classified_error.status_code})." ) - if classified_error.status_code == 429: + + # Handle rate limits with cooldown + if classified_error.status_code == 429 or classified_error.error_type in ["rate_limit", "quota_exceeded"]: cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration ) - lib_logger.warning( - f"IP-based rate limit detected for {provider} from generic exception. Starting a {cooldown_duration}-second global cooldown." - ) - if classified_error.error_type in [ - "invalid_request", - "context_window_exceeded", - "authentication", - ]: - # For these errors, we should not retry with other keys. + # Check if this error should trigger rotation + if not should_rotate_on_error(classified_error): + lib_logger.error( + f"Non-recoverable error ({classified_error.error_type}). Failing request." + ) raise last_exception await self.usage_manager.record_failure( @@ -1141,14 +1208,18 @@ async def _execute_with_retry( if key_acquired and current_cred: await self.usage_manager.release_key(current_cred, model) - if last_exception: - # Log the final error but do not raise it, as per the new requirement. - # The client should not see intermittent failures. - lib_logger.error( - f"Request failed after trying all keys or exceeding global timeout. Last error: {last_exception}" - ) + # Check if we exhausted all credentials or timed out + if time.time() >= deadline: + error_accumulator.timeout_occurred = True + + if error_accumulator.has_errors(): + # Log concise summary for server logs + lib_logger.error(error_accumulator.build_log_message()) + + # Return the structured error response for the client + return error_accumulator.build_client_error_response() - # Return None to indicate failure without propagating a disruptive exception. + # Return None to indicate failure without error details (shouldn't normally happen) return None async def _streaming_acompletion_with_retry( @@ -1259,6 +1330,11 @@ async def _streaming_acompletion_with_retry( f"Credential priorities for {provider}: {', '.join(f'P{p}={len([c for c in credentials_for_provider if credential_priorities.get(c)==p])}' for p in sorted(set(credential_priorities.values())))}" ) + # Initialize error accumulator for tracking errors across credential rotation + error_accumulator = RequestErrorAccumulator() + error_accumulator.model = model + error_accumulator.provider = provider + try: while ( len(tried_creds) < len(credentials_for_provider) @@ -1402,21 +1478,44 @@ async def _streaming_acompletion_with_retry( litellm.RateLimitError, httpx.HTTPStatusError, ) as e: - if ( - isinstance(e, httpx.HTTPStatusError) - and e.response.status_code != 429 - ): - raise e - last_exception = e # If the exception is our custom wrapper, unwrap the original error original_exc = getattr(e, "data", e) classified_error = classify_error(original_exc) + error_message = str(original_exc).split("\n")[0] + + log_failure( + api_key=current_cred, + model=model, + attempt=attempt + 1, + error=e, + request_headers=dict(request.headers) + if request + else {}, + ) + + # Record in accumulator for client reporting + error_accumulator.record_error(current_cred, classified_error, error_message) + + # Check if this error should trigger rotation + if not should_rotate_on_error(classified_error): + lib_logger.error( + f"Non-recoverable error ({classified_error.error_type}) during custom stream. Failing." + ) + raise last_exception + + # Handle rate limits with cooldown + if classified_error.error_type in ["rate_limit", "quota_exceeded"]: + cooldown_duration = classified_error.retry_after or 60 + await self.cooldown_manager.start_cooldown( + provider, cooldown_duration + ) + await self.usage_manager.record_failure( current_cred, model, classified_error ) lib_logger.warning( - f"Credential ...{current_cred[-6:]} encountered a recoverable error ({classified_error.error_type}) during custom provider stream. Rotating key." + f"Cred {mask_credential(current_cred)} {classified_error.error_type} (HTTP {classified_error.status_code}). Rotating." ) break @@ -1436,6 +1535,8 @@ async def _streaming_acompletion_with_retry( else {}, ) classified_error = classify_error(e) + error_message = str(e).split("\n")[0] + # Provider-level error: don't increment consecutive failures await self.usage_manager.record_failure( current_cred, model, classified_error, @@ -1443,8 +1544,9 @@ async def _streaming_acompletion_with_retry( ) if attempt >= self.max_retries - 1: + error_accumulator.record_error(current_cred, classified_error, error_message) lib_logger.warning( - f"Credential ...{current_cred[-6:]} failed after max retries for model {model} due to a server error. Rotating key." + f"Cred {mask_credential(current_cred)} failed after max retries. Rotating." ) break @@ -1453,14 +1555,14 @@ async def _streaming_acompletion_with_retry( ) + random.uniform(0, 1) remaining_budget = deadline - time.time() if wait_time > remaining_budget: + error_accumulator.record_error(current_cred, classified_error, error_message) lib_logger.warning( - f"Required retry wait time ({wait_time:.2f}s) exceeds remaining budget ({remaining_budget:.2f}s). Rotating key early." + f"Retry wait ({wait_time:.2f}s) exceeds budget. Rotating." ) break - error_message = str(e).split("\n")[0] lib_logger.warning( - f"Credential ...{current_cred[-6:]} encountered a server error for model {model}. Reason: '{error_message}'. Retrying in {wait_time:.2f}s." + f"Cred {mask_credential(current_cred)} server error. Retrying in {wait_time:.2f}s." ) await asyncio.sleep(wait_time) continue @@ -1477,15 +1579,22 @@ async def _streaming_acompletion_with_retry( else {}, ) classified_error = classify_error(e) + error_message = str(e).split("\n")[0] + + # Record in accumulator + error_accumulator.record_error(current_cred, classified_error, error_message) + lib_logger.warning( - f"Credential ...{current_cred[-6:]} failed with {classified_error.error_type} (Status: {classified_error.status_code}). Error: {str(e)}. Rotating key." + f"Cred {mask_credential(current_cred)} {classified_error.error_type} (HTTP {classified_error.status_code})." ) - if classified_error.error_type in [ - "invalid_request", - "context_window_exceeded", - "authentication", - ]: + + # Check if this error should trigger rotation + if not should_rotate_on_error(classified_error): + lib_logger.error( + f"Non-recoverable error ({classified_error.error_type}). Failing." + ) raise last_exception + await self.usage_manager.record_failure( current_cred, model, classified_error ) @@ -1590,7 +1699,7 @@ async def _streaming_acompletion_with_retry( yield chunk return - except (StreamedAPIError, litellm.RateLimitError) as e: + except (StreamedAPIError, litellm.RateLimitError, httpx.HTTPStatusError) as e: last_exception = e # This is the final, robust handler for streamed errors. @@ -1599,6 +1708,13 @@ async def _streaming_acompletion_with_retry( # The actual exception might be wrapped in our StreamedAPIError. original_exc = getattr(e, "data", e) classified_error = classify_error(original_exc) + + # Check if this error should trigger rotation + if not should_rotate_on_error(classified_error): + lib_logger.error( + f"Non-recoverable error ({classified_error.error_type}) during litellm stream. Failing." + ) + raise last_exception try: # The full error JSON is in the string representation of the exception. @@ -1606,18 +1722,13 @@ async def _streaming_acompletion_with_retry( r"(\{.*\})", str(original_exc), re.DOTALL ) if json_str_match: - # The string may contain byte-escaped characters (e.g., \\n). cleaned_str = codecs.decode( json_str_match.group(1), "unicode_escape" ) error_payload = json.loads(cleaned_str) except (json.JSONDecodeError, TypeError): - lib_logger.warning( - "Could not parse JSON details from streamed error exception." - ) error_payload = {} - # Now, log the failure with the extracted raw response. log_failure( api_key=current_cred, model=model, @@ -1631,20 +1742,19 @@ async def _streaming_acompletion_with_retry( error_details = error_payload.get("error", {}) error_status = error_details.get("status", "") - # Fallback to the full string if parsing fails. error_message_text = error_details.get( - "message", str(original_exc) + "message", str(original_exc).split("\n")[0] ) + + # Record in accumulator for client reporting + error_accumulator.record_error(current_cred, classified_error, error_message_text) if ( "quota" in error_message_text.lower() or "resource_exhausted" in error_status.lower() ): consecutive_quota_failures += 1 - lib_logger.warning( - f"Credential ...{current_cred[-6:]} hit a quota limit. This is consecutive failure #{consecutive_quota_failures} for this request." - ) - + quota_value = "N/A" quota_id = "N/A" if "details" in error_details and isinstance( @@ -1654,15 +1764,10 @@ async def _streaming_acompletion_with_retry( if isinstance(detail.get("violations"), list): for violation in detail["violations"]: if "quotaValue" in violation: - quota_value = violation[ - "quotaValue" - ] + quota_value = violation["quotaValue"] if "quotaId" in violation: quota_id = violation["quotaId"] - if ( - quota_value != "N/A" - and quota_id != "N/A" - ): + if quota_value != "N/A" and quota_id != "N/A": break await self.usage_manager.record_failure( @@ -1670,48 +1775,34 @@ async def _streaming_acompletion_with_retry( ) if consecutive_quota_failures >= 3: - console_log_message = ( - f"Terminating stream for credential ...{current_cred[-6:]} due to 3rd consecutive quota error. " - f"This is now considered a fatal input data error. ID: {quota_id}, Limit: {quota_value}." - ) + # Fatal: likely input data too large client_error_message = ( - "FATAL: Request failed after 3 consecutive quota errors, " - "indicating the input data is too large for the model's per-request limit. " - f"Last Error Message: '{error_message_text}'. Limit: {quota_value} (Quota ID: {quota_id})." + f"Request failed after 3 consecutive quota errors (input may be too large). " + f"Limit: {quota_value} (Quota ID: {quota_id})" + ) + lib_logger.error( + f"Fatal quota error for {mask_credential(current_cred)}. ID: {quota_id}, Limit: {quota_value}" ) - lib_logger.error(console_log_message) - yield f"data: {json.dumps({'error': {'message': client_error_message, 'type': 'proxy_fatal_quota_error'}})}\n\n" yield "data: [DONE]\n\n" return - else: - # [MODIFIED] Do not yield to the client. Just log and break to rotate the key. lib_logger.warning( - f"Quota error on credential ...{current_cred[-6:]} (failure {consecutive_quota_failures}/3). Rotating key silently." + f"Cred {mask_credential(current_cred)} quota error ({consecutive_quota_failures}/3). Rotating." ) break else: consecutive_quota_failures = 0 - # [MODIFIED] Do not yield to the client. Just log and break to rotate the key. lib_logger.warning( - f"Credential ...{current_cred[-6:]} encountered a recoverable error ({classified_error.error_type}) during stream. Rotating key silently." + f"Cred {mask_credential(current_cred)} {classified_error.error_type}. Rotating." ) - if ( - classified_error.error_type == "rate_limit" - and classified_error.status_code == 429 - ): - cooldown_duration = ( - classified_error.retry_after or 60 - ) + if classified_error.error_type in ["rate_limit", "quota_exceeded"]: + cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration ) - lib_logger.warning( - f"IP-based rate limit detected for {provider}. Starting a {cooldown_duration}-second global cooldown." - ) await self.usage_manager.record_failure( current_cred, model, classified_error @@ -1735,6 +1826,11 @@ async def _streaming_acompletion_with_retry( else {}, ) classified_error = classify_error(e) + error_message_text = str(e).split("\n")[0] + + # Record error in accumulator (server errors are abnormal) + error_accumulator.record_error(current_cred, classified_error, error_message_text) + # Provider-level error: don't increment consecutive failures await self.usage_manager.record_failure( current_cred, model, classified_error, @@ -1758,9 +1854,8 @@ async def _streaming_acompletion_with_retry( ) break - error_message = str(e).split("\n")[0] lib_logger.warning( - f"Credential ...{current_cred[-6:]} encountered a server error for model {model}. Reason: '{error_message}'. Retrying in {wait_time:.2f}s." + f"Credential ...{current_cred[-6:]} encountered a server error for model {model}. Reason: '{error_message_text}'. Retrying in {wait_time:.2f}s." ) await asyncio.sleep(wait_time) continue @@ -1778,49 +1873,66 @@ async def _streaming_acompletion_with_retry( else {}, ) classified_error = classify_error(e) + error_message_text = str(e).split("\n")[0] + + # Record error in accumulator + error_accumulator.record_error(current_cred, classified_error, error_message_text) lib_logger.warning( - f"Credential ...{current_cred[-6:]} failed with {classified_error.error_type} (Status: {classified_error.status_code}). Error: {str(e)}. Rotating key." + f"Credential ...{current_cred[-6:]} failed with {classified_error.error_type} (Status: {classified_error.status_code}). Error: {error_message_text}." ) - if classified_error.status_code == 429: + # Handle rate limits with cooldown + if classified_error.status_code == 429 or classified_error.error_type in ["rate_limit", "quota_exceeded"]: cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration ) lib_logger.warning( - f"IP-based rate limit detected for {provider} from generic stream exception. Starting a {cooldown_duration}-second global cooldown." + f"Rate limit detected for {provider}. Starting {cooldown_duration}s cooldown." ) - if classified_error.error_type in [ - "invalid_request", - "context_window_exceeded", - "authentication", - ]: + # Check if this error should trigger rotation + if not should_rotate_on_error(classified_error): + # Non-rotatable errors - fail immediately + lib_logger.error( + f"Non-recoverable error ({classified_error.error_type}). Failing request." + ) raise last_exception - # [MODIFIED] Do not yield to the client here. + # Record failure and rotate to next key await self.usage_manager.record_failure( current_cred, model, classified_error ) + lib_logger.info(f"Rotating to next key after {classified_error.error_type} error.") break finally: if key_acquired and current_cred: await self.usage_manager.release_key(current_cred, model) - final_error_message = "Failed to complete the streaming request: No available API keys after rotation or global timeout exceeded." - if last_exception: - final_error_message = f"Failed to complete the streaming request. Last error: {str(last_exception)}" - lib_logger.error( - f"Streaming request failed after trying all keys. Last error: {last_exception}" - ) + # Build detailed error response using error accumulator + error_accumulator.timeout_occurred = time.time() >= deadline + error_accumulator.model = model + error_accumulator.provider = provider + + if error_accumulator.has_errors(): + # Log concise summary for server logs + lib_logger.error(error_accumulator.build_log_message()) + + # Build structured error response for client + error_response = error_accumulator.build_client_error_response() + error_data = error_response else: + # Fallback if no errors were recorded (shouldn't happen) + final_error_message = "Request failed: No available API keys after rotation or timeout." + if last_exception: + final_error_message = f"Request failed. Last error: {str(last_exception)}" + error_data = { + "error": {"message": final_error_message, "type": "proxy_error"} + } lib_logger.error(final_error_message) - - error_data = { - "error": {"message": final_error_message, "type": "proxy_error"} - } + yield f"data: {json.dumps(error_data)}\n\n" yield "data: [DONE]\n\n" diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index a3775f7f..96a6cb73 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -65,6 +65,208 @@ class PreRequestCallbackError(Exception): pass +# ============================================================================= +# ERROR TRACKING FOR CLIENT REPORTING +# ============================================================================= + +# Abnormal errors that require attention and should always be reported to client +ABNORMAL_ERROR_TYPES = frozenset({ + "forbidden", # 403 - credential access issue + "authentication", # 401 - credential invalid/revoked + "pre_request_callback_error", # Internal proxy error +}) + +# Normal/expected errors during operation - only report if ALL credentials fail +NORMAL_ERROR_TYPES = frozenset({ + "rate_limit", # 429 - expected during high load + "quota_exceeded", # Expected when quota runs out + "server_error", # 5xx - transient provider issues + "api_connection", # Network issues - transient +}) + + +def is_abnormal_error(classified_error: "ClassifiedError") -> bool: + """ + Check if an error is abnormal and should be reported to the client. + + Abnormal errors indicate credential issues that need attention: + - 403 Forbidden: Credential doesn't have access + - 401 Unauthorized: Credential is invalid/revoked + + Normal errors are expected during operation: + - 429 Rate limit: Expected during high load + - 5xx Server errors: Transient provider issues + """ + return classified_error.error_type in ABNORMAL_ERROR_TYPES + + +def mask_credential(credential: str) -> str: + """ + Mask a credential for safe display in logs and error messages. + + - For API keys: shows last 6 characters (e.g., "...xyz123") + - For OAuth file paths: shows just the filename (e.g., "antigravity_oauth_1.json") + """ + import os + if os.path.isfile(credential): + return os.path.basename(credential) + elif len(credential) > 6: + return f"...{credential[-6:]}" + else: + return "***" + + +class RequestErrorAccumulator: + """ + Tracks errors encountered during a request's credential rotation cycle. + + Used to build informative error messages for clients when all credentials + are exhausted. Distinguishes between abnormal errors (that need attention) + and normal errors (expected during operation). + """ + + def __init__(self): + self.abnormal_errors: list = [] # 403, 401 - always report details + self.normal_errors: list = [] # 429, 5xx - summarize only + self.total_credentials_tried: int = 0 + self.timeout_occurred: bool = False + self.model: str = "" + self.provider: str = "" + + def record_error( + self, + credential: str, + classified_error: "ClassifiedError", + error_message: str + ): + """Record an error for a credential.""" + self.total_credentials_tried += 1 + masked_cred = mask_credential(credential) + + error_record = { + "credential": masked_cred, + "error_type": classified_error.error_type, + "status_code": classified_error.status_code, + "message": self._truncate_message(error_message, 150) + } + + if is_abnormal_error(classified_error): + self.abnormal_errors.append(error_record) + else: + self.normal_errors.append(error_record) + + def _truncate_message(self, message: str, max_length: int = 150) -> str: + """Truncate error message for readability.""" + # Take first line and truncate + first_line = message.split('\n')[0] + if len(first_line) > max_length: + return first_line[:max_length] + "..." + return first_line + + def has_errors(self) -> bool: + """Check if any errors were recorded.""" + return bool(self.abnormal_errors or self.normal_errors) + + def has_abnormal_errors(self) -> bool: + """Check if any abnormal errors were recorded.""" + return bool(self.abnormal_errors) + + def get_normal_error_summary(self) -> str: + """Get a summary of normal errors (not individual details).""" + if not self.normal_errors: + return "" + + # Count by type + counts = {} + for err in self.normal_errors: + err_type = err["error_type"] + counts[err_type] = counts.get(err_type, 0) + 1 + + # Build summary like "3 rate_limit, 1 server_error" + parts = [f"{count} {err_type}" for err_type, count in counts.items()] + return ", ".join(parts) + + def build_client_error_response(self) -> dict: + """ + Build a structured error response for the client. + + Returns a dict suitable for JSON serialization in the error response. + """ + # Determine the primary failure reason + if self.timeout_occurred: + error_type = "proxy_timeout" + base_message = f"Request timed out after trying {self.total_credentials_tried} credential(s)" + else: + error_type = "proxy_all_credentials_exhausted" + base_message = f"All {self.total_credentials_tried} credential(s) exhausted for {self.provider}" + + # Build human-readable message + message_parts = [base_message] + + if self.abnormal_errors: + message_parts.append("\n\nCredential issues (require attention):") + for err in self.abnormal_errors: + status = f"HTTP {err['status_code']}" if err['status_code'] else err['error_type'] + message_parts.append(f"\n • {err['credential']}: {status} - {err['message']}") + + normal_summary = self.get_normal_error_summary() + if normal_summary: + if self.abnormal_errors: + message_parts.append(f"\n\nAdditionally: {normal_summary} (expected during normal operation)") + else: + message_parts.append(f"\n\nAll failures were: {normal_summary}") + message_parts.append("\nThis is normal during high load - retry later or add more credentials.") + + response = { + "error": { + "message": "".join(message_parts), + "type": error_type, + "details": { + "model": self.model, + "provider": self.provider, + "credentials_tried": self.total_credentials_tried, + "timeout": self.timeout_occurred, + } + } + } + + # Only include abnormal errors in details (they need attention) + if self.abnormal_errors: + response["error"]["details"]["abnormal_errors"] = self.abnormal_errors + + # Include summary of normal errors + if normal_summary: + response["error"]["details"]["normal_error_summary"] = normal_summary + + return response + + def build_log_message(self) -> str: + """ + Build a concise log message for server-side logging. + + Shorter than client message, suitable for terminal display. + """ + parts = [] + + if self.timeout_occurred: + parts.append(f"TIMEOUT: {self.total_credentials_tried} creds tried for {self.model}") + else: + parts.append(f"ALL CREDS EXHAUSTED: {self.total_credentials_tried} tried for {self.model}") + + if self.abnormal_errors: + abnormal_summary = ", ".join( + f"{e['credential']}={e['status_code'] or e['error_type']}" + for e in self.abnormal_errors + ) + parts.append(f"ISSUES: {abnormal_summary}") + + normal_summary = self.get_normal_error_summary() + if normal_summary: + parts.append(f"Normal: {normal_summary}") + + return " | ".join(parts) + + class ClassifiedError: """A structured representation of a classified error.""" @@ -197,25 +399,74 @@ def classify_error(e: Exception) -> ClassifiedError: """ Classifies an exception into a structured ClassifiedError object. Now handles both litellm and httpx exceptions. + + Error types and their typical handling: + - rate_limit (429): Rotate key, may retry with backoff + - server_error (5xx): Retry with backoff, then rotate + - forbidden (403): Rotate key immediately (access denied for this credential) + - authentication (401): Rotate key, trigger re-auth if OAuth + - quota_exceeded: Rotate key (credential quota exhausted) + - invalid_request (400): Don't retry - client error in request + - context_window_exceeded: Don't retry - request too large + - api_connection: Retry with backoff, then rotate + - unknown: Rotate key (safer to try another) """ status_code = getattr(e, "status_code", None) + if isinstance(e, httpx.HTTPStatusError): # [NEW] Handle httpx errors first status_code = e.response.status_code + + # Try to get error body for better classification + try: + error_body = e.response.text.lower() if hasattr(e.response, 'text') else "" + except Exception: + error_body = "" + if status_code == 401: return ClassifiedError( error_type="authentication", original_exception=e, status_code=status_code, ) + if status_code == 403: + # 403 Forbidden - credential doesn't have access, should rotate + # Could be: IP restriction, account disabled, permission denied, etc. + return ClassifiedError( + error_type="forbidden", + original_exception=e, + status_code=status_code, + ) if status_code == 429: retry_after = get_retry_after(e) + # Check if this is a quota error vs rate limit + if "quota" in error_body or "resource_exhausted" in error_body: + return ClassifiedError( + error_type="quota_exceeded", + original_exception=e, + status_code=status_code, + retry_after=retry_after, + ) return ClassifiedError( error_type="rate_limit", original_exception=e, status_code=status_code, retry_after=retry_after, ) + if status_code == 400: + # Check for context window / token limit errors + if "context" in error_body or "token" in error_body or "too long" in error_body: + return ClassifiedError( + error_type="context_window_exceeded", + original_exception=e, + status_code=status_code, + ) + return ClassifiedError( + error_type="invalid_request", + original_exception=e, + status_code=status_code, + ) if 400 <= status_code < 500: + # Other 4xx errors - generally client errors return ClassifiedError( error_type="invalid_request", original_exception=e, @@ -313,6 +564,52 @@ def is_unrecoverable_error(e: Exception) -> bool: return isinstance(e, (InvalidRequestError, AuthenticationError, BadRequestError)) +def should_rotate_on_error(classified_error: ClassifiedError) -> bool: + """ + Determines if an error should trigger key rotation. + + Errors that SHOULD rotate (try another key): + - rate_limit: Current key is throttled + - quota_exceeded: Current key/account exhausted + - forbidden: Current credential denied access + - authentication: Current credential invalid + - server_error: Provider having issues (might work with different endpoint/key) + - api_connection: Network issues (might be transient) + - unknown: Safer to try another key + + Errors that should NOT rotate (fail immediately): + - invalid_request: Client error in request payload (won't help to retry) + - context_window_exceeded: Request too large (won't help to retry) + - pre_request_callback_error: Internal proxy error + + Returns: + True if should rotate to next key, False if should fail immediately + """ + non_rotatable_errors = { + "invalid_request", + "context_window_exceeded", + "pre_request_callback_error", + } + return classified_error.error_type not in non_rotatable_errors + + +def should_retry_same_key(classified_error: ClassifiedError) -> bool: + """ + Determines if an error should retry with the same key (with backoff). + + Only server errors and connection issues should retry the same key, + as these are often transient. + + Returns: + True if should retry same key, False if should rotate immediately + """ + retryable_errors = { + "server_error", + "api_connection", + } + return classified_error.error_type in retryable_errors + + class AllProviders: """ A class to handle provider-specific settings, such as custom API bases. diff --git a/src/rotator_library/failure_logger.py b/src/rotator_library/failure_logger.py index 3f92c8f3..8c4e043a 100644 --- a/src/rotator_library/failure_logger.py +++ b/src/rotator_library/failure_logger.py @@ -43,32 +43,95 @@ def format(self, record): # Get the main library logger for concise, propagated messages main_lib_logger = logging.getLogger('rotator_library') +def _extract_response_body(error: Exception) -> str: + """ + Extract the full response body from various error types. + + Handles: + - httpx.HTTPStatusError: response.text or response.content + - litellm exceptions: various response attributes + - Other exceptions: str(error) + """ + # Try to get response body from httpx errors + if hasattr(error, 'response') and error.response is not None: + response = error.response + # Try .text first (decoded) + if hasattr(response, 'text') and response.text: + return response.text + # Try .content (bytes) + if hasattr(response, 'content') and response.content: + try: + return response.content.decode('utf-8', errors='replace') + except Exception: + return str(response.content) + # Try reading response if it's a streaming response that was read + if hasattr(response, '_content') and response._content: + try: + return response._content.decode('utf-8', errors='replace') + except Exception: + return str(response._content) + + # Check for litellm's body attribute + if hasattr(error, 'body') and error.body: + return str(error.body) + + # Check for message attribute that might contain response + if hasattr(error, 'message') and error.message: + return str(error.message) + + return None + + def log_failure(api_key: str, model: str, attempt: int, error: Exception, request_headers: dict, raw_response_text: str = None): """ Logs a detailed failure message to a file and a concise summary to the main logger. + + Args: + api_key: The API key or credential path that was used + model: The model that was requested + attempt: The attempt number (1-based) + error: The exception that occurred + request_headers: Headers from the original request + raw_response_text: Optional pre-extracted response body (e.g., from streaming) """ # 1. Log the full, detailed error to the dedicated failures.log file # Prioritize the explicitly passed raw response text, as it may contain # reassembled data from a stream that is not available on the exception object. raw_response = raw_response_text - if not raw_response and hasattr(error, 'response') and hasattr(error.response, 'text'): - raw_response = error.response.text + if not raw_response: + raw_response = _extract_response_body(error) + # Get full error message (not truncated) + full_error_message = str(error) + + # Also capture any nested/wrapped exception info + error_chain = [] + current_error = error + while current_error: + error_chain.append({ + "type": type(current_error).__name__, + "message": str(current_error)[:2000] # Limit per-error message size + }) + current_error = getattr(current_error, '__cause__', None) or getattr(current_error, '__context__', None) + if len(error_chain) > 5: # Prevent infinite loops + break + detailed_log_data = { "timestamp": datetime.utcnow().isoformat(), - "api_key_ending": api_key[-4:], + "api_key_ending": api_key[-4:] if len(api_key) >= 4 else "****", "model": model, "attempt_number": attempt, "error_type": type(error).__name__, - "error_message": str(error), - "raw_response": raw_response, + "error_message": full_error_message[:5000], # Limit total size + "raw_response": raw_response[:10000] if raw_response else None, # Limit response size "request_headers": request_headers, + "error_chain": error_chain if len(error_chain) > 1 else None, } failure_logger.error(detailed_log_data) # 2. Log a concise summary to the main library logger, which will propagate summary_message = ( - f"API call failed for model {model} with key ...{api_key[-4:]}. " + f"API call failed for model {model} with key ...{api_key[-4:] if len(api_key) >= 4 else '****'}. " f"Error: {type(error).__name__}. See failures.log for details." ) main_lib_logger.error(summary_message) From d6e982eddfe7f23d5ae58d0b10c861f2ba168bc4 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Wed, 3 Dec 2025 23:47:56 +0100 Subject: [PATCH 061/221] =?UTF-8?q?refactor(provider):=20=F0=9F=94=A8=20re?= =?UTF-8?q?place=20hardcoded=20project=20generation=20with=20dynamic=20GCP?= =?UTF-8?q?=20resolution?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eliminate synthetic project ID generation in favor of real Google Cloud project lookup mechanism. The Antigravity provider previously generated random project identifiers for API requests. This implementation now queries actual GCP infrastructure to obtain legitimate project IDs through multiple fallback strategies. - Add credential-path-indexed memory store to avoid redundant lookups across requests - Create waterfall resolution: environment configuration → stored credentials → Cloud Code API probe → Resource Manager enumeration - Modify payload assembly to accept externally-resolved project parameter instead of generating random values - Inject resolution step into request pipeline before format transformation occurs - Store successfully discovered identifiers in credential metadata for subsequent invocations - Handle all network failures gracefully with 20-second timeout boundaries The transformation function signature now requires explicit project_id argument rather than computing it internally, shifting discovery responsibility to the caller context. --- .../providers/antigravity_provider.py | 96 ++++++++++++++++++- 1 file changed, 93 insertions(+), 3 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index dddbcefb..b3d51d8a 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -448,6 +448,7 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface): def __init__(self): super().__init__() self.model_definitions = ModelDefinitions() + self.project_id_cache: Dict[str, str] = {} # Cache project ID per credential path # Base URL management self._base_url_index = 0 @@ -587,6 +588,90 @@ def _generate_thinking_cache_key( return "thinking_" + "_".join(key_parts) if key_parts else None + # ========================================================================= + # PROJECT ID DISCOVERY + # ========================================================================= + + async def _discover_project_id(self, credential_path: str, litellm_params: Dict[str, Any]) -> str: + """ + Discovers the Google Cloud Project ID for Antigravity API. + + Priority: cache → env vars → persisted file → API discovery → GCP listing + """ + # Check cache + if credential_path in self.project_id_cache: + return self.project_id_cache[credential_path] + + # Check env vars + configured_project_id = ( + litellm_params.get("project_id") or + os.getenv("ANTIGRAVITY_PROJECT_ID") or + os.getenv("GOOGLE_CLOUD_PROJECT") + ) + if configured_project_id: + self.project_id_cache[credential_path] = configured_project_id + return configured_project_id + + # Try persisted file + try: + with open(credential_path, 'r') as f: + creds = json.load(f) + persisted = creds.get("_proxy_metadata", {}).get("project_id") + if persisted: + self.project_id_cache[credential_path] = persisted + return persisted + except: + pass + + # API discovery + access_token = await self.get_valid_token(credential_path) + headers = {'Authorization': f'Bearer {access_token}', 'Content-Type': 'application/json'} + + async with httpx.AsyncClient() as client: + try: + response = await client.post( + "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist", + headers=headers, + json={"cloudaicompanionProject": None, "metadata": {"ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI"}}, + timeout=20 + ) + response.raise_for_status() + server_project = response.json().get('cloudaicompanionProject') + if server_project: + self.project_id_cache[credential_path] = server_project + await self._persist_project_id(credential_path, server_project) + return server_project + except: + pass + + # GCP listing fallback + try: + async with httpx.AsyncClient() as client: + response = await client.get("https://cloudresourcemanager.googleapis.com/v1/projects", headers=headers, timeout=20) + response.raise_for_status() + active_projects = [p for p in response.json().get('projects', []) if p.get('lifecycleState') == 'ACTIVE'] + if active_projects: + project_id = active_projects[0]['projectId'] + self.project_id_cache[credential_path] = project_id + await self._persist_project_id(credential_path, project_id) + return project_id + except: + pass + + raise ValueError("Could not discover Google Cloud project ID for Antigravity. Set ANTIGRAVITY_PROJECT_ID or GOOGLE_CLOUD_PROJECT environment variable.") + + async def _persist_project_id(self, credential_path: str, project_id: str): + """Persist project ID to credential file.""" + try: + with open(credential_path, 'r') as f: + creds = json.load(f) + if "_proxy_metadata" not in creds: + creds["_proxy_metadata"] = {} + creds["_proxy_metadata"]["project_id"] = project_id + await self._save_credentials(credential_path, creds) + except: + pass + # ========================================================================= # THINKING MODE SANITIZATION # ========================================================================= @@ -1588,6 +1673,7 @@ def _transform_to_antigravity_format( self, gemini_payload: Dict[str, Any], model: str, + project_id: str, max_tokens: Optional[int] = None, reasoning_effort: Optional[str] = None, tool_choice: Optional[Union[str, Dict[str, Any]]] = None @@ -1620,7 +1706,7 @@ def _transform_to_antigravity_format( # Wrap in Antigravity envelope antigravity_payload = { - "project": _generate_project_id(), + "project": project_id, # Will be passed as parameter "userAgent": "antigravity", "requestId": _generate_request_id(), "model": internal_model, @@ -2158,8 +2244,12 @@ async def acompletion( self._claude_description_prompt ) - # Transform to Antigravity format - payload = self._transform_to_antigravity_format(gemini_payload, model, max_tokens, reasoning_effort, tool_choice) + # Discover real project ID + litellm_params = kwargs.get("litellm_params", {}) or {} + project_id = await self._discover_project_id(credential_path, litellm_params) + + # Transform to Antigravity format with real project ID + payload = self._transform_to_antigravity_format(gemini_payload, model, project_id, max_tokens, reasoning_effort, tool_choice) file_logger.log_request(payload) # Make API call From d2adf05133cda8e779adb98c2686dba0f5492b09 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Wed, 3 Dec 2025 23:57:41 +0100 Subject: [PATCH 062/221] =?UTF-8?q?feat(provider):=20=E2=9C=A8=20implement?= =?UTF-8?q?=20Google=20Cloud=20onboarding=20flow=20with=20automatic=20proj?= =?UTF-8?q?ect=20discovery?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces a comprehensive onboarding and project discovery system for the Antigravity provider, following the official Gemini CLI discovery flow adapted for API usage. Key changes: - Implement multi-stage project discovery with in-memory caching, environment variable overrides, persisted credential metadata, and API-based discovery - Add automatic user onboarding via the Code Assist `onboardUser` endpoint with long-running operation (LRO) polling support - Introduce tier-aware project management that distinguishes between free-tier (server-managed projects) and paid-tier (user-defined projects) workflows - Add comprehensive debug logging throughout the discovery process to aid troubleshooting - Implement metadata persistence for both project ID and tier information to speed up future startups - Add asyncio import for LRO polling delays - Enhance error messages with actionable guidance for common failure scenarios (403 Forbidden, missing API enablement, missing projects) - Update `_discover_project_id` signature to accept access_token as parameter, eliminating redundant token fetching - Fix `count_tokens` method to use discovered project_id instead of hardcoded generation - Add `project_tier_cache` dictionary for debugging and consistency with Gemini CLI behavior - Skip file-based persistence for environment-based credentials (env:// paths) The discovery flow prioritizes: 1. In-memory cache (fastest) 2. Configured project ID override (env vars or litellm_params) 3. Persisted metadata from credential file 4. Code Assist loadCodeAssist endpoint (checks existing session) 5. Automatic onboarding for new users (creates free-tier session) 6. GCP Resource Manager project listing (last resort fallback) This implementation ensures seamless first-run experience while maintaining compatibility with both free and paid tier users. --- .../providers/antigravity_provider.py | 421 +++++++++++++++--- 1 file changed, 364 insertions(+), 57 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index b3d51d8a..4adb1114 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -16,6 +16,7 @@ from __future__ import annotations +import asyncio import copy import hashlib import json @@ -449,6 +450,7 @@ def __init__(self): super().__init__() self.model_definitions = ModelDefinitions() self.project_id_cache: Dict[str, str] = {} # Cache project ID per credential path + self.project_tier_cache: Dict[str, str] = {} # Cache project tier per credential path (for debugging) # Base URL management self._base_url_index = 0 @@ -592,85 +594,385 @@ def _generate_thinking_cache_key( # PROJECT ID DISCOVERY # ========================================================================= - async def _discover_project_id(self, credential_path: str, litellm_params: Dict[str, Any]) -> str: + async def _discover_project_id(self, credential_path: str, access_token: str, litellm_params: Dict[str, Any]) -> str: """ - Discovers the Google Cloud Project ID for Antigravity API. - - Priority: cache → env vars → persisted file → API discovery → GCP listing + Discovers the Google Cloud Project ID, with caching and onboarding for new accounts. + + This follows the official Gemini CLI discovery flow adapted for Antigravity: + 1. Check in-memory cache + 2. Check configured project_id override (litellm_params or env var) + 3. Check persisted project_id in credential file + 4. Call loadCodeAssist to check if user is already known (has currentTier) + - If currentTier exists AND cloudaicompanionProject returned: use server's project + - If no currentTier: user needs onboarding + 5. Onboard user (FREE tier: pass cloudaicompanionProject=None for server-managed) + 6. Fallback to GCP Resource Manager project listing + + Note: Unlike GeminiCli, Antigravity doesn't use tier-based credential prioritization, + but we still cache tier info for debugging and consistency. """ - # Check cache + lib_logger.debug(f"Starting Antigravity project discovery for credential: {credential_path}") + + # Check in-memory cache first if credential_path in self.project_id_cache: - return self.project_id_cache[credential_path] - - # Check env vars + cached_project = self.project_id_cache[credential_path] + lib_logger.debug(f"Using cached project ID: {cached_project}") + return cached_project + + # Check for configured project ID override (from litellm_params or env var) configured_project_id = ( - litellm_params.get("project_id") or - os.getenv("ANTIGRAVITY_PROJECT_ID") or + litellm_params.get("project_id") or + os.getenv("ANTIGRAVITY_PROJECT_ID") or os.getenv("GOOGLE_CLOUD_PROJECT") ) if configured_project_id: - self.project_id_cache[credential_path] = configured_project_id - return configured_project_id - - # Try persisted file - try: - with open(credential_path, 'r') as f: - creds = json.load(f) - persisted = creds.get("_proxy_metadata", {}).get("project_id") - if persisted: - self.project_id_cache[credential_path] = persisted - return persisted - except: - pass - - # API discovery - access_token = await self.get_valid_token(credential_path) + lib_logger.debug(f"Found configured project_id override: {configured_project_id}") + + # Load credentials from file to check for persisted project_id and tier + # Skip for env:// paths (environment-based credentials don't persist to files) + credential_index = self._parse_env_credential_path(credential_path) + if credential_index is None: + # Only try to load from file if it's not an env:// path + try: + with open(credential_path, 'r') as f: + creds = json.load(f) + + metadata = creds.get("_proxy_metadata", {}) + persisted_project_id = metadata.get("project_id") + persisted_tier = metadata.get("tier") + + if persisted_project_id: + lib_logger.info(f"Loaded persisted project ID from credential file: {persisted_project_id}") + self.project_id_cache[credential_path] = persisted_project_id + + # Also load tier if available (for debugging/logging purposes) + if persisted_tier: + self.project_tier_cache[credential_path] = persisted_tier + lib_logger.debug(f"Loaded persisted tier: {persisted_tier}") + + return persisted_project_id + except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: + lib_logger.debug(f"Could not load persisted project ID from file: {e}") + + lib_logger.debug("No cached or configured project ID found, initiating discovery...") headers = {'Authorization': f'Bearer {access_token}', 'Content-Type': 'application/json'} - + + discovered_project_id = None + discovered_tier = None + + # Use production endpoint for loadCodeAssist (more reliable than sandbox URLs) + code_assist_endpoint = "https://cloudcode-pa.googleapis.com/v1internal" + async with httpx.AsyncClient() as client: + # 1. Try discovery endpoint with loadCodeAssist + lib_logger.debug("Attempting project discovery via Code Assist loadCodeAssist endpoint...") try: - response = await client.post( - "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist", - headers=headers, - json={"cloudaicompanionProject": None, "metadata": {"ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI"}}, - timeout=20 - ) + # Build metadata - include duetProject only if we have a configured project + core_client_metadata = { + "ideType": "IDE_UNSPECIFIED", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + } + if configured_project_id: + core_client_metadata["duetProject"] = configured_project_id + + # Build load request - pass configured_project_id if available, otherwise None + load_request = { + "cloudaicompanionProject": configured_project_id, # Can be None + "metadata": core_client_metadata, + } + + lib_logger.debug(f"Sending loadCodeAssist request with cloudaicompanionProject={configured_project_id}") + response = await client.post(f"{code_assist_endpoint}:loadCodeAssist", headers=headers, json=load_request, timeout=20) response.raise_for_status() - server_project = response.json().get('cloudaicompanionProject') - if server_project: - self.project_id_cache[credential_path] = server_project - await self._persist_project_id(credential_path, server_project) - return server_project - except: - pass - - # GCP listing fallback + data = response.json() + + # Log full response for debugging + lib_logger.debug(f"loadCodeAssist full response keys: {list(data.keys())}") + + # Extract tier information + allowed_tiers = data.get('allowedTiers', []) + current_tier = data.get('currentTier') + + lib_logger.debug(f"=== Tier Information ===") + lib_logger.debug(f"currentTier: {current_tier}") + lib_logger.debug(f"allowedTiers count: {len(allowed_tiers)}") + for i, tier in enumerate(allowed_tiers): + tier_id = tier.get('id', 'unknown') + is_default = tier.get('isDefault', False) + user_defined = tier.get('userDefinedCloudaicompanionProject', False) + lib_logger.debug(f" Tier {i+1}: id={tier_id}, isDefault={is_default}, userDefinedProject={user_defined}") + lib_logger.debug(f"========================") + + # Determine the current tier ID + current_tier_id = None + if current_tier: + current_tier_id = current_tier.get('id') + lib_logger.debug(f"User has currentTier: {current_tier_id}") + + # Check if user is already known to server (has currentTier) + if current_tier_id: + # User is already onboarded - check for project from server + server_project = data.get('cloudaicompanionProject') + + # Check if this tier requires user-defined project (paid tiers) + requires_user_project = any( + t.get('id') == current_tier_id and t.get('userDefinedCloudaicompanionProject', False) + for t in allowed_tiers + ) + is_free_tier = current_tier_id == 'free-tier' + + if server_project: + # Server returned a project - use it (server wins) + project_id = server_project + lib_logger.debug(f"Server returned project: {project_id}") + elif configured_project_id: + # No server project but we have configured one - use it + project_id = configured_project_id + lib_logger.debug(f"No server project, using configured: {project_id}") + elif is_free_tier: + # Free tier user without server project - try onboarding + lib_logger.debug("Free tier user with currentTier but no project - will try onboarding") + project_id = None + elif requires_user_project: + # Paid tier requires a project ID to be set + raise ValueError( + f"Paid tier '{current_tier_id}' requires setting ANTIGRAVITY_PROJECT_ID environment variable." + ) + else: + # Unknown tier without project - proceed to onboarding + lib_logger.warning(f"Tier '{current_tier_id}' has no project and none configured - will try onboarding") + project_id = None + + if project_id: + # Cache tier info + self.project_tier_cache[credential_path] = current_tier_id + discovered_tier = current_tier_id + + # Log appropriately based on tier + is_paid = current_tier_id and current_tier_id not in ['free-tier', 'legacy-tier', 'unknown'] + if is_paid: + lib_logger.info(f"Using Antigravity paid tier '{current_tier_id}' with project: {project_id}") + else: + lib_logger.info(f"Discovered Antigravity project ID via loadCodeAssist: {project_id}") + + self.project_id_cache[credential_path] = project_id + discovered_project_id = project_id + + # Persist to credential file + await self._persist_project_metadata(credential_path, project_id, discovered_tier) + + return project_id + + # 2. User needs onboarding - no currentTier or no project found + lib_logger.info("No existing Antigravity session found (no currentTier), attempting to onboard user...") + + # Determine which tier to onboard with + onboard_tier = None + for tier in allowed_tiers: + if tier.get('isDefault'): + onboard_tier = tier + break + + # Fallback to legacy tier if no default + if not onboard_tier and allowed_tiers: + for tier in allowed_tiers: + if tier.get('id') == 'legacy-tier': + onboard_tier = tier + break + if not onboard_tier: + onboard_tier = allowed_tiers[0] + + if not onboard_tier: + raise ValueError("No onboarding tiers available from server") + + tier_id = onboard_tier.get('id', 'free-tier') + requires_user_project = onboard_tier.get('userDefinedCloudaicompanionProject', False) + + lib_logger.debug(f"Onboarding with tier: {tier_id}, requiresUserProject: {requires_user_project}") + + # Build onboard request based on tier type + # FREE tier: cloudaicompanionProject = None (server-managed) + # PAID tier: cloudaicompanionProject = configured_project_id + is_free_tier = tier_id == 'free-tier' + + if is_free_tier: + # Free tier uses server-managed project + onboard_request = { + "tierId": tier_id, + "cloudaicompanionProject": None, # Server will create/manage + "metadata": core_client_metadata, + } + lib_logger.debug("Free tier onboarding: using server-managed project") + else: + # Paid/legacy tier requires user-provided project + if not configured_project_id and requires_user_project: + raise ValueError( + f"Tier '{tier_id}' requires setting ANTIGRAVITY_PROJECT_ID environment variable." + ) + onboard_request = { + "tierId": tier_id, + "cloudaicompanionProject": configured_project_id, + "metadata": { + **core_client_metadata, + "duetProject": configured_project_id, + } if configured_project_id else core_client_metadata, + } + lib_logger.debug(f"Paid tier onboarding: using project {configured_project_id}") + + lib_logger.debug("Initiating onboardUser request...") + lro_response = await client.post(f"{code_assist_endpoint}:onboardUser", headers=headers, json=onboard_request, timeout=30) + lro_response.raise_for_status() + lro_data = lro_response.json() + lib_logger.debug(f"Initial onboarding response: done={lro_data.get('done')}") + + # Poll for onboarding completion (up to 5 minutes) + for i in range(150): # 150 × 2s = 5 minutes + if lro_data.get('done'): + lib_logger.debug(f"Onboarding completed after {i} polling attempts") + break + await asyncio.sleep(2) + if (i + 1) % 15 == 0: # Log every 30 seconds + lib_logger.info(f"Still waiting for onboarding completion... ({(i+1)*2}s elapsed)") + lib_logger.debug(f"Polling onboarding status... (Attempt {i+1}/150)") + lro_response = await client.post(f"{code_assist_endpoint}:onboardUser", headers=headers, json=onboard_request, timeout=30) + lro_response.raise_for_status() + lro_data = lro_response.json() + + if not lro_data.get('done'): + lib_logger.error("Onboarding process timed out after 5 minutes") + raise ValueError("Onboarding process timed out after 5 minutes. Please try again or contact support.") + + # Extract project ID from LRO response + # Note: onboardUser returns response.cloudaicompanionProject as an object with .id + lro_response_data = lro_data.get('response', {}) + lro_project_obj = lro_response_data.get('cloudaicompanionProject', {}) + project_id = lro_project_obj.get('id') if isinstance(lro_project_obj, dict) else None + + # Fallback to configured project if LRO didn't return one + if not project_id and configured_project_id: + project_id = configured_project_id + lib_logger.debug(f"LRO didn't return project, using configured: {project_id}") + + if not project_id: + lib_logger.error("Onboarding completed but no project ID in response and none configured") + raise ValueError( + "Onboarding completed, but no project ID was returned. " + "For paid tiers, set ANTIGRAVITY_PROJECT_ID environment variable." + ) + + lib_logger.debug(f"Successfully extracted project ID from onboarding response: {project_id}") + + # Cache tier info + self.project_tier_cache[credential_path] = tier_id + discovered_tier = tier_id + lib_logger.debug(f"Cached tier information: {tier_id}") + + # Log concise message based on tier + is_paid = tier_id and tier_id not in ['free-tier', 'legacy-tier'] + if is_paid: + lib_logger.info(f"Using Antigravity paid tier '{tier_id}' with project: {project_id}") + else: + lib_logger.info(f"Successfully onboarded user and discovered project ID: {project_id}") + + self.project_id_cache[credential_path] = project_id + discovered_project_id = project_id + + # Persist to credential file + await self._persist_project_metadata(credential_path, project_id, discovered_tier) + + return project_id + + except httpx.HTTPStatusError as e: + error_body = "" + try: + error_body = e.response.text + except Exception: + pass + if e.response.status_code == 403: + lib_logger.error(f"Antigravity Code Assist API access denied (403). Response: {error_body}") + lib_logger.error("Possible causes: 1) cloudaicompanion.googleapis.com API not enabled, 2) Wrong project ID for paid tier, 3) Account lacks permissions") + elif e.response.status_code == 404: + lib_logger.warning(f"Antigravity Code Assist endpoint not found (404). Falling back to project listing.") + elif e.response.status_code == 412: + # Precondition Failed - often means wrong project for free tier onboarding + lib_logger.error(f"Precondition failed (412): {error_body}. This may mean the project ID is incompatible with the selected tier.") + else: + lib_logger.warning(f"Antigravity onboarding/discovery failed with status {e.response.status_code}: {error_body}. Falling back to project listing.") + except httpx.RequestError as e: + lib_logger.warning(f"Antigravity onboarding/discovery network error: {e}. Falling back to project listing.") + + # 3. Fallback to listing all available GCP projects (last resort) + lib_logger.debug("Attempting to discover project via GCP Resource Manager API...") try: async with httpx.AsyncClient() as client: + lib_logger.debug("Querying Cloud Resource Manager for available projects...") response = await client.get("https://cloudresourcemanager.googleapis.com/v1/projects", headers=headers, timeout=20) response.raise_for_status() - active_projects = [p for p in response.json().get('projects', []) if p.get('lifecycleState') == 'ACTIVE'] - if active_projects: + projects = response.json().get('projects', []) + lib_logger.debug(f"Found {len(projects)} total projects") + active_projects = [p for p in projects if p.get('lifecycleState') == 'ACTIVE'] + lib_logger.debug(f"Found {len(active_projects)} active projects") + + if not projects: + lib_logger.error("No GCP projects found for this account. Please create a project in Google Cloud Console.") + elif not active_projects: + lib_logger.error("No active GCP projects found. Please activate a project in Google Cloud Console.") + else: project_id = active_projects[0]['projectId'] + lib_logger.info(f"Discovered Antigravity project ID from active projects list: {project_id}") + lib_logger.debug(f"Selected first active project: {project_id} (out of {len(active_projects)} active projects)") self.project_id_cache[credential_path] = project_id - await self._persist_project_id(credential_path, project_id) + discovered_project_id = project_id + + # Persist to credential file (no tier info from resource manager) + await self._persist_project_metadata(credential_path, project_id, None) + return project_id - except: - pass - - raise ValueError("Could not discover Google Cloud project ID for Antigravity. Set ANTIGRAVITY_PROJECT_ID or GOOGLE_CLOUD_PROJECT environment variable.") + except httpx.HTTPStatusError as e: + if e.response.status_code == 403: + lib_logger.error("Failed to list GCP projects due to a 403 Forbidden error. The Cloud Resource Manager API may not be enabled, or your account lacks the 'resourcemanager.projects.list' permission.") + else: + lib_logger.error(f"Failed to list GCP projects with status {e.response.status_code}: {e}") + except httpx.RequestError as e: + lib_logger.error(f"Network error while listing GCP projects: {e}") + + raise ValueError( + "Could not auto-discover Antigravity project ID. Possible causes:\n" + " 1. The cloudaicompanion.googleapis.com API is not enabled (enable it in Google Cloud Console)\n" + " 2. No active GCP projects exist for this account (create one in Google Cloud Console)\n" + " 3. Account lacks necessary permissions\n" + "To manually specify a project, set ANTIGRAVITY_PROJECT_ID in your .env file." + ) - async def _persist_project_id(self, credential_path: str, project_id: str): - """Persist project ID to credential file.""" + async def _persist_project_metadata(self, credential_path: str, project_id: str, tier: Optional[str]): + """Persists project ID and tier to the credential file for faster future startups.""" + # Skip persistence for env:// paths (environment-based credentials) + credential_index = self._parse_env_credential_path(credential_path) + if credential_index is not None: + lib_logger.debug(f"Skipping project metadata persistence for env:// credential path: {credential_path}") + return + try: + # Load current credentials with open(credential_path, 'r') as f: creds = json.load(f) + + # Update metadata if "_proxy_metadata" not in creds: creds["_proxy_metadata"] = {} + creds["_proxy_metadata"]["project_id"] = project_id + if tier: + creds["_proxy_metadata"]["tier"] = tier + + # Save back using the existing save method (handles atomic writes and permissions) await self._save_credentials(credential_path, creds) - except: - pass + + lib_logger.debug(f"Persisted project_id and tier to credential file: {credential_path}") + except Exception as e: + lib_logger.warning(f"Failed to persist project metadata to credential file: {e}") + # Non-fatal - just means slower startup next time # ========================================================================= # THINKING MODE SANITIZATION @@ -2244,16 +2546,18 @@ async def acompletion( self._claude_description_prompt ) + # Get access token first (needed for project discovery) + token = await self.get_valid_token(credential_path) + # Discover real project ID litellm_params = kwargs.get("litellm_params", {}) or {} - project_id = await self._discover_project_id(credential_path, litellm_params) + project_id = await self._discover_project_id(credential_path, token, litellm_params) # Transform to Antigravity format with real project ID payload = self._transform_to_antigravity_format(gemini_payload, model, project_id, max_tokens, reasoning_effort, tool_choice) file_logger.log_request(payload) # Make API call - token = await self.get_valid_token(credential_path) base_url = self._get_base_url() endpoint = ":streamGenerateContent" if stream else ":generateContent" url = f"{base_url}{endpoint}" @@ -2409,13 +2713,16 @@ async def count_tokens( model: str, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, - _litellm_params: Optional[Dict[str, Any]] = None + litellm_params: Optional[Dict[str, Any]] = None ) -> Dict[str, int]: """Count tokens for the given prompt using Antigravity :countTokens endpoint.""" try: token = await self.get_valid_token(credential_path) internal_model = self._alias_to_internal(model) + # Discover project ID + project_id = await self._discover_project_id(credential_path, token, litellm_params or {}) + system_instruction, contents = self._transform_messages(messages, internal_model) contents = self._fix_tool_response_grouping(contents) @@ -2428,7 +2735,7 @@ async def count_tokens( gemini_payload["tools"] = gemini_tools antigravity_payload = { - "project": _generate_project_id(), + "project": project_id, "userAgent": "antigravity", "requestId": _generate_request_id(), "model": internal_model, From a1cc8752aeb76a1568b7898518eee0ca30553287 Mon Sep 17 00:00:00 2001 From: "mirrobot-agent[bot]" <2140342+mirrobot-agent@users.noreply.github.com> Date: Thu, 4 Dec 2025 00:54:42 +0000 Subject: [PATCH 063/221] fix: improve error handling implementation based on code review - Fix credential counting to track unique credentials (RequestErrorAccumulator) - Move import os to module level in mask_credential function - Fix status code check to use explicit 'is not None' comparison - Improve context window error detection with more specific patterns - Correct comment about server error classification - Remove redundant '1 *' in exponential backoff calculations - Add warning log for unreachable None return path - Remove redundant error_accumulator.model/provider assignments - Remove access to private _content attribute in failure_logger - Add circular reference detection in error chain loop - Reorder error recording to occur after should_rotate_on_error check These changes address issues identified in both mirrobot-agent and GitHub Copilot code reviews. --- src/rotator_library/client.py | 398 ++++++++++++++++---------- src/rotator_library/error_handler.py | 208 ++++++++------ src/rotator_library/failure_logger.py | 83 +++--- 3 files changed, 426 insertions(+), 263 deletions(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index ef322e6c..d603d463 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -71,7 +71,7 @@ def __init__( ): """ Initialize the RotatingClient with intelligent credential rotation. - + Args: api_keys: Dictionary mapping provider names to lists of API keys oauth_credentials: Dictionary mapping provider names to OAuth credential paths @@ -140,8 +140,7 @@ def __init__( self.global_timeout = global_timeout self.abort_on_callback_error = abort_on_callback_error self.usage_manager = UsageManager( - file_path=usage_file_path, - rotation_tolerance=rotation_tolerance + file_path=usage_file_path, rotation_tolerance=rotation_tolerance ) self._model_list_cache = {} self._provider_plugins = PROVIDER_PLUGINS @@ -160,7 +159,9 @@ def __init__( # Validate all values are >= 1 for provider, max_val in self.max_concurrent_requests_per_key.items(): if max_val < 1: - lib_logger.warning(f"Invalid max_concurrent for '{provider}': {max_val}. Setting to 1.") + lib_logger.warning( + f"Invalid max_concurrent for '{provider}': {max_val}. Setting to 1." + ) self.max_concurrent_requests_per_key[provider] = 1 def _is_model_ignored(self, provider: str, model_id: str) -> bool: @@ -368,7 +369,9 @@ def _convert_model_params_for_litellm(self, **kwargs) -> Dict[str, Any]: return kwargs - def _apply_default_safety_settings(self, litellm_kwargs: Dict[str, Any], provider: str): + def _apply_default_safety_settings( + self, litellm_kwargs: Dict[str, Any], provider: str + ): """ Ensure default Gemini safety settings are present when calling the Gemini provider. This will not override any explicit settings provided by the request. It accepts @@ -397,22 +400,33 @@ def _apply_default_safety_settings(self, litellm_kwargs: Dict[str, Any], provide ] # If generic form is present, ensure missing generic keys are filled in - if "safety_settings" in litellm_kwargs and isinstance(litellm_kwargs["safety_settings"], dict): + if "safety_settings" in litellm_kwargs and isinstance( + litellm_kwargs["safety_settings"], dict + ): for k, v in default_generic.items(): if k not in litellm_kwargs["safety_settings"]: litellm_kwargs["safety_settings"][k] = v return # If Gemini form is present, ensure missing gemini categories are appended - if "safetySettings" in litellm_kwargs and isinstance(litellm_kwargs["safetySettings"], list): - present = {item.get("category") for item in litellm_kwargs["safetySettings"] if isinstance(item, dict)} + if "safetySettings" in litellm_kwargs and isinstance( + litellm_kwargs["safetySettings"], list + ): + present = { + item.get("category") + for item in litellm_kwargs["safetySettings"] + if isinstance(item, dict) + } for d in default_gemini: if d["category"] not in present: litellm_kwargs["safetySettings"].append(d) return # Neither present: set generic defaults so provider conversion will translate them - if "safety_settings" not in litellm_kwargs and "safetySettings" not in litellm_kwargs: + if ( + "safety_settings" not in litellm_kwargs + and "safetySettings" not in litellm_kwargs + ): litellm_kwargs["safety_settings"] = default_generic.copy() def get_oauth_credentials(self) -> Dict[str, List[str]]: @@ -430,10 +444,10 @@ def _get_provider_instance(self, provider_name: str): """ Lazily initializes and returns a provider instance. Only initializes providers that have configured credentials. - + Args: provider_name: The name of the provider to get an instance for. - + Returns: Provider instance if credentials exist, None otherwise. """ @@ -443,7 +457,7 @@ def _get_provider_instance(self, provider_name: str): f"Skipping provider '{provider_name}' initialization: no credentials configured" ) return None - + if provider_name not in self._provider_instances: if provider_name in self._provider_plugins: self._provider_instances[provider_name] = self._provider_plugins[ @@ -465,46 +479,47 @@ def _get_provider_instance(self, provider_name: str): def _resolve_model_id(self, model: str, provider: str) -> str: """ Resolves the actual model ID to send to the provider. - + For custom models with name/ID mappings, returns the ID. Otherwise, returns the model name unchanged. - + Args: model: Full model string with provider (e.g., "iflow/DS-v3.2") provider: Provider name (e.g., "iflow") - + Returns: Full model string with ID (e.g., "iflow/deepseek-v3.2") """ # Extract model name from "provider/model_name" format - model_name = model.split('/')[-1] if '/' in model else model - + model_name = model.split("/")[-1] if "/" in model else model + # Try to get provider instance to check for model definitions provider_plugin = self._get_provider_instance(provider) - + # Check if provider has model definitions - if provider_plugin and hasattr(provider_plugin, 'model_definitions'): - model_id = provider_plugin.model_definitions.get_model_id(provider, model_name) + if provider_plugin and hasattr(provider_plugin, "model_definitions"): + model_id = provider_plugin.model_definitions.get_model_id( + provider, model_name + ) if model_id and model_id != model_name: # Return with provider prefix return f"{provider}/{model_id}" - + # Fallback: use client's own model definitions model_id = self.model_definitions.get_model_id(provider, model_name) if model_id and model_id != model_name: return f"{provider}/{model_id}" - + # No conversion needed, return original return model - async def _safe_streaming_wrapper( self, stream: Any, key: str, model: str, request: Optional[Any] = None ) -> AsyncGenerator[Any, None]: """ A hybrid wrapper for streaming that buffers fragmented JSON, handles client disconnections gracefully, and distinguishes between content and streamed errors. - + FINISH_REASON HANDLING: Providers just translate chunks - this wrapper handles ALL finish_reason logic: 1. Strip finish_reason from intermediate chunks (litellm defaults to "stop") @@ -541,7 +556,7 @@ async def _safe_streaming_wrapper( chunk_dict = chunk.model_dump() else: chunk_dict = chunk - + # === FINISH_REASON LOGIC === # Providers send raw chunks without finish_reason logic. # This wrapper determines finish_reason based on accumulated state. @@ -549,19 +564,19 @@ async def _safe_streaming_wrapper( choice = chunk_dict["choices"][0] delta = choice.get("delta", {}) usage = chunk_dict.get("usage", {}) - + # Track tool_calls across ALL chunks - if we ever see one, finish_reason must be tool_calls if delta.get("tool_calls"): has_tool_calls = True accumulated_finish_reason = "tool_calls" - + # Detect final chunk: has usage with completion_tokens > 0 has_completion_tokens = ( - usage and - isinstance(usage, dict) and - usage.get("completion_tokens", 0) > 0 + usage + and isinstance(usage, dict) + and usage.get("completion_tokens", 0) > 0 ) - + if has_completion_tokens: # FINAL CHUNK: Determine correct finish_reason if has_tool_calls: @@ -577,7 +592,7 @@ async def _safe_streaming_wrapper( # INTERMEDIATE CHUNK: Never emit finish_reason # (litellm.ModelResponse defaults to "stop" which is wrong) choice["finish_reason"] = None - + yield f"data: {json.dumps(chunk_dict)}\n\n" if hasattr(chunk, "usage") and chunk.usage: @@ -726,12 +741,13 @@ async def _execute_with_retry( # multiple keys have the same usage stats. credentials_for_provider = list(self.all_credentials[provider]) random.shuffle(credentials_for_provider) - + # Filter out credentials that are unavailable (queued for re-auth) provider_plugin = self._get_provider_instance(provider) - if provider_plugin and hasattr(provider_plugin, 'is_credential_available'): + if provider_plugin and hasattr(provider_plugin, "is_credential_available"): available_creds = [ - cred for cred in credentials_for_provider + cred + for cred in credentials_for_provider if provider_plugin.is_credential_available(cred) ] if available_creds: @@ -744,7 +760,7 @@ async def _execute_with_retry( kwargs = self._convert_model_params(**kwargs) # The main rotation loop. It continues as long as there are untried credentials and the global deadline has not been exceeded. - + # Resolve model ID early, before any credential operations # This ensures consistent model ID usage for acquisition, release, and tracking resolved_model = self._resolve_model_id(model, provider) @@ -752,10 +768,10 @@ async def _execute_with_retry( lib_logger.info(f"Resolved model '{model}' to '{resolved_model}'") model = resolved_model kwargs["model"] = model # Ensure kwargs has the resolved model for litellm - + # [NEW] Filter by model tier requirement and build priority map credential_priorities = None - if provider_plugin and hasattr(provider_plugin, 'get_model_tier_requirement'): + if provider_plugin and hasattr(provider_plugin, "get_model_tier_requirement"): required_tier = provider_plugin.get_model_tier_requirement(model) if required_tier is not None: # Filter OUT only credentials we KNOW are too low priority @@ -763,9 +779,9 @@ async def _execute_with_retry( incompatible_creds = [] compatible_creds = [] unknown_creds = [] - + for cred in credentials_for_provider: - if hasattr(provider_plugin, 'get_credential_priority'): + if hasattr(provider_plugin, "get_credential_priority"): priority = provider_plugin.get_credential_priority(cred) if priority is None: # Unknown priority - keep it, will be discovered on first use @@ -779,7 +795,7 @@ async def _execute_with_retry( else: # Provider doesn't support priorities - keep all unknown_creds.append(cred) - + # If we have any known-compatible or unknown credentials, use them tier_compatible_creds = compatible_creds + unknown_creds if tier_compatible_creds: @@ -806,18 +822,18 @@ async def _execute_with_retry( f"but all {len(incompatible_creds)} known credentials have priority > {required_tier}. " f"Request will likely fail." ) - + # Build priority map for usage_manager - if provider_plugin and hasattr(provider_plugin, 'get_credential_priority'): + if provider_plugin and hasattr(provider_plugin, "get_credential_priority"): credential_priorities = {} for cred in credentials_for_provider: priority = provider_plugin.get_credential_priority(cred) if priority is not None: credential_priorities[cred] = priority - + if credential_priorities: lib_logger.debug( - f"Credential priorities for {provider}: {', '.join(f'P{p}={len([c for c in credentials_for_provider if credential_priorities.get(c)==p])}' for p in sorted(set(credential_priorities.values())))}" + f"Credential priorities for {provider}: {', '.join(f'P{p}={len([c for c in credentials_for_provider if credential_priorities.get(c) == p])}' for p in sorted(set(credential_priorities.values())))}" ) # Initialize error accumulator for tracking errors across credential rotation @@ -861,9 +877,11 @@ async def _execute_with_retry( ) max_concurrent = self.max_concurrent_requests_per_key.get(provider, 1) current_cred = await self.usage_manager.acquire_key( - available_keys=creds_to_try, model=model, deadline=deadline, + available_keys=creds_to_try, + model=model, + deadline=deadline, max_concurrent=max_concurrent, - credential_priorities=credential_priorities + credential_priorities=credential_priorities, ) key_acquired = True tried_creds.add(current_cred) @@ -946,10 +964,14 @@ async def _execute_with_retry( if provider_instance: # Ensure default Gemini safety settings are present (without overriding request) try: - self._apply_default_safety_settings(litellm_kwargs, provider) + self._apply_default_safety_settings( + litellm_kwargs, provider + ) except Exception: # If anything goes wrong here, avoid breaking the request flow. - lib_logger.debug("Could not apply default safety settings; continuing.") + lib_logger.debug( + "Could not apply default safety settings; continuing." + ) if "safety_settings" in litellm_kwargs: converted_settings = ( @@ -1032,9 +1054,11 @@ async def _execute_with_retry( # Extract a clean error message for the user-facing log error_message = str(e).split("\n")[0] - + # Record in accumulator for client reporting - error_accumulator.record_error(current_cred, classified_error, error_message) + error_accumulator.record_error( + current_cred, classified_error, error_message + ) lib_logger.info( f"Key {mask_credential(current_cred)} hit rate limit for {model}. Rotating key." @@ -1068,16 +1092,20 @@ async def _execute_with_retry( ) classified_error = classify_error(e) error_message = str(e).split("\n")[0] - + # Provider-level error: don't increment consecutive failures await self.usage_manager.record_failure( - current_cred, model, classified_error, - increment_consecutive_failures=False + current_cred, + model, + classified_error, + increment_consecutive_failures=False, ) if attempt >= self.max_retries - 1: # Record in accumulator only on final failure for this key - error_accumulator.record_error(current_cred, classified_error, error_message) + error_accumulator.record_error( + current_cred, classified_error, error_message + ) lib_logger.warning( f"Key {mask_credential(current_cred)} failed after max retries due to server error. Rotating." ) @@ -1085,13 +1113,15 @@ async def _execute_with_retry( # For temporary errors, wait before retrying with the same key. wait_time = classified_error.retry_after or ( - 1 * (2**attempt) + 2**attempt ) + random.uniform(0, 1) remaining_budget = deadline - time.time() # If the required wait time exceeds the budget, don't wait; rotate to the next key immediately. if wait_time > remaining_budget: - error_accumulator.record_error(current_cred, classified_error, error_message) + error_accumulator.record_error( + current_cred, classified_error, error_message + ) lib_logger.warning( f"Retry wait ({wait_time:.2f}s) exceeds budget ({remaining_budget:.2f}s). Rotating key." ) @@ -1115,34 +1145,44 @@ async def _execute_with_retry( if request else {}, ) - + classified_error = classify_error(e) error_message = str(e).split("\n")[0] - - # Record in accumulator for client reporting - error_accumulator.record_error(current_cred, classified_error, error_message) - + lib_logger.warning( f"Key {mask_credential(current_cred)} HTTP {e.response.status_code} ({classified_error.error_type})." ) - + # Check if this error should trigger rotation if not should_rotate_on_error(classified_error): lib_logger.error( f"Non-recoverable error ({classified_error.error_type}). Failing request." ) raise last_exception - + + # Record in accumulator after confirming it's a rotatable error + error_accumulator.record_error( + current_cred, classified_error, error_message + ) + # Handle rate limits with cooldown - if classified_error.error_type in ["rate_limit", "quota_exceeded"]: + if classified_error.error_type in [ + "rate_limit", + "quota_exceeded", + ]: cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration ) - + # Check if we should retry same key (server errors with retries left) - if should_retry_same_key(classified_error) and attempt < self.max_retries - 1: - wait_time = classified_error.retry_after or (1 * (2**attempt)) + random.uniform(0, 1) + if ( + should_retry_same_key(classified_error) + and attempt < self.max_retries - 1 + ): + wait_time = classified_error.retry_after or ( + 2**attempt + ) + random.uniform(0, 1) remaining_budget = deadline - time.time() if wait_time <= remaining_budget: lib_logger.warning( @@ -1150,12 +1190,14 @@ async def _execute_with_retry( ) await asyncio.sleep(wait_time) continue - + # Record failure and rotate to next key await self.usage_manager.record_failure( current_cred, model, classified_error ) - lib_logger.info(f"Rotating to next key after {classified_error.error_type} error.") + lib_logger.info( + f"Rotating to next key after {classified_error.error_type} error." + ) break except Exception as e: @@ -1178,16 +1220,17 @@ async def _execute_with_retry( classified_error = classify_error(e) error_message = str(e).split("\n")[0] - - # Record in accumulator for client reporting - error_accumulator.record_error(current_cred, classified_error, error_message) - + lib_logger.warning( f"Key {mask_credential(current_cred)} {classified_error.error_type} (HTTP {classified_error.status_code})." ) - + # Handle rate limits with cooldown - if classified_error.status_code == 429 or classified_error.error_type in ["rate_limit", "quota_exceeded"]: + if ( + classified_error.status_code == 429 + or classified_error.error_type + in ["rate_limit", "quota_exceeded"] + ): cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration @@ -1200,6 +1243,11 @@ async def _execute_with_retry( ) raise last_exception + # Record in accumulator after confirming it's a rotatable error + error_accumulator.record_error( + current_cred, classified_error, error_message + ) + await self.usage_manager.record_failure( current_cred, model, classified_error ) @@ -1211,15 +1259,19 @@ async def _execute_with_retry( # Check if we exhausted all credentials or timed out if time.time() >= deadline: error_accumulator.timeout_occurred = True - + if error_accumulator.has_errors(): # Log concise summary for server logs lib_logger.error(error_accumulator.build_log_message()) - + # Return the structured error response for the client return error_accumulator.build_client_error_response() # Return None to indicate failure without error details (shouldn't normally happen) + lib_logger.warning( + "Unexpected state: request failed with no recorded errors. " + "This may indicate a logic error in error tracking." + ) return None async def _streaming_acompletion_with_retry( @@ -1235,12 +1287,13 @@ async def _streaming_acompletion_with_retry( # Create a mutable copy of the keys and shuffle it. credentials_for_provider = list(self.all_credentials[provider]) random.shuffle(credentials_for_provider) - + # Filter out credentials that are unavailable (queued for re-auth) provider_plugin = self._get_provider_instance(provider) - if provider_plugin and hasattr(provider_plugin, 'is_credential_available'): + if provider_plugin and hasattr(provider_plugin, "is_credential_available"): available_creds = [ - cred for cred in credentials_for_provider + cred + for cred in credentials_for_provider if provider_plugin.is_credential_available(cred) ] if available_creds: @@ -1262,10 +1315,10 @@ async def _streaming_acompletion_with_retry( lib_logger.info(f"Resolved model '{model}' to '{resolved_model}'") model = resolved_model kwargs["model"] = model # Ensure kwargs has the resolved model for litellm - + # [NEW] Filter by model tier requirement and build priority map credential_priorities = None - if provider_plugin and hasattr(provider_plugin, 'get_model_tier_requirement'): + if provider_plugin and hasattr(provider_plugin, "get_model_tier_requirement"): required_tier = provider_plugin.get_model_tier_requirement(model) if required_tier is not None: # Filter OUT only credentials we KNOW are too low priority @@ -1273,9 +1326,9 @@ async def _streaming_acompletion_with_retry( incompatible_creds = [] compatible_creds = [] unknown_creds = [] - + for cred in credentials_for_provider: - if hasattr(provider_plugin, 'get_credential_priority'): + if hasattr(provider_plugin, "get_credential_priority"): priority = provider_plugin.get_credential_priority(cred) if priority is None: # Unknown priority - keep it, will be discovered on first use @@ -1289,7 +1342,7 @@ async def _streaming_acompletion_with_retry( else: # Provider doesn't support priorities - keep all unknown_creds.append(cred) - + # If we have any known-compatible or unknown credentials, use them tier_compatible_creds = compatible_creds + unknown_creds if tier_compatible_creds: @@ -1316,18 +1369,18 @@ async def _streaming_acompletion_with_retry( f"but all {len(incompatible_creds)} known credentials have priority > {required_tier}. " f"Request will likely fail." ) - + # Build priority map for usage_manager - if provider_plugin and hasattr(provider_plugin, 'get_credential_priority'): + if provider_plugin and hasattr(provider_plugin, "get_credential_priority"): credential_priorities = {} for cred in credentials_for_provider: priority = provider_plugin.get_credential_priority(cred) if priority is not None: credential_priorities[cred] = priority - + if credential_priorities: lib_logger.debug( - f"Credential priorities for {provider}: {', '.join(f'P{p}={len([c for c in credentials_for_provider if credential_priorities.get(c)==p])}' for p in sorted(set(credential_priorities.values())))}" + f"Credential priorities for {provider}: {', '.join(f'P{p}={len([c for c in credentials_for_provider if credential_priorities.get(c) == p])}' for p in sorted(set(credential_priorities.values())))}" ) # Initialize error accumulator for tracking errors across credential rotation @@ -1370,11 +1423,15 @@ async def _streaming_acompletion_with_retry( lib_logger.info( f"Acquiring credential for model {model}. Tried credentials: {len(tried_creds)}/{len(credentials_for_provider)}" ) - max_concurrent = self.max_concurrent_requests_per_key.get(provider, 1) + max_concurrent = self.max_concurrent_requests_per_key.get( + provider, 1 + ) current_cred = await self.usage_manager.acquire_key( - available_keys=creds_to_try, model=model, deadline=deadline, + available_keys=creds_to_try, + model=model, + deadline=deadline, max_concurrent=max_concurrent, - credential_priorities=credential_priorities + credential_priorities=credential_priorities, ) key_acquired = True tried_creds.add(current_cred) @@ -1483,7 +1540,7 @@ async def _streaming_acompletion_with_retry( original_exc = getattr(e, "data", e) classified_error = classify_error(original_exc) error_message = str(original_exc).split("\n")[0] - + log_failure( api_key=current_cred, model=model, @@ -1493,24 +1550,31 @@ async def _streaming_acompletion_with_retry( if request else {}, ) - + # Record in accumulator for client reporting - error_accumulator.record_error(current_cred, classified_error, error_message) - + error_accumulator.record_error( + current_cred, classified_error, error_message + ) + # Check if this error should trigger rotation if not should_rotate_on_error(classified_error): lib_logger.error( f"Non-recoverable error ({classified_error.error_type}) during custom stream. Failing." ) raise last_exception - + # Handle rate limits with cooldown - if classified_error.error_type in ["rate_limit", "quota_exceeded"]: - cooldown_duration = classified_error.retry_after or 60 + if classified_error.error_type in [ + "rate_limit", + "quota_exceeded", + ]: + cooldown_duration = ( + classified_error.retry_after or 60 + ) await self.cooldown_manager.start_cooldown( provider, cooldown_duration ) - + await self.usage_manager.record_failure( current_cred, model, classified_error ) @@ -1536,26 +1600,32 @@ async def _streaming_acompletion_with_retry( ) classified_error = classify_error(e) error_message = str(e).split("\n")[0] - + # Provider-level error: don't increment consecutive failures await self.usage_manager.record_failure( - current_cred, model, classified_error, - increment_consecutive_failures=False + current_cred, + model, + classified_error, + increment_consecutive_failures=False, ) if attempt >= self.max_retries - 1: - error_accumulator.record_error(current_cred, classified_error, error_message) + error_accumulator.record_error( + current_cred, classified_error, error_message + ) lib_logger.warning( f"Cred {mask_credential(current_cred)} failed after max retries. Rotating." ) break wait_time = classified_error.retry_after or ( - 1 * (2**attempt) + 2**attempt ) + random.uniform(0, 1) remaining_budget = deadline - time.time() if wait_time > remaining_budget: - error_accumulator.record_error(current_cred, classified_error, error_message) + error_accumulator.record_error( + current_cred, classified_error, error_message + ) lib_logger.warning( f"Retry wait ({wait_time:.2f}s) exceeds budget. Rotating." ) @@ -1580,21 +1650,23 @@ async def _streaming_acompletion_with_retry( ) classified_error = classify_error(e) error_message = str(e).split("\n")[0] - + # Record in accumulator - error_accumulator.record_error(current_cred, classified_error, error_message) - + error_accumulator.record_error( + current_cred, classified_error, error_message + ) + lib_logger.warning( f"Cred {mask_credential(current_cred)} {classified_error.error_type} (HTTP {classified_error.status_code})." ) - + # Check if this error should trigger rotation if not should_rotate_on_error(classified_error): lib_logger.error( f"Non-recoverable error ({classified_error.error_type}). Failing." ) raise last_exception - + await self.usage_manager.record_failure( current_cred, model, classified_error ) @@ -1616,9 +1688,13 @@ async def _streaming_acompletion_with_retry( if provider_instance: # Ensure default Gemini safety settings are present (without overriding request) try: - self._apply_default_safety_settings(litellm_kwargs, provider) + self._apply_default_safety_settings( + litellm_kwargs, provider + ) except Exception: - lib_logger.debug("Could not apply default safety settings for streaming path; continuing.") + lib_logger.debug( + "Could not apply default safety settings for streaming path; continuing." + ) if "safety_settings" in litellm_kwargs: converted_settings = ( @@ -1699,7 +1775,11 @@ async def _streaming_acompletion_with_retry( yield chunk return - except (StreamedAPIError, litellm.RateLimitError, httpx.HTTPStatusError) as e: + except ( + StreamedAPIError, + litellm.RateLimitError, + httpx.HTTPStatusError, + ) as e: last_exception = e # This is the final, robust handler for streamed errors. @@ -1708,7 +1788,7 @@ async def _streaming_acompletion_with_retry( # The actual exception might be wrapped in our StreamedAPIError. original_exc = getattr(e, "data", e) classified_error = classify_error(original_exc) - + # Check if this error should trigger rotation if not should_rotate_on_error(classified_error): lib_logger.error( @@ -1745,16 +1825,18 @@ async def _streaming_acompletion_with_retry( error_message_text = error_details.get( "message", str(original_exc).split("\n")[0] ) - + # Record in accumulator for client reporting - error_accumulator.record_error(current_cred, classified_error, error_message_text) + error_accumulator.record_error( + current_cred, classified_error, error_message_text + ) if ( "quota" in error_message_text.lower() or "resource_exhausted" in error_status.lower() ): consecutive_quota_failures += 1 - + quota_value = "N/A" quota_id = "N/A" if "details" in error_details and isinstance( @@ -1764,10 +1846,15 @@ async def _streaming_acompletion_with_retry( if isinstance(detail.get("violations"), list): for violation in detail["violations"]: if "quotaValue" in violation: - quota_value = violation["quotaValue"] + quota_value = violation[ + "quotaValue" + ] if "quotaId" in violation: quota_id = violation["quotaId"] - if quota_value != "N/A" and quota_id != "N/A": + if ( + quota_value != "N/A" + and quota_id != "N/A" + ): break await self.usage_manager.record_failure( @@ -1798,8 +1885,13 @@ async def _streaming_acompletion_with_retry( f"Cred {mask_credential(current_cred)} {classified_error.error_type}. Rotating." ) - if classified_error.error_type in ["rate_limit", "quota_exceeded"]: - cooldown_duration = classified_error.retry_after or 60 + if classified_error.error_type in [ + "rate_limit", + "quota_exceeded", + ]: + cooldown_duration = ( + classified_error.retry_after or 60 + ) await self.cooldown_manager.start_cooldown( provider, cooldown_duration ) @@ -1827,14 +1919,18 @@ async def _streaming_acompletion_with_retry( ) classified_error = classify_error(e) error_message_text = str(e).split("\n")[0] - - # Record error in accumulator (server errors are abnormal) - error_accumulator.record_error(current_cred, classified_error, error_message_text) - + + # Record error in accumulator (server errors are transient, not abnormal) + error_accumulator.record_error( + current_cred, classified_error, error_message_text + ) + # Provider-level error: don't increment consecutive failures await self.usage_manager.record_failure( - current_cred, model, classified_error, - increment_consecutive_failures=False + current_cred, + model, + classified_error, + increment_consecutive_failures=False, ) if attempt >= self.max_retries - 1: @@ -1845,7 +1941,7 @@ async def _streaming_acompletion_with_retry( break wait_time = classified_error.retry_after or ( - 1 * (2**attempt) + 2**attempt ) + random.uniform(0, 1) remaining_budget = deadline - time.time() if wait_time > remaining_budget: @@ -1874,16 +1970,22 @@ async def _streaming_acompletion_with_retry( ) classified_error = classify_error(e) error_message_text = str(e).split("\n")[0] - + # Record error in accumulator - error_accumulator.record_error(current_cred, classified_error, error_message_text) + error_accumulator.record_error( + current_cred, classified_error, error_message_text + ) lib_logger.warning( f"Credential ...{current_cred[-6:]} failed with {classified_error.error_type} (Status: {classified_error.status_code}). Error: {error_message_text}." ) # Handle rate limits with cooldown - if classified_error.status_code == 429 or classified_error.error_type in ["rate_limit", "quota_exceeded"]: + if ( + classified_error.status_code == 429 + or classified_error.error_type + in ["rate_limit", "quota_exceeded"] + ): cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration @@ -1904,7 +2006,9 @@ async def _streaming_acompletion_with_retry( await self.usage_manager.record_failure( current_cred, model, classified_error ) - lib_logger.info(f"Rotating to next key after {classified_error.error_type} error.") + lib_logger.info( + f"Rotating to next key after {classified_error.error_type} error." + ) break finally: @@ -1913,26 +2017,28 @@ async def _streaming_acompletion_with_retry( # Build detailed error response using error accumulator error_accumulator.timeout_occurred = time.time() >= deadline - error_accumulator.model = model - error_accumulator.provider = provider - + if error_accumulator.has_errors(): # Log concise summary for server logs lib_logger.error(error_accumulator.build_log_message()) - + # Build structured error response for client error_response = error_accumulator.build_client_error_response() error_data = error_response else: # Fallback if no errors were recorded (shouldn't happen) - final_error_message = "Request failed: No available API keys after rotation or timeout." + final_error_message = ( + "Request failed: No available API keys after rotation or timeout." + ) if last_exception: - final_error_message = f"Request failed. Last error: {str(last_exception)}" + final_error_message = ( + f"Request failed. Last error: {str(last_exception)}" + ) error_data = { "error": {"message": final_error_message, "type": "proxy_error"} } lib_logger.error(final_error_message) - + yield f"data: {json.dumps(error_data)}\n\n" yield "data: [DONE]\n\n" @@ -1980,11 +2086,13 @@ def acompletion( # Handle iflow provider: remove stream_options to avoid HTTP 406 model = kwargs.get("model", "") provider = model.split("/")[0] if "/" in model else "" - + if provider == "iflow" and "stream_options" in kwargs: - lib_logger.debug("Removing stream_options for iflow provider to avoid HTTP 406") + lib_logger.debug( + "Removing stream_options for iflow provider to avoid HTTP 406" + ) kwargs.pop("stream_options", None) - + if kwargs.get("stream"): # Only add stream_options for providers that support it (excluding iflow) if provider != "iflow": @@ -1992,7 +2100,7 @@ def acompletion( kwargs["stream_options"] = {} if "include_usage" not in kwargs["stream_options"]: kwargs["stream_options"]["include_usage"] = True - + return self._streaming_acompletion_with_retry( request=request, pre_request_callback=pre_request_callback, **kwargs ) diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index 96a6cb73..76616c10 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -1,5 +1,6 @@ import re import json +import os from typing import Optional, Dict, Any import httpx @@ -20,20 +21,20 @@ def extract_retry_after_from_body(error_body: Optional[str]) -> Optional[int]: """ Extract the retry-after time from an API error response body. - + Handles various error formats including: - Gemini CLI: "Your quota will reset after 39s." - Generic: "quota will reset after 120s", "retry after 60s" - + Args: error_body: The raw error response body - + Returns: The retry time in seconds, or None if not found """ if not error_body: return None - + # Pattern to match various "reset after Xs" or "retry after Xs" formats patterns = [ r"quota will reset after\s*(\d+)s", @@ -41,7 +42,7 @@ def extract_retry_after_from_body(error_body: Optional[str]) -> Optional[int]: r"retry after\s*(\d+)s", r"try again in\s*(\d+)\s*seconds?", ] - + for pattern in patterns: match = re.search(pattern, error_body, re.IGNORECASE) if match: @@ -49,7 +50,7 @@ def extract_retry_after_from_body(error_body: Optional[str]) -> Optional[int]: return int(match.group(1)) except (ValueError, IndexError): continue - + return None @@ -70,29 +71,33 @@ class PreRequestCallbackError(Exception): # ============================================================================= # Abnormal errors that require attention and should always be reported to client -ABNORMAL_ERROR_TYPES = frozenset({ - "forbidden", # 403 - credential access issue - "authentication", # 401 - credential invalid/revoked - "pre_request_callback_error", # Internal proxy error -}) +ABNORMAL_ERROR_TYPES = frozenset( + { + "forbidden", # 403 - credential access issue + "authentication", # 401 - credential invalid/revoked + "pre_request_callback_error", # Internal proxy error + } +) # Normal/expected errors during operation - only report if ALL credentials fail -NORMAL_ERROR_TYPES = frozenset({ - "rate_limit", # 429 - expected during high load - "quota_exceeded", # Expected when quota runs out - "server_error", # 5xx - transient provider issues - "api_connection", # Network issues - transient -}) +NORMAL_ERROR_TYPES = frozenset( + { + "rate_limit", # 429 - expected during high load + "quota_exceeded", # Expected when quota runs out + "server_error", # 5xx - transient provider issues + "api_connection", # Network issues - transient + } +) def is_abnormal_error(classified_error: "ClassifiedError") -> bool: """ Check if an error is abnormal and should be reported to the client. - + Abnormal errors indicate credential issues that need attention: - 403 Forbidden: Credential doesn't have access - 401 Unauthorized: Credential is invalid/revoked - + Normal errors are expected during operation: - 429 Rate limit: Expected during high load - 5xx Server errors: Transient provider issues @@ -103,11 +108,10 @@ def is_abnormal_error(classified_error: "ClassifiedError") -> bool: def mask_credential(credential: str) -> str: """ Mask a credential for safe display in logs and error messages. - + - For API keys: shows last 6 characters (e.g., "...xyz123") - For OAuth file paths: shows just the filename (e.g., "antigravity_oauth_1.json") """ - import os if os.path.isfile(credential): return os.path.basename(credential) elif len(credential) > 6: @@ -119,77 +123,79 @@ def mask_credential(credential: str) -> str: class RequestErrorAccumulator: """ Tracks errors encountered during a request's credential rotation cycle. - + Used to build informative error messages for clients when all credentials are exhausted. Distinguishes between abnormal errors (that need attention) and normal errors (expected during operation). """ - + def __init__(self): self.abnormal_errors: list = [] # 403, 401 - always report details - self.normal_errors: list = [] # 429, 5xx - summarize only - self.total_credentials_tried: int = 0 + self.normal_errors: list = [] # 429, 5xx - summarize only + self._tried_credentials: set = set() # Track unique credentials self.timeout_occurred: bool = False self.model: str = "" self.provider: str = "" - + def record_error( - self, - credential: str, - classified_error: "ClassifiedError", - error_message: str + self, credential: str, classified_error: "ClassifiedError", error_message: str ): """Record an error for a credential.""" - self.total_credentials_tried += 1 + self._tried_credentials.add(credential) masked_cred = mask_credential(credential) - + error_record = { "credential": masked_cred, "error_type": classified_error.error_type, "status_code": classified_error.status_code, - "message": self._truncate_message(error_message, 150) + "message": self._truncate_message(error_message, 150), } - + if is_abnormal_error(classified_error): self.abnormal_errors.append(error_record) else: self.normal_errors.append(error_record) - + + @property + def total_credentials_tried(self) -> int: + """Return the number of unique credentials tried.""" + return len(self._tried_credentials) + def _truncate_message(self, message: str, max_length: int = 150) -> str: """Truncate error message for readability.""" # Take first line and truncate - first_line = message.split('\n')[0] + first_line = message.split("\n")[0] if len(first_line) > max_length: return first_line[:max_length] + "..." return first_line - + def has_errors(self) -> bool: """Check if any errors were recorded.""" return bool(self.abnormal_errors or self.normal_errors) - + def has_abnormal_errors(self) -> bool: """Check if any abnormal errors were recorded.""" return bool(self.abnormal_errors) - + def get_normal_error_summary(self) -> str: """Get a summary of normal errors (not individual details).""" if not self.normal_errors: return "" - + # Count by type counts = {} for err in self.normal_errors: err_type = err["error_type"] counts[err_type] = counts.get(err_type, 0) + 1 - + # Build summary like "3 rate_limit, 1 server_error" parts = [f"{count} {err_type}" for err_type, count in counts.items()] return ", ".join(parts) - + def build_client_error_response(self) -> dict: """ Build a structured error response for the client. - + Returns a dict suitable for JSON serialization in the error response. """ # Determine the primary failure reason @@ -199,24 +205,34 @@ def build_client_error_response(self) -> dict: else: error_type = "proxy_all_credentials_exhausted" base_message = f"All {self.total_credentials_tried} credential(s) exhausted for {self.provider}" - + # Build human-readable message message_parts = [base_message] - + if self.abnormal_errors: message_parts.append("\n\nCredential issues (require attention):") for err in self.abnormal_errors: - status = f"HTTP {err['status_code']}" if err['status_code'] else err['error_type'] - message_parts.append(f"\n • {err['credential']}: {status} - {err['message']}") - + status = ( + f"HTTP {err['status_code']}" + if err["status_code"] is not None + else err["error_type"] + ) + message_parts.append( + f"\n • {err['credential']}: {status} - {err['message']}" + ) + normal_summary = self.get_normal_error_summary() if normal_summary: if self.abnormal_errors: - message_parts.append(f"\n\nAdditionally: {normal_summary} (expected during normal operation)") + message_parts.append( + f"\n\nAdditionally: {normal_summary} (expected during normal operation)" + ) else: message_parts.append(f"\n\nAll failures were: {normal_summary}") - message_parts.append("\nThis is normal during high load - retry later or add more credentials.") - + message_parts.append( + "\nThis is normal during high load - retry later or add more credentials." + ) + response = { "error": { "message": "".join(message_parts), @@ -226,44 +242,48 @@ def build_client_error_response(self) -> dict: "provider": self.provider, "credentials_tried": self.total_credentials_tried, "timeout": self.timeout_occurred, - } + }, } } - + # Only include abnormal errors in details (they need attention) if self.abnormal_errors: response["error"]["details"]["abnormal_errors"] = self.abnormal_errors - + # Include summary of normal errors if normal_summary: response["error"]["details"]["normal_error_summary"] = normal_summary - + return response - + def build_log_message(self) -> str: """ Build a concise log message for server-side logging. - + Shorter than client message, suitable for terminal display. """ parts = [] - + if self.timeout_occurred: - parts.append(f"TIMEOUT: {self.total_credentials_tried} creds tried for {self.model}") + parts.append( + f"TIMEOUT: {self.total_credentials_tried} creds tried for {self.model}" + ) else: - parts.append(f"ALL CREDS EXHAUSTED: {self.total_credentials_tried} tried for {self.model}") - + parts.append( + f"ALL CREDS EXHAUSTED: {self.total_credentials_tried} tried for {self.model}" + ) + if self.abnormal_errors: abnormal_summary = ", ".join( f"{e['credential']}={e['status_code'] or e['error_type']}" for e in self.abnormal_errors ) parts.append(f"ISSUES: {abnormal_summary}") - + normal_summary = self.get_normal_error_summary() if normal_summary: parts.append(f"Normal: {normal_summary}") - + return " | ".join(parts) @@ -296,7 +316,7 @@ def get_retry_after(error: Exception) -> Optional[int]: if isinstance(error, httpx.HTTPStatusError): headers = error.response.headers # Check standard Retry-After header (case-insensitive) - retry_header = headers.get('retry-after') or headers.get('Retry-After') + retry_header = headers.get("retry-after") or headers.get("Retry-After") if retry_header: try: return int(retry_header) # Assumes seconds format @@ -304,10 +324,13 @@ def get_retry_after(error: Exception) -> Optional[int]: pass # Might be HTTP date format, skip for now # Check X-RateLimit-Reset header (Unix timestamp) - reset_header = headers.get('x-ratelimit-reset') or headers.get('X-RateLimit-Reset') + reset_header = headers.get("x-ratelimit-reset") or headers.get( + "X-RateLimit-Reset" + ) if reset_header: try: import time + reset_timestamp = int(reset_header) current_time = int(time.time()) wait_seconds = reset_timestamp - current_time @@ -357,16 +380,16 @@ def get_retry_after(error: Exception) -> Optional[int]: continue # 3. Handle duration formats like "60s", "2m", "1h" - duration_match = re.search(r'(\d+)\s*([smh])', error_str) + duration_match = re.search(r"(\d+)\s*([smh])", error_str) if duration_match: try: value = int(duration_match.group(1)) unit = duration_match.group(2) - if unit == 's': + if unit == "s": return value - elif unit == 'm': + elif unit == "m": return value * 60 - elif unit == 'h': + elif unit == "h": return value * 3600 except (ValueError, IndexError): pass @@ -381,15 +404,15 @@ def get_retry_after(error: Exception) -> Optional[int]: if value.isdigit(): return int(value) # Handle "60s", "2m" format in attribute - duration_match = re.search(r'(\d+)\s*([smh])', value.lower()) + duration_match = re.search(r"(\d+)\s*([smh])", value.lower()) if duration_match: val = int(duration_match.group(1)) unit = duration_match.group(2) - if unit == 's': + if unit == "s": return val - elif unit == 'm': + elif unit == "m": return val * 60 - elif unit == 'h': + elif unit == "h": return val * 3600 return None @@ -399,7 +422,7 @@ def classify_error(e: Exception) -> ClassifiedError: """ Classifies an exception into a structured ClassifiedError object. Now handles both litellm and httpx exceptions. - + Error types and their typical handling: - rate_limit (429): Rotate key, may retry with backoff - server_error (5xx): Retry with backoff, then rotate @@ -412,16 +435,16 @@ def classify_error(e: Exception) -> ClassifiedError: - unknown: Rotate key (safer to try another) """ status_code = getattr(e, "status_code", None) - + if isinstance(e, httpx.HTTPStatusError): # [NEW] Handle httpx errors first status_code = e.response.status_code - + # Try to get error body for better classification try: - error_body = e.response.text.lower() if hasattr(e.response, 'text') else "" + error_body = e.response.text.lower() if hasattr(e.response, "text") else "" except Exception: error_body = "" - + if status_code == 401: return ClassifiedError( error_type="authentication", @@ -453,8 +476,18 @@ def classify_error(e: Exception) -> ClassifiedError: retry_after=retry_after, ) if status_code == 400: - # Check for context window / token limit errors - if "context" in error_body or "token" in error_body or "too long" in error_body: + # Check for context window / token limit errors with more specific patterns + if any( + pattern in error_body + for pattern in [ + "context_length", + "max_tokens", + "token limit", + "context window", + "too many tokens", + "too long", + ] + ): return ClassifiedError( error_type="context_window_exceeded", original_exception=e, @@ -465,6 +498,11 @@ def classify_error(e: Exception) -> ClassifiedError: original_exception=e, status_code=status_code, ) + return ClassifiedError( + error_type="invalid_request", + original_exception=e, + status_code=status_code, + ) if 400 <= status_code < 500: # Other 4xx errors - generally client errors return ClassifiedError( @@ -567,7 +605,7 @@ def is_unrecoverable_error(e: Exception) -> bool: def should_rotate_on_error(classified_error: ClassifiedError) -> bool: """ Determines if an error should trigger key rotation. - + Errors that SHOULD rotate (try another key): - rate_limit: Current key is throttled - quota_exceeded: Current key/account exhausted @@ -576,12 +614,12 @@ def should_rotate_on_error(classified_error: ClassifiedError) -> bool: - server_error: Provider having issues (might work with different endpoint/key) - api_connection: Network issues (might be transient) - unknown: Safer to try another key - + Errors that should NOT rotate (fail immediately): - invalid_request: Client error in request payload (won't help to retry) - context_window_exceeded: Request too large (won't help to retry) - pre_request_callback_error: Internal proxy error - + Returns: True if should rotate to next key, False if should fail immediately """ @@ -596,10 +634,10 @@ def should_rotate_on_error(classified_error: ClassifiedError) -> bool: def should_retry_same_key(classified_error: ClassifiedError) -> bool: """ Determines if an error should retry with the same key (with backoff). - + Only server errors and connection issues should retry the same key, as these are often transient. - + Returns: True if should retry same key, False if should rotate immediately """ diff --git a/src/rotator_library/failure_logger.py b/src/rotator_library/failure_logger.py index 8c4e043a..b1dddfbc 100644 --- a/src/rotator_library/failure_logger.py +++ b/src/rotator_library/failure_logger.py @@ -4,6 +4,7 @@ import os from datetime import datetime + def setup_failure_logger(): """Sets up a dedicated JSON logger for writing detailed failure logs to a file.""" log_dir = "logs" @@ -12,15 +13,15 @@ def setup_failure_logger(): # Create a logger specifically for failures. # This logger will NOT propagate to the root logger. - logger = logging.getLogger('failure_logger') + logger = logging.getLogger("failure_logger") logger.setLevel(logging.INFO) logger.propagate = False # Use a rotating file handler handler = RotatingFileHandler( - os.path.join(log_dir, 'failures.log'), - maxBytes=5*1024*1024, # 5 MB - backupCount=2 + os.path.join(log_dir, "failures.log"), + maxBytes=5 * 1024 * 1024, # 5 MB + backupCount=2, ) # Custom JSON formatter for structured logs @@ -30,62 +31,65 @@ def format(self, record): return json.dumps(record.msg) handler.setFormatter(JsonFormatter()) - + # Add handler only if it hasn't been added before if not logger.handlers: logger.addHandler(handler) return logger + # Initialize the dedicated logger for detailed failure logs failure_logger = setup_failure_logger() # Get the main library logger for concise, propagated messages -main_lib_logger = logging.getLogger('rotator_library') +main_lib_logger = logging.getLogger("rotator_library") + def _extract_response_body(error: Exception) -> str: """ Extract the full response body from various error types. - + Handles: - httpx.HTTPStatusError: response.text or response.content - litellm exceptions: various response attributes - Other exceptions: str(error) """ # Try to get response body from httpx errors - if hasattr(error, 'response') and error.response is not None: + if hasattr(error, "response") and error.response is not None: response = error.response # Try .text first (decoded) - if hasattr(response, 'text') and response.text: + if hasattr(response, "text") and response.text: return response.text # Try .content (bytes) - if hasattr(response, 'content') and response.content: + if hasattr(response, "content") and response.content: try: - return response.content.decode('utf-8', errors='replace') + return response.content.decode("utf-8", errors="replace") except Exception: return str(response.content) - # Try reading response if it's a streaming response that was read - if hasattr(response, '_content') and response._content: - try: - return response._content.decode('utf-8', errors='replace') - except Exception: - return str(response._content) - + # Check for litellm's body attribute - if hasattr(error, 'body') and error.body: + if hasattr(error, "body") and error.body: return str(error.body) - + # Check for message attribute that might contain response - if hasattr(error, 'message') and error.message: + if hasattr(error, "message") and error.message: return str(error.message) - + return None -def log_failure(api_key: str, model: str, attempt: int, error: Exception, request_headers: dict, raw_response_text: str = None): +def log_failure( + api_key: str, + model: str, + attempt: int, + error: Exception, + request_headers: dict, + raw_response_text: str = None, +): """ Logs a detailed failure message to a file and a concise summary to the main logger. - + Args: api_key: The API key or credential path that was used model: The model that was requested @@ -103,19 +107,30 @@ def log_failure(api_key: str, model: str, attempt: int, error: Exception, reques # Get full error message (not truncated) full_error_message = str(error) - + # Also capture any nested/wrapped exception info error_chain = [] + visited = set() # Track visited exceptions to detect circular references current_error = error while current_error: - error_chain.append({ - "type": type(current_error).__name__, - "message": str(current_error)[:2000] # Limit per-error message size - }) - current_error = getattr(current_error, '__cause__', None) or getattr(current_error, '__context__', None) - if len(error_chain) > 5: # Prevent infinite loops + # Check for circular references + error_id = id(current_error) + if error_id in visited: break - + visited.add(error_id) + + error_chain.append( + { + "type": type(current_error).__name__, + "message": str(current_error)[:2000], # Limit per-error message size + } + ) + current_error = getattr(current_error, "__cause__", None) or getattr( + current_error, "__context__", None + ) + if len(error_chain) > 5: # Prevent excessive chain length + break + detailed_log_data = { "timestamp": datetime.utcnow().isoformat(), "api_key_ending": api_key[-4:] if len(api_key) >= 4 else "****", @@ -123,7 +138,9 @@ def log_failure(api_key: str, model: str, attempt: int, error: Exception, reques "attempt_number": attempt, "error_type": type(error).__name__, "error_message": full_error_message[:5000], # Limit total size - "raw_response": raw_response[:10000] if raw_response else None, # Limit response size + "raw_response": raw_response[:10000] + if raw_response + else None, # Limit response size "request_headers": request_headers, "error_chain": error_chain if len(error_chain) > 1 else None, } From 956bdbbffa813a623b911fc6ff61caf0dba00fbf Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 4 Dec 2025 04:59:15 +0100 Subject: [PATCH 064/221] =?UTF-8?q?fix(provider):=20=F0=9F=90=9B=20increas?= =?UTF-8?q?e=20timeout=20for=20antigravity=20API=20requests=20from=20120s?= =?UTF-8?q?=20to=20600s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous 120-second timeout was insufficient for long-running requests to the Antigravity provider API, causing premature request failures. This change increases the timeout to 600 seconds (10 minutes) for both streaming and non-streaming completion requests to accommodate longer processing times and prevent timeout errors during complex operations. --- src/rotator_library/providers/antigravity_provider.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 4adb1114..5751bba2 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -2620,7 +2620,7 @@ async def _handle_non_streaming( file_logger: Optional[AntigravityFileLogger] = None ) -> litellm.ModelResponse: """Handle non-streaming completion.""" - response = await client.post(url, headers=headers, json=payload, timeout=120.0) + response = await client.post(url, headers=headers, json=payload, timeout=600.0) response.raise_for_status() data = response.json() @@ -2652,7 +2652,7 @@ async def _handle_streaming( "is_complete": False # Track if we received usageMetadata } - async with client.stream("POST", url, headers=headers, json=payload, timeout=120.0) as response: + async with client.stream("POST", url, headers=headers, json=payload, timeout=600.0) as response: if response.status_code >= 400: try: error_body = await response.aread() From fce1762ff0e0bee2a09f369f6bdfdce903faf244 Mon Sep 17 00:00:00 2001 From: MasuRii Date: Thu, 4 Dec 2025 17:24:13 +0800 Subject: [PATCH 065/221] fix(logging): preserve full credential filenames in logs Resolved logging truncation issue where OAuth credential filenames were being aggressively abbreviated (e.g., `...6.json` instead of `antigravity_oauth_16.json`), causing ambiguity when debugging or auditing specific credentials. **Changes:** - Enhanced `mask_credential()` utility in error_handler.py: - Now explicitly detects `.json` file extensions - Returns full basename for file paths (e.g., `antigravity_oauth_16.json`) - Maintains security by masking API keys to last 6 characters (`...xyz123`) - Replaced all manual credential truncation with centralized `mask_credential()`: - client.py: 15 instances (stream handling, retry logging, model discovery) - usage_manager.py: 16 instances (key acquisition, release, cooldown tracking) - failure_logger.py: 2 instances (failure logging and summaries) - Code quality improvements: - Fixed indentation error in client.py during refactoring - Ensured consistent, safe credential logging across entire application - Configuration: - Added `oauth_creds/` to .gitignore to prevent accidental credential commits **Impact:** This standardizes credential display throughout the application, enabling accurate debugging and auditing while maintaining security for raw API keys. Logs now clearly distinguish between multiple OAuth files (e.g., `6.json` vs `16.json`) without exposing sensitive key material. **Files Modified:** - .gitignore (added oauth_creds exclusion) - src/rotator_library/client.py (15 replacements) - src/rotator_library/error_handler.py (enhanced mask_credential logic) - src/rotator_library/failure_logger.py (2 replacements) - src/rotator_library/usage_manager.py (16 replacements) --- .gitignore | 2 ++ src/rotator_library/client.py | 40 +++++++++++---------------- src/rotator_library/error_handler.py | 2 +- src/rotator_library/failure_logger.py | 5 ++-- src/rotator_library/usage_manager.py | 32 ++++++++++----------- 5 files changed, 38 insertions(+), 43 deletions(-) diff --git a/.gitignore b/.gitignore index 1a75e867..0c94208b 100644 --- a/.gitignore +++ b/.gitignore @@ -128,3 +128,5 @@ cache/antigravity/thought_signatures.json logs/ cache/ *.env + +oauth_creds diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index d603d463..5956d193 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -537,7 +537,7 @@ async def _safe_streaming_wrapper( while True: if request and await request.is_disconnected(): lib_logger.info( - f"Client disconnected. Aborting stream for credential ...{key[-6:]}." + f"Client disconnected. Aborting stream for credential {mask_credential(key)}." ) break @@ -695,7 +695,7 @@ async def _safe_streaming_wrapper( # Catch any other unexpected errors during streaming. lib_logger.error(f"Caught unexpected exception of type: {type(e).__name__}") lib_logger.error( - f"An unexpected error occurred during the stream for credential ...{key[-6:]}: {e}" + f"An unexpected error occurred during the stream for credential {mask_credential(key)}: {e}" ) # We still need to raise it so the client knows something went wrong. raise @@ -705,7 +705,7 @@ async def _safe_streaming_wrapper( # The primary goal is to ensure usage is always logged internally. await self.usage_manager.release_key(key, model) lib_logger.info( - f"STREAM FINISHED and lock released for credential ...{key[-6:]}." + f"STREAM FINISHED and lock released for credential {mask_credential(key)}." ) # Only send [DONE] if the stream completed naturally and the client is still there. @@ -1006,7 +1006,7 @@ async def _execute_with_retry( for attempt in range(self.max_retries): try: lib_logger.info( - f"Attempting call with credential ...{current_cred[-6:]} (Attempt {attempt + 1}/{self.max_retries})" + f"Attempting call with credential {mask_credential(current_cred)} (Attempt {attempt + 1}/{self.max_retries})" ) if pre_request_callback: @@ -1495,9 +1495,9 @@ async def _streaming_acompletion_with_retry( for attempt in range(self.max_retries): try: lib_logger.info( - f"Attempting stream with credential ...{current_cred[-6:]} (Attempt {attempt + 1}/{self.max_retries})" + f"Attempting stream with credential {mask_credential(current_cred)} (Attempt {attempt + 1}/{self.max_retries})" ) - + if pre_request_callback: try: await pre_request_callback( @@ -1518,7 +1518,7 @@ async def _streaming_acompletion_with_retry( ) lib_logger.info( - f"Stream connection established for credential ...{current_cred[-6:]}. Processing response." + f"Stream connection established for credential {mask_credential(current_cred)}. Processing response." ) key_acquired = False @@ -1735,7 +1735,7 @@ async def _streaming_acompletion_with_retry( for attempt in range(self.max_retries): try: lib_logger.info( - f"Attempting stream with credential ...{current_cred[-6:]} (Attempt {attempt + 1}/{self.max_retries})" + f"Attempting stream with credential {mask_credential(current_cred)} (Attempt {attempt + 1}/{self.max_retries})" ) if pre_request_callback: @@ -1763,7 +1763,7 @@ async def _streaming_acompletion_with_retry( ) lib_logger.info( - f"Stream connection established for credential ...{current_cred[-6:]}. Processing response." + f"Stream connection established for credential {mask_credential(current_cred)}. Processing response." ) key_acquired = False @@ -1935,7 +1935,7 @@ async def _streaming_acompletion_with_retry( if attempt >= self.max_retries - 1: lib_logger.warning( - f"Credential ...{current_cred[-6:]} failed after max retries for model {model} due to a server error. Rotating key silently." + f"Credential {mask_credential(current_cred)} failed after max retries for model {model} due to a server error. Rotating key silently." ) # [MODIFIED] Do not yield to the client here. break @@ -1951,7 +1951,7 @@ async def _streaming_acompletion_with_retry( break lib_logger.warning( - f"Credential ...{current_cred[-6:]} encountered a server error for model {model}. Reason: '{error_message_text}'. Retrying in {wait_time:.2f}s." + f"Credential {mask_credential(current_cred)} encountered a server error for model {model}. Reason: '{error_message_text}'. Retrying in {wait_time:.2f}s." ) await asyncio.sleep(wait_time) continue @@ -1977,7 +1977,7 @@ async def _streaming_acompletion_with_retry( ) lib_logger.warning( - f"Credential ...{current_cred[-6:]} failed with {classified_error.error_type} (Status: {classified_error.status_code}). Error: {error_message_text}." + f"Credential {mask_credential(current_cred)} failed with {classified_error.error_type} (Status: {classified_error.status_code}). Error: {error_message_text}." ) # Handle rate limits with cooldown @@ -2179,13 +2179,9 @@ async def get_available_models(self, provider: str) -> List[str]: for credential in shuffled_credentials: try: # Display last 6 chars for API keys, or the filename for OAuth paths - cred_display = ( - credential[-6:] - if not os.path.isfile(credential) - else os.path.basename(credential) - ) + cred_display = mask_credential(credential) lib_logger.debug( - f"Attempting to get models for {provider} with credential ...{cred_display}" + f"Attempting to get models for {provider} with credential {cred_display}" ) models = await provider_instance.get_models( credential, self.http_client @@ -2216,13 +2212,9 @@ async def get_available_models(self, provider: str) -> List[str]: return final_models except Exception as e: classified_error = classify_error(e) - cred_display = ( - credential[-6:] - if not os.path.isfile(credential) - else os.path.basename(credential) - ) + cred_display = mask_credential(credential) lib_logger.debug( - f"Failed to get models for provider {provider} with credential ...{cred_display}: {classified_error.error_type}. Trying next credential." + f"Failed to get models for provider {provider} with credential {cred_display}: {classified_error.error_type}. Trying next credential." ) continue # Try the next credential diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index 76616c10..ac4b8a1e 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -112,7 +112,7 @@ def mask_credential(credential: str) -> str: - For API keys: shows last 6 characters (e.g., "...xyz123") - For OAuth file paths: shows just the filename (e.g., "antigravity_oauth_1.json") """ - if os.path.isfile(credential): + if os.path.isfile(credential) or credential.endswith(".json"): return os.path.basename(credential) elif len(credential) > 6: return f"...{credential[-6:]}" diff --git a/src/rotator_library/failure_logger.py b/src/rotator_library/failure_logger.py index b1dddfbc..8f1848ae 100644 --- a/src/rotator_library/failure_logger.py +++ b/src/rotator_library/failure_logger.py @@ -3,6 +3,7 @@ from logging.handlers import RotatingFileHandler import os from datetime import datetime +from .error_handler import mask_credential def setup_failure_logger(): @@ -133,7 +134,7 @@ def log_failure( detailed_log_data = { "timestamp": datetime.utcnow().isoformat(), - "api_key_ending": api_key[-4:] if len(api_key) >= 4 else "****", + "api_key_ending": mask_credential(api_key), "model": model, "attempt_number": attempt, "error_type": type(error).__name__, @@ -148,7 +149,7 @@ def log_failure( # 2. Log a concise summary to the main library logger, which will propagate summary_message = ( - f"API call failed for model {model} with key ...{api_key[-4:] if len(api_key) >= 4 else '****'}. " + f"API call failed for model {model} with key {mask_credential(api_key)}. " f"Error: {type(error).__name__}. See failures.log for details." ) main_lib_logger.error(summary_message) diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 4ec2b825..76ee21e8 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -9,7 +9,7 @@ import aiofiles import litellm -from .error_handler import ClassifiedError, NoAvailableKeysError +from .error_handler import ClassifiedError, NoAvailableKeysError, mask_credential from .providers import PROVIDER_PLUGINS lib_logger = logging.getLogger("rotator_library") @@ -139,7 +139,7 @@ async def _reset_daily_stats_if_needed(self): last_reset_dt is None or last_reset_dt < reset_threshold_today <= now_utc ): - lib_logger.debug(f"Performing daily reset for key ...{key[-6:]}") + lib_logger.debug(f"Performing daily reset for key {mask_credential(key)}") needs_saving = True # Reset cooldowns @@ -237,7 +237,7 @@ def _select_weighted_random( if lib_logger.isEnabledFor(logging.DEBUG): total_weight = sum(weights) weight_info = ", ".join( - f"...{cred[-6:]}: w={w:.1f} ({w/total_weight*100:.1f}%)" + f"{mask_credential(cred)}: w={w:.1f} ({w/total_weight*100:.1f}%)" for (cred, _), w in zip(candidates, weights) ) #lib_logger.debug(f"Weighted selection candidates: {weight_info}") @@ -358,7 +358,7 @@ async def acquire_key( if not state["models_in_use"]: state["models_in_use"][model] = 1 lib_logger.info( - f"Acquired Priority-{priority_level} Tier-1 key ...{key[-6:]} for model {model} " + f"Acquired Priority-{priority_level} Tier-1 key {mask_credential(key)} for model {model} " f"(selection: {selection_method}, usage: {usage})" ) return key @@ -371,7 +371,7 @@ async def acquire_key( if current_count < max_concurrent: state["models_in_use"][model] = current_count + 1 lib_logger.info( - f"Acquired Priority-{priority_level} Tier-2 key ...{key[-6:]} for model {model} " + f"Acquired Priority-{priority_level} Tier-2 key {mask_credential(key)} for model {model} " f"(selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" ) return key @@ -452,7 +452,7 @@ async def acquire_key( if not state["models_in_use"]: state["models_in_use"][model] = 1 lib_logger.info( - f"Acquired Tier 1 key ...{key[-6:]} for model {model} " + f"Acquired Tier 1 key {mask_credential(key)} for model {model} " f"(selection: {selection_method}, usage: {usage})" ) return key @@ -465,7 +465,7 @@ async def acquire_key( if current_count < max_concurrent: state["models_in_use"][model] = current_count + 1 lib_logger.info( - f"Acquired Tier 2 key ...{key[-6:]} for model {model} " + f"Acquired Tier 2 key {mask_credential(key)} for model {model} " f"(selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" ) return key @@ -521,12 +521,12 @@ async def release_key(self, key: str, model: str): if remaining <= 0: del state["models_in_use"][model] # Clean up when count reaches 0 lib_logger.info( - f"Released credential ...{key[-6:]} from model {model} " + f"Released credential {mask_credential(key)} from model {model} " f"(remaining concurrent: {max(0, remaining)})" ) else: lib_logger.warning( - f"Attempted to release credential ...{key[-6:]} for model {model}, but it was not in use." + f"Attempted to release credential {mask_credential(key)} for model {model}, but it was not in use." ) # Notify all tasks waiting on this key's condition @@ -589,7 +589,7 @@ async def record_success( usage, "completion_tokens", 0 ) # Not present in embedding responses lib_logger.info( - f"Recorded usage from response object for key ...{key[-6:]}" + f"Recorded usage from response object for key {mask_credential(key)}" ) try: provider_name = model.split("/")[0] @@ -681,14 +681,14 @@ async def record_failure( # Rate limit errors: use retry_after if available, otherwise default to 60s cooldown_seconds = classified_error.retry_after or 60 lib_logger.info( - f"Rate limit error on key ...{key[-6:]} for model {model}. " + f"Rate limit error on key {mask_credential(key)} for model {model}. " f"Using {'provided' if classified_error.retry_after else 'default'} retry_after: {cooldown_seconds}s" ) elif classified_error.error_type == "authentication": # Apply a 5-minute key-level lockout for auth errors key_data["key_cooldown_until"] = time.time() + 300 lib_logger.warning( - f"Authentication error on key ...{key[-6:]}. Applying 5-minute key-level lockout." + f"Authentication error on key {mask_credential(key)}. Applying 5-minute key-level lockout." ) # Auth errors still use escalating backoff for the specific model cooldown_seconds = 300 # 5 minutes for model cooldown @@ -707,7 +707,7 @@ async def record_failure( backoff_tiers = {1: 10, 2: 30, 3: 60, 4: 120} cooldown_seconds = backoff_tiers.get(count, 7200) # Default to 2 hours for "spent" keys lib_logger.warning( - f"Failure #{count} for key ...{key[-6:]} with model {model}. " + f"Failure #{count} for key {mask_credential(key)} with model {model}. " f"Error type: {classified_error.error_type}" ) else: @@ -715,7 +715,7 @@ async def record_failure( if cooldown_seconds is None: cooldown_seconds = 30 # 30s cooldown for provider issues lib_logger.info( - f"Provider-level error ({classified_error.error_type}) for key ...{key[-6:]} with model {model}. " + f"Provider-level error ({classified_error.error_type}) for key {mask_credential(key)} with model {model}. " f"NOT incrementing consecutive failures. Applying {cooldown_seconds}s cooldown." ) @@ -723,7 +723,7 @@ async def record_failure( model_cooldowns = key_data.setdefault("model_cooldowns", {}) model_cooldowns[model] = time.time() + cooldown_seconds lib_logger.warning( - f"Cooldown applied for key ...{key[-6:]} with model {model}: {cooldown_seconds}s. " + f"Cooldown applied for key {mask_credential(key)} with model {model}: {cooldown_seconds}s. " f"Error type: {classified_error.error_type}" ) @@ -750,5 +750,5 @@ async def _check_key_lockout(self, key: str, key_data: Dict): if long_term_lockout_models >= 3: key_data["key_cooldown_until"] = now + 300 # 5-minute key lockout lib_logger.error( - f"Key ...{key[-6:]} has {long_term_lockout_models} models in long-term lockout. Applying 5-minute key-level lockout." + f"Key {mask_credential(key)} has {long_term_lockout_models} models in long-term lockout. Applying 5-minute key-level lockout." ) From 0dd6d21217d7417004183adec77354904d76f9c1 Mon Sep 17 00:00:00 2001 From: MasuRii Date: Thu, 4 Dec 2025 18:08:31 +0800 Subject: [PATCH 066/221] fix(rotator): prevent quota errors from global cooldown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves a critical issue where a single credential hitting quota limits triggered a provider-wide cooldown, causing denial of service for all remaining healthy credentials. **Problem:** When any credential encountered a 429 "Quota Exceeded" error, the system applied a global 60-second cooldown to the entire provider. This blocked ALL credentials from being used, even though quota errors are credential-specific, not provider-wide. **Root Cause:** The error classification system did not distinguish between: - Rate limit errors (IP/provider throttling) → affects all credentials - Quota errors (account/key limits) → affects only that credential Both were treated as `rate_limit` and triggered `cooldown_manager`, which pauses the entire provider. **Solution:** 1. **Enhanced Error Classification** (error_handler.py): - Parse RateLimitError messages for "quota"/"resource_exhausted" - Classify as `quota_exceeded` (not `rate_limit`) - Preserves retry_after headers for both types 2. **Separated Cooldown Logic** (client.py): - Global cooldowns now ONLY for `rate_limit` errors - `quota_exceeded` errors skip `cooldown_manager.start_cooldown()` - Quota failures still apply key-specific backoff via `usage_manager` 3. **Updated Logging** (usage_manager.py): - Recognizes both `rate_limit` and `quota_exceeded` for key backoff - Logs precise error type for debugging **Impact:** - ✅ Quota failures now immediately rotate to next credential - ✅ No more provider-wide DoS from single key quota exhaustion - ✅ Global cooldowns reserved for true rate limiting (IP/throttling) - ✅ Maintains per-key escalating backoff for quota errors **Changed Files:** - `src/rotator_library/client.py` (6 locations) - `src/rotator_library/error_handler.py` - `src/rotator_library/usage_manager.py` - `.gitignore` (added oauth_creds) --- .gitignore | 2 ++ src/rotator_library/client.py | 36 ++++++++++------------------ src/rotator_library/error_handler.py | 9 +++++++ src/rotator_library/usage_manager.py | 6 ++--- 4 files changed, 27 insertions(+), 26 deletions(-) diff --git a/.gitignore b/.gitignore index 1a75e867..16fefd3f 100644 --- a/.gitignore +++ b/.gitignore @@ -128,3 +128,5 @@ cache/antigravity/thought_signatures.json logs/ cache/ *.env + +oauth_creds \ No newline at end of file diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index d603d463..e1923c44 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -1064,7 +1064,8 @@ async def _execute_with_retry( f"Key {mask_credential(current_cred)} hit rate limit for {model}. Rotating key." ) - if classified_error.status_code == 429: + # Only trigger provider-wide cooldown for rate limits, not quota issues + if classified_error.status_code == 429 and classified_error.error_type != "quota_exceeded": cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration @@ -1165,11 +1166,8 @@ async def _execute_with_retry( current_cred, classified_error, error_message ) - # Handle rate limits with cooldown - if classified_error.error_type in [ - "rate_limit", - "quota_exceeded", - ]: + # Handle rate limits with cooldown (exclude quota_exceeded from provider-wide cooldown) + if classified_error.error_type == "rate_limit": cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration @@ -1225,11 +1223,10 @@ async def _execute_with_retry( f"Key {mask_credential(current_cred)} {classified_error.error_type} (HTTP {classified_error.status_code})." ) - # Handle rate limits with cooldown + # Handle rate limits with cooldown (exclude quota_exceeded from provider-wide cooldown) if ( - classified_error.status_code == 429 - or classified_error.error_type - in ["rate_limit", "quota_exceeded"] + (classified_error.status_code == 429 and classified_error.error_type != "quota_exceeded") + or classified_error.error_type == "rate_limit" ): cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( @@ -1563,11 +1560,8 @@ async def _streaming_acompletion_with_retry( ) raise last_exception - # Handle rate limits with cooldown - if classified_error.error_type in [ - "rate_limit", - "quota_exceeded", - ]: + # Handle rate limits with cooldown (exclude quota_exceeded) + if classified_error.error_type == "rate_limit": cooldown_duration = ( classified_error.retry_after or 60 ) @@ -1885,10 +1879,7 @@ async def _streaming_acompletion_with_retry( f"Cred {mask_credential(current_cred)} {classified_error.error_type}. Rotating." ) - if classified_error.error_type in [ - "rate_limit", - "quota_exceeded", - ]: + if classified_error.error_type == "rate_limit": cooldown_duration = ( classified_error.retry_after or 60 ) @@ -1980,11 +1971,10 @@ async def _streaming_acompletion_with_retry( f"Credential ...{current_cred[-6:]} failed with {classified_error.error_type} (Status: {classified_error.status_code}). Error: {error_message_text}." ) - # Handle rate limits with cooldown + # Handle rate limits with cooldown (exclude quota_exceeded) if ( - classified_error.status_code == 429 - or classified_error.error_type - in ["rate_limit", "quota_exceeded"] + (classified_error.status_code == 429 and classified_error.error_type != "quota_exceeded") + or classified_error.error_type == "rate_limit" ): cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index 76616c10..9605b05a 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -531,6 +531,15 @@ def classify_error(e: Exception) -> ClassifiedError: if isinstance(e, RateLimitError): retry_after = get_retry_after(e) + # Check if this is a quota error vs rate limit + error_msg = str(e).lower() + if "quota" in error_msg or "resource_exhausted" in error_msg: + return ClassifiedError( + error_type="quota_exceeded", + original_exception=e, + status_code=status_code or 429, + retry_after=retry_after, + ) return ClassifiedError( error_type="rate_limit", original_exception=e, diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 4ec2b825..71401463 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -677,11 +677,11 @@ async def record_failure( # Calculate cooldown duration based on error type cooldown_seconds = None - if classified_error.error_type == "rate_limit": - # Rate limit errors: use retry_after if available, otherwise default to 60s + if classified_error.error_type in ["rate_limit", "quota_exceeded"]: + # Rate limit / Quota errors: use retry_after if available, otherwise default to 60s cooldown_seconds = classified_error.retry_after or 60 lib_logger.info( - f"Rate limit error on key ...{key[-6:]} for model {model}. " + f"{classified_error.error_type} error on key ...{key[-6:]} for model {model}. " f"Using {'provided' if classified_error.retry_after else 'default'} retry_after: {cooldown_seconds}s" ) elif classified_error.error_type == "authentication": From 96e1b9763b9e18a0c1617e80ad526d241edb586e Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 4 Dec 2025 17:40:44 +0100 Subject: [PATCH 067/221] feat(provider): add support for Claude Opus 4.5 model Add claude-opus-4-5 to available models in Antigravity provider with proper mapping to thinking variant when reasoning_effort is provided. Changes: - Add claude-opus-4-5 to AVAILABLE_MODELS list - Update docstrings to include Claude Opus 4.5 - Extend reasoning_effort mapping to support both Sonnet and Opus models - Apply Claude tool schema transformation for claude-opus-* prefix Cherry-picked from PR #15 Co-authored-by: JoeGrimes123 --- src/rotator_library/providers/antigravity_provider.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 5751bba2..be5ef893 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -6,6 +6,7 @@ - Gemini 2.5 (Pro/Flash) with thinkingBudget - Gemini 3 (Pro/Image) with thinkingLevel - Claude (Sonnet 4.5) via Antigravity proxy +- Claude (Opus 4.5) via Antigravity proxy Key Features: - Unified streaming/non-streaming handling @@ -62,6 +63,7 @@ #"gemini-3-pro-image-preview", #"gemini-2.5-computer-use-preview-10-2025", "claude-sonnet-4-5", # Internally mapped to -thinking variant when reasoning_effort is provided + "claude-opus-4-5", # Internally mapped to -thinking variant when reasoning_effort is provided ] # Default max output tokens (including thinking) - can be overridden per request @@ -436,6 +438,7 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface): - Gemini 2.5 (Pro/Flash) with thinkingBudget - Gemini 3 (Pro/Image) with thinkingLevel - Claude Sonnet 4.5 via Antigravity proxy + - Claude Opus 4.5 via Antigravity proxy Features: - Unified streaming/non-streaming handling @@ -1993,8 +1996,8 @@ def _transform_to_antigravity_format( # Map base Claude model to -thinking variant when reasoning_effort is provided if self._is_claude(internal_model) and reasoning_effort: - if internal_model == "claude-sonnet-4-5" and not internal_model.endswith("-thinking"): - internal_model = "claude-sonnet-4-5-thinking" + if internal_model in ["claude-sonnet-4-5", "claude-opus-4-5"] and not internal_model.endswith("-thinking"): + internal_model = f"{internal_model}-thinking" # Map gemini-3-pro-preview to -low/-high variant based on thinking config if model == "gemini-3-pro-preview" or internal_model == "gemini-3-pro-preview": @@ -2070,7 +2073,7 @@ def _transform_to_antigravity_format( # Subsequent parallel calls: leave as-is (no signature) # Claude-specific tool schema transformation - if internal_model.startswith("claude-sonnet-"): + if internal_model.startswith("claude-sonnet-") or internal_model.startswith("claude-opus-"): self._apply_claude_tool_transform(antigravity_payload) return antigravity_payload From 08893078c20c73c60bc0396dd20c2f39d3bf118b Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Thu, 4 Dec 2025 17:51:52 +0100 Subject: [PATCH 068/221] =?UTF-8?q?docs:=20=F0=9F=93=9A=20update=20documen?= =?UTF-8?q?tation=20for=20Claude=20Opus=204.5=20support=20and=20fix=20giti?= =?UTF-8?q?gnore?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add comprehensive documentation for the newly supported Claude Opus 4.5 model via Antigravity provider across README and DOCUMENTATION files. - Document Claude Opus 4.5 as Anthropic's most powerful model now available through Antigravity - Add technical details about `claude-opus-4-5-thinking` internal model name - Highlight `thinkingBudget` parameter support and thinking preservation features - Update feature lists to emphasize Claude Opus 4.5 availability alongside existing Gemini 3 and Sonnet 4.5 - Generalize Claude-specific notes from "Sonnet 4.5" to "Claude" models for broader applicability - Fix `.gitignore` entry to correctly ignore `oauth_creds/` directory instead of file --- .gitignore | 2 +- DOCUMENTATION.md | 9 ++++++++- README.md | 9 ++++++--- 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/.gitignore b/.gitignore index 0c94208b..33e03301 100644 --- a/.gitignore +++ b/.gitignore @@ -129,4 +129,4 @@ logs/ cache/ *.env -oauth_creds +oauth_creds/ diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index b5a94938..29ea7838 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -391,7 +391,7 @@ A modular, shared caching system for providers to persist conversation state acr ### 3.5. Antigravity (`antigravity_provider.py`) -The most sophisticated provider implementation, supporting Google's internal Antigravity API for Gemini and Claude models. +The most sophisticated provider implementation, supporting Google's internal Antigravity API for Gemini and Claude models (including **Claude Opus 4.5**, Anthropic's most powerful model). #### Architecture @@ -418,6 +418,13 @@ The most sophisticated provider implementation, supporting Google's internal Ant - Automatic injection into functionCalls for multi-turn conversations - Fallback to bypass value if signature unavailable +**Claude Opus 4.5 (NEW!):** +- Anthropic's most powerful model, now available via Antigravity proxy +- Uses internal model name `claude-opus-4-5-thinking` when reasoning is enabled +- Uses `thinkingBudget` parameter for extended thinking control +- Full support for tool use with schema cleaning +- Same thinking preservation and sanitization features as Sonnet + **Claude Sonnet 4.5:** - Proxied through Antigravity API (uses internal model name `claude-sonnet-4-5-thinking`) - Uses `thinkingBudget` parameter like Gemini 2.5 diff --git a/README.md b/README.md index 51399bd2..91971102 100644 --- a/README.md +++ b/README.md @@ -28,11 +28,13 @@ This project provides a powerful solution for developers building complex applic - **OpenAI-Compatible Proxy**: Offers a familiar API interface with additional endpoints for model and provider discovery. - **Advanced Model Filtering**: Supports both blacklists and whitelists to give you fine-grained control over which models are available through the proxy. -- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 2.5, Gemini 3, and Claude Sonnet 4.5 models with advanced features: +- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 2.5, Gemini 3, and Claude models with advanced features: + - **🚀 NEW: Claude Opus 4.5** - Anthropic's most powerful model, now available via Antigravity! + - Claude Sonnet 4.5 with extended thinking support - Thought signature caching for multi-turn conversations - Tool hallucination prevention via parameter signature injection - Automatic thinking block sanitization for Claude models - - Note: Claude Sonnet 4.5 thinking mode requires careful conversation state management (see [Antigravity documentation](DOCUMENTATION.md#antigravity-claude-extended-thinking-sanitization) for details) + - Note: Claude thinking mode requires careful conversation state management (see [Antigravity documentation](DOCUMENTATION.md#antigravity-claude-extended-thinking-sanitization) for details) - **🆕 Credential Prioritization**: Automatic tier detection and priority-based credential selection ensures paid-tier credentials are used for premium models that require them. - **🆕 Weighted Random Rotation**: Configurable credential rotation strategy - choose between deterministic (perfect balance) or weighted random (unpredictable, harder to fingerprint) selection. - **🆕 Enhanced Gemini CLI**: Improved project discovery, paid vs free tier detection, and Gemini 3 support with thoughtSignature caching. @@ -504,12 +506,13 @@ The following advanced settings can be added to your `.env` file (or configured SKIP_OAUTH_INIT_CHECK=true -#### **Antigravity (Advanced - Gemini 3 \Claude 4.5 Access)** +#### **Antigravity (Advanced - Gemini 3 \ Claude Opus 4.5 / Sonnet 4.5 Access)** The newest and most sophisticated provider, offering access to cutting-edge models via Google's internal Antigravity API. **Supported Models:** - Gemini 2.5 (Pro/Flash) with `thinkingBudget` parameter - **Gemini 3 Pro (High/Low)** - Latest preview models +- **🆕 Claude Opus 4.5 + Thinking** - Anthropic's most powerful model via Antigravity proxy - **Claude Sonnet 4.5 + Thinking** via Antigravity proxy **Advanced Features:** From 8aec88be536078fa543946287fd403cc1dd125e3 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 00:14:09 +0100 Subject: [PATCH 069/221] =?UTF-8?q?fix(provider):=20=F0=9F=90=9B=20ensure?= =?UTF-8?q?=20claude-opus-4-5=20always=20uses=20thinking=20variant?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Update Claude Opus 4.5 model handling to always map to the -thinking variant, as the non-thinking variant does not exist. The logic now differentiates between Opus 4.5 (always thinking) and Sonnet 4.5 (thinking only when reasoning_effort is provided). - Increase DEFAULT_MAX_OUTPUT_TOKENS from 32384 to 64000 to accommodate thinking token output - Add explicit condition to always append -thinking suffix for claude-opus-4-5 - Refactor model mapping logic with detailed comments explaining variant selection - Update inline comment in AVAILABLE_MODELS to clarify Opus 4.5 behavior --- .../providers/antigravity_provider.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index be5ef893..731a10fa 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -63,11 +63,11 @@ #"gemini-3-pro-image-preview", #"gemini-2.5-computer-use-preview-10-2025", "claude-sonnet-4-5", # Internally mapped to -thinking variant when reasoning_effort is provided - "claude-opus-4-5", # Internally mapped to -thinking variant when reasoning_effort is provided + "claude-opus-4-5", # ALWAYS uses -thinking variant (non-thinking doesn't exist) ] # Default max output tokens (including thinking) - can be overridden per request -DEFAULT_MAX_OUTPUT_TOKENS = 32384 +DEFAULT_MAX_OUTPUT_TOKENS = 64000 # Model alias mappings (internal ↔ public) MODEL_ALIAS_MAP = { @@ -1994,10 +1994,16 @@ def _transform_to_antigravity_format( """ internal_model = self._alias_to_internal(model) - # Map base Claude model to -thinking variant when reasoning_effort is provided - if self._is_claude(internal_model) and reasoning_effort: - if internal_model in ["claude-sonnet-4-5", "claude-opus-4-5"] and not internal_model.endswith("-thinking"): - internal_model = f"{internal_model}-thinking" + # Map Claude models to their -thinking variant + # claude-opus-4-5: ALWAYS use -thinking (non-thinking variant doesn't exist) + # claude-sonnet-4-5: only use -thinking when reasoning_effort is provided + if self._is_claude(internal_model) and not internal_model.endswith("-thinking"): + if internal_model == "claude-opus-4-5": + # Opus 4.5 ALWAYS requires -thinking variant + internal_model = "claude-opus-4-5-thinking" + elif internal_model == "claude-sonnet-4-5" and reasoning_effort: + # Sonnet 4.5 uses -thinking only when reasoning_effort is provided + internal_model = "claude-sonnet-4-5-thinking" # Map gemini-3-pro-preview to -low/-high variant based on thinking config if model == "gemini-3-pro-preview" or internal_model == "gemini-3-pro-preview": From 1450294685b124a254d39a094506a7eb85cc82c9 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 02:42:26 +0100 Subject: [PATCH 070/221] refactor(antigravity-claude): refactor the claude sanitization logic to prevent errors on compaction, model switching, and allow thinking. --- .../providers/antigravity_provider.py | 2026 +++++++++++------ 1 file changed, 1286 insertions(+), 740 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 731a10fa..a1c66152 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -44,24 +44,24 @@ # CONFIGURATION CONSTANTS # ============================================================================= -lib_logger = logging.getLogger('rotator_library') +lib_logger = logging.getLogger("rotator_library") # Antigravity base URLs with fallback order # Priority: daily (sandbox) → autopush (sandbox) → production BASE_URLS = [ "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal", "https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal", - "https://cloudcode-pa.googleapis.com/v1internal", # Production fallback + "https://cloudcode-pa.googleapis.com/v1internal", # Production fallback ] # Available models via Antigravity AVAILABLE_MODELS = [ - #"gemini-2.5-pro", - #"gemini-2.5-flash", - #"gemini-2.5-flash-lite", + # "gemini-2.5-pro", + # "gemini-2.5-flash", + # "gemini-2.5-flash-lite", "gemini-3-pro-preview", # Internally mapped to -low/-high variant based on thinkingLevel - #"gemini-3-pro-image-preview", - #"gemini-2.5-computer-use-preview-10-2025", + # "gemini-3-pro-image-preview", + # "gemini-2.5-computer-use-preview-10-2025", "claude-sonnet-4-5", # Internally mapped to -thinking variant when reasoning_effort is provided "claude-opus-4-5", # ALWAYS uses -thinking variant (non-thinking doesn't exist) ] @@ -79,7 +79,12 @@ MODEL_ALIAS_REVERSE = {v: k for k, v in MODEL_ALIAS_MAP.items()} # Models to exclude from dynamic discovery -EXCLUDED_MODELS = {"chat_20706", "chat_23310", "gemini-2.5-flash-thinking", "gemini-2.5-pro"} +EXCLUDED_MODELS = { + "chat_20706", + "chat_23310", + "gemini-2.5-flash-thinking", + "gemini-2.5-pro", +} # Gemini finish reason mapping FINISH_REASON_MAP = { @@ -182,6 +187,7 @@ # HELPER FUNCTIONS # ============================================================================= + def _env_bool(key: str, default: bool = False) -> bool: """Get boolean from environment variable.""" return os.getenv(key, str(default).lower()).lower() in ("true", "1", "yes") @@ -232,10 +238,10 @@ def _normalize_type_arrays(schema: Any) -> Any: def _recursively_parse_json_strings(obj: Any) -> Any: """ Recursively parse JSON strings in nested data structures. - + Antigravity sometimes returns tool arguments with JSON-stringified values: {"files": "[{...}]"} instead of {"files": [{...}]}. - + Additionally handles: - Malformed double-encoded JSON (extra trailing '}' or ']') - Escaped string content (\n, \t, \", etc.) @@ -246,10 +252,10 @@ def _recursively_parse_json_strings(obj: Any) -> Any: return [_recursively_parse_json_strings(item) for item in obj] elif isinstance(obj, str): stripped = obj.strip() - + # Check if string contains common escape sequences that need unescaping # This handles cases where diff content or other text has literal \n instead of newlines - if '\\n' in obj or '\\t' in obj or '\\"' in obj or '\\\\' in obj: + if "\\n" in obj or "\\t" in obj or '\\"' in obj or "\\\\" in obj: try: # Use json.loads with quotes to properly unescape the string # This converts \n -> newline, \t -> tab, \" -> quote, etc. @@ -262,26 +268,27 @@ def _recursively_parse_json_strings(obj: Any) -> Any: except (json.JSONDecodeError, ValueError): # If unescaping fails, continue with original processing pass - + # Check if it looks like JSON (starts with { or [) - if stripped and stripped[0] in ('{', '['): + if stripped and stripped[0] in ("{", "["): # Try standard parsing first - if (stripped.startswith('{') and stripped.endswith('}')) or \ - (stripped.startswith('[') and stripped.endswith(']')): + if (stripped.startswith("{") and stripped.endswith("}")) or ( + stripped.startswith("[") and stripped.endswith("]") + ): try: parsed = json.loads(obj) return _recursively_parse_json_strings(parsed) except (json.JSONDecodeError, ValueError): pass - + # Handle malformed JSON: array that doesn't end with ] # e.g., '[{"path": "..."}]}' instead of '[{"path": "..."}]' - if stripped.startswith('[') and not stripped.endswith(']'): + if stripped.startswith("[") and not stripped.endswith("]"): try: # Find the last ] and truncate there - last_bracket = stripped.rfind(']') + last_bracket = stripped.rfind("]") if last_bracket > 0: - cleaned = stripped[:last_bracket+1] + cleaned = stripped[: last_bracket + 1] parsed = json.loads(cleaned) lib_logger.warning( f"[Antigravity] Auto-corrected malformed JSON string: " @@ -290,14 +297,14 @@ def _recursively_parse_json_strings(obj: Any) -> Any: return _recursively_parse_json_strings(parsed) except (json.JSONDecodeError, ValueError): pass - + # Handle malformed JSON: object that doesn't end with } - if stripped.startswith('{') and not stripped.endswith('}'): + if stripped.startswith("{") and not stripped.endswith("}"): try: # Find the last } and truncate there - last_brace = stripped.rfind('}') + last_brace = stripped.rfind("}") if last_brace > 0: - cleaned = stripped[:last_brace+1] + cleaned = stripped[: last_brace + 1] parsed = json.loads(cleaned) lib_logger.warning( f"[Antigravity] Auto-corrected malformed JSON string: " @@ -318,48 +325,73 @@ def _clean_claude_schema(schema: Any) -> Any: """ if not isinstance(schema, dict): return schema - + # Fields not supported by Antigravity/Google's Proto-based API # Note: Claude via Antigravity rejects JSON Schema draft 2020-12 validation keywords incompatible = { - '$schema', 'additionalProperties', 'minItems', 'maxItems', 'pattern', - 'minLength', 'maxLength', 'minimum', 'maximum', 'default', - 'exclusiveMinimum', 'exclusiveMaximum', 'multipleOf', 'format', - 'minProperties', 'maxProperties', 'uniqueItems', 'contentEncoding', - 'contentMediaType', 'contentSchema', 'deprecated', 'readOnly', 'writeOnly', - 'examples', '$id', '$ref', '$defs', 'definitions', 'title', + "$schema", + "additionalProperties", + "minItems", + "maxItems", + "pattern", + "minLength", + "maxLength", + "minimum", + "maximum", + "default", + "exclusiveMinimum", + "exclusiveMaximum", + "multipleOf", + "format", + "minProperties", + "maxProperties", + "uniqueItems", + "contentEncoding", + "contentMediaType", + "contentSchema", + "deprecated", + "readOnly", + "writeOnly", + "examples", + "$id", + "$ref", + "$defs", + "definitions", + "title", } - + # Handle 'anyOf' by taking the first option (Claude doesn't support anyOf) - if 'anyOf' in schema and isinstance(schema['anyOf'], list) and schema['anyOf']: - first_option = _clean_claude_schema(schema['anyOf'][0]) + if "anyOf" in schema and isinstance(schema["anyOf"], list) and schema["anyOf"]: + first_option = _clean_claude_schema(schema["anyOf"][0]) if isinstance(first_option, dict): return first_option - + # Handle 'oneOf' similarly - if 'oneOf' in schema and isinstance(schema['oneOf'], list) and schema['oneOf']: - first_option = _clean_claude_schema(schema['oneOf'][0]) + if "oneOf" in schema and isinstance(schema["oneOf"], list) and schema["oneOf"]: + first_option = _clean_claude_schema(schema["oneOf"][0]) if isinstance(first_option, dict): return first_option - cleaned = {} - + # Handle 'const' by converting to 'enum' with single value - if 'const' in schema: - const_value = schema['const'] - cleaned['enum'] = [const_value] - + if "const" in schema: + const_value = schema["const"] + cleaned["enum"] = [const_value] + for key, value in schema.items(): - if key in incompatible or key == 'const': + if key in incompatible or key == "const": continue if isinstance(value, dict): cleaned[key] = _clean_claude_schema(value) elif isinstance(value, list): - cleaned[key] = [_clean_claude_schema(item) if isinstance(item, dict) else item for item in value] + cleaned[key] = [ + _clean_claude_schema(item) if isinstance(item, dict) else item + for item in value + ] else: cleaned[key] = value - + return cleaned @@ -367,44 +399,47 @@ def _clean_claude_schema(schema: Any) -> Any: # FILE LOGGER # ============================================================================= + class AntigravityFileLogger: """Transaction file logger for debugging Antigravity requests/responses.""" - - __slots__ = ('enabled', 'log_dir') - + + __slots__ = ("enabled", "log_dir") + def __init__(self, model_name: str, enabled: bool = True): self.enabled = enabled self.log_dir: Optional[Path] = None - + if not enabled: return - + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") - safe_model = model_name.replace('/', '_').replace(':', '_') + safe_model = model_name.replace("/", "_").replace(":", "_") self.log_dir = LOGS_DIR / f"{timestamp}_{safe_model}_{uuid.uuid4()}" - + try: self.log_dir.mkdir(parents=True, exist_ok=True) except Exception as e: lib_logger.error(f"Failed to create log directory: {e}") self.enabled = False - + def log_request(self, payload: Dict[str, Any]) -> None: """Log the request payload.""" self._write_json("request_payload.json", payload) - + def log_response_chunk(self, chunk: str) -> None: """Append a raw chunk to the response stream log.""" self._append_text("response_stream.log", chunk) - + def log_error(self, error_message: str) -> None: """Log an error message.""" - self._append_text("error.log", f"[{datetime.utcnow().isoformat()}] {error_message}") - + self._append_text( + "error.log", f"[{datetime.utcnow().isoformat()}] {error_message}" + ) + def log_final_response(self, response: Dict[str, Any]) -> None: """Log the final response.""" self._write_json("final_response.json", response) - + def _write_json(self, filename: str, data: Dict[str, Any]) -> None: if not self.enabled or not self.log_dir: return @@ -413,7 +448,7 @@ def _write_json(self, filename: str, data: Dict[str, Any]) -> None: json.dump(data, f, indent=2, ensure_ascii=False) except Exception as e: lib_logger.error(f"Failed to write {filename}: {e}") - + def _append_text(self, filename: str, text: str) -> None: if not self.enabled or not self.log_dir: return @@ -424,88 +459,104 @@ def _append_text(self, filename: str, text: str) -> None: lib_logger.error(f"Failed to append to {filename}: {e}") - - # ============================================================================= # MAIN PROVIDER CLASS # ============================================================================= + class AntigravityProvider(AntigravityAuthBase, ProviderInterface): """ Antigravity provider for Gemini and Claude models via Google's internal API. - + Supports: - Gemini 2.5 (Pro/Flash) with thinkingBudget - - Gemini 3 (Pro/Image) with thinkingLevel + - Gemini 3 (Pro/Image) with thinkingLevel - Claude Sonnet 4.5 via Antigravity proxy - Claude Opus 4.5 via Antigravity proxy - + Features: - Unified streaming/non-streaming handling - ThoughtSignature caching for multi-turn conversations - Automatic base URL fallback - Gemini 3 tool hallucination prevention """ - + skip_cost_calculation = True - + def __init__(self): super().__init__() self.model_definitions = ModelDefinitions() - self.project_id_cache: Dict[str, str] = {} # Cache project ID per credential path - self.project_tier_cache: Dict[str, str] = {} # Cache project tier per credential path (for debugging) - + self.project_id_cache: Dict[ + str, str + ] = {} # Cache project ID per credential path + self.project_tier_cache: Dict[ + str, str + ] = {} # Cache project tier per credential path (for debugging) + # Base URL management self._base_url_index = 0 self._current_base_url = BASE_URLS[0] - + # Configuration from environment memory_ttl = _env_int("ANTIGRAVITY_SIGNATURE_CACHE_TTL", 3600) disk_ttl = _env_int("ANTIGRAVITY_SIGNATURE_DISK_TTL", 86400) - + # Initialize caches using shared ProviderCache self._signature_cache = ProviderCache( - GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl, - env_prefix="ANTIGRAVITY_SIGNATURE" + GEMINI3_SIGNATURE_CACHE_FILE, + memory_ttl, + disk_ttl, + env_prefix="ANTIGRAVITY_SIGNATURE", ) self._thinking_cache = ProviderCache( - CLAUDE_THINKING_CACHE_FILE, memory_ttl, disk_ttl, - env_prefix="ANTIGRAVITY_THINKING" + CLAUDE_THINKING_CACHE_FILE, + memory_ttl, + disk_ttl, + env_prefix="ANTIGRAVITY_THINKING", ) - + # Feature flags - self._preserve_signatures_in_client = _env_bool("ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES", True) - self._enable_signature_cache = _env_bool("ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", True) - self._enable_dynamic_models = _env_bool("ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", False) + self._preserve_signatures_in_client = _env_bool( + "ANTIGRAVITY_PRESERVE_THOUGHT_SIGNATURES", True + ) + self._enable_signature_cache = _env_bool( + "ANTIGRAVITY_ENABLE_SIGNATURE_CACHE", True + ) + self._enable_dynamic_models = _env_bool( + "ANTIGRAVITY_ENABLE_DYNAMIC_MODELS", False + ) self._enable_gemini3_tool_fix = _env_bool("ANTIGRAVITY_GEMINI3_TOOL_FIX", True) self._enable_claude_tool_fix = _env_bool("ANTIGRAVITY_CLAUDE_TOOL_FIX", True) - self._enable_thinking_sanitization = _env_bool("ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION", True) - + self._enable_thinking_sanitization = _env_bool( + "ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION", True + ) + # Gemini 3 tool fix configuration - self._gemini3_tool_prefix = os.getenv("ANTIGRAVITY_GEMINI3_TOOL_PREFIX", "gemini3_") + self._gemini3_tool_prefix = os.getenv( + "ANTIGRAVITY_GEMINI3_TOOL_PREFIX", "gemini3_" + ) self._gemini3_description_prompt = os.getenv( "ANTIGRAVITY_GEMINI3_DESCRIPTION_PROMPT", - "\n\n⚠️ STRICT PARAMETERS (use EXACTLY as shown): {params}. Do NOT use parameters from your training data - use ONLY these parameter names." + "\n\n⚠️ STRICT PARAMETERS (use EXACTLY as shown): {params}. Do NOT use parameters from your training data - use ONLY these parameter names.", + ) + self._gemini3_enforce_strict_schema = _env_bool( + "ANTIGRAVITY_GEMINI3_STRICT_SCHEMA", True ) - self._gemini3_enforce_strict_schema = _env_bool("ANTIGRAVITY_GEMINI3_STRICT_SCHEMA", True) self._gemini3_system_instruction = os.getenv( - "ANTIGRAVITY_GEMINI3_SYSTEM_INSTRUCTION", - DEFAULT_GEMINI3_SYSTEM_INSTRUCTION + "ANTIGRAVITY_GEMINI3_SYSTEM_INSTRUCTION", DEFAULT_GEMINI3_SYSTEM_INSTRUCTION ) - + # Claude tool fix configuration (separate from Gemini 3) self._claude_description_prompt = os.getenv( - "ANTIGRAVITY_CLAUDE_DESCRIPTION_PROMPT", - "\n\nSTRICT PARAMETERS: {params}." + "ANTIGRAVITY_CLAUDE_DESCRIPTION_PROMPT", "\n\nSTRICT PARAMETERS: {params}." ) self._claude_system_instruction = os.getenv( - "ANTIGRAVITY_CLAUDE_SYSTEM_INSTRUCTION", - DEFAULT_CLAUDE_SYSTEM_INSTRUCTION + "ANTIGRAVITY_CLAUDE_SYSTEM_INSTRUCTION", DEFAULT_CLAUDE_SYSTEM_INSTRUCTION ) - + # Log configuration self._log_config() - + def _log_config(self) -> None: """Log provider configuration.""" lib_logger.debug( @@ -514,42 +565,42 @@ def _log_config(self) -> None: f"gemini3_fix={self._enable_gemini3_tool_fix}, gemini3_strict_schema={self._gemini3_enforce_strict_schema}, " f"claude_fix={self._enable_claude_tool_fix}, thinking_sanitization={self._enable_thinking_sanitization}" ) - + # ========================================================================= # MODEL UTILITIES # ========================================================================= - + def _alias_to_internal(self, alias: str) -> str: """Convert public alias to internal model name.""" return MODEL_ALIAS_REVERSE.get(alias, alias) - + def _internal_to_alias(self, internal: str) -> str: """Convert internal model name to public alias.""" if internal in EXCLUDED_MODELS: return "" return MODEL_ALIAS_MAP.get(internal, internal) - + def _is_gemini_3(self, model: str) -> bool: """Check if model is Gemini 3 (requires special handling).""" internal = self._alias_to_internal(model) return internal.startswith("gemini-3-") or model.startswith("gemini-3-") - + def _is_claude(self, model: str) -> bool: """Check if model is Claude.""" return "claude" in model.lower() - + def _strip_provider_prefix(self, model: str) -> str: """Strip provider prefix from model name.""" return model.split("/")[-1] if "/" in model else model - + # ========================================================================= # BASE URL MANAGEMENT # ========================================================================= - + def _get_base_url(self) -> str: """Get current base URL.""" return self._current_base_url - + def _try_next_base_url(self) -> bool: """Switch to next base URL in fallback list. Returns True if successful.""" if self._base_url_index < len(BASE_URLS) - 1: @@ -558,49 +609,49 @@ def _try_next_base_url(self) -> bool: lib_logger.info(f"Switching to fallback URL: {self._current_base_url}") return True return False - + def _reset_base_url(self) -> None: """Reset to primary base URL.""" self._base_url_index = 0 self._current_base_url = BASE_URLS[0] - + # ========================================================================= # THINKING CACHE KEY GENERATION # ========================================================================= - + def _generate_thinking_cache_key( - self, - text_content: str, - tool_calls: List[Dict] + self, text_content: str, tool_calls: List[Dict] ) -> Optional[str]: """ Generate stable cache key from response content for Claude thinking preservation. - + Uses composite key: - Tool call IDs (most stable) - Text hash (for text-only responses) """ key_parts = [] - + if tool_calls: first_id = tool_calls[0].get("id", "") if first_id: key_parts.append(f"tool_{first_id.replace('call_', '')}") - + if text_content: text_hash = hashlib.md5(text_content[:200].encode()).hexdigest()[:16] key_parts.append(f"text_{text_hash}") - + return "thinking_" + "_".join(key_parts) if key_parts else None - + # ========================================================================= # PROJECT ID DISCOVERY # ========================================================================= - - async def _discover_project_id(self, credential_path: str, access_token: str, litellm_params: Dict[str, Any]) -> str: + + async def _discover_project_id( + self, credential_path: str, access_token: str, litellm_params: Dict[str, Any] + ) -> str: """ Discovers the Google Cloud Project ID, with caching and onboarding for new accounts. - + This follows the official Gemini CLI discovery flow adapted for Antigravity: 1. Check in-memory cache 2. Check configured project_id override (litellm_params or env var) @@ -610,11 +661,13 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li - If no currentTier: user needs onboarding 5. Onboard user (FREE tier: pass cloudaicompanionProject=None for server-managed) 6. Fallback to GCP Resource Manager project listing - + Note: Unlike GeminiCli, Antigravity doesn't use tier-based credential prioritization, but we still cache tier info for debugging and consistency. """ - lib_logger.debug(f"Starting Antigravity project discovery for credential: {credential_path}") + lib_logger.debug( + f"Starting Antigravity project discovery for credential: {credential_path}" + ) # Check in-memory cache first if credential_path in self.project_id_cache: @@ -624,12 +677,14 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li # Check for configured project ID override (from litellm_params or env var) configured_project_id = ( - litellm_params.get("project_id") or - os.getenv("ANTIGRAVITY_PROJECT_ID") or - os.getenv("GOOGLE_CLOUD_PROJECT") + litellm_params.get("project_id") + or os.getenv("ANTIGRAVITY_PROJECT_ID") + or os.getenv("GOOGLE_CLOUD_PROJECT") ) if configured_project_id: - lib_logger.debug(f"Found configured project_id override: {configured_project_id}") + lib_logger.debug( + f"Found configured project_id override: {configured_project_id}" + ) # Load credentials from file to check for persisted project_id and tier # Skip for env:// paths (environment-based credentials don't persist to files) @@ -637,28 +692,35 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li if credential_index is None: # Only try to load from file if it's not an env:// path try: - with open(credential_path, 'r') as f: + with open(credential_path, "r") as f: creds = json.load(f) - + metadata = creds.get("_proxy_metadata", {}) persisted_project_id = metadata.get("project_id") persisted_tier = metadata.get("tier") - + if persisted_project_id: - lib_logger.info(f"Loaded persisted project ID from credential file: {persisted_project_id}") + lib_logger.info( + f"Loaded persisted project ID from credential file: {persisted_project_id}" + ) self.project_id_cache[credential_path] = persisted_project_id - + # Also load tier if available (for debugging/logging purposes) if persisted_tier: self.project_tier_cache[credential_path] = persisted_tier lib_logger.debug(f"Loaded persisted tier: {persisted_tier}") - + return persisted_project_id except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: lib_logger.debug(f"Could not load persisted project ID from file: {e}") - lib_logger.debug("No cached or configured project ID found, initiating discovery...") - headers = {'Authorization': f'Bearer {access_token}', 'Content-Type': 'application/json'} + lib_logger.debug( + "No cached or configured project ID found, initiating discovery..." + ) + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + } discovered_project_id = None discovered_tier = None @@ -668,7 +730,9 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li async with httpx.AsyncClient() as client: # 1. Try discovery endpoint with loadCodeAssist - lib_logger.debug("Attempting project discovery via Code Assist loadCodeAssist endpoint...") + lib_logger.debug( + "Attempting project discovery via Code Assist loadCodeAssist endpoint..." + ) try: # Build metadata - include duetProject only if we have a configured project core_client_metadata = { @@ -678,53 +742,65 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li } if configured_project_id: core_client_metadata["duetProject"] = configured_project_id - + # Build load request - pass configured_project_id if available, otherwise None load_request = { "cloudaicompanionProject": configured_project_id, # Can be None "metadata": core_client_metadata, } - - lib_logger.debug(f"Sending loadCodeAssist request with cloudaicompanionProject={configured_project_id}") - response = await client.post(f"{code_assist_endpoint}:loadCodeAssist", headers=headers, json=load_request, timeout=20) + + lib_logger.debug( + f"Sending loadCodeAssist request with cloudaicompanionProject={configured_project_id}" + ) + response = await client.post( + f"{code_assist_endpoint}:loadCodeAssist", + headers=headers, + json=load_request, + timeout=20, + ) response.raise_for_status() data = response.json() # Log full response for debugging - lib_logger.debug(f"loadCodeAssist full response keys: {list(data.keys())}") + lib_logger.debug( + f"loadCodeAssist full response keys: {list(data.keys())}" + ) # Extract tier information - allowed_tiers = data.get('allowedTiers', []) - current_tier = data.get('currentTier') - + allowed_tiers = data.get("allowedTiers", []) + current_tier = data.get("currentTier") + lib_logger.debug(f"=== Tier Information ===") lib_logger.debug(f"currentTier: {current_tier}") lib_logger.debug(f"allowedTiers count: {len(allowed_tiers)}") for i, tier in enumerate(allowed_tiers): - tier_id = tier.get('id', 'unknown') - is_default = tier.get('isDefault', False) - user_defined = tier.get('userDefinedCloudaicompanionProject', False) - lib_logger.debug(f" Tier {i+1}: id={tier_id}, isDefault={is_default}, userDefinedProject={user_defined}") + tier_id = tier.get("id", "unknown") + is_default = tier.get("isDefault", False) + user_defined = tier.get("userDefinedCloudaicompanionProject", False) + lib_logger.debug( + f" Tier {i + 1}: id={tier_id}, isDefault={is_default}, userDefinedProject={user_defined}" + ) lib_logger.debug(f"========================") # Determine the current tier ID current_tier_id = None if current_tier: - current_tier_id = current_tier.get('id') + current_tier_id = current_tier.get("id") lib_logger.debug(f"User has currentTier: {current_tier_id}") # Check if user is already known to server (has currentTier) if current_tier_id: # User is already onboarded - check for project from server - server_project = data.get('cloudaicompanionProject') - + server_project = data.get("cloudaicompanionProject") + # Check if this tier requires user-defined project (paid tiers) requires_user_project = any( - t.get('id') == current_tier_id and t.get('userDefinedCloudaicompanionProject', False) + t.get("id") == current_tier_id + and t.get("userDefinedCloudaicompanionProject", False) for t in allowed_tiers ) - is_free_tier = current_tier_id == 'free-tier' - + is_free_tier = current_tier_id == "free-tier" + if server_project: # Server returned a project - use it (server wins) project_id = server_project @@ -732,10 +808,14 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li elif configured_project_id: # No server project but we have configured one - use it project_id = configured_project_id - lib_logger.debug(f"No server project, using configured: {project_id}") + lib_logger.debug( + f"No server project, using configured: {project_id}" + ) elif is_free_tier: # Free tier user without server project - try onboarding - lib_logger.debug("Free tier user with currentTier but no project - will try onboarding") + lib_logger.debug( + "Free tier user with currentTier but no project - will try onboarding" + ) project_id = None elif requires_user_project: # Paid tier requires a project ID to be set @@ -744,7 +824,9 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li ) else: # Unknown tier without project - proceed to onboarding - lib_logger.warning(f"Tier '{current_tier_id}' has no project and none configured - will try onboarding") + lib_logger.warning( + f"Tier '{current_tier_id}' has no project and none configured - will try onboarding" + ) project_id = None if project_id: @@ -753,52 +835,68 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li discovered_tier = current_tier_id # Log appropriately based on tier - is_paid = current_tier_id and current_tier_id not in ['free-tier', 'legacy-tier', 'unknown'] + is_paid = current_tier_id and current_tier_id not in [ + "free-tier", + "legacy-tier", + "unknown", + ] if is_paid: - lib_logger.info(f"Using Antigravity paid tier '{current_tier_id}' with project: {project_id}") + lib_logger.info( + f"Using Antigravity paid tier '{current_tier_id}' with project: {project_id}" + ) else: - lib_logger.info(f"Discovered Antigravity project ID via loadCodeAssist: {project_id}") + lib_logger.info( + f"Discovered Antigravity project ID via loadCodeAssist: {project_id}" + ) self.project_id_cache[credential_path] = project_id discovered_project_id = project_id - + # Persist to credential file - await self._persist_project_metadata(credential_path, project_id, discovered_tier) - + await self._persist_project_metadata( + credential_path, project_id, discovered_tier + ) + return project_id - + # 2. User needs onboarding - no currentTier or no project found - lib_logger.info("No existing Antigravity session found (no currentTier), attempting to onboard user...") - + lib_logger.info( + "No existing Antigravity session found (no currentTier), attempting to onboard user..." + ) + # Determine which tier to onboard with onboard_tier = None for tier in allowed_tiers: - if tier.get('isDefault'): + if tier.get("isDefault"): onboard_tier = tier break - + # Fallback to legacy tier if no default if not onboard_tier and allowed_tiers: for tier in allowed_tiers: - if tier.get('id') == 'legacy-tier': + if tier.get("id") == "legacy-tier": onboard_tier = tier break if not onboard_tier: onboard_tier = allowed_tiers[0] - + if not onboard_tier: raise ValueError("No onboarding tiers available from server") - - tier_id = onboard_tier.get('id', 'free-tier') - requires_user_project = onboard_tier.get('userDefinedCloudaicompanionProject', False) - - lib_logger.debug(f"Onboarding with tier: {tier_id}, requiresUserProject: {requires_user_project}") - + + tier_id = onboard_tier.get("id", "free-tier") + requires_user_project = onboard_tier.get( + "userDefinedCloudaicompanionProject", False + ) + + lib_logger.debug( + f"Onboarding with tier: {tier_id}, requiresUserProject: {requires_user_project}" + ) + # Build onboard request based on tier type # FREE tier: cloudaicompanionProject = None (server-managed) # PAID tier: cloudaicompanionProject = configured_project_id - is_free_tier = tier_id == 'free-tier' - + is_free_tier = tier_id == "free-tier" + if is_free_tier: # Free tier uses server-managed project onboard_request = { @@ -806,7 +904,9 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li "cloudaicompanionProject": None, # Server will create/manage "metadata": core_client_metadata, } - lib_logger.debug("Free tier onboarding: using server-managed project") + lib_logger.debug( + "Free tier onboarding: using server-managed project" + ) else: # Paid/legacy tier requires user-provided project if not configured_project_id and requires_user_project: @@ -819,52 +919,86 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li "metadata": { **core_client_metadata, "duetProject": configured_project_id, - } if configured_project_id else core_client_metadata, + } + if configured_project_id + else core_client_metadata, } - lib_logger.debug(f"Paid tier onboarding: using project {configured_project_id}") + lib_logger.debug( + f"Paid tier onboarding: using project {configured_project_id}" + ) lib_logger.debug("Initiating onboardUser request...") - lro_response = await client.post(f"{code_assist_endpoint}:onboardUser", headers=headers, json=onboard_request, timeout=30) + lro_response = await client.post( + f"{code_assist_endpoint}:onboardUser", + headers=headers, + json=onboard_request, + timeout=30, + ) lro_response.raise_for_status() lro_data = lro_response.json() - lib_logger.debug(f"Initial onboarding response: done={lro_data.get('done')}") + lib_logger.debug( + f"Initial onboarding response: done={lro_data.get('done')}" + ) # Poll for onboarding completion (up to 5 minutes) for i in range(150): # 150 × 2s = 5 minutes - if lro_data.get('done'): - lib_logger.debug(f"Onboarding completed after {i} polling attempts") + if lro_data.get("done"): + lib_logger.debug( + f"Onboarding completed after {i} polling attempts" + ) break await asyncio.sleep(2) if (i + 1) % 15 == 0: # Log every 30 seconds - lib_logger.info(f"Still waiting for onboarding completion... ({(i+1)*2}s elapsed)") - lib_logger.debug(f"Polling onboarding status... (Attempt {i+1}/150)") - lro_response = await client.post(f"{code_assist_endpoint}:onboardUser", headers=headers, json=onboard_request, timeout=30) + lib_logger.info( + f"Still waiting for onboarding completion... ({(i + 1) * 2}s elapsed)" + ) + lib_logger.debug( + f"Polling onboarding status... (Attempt {i + 1}/150)" + ) + lro_response = await client.post( + f"{code_assist_endpoint}:onboardUser", + headers=headers, + json=onboard_request, + timeout=30, + ) lro_response.raise_for_status() lro_data = lro_response.json() - if not lro_data.get('done'): + if not lro_data.get("done"): lib_logger.error("Onboarding process timed out after 5 minutes") - raise ValueError("Onboarding process timed out after 5 minutes. Please try again or contact support.") + raise ValueError( + "Onboarding process timed out after 5 minutes. Please try again or contact support." + ) # Extract project ID from LRO response # Note: onboardUser returns response.cloudaicompanionProject as an object with .id - lro_response_data = lro_data.get('response', {}) - lro_project_obj = lro_response_data.get('cloudaicompanionProject', {}) - project_id = lro_project_obj.get('id') if isinstance(lro_project_obj, dict) else None - + lro_response_data = lro_data.get("response", {}) + lro_project_obj = lro_response_data.get("cloudaicompanionProject", {}) + project_id = ( + lro_project_obj.get("id") + if isinstance(lro_project_obj, dict) + else None + ) + # Fallback to configured project if LRO didn't return one if not project_id and configured_project_id: project_id = configured_project_id - lib_logger.debug(f"LRO didn't return project, using configured: {project_id}") - + lib_logger.debug( + f"LRO didn't return project, using configured: {project_id}" + ) + if not project_id: - lib_logger.error("Onboarding completed but no project ID in response and none configured") + lib_logger.error( + "Onboarding completed but no project ID in response and none configured" + ) raise ValueError( "Onboarding completed, but no project ID was returned. " "For paid tiers, set ANTIGRAVITY_PROJECT_ID environment variable." ) - lib_logger.debug(f"Successfully extracted project ID from onboarding response: {project_id}") + lib_logger.debug( + f"Successfully extracted project ID from onboarding response: {project_id}" + ) # Cache tier info self.project_tier_cache[credential_path] = tier_id @@ -872,18 +1006,24 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li lib_logger.debug(f"Cached tier information: {tier_id}") # Log concise message based on tier - is_paid = tier_id and tier_id not in ['free-tier', 'legacy-tier'] + is_paid = tier_id and tier_id not in ["free-tier", "legacy-tier"] if is_paid: - lib_logger.info(f"Using Antigravity paid tier '{tier_id}' with project: {project_id}") + lib_logger.info( + f"Using Antigravity paid tier '{tier_id}' with project: {project_id}" + ) else: - lib_logger.info(f"Successfully onboarded user and discovered project ID: {project_id}") + lib_logger.info( + f"Successfully onboarded user and discovered project ID: {project_id}" + ) self.project_id_cache[credential_path] = project_id discovered_project_id = project_id - + # Persist to credential file - await self._persist_project_metadata(credential_path, project_id, discovered_tier) - + await self._persist_project_metadata( + credential_path, project_id, discovered_tier + ) + return project_id except httpx.HTTPStatusError as e: @@ -893,50 +1033,86 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li except Exception: pass if e.response.status_code == 403: - lib_logger.error(f"Antigravity Code Assist API access denied (403). Response: {error_body}") - lib_logger.error("Possible causes: 1) cloudaicompanion.googleapis.com API not enabled, 2) Wrong project ID for paid tier, 3) Account lacks permissions") + lib_logger.error( + f"Antigravity Code Assist API access denied (403). Response: {error_body}" + ) + lib_logger.error( + "Possible causes: 1) cloudaicompanion.googleapis.com API not enabled, 2) Wrong project ID for paid tier, 3) Account lacks permissions" + ) elif e.response.status_code == 404: - lib_logger.warning(f"Antigravity Code Assist endpoint not found (404). Falling back to project listing.") + lib_logger.warning( + f"Antigravity Code Assist endpoint not found (404). Falling back to project listing." + ) elif e.response.status_code == 412: # Precondition Failed - often means wrong project for free tier onboarding - lib_logger.error(f"Precondition failed (412): {error_body}. This may mean the project ID is incompatible with the selected tier.") + lib_logger.error( + f"Precondition failed (412): {error_body}. This may mean the project ID is incompatible with the selected tier." + ) else: - lib_logger.warning(f"Antigravity onboarding/discovery failed with status {e.response.status_code}: {error_body}. Falling back to project listing.") + lib_logger.warning( + f"Antigravity onboarding/discovery failed with status {e.response.status_code}: {error_body}. Falling back to project listing." + ) except httpx.RequestError as e: - lib_logger.warning(f"Antigravity onboarding/discovery network error: {e}. Falling back to project listing.") + lib_logger.warning( + f"Antigravity onboarding/discovery network error: {e}. Falling back to project listing." + ) # 3. Fallback to listing all available GCP projects (last resort) - lib_logger.debug("Attempting to discover project via GCP Resource Manager API...") + lib_logger.debug( + "Attempting to discover project via GCP Resource Manager API..." + ) try: async with httpx.AsyncClient() as client: - lib_logger.debug("Querying Cloud Resource Manager for available projects...") - response = await client.get("https://cloudresourcemanager.googleapis.com/v1/projects", headers=headers, timeout=20) + lib_logger.debug( + "Querying Cloud Resource Manager for available projects..." + ) + response = await client.get( + "https://cloudresourcemanager.googleapis.com/v1/projects", + headers=headers, + timeout=20, + ) response.raise_for_status() - projects = response.json().get('projects', []) + projects = response.json().get("projects", []) lib_logger.debug(f"Found {len(projects)} total projects") - active_projects = [p for p in projects if p.get('lifecycleState') == 'ACTIVE'] + active_projects = [ + p for p in projects if p.get("lifecycleState") == "ACTIVE" + ] lib_logger.debug(f"Found {len(active_projects)} active projects") if not projects: - lib_logger.error("No GCP projects found for this account. Please create a project in Google Cloud Console.") + lib_logger.error( + "No GCP projects found for this account. Please create a project in Google Cloud Console." + ) elif not active_projects: - lib_logger.error("No active GCP projects found. Please activate a project in Google Cloud Console.") + lib_logger.error( + "No active GCP projects found. Please activate a project in Google Cloud Console." + ) else: - project_id = active_projects[0]['projectId'] - lib_logger.info(f"Discovered Antigravity project ID from active projects list: {project_id}") - lib_logger.debug(f"Selected first active project: {project_id} (out of {len(active_projects)} active projects)") + project_id = active_projects[0]["projectId"] + lib_logger.info( + f"Discovered Antigravity project ID from active projects list: {project_id}" + ) + lib_logger.debug( + f"Selected first active project: {project_id} (out of {len(active_projects)} active projects)" + ) self.project_id_cache[credential_path] = project_id discovered_project_id = project_id - + # Persist to credential file (no tier info from resource manager) - await self._persist_project_metadata(credential_path, project_id, None) - + await self._persist_project_metadata( + credential_path, project_id, None + ) + return project_id except httpx.HTTPStatusError as e: if e.response.status_code == 403: - lib_logger.error("Failed to list GCP projects due to a 403 Forbidden error. The Cloud Resource Manager API may not be enabled, or your account lacks the 'resourcemanager.projects.list' permission.") + lib_logger.error( + "Failed to list GCP projects due to a 403 Forbidden error. The Cloud Resource Manager API may not be enabled, or your account lacks the 'resourcemanager.projects.list' permission." + ) else: - lib_logger.error(f"Failed to list GCP projects with status {e.response.status_code}: {e}") + lib_logger.error( + f"Failed to list GCP projects with status {e.response.status_code}: {e}" + ) except httpx.RequestError as e: lib_logger.error(f"Network error while listing GCP projects: {e}") @@ -947,20 +1123,24 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li " 3. Account lacks necessary permissions\n" "To manually specify a project, set ANTIGRAVITY_PROJECT_ID in your .env file." ) - - async def _persist_project_metadata(self, credential_path: str, project_id: str, tier: Optional[str]): + + async def _persist_project_metadata( + self, credential_path: str, project_id: str, tier: Optional[str] + ): """Persists project ID and tier to the credential file for faster future startups.""" # Skip persistence for env:// paths (environment-based credentials) credential_index = self._parse_env_credential_path(credential_path) if credential_index is not None: - lib_logger.debug(f"Skipping project metadata persistence for env:// credential path: {credential_path}") + lib_logger.debug( + f"Skipping project metadata persistence for env:// credential path: {credential_path}" + ) return - + try: # Load current credentials - with open(credential_path, 'r') as f: + with open(credential_path, "r") as f: creds = json.load(f) - + # Update metadata if "_proxy_metadata" not in creds: creds["_proxy_metadata"] = {} @@ -968,29 +1148,38 @@ async def _persist_project_metadata(self, credential_path: str, project_id: str, creds["_proxy_metadata"]["project_id"] = project_id if tier: creds["_proxy_metadata"]["tier"] = tier - + # Save back using the existing save method (handles atomic writes and permissions) await self._save_credentials(credential_path, creds) - - lib_logger.debug(f"Persisted project_id and tier to credential file: {credential_path}") + + lib_logger.debug( + f"Persisted project_id and tier to credential file: {credential_path}" + ) except Exception as e: - lib_logger.warning(f"Failed to persist project metadata to credential file: {e}") + lib_logger.warning( + f"Failed to persist project metadata to credential file: {e}" + ) # Non-fatal - just means slower startup next time # ========================================================================= # THINKING MODE SANITIZATION # ========================================================================= - + def _analyze_conversation_state( - self, - messages: List[Dict[str, Any]] + self, messages: List[Dict[str, Any]] ) -> Dict[str, Any]: """ Analyze conversation state to detect tool use loops and thinking mode issues. - + + Key insight: A "turn" can span multiple assistant messages in a tool-use loop. + We need to find the TURN START (first assistant message after last real user message) + and check if THAT message had thinking, not just the last assistant message. + Returns: { "in_tool_loop": bool - True if we're in an incomplete tool use loop + "turn_start_idx": int - Index of first assistant message in current turn + "turn_has_thinking": bool - Whether the TURN started with thinking "last_assistant_idx": int - Index of last assistant message "last_assistant_has_thinking": bool - Whether last assistant msg has thinking "last_assistant_has_tool_calls": bool - Whether last assistant msg has tool calls @@ -1000,73 +1189,112 @@ def _analyze_conversation_state( """ state = { "in_tool_loop": False, + "turn_start_idx": -1, + "turn_has_thinking": False, "last_assistant_idx": -1, "last_assistant_has_thinking": False, "last_assistant_has_tool_calls": False, "pending_tool_results": False, "thinking_block_indices": [], } - - # Find last assistant message and analyze the conversation + + # First pass: Find the last "real" user message (not a tool result) + # A real user message is one that doesn't immediately follow an assistant with tool_calls + last_real_user_idx = -1 for i, msg in enumerate(messages): role = msg.get("role") - - if role == "assistant": - state["last_assistant_idx"] = i - state["last_assistant_has_tool_calls"] = bool(msg.get("tool_calls")) - # Check for thinking/reasoning content - has_thinking = bool(msg.get("reasoning_content")) - # Also check for thinking in content array (some formats) + if role == "user": + # Check if this is a real user message or just follows tool results + # Tool messages have role="tool", so if this is role="user" and + # it's not just a tool_result container, it's a real user message. + # However, we need to be careful: the client might format tool results + # as user messages with tool_result content. Check the content. content = msg.get("content") + + # If content is a list with tool_result items, it's a tool response + is_tool_result_msg = False if isinstance(content, list): for item in content: - if isinstance(item, dict) and item.get("type") == "thinking": - has_thinking = True + if isinstance(item, dict) and item.get("type") == "tool_result": + is_tool_result_msg = True break + + if not is_tool_result_msg: + last_real_user_idx = i + + # Second pass: Analyze conversation and find turn boundaries + for i, msg in enumerate(messages): + role = msg.get("role") + + if role == "assistant": + # Check for thinking/reasoning content + has_thinking = self._message_has_thinking(msg) + + # Track if this is the turn start + if i > last_real_user_idx and state["turn_start_idx"] == -1: + state["turn_start_idx"] = i + state["turn_has_thinking"] = has_thinking + + state["last_assistant_idx"] = i + state["last_assistant_has_tool_calls"] = bool(msg.get("tool_calls")) state["last_assistant_has_thinking"] = has_thinking + if has_thinking: state["thinking_block_indices"].append(i) + elif role == "tool": # Tool result after an assistant message with tool calls = in tool loop if state["last_assistant_has_tool_calls"]: state["pending_tool_results"] = True - + # We're in a tool loop if: - # 1. Last assistant message had tool calls - # 2. There are tool results after it - # 3. There's no final text response yet (the conversation ends with tool results) + # 1. There are pending tool results + # 2. The conversation ends with tool results (last message is "tool" role) if state["pending_tool_results"] and messages: last_msg = messages[-1] if last_msg.get("role") == "tool": state["in_tool_loop"] = True - + return state - + + def _message_has_thinking(self, msg: Dict[str, Any]) -> bool: + """Check if an assistant message contains thinking/reasoning content.""" + # Check reasoning_content field (OpenAI format) + if msg.get("reasoning_content"): + return True + + # Check for thinking in content array (some formats) + content = msg.get("content") + if isinstance(content, list): + for item in content: + if isinstance(item, dict) and item.get("type") == "thinking": + return True + + return False + def _sanitize_thinking_for_claude( - self, - messages: List[Dict[str, Any]], - thinking_enabled: bool + self, messages: List[Dict[str, Any]], thinking_enabled: bool ) -> Tuple[List[Dict[str, Any]], bool]: """ Sanitize thinking blocks in conversation history for Claude compatibility. - + Handles the following scenarios per Claude docs: 1. If thinking is disabled, remove all thinking blocks from conversation 2. If thinking is enabled: a. In a tool use loop WITH thinking: preserve it (same mode continues) b. In a tool use loop WITHOUT thinking: this is INVALID toggle - force disable c. Not in tool loop: strip old thinking, new response adds thinking naturally - + Per Claude docs: - "If thinking is enabled, the final assistant turn must start with a thinking block" - "If thinking is disabled, the final assistant turn must not contain any thinking blocks" - Tool use loops are part of a single assistant turn - You CANNOT toggle thinking mid-turn - + The key insight: We only force-disable thinking when TOGGLING it ON mid-turn. If thinking was already enabled (assistant has thinking), we preserve. If thinking was disabled (assistant has no thinking), enabling it now is invalid. - + Returns: Tuple of (sanitized_messages, force_disable_thinking) - sanitized_messages: The cleaned message list @@ -1074,86 +1302,179 @@ def _sanitize_thinking_for_claude( """ messages = copy.deepcopy(messages) state = self._analyze_conversation_state(messages) - + lib_logger.debug( f"[Thinking Sanitization] thinking_enabled={thinking_enabled}, " f"in_tool_loop={state['in_tool_loop']}, " + f"turn_has_thinking={state['turn_has_thinking']}, " + f"turn_start_idx={state['turn_start_idx']}, " f"last_assistant_has_thinking={state['last_assistant_has_thinking']}, " f"last_assistant_has_tool_calls={state['last_assistant_has_tool_calls']}" ) - + if not thinking_enabled: # CASE 1: Thinking is disabled - strip ALL thinking blocks return self._strip_all_thinking_blocks(messages), False - + # CASE 2: Thinking is enabled if state["in_tool_loop"]: # We're in a tool use loop (conversation ends with tool_result) # Per Claude docs: entire assistant turn must operate in single thinking mode - - if state["last_assistant_has_thinking"]: - # Last assistant turn HAD thinking - this is valid! + # + # KEY FIX: Check turn_has_thinking (thinking at turn START), not last_assistant_has_thinking. + # In multi-message tool loops, thinking is at the FIRST assistant message of the turn, + # not necessarily the last one (which might just have tool_calls). + + if state["turn_has_thinking"]: + # The TURN started with thinking - this is valid! # Thinking was enabled when tool was called, continue with thinking enabled. - # Only keep thinking for the current turn (last assistant + following tools) + # Preserve thinking for the turn start message. lib_logger.debug( - "[Thinking Sanitization] Tool loop with existing thinking - preserving." + "[Thinking Sanitization] Tool loop with thinking at turn start - preserving. " + f"turn_start_idx={state['turn_start_idx']}, last_assistant_idx={state['last_assistant_idx']}" ) - return self._preserve_current_turn_thinking( - messages, state["last_assistant_idx"] + return self._preserve_turn_start_thinking( + messages, state["turn_start_idx"] ), False else: - # Last assistant turn DID NOT have thinking, but thinking is NOW enabled + # The TURN did NOT start with thinking, but thinking is NOW enabled # This is the INVALID case: toggling thinking ON mid-turn - # + # # Per Claude docs, this causes: # "Expected `thinking` or `redacted_thinking`, but found `tool_use`." # - # SOLUTION: Inject a synthetic assistant message to CLOSE the tool loop. - # This allows Claude to start a fresh turn WITH thinking. - # - # The synthetic message summarizes the tool results, allowing the model - # to respond naturally with thinking enabled on what is now a "new" turn. - lib_logger.info( - "[Thinking Sanitization] Closing tool loop with synthetic response. " - "This allows thinking to be enabled on the new turn." + # There are TWO possible scenarios: + # 1. Original turn was made WITHOUT thinking (e.g., by Gemini or non-thinking Claude) + # → Solution: Close the tool loop with synthetic message + # 2. Original turn HAD thinking but compaction stripped it + # → Solution: Try to inject cached thinking, fallback to synthetic closure + + turn_start_msg = ( + messages[state["turn_start_idx"]] + if state["turn_start_idx"] >= 0 + else None ) - return self._close_tool_loop_for_thinking(messages), False + + # Check if this looks like a compacted thinking turn + if turn_start_msg and self._looks_like_compacted_thinking_turn( + turn_start_msg + ): + # Try to recover cached thinking block + recovered = self._try_recover_thinking_from_cache( + messages, state["turn_start_idx"] + ) + if recovered: + lib_logger.info( + "[Thinking Sanitization] Recovered thinking from cache for compacted turn." + ) + return self._preserve_turn_start_thinking( + messages, state["turn_start_idx"] + ), False + else: + # Can't recover from cache - close the loop with synthetic messages + # This allows Claude to start a fresh turn with thinking + lib_logger.info( + "[Thinking Sanitization] Compacted thinking turn detected in tool loop. " + "Cache miss - closing loop with synthetic messages to enable fresh thinking turn." + ) + return self._close_tool_loop_for_thinking(messages), False + else: + # Not a compacted turn - genuinely no thinking. Close the loop. + lib_logger.info( + "[Thinking Sanitization] Closing tool loop with synthetic response. " + "Turn did not start with thinking (turn_has_thinking=False). " + "This allows thinking to be enabled on the new turn." + ) + return self._close_tool_loop_for_thinking(messages), False else: # Not in a tool loop - this is the simple case # The conversation doesn't end with tool_result, so we're starting fresh. - # Strip thinking from old turns (API ignores them anyway). - # New response will include thinking naturally. - - if state["last_assistant_idx"] >= 0 and not state["last_assistant_has_thinking"]: - if state["last_assistant_has_tool_calls"]: - # Last assistant made tool calls but no thinking - # This could be from context compression, model switch, or - # the assistant responded after tool results (completing the turn) - lib_logger.debug( - "[Thinking Sanitization] Last assistant has completed tool_calls but no thinking. " - "This is likely from context compression or completed tool loop. " - "New response will include thinking." + # + # HOWEVER, there's a special case: compaction might have removed the thinking + # block from the turn start, but Claude still expects it. + # We detect this by checking if there's an assistant message with tool_calls + # but no thinking, and the conversation structure suggests thinking was expected. + + # Check if we need to inject a fake thinking block for compaction recovery + if state["last_assistant_idx"] >= 0: + last_assistant = messages[state["last_assistant_idx"]] + + if ( + state["last_assistant_has_tool_calls"] + and not state["turn_has_thinking"] + ): + # The turn has tool_calls but no thinking at turn start. + # This could be: + # 1. Compaction removed the thinking block + # 2. The original call was made without thinking + # + # For case 1, we need to close the turn and start fresh. + # For case 2, we let the model respond naturally. + # + # We can detect case 1 if there's evidence thinking was expected: + # - The turn_start message has tool_calls (typical thinking-enabled flow) + # - The content structure suggests a thinking block was stripped + + # Check if turn_start has the hallmarks of a compacted thinking response + turn_start_msg = ( + messages[state["turn_start_idx"]] + if state["turn_start_idx"] >= 0 + else None ) - + if turn_start_msg and self._looks_like_compacted_thinking_turn( + turn_start_msg + ): + # Try cache recovery first + recovered = self._try_recover_thinking_from_cache( + messages, state["turn_start_idx"] + ) + if recovered: + lib_logger.info( + "[Thinking Sanitization] Recovered thinking from cache for compacted turn (not in tool loop)." + ) + return self._strip_old_turn_thinking( + messages, state["turn_start_idx"] + ), False + else: + # Can't recover - add synthetic user to start fresh turn + lib_logger.info( + "[Thinking Sanitization] Detected compacted turn missing thinking block. " + "Adding synthetic user message to start fresh thinking turn." + ) + # Add synthetic user message to trigger new turn with thinking + synthetic_user = {"role": "user", "content": "[Continue]"} + messages.append(synthetic_user) + return self._strip_all_thinking_blocks(messages), False + else: + lib_logger.debug( + "[Thinking Sanitization] Last assistant has tool_calls but no thinking. " + "This is likely from context compression or non-thinking model. " + "New response will include thinking naturally." + ) + # Strip thinking from old turns, let new response add thinking naturally - return self._strip_old_turn_thinking(messages, state["last_assistant_idx"]), False - + return self._strip_old_turn_thinking( + messages, state["last_assistant_idx"] + ), False + def _strip_all_thinking_blocks( - self, - messages: List[Dict[str, Any]] + self, messages: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: """Remove all thinking/reasoning content from messages.""" for msg in messages: if msg.get("role") == "assistant": # Remove reasoning_content field msg.pop("reasoning_content", None) - + # Remove thinking blocks from content array content = msg.get("content") if isinstance(content, list): filtered = [ - item for item in content - if not (isinstance(item, dict) and item.get("type") == "thinking") + item + for item in content + if not ( + isinstance(item, dict) and item.get("type") == "thinking" + ) ] # If filtering leaves empty list, we need to preserve message structure # to maintain user/assistant alternation. Use empty string as placeholder @@ -1163,19 +1484,19 @@ def _strip_all_thinking_blocks( if not msg.get("tool_calls"): msg["content"] = "" else: - msg["content"] = None # tool_calls exist, content not needed + msg["content"] = ( + None # tool_calls exist, content not needed + ) else: msg["content"] = filtered return messages - + def _strip_old_turn_thinking( - self, - messages: List[Dict[str, Any]], - last_assistant_idx: int + self, messages: List[Dict[str, Any]], last_assistant_idx: int ) -> List[Dict[str, Any]]: """ Strip thinking from old turns but preserve for the last assistant turn. - + Per Claude docs: "thinking blocks from previous turns are removed from context" This mimics the API behavior and prevents issues. """ @@ -1186,8 +1507,11 @@ def _strip_old_turn_thinking( content = msg.get("content") if isinstance(content, list): filtered = [ - item for item in content - if not (isinstance(item, dict) and item.get("type") == "thinking") + item + for item in content + if not ( + isinstance(item, dict) and item.get("type") == "thinking" + ) ] # Preserve message structure with empty string if needed if not filtered: @@ -1195,11 +1519,9 @@ def _strip_old_turn_thinking( else: msg["content"] = filtered return messages - + def _preserve_current_turn_thinking( - self, - messages: List[Dict[str, Any]], - last_assistant_idx: int + self, messages: List[Dict[str, Any]], last_assistant_idx: int ) -> List[Dict[str, Any]]: """ Preserve thinking only for the current (last) assistant turn. @@ -1207,29 +1529,169 @@ def _preserve_current_turn_thinking( """ # Same as strip_old_turn_thinking - we keep the last turn intact return self._strip_old_turn_thinking(messages, last_assistant_idx) - + + def _preserve_turn_start_thinking( + self, messages: List[Dict[str, Any]], turn_start_idx: int + ) -> List[Dict[str, Any]]: + """ + Preserve thinking at the turn start message. + + In multi-message tool loops, the thinking block is at the FIRST assistant + message of the turn (turn_start_idx), not the last one. We need to preserve + thinking from the turn start, and strip it from all older turns. + """ + for i, msg in enumerate(messages): + if msg.get("role") == "assistant" and i < turn_start_idx: + # Old turn - strip thinking + msg.pop("reasoning_content", None) + content = msg.get("content") + if isinstance(content, list): + filtered = [ + item + for item in content + if not ( + isinstance(item, dict) and item.get("type") == "thinking" + ) + ] + if not filtered: + msg["content"] = "" if not msg.get("tool_calls") else None + else: + msg["content"] = filtered + return messages + + def _looks_like_compacted_thinking_turn(self, msg: Dict[str, Any]) -> bool: + """ + Detect if a message looks like it was compacted from a thinking-enabled turn. + + Heuristics: + 1. Has tool_calls (typical thinking flow produces tool calls) + 2. Content structure suggests stripped thinking (e.g., starts with tool_use directly) + 3. No text content before tool_use (thinking responses usually have text) + + This is imperfect but helps catch common compaction scenarios. + """ + if not msg.get("tool_calls"): + return False + + content = msg.get("content") + + # If content is just tool_use blocks with no text, it might be compacted + if isinstance(content, list): + has_text = any( + isinstance(item, dict) + and item.get("type") == "text" + and item.get("text", "").strip() + for item in content + ) + has_tool_use = any( + isinstance(item, dict) and item.get("type") == "tool_use" + for item in content + ) + + # Typical compacted thinking: tool_use without preceding text + # Normal non-thinking response would have explanatory text + if has_tool_use and not has_text: + return True + + # If content is empty/None but has tool_calls, likely compacted + if not content and msg.get("tool_calls"): + return True + + return False + + def _try_recover_thinking_from_cache( + self, messages: List[Dict[str, Any]], turn_start_idx: int + ) -> bool: + """ + Try to recover thinking content from cache for a compacted turn. + + Returns True if thinking was successfully recovered and injected, False otherwise. + """ + if turn_start_idx < 0 or turn_start_idx >= len(messages): + return False + + msg = messages[turn_start_idx] + + # Extract tool_calls for cache key lookup + tool_calls = msg.get("tool_calls", []) + content = msg.get("content", "") + text_content = content if isinstance(content, str) else "" + + # Generate cache key and try to retrieve + cache_key = self._generate_thinking_cache_key(text_content, tool_calls) + if not cache_key: + return False + + cached_json = self._thinking_cache.retrieve(cache_key) + if not cached_json: + lib_logger.debug( + f"[Thinking Sanitization] No cached thinking found for key: {cache_key}" + ) + return False + + try: + thinking_data = json.loads(cached_json) + thinking_text = thinking_data.get("thinking_text", "") + signature = thinking_data.get("thought_signature", "") + + if not thinking_text or not signature: + lib_logger.debug( + "[Thinking Sanitization] Cached thinking missing text or signature" + ) + return False + + # Inject the recovered thinking block + thinking_block = { + "type": "thinking", + "thinking": thinking_text, + "signature": signature, + } + + if isinstance(content, list): + msg["content"] = [thinking_block] + content + elif isinstance(content, str): + msg["content"] = [thinking_block, {"type": "text", "text": content}] + else: + msg["content"] = [thinking_block] + + lib_logger.debug( + f"[Thinking Sanitization] Recovered thinking from cache: {len(thinking_text)} chars" + ) + return True + + except json.JSONDecodeError: + lib_logger.warning( + f"[Thinking Sanitization] Failed to parse cached thinking" + ) + return False + def _close_tool_loop_for_thinking( - self, - messages: List[Dict[str, Any]] + self, messages: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: """ - Close an incomplete tool loop by injecting a synthetic assistant response. - + Close an incomplete tool loop by injecting synthetic messages to start a new turn. + This is used when: - We're in a tool loop (conversation ends with tool_result) - - The tool call was made WITHOUT thinking (e.g., by Gemini or non-thinking Claude) + - The tool call was made WITHOUT thinking (e.g., by Gemini, non-thinking Claude, or compaction stripped it) - We NOW want to enable thinking - - By injecting a synthetic response that "closes" the previous turn, - Claude can start a fresh turn with thinking enabled. - - The synthetic message is minimal and factual - it just acknowledges - the tool results were received, allowing the model to process them - with thinking on the new turn. + + Per Claude docs on toggling thinking modes: + - "If thinking is enabled, the final assistant turn must start with a thinking block" + - "To toggle thinking, you must complete the assistant turn first" + - A non-tool-result user message ends the turn and allows a fresh start + + Solution: + 1. Add synthetic ASSISTANT message to complete the non-thinking turn + 2. Add synthetic USER message to start a NEW turn + 3. Claude will generate thinking for its response to the new turn + + The synthetic messages are minimal and unobtrusive - they just satisfy the + turn structure requirements without influencing model behavior. """ # Strip any old thinking first messages = self._strip_all_thinking_blocks(messages) - + # Collect tool results from the end of the conversation tool_results = [] for msg in reversed(messages): @@ -1237,9 +1699,9 @@ def _close_tool_loop_for_thinking( tool_results.append(msg) elif msg.get("role") == "assistant": break # Stop at the assistant that made the tool calls - + tool_results.reverse() # Put back in order - + # Safety check: if no tool results found, this shouldn't have been called # But handle gracefully with a generic message if not tool_results: @@ -1247,38 +1709,45 @@ def _close_tool_loop_for_thinking( "[Thinking Sanitization] _close_tool_loop_for_thinking called but no tool results found. " "This may indicate malformed conversation history." ) - synthetic_content = "[Processing previous context.]" + synthetic_assistant_content = "[Processing previous context.]" elif len(tool_results) == 1: - synthetic_content = "[Tool execution completed. Processing results.]" + synthetic_assistant_content = "[Tool execution completed.]" else: - synthetic_content = f"[{len(tool_results)} tool executions completed. Processing results.]" - - # Inject the synthetic assistant message to close the loop - synthetic_msg = { + synthetic_assistant_content = ( + f"[{len(tool_results)} tool executions completed.]" + ) + + # Step 1: Inject synthetic ASSISTANT message to complete the non-thinking turn + synthetic_assistant = { "role": "assistant", - "content": synthetic_content + "content": synthetic_assistant_content, } - messages.append(synthetic_msg) - - lib_logger.debug( - f"[Thinking Sanitization] Injected synthetic closure: '{synthetic_content}'" + messages.append(synthetic_assistant) + + # Step 2: Inject synthetic USER message to start a NEW turn + # This allows Claude to generate thinking for its response + # The message is minimal and unobtrusive - just triggers a new turn + synthetic_user = {"role": "user", "content": "[Continue]"} + messages.append(synthetic_user) + + lib_logger.info( + f"[Thinking Sanitization] Closed tool loop with synthetic messages. " + f"Assistant: '{synthetic_assistant_content}', User: '[Continue]'. " + f"Claude will now start a fresh turn with thinking enabled." ) - + return messages - + # ========================================================================= # REASONING CONFIGURATION # ========================================================================= - + def _get_thinking_config( - self, - reasoning_effort: Optional[str], - model: str, - custom_budget: bool = False + self, reasoning_effort: Optional[str], model: str, custom_budget: bool = False ) -> Optional[Dict[str, Any]]: """ Map reasoning_effort to thinking configuration. - + - Gemini 2.5 & Claude: thinkingBudget (integer tokens) - Gemini 3: thinkingLevel (string: "low"/"high") """ @@ -1286,23 +1755,23 @@ def _get_thinking_config( is_gemini_25 = "gemini-2.5" in model is_gemini_3 = internal.startswith("gemini-3-") is_claude = self._is_claude(model) - + if not (is_gemini_25 or is_gemini_3 or is_claude): return None - + # Gemini 3: String-based thinkingLevel if is_gemini_3: if reasoning_effort == "low": return {"thinkingLevel": "low", "include_thoughts": True} return {"thinkingLevel": "high", "include_thoughts": True} - + # Gemini 2.5 & Claude: Integer thinkingBudget if not reasoning_effort: return {"thinkingBudget": -1, "include_thoughts": True} # Auto - + if reasoning_effort == "disable": return {"thinkingBudget": 0, "include_thoughts": False} - + # Model-specific budgets if "gemini-2.5-pro" in model or is_claude: budgets = {"low": 8192, "medium": 16384, "high": 32768} @@ -1310,25 +1779,23 @@ def _get_thinking_config( budgets = {"low": 6144, "medium": 12288, "high": 24576} else: budgets = {"low": 1024, "medium": 2048, "high": 4096} - + budget = budgets.get(reasoning_effort, -1) if not custom_budget: budget = budget // 4 # Default to 25% of max output tokens - + return {"thinkingBudget": budget, "include_thoughts": True} - + # ========================================================================= # MESSAGE TRANSFORMATION (OpenAI → Gemini) # ========================================================================= - + def _transform_messages( - self, - messages: List[Dict[str, Any]], - model: str + self, messages: List[Dict[str, Any]], model: str ) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: """ Transform OpenAI messages to Gemini CLI format. - + Handles: - System instruction extraction - Multi-part content (text, images) @@ -1339,15 +1806,17 @@ def _transform_messages( messages = copy.deepcopy(messages) system_instruction = None gemini_contents = [] - + # Extract system prompt - if messages and messages[0].get('role') == 'system': - system_content = messages.pop(0).get('content', '') + if messages and messages[0].get("role") == "system": + system_content = messages.pop(0).get("content", "") if system_content: - system_parts = self._parse_content_parts(system_content, _strip_cache_control=True) + system_parts = self._parse_content_parts( + system_content, _strip_cache_control=True + ) if system_parts: system_instruction = {"role": "user", "parts": system_parts} - + # Build tool_call_id → name mapping tool_id_to_name = {} for msg in messages: @@ -1357,22 +1826,22 @@ def _transform_messages( tc_id = tc["id"] tc_name = tc["function"]["name"] tool_id_to_name[tc_id] = tc_name - #lib_logger.debug(f"[ID Mapping] Registered tool_call: id={tc_id}, name={tc_name}") - + # lib_logger.debug(f"[ID Mapping] Registered tool_call: id={tc_id}, name={tc_name}") + # Convert each message, consolidating consecutive tool responses # Per Gemini docs: parallel function responses must be in a single user message pending_tool_parts = [] - + for msg in messages: role = msg.get("role") content = msg.get("content") parts = [] - + # Flush pending tool parts before non-tool message if pending_tool_parts and role != "tool": gemini_contents.append({"role": "user", "parts": pending_tool_parts}) pending_tool_parts = [] - + if role == "user": parts = self._transform_user_message(content) elif role == "assistant": @@ -1382,25 +1851,23 @@ def _transform_messages( # Accumulate tool responses instead of adding individually pending_tool_parts.extend(tool_parts) continue - + if parts: gemini_role = "model" if role == "assistant" else "user" gemini_contents.append({"role": gemini_role, "parts": parts}) - + # Flush any remaining tool parts if pending_tool_parts: gemini_contents.append({"role": "user", "parts": pending_tool_parts}) - + return system_instruction, gemini_contents - + def _parse_content_parts( - self, - content: Any, - _strip_cache_control: bool = False + self, content: Any, _strip_cache_control: bool = False ) -> List[Dict[str, Any]]: """Parse content into Gemini parts format.""" parts = [] - + if isinstance(content, str): if content: parts.append({"text": content}) @@ -1414,15 +1881,15 @@ def _parse_content_parts( image_part = self._parse_image_url(item.get("image_url", {})) if image_part: parts.append(image_part) - + return parts - + def _parse_image_url(self, image_url: Dict[str, Any]) -> Optional[Dict[str, Any]]: """Parse image URL into Gemini inlineData format.""" url = image_url.get("url", "") if not url.startswith("data:"): return None - + try: header, data = url.split(",", 1) mime_type = header.split(":")[1].split(";")[0] @@ -1430,23 +1897,20 @@ def _parse_image_url(self, image_url: Dict[str, Any]) -> Optional[Dict[str, Any] except Exception as e: lib_logger.warning(f"Failed to parse image URL: {e}") return None - + def _transform_user_message(self, content: Any) -> List[Dict[str, Any]]: """Transform user message content to Gemini parts.""" return self._parse_content_parts(content) - + def _transform_assistant_message( - self, - msg: Dict[str, Any], - model: str, - _tool_id_to_name: Dict[str, str] + self, msg: Dict[str, Any], model: str, _tool_id_to_name: Dict[str, str] ) -> List[Dict[str, Any]]: """Transform assistant message including tool calls and thinking injection.""" parts = [] content = msg.get("content") tool_calls = msg.get("tool_calls", []) reasoning_content = msg.get("reasoning_content") - + # Handle reasoning_content if present (from original Claude response with thinking) if reasoning_content and self._is_claude(model): # Add thinking part with cached signature @@ -1456,8 +1920,7 @@ def _transform_assistant_message( } # Try to get signature from cache cache_key = self._generate_thinking_cache_key( - content if isinstance(content, str) else "", - tool_calls + content if isinstance(content, str) else "", tool_calls ) cached_sig = None if cache_key: @@ -1468,11 +1931,13 @@ def _transform_assistant_message( cached_sig = cached_data.get("thought_signature", "") except json.JSONDecodeError: pass - + if cached_sig: thinking_part["thoughtSignature"] = cached_sig parts.append(thinking_part) - lib_logger.debug(f"Added reasoning_content with cached signature ({len(reasoning_content)} chars)") + lib_logger.debug( + f"Added reasoning_content with cached signature ({len(reasoning_content)} chars)" + ) else: # No cached signature - skip the thinking block # This can happen if context was compressed and signature was lost @@ -1480,15 +1945,19 @@ def _transform_assistant_message( f"Skipping reasoning_content - no valid signature found. " f"This may cause issues if thinking is enabled." ) - elif self._is_claude(model) and self._enable_signature_cache and not reasoning_content: + elif ( + self._is_claude(model) + and self._enable_signature_cache + and not reasoning_content + ): # Fallback: Try to inject cached thinking for Claude (original behavior) thinking_parts = self._get_cached_thinking(content, tool_calls) parts.extend(thinking_parts) - + # Add regular content if isinstance(content, str) and content: parts.append({"text": content}) - + # Add tool calls # Track if we've seen the first function call in this message # Per Gemini docs: Only the FIRST parallel function call gets a signature @@ -1496,32 +1965,28 @@ def _transform_assistant_message( for tc in tool_calls: if tc.get("type") != "function": continue - + try: args = json.loads(tc["function"]["arguments"]) except (json.JSONDecodeError, TypeError): args = {} - + tool_id = tc.get("id", "") func_name = tc["function"]["name"] - - #lib_logger.debug( + + # lib_logger.debug( # f"[ID Transform] Converting assistant tool_call to functionCall: " # f"id={tool_id}, name={func_name}" - #) + # ) # Add prefix for Gemini 3 if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: func_name = f"{self._gemini3_tool_prefix}{func_name}" - + func_part = { - "functionCall": { - "name": func_name, - "args": args, - "id": tool_id - } + "functionCall": {"name": func_name, "args": args, "id": tool_id} } - + # Add thoughtSignature for Gemini 3 # Per Gemini docs: Only the FIRST parallel function call gets a signature. # Subsequent parallel calls should NOT have a thoughtSignature field. @@ -1529,19 +1994,21 @@ def _transform_assistant_message( sig = tc.get("thought_signature") if not sig and tool_id and self._enable_signature_cache: sig = self._signature_cache.retrieve(tool_id) - + if sig: func_part["thoughtSignature"] = sig elif first_func_in_msg: # Only add bypass to the first function call if no sig available func_part["thoughtSignature"] = "skip_thought_signature_validator" - lib_logger.warning(f"Missing thoughtSignature for first func call {tool_id}, using bypass") + lib_logger.warning( + f"Missing thoughtSignature for first func call {tool_id}, using bypass" + ) # Subsequent parallel calls: no signature field at all - + first_func_in_msg = False - + parts.append(func_part) - + # Safety: ensure we return at least one part to maintain role alternation # This handles edge cases like assistant messages that had only thinking content # which got stripped, leaving the message otherwise empty @@ -1551,107 +2018,103 @@ def _transform_assistant_message( lib_logger.debug( "[Transform] Added empty text part to maintain role alternation" ) - + return parts - + def _get_cached_thinking( - self, - content: Any, - tool_calls: List[Dict] + self, content: Any, tool_calls: List[Dict] ) -> List[Dict[str, Any]]: """Retrieve and format cached thinking content for Claude.""" parts = [] msg_text = content if isinstance(content, str) else "" cache_key = self._generate_thinking_cache_key(msg_text, tool_calls) - + if not cache_key: return parts - + cached_json = self._thinking_cache.retrieve(cache_key) if not cached_json: return parts - + try: thinking_data = json.loads(cached_json) thinking_text = thinking_data.get("thinking_text", "") sig = thinking_data.get("thought_signature", "") - + if thinking_text: thinking_part = { "text": thinking_text, "thought": True, - "thoughtSignature": sig or "skip_thought_signature_validator" + "thoughtSignature": sig or "skip_thought_signature_validator", } parts.append(thinking_part) lib_logger.debug(f"Injected {len(thinking_text)} chars of thinking") except json.JSONDecodeError: lib_logger.warning(f"Failed to parse cached thinking: {cache_key}") - + return parts - + def _transform_tool_message( - self, - msg: Dict[str, Any], - model: str, - tool_id_to_name: Dict[str, str] + self, msg: Dict[str, Any], model: str, tool_id_to_name: Dict[str, str] ) -> List[Dict[str, Any]]: """Transform tool response message.""" tool_id = msg.get("tool_call_id", "") func_name = tool_id_to_name.get(tool_id, "unknown_function") content = msg.get("content", "{}") - + # Log ID lookup if tool_id not in tool_id_to_name: lib_logger.warning( f"[ID Mismatch] Tool response has ID '{tool_id}' which was not found in tool_id_to_name map. " f"Available IDs: {list(tool_id_to_name.keys())}" ) - #else: - #lib_logger.debug(f"[ID Mapping] Tool response matched: id={tool_id}, name={func_name}") - + # else: + # lib_logger.debug(f"[ID Mapping] Tool response matched: id={tool_id}, name={func_name}") + # Add prefix for Gemini 3 if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: func_name = f"{self._gemini3_tool_prefix}{func_name}" - + try: parsed_content = json.loads(content) except (json.JSONDecodeError, TypeError): parsed_content = content - - return [{ - "functionResponse": { - "name": func_name, - "response": {"result": parsed_content}, - "id": tool_id + + return [ + { + "functionResponse": { + "name": func_name, + "response": {"result": parsed_content}, + "id": tool_id, + } } - }] - + ] + # ========================================================================= # TOOL RESPONSE GROUPING # ========================================================================= - + def _fix_tool_response_grouping( - self, - contents: List[Dict[str, Any]] + self, contents: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: """ Group function calls with their responses for Antigravity compatibility. - + Converts linear format (call, response, call, response) to grouped format (model with calls, user with all responses). - + IMPORTANT: Preserves ID-based pairing to prevent mismatches. """ new_contents = [] pending_groups = [] # List of {"ids": [id1, id2, ...], "call_indices": [...]} collected_responses = {} # Dict mapping ID -> response_part - + for content in contents: role = content.get("role") parts = content.get("parts", []) - + response_parts = [p for p in parts if "functionResponse" in p] - + if response_parts: # Collect responses by ID (ignore duplicates - keep first occurrence) for resp in response_parts: @@ -1663,45 +2126,56 @@ def _fix_tool_response_grouping( f"Ignoring duplicate - this may indicate malformed conversation history." ) continue - #lib_logger.debug(f"[Grouping] Collected response for ID: {resp_id}") + # lib_logger.debug(f"[Grouping] Collected response for ID: {resp_id}") collected_responses[resp_id] = resp - + # Try to satisfy pending groups (newest first) for i in range(len(pending_groups) - 1, -1, -1): group = pending_groups[i] group_ids = group["ids"] - + # Check if we have ALL responses for this group if all(gid in collected_responses for gid in group_ids): # Extract responses in the same order as the function calls - group_responses = [collected_responses.pop(gid) for gid in group_ids] + group_responses = [ + collected_responses.pop(gid) for gid in group_ids + ] new_contents.append({"parts": group_responses, "role": "user"}) - #lib_logger.debug( + # lib_logger.debug( # f"[Grouping] Satisfied group with {len(group_responses)} responses: " # f"ids={group_ids}" - #) + # ) pending_groups.pop(i) break continue - + if role == "model": func_calls = [p for p in parts if "functionCall" in p] new_contents.append(content) if func_calls: - call_ids = [fc.get("functionCall", {}).get("id", "") for fc in func_calls] + call_ids = [ + fc.get("functionCall", {}).get("id", "") for fc in func_calls + ] call_ids = [cid for cid in call_ids if cid] # Filter empty IDs if call_ids: - lib_logger.debug(f"[Grouping] Created pending group expecting {len(call_ids)} responses: ids={call_ids}") - pending_groups.append({"ids": call_ids, "call_indices": list(range(len(func_calls)))}) + lib_logger.debug( + f"[Grouping] Created pending group expecting {len(call_ids)} responses: ids={call_ids}" + ) + pending_groups.append( + { + "ids": call_ids, + "call_indices": list(range(len(func_calls))), + } + ) else: new_contents.append(content) - + # Handle remaining groups (shouldn't happen in well-formed conversations) # Attempt recovery by matching orphans to unsatisfied calls for group in pending_groups: group_ids = group["ids"] group_responses = [] - + for expected_id in group_ids: if expected_id in collected_responses: group_responses.append(collected_responses.pop(expected_id)) @@ -1711,151 +2185,155 @@ def _fix_tool_response_grouping( # Get the first available orphan ID to maintain order orphan_id = next(iter(collected_responses)) orphan_resp = collected_responses.pop(orphan_id) - + # Fix the ID in the response to match the call orphan_resp["functionResponse"]["id"] = expected_id - + lib_logger.warning( f"[Grouping] Auto-repaired ID mismatch: mapped response '{orphan_id}' " f"to call '{expected_id}'" ) group_responses.append(orphan_resp) - + if group_responses: new_contents.append({"parts": group_responses, "role": "user"}) - + if len(group_responses) != len(group_ids): lib_logger.warning( f"[Grouping] Partial group satisfaction after repair: " f"expected {len(group_ids)}, got {len(group_responses)} responses" ) - + # Warn about unmatched responses if collected_responses: lib_logger.warning( f"[Grouping] {len(collected_responses)} unmatched responses remaining: " f"ids={list(collected_responses.keys())}" ) - + return new_contents - + # ========================================================================= # GEMINI 3 TOOL TRANSFORMATIONS # ========================================================================= - + def _apply_gemini3_namespace( - self, - tools: List[Dict[str, Any]] + self, tools: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: """Add namespace prefix to tool names for Gemini 3.""" if not tools: return tools - + modified = copy.deepcopy(tools) for tool in modified: for func_decl in tool.get("functionDeclarations", []): name = func_decl.get("name", "") if name: func_decl["name"] = f"{self._gemini3_tool_prefix}{name}" - + return modified - - def _enforce_strict_schema(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + + def _enforce_strict_schema( + self, tools: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: """ Enforce strict JSON schema for Gemini 3 to prevent hallucinated parameters. - + Adds 'additionalProperties: false' recursively to all object schemas, which tells the model it CANNOT add properties not in the schema. """ if not tools: return tools - + def enforce_strict(schema: Any) -> Any: if not isinstance(schema, dict): return schema - + result = {} for key, value in schema.items(): if isinstance(value, dict): result[key] = enforce_strict(value) elif isinstance(value, list): - result[key] = [enforce_strict(item) if isinstance(item, dict) else item for item in value] + result[key] = [ + enforce_strict(item) if isinstance(item, dict) else item + for item in value + ] else: result[key] = value - + # Add additionalProperties: false to object schemas if result.get("type") == "object" and "properties" in result: result["additionalProperties"] = False - + return result - + modified = copy.deepcopy(tools) for tool in modified: for func_decl in tool.get("functionDeclarations", []): if "parametersJsonSchema" in func_decl: - func_decl["parametersJsonSchema"] = enforce_strict(func_decl["parametersJsonSchema"]) - + func_decl["parametersJsonSchema"] = enforce_strict( + func_decl["parametersJsonSchema"] + ) + return modified - + def _inject_signature_into_descriptions( - self, - tools: List[Dict[str, Any]], - description_prompt: Optional[str] = None + self, tools: List[Dict[str, Any]], description_prompt: Optional[str] = None ) -> List[Dict[str, Any]]: """Inject parameter signatures into tool descriptions for Gemini 3 & Claude.""" if not tools: return tools - + # Use provided prompt or default to Gemini 3 prompt prompt_template = description_prompt or self._gemini3_description_prompt - + modified = copy.deepcopy(tools) for tool in modified: for func_decl in tool.get("functionDeclarations", []): schema = func_decl.get("parametersJsonSchema", {}) if not schema: continue - + required = schema.get("required", []) properties = schema.get("properties", {}) - + if not properties: continue - + param_list = [] for prop_name, prop_data in properties.items(): if not isinstance(prop_data, dict): continue - + type_hint = self._format_type_hint(prop_data) is_required = prop_name in required param_list.append( f"{prop_name} ({type_hint}{', REQUIRED' if is_required else ''})" ) - + if param_list: - sig_str = prompt_template.replace( - "{params}", ", ".join(param_list) + sig_str = prompt_template.replace("{params}", ", ".join(param_list)) + func_decl["description"] = ( + func_decl.get("description", "") + sig_str ) - func_decl["description"] = func_decl.get("description", "") + sig_str - + return modified - + def _format_type_hint(self, prop_data: Dict[str, Any], depth: int = 0) -> str: """Format a detailed type hint for a property schema.""" type_hint = prop_data.get("type", "unknown") - + # Handle enum values - show allowed options if "enum" in prop_data: enum_vals = prop_data["enum"] if len(enum_vals) <= 5: return f"string ENUM[{', '.join(repr(v) for v in enum_vals)}]" return f"string ENUM[{len(enum_vals)} options]" - + # Handle const values if "const" in prop_data: return f"string CONST={repr(prop_data['const'])}" - + if type_hint == "array": items = prop_data.get("items", {}) if isinstance(items, dict): @@ -1878,7 +2356,7 @@ def _format_type_hint(self, prop_data: Dict[str, Any], depth: int = 0) -> str: return "ARRAY_OF_OBJECTS" return f"ARRAY_OF_{item_type.upper()}" return "ARRAY" - + if type_hint == "object": nested_props = prop_data.get("properties", {}) nested_req = prop_data.get("required", []) @@ -1890,16 +2368,18 @@ def _format_type_hint(self, prop_data: Dict[str, Any], depth: int = 0) -> str: req = " REQUIRED" if n in nested_req else "" nested_list.append(f"{n}: {t}{req}") return f"object{{{', '.join(nested_list)}}}" - + return type_hint - + def _strip_gemini3_prefix(self, name: str) -> str: """Strip the Gemini 3 namespace prefix from a tool name.""" if name and name.startswith(self._gemini3_tool_prefix): - return name[len(self._gemini3_tool_prefix):] + return name[len(self._gemini3_tool_prefix) :] return name - - def _translate_tool_choice(self, tool_choice: Union[str, Dict[str, Any]], model: str = "") -> Optional[Dict[str, Any]]: + + def _translate_tool_choice( + self, tool_choice: Union[str, Dict[str, Any]], model: str = "" + ) -> Optional[Dict[str, Any]]: """ Translates OpenAI's `tool_choice` to Gemini's `toolConfig`. Handles Gemini 3 namespace prefixes for specific tool selection. @@ -1924,43 +2404,41 @@ def _translate_tool_choice(self, tool_choice: Union[str, Dict[str, Any]], model: # Add Gemini 3 prefix if needed if is_gemini_3 and self._enable_gemini3_tool_fix: function_name = f"{self._gemini3_tool_prefix}{function_name}" - + mode = "ANY" # Force a call, but only to this function config["functionCallingConfig"] = { "mode": mode, - "allowedFunctionNames": [function_name] + "allowedFunctionNames": [function_name], } return config config["functionCallingConfig"] = {"mode": mode} return config - + # ========================================================================= # REQUEST TRANSFORMATION # ========================================================================= - + def _build_tools_payload( - self, - tools: Optional[List[Dict[str, Any]]], - _model: str + self, tools: Optional[List[Dict[str, Any]]], _model: str ) -> Optional[List[Dict[str, Any]]]: """Build Gemini-format tools from OpenAI tools.""" if not tools: return None - + gemini_tools = [] for tool in tools: if tool.get("type") != "function": continue - + func = tool.get("function", {}) params = func.get("parameters") - + func_decl = { "name": func.get("name", ""), - "description": func.get("description", "") + "description": func.get("description", ""), } - + if params and isinstance(params, dict): schema = dict(params) schema.pop("$schema", None) @@ -1969,11 +2447,11 @@ def _build_tools_payload( func_decl["parametersJsonSchema"] = schema else: func_decl["parametersJsonSchema"] = {"type": "object", "properties": {}} - + gemini_tools.append({"functionDeclarations": [func_decl]}) - + return gemini_tools or None - + def _transform_to_antigravity_format( self, gemini_payload: Dict[str, Any], @@ -1981,11 +2459,11 @@ def _transform_to_antigravity_format( project_id: str, max_tokens: Optional[int] = None, reasoning_effort: Optional[str] = None, - tool_choice: Optional[Union[str, Dict[str, Any]]] = None + tool_choice: Optional[Union[str, Dict[str, Any]]] = None, ) -> Dict[str, Any]: """ Transform Gemini CLI payload to complete Antigravity format. - + Args: gemini_payload: Request in Gemini CLI format model: Model name (public alias) @@ -1993,7 +2471,7 @@ def _transform_to_antigravity_format( reasoning_effort: Reasoning effort level (determines -thinking variant for Claude) """ internal_model = self._alias_to_internal(model) - + # Map Claude models to their -thinking variant # claude-opus-4-5: ALWAYS use -thinking (non-thinking variant doesn't exist) # claude-sonnet-4-5: only use -thinking when reasoning_effort is provided @@ -2004,38 +2482,42 @@ def _transform_to_antigravity_format( elif internal_model == "claude-sonnet-4-5" and reasoning_effort: # Sonnet 4.5 uses -thinking only when reasoning_effort is provided internal_model = "claude-sonnet-4-5-thinking" - + # Map gemini-3-pro-preview to -low/-high variant based on thinking config if model == "gemini-3-pro-preview" or internal_model == "gemini-3-pro-preview": # Check thinking config to determine variant - thinking_config = gemini_payload.get("generationConfig", {}).get("thinkingConfig", {}) + thinking_config = gemini_payload.get("generationConfig", {}).get( + "thinkingConfig", {} + ) thinking_level = thinking_config.get("thinkingLevel", "high") if thinking_level == "low": internal_model = "gemini-3-pro-low" else: internal_model = "gemini-3-pro-high" - + # Wrap in Antigravity envelope antigravity_payload = { "project": project_id, # Will be passed as parameter "userAgent": "antigravity", "requestId": _generate_request_id(), "model": internal_model, - "request": copy.deepcopy(gemini_payload) + "request": copy.deepcopy(gemini_payload), } - + # Add session ID antigravity_payload["request"]["sessionId"] = _generate_session_id() - + # Add default safety settings to prevent content filtering # Only add if not already present in the payload if "safetySettings" not in antigravity_payload["request"]: - antigravity_payload["request"]["safetySettings"] = copy.deepcopy(DEFAULT_SAFETY_SETTINGS) - + antigravity_payload["request"]["safetySettings"] = copy.deepcopy( + DEFAULT_SAFETY_SETTINGS + ) + # Handle max_tokens - only apply to Claude, or if explicitly set for others gen_config = antigravity_payload["request"].get("generationConfig", {}) is_claude = self._is_claude(model) - + if max_tokens is not None: # Explicitly set in request - apply to all models gen_config["maxOutputTokens"] = max_tokens @@ -2043,9 +2525,9 @@ def _transform_to_antigravity_format( # Claude model without explicit max_tokens - use default gen_config["maxOutputTokens"] = DEFAULT_MAX_OUTPUT_TOKENS # For non-Claude models without explicit max_tokens, don't set it - + antigravity_payload["request"]["generationConfig"] = gen_config - + # Set toolConfig based on tool_choice parameter tool_config_result = self._translate_tool_choice(tool_choice, model) if tool_config_result: @@ -2055,14 +2537,14 @@ def _transform_to_antigravity_format( tool_config = antigravity_payload["request"].setdefault("toolConfig", {}) func_config = tool_config.setdefault("functionCallingConfig", {}) func_config["mode"] = "AUTO" - + # Handle Gemini 3 thinking logic if not internal_model.startswith("gemini-3-"): thinking_config = gen_config.get("thinkingConfig", {}) if "thinkingLevel" in thinking_config: del thinking_config["thinkingLevel"] thinking_config["thinkingBudget"] = -1 - + # Ensure first function call in each model message has a thoughtSignature for Gemini 3 # Per Gemini docs: Only the FIRST parallel function call gets a signature if internal_model.startswith("gemini-3-"): @@ -2074,16 +2556,20 @@ def _transform_to_antigravity_format( if not first_func_seen: # First function call in this message - needs a signature if "thoughtSignature" not in part: - part["thoughtSignature"] = "skip_thought_signature_validator" + part["thoughtSignature"] = ( + "skip_thought_signature_validator" + ) first_func_seen = True # Subsequent parallel calls: leave as-is (no signature) - + # Claude-specific tool schema transformation - if internal_model.startswith("claude-sonnet-") or internal_model.startswith("claude-opus-"): + if internal_model.startswith("claude-sonnet-") or internal_model.startswith( + "claude-opus-" + ): self._apply_claude_tool_transform(antigravity_payload) - + return antigravity_payload - + def _apply_claude_tool_transform(self, payload: Dict[str, Any]) -> None: """Apply Claude-specific tool schema transformations.""" tools = payload["request"].get("tools", []) @@ -2091,27 +2577,31 @@ def _apply_claude_tool_transform(self, payload: Dict[str, Any]) -> None: for func_decl in tool.get("functionDeclarations", []): if "parametersJsonSchema" in func_decl: params = func_decl["parametersJsonSchema"] - params = _clean_claude_schema(params) if isinstance(params, dict) else params + params = ( + _clean_claude_schema(params) + if isinstance(params, dict) + else params + ) func_decl["parameters"] = params del func_decl["parametersJsonSchema"] - + # ========================================================================= # RESPONSE TRANSFORMATION # ========================================================================= - + def _unwrap_response(self, response: Dict[str, Any]) -> Dict[str, Any]: """Extract Gemini response from Antigravity envelope.""" return response.get("response", response) - + def _gemini_to_openai_chunk( self, chunk: Dict[str, Any], model: str, - accumulator: Optional[Dict[str, Any]] = None + accumulator: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """ Convert Gemini response chunk to OpenAI streaming format. - + Args: chunk: Gemini API response chunk model: Model name @@ -2120,30 +2610,33 @@ def _gemini_to_openai_chunk( candidates = chunk.get("candidates", []) if not candidates: return {} - + candidate = candidates[0] content_parts = candidate.get("content", {}).get("parts", []) - + text_content = "" reasoning_content = "" tool_calls = [] # Use accumulator's tool_idx if available, otherwise use local counter tool_idx = accumulator.get("tool_idx", 0) if accumulator else 0 - + for part in content_parts: has_func = "functionCall" in part has_text = "text" in part has_sig = bool(part.get("thoughtSignature")) - is_thought = part.get("thought") is True or str(part.get("thought")).lower() == 'true' - + is_thought = ( + part.get("thought") is True + or str(part.get("thought")).lower() == "true" + ) + # Accumulate signature for Claude caching if has_sig and is_thought and accumulator is not None: accumulator["thought_signature"] = part["thoughtSignature"] - + # Skip standalone signature parts if has_sig and not has_func and (not has_text or not part.get("text")): continue - + if has_text: text = part["text"] if is_thought: @@ -2154,17 +2647,17 @@ def _gemini_to_openai_chunk( text_content += text if accumulator is not None: accumulator["text_content"] += text - + if has_func: tool_call = self._extract_tool_call(part, model, tool_idx, accumulator) - + # Store signature for each tool call (needed for parallel tool calls) if has_sig: self._handle_tool_signature(tool_call, part["thoughtSignature"]) - + tool_calls.append(tool_call) tool_idx += 1 - + # Build delta delta = {} if text_content: @@ -2179,80 +2672,87 @@ def _gemini_to_openai_chunk( accumulator["tool_idx"] = tool_idx elif text_content or reasoning_content: delta["role"] = "assistant" - + # Build usage if present usage = self._build_usage(chunk.get("usageMetadata", {})) - + # Mark completion when we see usageMetadata if chunk.get("usageMetadata") and accumulator is not None: accumulator["is_complete"] = True - + # Build choice - just translate, don't include finish_reason # Client will handle finish_reason logic choice = {"index": 0, "delta": delta} - + response = { "id": chunk.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), "object": "chat.completion.chunk", "created": int(time.time()), "model": model, - "choices": [choice] + "choices": [choice], } - + if usage: response["usage"] = usage - + return response - + def _gemini_to_openai_non_streaming( - self, - response: Dict[str, Any], - model: str + self, response: Dict[str, Any], model: str ) -> Dict[str, Any]: """Convert Gemini response to OpenAI non-streaming format.""" candidates = response.get("candidates", []) if not candidates: return {} - + candidate = candidates[0] content_parts = candidate.get("content", {}).get("parts", []) - + text_content = "" reasoning_content = "" tool_calls = [] thought_sig = "" - + for part in content_parts: has_func = "functionCall" in part has_text = "text" in part has_sig = bool(part.get("thoughtSignature")) - is_thought = part.get("thought") is True or str(part.get("thought")).lower() == 'true' - + is_thought = ( + part.get("thought") is True + or str(part.get("thought")).lower() == "true" + ) + if has_sig and is_thought: thought_sig = part["thoughtSignature"] - + if has_sig and not has_func and (not has_text or not part.get("text")): continue - + if has_text: if is_thought: reasoning_content += part["text"] else: text_content += part["text"] - + if has_func: tool_call = self._extract_tool_call(part, model, len(tool_calls)) - + # Store signature for each tool call (needed for parallel tool calls) if has_sig: self._handle_tool_signature(tool_call, part["thoughtSignature"]) - + tool_calls.append(tool_call) - + # Cache Claude thinking - if reasoning_content and self._is_claude(model) and self._enable_signature_cache: - self._cache_thinking(reasoning_content, thought_sig, text_content, tool_calls) - + if ( + reasoning_content + and self._is_claude(model) + and self._enable_signature_cache + ): + self._cache_thinking( + reasoning_content, thought_sig, text_content, tool_calls + ) + # Build message message = {"role": "assistant"} if text_content: @@ -2264,172 +2764,169 @@ def _gemini_to_openai_non_streaming( if tool_calls: message["tool_calls"] = tool_calls message.pop("content", None) - - finish_reason = self._map_finish_reason(candidate.get("finishReason"), bool(tool_calls)) + + finish_reason = self._map_finish_reason( + candidate.get("finishReason"), bool(tool_calls) + ) usage = self._build_usage(response.get("usageMetadata", {})) - + # For non-streaming, always include finish_reason (should always be present) result = { "id": response.get("responseId", f"chatcmpl-{uuid.uuid4().hex[:24]}"), "object": "chat.completion", "created": int(time.time()), "model": model, - "choices": [{"index": 0, "message": message, "finish_reason": finish_reason or "stop"}] + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": finish_reason or "stop", + } + ], } - + if usage: result["usage"] = usage - + return result - + def _extract_tool_call( self, part: Dict[str, Any], model: str, index: int, - accumulator: Optional[Dict[str, Any]] = None + accumulator: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Extract and format a tool call from a response part.""" func_call = part["functionCall"] tool_id = func_call.get("id") or f"call_{uuid.uuid4().hex[:24]}" - - #lib_logger.debug(f"[ID Extraction] Extracting tool call: id={tool_id}, raw_id={func_call.get('id')}") - + + # lib_logger.debug(f"[ID Extraction] Extracting tool call: id={tool_id}, raw_id={func_call.get('id')}") + tool_name = func_call.get("name", "") if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: tool_name = self._strip_gemini3_prefix(tool_name) - + raw_args = func_call.get("args", {}) parsed_args = _recursively_parse_json_strings(raw_args) - + tool_call = { "id": tool_id, "type": "function", "index": index, - "function": { - "name": tool_name, - "arguments": json.dumps(parsed_args) - } + "function": {"name": tool_name, "arguments": json.dumps(parsed_args)}, } - + if accumulator is not None: accumulator["tool_calls"].append(tool_call) - + return tool_call - + def _handle_tool_signature(self, tool_call: Dict, signature: str) -> None: """Handle thoughtSignature for a tool call.""" tool_id = tool_call["id"] - + if self._enable_signature_cache: self._signature_cache.store(tool_id, signature) lib_logger.debug(f"Stored signature for {tool_id}") - + if self._preserve_signatures_in_client: tool_call["thought_signature"] = signature - + def _map_finish_reason( - self, - gemini_reason: Optional[str], - has_tool_calls: bool + self, gemini_reason: Optional[str], has_tool_calls: bool ) -> Optional[str]: """Map Gemini finish reason to OpenAI format.""" if not gemini_reason: return None reason = FINISH_REASON_MAP.get(gemini_reason, "stop") return "tool_calls" if has_tool_calls else reason - + def _build_usage(self, metadata: Dict[str, Any]) -> Optional[Dict[str, Any]]: """Build usage dict from Gemini usage metadata.""" if not metadata: return None - + prompt = metadata.get("promptTokenCount", 0) thoughts = metadata.get("thoughtsTokenCount", 0) completion = metadata.get("candidatesTokenCount", 0) - + usage = { "prompt_tokens": prompt + thoughts, "completion_tokens": completion, - "total_tokens": metadata.get("totalTokenCount", 0) + "total_tokens": metadata.get("totalTokenCount", 0), } - + if thoughts > 0: usage["completion_tokens_details"] = {"reasoning_tokens": thoughts} - + return usage - + def _cache_thinking( - self, - reasoning: str, - signature: str, - text: str, - tool_calls: List[Dict] + self, reasoning: str, signature: str, text: str, tool_calls: List[Dict] ) -> None: """Cache Claude thinking content.""" cache_key = self._generate_thinking_cache_key(text, tool_calls) if not cache_key: return - + data = { "thinking_text": reasoning, "thought_signature": signature, "text_preview": text[:100] if text else "", "tool_ids": [tc.get("id", "") for tc in tool_calls], - "timestamp": time.time() + "timestamp": time.time(), } - + self._thinking_cache.store(cache_key, json.dumps(data)) lib_logger.info(f"Cached thinking: {cache_key[:50]}...") - + # ========================================================================= # PROVIDER INTERFACE IMPLEMENTATION # ========================================================================= - + async def get_valid_token(self, credential_identifier: str) -> str: """Get a valid access token for the credential.""" creds = await self._load_credentials(credential_identifier) if self._is_token_expired(creds): creds = await self._refresh_token(credential_identifier, creds) - return creds['access_token'] - + return creds["access_token"] + def has_custom_logic(self) -> bool: """Antigravity uses custom translation logic.""" return True - + async def get_auth_header(self, credential_identifier: str) -> Dict[str, str]: """Get OAuth authorization header.""" token = await self.get_valid_token(credential_identifier) return {"Authorization": f"Bearer {token}"} - - async def get_models( - self, - api_key: str, - client: httpx.AsyncClient - ) -> List[str]: + + async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: """Fetch available models from Antigravity.""" if not self._enable_dynamic_models: lib_logger.debug("Using hardcoded model list") return [f"antigravity/{m}" for m in AVAILABLE_MODELS] - + try: token = await self.get_valid_token(api_key) url = f"{self._get_base_url()}/fetchAvailableModels" - + headers = { "Authorization": f"Bearer {token}", - "Content-Type": "application/json" + "Content-Type": "application/json", } payload = { "project": _generate_project_id(), "requestId": _generate_request_id(), - "userAgent": "antigravity" + "userAgent": "antigravity", } - - response = await client.post(url, json=payload, headers=headers, timeout=30.0) + + response = await client.post( + url, json=payload, headers=headers, timeout=30.0 + ) response.raise_for_status() data = response.json() - + models = [] for model_info in data.get("models", []): internal = model_info.get("name", "").replace("models/", "") @@ -2437,23 +2934,21 @@ async def get_models( public = self._internal_to_alias(internal) if public: models.append(f"antigravity/{public}") - + if models: lib_logger.info(f"Discovered {len(models)} models") return models except Exception as e: lib_logger.warning(f"Dynamic model discovery failed: {e}") - + return [f"antigravity/{m}" for m in AVAILABLE_MODELS] - + async def acompletion( - self, - client: httpx.AsyncClient, - **kwargs + self, client: httpx.AsyncClient, **kwargs ) -> Union[litellm.ModelResponse, AsyncGenerator[litellm.ModelResponse, None]]: """ Handle completion requests for Antigravity. - + Main entry point that: 1. Extracts parameters and transforms messages 2. Builds Antigravity request payload @@ -2473,140 +2968,168 @@ async def acompletion( max_tokens = kwargs.get("max_tokens") custom_budget = kwargs.get("custom_reasoning_budget", False) enable_logging = kwargs.pop("enable_request_logging", False) - + # Create logger file_logger = AntigravityFileLogger(model, enable_logging) - + # Determine if thinking is enabled for this request # Thinking is enabled if reasoning_effort is set (and not "disable") for Claude thinking_enabled = False if self._is_claude(model): # For Claude, thinking is enabled when reasoning_effort is provided and not "disable" - thinking_enabled = reasoning_effort is not None and reasoning_effort != "disable" - + thinking_enabled = ( + reasoning_effort is not None and reasoning_effort != "disable" + ) + # Sanitize thinking blocks for Claude to prevent 400 errors # This handles: context compression, model switching, mid-turn thinking toggle # Returns (sanitized_messages, force_disable_thinking) force_disable_thinking = False if self._is_claude(model) and self._enable_thinking_sanitization: - messages, force_disable_thinking = self._sanitize_thinking_for_claude(messages, thinking_enabled) - + messages, force_disable_thinking = self._sanitize_thinking_for_claude( + messages, thinking_enabled + ) + # If we're in a mid-turn thinking toggle situation, we MUST disable thinking # for this request. Thinking will naturally resume on the next turn. if force_disable_thinking: thinking_enabled = False reasoning_effort = "disable" # Force disable for this request - + # Transform messages system_instruction, gemini_contents = self._transform_messages(messages, model) gemini_contents = self._fix_tool_response_grouping(gemini_contents) - + # Build payload gemini_payload = {"contents": gemini_contents} - + if system_instruction: gemini_payload["system_instruction"] = system_instruction - + # Inject tool usage hardening system instructions if tools: if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: - self._inject_tool_hardening_instruction(gemini_payload, self._gemini3_system_instruction) + self._inject_tool_hardening_instruction( + gemini_payload, self._gemini3_system_instruction + ) elif self._is_claude(model) and self._enable_claude_tool_fix: - self._inject_tool_hardening_instruction(gemini_payload, self._claude_system_instruction) - + self._inject_tool_hardening_instruction( + gemini_payload, self._claude_system_instruction + ) + # Add generation config gen_config = {} if top_p is not None: gen_config["topP"] = top_p - + # Handle temperature - Gemini 3 defaults to 1 if not explicitly set if temperature is not None: gen_config["temperature"] = temperature elif self._is_gemini_3(model): # Gemini 3 performs better with temperature=1 for tool use gen_config["temperature"] = 1.0 - - thinking_config = self._get_thinking_config(reasoning_effort, model, custom_budget) + + thinking_config = self._get_thinking_config( + reasoning_effort, model, custom_budget + ) if thinking_config: gen_config.setdefault("thinkingConfig", {}).update(thinking_config) - + if gen_config: gemini_payload["generationConfig"] = gen_config - + # Add tools gemini_tools = self._build_tools_payload(tools, model) if gemini_tools: gemini_payload["tools"] = gemini_tools - + # Apply tool transformations if self._is_gemini_3(model) and self._enable_gemini3_tool_fix: # Gemini 3: namespace prefix + strict schema + parameter signatures - gemini_payload["tools"] = self._apply_gemini3_namespace(gemini_payload["tools"]) + gemini_payload["tools"] = self._apply_gemini3_namespace( + gemini_payload["tools"] + ) if self._gemini3_enforce_strict_schema: - gemini_payload["tools"] = self._enforce_strict_schema(gemini_payload["tools"]) + gemini_payload["tools"] = self._enforce_strict_schema( + gemini_payload["tools"] + ) gemini_payload["tools"] = self._inject_signature_into_descriptions( - gemini_payload["tools"], - self._gemini3_description_prompt + gemini_payload["tools"], self._gemini3_description_prompt ) elif self._is_claude(model) and self._enable_claude_tool_fix: # Claude: parameter signatures only (no namespace prefix) gemini_payload["tools"] = self._inject_signature_into_descriptions( - gemini_payload["tools"], - self._claude_description_prompt + gemini_payload["tools"], self._claude_description_prompt ) - + # Get access token first (needed for project discovery) token = await self.get_valid_token(credential_path) - + # Discover real project ID litellm_params = kwargs.get("litellm_params", {}) or {} - project_id = await self._discover_project_id(credential_path, token, litellm_params) + project_id = await self._discover_project_id( + credential_path, token, litellm_params + ) # Transform to Antigravity format with real project ID - payload = self._transform_to_antigravity_format(gemini_payload, model, project_id, max_tokens, reasoning_effort, tool_choice) + payload = self._transform_to_antigravity_format( + gemini_payload, model, project_id, max_tokens, reasoning_effort, tool_choice + ) file_logger.log_request(payload) - + # Make API call base_url = self._get_base_url() endpoint = ":streamGenerateContent" if stream else ":generateContent" url = f"{base_url}{endpoint}" - + if stream: url = f"{url}?alt=sse" - + parsed = urlparse(base_url) - host = parsed.netloc or base_url.replace("https://", "").replace("http://", "").rstrip("/") - + host = parsed.netloc or base_url.replace("https://", "").replace( + "http://", "" + ).rstrip("/") + headers = { "Authorization": f"Bearer {token}", "Content-Type": "application/json", "Host": host, "User-Agent": "antigravity/1.11.9", - "Accept": "text/event-stream" if stream else "application/json" + "Accept": "text/event-stream" if stream else "application/json", } - + try: if stream: - return self._handle_streaming(client, url, headers, payload, model, file_logger) + return self._handle_streaming( + client, url, headers, payload, model, file_logger + ) else: - return await self._handle_non_streaming(client, url, headers, payload, model, file_logger) + return await self._handle_non_streaming( + client, url, headers, payload, model, file_logger + ) except Exception as e: if self._try_next_base_url(): lib_logger.warning(f"Retrying with fallback URL: {e}") url = f"{self._get_base_url()}{endpoint}" if stream: - return self._handle_streaming(client, url, headers, payload, model, file_logger) + return self._handle_streaming( + client, url, headers, payload, model, file_logger + ) else: - return await self._handle_non_streaming(client, url, headers, payload, model, file_logger) + return await self._handle_non_streaming( + client, url, headers, payload, model, file_logger + ) raise - - def _inject_tool_hardening_instruction(self, payload: Dict[str, Any], instruction_text: str) -> None: + + def _inject_tool_hardening_instruction( + self, payload: Dict[str, Any], instruction_text: str + ) -> None: """Inject tool usage hardening system instruction for Gemini 3 & Claude.""" if not instruction_text: return - + instruction_part = {"text": instruction_text} - + if "system_instruction" in payload: existing = payload["system_instruction"] if isinstance(existing, dict) and "parts" in existing: @@ -2614,11 +3137,14 @@ def _inject_tool_hardening_instruction(self, payload: Dict[str, Any], instructio else: payload["system_instruction"] = { "role": "user", - "parts": [instruction_part, {"text": str(existing)}] + "parts": [instruction_part, {"text": str(existing)}], } else: - payload["system_instruction"] = {"role": "user", "parts": [instruction_part]} - + payload["system_instruction"] = { + "role": "user", + "parts": [instruction_part], + } + async def _handle_non_streaming( self, client: httpx.AsyncClient, @@ -2626,21 +3152,21 @@ async def _handle_non_streaming( headers: Dict[str, str], payload: Dict[str, Any], model: str, - file_logger: Optional[AntigravityFileLogger] = None + file_logger: Optional[AntigravityFileLogger] = None, ) -> litellm.ModelResponse: """Handle non-streaming completion.""" response = await client.post(url, headers=headers, json=payload, timeout=600.0) response.raise_for_status() - + data = response.json() if file_logger: file_logger.log_final_response(data) - + gemini_response = self._unwrap_response(data) openai_response = self._gemini_to_openai_non_streaming(gemini_response, model) - + return litellm.ModelResponse(**openai_response) - + async def _handle_streaming( self, client: httpx.AsyncClient, @@ -2648,7 +3174,7 @@ async def _handle_streaming( headers: Dict[str, str], payload: Dict[str, Any], model: str, - file_logger: Optional[AntigravityFileLogger] = None + file_logger: Optional[AntigravityFileLogger] = None, ) -> AsyncGenerator[litellm.ModelResponse, None]: """Handle streaming completion.""" # Accumulator tracks state across chunks for caching and tool indexing @@ -2658,39 +3184,45 @@ async def _handle_streaming( "text_content": "", "tool_calls": [], "tool_idx": 0, # Track tool call index across chunks - "is_complete": False # Track if we received usageMetadata + "is_complete": False, # Track if we received usageMetadata } - - async with client.stream("POST", url, headers=headers, json=payload, timeout=600.0) as response: + + async with client.stream( + "POST", url, headers=headers, json=payload, timeout=600.0 + ) as response: if response.status_code >= 400: try: error_body = await response.aread() - lib_logger.error(f"API error {response.status_code}: {error_body.decode()}") + lib_logger.error( + f"API error {response.status_code}: {error_body.decode()}" + ) except Exception: pass - + response.raise_for_status() - + async for line in response.aiter_lines(): if file_logger: file_logger.log_response_chunk(line) - + if line.startswith("data: "): data_str = line[6:] if data_str == "[DONE]": break - + try: chunk = json.loads(data_str) gemini_chunk = self._unwrap_response(chunk) - openai_chunk = self._gemini_to_openai_chunk(gemini_chunk, model, accumulator) - + openai_chunk = self._gemini_to_openai_chunk( + gemini_chunk, model, accumulator + ) + yield litellm.ModelResponse(**openai_chunk) except json.JSONDecodeError: if file_logger: file_logger.log_error(f"Parse error: {data_str[:100]}") continue - + # If stream ended without usageMetadata chunk, emit a final chunk with finish_reason # Emit final chunk if stream ended without usageMetadata # Client will determine the correct finish_reason based on accumulated state @@ -2702,19 +3234,27 @@ async def _handle_streaming( "model": model, "choices": [{"index": 0, "delta": {}, "finish_reason": None}], # Include minimal usage to signal this is the final chunk - "usage": {"prompt_tokens": 0, "completion_tokens": 1, "total_tokens": 1} + "usage": { + "prompt_tokens": 0, + "completion_tokens": 1, + "total_tokens": 1, + }, } yield litellm.ModelResponse(**final_chunk) - + # Cache Claude thinking after stream completes - if self._is_claude(model) and self._enable_signature_cache and accumulator.get("reasoning_content"): + if ( + self._is_claude(model) + and self._enable_signature_cache + and accumulator.get("reasoning_content") + ): self._cache_thinking( accumulator["reasoning_content"], accumulator["thought_signature"], accumulator["text_content"], - accumulator["tool_calls"] + accumulator["tool_calls"], ) - + async def count_tokens( self, client: httpx.AsyncClient, @@ -2722,49 +3262,55 @@ async def count_tokens( model: str, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, - litellm_params: Optional[Dict[str, Any]] = None + litellm_params: Optional[Dict[str, Any]] = None, ) -> Dict[str, int]: """Count tokens for the given prompt using Antigravity :countTokens endpoint.""" try: token = await self.get_valid_token(credential_path) internal_model = self._alias_to_internal(model) - + # Discover project ID - project_id = await self._discover_project_id(credential_path, token, litellm_params or {}) - - system_instruction, contents = self._transform_messages(messages, internal_model) + project_id = await self._discover_project_id( + credential_path, token, litellm_params or {} + ) + + system_instruction, contents = self._transform_messages( + messages, internal_model + ) contents = self._fix_tool_response_grouping(contents) - + gemini_payload = {"contents": contents} if system_instruction: gemini_payload["systemInstruction"] = system_instruction - + gemini_tools = self._build_tools_payload(tools, model) if gemini_tools: gemini_payload["tools"] = gemini_tools - + antigravity_payload = { "project": project_id, "userAgent": "antigravity", "requestId": _generate_request_id(), "model": internal_model, - "request": gemini_payload + "request": gemini_payload, } - + url = f"{self._get_base_url()}:countTokens" headers = { "Authorization": f"Bearer {token}", - "Content-Type": "application/json" + "Content-Type": "application/json", } - - response = await client.post(url, headers=headers, json=antigravity_payload, timeout=30) + + response = await client.post( + url, headers=headers, json=antigravity_payload, timeout=30 + ) response.raise_for_status() - + data = response.json() unwrapped = self._unwrap_response(data) - total = unwrapped.get('totalTokens', 0) - - return {'prompt_tokens': total, 'total_tokens': total} + total = unwrapped.get("totalTokens", 0) + + return {"prompt_tokens": total, "total_tokens": total} except Exception as e: lib_logger.error(f"Token counting failed: {e}") - return {'prompt_tokens': 0, 'total_tokens': 0} \ No newline at end of file + return {"prompt_tokens": 0, "total_tokens": 0} From bccb879ce836e86741523d8e681cd1f2d16df797 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 03:53:38 +0100 Subject: [PATCH 071/221] =?UTF-8?q?refactor(antigravity):=20=F0=9F=94=A8?= =?UTF-8?q?=20migrate=20thinking=20sanitization=20to=20gemini=20message=20?= =?UTF-8?q?format?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit refactors the Claude thinking sanitization logic to operate on Gemini-format messages after transformation, rather than before. This enables the sanitization to work with the full message context including thinking blocks that were restored from cache during the transformation process. Key changes: - Move `_sanitize_thinking_for_claude` call to after `_transform_messages` instead of before - Update all thinking detection and manipulation methods to work with Gemini format (role "model", "parts" array with "thought": true) - Refactor `_analyze_turn_state` to detect tool results as user messages with "functionResponse" parts - Update `_message_has_thinking` to check for "thought": true in parts array - Add new `_message_has_tool_calls` helper for Gemini format detection - Refactor `_strip_all_thinking_blocks` to filter parts with "thought": true - Update `_strip_old_turn_thinking` and `_preserve_turn_start_thinking` for Gemini format - Refactor `_looks_like_compacted_thinking_turn` to detect functionCall parts without thinking - Update `_recover_thinking_from_cache` to inject thinking as Gemini-format part with "thought": true - Refactor `_close_tool_loop_for_thinking` to use Gemini message structure - Update all docstrings and comments to reflect "model" role instead of "assistant" This change fixes issues where context compression or client-side stripping of reasoning_content would prevent proper thinking sanitization, as the sanitization now occurs after the transformation has restored thinking from cache. --- .../providers/antigravity_provider.py | 401 ++++++++++-------- 1 file changed, 229 insertions(+), 172 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index a1c66152..d0d46457 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -1178,20 +1178,27 @@ def _analyze_conversation_state( Returns: { "in_tool_loop": bool - True if we're in an incomplete tool use loop - "turn_start_idx": int - Index of first assistant message in current turn + "turn_start_idx": int - Index of first model message in current turn "turn_has_thinking": bool - Whether the TURN started with thinking - "last_assistant_idx": int - Index of last assistant message - "last_assistant_has_thinking": bool - Whether last assistant msg has thinking - "last_assistant_has_tool_calls": bool - Whether last assistant msg has tool calls - "pending_tool_results": bool - Whether there are tool results after last assistant + "last_model_idx": int - Index of last model message + "last_model_has_thinking": bool - Whether last model msg has thinking + "last_model_has_tool_calls": bool - Whether last model msg has tool calls + "pending_tool_results": bool - Whether there are tool results after last model "thinking_block_indices": List[int] - Indices of messages with thinking/reasoning } + + NOTE: This now operates on Gemini-format messages (after transformation): + - Role "model" instead of "assistant" + - Role "user" for both user messages AND tool results (with functionResponse) + - "parts" array with "thought": true for thinking + - "parts" array with "functionCall" for tool calls + - "parts" array with "functionResponse" for tool results """ state = { "in_tool_loop": False, "turn_start_idx": -1, "turn_has_thinking": False, - "last_assistant_idx": -1, + "last_assistant_idx": -1, # Keep name for compatibility "last_assistant_has_thinking": False, "last_assistant_has_tool_calls": False, "pending_tool_results": False, @@ -1199,25 +1206,16 @@ def _analyze_conversation_state( } # First pass: Find the last "real" user message (not a tool result) - # A real user message is one that doesn't immediately follow an assistant with tool_calls + # In Gemini format, tool results are "user" role with functionResponse parts last_real_user_idx = -1 for i, msg in enumerate(messages): role = msg.get("role") if role == "user": - # Check if this is a real user message or just follows tool results - # Tool messages have role="tool", so if this is role="user" and - # it's not just a tool_result container, it's a real user message. - # However, we need to be careful: the client might format tool results - # as user messages with tool_result content. Check the content. - content = msg.get("content") - - # If content is a list with tool_result items, it's a tool response - is_tool_result_msg = False - if isinstance(content, list): - for item in content: - if isinstance(item, dict) and item.get("type") == "tool_result": - is_tool_result_msg = True - break + # Check if this is a real user message or a tool result container + parts = msg.get("parts", []) + is_tool_result_msg = any( + isinstance(p, dict) and "functionResponse" in p for p in parts + ) if not is_tool_result_msg: last_real_user_idx = i @@ -1226,52 +1224,71 @@ def _analyze_conversation_state( for i, msg in enumerate(messages): role = msg.get("role") - if role == "assistant": - # Check for thinking/reasoning content + if role == "model": + # Check for thinking/reasoning content (Gemini format) has_thinking = self._message_has_thinking(msg) + # Check for tool calls (functionCall in parts) + parts = msg.get("parts", []) + has_tool_calls = any( + isinstance(p, dict) and "functionCall" in p for p in parts + ) + # Track if this is the turn start if i > last_real_user_idx and state["turn_start_idx"] == -1: state["turn_start_idx"] = i state["turn_has_thinking"] = has_thinking state["last_assistant_idx"] = i - state["last_assistant_has_tool_calls"] = bool(msg.get("tool_calls")) + state["last_assistant_has_tool_calls"] = has_tool_calls state["last_assistant_has_thinking"] = has_thinking if has_thinking: state["thinking_block_indices"].append(i) - elif role == "tool": - # Tool result after an assistant message with tool calls = in tool loop - if state["last_assistant_has_tool_calls"]: + elif role == "user": + # Check if this is a tool result (functionResponse in parts) + parts = msg.get("parts", []) + is_tool_result = any( + isinstance(p, dict) and "functionResponse" in p for p in parts + ) + + if is_tool_result and state["last_assistant_has_tool_calls"]: state["pending_tool_results"] = True # We're in a tool loop if: # 1. There are pending tool results - # 2. The conversation ends with tool results (last message is "tool" role) + # 2. The conversation ends with tool results (last message is user with functionResponse) if state["pending_tool_results"] and messages: last_msg = messages[-1] - if last_msg.get("role") == "tool": - state["in_tool_loop"] = True + if last_msg.get("role") == "user": + parts = last_msg.get("parts", []) + ends_with_tool_result = any( + isinstance(p, dict) and "functionResponse" in p for p in parts + ) + if ends_with_tool_result: + state["in_tool_loop"] = True return state def _message_has_thinking(self, msg: Dict[str, Any]) -> bool: - """Check if an assistant message contains thinking/reasoning content.""" - # Check reasoning_content field (OpenAI format) - if msg.get("reasoning_content"): - return True - - # Check for thinking in content array (some formats) - content = msg.get("content") - if isinstance(content, list): - for item in content: - if isinstance(item, dict) and item.get("type") == "thinking": - return True + """ + Check if a message contains thinking/reasoning content. + Handles GEMINI format (after transformation): + - "parts" array with items having "thought": true + """ + parts = msg.get("parts", []) + for part in parts: + if isinstance(part, dict) and part.get("thought") is True: + return True return False + def _message_has_tool_calls(self, msg: Dict[str, Any]) -> bool: + """Check if a message contains tool calls (Gemini format).""" + parts = msg.get("parts", []) + return any(isinstance(p, dict) and "functionCall" in p for p in parts) + def _sanitize_thinking_for_claude( self, messages: List[Dict[str, Any]], thinking_enabled: bool ) -> Tuple[List[Dict[str, Any]], bool]: @@ -1403,7 +1420,7 @@ def _sanitize_thinking_for_claude( state["last_assistant_has_tool_calls"] and not state["turn_has_thinking"] ): - # The turn has tool_calls but no thinking at turn start. + # The turn has functionCall but no thinking at turn start. # This could be: # 1. Compaction removed the thinking block # 2. The original call was made without thinking @@ -1412,7 +1429,7 @@ def _sanitize_thinking_for_claude( # For case 2, we let the model respond naturally. # # We can detect case 1 if there's evidence thinking was expected: - # - The turn_start message has tool_calls (typical thinking-enabled flow) + # - The turn_start message has functionCall (typical thinking-enabled flow) # - The content structure suggests a thinking block was stripped # Check if turn_start has the hallmarks of a compacted thinking response @@ -1436,18 +1453,21 @@ def _sanitize_thinking_for_claude( messages, state["turn_start_idx"] ), False else: - # Can't recover - add synthetic user to start fresh turn + # Can't recover - add synthetic user to start fresh turn (Gemini format) lib_logger.info( "[Thinking Sanitization] Detected compacted turn missing thinking block. " "Adding synthetic user message to start fresh thinking turn." ) # Add synthetic user message to trigger new turn with thinking - synthetic_user = {"role": "user", "content": "[Continue]"} + synthetic_user = { + "role": "user", + "parts": [{"text": "[Continue]"}], + } messages.append(synthetic_user) return self._strip_all_thinking_blocks(messages), False else: lib_logger.debug( - "[Thinking Sanitization] Last assistant has tool_calls but no thinking. " + "[Thinking Sanitization] Last model has functionCall but no thinking. " "This is likely from context compression or non-thinking model. " "New response will include thinking naturally." ) @@ -1460,75 +1480,80 @@ def _sanitize_thinking_for_claude( def _strip_all_thinking_blocks( self, messages: List[Dict[str, Any]] ) -> List[Dict[str, Any]]: - """Remove all thinking/reasoning content from messages.""" - for msg in messages: - if msg.get("role") == "assistant": - # Remove reasoning_content field - msg.pop("reasoning_content", None) + """ + Remove all thinking/reasoning content from messages. - # Remove thinking blocks from content array - content = msg.get("content") - if isinstance(content, list): + Handles GEMINI format (after transformation): + - Role "model" instead of "assistant" + - "parts" array with "thought": true for thinking + """ + for msg in messages: + if msg.get("role") == "model": + parts = msg.get("parts", []) + if parts: + # Filter out thinking parts (those with "thought": true) filtered = [ - item - for item in content - if not ( - isinstance(item, dict) and item.get("type") == "thinking" - ) + p + for p in parts + if not (isinstance(p, dict) and p.get("thought") is True) ] - # If filtering leaves empty list, we need to preserve message structure - # to maintain user/assistant alternation. Use empty string as placeholder - # (will result in empty "text" part which is valid). + + # Check if there are still functionCalls remaining + has_function_calls = any( + isinstance(p, dict) and "functionCall" in p for p in filtered + ) + if not filtered: - # Only if there are no tool_calls either - otherwise message is valid - if not msg.get("tool_calls"): - msg["content"] = "" + # All parts were thinking - need placeholder for valid structure + if not has_function_calls: + msg["parts"] = [{"text": ""}] else: - msg["content"] = ( - None # tool_calls exist, content not needed - ) + msg["parts"] = [] # Will be invalid, but shouldn't happen else: - msg["content"] = filtered + msg["parts"] = filtered return messages def _strip_old_turn_thinking( - self, messages: List[Dict[str, Any]], last_assistant_idx: int + self, messages: List[Dict[str, Any]], last_model_idx: int ) -> List[Dict[str, Any]]: """ - Strip thinking from old turns but preserve for the last assistant turn. + Strip thinking from old turns but preserve for the last model turn. Per Claude docs: "thinking blocks from previous turns are removed from context" This mimics the API behavior and prevents issues. + + Handles GEMINI format: role "model", "parts" with "thought": true """ for i, msg in enumerate(messages): - if msg.get("role") == "assistant" and i < last_assistant_idx: - # Old turn - strip thinking - msg.pop("reasoning_content", None) - content = msg.get("content") - if isinstance(content, list): + if msg.get("role") == "model" and i < last_model_idx: + # Old turn - strip thinking parts + parts = msg.get("parts", []) + if parts: filtered = [ - item - for item in content - if not ( - isinstance(item, dict) and item.get("type") == "thinking" - ) + p + for p in parts + if not (isinstance(p, dict) and p.get("thought") is True) ] - # Preserve message structure with empty string if needed + + has_function_calls = any( + isinstance(p, dict) and "functionCall" in p for p in filtered + ) + if not filtered: - msg["content"] = "" if not msg.get("tool_calls") else None + msg["parts"] = [{"text": ""}] if not has_function_calls else [] else: - msg["content"] = filtered + msg["parts"] = filtered return messages def _preserve_current_turn_thinking( - self, messages: List[Dict[str, Any]], last_assistant_idx: int + self, messages: List[Dict[str, Any]], last_model_idx: int ) -> List[Dict[str, Any]]: """ - Preserve thinking only for the current (last) assistant turn. + Preserve thinking only for the current (last) model turn. Strip from all previous turns. """ # Same as strip_old_turn_thinking - we keep the last turn intact - return self._strip_old_turn_thinking(messages, last_assistant_idx) + return self._strip_old_turn_thinking(messages, last_model_idx) def _preserve_turn_start_thinking( self, messages: List[Dict[str, Any]], turn_start_idx: int @@ -1536,65 +1561,66 @@ def _preserve_turn_start_thinking( """ Preserve thinking at the turn start message. - In multi-message tool loops, the thinking block is at the FIRST assistant + In multi-message tool loops, the thinking block is at the FIRST model message of the turn (turn_start_idx), not the last one. We need to preserve thinking from the turn start, and strip it from all older turns. + + Handles GEMINI format: role "model", "parts" with "thought": true """ for i, msg in enumerate(messages): - if msg.get("role") == "assistant" and i < turn_start_idx: - # Old turn - strip thinking - msg.pop("reasoning_content", None) - content = msg.get("content") - if isinstance(content, list): + if msg.get("role") == "model" and i < turn_start_idx: + # Old turn - strip thinking parts + parts = msg.get("parts", []) + if parts: filtered = [ - item - for item in content - if not ( - isinstance(item, dict) and item.get("type") == "thinking" - ) + p + for p in parts + if not (isinstance(p, dict) and p.get("thought") is True) ] + + has_function_calls = any( + isinstance(p, dict) and "functionCall" in p for p in filtered + ) + if not filtered: - msg["content"] = "" if not msg.get("tool_calls") else None + msg["parts"] = [{"text": ""}] if not has_function_calls else [] else: - msg["content"] = filtered + msg["parts"] = filtered return messages def _looks_like_compacted_thinking_turn(self, msg: Dict[str, Any]) -> bool: """ Detect if a message looks like it was compacted from a thinking-enabled turn. - Heuristics: - 1. Has tool_calls (typical thinking flow produces tool calls) - 2. Content structure suggests stripped thinking (e.g., starts with tool_use directly) - 3. No text content before tool_use (thinking responses usually have text) + Heuristics (GEMINI format): + 1. Has functionCall parts (typical thinking flow produces tool calls) + 2. No thinking parts (thought: true) + 3. No text content before functionCall (thinking responses usually have text) This is imperfect but helps catch common compaction scenarios. """ - if not msg.get("tool_calls"): + parts = msg.get("parts", []) + if not parts: return False - content = msg.get("content") + has_function_call = any( + isinstance(p, dict) and "functionCall" in p for p in parts + ) - # If content is just tool_use blocks with no text, it might be compacted - if isinstance(content, list): - has_text = any( - isinstance(item, dict) - and item.get("type") == "text" - and item.get("text", "").strip() - for item in content - ) - has_tool_use = any( - isinstance(item, dict) and item.get("type") == "tool_use" - for item in content - ) + if not has_function_call: + return False - # Typical compacted thinking: tool_use without preceding text - # Normal non-thinking response would have explanatory text - if has_tool_use and not has_text: - return True + # Check for text content (not thinking) + has_text = any( + isinstance(p, dict) + and "text" in p + and p.get("text", "").strip() + and not p.get("thought") # Exclude thinking text + for p in parts + ) - # If content is empty/None but has tool_calls, likely compacted - if not content and msg.get("tool_calls"): + # If we have functionCall but no non-thinking text, likely compacted + if not has_text: return True return False @@ -1605,17 +1631,38 @@ def _try_recover_thinking_from_cache( """ Try to recover thinking content from cache for a compacted turn. + Handles GEMINI format: extracts functionCall for cache key lookup, + injects thinking as a part with thought: true. + Returns True if thinking was successfully recovered and injected, False otherwise. """ if turn_start_idx < 0 or turn_start_idx >= len(messages): return False msg = messages[turn_start_idx] + parts = msg.get("parts", []) - # Extract tool_calls for cache key lookup - tool_calls = msg.get("tool_calls", []) - content = msg.get("content", "") - text_content = content if isinstance(content, str) else "" + # Extract text content and build tool_calls structure for cache key lookup + text_content = "" + tool_calls = [] + + for part in parts: + if isinstance(part, dict): + if "text" in part and not part.get("thought"): + text_content = part["text"] + elif "functionCall" in part: + fc = part["functionCall"] + # Convert to OpenAI tool_calls format for cache key compatibility + tool_calls.append( + { + "id": fc.get("id", ""), + "type": "function", + "function": { + "name": fc.get("name", ""), + "arguments": json.dumps(fc.get("args", {})), + }, + } + ) # Generate cache key and try to retrieve cache_key = self._generate_thinking_cache_key(text_content, tool_calls) @@ -1640,19 +1687,14 @@ def _try_recover_thinking_from_cache( ) return False - # Inject the recovered thinking block - thinking_block = { - "type": "thinking", - "thinking": thinking_text, - "signature": signature, + # Inject the recovered thinking part at the beginning (Gemini format) + thinking_part = { + "text": thinking_text, + "thought": True, + "thoughtSignature": signature, } - if isinstance(content, list): - msg["content"] = [thinking_block] + content - elif isinstance(content, str): - msg["content"] = [thinking_block, {"type": "text", "text": content}] - else: - msg["content"] = [thinking_block] + msg["parts"] = [thinking_part] + parts lib_logger.debug( f"[Thinking Sanitization] Recovered thinking from cache: {len(thinking_text)} chars" @@ -1672,7 +1714,7 @@ def _close_tool_loop_for_thinking( Close an incomplete tool loop by injecting synthetic messages to start a new turn. This is used when: - - We're in a tool loop (conversation ends with tool_result) + - We're in a tool loop (conversation ends with functionResponse) - The tool call was made WITHOUT thinking (e.g., by Gemini, non-thinking Claude, or compaction stripped it) - We NOW want to enable thinking @@ -1681,8 +1723,8 @@ def _close_tool_loop_for_thinking( - "To toggle thinking, you must complete the assistant turn first" - A non-tool-result user message ends the turn and allows a fresh start - Solution: - 1. Add synthetic ASSISTANT message to complete the non-thinking turn + Solution (GEMINI format): + 1. Add synthetic MODEL message to complete the non-thinking turn 2. Add synthetic USER message to start a NEW turn 3. Claude will generate thinking for its response to the new turn @@ -1692,47 +1734,61 @@ def _close_tool_loop_for_thinking( # Strip any old thinking first messages = self._strip_all_thinking_blocks(messages) - # Collect tool results from the end of the conversation - tool_results = [] + # Count tool results from the end of the conversation (Gemini format) + tool_result_count = 0 for msg in reversed(messages): - if msg.get("role") == "tool": - tool_results.append(msg) - elif msg.get("role") == "assistant": - break # Stop at the assistant that made the tool calls - - tool_results.reverse() # Put back in order + if msg.get("role") == "user": + parts = msg.get("parts", []) + has_function_response = any( + isinstance(p, dict) and "functionResponse" in p for p in parts + ) + if has_function_response: + tool_result_count += len( + [ + p + for p in parts + if isinstance(p, dict) and "functionResponse" in p + ] + ) + else: + break # Real user message, stop counting + elif msg.get("role") == "model": + break # Stop at the model that made the tool calls # Safety check: if no tool results found, this shouldn't have been called # But handle gracefully with a generic message - if not tool_results: + if tool_result_count == 0: lib_logger.warning( "[Thinking Sanitization] _close_tool_loop_for_thinking called but no tool results found. " "This may indicate malformed conversation history." ) - synthetic_assistant_content = "[Processing previous context.]" - elif len(tool_results) == 1: - synthetic_assistant_content = "[Tool execution completed.]" + synthetic_model_content = "[Processing previous context.]" + elif tool_result_count == 1: + synthetic_model_content = "[Tool execution completed.]" else: - synthetic_assistant_content = ( - f"[{len(tool_results)} tool executions completed.]" + synthetic_model_content = ( + f"[{tool_result_count} tool executions completed.]" ) - # Step 1: Inject synthetic ASSISTANT message to complete the non-thinking turn - synthetic_assistant = { - "role": "assistant", - "content": synthetic_assistant_content, + # Step 1: Inject synthetic MODEL message to complete the non-thinking turn (Gemini format) + synthetic_model = { + "role": "model", + "parts": [{"text": synthetic_model_content}], } - messages.append(synthetic_assistant) + messages.append(synthetic_model) - # Step 2: Inject synthetic USER message to start a NEW turn + # Step 2: Inject synthetic USER message to start a NEW turn (Gemini format) # This allows Claude to generate thinking for its response # The message is minimal and unobtrusive - just triggers a new turn - synthetic_user = {"role": "user", "content": "[Continue]"} + synthetic_user = { + "role": "user", + "parts": [{"text": "[Continue]"}], + } messages.append(synthetic_user) lib_logger.info( f"[Thinking Sanitization] Closed tool loop with synthetic messages. " - f"Assistant: '{synthetic_assistant_content}', User: '[Continue]'. " + f"Model: '{synthetic_model_content}', User: '[Continue]'. " f"Claude will now start a fresh turn with thinking enabled." ) @@ -2981,13 +3037,18 @@ async def acompletion( reasoning_effort is not None and reasoning_effort != "disable" ) - # Sanitize thinking blocks for Claude to prevent 400 errors + # Transform messages to Gemini format FIRST + # This restores thinking from cache if reasoning_content was stripped by client + system_instruction, gemini_contents = self._transform_messages(messages, model) + gemini_contents = self._fix_tool_response_grouping(gemini_contents) + + # Sanitize thinking blocks for Claude AFTER transformation + # Now we can see the full picture including cached thinking that was restored # This handles: context compression, model switching, mid-turn thinking toggle - # Returns (sanitized_messages, force_disable_thinking) force_disable_thinking = False if self._is_claude(model) and self._enable_thinking_sanitization: - messages, force_disable_thinking = self._sanitize_thinking_for_claude( - messages, thinking_enabled + gemini_contents, force_disable_thinking = ( + self._sanitize_thinking_for_claude(gemini_contents, thinking_enabled) ) # If we're in a mid-turn thinking toggle situation, we MUST disable thinking @@ -2996,10 +3057,6 @@ async def acompletion( thinking_enabled = False reasoning_effort = "disable" # Force disable for this request - # Transform messages - system_instruction, gemini_contents = self._transform_messages(messages, model) - gemini_contents = self._fix_tool_response_grouping(gemini_contents) - # Build payload gemini_payload = {"contents": gemini_contents} From ba6dcaa2d096eda9fae0b9e6c4b38ed59396c6d7 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 04:47:58 +0100 Subject: [PATCH 072/221] =?UTF-8?q?fix(antigravity):=20=F0=9F=90=9B=20impr?= =?UTF-8?q?ove=20function=20call=20response=20pairing=20with=20recovery=20?= =?UTF-8?q?strategies?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enhanced the response grouping logic in AntigravityProvider to handle ID mismatches between function calls and their responses more robustly. - Added three-tier matching strategy: direct ID match, function name match, then order-based fallback - Function names are now tracked alongside IDs for orphan response recovery - Responses with "unknown_function" can now be repaired with correct function names - Placeholder responses are automatically created for completely missing tool responses - Fixed insertion position tracking to ensure responses are added immediately after their corresponding model message - Pending groups are now processed in reverse order to prevent index shifting during insertion - Re-enabled debug logging for response collection and group satisfaction - Added comprehensive recovery logging for troubleshooting pairing issues This prevents conversation history corruption when client/proxy systems mutate response IDs or when responses are lost during context processing. --- .../providers/antigravity_provider.py | 154 +++++++++++++++--- 1 file changed, 128 insertions(+), 26 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index d0d46457..e9a081d0 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -2160,9 +2160,18 @@ def _fix_tool_response_grouping( to grouped format (model with calls, user with all responses). IMPORTANT: Preserves ID-based pairing to prevent mismatches. + When IDs don't match, attempts recovery by: + 1. Matching by function name first + 2. Matching by order if names don't match + 3. Inserting placeholder responses if responses are missing + 4. Inserting responses at the CORRECT position (after their corresponding call) """ new_contents = [] - pending_groups = [] # List of {"ids": [id1, id2, ...], "call_indices": [...]} + # Each pending group tracks: + # - ids: expected response IDs + # - func_names: expected function names (for orphan matching) + # - insert_after_idx: position in new_contents where model message was added + pending_groups = [] collected_responses = {} # Dict mapping ID -> response_part for content in contents: @@ -2182,7 +2191,9 @@ def _fix_tool_response_grouping( f"Ignoring duplicate - this may indicate malformed conversation history." ) continue - # lib_logger.debug(f"[Grouping] Collected response for ID: {resp_id}") + lib_logger.debug( + f"[Grouping] Collected response for ID: {resp_id}" + ) collected_responses[resp_id] = resp # Try to satisfy pending groups (newest first) @@ -2197,10 +2208,10 @@ def _fix_tool_response_grouping( collected_responses.pop(gid) for gid in group_ids ] new_contents.append({"parts": group_responses, "role": "user"}) - # lib_logger.debug( - # f"[Grouping] Satisfied group with {len(group_responses)} responses: " - # f"ids={group_ids}" - # ) + lib_logger.debug( + f"[Grouping] Satisfied group with {len(group_responses)} responses: " + f"ids={group_ids}" + ) pending_groups.pop(i) break continue @@ -2213,14 +2224,22 @@ def _fix_tool_response_grouping( fc.get("functionCall", {}).get("id", "") for fc in func_calls ] call_ids = [cid for cid in call_ids if cid] # Filter empty IDs + + # Also extract function names for orphan matching + func_names = [ + fc.get("functionCall", {}).get("name", "") for fc in func_calls + ] + if call_ids: lib_logger.debug( - f"[Grouping] Created pending group expecting {len(call_ids)} responses: ids={call_ids}" + f"[Grouping] Created pending group expecting {len(call_ids)} responses: " + f"ids={call_ids}, names={func_names}" ) pending_groups.append( { "ids": call_ids, - "call_indices": list(range(len(func_calls))), + "func_names": func_names, + "insert_after_idx": len(new_contents) - 1, } ) else: @@ -2228,37 +2247,120 @@ def _fix_tool_response_grouping( # Handle remaining groups (shouldn't happen in well-formed conversations) # Attempt recovery by matching orphans to unsatisfied calls + # Process in REVERSE order of insert_after_idx so insertions don't shift indices + pending_groups.sort(key=lambda g: g["insert_after_idx"], reverse=True) + for group in pending_groups: group_ids = group["ids"] + group_func_names = group.get("func_names", []) + insert_idx = group["insert_after_idx"] + 1 group_responses = [] - for expected_id in group_ids: + lib_logger.debug( + f"[Grouping Recovery] Processing unsatisfied group: " + f"ids={group_ids}, names={group_func_names}, insert_at={insert_idx}" + ) + + for i, expected_id in enumerate(group_ids): + expected_name = group_func_names[i] if i < len(group_func_names) else "" + if expected_id in collected_responses: + # Direct ID match group_responses.append(collected_responses.pop(expected_id)) + lib_logger.debug( + f"[Grouping Recovery] Direct ID match for '{expected_id}'" + ) elif collected_responses: - # Recovery: Match with an orphan response - # This handles cases where client/proxy mutates IDs (e.g. toolu_ -> call_) - # Get the first available orphan ID to maintain order - orphan_id = next(iter(collected_responses)) - orphan_resp = collected_responses.pop(orphan_id) + # Try to find orphan with matching function name first + matched_orphan_id = None - # Fix the ID in the response to match the call - orphan_resp["functionResponse"]["id"] = expected_id + # First pass: match by function name + for orphan_id, orphan_resp in collected_responses.items(): + orphan_name = orphan_resp.get("functionResponse", {}).get( + "name", "" + ) + # Match if names are equal, or if orphan has "unknown_function" (can be fixed) + if orphan_name == expected_name: + matched_orphan_id = orphan_id + lib_logger.debug( + f"[Grouping Recovery] Matched orphan '{orphan_id}' by name '{orphan_name}'" + ) + break - lib_logger.warning( - f"[Grouping] Auto-repaired ID mismatch: mapped response '{orphan_id}' " - f"to call '{expected_id}'" - ) - group_responses.append(orphan_resp) + # Second pass: if no name match, try "unknown_function" orphans + if not matched_orphan_id: + for orphan_id, orphan_resp in collected_responses.items(): + orphan_name = orphan_resp.get("functionResponse", {}).get( + "name", "" + ) + if orphan_name == "unknown_function": + matched_orphan_id = orphan_id + lib_logger.debug( + f"[Grouping Recovery] Matched unknown_function orphan '{orphan_id}' " + f"to expected '{expected_name}'" + ) + break + + # Third pass: if still no match, take first available (order-based) + if not matched_orphan_id: + matched_orphan_id = next(iter(collected_responses)) + lib_logger.debug( + f"[Grouping Recovery] No name match, using first available orphan '{matched_orphan_id}'" + ) - if group_responses: - new_contents.append({"parts": group_responses, "role": "user"}) + if matched_orphan_id: + orphan_resp = collected_responses.pop(matched_orphan_id) + + # Fix the ID in the response to match the call + old_id = orphan_resp["functionResponse"].get("id", "") + orphan_resp["functionResponse"]["id"] = expected_id - if len(group_responses) != len(group_ids): + # Fix the name if it was "unknown_function" + if ( + orphan_resp["functionResponse"].get("name") + == "unknown_function" + and expected_name + ): + orphan_resp["functionResponse"]["name"] = expected_name + lib_logger.info( + f"[Grouping Recovery] Fixed function name from 'unknown_function' to '{expected_name}'" + ) + + lib_logger.warning( + f"[Grouping] Auto-repaired ID mismatch: mapped response '{old_id}' " + f"to call '{expected_id}' (function: {expected_name})" + ) + group_responses.append(orphan_resp) + else: + # No responses available - create placeholder + placeholder_resp = { + "functionResponse": { + "name": expected_name or "unknown_function", + "response": { + "result": { + "error": "Tool response was lost during context processing. " + "This is a recovered placeholder.", + "recovered": True, + } + }, + "id": expected_id, + } + } lib_logger.warning( - f"[Grouping] Partial group satisfaction after repair: " - f"expected {len(group_ids)}, got {len(group_responses)} responses" + f"[Grouping Recovery] Created placeholder response for missing tool: " + f"id='{expected_id}', name='{expected_name}'" ) + group_responses.append(placeholder_resp) + + if group_responses: + # Insert at the correct position (right after the model message with the calls) + new_contents.insert( + insert_idx, {"parts": group_responses, "role": "user"} + ) + lib_logger.info( + f"[Grouping Recovery] Inserted {len(group_responses)} responses at position {insert_idx} " + f"(expected {len(group_ids)})" + ) # Warn about unmatched responses if collected_responses: From 64f7fc091c0e50d014d1e06375bdf6e7ed03b770 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 05:04:04 +0100 Subject: [PATCH 073/221] =?UTF-8?q?docs:=20=F0=9F=93=9A=20update=20documen?= =?UTF-8?q?tation=20for=20enhanced=20claude=20thinking=20sanitization=20an?= =?UTF-8?q?d=20remove=20obsolete=20todo=20file?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit comprehensively updates documentation to reflect the improved Claude extended thinking sanitization system and removes the completed todo.md file. - Enhanced DOCUMENTATION.md with detailed explanations of the robust thinking sanitization system, including: - Clarification that Claude Opus 4.5 always uses the thinking variant (non-thinking version doesn't exist) - Complete sanitization scenario table with new edge cases (function call ID mismatch, missing tool responses, cached conversations) - Detailed implementation notes on Gemini-format message processing and turn state analysis - Three-tier function call response pairing strategy (ID match → name match → fallback) - Recovery mechanisms for cache post-transformation - Increased default max output tokens to 64000 for thinking output - Updated README.md to mention improved function call response pairing with three-tier matching strategy - Removed todo.md as tasks have been completed (thinking sanitization refinements and function call pairing improvements are now implemented) --- DOCUMENTATION.md | 55 ++++++++++++++++-------------- README.md | 3 +- src/rotator_library/pyproject.toml | 2 +- todo.md | 7 ---- 4 files changed, 33 insertions(+), 34 deletions(-) delete mode 100644 todo.md diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index 29ea7838..39b266b0 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -420,10 +420,11 @@ The most sophisticated provider implementation, supporting Google's internal Ant **Claude Opus 4.5 (NEW!):** - Anthropic's most powerful model, now available via Antigravity proxy -- Uses internal model name `claude-opus-4-5-thinking` when reasoning is enabled -- Uses `thinkingBudget` parameter for extended thinking control +- **Always uses thinking variant** - `claude-opus-4-5-thinking` is the only available variant (non-thinking version doesn't exist) +- Uses `thinkingBudget` parameter for extended thinking control (-1 for auto, 0 to disable, or specific token count) - Full support for tool use with schema cleaning - Same thinking preservation and sanitization features as Sonnet +- Increased default max output tokens to 64000 to accommodate thinking output **Claude Sonnet 4.5:** - Proxied through Antigravity API (uses internal model name `claude-sonnet-4-5-thinking`) @@ -475,7 +476,7 @@ ANTIGRAVITY_GEMINI3_SYSTEM_INSTRUCTION="..." # Full system prompt #### Claude Extended Thinking Sanitization -The provider includes automatic sanitization for Claude's extended thinking mode, handling common error scenarios: +The provider now includes robust automatic sanitization for Claude's extended thinking mode, handling all common error scenarios with conversation history. **Problem**: Claude's extended thinking API requires strict consistency in thinking blocks: - If thinking is enabled, the final assistant turn must start with a thinking block @@ -491,38 +492,42 @@ The provider includes automatic sanitization for Claude's extended thinking mode | Tool loop WITHOUT thinking + thinking enabled | **Inject synthetic closure** to start fresh turn with thinking | | Thinking disabled | Strip all thinking blocks | | Normal conversation (no tool loop) | Strip old thinking, new response adds thinking naturally | +| Function call ID mismatch | Three-tier recovery: ID match → name match → fallback | +| Missing tool responses | Automatic placeholder injection | +| Compacted/cached conversations | Recover thinking from cache post-transformation | -**Solution**: The `_sanitize_thinking_for_claude()` method: -- Analyzes conversation state to detect incomplete tool use loops -- When enabling thinking in a tool loop that started without thinking: - - Injects a minimal synthetic assistant message: `"[Tool execution completed. Processing results.]"` - - This **closes** the previous turn, allowing Claude to start a **fresh turn with thinking** -- Strips thinking from old turns (Claude API ignores them anyway) -- Preserves thinking when the turn was started with thinking enabled +**Key Implementation Details**: -**Key Insight**: Instead of force-disabling thinking, we close the tool loop with a synthetic message. This allows seamless model switching (e.g., Gemini → Claude with thinking) without losing the ability to think. +The `_sanitize_thinking_for_claude()` method now: +- Operates on Gemini-format messages (`parts[]` with `"thought": true` markers) +- Detects tool results as user messages with `functionResponse` parts +- Uses `_analyze_turn_state()` to classify conversation state on Gemini format +- Recovers thinking from cache when client strips reasoning_content +- When enabling thinking in a tool loop started without thinking: + - Injects synthetic assistant message to close the previous turn + - Allows Claude to start fresh turn with thinking capability -**Example**: +**Function Call Response Grouping**: + +The enhanced pairing system ensures conversation history integrity: ``` -Before sanitization: - User: "What's the weather?" - Assistant: [tool_use: get_weather] ← Made by Gemini (no thinking) - User: [tool_result: "20C sunny"] - -After sanitization (thinking enabled): - User: "What's the weather?" - Assistant: [tool_use: get_weather] - User: [tool_result: "20C sunny"] - Assistant: "[Tool execution completed. Processing results.]" ← INJECTED - - → Claude now starts a NEW turn and CAN think! +Problem: Client/proxy may mutate response IDs or lose responses during context processing + +Solution: +1. Try direct ID match (tool_call_id == response.id) +2. If no match, try function name match (tool.name == response.name) +3. If still no match, use order-based fallback (nth tool → nth response) +4. Repair "unknown_function" responses with correct names +5. Create placeholders for completely missing responses ``` **Configuration**: ```env -ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION=true # Enable/disable auto-correction +ANTIGRAVITY_CLAUDE_THINKING_SANITIZATION=true # Enable/disable auto-correction (default: true) ``` +**Note**: These fixes ensure Claude thinking mode works seamlessly with tool use, model switching, context compression, and cached conversations. No manual intervention required. + #### File Logging Optional transaction logging for debugging: diff --git a/README.md b/README.md index 91971102..85df3b70 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,8 @@ This project provides a powerful solution for developers building complex applic - Claude Sonnet 4.5 with extended thinking support - Thought signature caching for multi-turn conversations - Tool hallucination prevention via parameter signature injection - - Automatic thinking block sanitization for Claude models + - Automatic thinking block sanitization for Claude models (with recovery strategies) + - Improved function call response pairing with three-tier matching strategy - Note: Claude thinking mode requires careful conversation state management (see [Antigravity documentation](DOCUMENTATION.md#antigravity-claude-extended-thinking-sanitization) for details) - **🆕 Credential Prioritization**: Automatic tier detection and priority-based credential selection ensures paid-tier credentials are used for premium models that require them. - **🆕 Weighted Random Rotation**: Configurable credential rotation strategy - choose between deterministic (perfect balance) or weighted random (unpredictable, harder to fingerprint) selection. diff --git a/src/rotator_library/pyproject.toml b/src/rotator_library/pyproject.toml index 4cfa41a3..1ad55af7 100644 --- a/src/rotator_library/pyproject.toml +++ b/src/rotator_library/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "rotator_library" -version = "0.95" +version = "1.0" authors = [ { name="Mirrowel", email="nuh@uh.com" }, ] diff --git a/todo.md b/todo.md deleted file mode 100644 index 5966e4b1..00000000 --- a/todo.md +++ /dev/null @@ -1,7 +0,0 @@ -~~Refine claude injection to inject even if we have correct thinking - to force it to think if we made ultrathink prompt. If last msg is tool use and you prompt - it never thinks again.~~ Maybe done - -Anthropic translation and anthropic compatible endpoint. - -Refine for deployment. - - From 42bd5aeb74855ff82d0d06787192cf87d4ac3982 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 06:12:50 +0100 Subject: [PATCH 074/221] =?UTF-8?q?fix(antigravity):=20=F0=9F=90=9B=20prev?= =?UTF-8?q?ent=20unescaping=20of=20intentional=20quotes=20and=20backslashe?= =?UTF-8?q?s=20in=20strings?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recursive JSON string parser was incorrectly unescaping all strings containing escape sequences, including those with intentional \" and \\ escapes. This corrupted content like JSON embedded in YAML configurations, causing oldString and newString to become identical when they should differ. - Added logic to differentiate between control character escapes (\n, \t) and intentional escapes (\", \\) - Only unescape strings with control character escapes if they don't contain intentional escapes - Enhanced debug logging with string snippets for better troubleshooting - Updated comments to clarify the reasoning and provide concrete examples --- .../providers/antigravity_provider.py | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index e9a081d0..fb63a5d9 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -253,16 +253,27 @@ def _recursively_parse_json_strings(obj: Any) -> Any: elif isinstance(obj, str): stripped = obj.strip() - # Check if string contains common escape sequences that need unescaping - # This handles cases where diff content or other text has literal \n instead of newlines - if "\\n" in obj or "\\t" in obj or '\\"' in obj or "\\\\" in obj: + # Check if string contains control character escape sequences that need unescaping + # This handles cases where diff content has literal \n or \t instead of actual newlines/tabs + # + # IMPORTANT: We intentionally do NOT unescape strings containing \" or \\ + # because these are typically intentional escapes in code/config content + # (e.g., JSON embedded in YAML: BOT_NAMES_JSON: '["mirrobot", ...]') + # Unescaping these would corrupt the content and cause issues like + # oldString and newString becoming identical when they should differ. + has_control_char_escapes = "\\n" in obj or "\\t" in obj + has_intentional_escapes = '\\"' in obj or "\\\\" in obj + + if has_control_char_escapes and not has_intentional_escapes: try: # Use json.loads with quotes to properly unescape the string - # This converts \n -> newline, \t -> tab, \" -> quote, etc. + # This converts \n -> newline, \t -> tab unescaped = json.loads(f'"{obj}"') + # Log the fix with a snippet for debugging + snippet = obj[:80] + "..." if len(obj) > 80 else obj lib_logger.debug( - f"[Antigravity] Unescaped string content: " - f"{len(obj) - len(unescaped)} chars changed" + f"[Antigravity] Unescaped control chars in string: " + f"{len(obj) - len(unescaped)} chars changed. Snippet: {snippet!r}" ) return unescaped except (json.JSONDecodeError, ValueError): From edef5b9f3a7b90300f857864b325727e7e5a570c Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 06:19:23 +0100 Subject: [PATCH 075/221] ci: Agent compliance check fix --- .github/workflows/compliance-check.yml | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/.github/workflows/compliance-check.yml b/.github/workflows/compliance-check.yml index 936eb270..876c87a0 100644 --- a/.github/workflows/compliance-check.yml +++ b/.github/workflows/compliance-check.yml @@ -87,7 +87,7 @@ jobs: # BASIC CONFIGURATION # ----------------------------------------------------------------------- PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number || inputs.pr_number || github.event.workflow_run.pull_requests[0].number }} - BOT_NAMES_JSON: '[\"mirrobot\", \"mirrobot-agent\", \"mirrobot-agent[bot]\"]' + BOT_NAMES_JSON: '["mirrobot", "mirrobot-agent", "mirrobot-agent[bot]"]' # ----------------------------------------------------------------------- # FEATURE TOGGLES @@ -179,7 +179,7 @@ jobs: opencode-fast-model: ${{ secrets.OPENCODE_FAST_MODEL }} custom-providers-json: ${{ secrets.CUSTOM_PROVIDERS_JSON }} - # ====================================================================== + # ====================================================================== # CONDITIONAL WAIT: Wait for PR Review to Complete # ====================================================================== # Only wait when triggered by ready_for_review event @@ -241,7 +241,10 @@ jobs: echo "head_sha=$(echo "$pr_json" | jq -r .headRefOid)" >> $GITHUB_OUTPUT echo "pr_title=$(echo "$pr_json" | jq -r .title)" >> $GITHUB_OUTPUT - echo "pr_author=$(echo "$pr_json" | jq -r .author.login)" >> $GITHUB_OUTPUT + + # Extract author to shell variable first (can't self-reference step outputs) + pr_author=$(echo "$pr_json" | jq -r .author.login) + echo "pr_author=$pr_author" >> $GITHUB_OUTPUT pr_body=$(echo "$pr_json" | jq -r '.body // ""') echo "pr_body<> $GITHUB_OUTPUT @@ -262,7 +265,7 @@ jobs: # Requested reviewers for mentions reviewers=$(echo "$pr_json" | jq -r '.reviewRequests[]? | .login' | tr '\n' ' ') - mentions="@${{ steps.pr_info.outputs.pr_author }}" + mentions="@$pr_author" if [ -n "$reviewers" ]; then for reviewer in $reviewers; do mentions="$mentions @$reviewer" From cd3d0e6992c285c492ef9aa04b68a5fb78e68afd Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 06:27:46 +0100 Subject: [PATCH 076/221] typo fix --- .github/workflows/compliance-check.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/compliance-check.yml b/.github/workflows/compliance-check.yml index 876c87a0..c3d403c4 100644 --- a/.github/workflows/compliance-check.yml +++ b/.github/workflows/compliance-check.yml @@ -115,7 +115,7 @@ jobs: "description": "When code changes affect the build or CI process, verify build.yml is updated with new steps, jobs, or release configurations. Check that code changes are reflected in build matrix, deploy steps, and CI/CD pipeline.", "files": [ ".github/workflows/build.yml", - ".github/workflows/cleanup.yml", + ".github/workflows/cleanup.yml" ] }, { From 7d43e9832869373385cdce778d9b48e74d3c6d49 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 06:33:16 +0100 Subject: [PATCH 077/221] ci: Guess what? yet another fix --- .github/workflows/compliance-check.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/compliance-check.yml b/.github/workflows/compliance-check.yml index c3d403c4..4b561940 100644 --- a/.github/workflows/compliance-check.yml +++ b/.github/workflows/compliance-check.yml @@ -391,7 +391,7 @@ jobs: echo "" >> /tmp/file_groups.txt # Parse JSON and format for prompt - echo '${{ env.FILE_GROUPS_JSON }}' | jq -r '.[] | + echo "$FILE_GROUPS_JSON" | jq -r '.[] | "Group: \(.name)\n" + "Description: \(.description)\n" + "Files:\n" + From 81e9ff5e527814b69e427b6b5da7e01f59ab0037 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 07:22:41 +0100 Subject: [PATCH 078/221] =?UTF-8?q?fix(oauth):=20=F0=9F=90=9B=20escape=20r?= =?UTF-8?q?ich=20markup=20in=20oauth=20authorization=20urls?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prevent Rich markup interpretation issues when displaying OAuth authorization URLs in terminal output. - Import `rich.markup.escape` to properly escape special characters (=, &, etc.) in URLs - Add extensive inline documentation explaining the escaping rationale and known terminal compatibility issues - Apply URL escaping to authorization URLs in Google OAuth, iFlow, and Qwen Code providers - Refine headless environment detection to exclude macOS from DISPLAY checks (macOS uses Quartz, not X11) - Improve code formatting consistency (string quotes, line wrapping) across OAuth providers The escaped URLs display correctly in all terminal configurations while remaining clickable in supported terminals (iTerm2, Windows Terminal, etc.). --- .../providers/google_oauth_base.py | 431 ++++++++++++------ .../providers/iflow_auth_base.py | 378 ++++++++++----- .../providers/qwen_auth_base.py | 402 +++++++++++----- .../utils/headless_detection.py | 54 ++- 4 files changed, 883 insertions(+), 382 deletions(-) diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index 3f1ed9d6..0b34153b 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -16,35 +16,37 @@ from rich.console import Console from rich.panel import Panel from rich.text import Text +from rich.markup import escape as rich_escape from ..utils.headless_detection import is_headless_environment -lib_logger = logging.getLogger('rotator_library') +lib_logger = logging.getLogger("rotator_library") console = Console() + class GoogleOAuthBase: """ Base class for Google OAuth2 authentication providers. - + Subclasses must override: - CLIENT_ID: OAuth client ID - CLIENT_SECRET: OAuth client secret - OAUTH_SCOPES: List of OAuth scopes - ENV_PREFIX: Prefix for environment variables (e.g., "GEMINI_CLI", "ANTIGRAVITY") - + Subclasses may optionally override: - CALLBACK_PORT: Local OAuth callback server port (default: 8085) - CALLBACK_PATH: OAuth callback path (default: "/oauth2callback") - REFRESH_EXPIRY_BUFFER_SECONDS: Time buffer before token expiry (default: 30 minutes) """ - + # Subclasses MUST override these CLIENT_ID: str = None CLIENT_SECRET: str = None OAUTH_SCOPES: list = None ENV_PREFIX: str = None - + # Subclasses MAY override these TOKEN_URI: str = "https://oauth2.googleapis.com/token" USER_INFO_URI: str = "https://www.googleapis.com/oauth2/v1/userinfo" @@ -57,49 +59,65 @@ def __init__(self): if self.CLIENT_ID is None: raise NotImplementedError(f"{self.__class__.__name__} must set CLIENT_ID") if self.CLIENT_SECRET is None: - raise NotImplementedError(f"{self.__class__.__name__} must set CLIENT_SECRET") + raise NotImplementedError( + f"{self.__class__.__name__} must set CLIENT_SECRET" + ) if self.OAUTH_SCOPES is None: - raise NotImplementedError(f"{self.__class__.__name__} must set OAUTH_SCOPES") + raise NotImplementedError( + f"{self.__class__.__name__} must set OAUTH_SCOPES" + ) if self.ENV_PREFIX is None: raise NotImplementedError(f"{self.__class__.__name__} must set ENV_PREFIX") - + self._credentials_cache: Dict[str, Dict[str, Any]] = {} self._refresh_locks: Dict[str, asyncio.Lock] = {} - self._locks_lock = asyncio.Lock() # Protects the locks dict from race conditions + self._locks_lock = ( + asyncio.Lock() + ) # Protects the locks dict from race conditions # [BACKOFF TRACKING] Track consecutive failures per credential - self._refresh_failures: Dict[str, int] = {} # Track consecutive failures per credential - self._next_refresh_after: Dict[str, float] = {} # Track backoff timers (Unix timestamp) - + self._refresh_failures: Dict[ + str, int + ] = {} # Track consecutive failures per credential + self._next_refresh_after: Dict[ + str, float + ] = {} # Track backoff timers (Unix timestamp) + # [QUEUE SYSTEM] Sequential refresh processing self._refresh_queue: asyncio.Queue = asyncio.Queue() self._queued_credentials: set = set() # Track credentials already in queue - self._unavailable_credentials: set = set() # Mark credentials unavailable during re-auth + self._unavailable_credentials: set = ( + set() + ) # Mark credentials unavailable during re-auth self._queue_tracking_lock = asyncio.Lock() # Protects queue sets - self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task + self._queue_processor_task: Optional[asyncio.Task] = ( + None # Background worker task + ) def _parse_env_credential_path(self, path: str) -> Optional[str]: """ Parse a virtual env:// path and return the credential index. - + Supported formats: - "env://provider/0" - Legacy single credential (no index in env var names) - "env://provider/1" - First numbered credential (PROVIDER_1_ACCESS_TOKEN) - "env://provider/2" - Second numbered credential, etc. - + Returns: The credential index as string ("0" for legacy, "1", "2", etc. for numbered) or None if path is not an env:// path """ if not path.startswith("env://"): return None - + # Parse: env://provider/index parts = path[6:].split("/") # Remove "env://" prefix if len(parts) >= 2: return parts[1] # Return the index return "0" # Default to legacy format - def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dict[str, Any]]: + def _load_from_env( + self, credential_index: Optional[str] = None + ) -> Optional[Dict[str, Any]]: """ Load OAuth credentials from environment variables for stateless deployments. @@ -133,7 +151,7 @@ def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dic # Legacy format: PROVIDER_ACCESS_TOKEN prefix = self.ENV_PREFIX default_email = "env-user" - + access_token = os.getenv(f"{prefix}_ACCESS_TOKEN") refresh_token = os.getenv(f"{prefix}_REFRESH_TOKEN") @@ -148,7 +166,9 @@ def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dic try: expiry_date = float(expiry_str) except ValueError: - lib_logger.warning(f"Invalid {prefix}_EXPIRY_DATE value: {expiry_str}, using 0") + lib_logger.warning( + f"Invalid {prefix}_EXPIRY_DATE value: {expiry_str}, using 0" + ) expiry_date = 0 creds = { @@ -163,15 +183,16 @@ def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dic "email": os.getenv(f"{prefix}_EMAIL", default_email), "last_check_timestamp": time.time(), "loaded_from_env": True, # Flag to indicate env-based credentials - "env_credential_index": credential_index or "0" # Track which env credential this is - } + "env_credential_index": credential_index + or "0", # Track which env credential this is + }, } # Add project_id if provided project_id = os.getenv(f"{prefix}_PROJECT_ID") if project_id: creds["_proxy_metadata"]["project_id"] = project_id - + # Add tier if provided tier = os.getenv(f"{prefix}_TIER") if tier: @@ -193,24 +214,32 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: # Load from environment variables with specific index env_creds = self._load_from_env(credential_index) if env_creds: - lib_logger.info(f"Using {self.ENV_PREFIX} credentials from environment variables (index: {credential_index})") + lib_logger.info( + f"Using {self.ENV_PREFIX} credentials from environment variables (index: {credential_index})" + ) self._credentials_cache[path] = env_creds return env_creds else: - raise IOError(f"Environment variables for {self.ENV_PREFIX} credential index {credential_index} not found") + raise IOError( + f"Environment variables for {self.ENV_PREFIX} credential index {credential_index} not found" + ) # For file paths, first try loading from legacy env vars (for backwards compatibility) env_creds = self._load_from_env() if env_creds: - lib_logger.info(f"Using {self.ENV_PREFIX} credentials from environment variables") + lib_logger.info( + f"Using {self.ENV_PREFIX} credentials from environment variables" + ) # Cache env-based credentials using the path as key self._credentials_cache[path] = env_creds return env_creds # Fall back to file-based loading try: - lib_logger.debug(f"Loading {self.ENV_PREFIX} credentials from file: {path}") - with open(path, 'r') as f: + lib_logger.debug( + f"Loading {self.ENV_PREFIX} credentials from file: {path}" + ) + with open(path, "r") as f: creds = json.load(f) # Handle gcloud-style creds file which nest tokens under "credential" if "credential" in creds: @@ -218,11 +247,17 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: self._credentials_cache[path] = creds return creds except FileNotFoundError: - raise IOError(f"{self.ENV_PREFIX} OAuth credential file not found at '{path}'") + raise IOError( + f"{self.ENV_PREFIX} OAuth credential file not found at '{path}'" + ) except Exception as e: - raise IOError(f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}") + raise IOError( + f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}" + ) except Exception as e: - raise IOError(f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}") + raise IOError( + f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}" + ) async def _save_credentials(self, path: str, creds: Dict[str, Any]): # Don't save to file if credentials were loaded from environment @@ -241,10 +276,12 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): tmp_path = None try: # Create temp file in same directory as target (ensures same filesystem) - tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json', text=True) + tmp_fd, tmp_path = tempfile.mkstemp( + dir=parent_dir, prefix=".tmp_", suffix=".json", text=True + ) # Write JSON to temp file - with os.fdopen(tmp_fd, 'w') as f: + with os.fdopen(tmp_fd, "w") as f: json.dump(creds, f, indent=2) tmp_fd = None # fdopen closes the fd @@ -261,10 +298,14 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): # Update cache AFTER successful file write (prevents cache/file inconsistency) self._credentials_cache[path] = creds - lib_logger.debug(f"Saved updated {self.ENV_PREFIX} OAuth credentials to '{path}' (atomic write).") + lib_logger.debug( + f"Saved updated {self.ENV_PREFIX} OAuth credentials to '{path}' (atomic write)." + ) except Exception as e: - lib_logger.error(f"Failed to save updated {self.ENV_PREFIX} OAuth credentials to '{path}': {e}") + lib_logger.error( + f"Failed to save updated {self.ENV_PREFIX} OAuth credentials to '{path}': {e}" + ) # Clean up temp file if it still exists if tmp_fd is not None: try: @@ -279,20 +320,26 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): raise def _is_token_expired(self, creds: Dict[str, Any]) -> bool: - expiry = creds.get("token_expiry") # gcloud format - if not expiry: # gemini-cli format - expiry_timestamp = creds.get("expiry_date", 0) / 1000 + expiry = creds.get("token_expiry") # gcloud format + if not expiry: # gemini-cli format + expiry_timestamp = creds.get("expiry_date", 0) / 1000 else: expiry_timestamp = time.mktime(time.strptime(expiry, "%Y-%m-%dT%H:%M:%SZ")) return expiry_timestamp < time.time() + self.REFRESH_EXPIRY_BUFFER_SECONDS - async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = False) -> Dict[str, Any]: + async def _refresh_token( + self, path: str, creds: Dict[str, Any], force: bool = False + ) -> Dict[str, Any]: async with await self._get_lock(path): # Skip the expiry check if a refresh is being forced - if not force and not self._is_token_expired(self._credentials_cache.get(path, creds)): + if not force and not self._is_token_expired( + self._credentials_cache.get(path, creds) + ): return self._credentials_cache.get(path, creds) - lib_logger.debug(f"Refreshing {self.ENV_PREFIX} OAuth token for '{Path(path).name}' (forced: {force})...") + lib_logger.debug( + f"Refreshing {self.ENV_PREFIX} OAuth token for '{Path(path).name}' (forced: {force})..." + ) refresh_token = creds.get("refresh_token") if not refresh_token: raise ValueError("No refresh_token found in credentials file.") @@ -306,12 +353,18 @@ async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = F async with httpx.AsyncClient() as client: for attempt in range(max_retries): try: - response = await client.post(self.TOKEN_URI, data={ - "client_id": creds.get("client_id", self.CLIENT_ID), - "client_secret": creds.get("client_secret", self.CLIENT_SECRET), - "refresh_token": refresh_token, - "grant_type": "refresh_token", - }, timeout=30.0) + response = await client.post( + self.TOKEN_URI, + data={ + "client_id": creds.get("client_id", self.CLIENT_ID), + "client_secret": creds.get( + "client_secret", self.CLIENT_SECRET + ), + "refresh_token": refresh_token, + "grant_type": "refresh_token", + }, + timeout=30.0, + ) response.raise_for_status() new_token_data = response.json() break # Success, exit retry loop @@ -332,7 +385,9 @@ async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = F elif status_code == 429: # Rate limit - honor Retry-After header if present retry_after = int(e.response.headers.get("Retry-After", 60)) - lib_logger.warning(f"Rate limited (HTTP 429), retry after {retry_after}s") + lib_logger.warning( + f"Rate limited (HTTP 429), retry after {retry_after}s" + ) if attempt < max_retries - 1: await asyncio.sleep(retry_after) continue @@ -341,8 +396,10 @@ async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = F elif status_code >= 500 and status_code < 600: # Server error - retry with exponential backoff if attempt < max_retries - 1: - wait_time = 2 ** attempt # 1s, 2s, 4s - lib_logger.warning(f"Server error (HTTP {status_code}), retry {attempt + 1}/{max_retries} in {wait_time}s") + wait_time = 2**attempt # 1s, 2s, 4s + lib_logger.warning( + f"Server error (HTTP {status_code}), retry {attempt + 1}/{max_retries} in {wait_time}s" + ) await asyncio.sleep(wait_time) continue raise # Final attempt failed @@ -355,22 +412,30 @@ async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = F # Network errors - retry with backoff last_error = e if attempt < max_retries - 1: - wait_time = 2 ** attempt - lib_logger.warning(f"Network error during refresh: {e}, retry {attempt + 1}/{max_retries} in {wait_time}s") + wait_time = 2**attempt + lib_logger.warning( + f"Network error during refresh: {e}, retry {attempt + 1}/{max_retries} in {wait_time}s" + ) await asyncio.sleep(wait_time) continue raise # [INVALID GRANT RE-AUTH] Trigger OAuth flow if refresh token is invalid if needs_reauth: - lib_logger.info(f"Starting re-authentication for '{Path(path).name}'...") + lib_logger.info( + f"Starting re-authentication for '{Path(path).name}'..." + ) try: # Call initialize_token to trigger OAuth flow new_creds = await self.initialize_token(path) return new_creds except Exception as reauth_error: - lib_logger.error(f"Re-authentication failed for '{Path(path).name}': {reauth_error}") - raise ValueError(f"Refresh token invalid and re-authentication failed: {reauth_error}") + lib_logger.error( + f"Re-authentication failed for '{Path(path).name}': {reauth_error}" + ) + raise ValueError( + f"Refresh token invalid and re-authentication failed: {reauth_error}" + ) # If we exhausted retries without success if new_token_data is None: @@ -379,7 +444,7 @@ async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = F # [FIX 1] Update OAuth token fields from response creds["access_token"] = new_token_data["access_token"] expiry_timestamp = time.time() + new_token_data["expires_in"] - creds["expiry_date"] = expiry_timestamp * 1000 # gemini-cli format + creds["expiry_date"] = expiry_timestamp * 1000 # gemini-cli format # [FIX 2] Update refresh_token if server provided a new one (rare but possible with Google OAuth) if "refresh_token" in new_token_data: @@ -405,10 +470,20 @@ async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = F creds["_proxy_metadata"]["last_check_timestamp"] = time.time() # [VALIDATION] Verify refreshed credentials have all required fields - required_fields = ["access_token", "refresh_token", "client_id", "client_secret", "token_uri"] - missing_fields = [field for field in required_fields if not creds.get(field)] + required_fields = [ + "access_token", + "refresh_token", + "client_id", + "client_secret", + "token_uri", + ] + missing_fields = [ + field for field in required_fields if not creds.get(field) + ] if missing_fields: - raise ValueError(f"Refreshed credentials missing required fields: {missing_fields}") + raise ValueError( + f"Refreshed credentials missing required fields: {missing_fields}" + ) # [VALIDATION] Optional: Test that the refreshed token is actually usable try: @@ -416,17 +491,23 @@ async def _refresh_token(self, path: str, creds: Dict[str, Any], force: bool = F test_response = await client.get( self.USER_INFO_URI, headers={"Authorization": f"Bearer {creds['access_token']}"}, - timeout=5.0 + timeout=5.0, ) test_response.raise_for_status() - lib_logger.debug(f"Token validation successful for '{Path(path).name}'") + lib_logger.debug( + f"Token validation successful for '{Path(path).name}'" + ) except Exception as e: - lib_logger.warning(f"Refreshed token validation failed for '{Path(path).name}': {e}") + lib_logger.warning( + f"Refreshed token validation failed for '{Path(path).name}': {e}" + ) # Don't fail the refresh - the token might still work for other endpoints # But log it for debugging purposes await self._save_credentials(path, creds) - lib_logger.debug(f"Successfully refreshed {self.ENV_PREFIX} OAuth token for '{Path(path).name}'.") + lib_logger.debug( + f"Successfully refreshed {self.ENV_PREFIX} OAuth token for '{Path(path).name}'." + ) return creds async def proactively_refresh(self, credential_path: str): @@ -451,11 +532,15 @@ def is_credential_available(self, path: str) -> bool: async def _ensure_queue_processor_running(self): """Lazily starts the queue processor if not already running.""" if self._queue_processor_task is None or self._queue_processor_task.done(): - self._queue_processor_task = asyncio.create_task(self._process_refresh_queue()) + self._queue_processor_task = asyncio.create_task( + self._process_refresh_queue() + ) - async def _queue_refresh(self, path: str, force: bool = False, needs_reauth: bool = False): + async def _queue_refresh( + self, path: str, force: bool = False, needs_reauth: bool = False + ): """Add a credential to the refresh queue if not already queued. - + Args: path: Credential file path force: Force refresh even if not expired @@ -470,9 +555,11 @@ async def _queue_refresh(self, path: str, force: bool = False, needs_reauth: boo if now < backoff_until: # Credential is in backoff for automated refresh, do not queue remaining = int(backoff_until - now) - lib_logger.debug(f"Skipping automated refresh for '{Path(path).name}' (in backoff for {remaining}s)") + lib_logger.debug( + f"Skipping automated refresh for '{Path(path).name}' (in backoff for {remaining}s)" + ) return - + async with self._queue_tracking_lock: if path not in self._queued_credentials: self._queued_credentials.add(path) @@ -488,14 +575,13 @@ async def _process_refresh_queue(self): # Wait for an item with timeout to allow graceful shutdown try: path, force, needs_reauth = await asyncio.wait_for( - self._refresh_queue.get(), - timeout=60.0 + self._refresh_queue.get(), timeout=60.0 ) except asyncio.TimeoutError: # No items for 60s, exit to save resources self._queue_processor_task = None return - + try: # Perform the actual refresh (still using per-credential lock) async with await self._get_lock(path): @@ -506,16 +592,16 @@ async def _process_refresh_queue(self): async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) continue - + # Perform refresh if not creds: creds = await self._load_credentials(path) await self._refresh_token(path, creds, force=force) - + # SUCCESS: Mark as available again async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) - + finally: # Remove from queued set async with self._queue_tracking_lock: @@ -530,18 +616,26 @@ async def _process_refresh_queue(self): async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) - async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + async def initialize_token( + self, creds_or_path: Union[Dict[str, Any], str] + ) -> Dict[str, Any]: path = creds_or_path if isinstance(creds_or_path, str) else None # Get display name from metadata if available, otherwise derive from path if isinstance(creds_or_path, dict): - display_name = creds_or_path.get("_proxy_metadata", {}).get("display_name", "in-memory object") + display_name = creds_or_path.get("_proxy_metadata", {}).get( + "display_name", "in-memory object" + ) else: display_name = Path(path).name if path else "in-memory object" - lib_logger.debug(f"Initializing {self.ENV_PREFIX} token for '{display_name}'...") + lib_logger.debug( + f"Initializing {self.ENV_PREFIX} token for '{display_name}'..." + ) try: - creds = await self._load_credentials(creds_or_path) if path else creds_or_path + creds = ( + await self._load_credentials(creds_or_path) if path else creds_or_path + ) reason = "" if not creds.get("refresh_token"): reason = "refresh token is missing" @@ -553,34 +647,51 @@ async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> D try: return await self._refresh_token(path, creds) except Exception as e: - lib_logger.warning(f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login.") + lib_logger.warning( + f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login." + ) + + lib_logger.warning( + f"{self.ENV_PREFIX} OAuth token for '{display_name}' needs setup: {reason}." + ) - lib_logger.warning(f"{self.ENV_PREFIX} OAuth token for '{display_name}' needs setup: {reason}.") - # [HEADLESS DETECTION] Check if running in headless environment is_headless = is_headless_environment() - + auth_code_future = asyncio.get_event_loop().create_future() server = None async def handle_callback(reader, writer): try: request_line_bytes = await reader.readline() - if not request_line_bytes: return - path_str = request_line_bytes.decode('utf-8').strip().split(' ')[1] - while await reader.readline() != b'\r\n': pass + if not request_line_bytes: + return + path_str = ( + request_line_bytes.decode("utf-8").strip().split(" ")[1] + ) + while await reader.readline() != b"\r\n": + pass from urllib.parse import urlparse, parse_qs + query_params = parse_qs(urlparse(path_str).query) - writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n") - if 'code' in query_params: + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n" + ) + if "code" in query_params: if not auth_code_future.done(): - auth_code_future.set_result(query_params['code'][0]) - writer.write(b"

Authentication successful!

You can close this window.

") + auth_code_future.set_result(query_params["code"][0]) + writer.write( + b"

Authentication successful!

You can close this window.

" + ) else: - error = query_params.get('error', ['Unknown error'])[0] + error = query_params.get("error", ["Unknown error"])[0] if not auth_code_future.done(): - auth_code_future.set_exception(Exception(f"OAuth failed: {error}")) - writer.write(f"

Authentication Failed

Error: {error}. Please try again.

".encode()) + auth_code_future.set_exception( + Exception(f"OAuth failed: {error}") + ) + writer.write( + f"

Authentication Failed

Error: {error}. Please try again.

".encode() + ) await writer.drain() except Exception as e: lib_logger.error(f"Error in OAuth callback handler: {e}") @@ -588,15 +699,25 @@ async def handle_callback(reader, writer): writer.close() try: - server = await asyncio.start_server(handle_callback, '127.0.0.1', self.CALLBACK_PORT) + server = await asyncio.start_server( + handle_callback, "127.0.0.1", self.CALLBACK_PORT + ) from urllib.parse import urlencode - auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode({ - "client_id": self.CLIENT_ID, - "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", - "scope": " ".join(self.OAUTH_SCOPES), - "access_type": "offline", "response_type": "code", "prompt": "consent" - }) - + + auth_url = ( + "https://accounts.google.com/o/oauth2/v2/auth?" + + urlencode( + { + "client_id": self.CLIENT_ID, + "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", + "scope": " ".join(self.OAUTH_SCOPES), + "access_type": "offline", + "response_type": "code", + "prompt": "consent", + } + ) + ) + # [HEADLESS SUPPORT] Display appropriate instructions if is_headless: auth_panel_text = Text.from_markup( @@ -606,68 +727,118 @@ async def handle_callback(reader, writer): else: auth_panel_text = Text.from_markup( "1. Your browser will now open to log in and authorize the application.\n" - "2. If it doesn't open automatically, please open the URL below manually." + "2. If it doesn't open automatically, please open the URL below manually." + ) + + console.print( + Panel( + auth_panel_text, + title=f"{self.ENV_PREFIX} OAuth Setup for [bold yellow]{display_name}[/bold yellow]", + style="bold blue", ) - - console.print(Panel(auth_panel_text, title=f"{self.ENV_PREFIX} OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue")) - console.print(f"[bold]URL:[/bold] [link={auth_url}]{auth_url}[/link]\n") - + ) + # [URL DISPLAY] Print URL with proper escaping to prevent Rich markup issues. + # IMPORTANT: OAuth URLs contain special characters (=, &, etc.) that Rich might + # interpret as markup in some terminal configurations. We escape the URL to + # ensure it displays correctly. + # + # KNOWN ISSUE: If Rich rendering fails entirely (e.g., terminal doesn't support + # ANSI codes, or output is piped), the escaped URL should still be valid. + # However, if the terminal strips or mangles the output, users should copy + # the URL directly from logs or use --verbose to see the raw URL. + # + # The [link=...] markup creates a clickable hyperlink in supported terminals + # (iTerm2, Windows Terminal, etc.), but the displayed text is the escaped URL + # which can be safely copied even if the hyperlink doesn't work. + escaped_url = rich_escape(auth_url) + console.print( + f"[bold]URL:[/bold] [link={auth_url}]{escaped_url}[/link]\n" + ) + # [HEADLESS SUPPORT] Only attempt browser open if NOT headless if not is_headless: try: webbrowser.open(auth_url) - lib_logger.info("Browser opened successfully for OAuth flow") + lib_logger.info( + "Browser opened successfully for OAuth flow" + ) except Exception as e: - lib_logger.warning(f"Failed to open browser automatically: {e}. Please open the URL manually.") - - with console.status(f"[bold green]Waiting for you to complete authentication in the browser...[/bold green]", spinner="dots"): - auth_code = await asyncio.wait_for(auth_code_future, timeout=300) + lib_logger.warning( + f"Failed to open browser automatically: {e}. Please open the URL manually." + ) + + with console.status( + f"[bold green]Waiting for you to complete authentication in the browser...[/bold green]", + spinner="dots", + ): + auth_code = await asyncio.wait_for( + auth_code_future, timeout=300 + ) except asyncio.TimeoutError: raise Exception("OAuth flow timed out. Please try again.") finally: if server: server.close() await server.wait_closed() - - lib_logger.info(f"Attempting to exchange authorization code for tokens...") + + lib_logger.info( + f"Attempting to exchange authorization code for tokens..." + ) async with httpx.AsyncClient() as client: - response = await client.post(self.TOKEN_URI, data={ - "code": auth_code.strip(), "client_id": self.CLIENT_ID, "client_secret": self.CLIENT_SECRET, - "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", "grant_type": "authorization_code" - }) + response = await client.post( + self.TOKEN_URI, + data={ + "code": auth_code.strip(), + "client_id": self.CLIENT_ID, + "client_secret": self.CLIENT_SECRET, + "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", + "grant_type": "authorization_code", + }, + ) response.raise_for_status() token_data = response.json() # Start with the full token data from the exchange creds = token_data.copy() - + # Convert 'expires_in' to 'expiry_date' in milliseconds - creds["expiry_date"] = (time.time() + creds.pop("expires_in")) * 1000 - + creds["expiry_date"] = ( + time.time() + creds.pop("expires_in") + ) * 1000 + # Ensure client_id and client_secret are present creds["client_id"] = self.CLIENT_ID creds["client_secret"] = self.CLIENT_SECRET creds["token_uri"] = self.TOKEN_URI creds["universe_domain"] = "googleapis.com" - + # Fetch user info and add metadata - user_info_response = await client.get(self.USER_INFO_URI, headers={"Authorization": f"Bearer {creds['access_token']}"}) + user_info_response = await client.get( + self.USER_INFO_URI, + headers={"Authorization": f"Bearer {creds['access_token']}"}, + ) user_info_response.raise_for_status() user_info = user_info_response.json() creds["_proxy_metadata"] = { "email": user_info.get("email"), - "last_check_timestamp": time.time() + "last_check_timestamp": time.time(), } if path: await self._save_credentials(path, creds) - lib_logger.info(f"{self.ENV_PREFIX} OAuth initialized successfully for '{display_name}'.") + lib_logger.info( + f"{self.ENV_PREFIX} OAuth initialized successfully for '{display_name}'." + ) return creds - lib_logger.info(f"{self.ENV_PREFIX} OAuth token at '{display_name}' is valid.") + lib_logger.info( + f"{self.ENV_PREFIX} OAuth token at '{display_name}' is valid." + ) return creds except Exception as e: - raise ValueError(f"Failed to initialize {self.ENV_PREFIX} OAuth for '{path}': {e}") + raise ValueError( + f"Failed to initialize {self.ENV_PREFIX} OAuth for '{path}': {e}" + ) async def get_auth_header(self, credential_path: str) -> Dict[str, str]: creds = await self._load_credentials(credential_path) @@ -675,13 +846,15 @@ async def get_auth_header(self, credential_path: str) -> Dict[str, str]: creds = await self._refresh_token(credential_path, creds) return {"Authorization": f"Bearer {creds['access_token']}"} - async def get_user_info(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + async def get_user_info( + self, creds_or_path: Union[Dict[str, Any], str] + ) -> Dict[str, Any]: path = creds_or_path if isinstance(creds_or_path, str) else None creds = await self._load_credentials(creds_or_path) if path else creds_or_path if path and self._is_token_expired(creds): creds = await self._refresh_token(path, creds) - + # Prefer locally stored metadata if creds.get("_proxy_metadata", {}).get("email"): if path: @@ -695,11 +868,11 @@ async def get_user_info(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict response = await client.get(self.USER_INFO_URI, headers=headers) response.raise_for_status() user_info = response.json() - + # Save the retrieved info for future use creds["_proxy_metadata"] = { "email": user_info.get("email"), - "last_check_timestamp": time.time() + "last_check_timestamp": time.time(), } if path: await self._save_credentials(path, creds) diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index cae85928..021c3100 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -21,9 +21,10 @@ from rich.panel import Panel from rich.prompt import Prompt from rich.text import Text +from rich.markup import escape as rich_escape from ..utils.headless_detection import is_headless_environment -lib_logger = logging.getLogger('rotator_library') +lib_logger = logging.getLogger("rotator_library") IFLOW_OAUTH_AUTHORIZE_ENDPOINT = "https://iflow.cn/oauth" IFLOW_OAUTH_TOKEN_ENDPOINT = "https://iflow.cn/oauth/token" @@ -61,7 +62,7 @@ def _is_port_available(self) -> bool: """Checks if the callback port is available.""" try: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.bind(('', self.port)) + sock.bind(("", self.port)) sock.close() return True except OSError: @@ -76,12 +77,12 @@ async def start(self, expected_state: str): self.result_future = asyncio.Future() # Setup route - self.app.router.add_get('/oauth2callback', self._handle_callback) + self.app.router.add_get("/oauth2callback", self._handle_callback) # Start server self.runner = web.AppRunner(self.app) await self.runner.setup() - self.site = web.TCPSite(self.runner, 'localhost', self.port) + self.site = web.TCPSite(self.runner, "localhost", self.port) await self.site.start() lib_logger.debug(f"iFlow OAuth callback server started on port {self.port}") @@ -99,34 +100,46 @@ async def _handle_callback(self, request: web.Request) -> web.Response: query = request.query # Check for error parameter - if 'error' in query: - error = query.get('error', 'unknown_error') + if "error" in query: + error = query.get("error", "unknown_error") lib_logger.error(f"iFlow OAuth callback received error: {error}") if not self.result_future.done(): self.result_future.set_exception(ValueError(f"OAuth error: {error}")) - return web.Response(status=302, headers={'Location': IFLOW_ERROR_REDIRECT_URL}) + return web.Response( + status=302, headers={"Location": IFLOW_ERROR_REDIRECT_URL} + ) # Check for authorization code - code = query.get('code') + code = query.get("code") if not code: lib_logger.error("iFlow OAuth callback missing authorization code") if not self.result_future.done(): - self.result_future.set_exception(ValueError("Missing authorization code")) - return web.Response(status=302, headers={'Location': IFLOW_ERROR_REDIRECT_URL}) + self.result_future.set_exception( + ValueError("Missing authorization code") + ) + return web.Response( + status=302, headers={"Location": IFLOW_ERROR_REDIRECT_URL} + ) # Validate state parameter - state = query.get('state', '') + state = query.get("state", "") if state != self.expected_state: - lib_logger.error(f"iFlow OAuth state mismatch. Expected: {self.expected_state}, Got: {state}") + lib_logger.error( + f"iFlow OAuth state mismatch. Expected: {self.expected_state}, Got: {state}" + ) if not self.result_future.done(): self.result_future.set_exception(ValueError("State parameter mismatch")) - return web.Response(status=302, headers={'Location': IFLOW_ERROR_REDIRECT_URL}) + return web.Response( + status=302, headers={"Location": IFLOW_ERROR_REDIRECT_URL} + ) # Success - set result and redirect to success page if not self.result_future.done(): self.result_future.set_result(code) - return web.Response(status=302, headers={'Location': IFLOW_SUCCESS_REDIRECT_URL}) + return web.Response( + status=302, headers={"Location": IFLOW_SUCCESS_REDIRECT_URL} + ) async def wait_for_callback(self, timeout: float = 300.0) -> str: """Waits for the OAuth callback and returns the authorization code.""" @@ -146,38 +159,50 @@ class IFlowAuthBase: def __init__(self): self._credentials_cache: Dict[str, Dict[str, Any]] = {} self._refresh_locks: Dict[str, asyncio.Lock] = {} - self._locks_lock = asyncio.Lock() # Protects the locks dict from race conditions + self._locks_lock = ( + asyncio.Lock() + ) # Protects the locks dict from race conditions # [BACKOFF TRACKING] Track consecutive failures per credential - self._refresh_failures: Dict[str, int] = {} # Track consecutive failures per credential - self._next_refresh_after: Dict[str, float] = {} # Track backoff timers (Unix timestamp) - + self._refresh_failures: Dict[ + str, int + ] = {} # Track consecutive failures per credential + self._next_refresh_after: Dict[ + str, float + ] = {} # Track backoff timers (Unix timestamp) + # [QUEUE SYSTEM] Sequential refresh processing self._refresh_queue: asyncio.Queue = asyncio.Queue() self._queued_credentials: set = set() # Track credentials already in queue - self._unavailable_credentials: set = set() # Mark credentials unavailable during re-auth + self._unavailable_credentials: set = ( + set() + ) # Mark credentials unavailable during re-auth self._queue_tracking_lock = asyncio.Lock() # Protects queue sets - self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task + self._queue_processor_task: Optional[asyncio.Task] = ( + None # Background worker task + ) def _parse_env_credential_path(self, path: str) -> Optional[str]: """ Parse a virtual env:// path and return the credential index. - + Supported formats: - "env://provider/0" - Legacy single credential (no index in env var names) - "env://provider/1" - First numbered credential (IFLOW_1_ACCESS_TOKEN) - + Returns: The credential index as string, or None if path is not an env:// path """ if not path.startswith("env://"): return None - + parts = path[6:].split("/") if len(parts) >= 2: return parts[1] return "0" - def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dict[str, Any]]: + def _load_from_env( + self, credential_index: Optional[str] = None + ) -> Optional[Dict[str, Any]]: """ Load OAuth credentials from environment variables for stateless deployments. @@ -204,7 +229,7 @@ def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dic else: prefix = "IFLOW" default_email = "env-user" - + access_token = os.getenv(f"{prefix}_ACCESS_TOKEN") refresh_token = os.getenv(f"{prefix}_REFRESH_TOKEN") api_key = os.getenv(f"{prefix}_API_KEY") @@ -213,7 +238,9 @@ def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dic if not (access_token and refresh_token and api_key): return None - lib_logger.debug(f"Loading iFlow credentials from environment variables (prefix: {prefix})") + lib_logger.debug( + f"Loading iFlow credentials from environment variables (prefix: {prefix})" + ) # Parse expiry_date as string (ISO 8601 format) expiry_str = os.getenv(f"{prefix}_EXPIRY_DATE", "") @@ -230,8 +257,8 @@ def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dic "email": os.getenv(f"{prefix}_EMAIL", default_email), "last_check_timestamp": time.time(), "loaded_from_env": True, - "env_credential_index": credential_index or "0" - } + "env_credential_index": credential_index or "0", + }, } return creds @@ -240,7 +267,7 @@ async def _read_creds_from_file(self, path: str) -> Dict[str, Any]: """Reads credentials from file and populates the cache. No locking.""" try: lib_logger.debug(f"Reading iFlow credentials from file: {path}") - with open(path, 'r') as f: + with open(path, "r") as f: creds = json.load(f) self._credentials_cache[path] = creds return creds @@ -264,11 +291,15 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: if credential_index is not None: env_creds = self._load_from_env(credential_index) if env_creds: - lib_logger.info(f"Using iFlow credentials from environment variables (index: {credential_index})") + lib_logger.info( + f"Using iFlow credentials from environment variables (index: {credential_index})" + ) self._credentials_cache[path] = env_creds return env_creds else: - raise IOError(f"Environment variables for iFlow credential index {credential_index} not found") + raise IOError( + f"Environment variables for iFlow credential index {credential_index} not found" + ) # For file paths, try loading from legacy env vars first env_creds = self._load_from_env() @@ -298,10 +329,12 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): tmp_path = None try: # Create temp file in same directory as target (ensures same filesystem) - tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json', text=True) + tmp_fd, tmp_path = tempfile.mkstemp( + dir=parent_dir, prefix=".tmp_", suffix=".json", text=True + ) # Write JSON to temp file - with os.fdopen(tmp_fd, 'w') as f: + with os.fdopen(tmp_fd, "w") as f: json.dump(creds, f, indent=2) tmp_fd = None # fdopen closes the fd @@ -318,10 +351,14 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): # Update cache AFTER successful file write self._credentials_cache[path] = creds - lib_logger.debug(f"Saved updated iFlow OAuth credentials to '{path}' (atomic write).") + lib_logger.debug( + f"Saved updated iFlow OAuth credentials to '{path}' (atomic write)." + ) except Exception as e: - lib_logger.error(f"Failed to save updated iFlow OAuth credentials to '{path}': {e}") + lib_logger.error( + f"Failed to save updated iFlow OAuth credentials to '{path}': {e}" + ) # Clean up temp file if it still exists if tmp_fd is not None: try: @@ -345,7 +382,8 @@ def _is_token_expired(self, creds: Dict[str, Any]) -> bool: try: # Parse ISO 8601 format (e.g., "2025-01-17T12:00:00Z") from datetime import datetime - expiry_dt = datetime.fromisoformat(expiry_str.replace('Z', '+00:00')) + + expiry_dt = datetime.fromisoformat(expiry_str.replace("Z", "+00:00")) expiry_timestamp = expiry_dt.timestamp() except (ValueError, AttributeError): # Fallback: treat as numeric timestamp @@ -389,7 +427,9 @@ async def _fetch_user_info(self, access_token: str) -> Dict[str, Any]: return {"api_key": api_key, "email": email} - async def _exchange_code_for_tokens(self, code: str, redirect_uri: str) -> Dict[str, Any]: + async def _exchange_code_for_tokens( + self, code: str, redirect_uri: str + ) -> Dict[str, Any]: """ Exchanges authorization code for access and refresh tokens. Uses Basic Auth with client credentials. @@ -401,7 +441,7 @@ async def _exchange_code_for_tokens(self, code: str, redirect_uri: str) -> Dict[ headers = { "Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json", - "Authorization": f"Basic {basic_auth}" + "Authorization": f"Basic {basic_auth}", } data = { @@ -409,16 +449,22 @@ async def _exchange_code_for_tokens(self, code: str, redirect_uri: str) -> Dict[ "code": code, "redirect_uri": redirect_uri, "client_id": IFLOW_CLIENT_ID, - "client_secret": IFLOW_CLIENT_SECRET + "client_secret": IFLOW_CLIENT_SECRET, } async with httpx.AsyncClient(timeout=30.0) as client: - response = await client.post(IFLOW_OAUTH_TOKEN_ENDPOINT, headers=headers, data=data) + response = await client.post( + IFLOW_OAUTH_TOKEN_ENDPOINT, headers=headers, data=data + ) if response.status_code != 200: error_text = response.text - lib_logger.error(f"iFlow token exchange failed: {response.status_code} {error_text}") - raise ValueError(f"Token exchange failed: {response.status_code} {error_text}") + lib_logger.error( + f"iFlow token exchange failed: {response.status_code} {error_text}" + ) + raise ValueError( + f"Token exchange failed: {response.status_code} {error_text}" + ) token_data = response.json() @@ -436,7 +482,10 @@ async def _exchange_code_for_tokens(self, code: str, redirect_uri: str) -> Dict[ # Calculate expiry date from datetime import datetime, timedelta - expiry_date = (datetime.utcnow() + timedelta(seconds=expires_in)).isoformat() + 'Z' + + expiry_date = ( + datetime.utcnow() + timedelta(seconds=expires_in) + ).isoformat() + "Z" return { "access_token": access_token, @@ -445,7 +494,7 @@ async def _exchange_code_for_tokens(self, code: str, redirect_uri: str) -> Dict[ "email": user_info["email"], "expiry_date": expiry_date, "token_type": token_type, - "scope": scope + "scope": scope, } async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any]: @@ -482,20 +531,22 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] headers = { "Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json", - "Authorization": f"Basic {basic_auth}" + "Authorization": f"Basic {basic_auth}", } data = { "grant_type": "refresh_token", "refresh_token": refresh_token, "client_id": IFLOW_CLIENT_ID, - "client_secret": IFLOW_CLIENT_SECRET + "client_secret": IFLOW_CLIENT_SECRET, } async with httpx.AsyncClient(timeout=30.0) as client: for attempt in range(max_retries): try: - response = await client.post(IFLOW_OAUTH_TOKEN_ENDPOINT, headers=headers, data=data) + response = await client.post( + IFLOW_OAUTH_TOKEN_ENDPOINT, headers=headers, data=data + ) response.raise_for_status() new_token_data = response.json() break # Success @@ -505,7 +556,9 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] status_code = e.response.status_code error_body = e.response.text - lib_logger.error(f"[REFRESH HTTP ERROR] HTTP {status_code} for '{Path(path).name}': {error_body}") + lib_logger.error( + f"[REFRESH HTTP ERROR] HTTP {status_code} for '{Path(path).name}': {error_body}" + ) # [STATUS CODE HANDLING] # [INVALID GRANT HANDLING] Handle 401/403 by triggering re-authentication @@ -519,7 +572,9 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] elif status_code == 429: retry_after = int(e.response.headers.get("Retry-After", 60)) - lib_logger.warning(f"Rate limited (HTTP 429), retry after {retry_after}s") + lib_logger.warning( + f"Rate limited (HTTP 429), retry after {retry_after}s" + ) if attempt < max_retries - 1: await asyncio.sleep(retry_after) continue @@ -527,8 +582,10 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] elif 500 <= status_code < 600: if attempt < max_retries - 1: - wait_time = 2 ** attempt - lib_logger.warning(f"Server error (HTTP {status_code}), retry {attempt + 1}/{max_retries} in {wait_time}s") + wait_time = 2**attempt + lib_logger.warning( + f"Server error (HTTP {status_code}), retry {attempt + 1}/{max_retries} in {wait_time}s" + ) await asyncio.sleep(wait_time) continue raise @@ -539,15 +596,19 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] except (httpx.RequestError, httpx.TimeoutException) as e: last_error = e if attempt < max_retries - 1: - wait_time = 2 ** attempt - lib_logger.warning(f"Network error during refresh: {e}, retry {attempt + 1}/{max_retries} in {wait_time}s") + wait_time = 2**attempt + lib_logger.warning( + f"Network error during refresh: {e}, retry {attempt + 1}/{max_retries} in {wait_time}s" + ) await asyncio.sleep(wait_time) continue raise # [INVALID GRANT RE-AUTH] Trigger OAuth flow if refresh token is invalid if needs_reauth: - lib_logger.info(f"Starting re-authentication for '{Path(path).name}'...") + lib_logger.info( + f"Starting re-authentication for '{Path(path).name}'..." + ) try: # Call initialize_token to trigger OAuth flow new_creds = await self.initialize_token(path) @@ -556,20 +617,34 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] self._next_refresh_after.pop(path, None) return new_creds except Exception as reauth_error: - lib_logger.error(f"Re-authentication failed for '{Path(path).name}': {reauth_error}") + lib_logger.error( + f"Re-authentication failed for '{Path(path).name}': {reauth_error}" + ) # [BACKOFF TRACKING] Increment failure count and set backoff timer - self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 - backoff_seconds = min(300, 30 * (2 ** self._refresh_failures[path])) # Max 5 min backoff + self._refresh_failures[path] = ( + self._refresh_failures.get(path, 0) + 1 + ) + backoff_seconds = min( + 300, 30 * (2 ** self._refresh_failures[path]) + ) # Max 5 min backoff self._next_refresh_after[path] = time.time() + backoff_seconds - lib_logger.debug(f"Setting backoff for '{Path(path).name}': {backoff_seconds}s") - raise ValueError(f"Refresh token invalid and re-authentication failed: {reauth_error}") + lib_logger.debug( + f"Setting backoff for '{Path(path).name}': {backoff_seconds}s" + ) + raise ValueError( + f"Refresh token invalid and re-authentication failed: {reauth_error}" + ) if new_token_data is None: # [BACKOFF TRACKING] Increment failure count and set backoff timer self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 - backoff_seconds = min(300, 30 * (2 ** self._refresh_failures[path])) # Max 5 min backoff + backoff_seconds = min( + 300, 30 * (2 ** self._refresh_failures[path]) + ) # Max 5 min backoff self._next_refresh_after[path] = time.time() + backoff_seconds - lib_logger.debug(f"Setting backoff for '{Path(path).name}': {backoff_seconds}s") + lib_logger.debug( + f"Setting backoff for '{Path(path).name}': {backoff_seconds}s" + ) raise last_error or Exception("Token refresh failed after all retries") # Update tokens @@ -578,14 +653,23 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] raise ValueError("Missing access_token in refresh response") creds_from_file["access_token"] = access_token - creds_from_file["refresh_token"] = new_token_data.get("refresh_token", creds_from_file["refresh_token"]) + creds_from_file["refresh_token"] = new_token_data.get( + "refresh_token", creds_from_file["refresh_token"] + ) expires_in = new_token_data.get("expires_in", 3600) from datetime import datetime, timedelta - creds_from_file["expiry_date"] = (datetime.utcnow() + timedelta(seconds=expires_in)).isoformat() + 'Z' - creds_from_file["token_type"] = new_token_data.get("token_type", creds_from_file.get("token_type", "Bearer")) - creds_from_file["scope"] = new_token_data.get("scope", creds_from_file.get("scope", "")) + creds_from_file["expiry_date"] = ( + datetime.utcnow() + timedelta(seconds=expires_in) + ).isoformat() + "Z" + + creds_from_file["token_type"] = new_token_data.get( + "token_type", creds_from_file.get("token_type", "Bearer") + ) + creds_from_file["scope"] = new_token_data.get( + "scope", creds_from_file.get("scope", "") + ) # CRITICAL: Re-fetch user info to get potentially updated API key try: @@ -595,7 +679,9 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] if user_info.get("email"): creds_from_file["email"] = user_info["email"] except Exception as e: - lib_logger.warning(f"Failed to update API key during token refresh: {e}") + lib_logger.warning( + f"Failed to update API key during token refresh: {e}" + ) # Ensure _proxy_metadata exists and update timestamp if "_proxy_metadata" not in creds_from_file: @@ -604,16 +690,22 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] # [VALIDATION] Verify required fields exist after refresh required_fields = ["access_token", "refresh_token", "api_key"] - missing_fields = [field for field in required_fields if not creds_from_file.get(field)] + missing_fields = [ + field for field in required_fields if not creds_from_file.get(field) + ] if missing_fields: - raise ValueError(f"Refreshed credentials missing required fields: {missing_fields}") + raise ValueError( + f"Refreshed credentials missing required fields: {missing_fields}" + ) # [BACKOFF TRACKING] Clear failure count on successful refresh self._refresh_failures.pop(path, None) self._next_refresh_after.pop(path, None) await self._save_credentials(path, creds_from_file) - lib_logger.debug(f"Successfully refreshed iFlow OAuth token for '{Path(path).name}'.") + lib_logger.debug( + f"Successfully refreshed iFlow OAuth token for '{Path(path).name}'." + ) return creds_from_file async def get_api_details(self, credential_identifier: str) -> Tuple[str, str]: @@ -628,7 +720,9 @@ async def get_api_details(self, credential_identifier: str) -> Tuple[str, str]: # Detect credential type if os.path.isfile(credential_identifier): # OAuth credential: file path to JSON - lib_logger.debug(f"Using OAuth credentials from file: {credential_identifier}") + lib_logger.debug( + f"Using OAuth credentials from file: {credential_identifier}" + ) creds = await self._load_credentials(credential_identifier) # Check if token needs refresh @@ -653,7 +747,7 @@ async def proactively_refresh(self, credential_identifier: str): """ # Check if it's an env:// virtual path (OAuth credentials from environment) is_env_path = credential_identifier.startswith("env://") - + # Only refresh if it's an OAuth credential (file path or env:// path) if not is_env_path and not os.path.isfile(credential_identifier): return # Direct API key, no refresh needed @@ -661,7 +755,9 @@ async def proactively_refresh(self, credential_identifier: str): creds = await self._load_credentials(credential_identifier) if self._is_token_expired(creds): # Queue for refresh with needs_reauth=False (automated refresh) - await self._queue_refresh(credential_identifier, force=False, needs_reauth=False) + await self._queue_refresh( + credential_identifier, force=False, needs_reauth=False + ) async def _get_lock(self, path: str) -> asyncio.Lock: """Gets or creates a lock for the given credential path.""" @@ -678,11 +774,15 @@ def is_credential_available(self, path: str) -> bool: async def _ensure_queue_processor_running(self): """Lazily starts the queue processor if not already running.""" if self._queue_processor_task is None or self._queue_processor_task.done(): - self._queue_processor_task = asyncio.create_task(self._process_refresh_queue()) + self._queue_processor_task = asyncio.create_task( + self._process_refresh_queue() + ) - async def _queue_refresh(self, path: str, force: bool = False, needs_reauth: bool = False): + async def _queue_refresh( + self, path: str, force: bool = False, needs_reauth: bool = False + ): """Add a credential to the refresh queue if not already queued. - + Args: path: Credential file path force: Force refresh even if not expired @@ -697,9 +797,11 @@ async def _queue_refresh(self, path: str, force: bool = False, needs_reauth: boo if now < backoff_until: # Credential is in backoff for automated refresh, do not queue remaining = int(backoff_until - now) - lib_logger.debug(f"Skipping automated refresh for '{Path(path).name}' (in backoff for {remaining}s)") + lib_logger.debug( + f"Skipping automated refresh for '{Path(path).name}' (in backoff for {remaining}s)" + ) return - + async with self._queue_tracking_lock: if path not in self._queued_credentials: self._queued_credentials.add(path) @@ -715,14 +817,13 @@ async def _process_refresh_queue(self): # Wait for an item with timeout to allow graceful shutdown try: path, force, needs_reauth = await asyncio.wait_for( - self._refresh_queue.get(), - timeout=60.0 + self._refresh_queue.get(), timeout=60.0 ) except asyncio.TimeoutError: # No items for 60s, exit to save resources self._queue_processor_task = None return - + try: # Perform the actual refresh (still using per-credential lock) async with await self._get_lock(path): @@ -733,16 +834,16 @@ async def _process_refresh_queue(self): async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) continue - + # Perform refresh if not creds: creds = await self._load_credentials(path) await self._refresh_token(path, force=force) - + # SUCCESS: Mark as available again async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) - + finally: # Remove from queued set async with self._queue_tracking_lock: @@ -757,7 +858,9 @@ async def _process_refresh_queue(self): async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) - async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + async def initialize_token( + self, creds_or_path: Union[Dict[str, Any], str] + ) -> Dict[str, Any]: """ Initiates OAuth authorization code flow if tokens are missing or invalid. Uses local callback server to receive authorization code. @@ -766,14 +869,18 @@ async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> D # Get display name from metadata if available, otherwise derive from path if isinstance(creds_or_path, dict): - display_name = creds_or_path.get("_proxy_metadata", {}).get("display_name", "in-memory object") + display_name = creds_or_path.get("_proxy_metadata", {}).get( + "display_name", "in-memory object" + ) else: display_name = Path(path).name if path else "in-memory object" lib_logger.debug(f"Initializing iFlow token for '{display_name}'...") try: - creds = await self._load_credentials(creds_or_path) if path else creds_or_path + creds = ( + await self._load_credentials(creds_or_path) if path else creds_or_path + ) reason = "" if not creds.get("refresh_token"): @@ -787,11 +894,15 @@ async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> D try: return await self._refresh_token(path) except Exception as e: - lib_logger.warning(f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login.") + lib_logger.warning( + f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login." + ) # Interactive OAuth flow - lib_logger.warning(f"iFlow OAuth token for '{display_name}' needs setup: {reason}.") - + lib_logger.warning( + f"iFlow OAuth token for '{display_name}' needs setup: {reason}." + ) + # [HEADLESS DETECTION] Check if running in headless environment is_headless = is_headless_environment() @@ -805,7 +916,7 @@ async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> D "type": "phone", "redirect": redirect_uri, "state": state, - "client_id": IFLOW_CLIENT_ID + "client_id": IFLOW_CLIENT_ID, } auth_url = f"{IFLOW_OAUTH_AUTHORIZE_ENDPOINT}?{urlencode(auth_params)}" @@ -829,49 +940,86 @@ async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> D "2. [bold]Authorize the application[/bold] to access your account.\n" "3. You will be automatically redirected after authorization." ) - - console.print(Panel(auth_panel_text, title=f"iFlow OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue")) - console.print(f"[bold]URL:[/bold] [link={auth_url}]{auth_url}[/link]\n") + + console.print( + Panel( + auth_panel_text, + title=f"iFlow OAuth Setup for [bold yellow]{display_name}[/bold yellow]", + style="bold blue", + ) + ) + # [URL DISPLAY] Print URL with proper escaping to prevent Rich markup issues. + # IMPORTANT: OAuth URLs contain special characters (=, &, etc.) that Rich might + # interpret as markup in some terminal configurations. We escape the URL to + # ensure it displays correctly. + # + # KNOWN ISSUE: If Rich rendering fails entirely (e.g., terminal doesn't support + # ANSI codes, or output is piped), the escaped URL should still be valid. + # However, if the terminal strips or mangles the output, users should copy + # the URL directly from logs or use --verbose to see the raw URL. + # + # The [link=...] markup creates a clickable hyperlink in supported terminals + # (iTerm2, Windows Terminal, etc.), but the displayed text is the escaped URL + # which can be safely copied even if the hyperlink doesn't work. + escaped_url = rich_escape(auth_url) + console.print( + f"[bold]URL:[/bold] [link={auth_url}]{escaped_url}[/link]\n" + ) # [HEADLESS SUPPORT] Only attempt browser open if NOT headless if not is_headless: try: webbrowser.open(auth_url) - lib_logger.info("Browser opened successfully for iFlow OAuth flow") + lib_logger.info( + "Browser opened successfully for iFlow OAuth flow" + ) except Exception as e: - lib_logger.warning(f"Failed to open browser automatically: {e}. Please open the URL manually.") + lib_logger.warning( + f"Failed to open browser automatically: {e}. Please open the URL manually." + ) # Wait for callback - with console.status("[bold green]Waiting for authorization in the browser...[/bold green]", spinner="dots"): + with console.status( + "[bold green]Waiting for authorization in the browser...[/bold green]", + spinner="dots", + ): code = await callback_server.wait_for_callback(timeout=300.0) - lib_logger.info("Received authorization code, exchanging for tokens...") + lib_logger.info( + "Received authorization code, exchanging for tokens..." + ) # Exchange code for tokens and API key - token_data = await self._exchange_code_for_tokens(code, redirect_uri) + token_data = await self._exchange_code_for_tokens( + code, redirect_uri + ) # Update credentials - creds.update({ - "access_token": token_data["access_token"], - "refresh_token": token_data["refresh_token"], - "api_key": token_data["api_key"], - "email": token_data["email"], - "expiry_date": token_data["expiry_date"], - "token_type": token_data["token_type"], - "scope": token_data["scope"] - }) + creds.update( + { + "access_token": token_data["access_token"], + "refresh_token": token_data["refresh_token"], + "api_key": token_data["api_key"], + "email": token_data["email"], + "expiry_date": token_data["expiry_date"], + "token_type": token_data["token_type"], + "scope": token_data["scope"], + } + ) # Create metadata object if not creds.get("_proxy_metadata"): creds["_proxy_metadata"] = { "email": token_data["email"], - "last_check_timestamp": time.time() + "last_check_timestamp": time.time(), } if path: await self._save_credentials(path, creds) - lib_logger.info(f"iFlow OAuth initialized successfully for '{display_name}'.") + lib_logger.info( + f"iFlow OAuth initialized successfully for '{display_name}'." + ) return creds finally: @@ -898,11 +1046,15 @@ async def get_auth_header(self, credential_path: str) -> Dict[str, str]: return {"Authorization": f"Bearer {api_key}"} - async def get_user_info(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + async def get_user_info( + self, creds_or_path: Union[Dict[str, Any], str] + ) -> Dict[str, Any]: """Retrieves user info from the _proxy_metadata in the credential file.""" try: path = creds_or_path if isinstance(creds_or_path, str) else None - creds = await self._load_credentials(creds_or_path) if path else creds_or_path + creds = ( + await self._load_credentials(creds_or_path) if path else creds_or_path + ) # Ensure the token is valid if path: @@ -912,7 +1064,9 @@ async def get_user_info(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict email = creds.get("email") or creds.get("_proxy_metadata", {}).get("email") if not email: - lib_logger.warning(f"No email found in iFlow credentials for '{path or 'in-memory object'}'.") + lib_logger.warning( + f"No email found in iFlow credentials for '{path or 'in-memory object'}'." + ) # Update timestamp on check if path and "_proxy_metadata" in creds: diff --git a/src/rotator_library/providers/qwen_auth_base.py b/src/rotator_library/providers/qwen_auth_base.py index 589e6bef..66e1d685 100644 --- a/src/rotator_library/providers/qwen_auth_base.py +++ b/src/rotator_library/providers/qwen_auth_base.py @@ -19,54 +19,70 @@ from rich.panel import Panel from rich.prompt import Prompt from rich.text import Text +from rich.markup import escape as rich_escape from ..utils.headless_detection import is_headless_environment -lib_logger = logging.getLogger('rotator_library') +lib_logger = logging.getLogger("rotator_library") -CLIENT_ID = "f0304373b74a44d2b584a3fb70ca9e56" #https://api.kilocode.ai/extension-config.json +CLIENT_ID = ( + "f0304373b74a44d2b584a3fb70ca9e56" # https://api.kilocode.ai/extension-config.json +) SCOPE = "openid profile email model.completion" TOKEN_ENDPOINT = "https://chat.qwen.ai/api/v1/oauth2/token" REFRESH_EXPIRY_BUFFER_SECONDS = 3 * 60 * 60 # 3 hours buffer before expiry console = Console() + class QwenAuthBase: def __init__(self): self._credentials_cache: Dict[str, Dict[str, Any]] = {} self._refresh_locks: Dict[str, asyncio.Lock] = {} - self._locks_lock = asyncio.Lock() # Protects the locks dict from race conditions + self._locks_lock = ( + asyncio.Lock() + ) # Protects the locks dict from race conditions # [BACKOFF TRACKING] Track consecutive failures per credential - self._refresh_failures: Dict[str, int] = {} # Track consecutive failures per credential - self._next_refresh_after: Dict[str, float] = {} # Track backoff timers (Unix timestamp) - + self._refresh_failures: Dict[ + str, int + ] = {} # Track consecutive failures per credential + self._next_refresh_after: Dict[ + str, float + ] = {} # Track backoff timers (Unix timestamp) + # [QUEUE SYSTEM] Sequential refresh processing self._refresh_queue: asyncio.Queue = asyncio.Queue() self._queued_credentials: set = set() # Track credentials already in queue - self._unavailable_credentials: set = set() # Mark credentials unavailable during re-auth + self._unavailable_credentials: set = ( + set() + ) # Mark credentials unavailable during re-auth self._queue_tracking_lock = asyncio.Lock() # Protects queue sets - self._queue_processor_task: Optional[asyncio.Task] = None # Background worker task + self._queue_processor_task: Optional[asyncio.Task] = ( + None # Background worker task + ) def _parse_env_credential_path(self, path: str) -> Optional[str]: """ Parse a virtual env:// path and return the credential index. - + Supported formats: - "env://provider/0" - Legacy single credential (no index in env var names) - "env://provider/1" - First numbered credential (QWEN_CODE_1_ACCESS_TOKEN) - + Returns: The credential index as string, or None if path is not an env:// path """ if not path.startswith("env://"): return None - + parts = path[6:].split("/") if len(parts) >= 2: return parts[1] return "0" - def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dict[str, Any]]: + def _load_from_env( + self, credential_index: Optional[str] = None + ) -> Optional[Dict[str, Any]]: """ Load OAuth credentials from environment variables for stateless deployments. @@ -91,7 +107,7 @@ def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dic else: prefix = "QWEN_CODE" default_email = "env-user" - + access_token = os.getenv(f"{prefix}_ACCESS_TOKEN") refresh_token = os.getenv(f"{prefix}_REFRESH_TOKEN") @@ -99,27 +115,33 @@ def _load_from_env(self, credential_index: Optional[str] = None) -> Optional[Dic if not (access_token and refresh_token): return None - lib_logger.debug(f"Loading Qwen Code credentials from environment variables (prefix: {prefix})") + lib_logger.debug( + f"Loading Qwen Code credentials from environment variables (prefix: {prefix})" + ) # Parse expiry_date as float, default to 0 if not present expiry_str = os.getenv(f"{prefix}_EXPIRY_DATE", "0") try: expiry_date = float(expiry_str) except ValueError: - lib_logger.warning(f"Invalid {prefix}_EXPIRY_DATE value: {expiry_str}, using 0") + lib_logger.warning( + f"Invalid {prefix}_EXPIRY_DATE value: {expiry_str}, using 0" + ) expiry_date = 0 creds = { "access_token": access_token, "refresh_token": refresh_token, "expiry_date": expiry_date, - "resource_url": os.getenv(f"{prefix}_RESOURCE_URL", "https://portal.qwen.ai/v1"), + "resource_url": os.getenv( + f"{prefix}_RESOURCE_URL", "https://portal.qwen.ai/v1" + ), "_proxy_metadata": { "email": os.getenv(f"{prefix}_EMAIL", default_email), "last_check_timestamp": time.time(), "loaded_from_env": True, - "env_credential_index": credential_index or "0" - } + "env_credential_index": credential_index or "0", + }, } return creds @@ -128,7 +150,7 @@ async def _read_creds_from_file(self, path: str) -> Dict[str, Any]: """Reads credentials from file and populates the cache. No locking.""" try: lib_logger.debug(f"Reading Qwen credentials from file: {path}") - with open(path, 'r') as f: + with open(path, "r") as f: creds = json.load(f) self._credentials_cache[path] = creds return creds @@ -152,16 +174,22 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: if credential_index is not None: env_creds = self._load_from_env(credential_index) if env_creds: - lib_logger.info(f"Using Qwen Code credentials from environment variables (index: {credential_index})") + lib_logger.info( + f"Using Qwen Code credentials from environment variables (index: {credential_index})" + ) self._credentials_cache[path] = env_creds return env_creds else: - raise IOError(f"Environment variables for Qwen Code credential index {credential_index} not found") + raise IOError( + f"Environment variables for Qwen Code credential index {credential_index} not found" + ) # For file paths, try loading from legacy env vars first env_creds = self._load_from_env() if env_creds: - lib_logger.info("Using Qwen Code credentials from environment variables") + lib_logger.info( + "Using Qwen Code credentials from environment variables" + ) self._credentials_cache[path] = env_creds return env_creds @@ -184,10 +212,12 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): tmp_path = None try: # Create temp file in same directory as target (ensures same filesystem) - tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json', text=True) + tmp_fd, tmp_path = tempfile.mkstemp( + dir=parent_dir, prefix=".tmp_", suffix=".json", text=True + ) # Write JSON to temp file - with os.fdopen(tmp_fd, 'w') as f: + with os.fdopen(tmp_fd, "w") as f: json.dump(creds, f, indent=2) tmp_fd = None # fdopen closes the fd @@ -204,10 +234,14 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): # Update cache AFTER successful file write self._credentials_cache[path] = creds - lib_logger.debug(f"Saved updated Qwen OAuth credentials to '{path}' (atomic write).") + lib_logger.debug( + f"Saved updated Qwen OAuth credentials to '{path}' (atomic write)." + ) except Exception as e: - lib_logger.error(f"Failed to save updated Qwen OAuth credentials to '{path}': {e}") + lib_logger.error( + f"Failed to save updated Qwen OAuth credentials to '{path}': {e}" + ) # Clean up temp file if it still exists if tmp_fd is not None: try: @@ -252,17 +286,22 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] headers = { "Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json", - "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36" + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", } async with httpx.AsyncClient() as client: for attempt in range(max_retries): try: - response = await client.post(TOKEN_ENDPOINT, headers=headers, data={ - "grant_type": "refresh_token", - "refresh_token": refresh_token, - "client_id": CLIENT_ID, - }, timeout=30.0) + response = await client.post( + TOKEN_ENDPOINT, + headers=headers, + data={ + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": CLIENT_ID, + }, + timeout=30.0, + ) response.raise_for_status() new_token_data = response.json() break # Success @@ -271,7 +310,9 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] last_error = e status_code = e.response.status_code error_body = e.response.text - lib_logger.error(f"HTTP {status_code} for '{Path(path).name}': {error_body}") + lib_logger.error( + f"HTTP {status_code} for '{Path(path).name}': {error_body}" + ) # [INVALID GRANT HANDLING] Handle 401/403 by triggering re-authentication if status_code in (401, 403): @@ -284,7 +325,9 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] elif status_code == 429: retry_after = int(e.response.headers.get("Retry-After", 60)) - lib_logger.warning(f"Rate limited (HTTP 429), retry after {retry_after}s") + lib_logger.warning( + f"Rate limited (HTTP 429), retry after {retry_after}s" + ) if attempt < max_retries - 1: await asyncio.sleep(retry_after) continue @@ -292,8 +335,10 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] elif 500 <= status_code < 600: if attempt < max_retries - 1: - wait_time = 2 ** attempt - lib_logger.warning(f"Server error (HTTP {status_code}), retry {attempt + 1}/{max_retries} in {wait_time}s") + wait_time = 2**attempt + lib_logger.warning( + f"Server error (HTTP {status_code}), retry {attempt + 1}/{max_retries} in {wait_time}s" + ) await asyncio.sleep(wait_time) continue raise @@ -304,15 +349,19 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] except (httpx.RequestError, httpx.TimeoutException) as e: last_error = e if attempt < max_retries - 1: - wait_time = 2 ** attempt - lib_logger.warning(f"Network error during refresh: {e}, retry {attempt + 1}/{max_retries} in {wait_time}s") + wait_time = 2**attempt + lib_logger.warning( + f"Network error during refresh: {e}, retry {attempt + 1}/{max_retries} in {wait_time}s" + ) await asyncio.sleep(wait_time) continue raise # [INVALID GRANT RE-AUTH] Trigger OAuth flow if refresh token is invalid if needs_reauth: - lib_logger.info(f"Starting re-authentication for '{Path(path).name}'...") + lib_logger.info( + f"Starting re-authentication for '{Path(path).name}'..." + ) try: # Call initialize_token to trigger OAuth flow new_creds = await self.initialize_token(path) @@ -321,26 +370,46 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] self._next_refresh_after.pop(path, None) return new_creds except Exception as reauth_error: - lib_logger.error(f"Re-authentication failed for '{Path(path).name}': {reauth_error}") + lib_logger.error( + f"Re-authentication failed for '{Path(path).name}': {reauth_error}" + ) # [BACKOFF TRACKING] Increment failure count and set backoff timer - self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 - backoff_seconds = min(300, 30 * (2 ** self._refresh_failures[path])) # Max 5 min backoff + self._refresh_failures[path] = ( + self._refresh_failures.get(path, 0) + 1 + ) + backoff_seconds = min( + 300, 30 * (2 ** self._refresh_failures[path]) + ) # Max 5 min backoff self._next_refresh_after[path] = time.time() + backoff_seconds - lib_logger.debug(f"Setting backoff for '{Path(path).name}': {backoff_seconds}s") - raise ValueError(f"Refresh token invalid and re-authentication failed: {reauth_error}") + lib_logger.debug( + f"Setting backoff for '{Path(path).name}': {backoff_seconds}s" + ) + raise ValueError( + f"Refresh token invalid and re-authentication failed: {reauth_error}" + ) if new_token_data is None: # [BACKOFF TRACKING] Increment failure count and set backoff timer self._refresh_failures[path] = self._refresh_failures.get(path, 0) + 1 - backoff_seconds = min(300, 30 * (2 ** self._refresh_failures[path])) # Max 5 min backoff + backoff_seconds = min( + 300, 30 * (2 ** self._refresh_failures[path]) + ) # Max 5 min backoff self._next_refresh_after[path] = time.time() + backoff_seconds - lib_logger.debug(f"Setting backoff for '{Path(path).name}': {backoff_seconds}s") + lib_logger.debug( + f"Setting backoff for '{Path(path).name}': {backoff_seconds}s" + ) raise last_error or Exception("Token refresh failed after all retries") creds_from_file["access_token"] = new_token_data["access_token"] - creds_from_file["refresh_token"] = new_token_data.get("refresh_token", creds_from_file["refresh_token"]) - creds_from_file["expiry_date"] = (time.time() + new_token_data["expires_in"]) * 1000 - creds_from_file["resource_url"] = new_token_data.get("resource_url", creds_from_file.get("resource_url")) + creds_from_file["refresh_token"] = new_token_data.get( + "refresh_token", creds_from_file["refresh_token"] + ) + creds_from_file["expiry_date"] = ( + time.time() + new_token_data["expires_in"] + ) * 1000 + creds_from_file["resource_url"] = new_token_data.get( + "resource_url", creds_from_file.get("resource_url") + ) # Ensure _proxy_metadata exists and update timestamp if "_proxy_metadata" not in creds_from_file: @@ -349,16 +418,22 @@ async def _refresh_token(self, path: str, force: bool = False) -> Dict[str, Any] # [VALIDATION] Verify required fields exist after refresh required_fields = ["access_token", "refresh_token"] - missing_fields = [field for field in required_fields if not creds_from_file.get(field)] + missing_fields = [ + field for field in required_fields if not creds_from_file.get(field) + ] if missing_fields: - raise ValueError(f"Refreshed credentials missing required fields: {missing_fields}") + raise ValueError( + f"Refreshed credentials missing required fields: {missing_fields}" + ) # [BACKOFF TRACKING] Clear failure count on successful refresh self._refresh_failures.pop(path, None) self._next_refresh_after.pop(path, None) await self._save_credentials(path, creds_from_file) - lib_logger.debug(f"Successfully refreshed Qwen OAuth token for '{Path(path).name}'.") + lib_logger.debug( + f"Successfully refreshed Qwen OAuth token for '{Path(path).name}'." + ) return creds_from_file async def get_api_details(self, credential_identifier: str) -> Tuple[str, str]: @@ -372,12 +447,14 @@ async def get_api_details(self, credential_identifier: str) -> Tuple[str, str]: # Detect credential type if os.path.isfile(credential_identifier): # OAuth credential: file path to JSON - lib_logger.debug(f"Using OAuth credentials from file: {credential_identifier}") + lib_logger.debug( + f"Using OAuth credentials from file: {credential_identifier}" + ) creds = await self._load_credentials(credential_identifier) if self._is_token_expired(creds): creds = await self._refresh_token(credential_identifier) - + base_url = creds.get("resource_url", "https://portal.qwen.ai/v1") if not base_url.startswith("http"): base_url = f"https://{base_url}" @@ -397,7 +474,7 @@ async def proactively_refresh(self, credential_identifier: str): """ # Check if it's an env:// virtual path (OAuth credentials from environment) is_env_path = credential_identifier.startswith("env://") - + # Only refresh if it's an OAuth credential (file path or env:// path) if not is_env_path and not os.path.isfile(credential_identifier): return # Direct API key, no refresh needed @@ -405,7 +482,9 @@ async def proactively_refresh(self, credential_identifier: str): creds = await self._load_credentials(credential_identifier) if self._is_token_expired(creds): # Queue for refresh with needs_reauth=False (automated refresh) - await self._queue_refresh(credential_identifier, force=False, needs_reauth=False) + await self._queue_refresh( + credential_identifier, force=False, needs_reauth=False + ) async def _get_lock(self, path: str) -> asyncio.Lock: # [FIX RACE CONDITION] Protect lock creation with a master lock @@ -421,11 +500,15 @@ def is_credential_available(self, path: str) -> bool: async def _ensure_queue_processor_running(self): """Lazily starts the queue processor if not already running.""" if self._queue_processor_task is None or self._queue_processor_task.done(): - self._queue_processor_task = asyncio.create_task(self._process_refresh_queue()) + self._queue_processor_task = asyncio.create_task( + self._process_refresh_queue() + ) - async def _queue_refresh(self, path: str, force: bool = False, needs_reauth: bool = False): + async def _queue_refresh( + self, path: str, force: bool = False, needs_reauth: bool = False + ): """Add a credential to the refresh queue if not already queued. - + Args: path: Credential file path force: Force refresh even if not expired @@ -440,9 +523,11 @@ async def _queue_refresh(self, path: str, force: bool = False, needs_reauth: boo if now < backoff_until: # Credential is in backoff for automated refresh, do not queue remaining = int(backoff_until - now) - lib_logger.debug(f"Skipping automated refresh for '{Path(path).name}' (in backoff for {remaining}s)") + lib_logger.debug( + f"Skipping automated refresh for '{Path(path).name}' (in backoff for {remaining}s)" + ) return - + async with self._queue_tracking_lock: if path not in self._queued_credentials: self._queued_credentials.add(path) @@ -458,14 +543,13 @@ async def _process_refresh_queue(self): # Wait for an item with timeout to allow graceful shutdown try: path, force, needs_reauth = await asyncio.wait_for( - self._refresh_queue.get(), - timeout=60.0 + self._refresh_queue.get(), timeout=60.0 ) except asyncio.TimeoutError: # No items for 60s, exit to save resources self._queue_processor_task = None return - + try: # Perform the actual refresh (still using per-credential lock) async with await self._get_lock(path): @@ -476,16 +560,16 @@ async def _process_refresh_queue(self): async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) continue - + # Perform refresh if not creds: creds = await self._load_credentials(path) await self._refresh_token(path, force=force) - + # SUCCESS: Mark as available again async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) - + finally: # Remove from queued set async with self._queue_tracking_lock: @@ -500,19 +584,25 @@ async def _process_refresh_queue(self): async with self._queue_tracking_lock: self._unavailable_credentials.discard(path) - async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + async def initialize_token( + self, creds_or_path: Union[Dict[str, Any], str] + ) -> Dict[str, Any]: """Initiates device flow if tokens are missing or invalid.""" path = creds_or_path if isinstance(creds_or_path, str) else None # Get display name from metadata if available, otherwise derive from path if isinstance(creds_or_path, dict): - display_name = creds_or_path.get("_proxy_metadata", {}).get("display_name", "in-memory object") + display_name = creds_or_path.get("_proxy_metadata", {}).get( + "display_name", "in-memory object" + ) else: display_name = Path(path).name if path else "in-memory object" lib_logger.debug(f"Initializing Qwen token for '{display_name}'...") try: - creds = await self._load_credentials(creds_or_path) if path else creds_or_path + creds = ( + await self._load_credentials(creds_or_path) if path else creds_or_path + ) reason = "" if not creds.get("refresh_token"): @@ -525,44 +615,58 @@ async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> D try: return await self._refresh_token(path) except Exception as e: - lib_logger.warning(f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login.") + lib_logger.warning( + f"Automatic token refresh for '{display_name}' failed: {e}. Proceeding to interactive login." + ) + + lib_logger.warning( + f"Qwen OAuth token for '{display_name}' needs setup: {reason}." + ) - lib_logger.warning(f"Qwen OAuth token for '{display_name}' needs setup: {reason}.") - # [HEADLESS DETECTION] Check if running in headless environment is_headless = is_headless_environment() - - code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode('utf-8').rstrip('=') - code_challenge = base64.urlsafe_b64encode( - hashlib.sha256(code_verifier.encode('utf-8')).digest() - ).decode('utf-8').rstrip('=') - + + code_verifier = ( + base64.urlsafe_b64encode(secrets.token_bytes(32)) + .decode("utf-8") + .rstrip("=") + ) + code_challenge = ( + base64.urlsafe_b64encode( + hashlib.sha256(code_verifier.encode("utf-8")).digest() + ) + .decode("utf-8") + .rstrip("=") + ) + headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", "Content-Type": "application/x-www-form-urlencoded", - "Accept": "application/json" + "Accept": "application/json", } async with httpx.AsyncClient() as client: request_data = { "client_id": CLIENT_ID, "scope": SCOPE, "code_challenge": code_challenge, - "code_challenge_method": "S256" + "code_challenge_method": "S256", } lib_logger.debug(f"Qwen device code request data: {request_data}") try: dev_response = await client.post( "https://chat.qwen.ai/api/v1/oauth2/device/code", headers=headers, - data=request_data + data=request_data, ) dev_response.raise_for_status() dev_data = dev_response.json() lib_logger.debug(f"Qwen device auth response: {dev_data}") except httpx.HTTPStatusError as e: - lib_logger.error(f"Qwen device code request failed with status {e.response.status_code}: {e.response.text}") + lib_logger.error( + f"Qwen device code request failed with status {e.response.status_code}: {e.response.text}" + ) raise e - + # [HEADLESS SUPPORT] Display appropriate instructions if is_headless: auth_panel_text = Text.from_markup( @@ -578,33 +682,63 @@ async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> D "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\n" "3. You will be prompted to enter your identifier after authorization." ) - - console.print(Panel(auth_panel_text, title=f"Qwen OAuth Setup for [bold yellow]{display_name}[/bold yellow]", style="bold blue")) - console.print(f"[bold]URL:[/bold] [link={dev_data['verification_uri_complete']}]{dev_data['verification_uri_complete']}[/link]\n") - + + console.print( + Panel( + auth_panel_text, + title=f"Qwen OAuth Setup for [bold yellow]{display_name}[/bold yellow]", + style="bold blue", + ) + ) + # [URL DISPLAY] Print URL with proper escaping to prevent Rich markup issues. + # IMPORTANT: OAuth URLs contain special characters (=, &, etc.) that Rich might + # interpret as markup in some terminal configurations. We escape the URL to + # ensure it displays correctly. + # + # KNOWN ISSUE: If Rich rendering fails entirely (e.g., terminal doesn't support + # ANSI codes, or output is piped), the escaped URL should still be valid. + # However, if the terminal strips or mangles the output, users should copy + # the URL directly from logs or use --verbose to see the raw URL. + # + # The [link=...] markup creates a clickable hyperlink in supported terminals + # (iTerm2, Windows Terminal, etc.), but the displayed text is the escaped URL + # which can be safely copied even if the hyperlink doesn't work. + verification_url = dev_data["verification_uri_complete"] + escaped_url = rich_escape(verification_url) + console.print( + f"[bold]URL:[/bold] [link={verification_url}]{escaped_url}[/link]\n" + ) + # [HEADLESS SUPPORT] Only attempt browser open if NOT headless if not is_headless: try: - webbrowser.open(dev_data['verification_uri_complete']) - lib_logger.info("Browser opened successfully for Qwen OAuth flow") + webbrowser.open(dev_data["verification_uri_complete"]) + lib_logger.info( + "Browser opened successfully for Qwen OAuth flow" + ) except Exception as e: - lib_logger.warning(f"Failed to open browser automatically: {e}. Please open the URL manually.") - + lib_logger.warning( + f"Failed to open browser automatically: {e}. Please open the URL manually." + ) + token_data = None start_time = time.time() - interval = dev_data.get('interval', 5) + interval = dev_data.get("interval", 5) - with console.status("[bold green]Polling for token, please complete authentication in the browser...[/bold green]", spinner="dots") as status: - while time.time() - start_time < dev_data['expires_in']: + with console.status( + "[bold green]Polling for token, please complete authentication in the browser...[/bold green]", + spinner="dots", + ) as status: + while time.time() - start_time < dev_data["expires_in"]: poll_response = await client.post( TOKEN_ENDPOINT, headers=headers, data={ "grant_type": "urn:ietf:params:oauth:grant-type:device_code", - "device_code": dev_data['device_code'], + "device_code": dev_data["device_code"], "client_id": CLIENT_ID, - "code_verifier": code_verifier - } + "code_verifier": code_verifier, + }, ) if poll_response.status_code == 200: token_data = poll_response.json() @@ -614,45 +748,63 @@ async def initialize_token(self, creds_or_path: Union[Dict[str, Any], str]) -> D poll_data = poll_response.json() error_type = poll_data.get("error") if error_type == "authorization_pending": - lib_logger.debug(f"Polling status: {error_type}, waiting {interval}s") + lib_logger.debug( + f"Polling status: {error_type}, waiting {interval}s" + ) elif error_type == "slow_down": interval = int(interval * 1.5) if interval > 10: interval = 10 - lib_logger.debug(f"Polling status: {error_type}, waiting {interval}s") + lib_logger.debug( + f"Polling status: {error_type}, waiting {interval}s" + ) else: - raise ValueError(f"Token polling failed: {poll_data.get('error_description', error_type)}") + raise ValueError( + f"Token polling failed: {poll_data.get('error_description', error_type)}" + ) else: poll_response.raise_for_status() - + await asyncio.sleep(interval) - + if not token_data: raise TimeoutError("Qwen device flow timed out.") - - creds.update({ - "access_token": token_data["access_token"], - "refresh_token": token_data.get("refresh_token"), - "expiry_date": (time.time() + token_data["expires_in"]) * 1000, - "resource_url": token_data.get("resource_url") - }) + + creds.update( + { + "access_token": token_data["access_token"], + "refresh_token": token_data.get("refresh_token"), + "expiry_date": (time.time() + token_data["expires_in"]) + * 1000, + "resource_url": token_data.get("resource_url"), + } + ) # Prompt for user identifier and create metadata object if needed if not creds.get("_proxy_metadata", {}).get("email"): try: - prompt_text = Text.from_markup(f"\\n[bold]Please enter your email or a unique identifier for [yellow]'{display_name}'[/yellow][/bold]") + prompt_text = Text.from_markup( + f"\\n[bold]Please enter your email or a unique identifier for [yellow]'{display_name}'[/yellow][/bold]" + ) email = Prompt.ask(prompt_text) creds["_proxy_metadata"] = { "email": email.strip(), - "last_check_timestamp": time.time() + "last_check_timestamp": time.time(), } except (EOFError, KeyboardInterrupt): - console.print("\\n[bold yellow]No identifier provided. Deduplication will not be possible.[/bold yellow]") - creds["_proxy_metadata"] = {"email": None, "last_check_timestamp": time.time()} + console.print( + "\\n[bold yellow]No identifier provided. Deduplication will not be possible.[/bold yellow]" + ) + creds["_proxy_metadata"] = { + "email": None, + "last_check_timestamp": time.time(), + } if path: await self._save_credentials(path, creds) - lib_logger.info(f"Qwen OAuth initialized successfully for '{display_name}'.") + lib_logger.info( + f"Qwen OAuth initialized successfully for '{display_name}'." + ) return creds lib_logger.info(f"Qwen OAuth token at '{display_name}' is valid.") @@ -666,24 +818,32 @@ async def get_auth_header(self, credential_path: str) -> Dict[str, str]: creds = await self._refresh_token(credential_path) return {"Authorization": f"Bearer {creds['access_token']}"} - async def get_user_info(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict[str, Any]: + async def get_user_info( + self, creds_or_path: Union[Dict[str, Any], str] + ) -> Dict[str, Any]: """ Retrieves user info from the _proxy_metadata in the credential file. """ try: path = creds_or_path if isinstance(creds_or_path, str) else None - creds = await self._load_credentials(creds_or_path) if path else creds_or_path - + creds = ( + await self._load_credentials(creds_or_path) if path else creds_or_path + ) + # This will ensure the token is valid and metadata exists if the flow was just run if path: await self.initialize_token(path) - creds = await self._load_credentials(path) # Re-load after potential init + creds = await self._load_credentials( + path + ) # Re-load after potential init metadata = creds.get("_proxy_metadata", {"email": None}) email = metadata.get("email") if not email: - lib_logger.warning(f"No email found in _proxy_metadata for '{path or 'in-memory object'}'.") + lib_logger.warning( + f"No email found in _proxy_metadata for '{path or 'in-memory object'}'." + ) # Update timestamp on check and save if it's a file-based credential if path and "_proxy_metadata" in creds: @@ -693,4 +853,4 @@ async def get_user_info(self, creds_or_path: Union[Dict[str, Any], str]) -> Dict return {"email": email} except Exception as e: lib_logger.error(f"Failed to get Qwen user info from credentials: {e}") - return {"email": None} \ No newline at end of file + return {"email": None} diff --git a/src/rotator_library/utils/headless_detection.py b/src/rotator_library/utils/headless_detection.py index ace75fb1..3fc5d274 100644 --- a/src/rotator_library/utils/headless_detection.py +++ b/src/rotator_library/utils/headless_detection.py @@ -1,24 +1,27 @@ # src/rotator_library/utils/headless_detection.py import os +import sys import logging -lib_logger = logging.getLogger('rotator_library') +lib_logger = logging.getLogger("rotator_library") # Import console for user-visible output try: from rich.console import Console + console = Console() except ImportError: console = None + def is_headless_environment() -> bool: """ Detects if the current environment is headless (no GUI available). - + Returns: True if headless environment is detected, False otherwise - + Detection logic: - Linux/Unix: Check DISPLAY environment variable - SSH detection: Check SSH_CONNECTION or SSH_CLIENT @@ -26,17 +29,20 @@ def is_headless_environment() -> bool: - Windows: Check SESSIONNAME for service/headless indicators """ headless_indicators = [] - - # Check DISPLAY for Linux/Unix GUI availability (skip on Windows) - if os.name != 'nt': # Only check DISPLAY on non-Windows systems + + # Check DISPLAY for Linux GUI availability (skip on Windows and macOS) + # NOTE: DISPLAY is an X11 (X Window System) variable used on Linux. + # macOS uses its native Quartz windowing system, NOT X11, so DISPLAY is + # typically unset on macOS even with a full GUI. Only check DISPLAY on Linux. + if os.name != "nt" and sys.platform != "darwin": # Linux only display = os.getenv("DISPLAY") if display is None or display.strip() == "": - headless_indicators.append("No DISPLAY variable (Linux/Unix headless)") - + headless_indicators.append("No DISPLAY variable (Linux headless)") + # Check for SSH connection if os.getenv("SSH_CONNECTION") or os.getenv("SSH_CLIENT") or os.getenv("SSH_TTY"): headless_indicators.append("SSH connection detected") - + # Check for CI environments ci_vars = [ "CI", # Generic CI indicator @@ -55,30 +61,38 @@ def is_headless_environment() -> bool: if os.getenv(var): headless_indicators.append(f"CI environment detected ({var})") break - + # Check Windows session type - if os.name == 'nt': # Windows + if os.name == "nt": # Windows session_name = os.getenv("SESSIONNAME", "").lower() if session_name in ["services", "rdp-tcp"]: headless_indicators.append(f"Windows headless session ({session_name})") - + # Detect Docker/container environment if os.path.exists("/.dockerenv") or os.path.exists("/run/.containerenv"): headless_indicators.append("Container environment detected") - + # Determine if headless is_headless = len(headless_indicators) > 0 - + if is_headless: # Log to logger - lib_logger.info(f"Headless environment detected: {'; '.join(headless_indicators)}") - + lib_logger.info( + f"Headless environment detected: {'; '.join(headless_indicators)}" + ) + # Print to console for user visibility if console: - console.print(f"[yellow]ℹ Headless environment detected:[/yellow] {'; '.join(headless_indicators)}") - console.print("[yellow]→ Browser will NOT open automatically. Please use the URL below.[/yellow]\n") + console.print( + f"[yellow]ℹ Headless environment detected:[/yellow] {'; '.join(headless_indicators)}" + ) + console.print( + "[yellow]→ Browser will NOT open automatically. Please use the URL below.[/yellow]\n" + ) else: # Only log to debug, no console output - lib_logger.debug("GUI environment detected, browser auto-open will be attempted") - + lib_logger.debug( + "GUI environment detected, browser auto-open will be attempted" + ) + return is_headless From 36e6348663d987a8002c2524148ad57e74556d5c Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 07:35:02 +0100 Subject: [PATCH 079/221] ci: Compliance check prompt update --- .github/prompts/compliance-check.md | 213 ++++++++++++++++++++++------ 1 file changed, 172 insertions(+), 41 deletions(-) diff --git a/.github/prompts/compliance-check.md b/.github/prompts/compliance-check.md index 7c6d8a9e..32346966 100644 --- a/.github/prompts/compliance-check.md +++ b/.github/prompts/compliance-check.md @@ -37,47 +37,49 @@ A PR is **BLOCKED** when: ## Agentic Environment Expectations -**YOU ARE OPERATING IN AN AGENTIC SYSTEM WHERE MULTIPLE TURNS ARE EXPECTED, REQUIRED, AND DESIRED.** +**YOU ARE OPERATING IN A SELF-DRIVEN AGENTIC SYSTEM WHERE YOU CONTROL YOUR OWN WORKFLOW.** -This is NOT a "complete everything in one response" environment. The system is specifically designed for you to: -- Take MULTIPLE TURNS to complete your work -- Review ONE file (or issue) PER TURN -- State findings after EACH turn -- STOP and wait for the next turn before proceeding +This is NOT a "complete everything in one response" environment. The system is designed for you to: +- Work through MULTIPLE ITERATIONS to complete your analysis +- Focus on ONE file (or issue) PER ITERATION for thorough review +- State findings after EACH iteration +- Then PROCEED to the next item automatically + +**CRITICAL**: You drive the workflow. There is no external system managing "turns" - you simply proceed from one item to the next until all items are reviewed, then produce the final report. **ATTEMPTING TO COMPLETE EVERYTHING IN ONE RESPONSE IS WRONG AND DEFEATS THE PURPOSE OF THIS SYSTEM.** The agentic environment provides focused attention on individual items. Bundling reviews or trying to be "efficient" by processing multiple files at once will result in superficial analysis and missed issues. -**EXPECTATION**: You will take 5-20+ turns to complete a compliance check, depending on PR size. This is normal and correct. +**EXPECTATION**: You will go through 5-20+ iterations to complete a compliance check, depending on PR size. This is normal and correct. For very large PRs, use subtasks to parallelize work (see Section 5.5). -## Mandatory Turn-Based Protocol +## Sequential Analysis Protocol -You MUST follow this strict protocol. Deviation is unacceptable. +You MUST follow this protocol. Deviation is unacceptable. ### Phase 1: Review Previous Issues (if any exist) If `${PREVIOUS_REVIEWS}` is not empty, you MUST check each previously flagged issue individually: -**Turn 1:** +**Iteration 1:** - Focus: Previous Issue #1 ONLY - Action: Check current PR state → Is this issue fixed, still present, or partially fixed? - Output: State your finding clearly -- **STOP** - Do NOT proceed to the next issue +- Then proceed to the next issue -**Turn 2:** +**Iteration 2:** - Focus: Previous Issue #2 ONLY - Action: Check current PR state - Output: State your finding -- **STOP** +- Then proceed to the next issue -Continue this pattern until ALL previous issues are reviewed. One issue per turn. No exceptions. +Continue this pattern until ALL previous issues are reviewed. One issue per iteration. No exceptions. ### Phase 2: Review Files from Affected Groups After previous issues (if any), review each file individually: -**Turn N:** +**Iteration N:** - Focus: File #1 from affected groups - Action: Examine changes for THIS FILE ONLY - Verify: Is this file updated correctly AND completely? @@ -86,21 +88,21 @@ After previous issues (if any), review each file individually: - Provider files: Are ALL necessary changes present? - DOCUMENTATION.md: Does the technical documentation include proper details? - Output: State your findings for THIS FILE -- **STOP** - Do NOT proceed to the next file +- Then proceed to the next file -**Turn N+1:** +**Iteration N+1:** - Focus: File #2 from affected groups - Action: Examine changes for THIS FILE ONLY - Verify: Correctness and completeness - Output: State your findings -- **STOP** +- Then proceed to the next file -Continue until ALL files in affected groups are reviewed. One file per turn. +Continue until ALL files in affected groups are reviewed. One file per iteration. ### Phase 3: Final Report Only after completing Phases 1 and 2: -- Aggregate all your findings from previous turns +- Aggregate all your findings from previous iterations - Fill in the report template - Set GitHub status check - Post the compliance report @@ -108,10 +110,9 @@ Only after completing Phases 1 and 2: ## Forbidden Actions **YOU MUST NOT:** -- Review multiple files in a single turn -- Review multiple previous issues in a single turn +- Review multiple files in a single iteration (unless they are trivially small) +- Review multiple previous issues in a single iteration - Skip stating findings for any item -- Proceed to the next item without explicit turn completion - Bundle reviews "for efficiency" - Try to complete the entire compliance check in one response @@ -160,7 +161,7 @@ If `${PREVIOUS_REVIEWS}` exists, you MUST review each flagged issue individually 2. Compare against current PR state (using the diff you already examined) 3. Determine: Fixed / Still Present / Partially Fixed 4. State your finding with **detailed self-contained description** -5. **STOP** - wait for next turn +5. Proceed to the next issue **CRITICAL: Write Detailed Issue Descriptions** @@ -184,13 +185,13 @@ README incomplete **Why This Matters:** Future compliance checks will re-read these issue descriptions. They need enough detail to understand the problem WITHOUT examining old file states or diffs. You're writing to your future self. -Do NOT review multiple previous issues in one turn. +Do NOT review multiple previous issues in one iteration. ## Step 3: Review Files One-By-One For each file in the affected groups: -**Single Turn Process:** +**Single Iteration Process:** 1. Focus on THIS FILE ONLY 2. Analyze the changes (from the diff you already read) against the group's description guidance 3. Verify correctness: Are the changes appropriate? @@ -200,13 +201,13 @@ For each file in the affected groups: - CHANGELOG: Entry has proper details? - Build script: All necessary updates? 5. State your findings for THIS FILE with detailed description -6. **STOP** - wait for next turn before proceeding to the next file +6. Proceed to the next file ## Step 4: Aggregate and Report After ALL reviews complete: -1. Aggregate findings from all your previous turns +1. Aggregate findings from all your previous iterations 2. Categorize by severity: - ❌ **BLOCKED**: Critical issues (missing documentation, incomplete feature coverage) - ⚠️ **WARNINGS**: Non-blocking concerns (minor missing details) @@ -303,6 +304,100 @@ ${REPORT_TEMPLATE} **Why**: Compliance checking verifies file completeness and correctness, not code quality. +## Parallel Analysis with Subtasks + +For large or complex PRs, use OpenCode's task/subtask capability to parallelize your analysis and avoid context overflow. + +### When to Use Subtasks + +Consider spawning subtasks when: +- **Many files changed**: PR modifies more than 15-20 files across multiple groups +- **Large total diff**: Changes exceed ~2000 lines spread across many files +- **Multiple independent groups**: Several file groups are affected and can be analyzed in parallel +- **Deep analysis needed**: You need to read full file contents (not just diff) to verify completeness + +**Rule of thumb**: A single agent can handle ~2000 lines of changes in one file without subtasks. But 2000 lines spread across 50+ files benefits greatly from parallelization. + +### How to Use Subtasks + +1. **Identify independent work units** - typically one subtask per affected file group +2. **Spawn subtasks in parallel** for each group +3. Each subtask performs deep analysis of its assigned group: + - Read the full file content when needed (not just diff) + - Check cross-references between files in the group + - Verify completeness of documentation, configurations, etc. +4. **Collect subtask reports** with structured findings +5. **Aggregate** all subtask findings into your single compliance report + +### Subtask Instructions Template + +When spawning a subtask, provide clear instructions: + +``` +Analyze the "[Group Name]" file group for compliance. + +Files in this group: +- file1.py +- file2.md + +PR Context: +- PR #${PR_NUMBER}: ${PR_TITLE} +- Changed files in this group: [list relevant files] + +Your task: +1. Read the diff for files in this group +2. Read full file contents where needed for context +3. Verify each file is updated correctly AND completely +4. Check cross-references (e.g., new code is documented, dependencies are listed) + +Return a structured report: +- Group name +- Files reviewed +- Finding per file: COMPLIANT / WARNING / BLOCKED +- Detailed issue descriptions (if any) +- Recommendations +``` + +### Subtask Report Structure + +Each subtask should return: +``` +GROUP: [Group Name] +FILES REVIEWED: file1.py, file2.md +FINDINGS: + - file1.py: ✅ COMPLIANT - [brief reason] + - file2.md: ❌ BLOCKED - [detailed issue description] +ISSUES: + - [Detailed, self-contained issue description for any non-compliant files] +RECOMMENDATIONS: + - [Actionable next steps] +``` + +### Benefits of Subtasks + +- **Reduces context overflow** on large PRs +- **Enables deeper analysis** - subtasks can read full files, not just diffs +- **Parallelizes independent work** - faster overall completion +- **Maintains focused attention** on each group +- **Scales with PR size** - spawn more subtasks for larger PRs + +### Example Workflow + +``` +Main agent identifies 4 affected groups, spawns: + ├── Subtask 1: "Documentation" group → Returns findings + ├── Subtask 2: "Python Dependencies" group → Returns findings + ├── Subtask 3: "Provider Configuration" group → Returns findings + └── Subtask 4: "Proxy Application" group → Returns findings + +Main agent: + 1. Waits for all subtasks to complete + 2. Aggregates findings from all subtasks + 3. Posts single unified compliance report +``` + +**Important**: Avoid copying large code excerpts in subtask reports. Cite file paths, function names, and line ranges instead. + --- # 6. [OUTPUT REQUIREMENTS] @@ -409,11 +504,11 @@ Clear, actionable guidance for the author: # 7. [REFERENCE] -## Example Turn-By-Turn Workflow +## Example Sequential Workflow Here's what a proper compliance check looks like: -**Turn 0 (FIRST ACTION):** +**Iteration 0 (FIRST ACTION):** ```bash # Examine the diff file cat ${DIFF_PATH} @@ -425,37 +520,41 @@ cat ${DIFF_PATH} # - DOCUMENTATION.md: not updated with new provider details ``` -**Turn 1:** +**Iteration 1:** ``` Checking Previous Issue #1: "Missing provider documentation" Current status: FIXED - DOCUMENTATION.md now includes NewProvider section +→ Proceeding to next issue... ``` -**Turn 2:** +**Iteration 2:** ``` Checking Previous Issue #2: "Incomplete README setup instructions" Current status: STILL PRESENT - Quick Start section still doesn't mention NewProvider setup The README lists the provider in Features but Quick Start lacks configuration steps. This was flagged in previous review and has not been addressed. +→ Proceeding to file reviews... ``` -**Turn 3:** +**Iteration 3:** ``` Reviewing File: requirements.txt (Python Dependencies group) Analysis: Added 'aiohttp' dependency for async HTTP support Verification: Dependency listed ✓, compatible with existing packages ✓ Finding: COMPLIANT +→ Proceeding to next file... ``` -**Turn 4:** +**Iteration 4:** ``` Reviewing File: src/rotator_library/pyproject.toml (Python Dependencies group) Analysis: No changes to pyproject.toml Verification: pyproject.toml doesn't need update for this change (aiohttp is a proxy_app dependency, not rotator_library) Finding: COMPLIANT +→ Proceeding to next file... ``` -**Turn 5:** +**Iteration 5:** ``` Reviewing File: README.md (Documentation group) Analysis: Added NewProvider to Features list with brief description @@ -463,17 +562,19 @@ Verification: - Feature mentioned ✓ - Quick Start section: INCOMPLETE ✗ - No setup instructions for NewProvider credentials Finding: BLOCKED - Setup instructions incomplete (pre-existing issue not fixed) +→ Proceeding to next file... ``` -**Turn 6:** +**Iteration 6:** ``` Reviewing File: DOCUMENTATION.md (Documentation group) Analysis: Added NewProvider section with API reference and configuration options Verification: Provider documented ✓, configuration examples provided ✓, architecture notes present ✓ Finding: COMPLIANT +→ All files reviewed. Proceeding to final report... ``` -**Turn 7:** +**Final Iteration:** ``` All reviews complete. Aggregating findings... @@ -486,15 +587,45 @@ Overall: ❌ BLOCKED Setting status to 'failure' and posting detailed report... ``` +## Example: Using Subtasks for Large PRs + +For a large PR with 40+ files across 4 groups: + +**Main Agent:** +``` +This PR modifies 45 files across 4 file groups with ~2500 lines changed. +Spawning parallel subtasks for efficient analysis... + +Subtask 1: Analyze "Documentation" group (README.md, DOCUMENTATION.md, Deployment guide.md) +Subtask 2: Analyze "Python Dependencies" group (requirements.txt, pyproject.toml) +Subtask 3: Analyze "Provider Configuration" group (15 provider files) +Subtask 4: Analyze "Proxy Application" group (5 application files) +``` + +**After subtasks complete:** +``` +Received reports from all 4 subtasks. Aggregating findings... + +Subtask 1 (Documentation): ⚠️ WARNING - Minor gaps in Deployment guide.md +Subtask 2 (Python Dependencies): ✅ COMPLIANT +Subtask 3 (Provider Configuration): ❌ BLOCKED - New provider missing from model_definitions.py +Subtask 4 (Proxy Application): ✅ COMPLIANT + +Overall: ❌ BLOCKED + +Posting unified compliance report with all findings... +``` + ## Critical Reminders 1. **READ DIFF ONCE**: Examine `${DIFF_PATH}` at the very beginning for full context -2. **ONE ITEM PER TURN**: Review exactly one file or one previous issue per turn -3. **STATE FINDINGS**: Always output your finding before stopping +2. **ONE ITEM PER ITERATION**: Review exactly one file or one previous issue per iteration +3. **STATE FINDINGS**: Always output your finding before proceeding 4. **DETAILED DESCRIPTIONS**: Write issue descriptions for your future self - be specific and complete -5. **MULTIPLE TURNS EXPECTED**: This system REQUIRES multiple turns - do not try to complete in one +5. **SELF-DRIVEN WORKFLOW**: You control the flow - proceed through all items, then produce the final report 6. **VERIFY COMPLETELY**: Check that files are not just touched, but updated correctly AND completely 7. **FOCUS ATTENTION**: Single-file review ensures you catch missing steps, incomplete documentation, etc. +8. **USE SUBTASKS FOR LARGE PRS**: When PR has many files across groups, parallelize with subtasks --- @@ -502,4 +633,4 @@ Setting status to 'failure' and posting detailed report... **First action:** Read `${DIFF_PATH}` to understand all changes. -Then analyze the PR context above, identify affected file groups, and start your turn-by-turn review. Remember: ONE item at a time, state detailed findings, STOP, wait for next turn. +Then analyze the PR context above, identify affected file groups, and proceed through your sequential review. For large PRs (many files, large diffs), consider using subtasks to parallelize analysis by group. Remember: focus on ONE item at a time, state detailed findings, then continue to the next item until all reviews are complete. Finally, aggregate findings and post the compliance report. From d389837afaf4a86d0ef3533945ff0b25f5d4c1e8 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 08:25:03 +0100 Subject: [PATCH 080/221] =?UTF-8?q?feat(antigravity):=20=E2=9C=A8=20implem?= =?UTF-8?q?ent=20credential=20prioritization=20for=20tier-based=20routing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add two new methods to AntigravityProvider to support credential prioritization based on account tier: - `get_credential_priority()`: Returns priority levels (1-10) based on Antigravity tier, with paid tiers getting highest priority (1), free tier getting medium priority (2), and legacy/unknown getting lowest priority (10) - `get_model_tier_requirement()`: Returns None for all models since Antigravity has no model-tier restrictions This enables the credential rotation system to intelligently prioritize paid tier credentials over free tier credentials when selecting accounts for API requests. --- .../providers/antigravity_provider.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index fb63a5d9..b17b21d9 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -577,6 +577,51 @@ def _log_config(self) -> None: f"claude_fix={self._enable_claude_tool_fix}, thinking_sanitization={self._enable_thinking_sanitization}" ) + # ========================================================================= + # CREDENTIAL PRIORITIZATION + # ========================================================================= + + def get_credential_priority(self, credential: str) -> Optional[int]: + """ + Returns priority based on Antigravity tier. + Paid tiers: priority 1 (highest) + Free tier: priority 2 + Legacy/Unknown: priority 10 (lowest) + + Args: + credential: The credential path + + Returns: + Priority level (1-10) or None if tier not yet discovered + """ + tier = self.project_tier_cache.get(credential) + if not tier: + return None # Not yet discovered + + # Paid tiers get highest priority + if tier not in ["free-tier", "legacy-tier", "unknown"]: + return 1 + + # Free tier gets lower priority + if tier == "free-tier": + return 2 + + # Legacy and unknown get even lower + return 10 + + def get_model_tier_requirement(self, model: str) -> Optional[int]: + """ + Returns the minimum priority tier required for a model. + Antigravity has no model-tier restrictions - all models work on all tiers. + + Args: + model: The model name (with or without provider prefix) + + Returns: + None - no restrictions for any model + """ + return None + # ========================================================================= # MODEL UTILITIES # ========================================================================= From df7a7566e4ac1aee2dfa6bcd6c8d273cfa034abd Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 08:25:41 +0100 Subject: [PATCH 081/221] =?UTF-8?q?docs(antigravity):=20=F0=9F=93=9A=20upd?= =?UTF-8?q?ate=20documentation=20for=20credential=20prioritization=20and?= =?UTF-8?q?=20model=20support=20changes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add credential prioritization details for automatic detection of paid vs free tier credentials - Update model support list to reflect current Gemini 3 Pro and Claude 4.5 variants - Remove outdated Gemini 2.5 references from thinking support section - Clarify Claude Sonnet 4.5 supports both thinking and non-thinking modes - Document Claude Opus 4.5 as thinking-only variant - Expand tool hallucination prevention to include Claude models --- src/rotator_library/README.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/rotator_library/README.md b/src/rotator_library/README.md index 2050f1ba..872e80e3 100644 --- a/src/rotator_library/README.md +++ b/src/rotator_library/README.md @@ -215,13 +215,14 @@ Use this tool to: ### Antigravity - **Auth**: Uses OAuth 2.0 flow similar to Gemini CLI, with Antigravity-specific credentials and scopes. -- **Models**: Supports Gemini 2.5 (Pro/Flash), Gemini 3 (Pro/Image), and Claude Sonnet 4.5 via Google's internal Antigravity API. +- **Credential Prioritization**: Automatic detection and prioritization of paid vs free tier credentials (paid tier resets every 5 hours, free tier resets weekly). +- **Models**: Supports Gemini 3 Pro, Claude Sonnet 4.5 (with/without thinking), and Claude Opus 4.5 (thinking only) via Google's internal Antigravity API. - **Thought Signature Caching**: Server-side caching of `thoughtSignature` data for multi-turn conversations with Gemini 3 models. -- **Tool Hallucination Prevention**: Automatic injection of system instructions and parameter signatures for Gemini 3 to prevent tool parameter hallucination. +- **Tool Hallucination Prevention**: Automatic injection of system instructions and parameter signatures for Gemini 3 and Claude to prevent tool parameter hallucination. - **Thinking Support**: - - Gemini 2.5: Uses `thinkingBudget` (integer tokens) - Gemini 3: Uses `thinkingLevel` (string: "low"/"high") - - Claude: Uses `thinkingBudget` via Antigravity proxy + - Claude Sonnet 4.5: Uses `thinkingBudget` (optional - supports both thinking and non-thinking modes) + - Claude Opus 4.5: Uses `thinkingBudget` (always uses thinking variant) - **Base URL Fallback**: Automatic fallback between sandbox and production endpoints. ## Error Handling and Cooldowns From 1d1a62be14143a805f80087992df594041318d66 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 08:26:56 +0100 Subject: [PATCH 082/221] =?UTF-8?q?docs:=20=F0=9F=93=9A=20update=20documen?= =?UTF-8?q?tation=20to=20reflect=20gemini=202.5=20removal=20and=20claude?= =?UTF-8?q?=20sonnet=20dual-mode=20support?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit updates both README.md and DOCUMENTATION.md to accurately reflect recent changes to the Antigravity provider: - Remove all references to Gemini 2.5 models (Pro/Flash) as they are no longer supported - Document Claude Sonnet 4.5's dual-mode capability (thinking and non-thinking variants) - Add provider support section explaining credential prioritization implementation for both Gemini CLI and Antigravity providers - Clarify that Claude Opus 4.5 only supports thinking mode - Update model-specific logic documentation to reflect current architecture (Gemini 3, Claude Sonnet, Claude Opus) - Add credential tier reset timing details (paid tier: 5 hours, free tier: weekly) - Remove outdated "NEW" badges and function call response pairing references --- DOCUMENTATION.md | 25 +++++++++++++++---------- README.md | 9 +++++---- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index 39b266b0..cf985326 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -361,6 +361,13 @@ def get_model_tier_requirement(self, model: str) -> Optional[int]: return None # All other models have no restrictions ``` +**Provider Support:** + +The following providers implement credential prioritization: + +- **Gemini CLI**: Paid tier (priority 1), Free tier (priority 2), Legacy/Unknown (priority 10). Gemini 3 models require paid tier. +- **Antigravity**: Same priority system as Gemini CLI. No model-tier restrictions (all models work on all tiers). Paid tier resets every 5 hours, free tier resets weekly. + **Usage Manager Integration:** The `acquire_key()` method has been enhanced to: @@ -391,22 +398,18 @@ A modular, shared caching system for providers to persist conversation state acr ### 3.5. Antigravity (`antigravity_provider.py`) -The most sophisticated provider implementation, supporting Google's internal Antigravity API for Gemini and Claude models (including **Claude Opus 4.5**, Anthropic's most powerful model). +The most sophisticated provider implementation, supporting Google's internal Antigravity API for Gemini 3 and Claude models (including **Claude Opus 4.5**, Anthropic's most powerful model). #### Architecture - **Unified Streaming/Non-Streaming**: Single code path handles both response types with optimal transformations - **Thought Signature Caching**: Server-side caching of encrypted signatures for multi-turn Gemini 3 conversations -- **Model-Specific Logic**: Automatic configuration based on model type (Gemini 2.5, Gemini 3, Claude) +- **Model-Specific Logic**: Automatic configuration based on model type (Gemini 3, Claude Sonnet, Claude Opus) +- **Credential Prioritization**: Automatic tier detection with paid credentials prioritized over free (paid tier resets every 5 hours, free tier resets weekly) #### Model Support -**Gemini 2.5 (Pro/Flash):** -- Uses `thinkingBudget` parameter (integer tokens: -1 for auto, 0 to disable, or specific value) -- Standard safety settings and toolConfig -- Stream processing with thinking content separation - -**Gemini 3 (Pro/Image):** +**Gemini 3 Pro:** - Uses `thinkingLevel` parameter (string: "low" or "high") - **Tool Hallucination Prevention**: - Automatic system instruction injection explaining custom tool schema rules @@ -427,8 +430,10 @@ The most sophisticated provider implementation, supporting Google's internal Ant - Increased default max output tokens to 64000 to accommodate thinking output **Claude Sonnet 4.5:** -- Proxied through Antigravity API (uses internal model name `claude-sonnet-4-5-thinking`) -- Uses `thinkingBudget` parameter like Gemini 2.5 +- Proxied through Antigravity API +- **Supports both thinking and non-thinking modes**: + - With `reasoning_effort`: Uses `claude-sonnet-4-5-thinking` variant with `thinkingBudget` + - Without `reasoning_effort`: Uses standard `claude-sonnet-4-5` variant - **Thinking Preservation**: Caches thinking content using composite keys (tool_call_id + text_hash) - **Schema Cleaning**: Removes unsupported properties (`$schema`, `additionalProperties`, `const` → `enum`) diff --git a/README.md b/README.md index 85df3b70..9c3e3809 100644 --- a/README.md +++ b/README.md @@ -28,13 +28,14 @@ This project provides a powerful solution for developers building complex applic - **OpenAI-Compatible Proxy**: Offers a familiar API interface with additional endpoints for model and provider discovery. - **Advanced Model Filtering**: Supports both blacklists and whitelists to give you fine-grained control over which models are available through the proxy. -- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 2.5, Gemini 3, and Claude models with advanced features: - - **🚀 NEW: Claude Opus 4.5** - Anthropic's most powerful model, now available via Antigravity! - - Claude Sonnet 4.5 with extended thinking support +- **🆕 Antigravity Provider**: Full support for Google's internal Antigravity API, providing access to Gemini 3 and Claude models with advanced features: + - **🚀 Claude Opus 4.5** - Anthropic's most powerful model (thinking mode only) + - **Claude Sonnet 4.5** - Supports both thinking and non-thinking modes + - **Gemini 3 Pro** - With thinkingLevel support (low/high) + - Credential prioritization with automatic paid/free tier detection - Thought signature caching for multi-turn conversations - Tool hallucination prevention via parameter signature injection - Automatic thinking block sanitization for Claude models (with recovery strategies) - - Improved function call response pairing with three-tier matching strategy - Note: Claude thinking mode requires careful conversation state management (see [Antigravity documentation](DOCUMENTATION.md#antigravity-claude-extended-thinking-sanitization) for details) - **🆕 Credential Prioritization**: Automatic tier detection and priority-based credential selection ensures paid-tier credentials are used for premium models that require them. - **🆕 Weighted Random Rotation**: Configurable credential rotation strategy - choose between deterministic (perfect balance) or weighted random (unpredictable, harder to fingerprint) selection. From fa51b1ad541546866e80aa2f009e5f6145500710 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Fri, 5 Dec 2025 23:59:06 +0100 Subject: [PATCH 083/221] =?UTF-8?q?fix(error-handler):=20=F0=9F=90=9B=20ha?= =?UTF-8?q?ndle=20compound=20duration=20formats=20in=20retry-after=20parsi?= =?UTF-8?q?ng?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Added support for parsing complex duration strings commonly returned by Antigravity and Google APIs, such as "156h14m36.752463453s" or "562476.752463453s". - Introduced `_parse_duration_string()` helper to parse compound duration formats (hours, minutes, seconds with decimals) - Updated `extract_retry_after_from_body()` to handle both simple and compound duration strings - Enhanced `get_retry_after()` to iterate through all error details (not just first item) and check both RetryInfo and ErrorInfo metadata - Added `httpx.HTTPStatusError` to exception handling in client retry logic - Fixed formatting inconsistencies in conditional statements for rate limit handling --- src/rotator_library/client.py | 22 ++-- src/rotator_library/error_handler.py | 161 ++++++++++++++++++--------- 2 files changed, 121 insertions(+), 62 deletions(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 30021d0b..cf1bb1cf 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -620,8 +620,9 @@ async def _safe_streaming_wrapper( litellm.ServiceUnavailableError, litellm.InternalServerError, APIConnectionError, + httpx.HTTPStatusError, ) as e: - # This is a critical, typed error from litellm that signals a key failure. + # This is a critical, typed error from litellm or httpx that signals a key failure. # We do not try to parse it here. We wrap it and raise it immediately # for the outer retry loop to handle. lib_logger.warning( @@ -1065,7 +1066,10 @@ async def _execute_with_retry( ) # Only trigger provider-wide cooldown for rate limits, not quota issues - if classified_error.status_code == 429 and classified_error.error_type != "quota_exceeded": + if ( + classified_error.status_code == 429 + and classified_error.error_type != "quota_exceeded" + ): cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration @@ -1225,9 +1229,9 @@ async def _execute_with_retry( # Handle rate limits with cooldown (exclude quota_exceeded from provider-wide cooldown) if ( - (classified_error.status_code == 429 and classified_error.error_type != "quota_exceeded") - or classified_error.error_type == "rate_limit" - ): + classified_error.status_code == 429 + and classified_error.error_type != "quota_exceeded" + ) or classified_error.error_type == "rate_limit": cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration @@ -1494,7 +1498,7 @@ async def _streaming_acompletion_with_retry( lib_logger.info( f"Attempting stream with credential {mask_credential(current_cred)} (Attempt {attempt + 1}/{self.max_retries})" ) - + if pre_request_callback: try: await pre_request_callback( @@ -1973,9 +1977,9 @@ async def _streaming_acompletion_with_retry( # Handle rate limits with cooldown (exclude quota_exceeded) if ( - (classified_error.status_code == 429 and classified_error.error_type != "quota_exceeded") - or classified_error.error_type == "rate_limit" - ): + classified_error.status_code == 429 + and classified_error.error_type != "quota_exceeded" + ) or classified_error.error_type == "rate_limit": cooldown_duration = classified_error.retry_after or 60 await self.cooldown_manager.start_cooldown( provider, cooldown_duration diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index fa24d4af..038d4f19 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -18,12 +18,60 @@ ) +def _parse_duration_string(duration_str: str) -> Optional[int]: + """ + Parse duration strings in various formats to total seconds. + + Handles: + - Compound durations: '156h14m36.752463453s', '2h30m', '45m30s' + - Simple durations: '562476.752463453s', '3600s', '60m', '2h' + - Plain seconds (no unit): '562476' + + Args: + duration_str: Duration string to parse + + Returns: + Total seconds as integer, or None if parsing fails + """ + if not duration_str: + return None + + total_seconds = 0 + remaining = duration_str.strip().lower() + + # Try parsing as plain number first (no units) + try: + return int(float(remaining)) + except ValueError: + pass + + # Parse hours component + hour_match = re.match(r"(\d+)h", remaining) + if hour_match: + total_seconds += int(hour_match.group(1)) * 3600 + remaining = remaining[hour_match.end() :] + + # Parse minutes component + min_match = re.match(r"(\d+)m", remaining) + if min_match: + total_seconds += int(min_match.group(1)) * 60 + remaining = remaining[min_match.end() :] + + # Parse seconds component (including decimals like 36.752463453s) + sec_match = re.match(r"([\d.]+)s", remaining) + if sec_match: + total_seconds += int(float(sec_match.group(1))) + + return total_seconds if total_seconds > 0 else None + + def extract_retry_after_from_body(error_body: Optional[str]) -> Optional[int]: """ Extract the retry-after time from an API error response body. Handles various error formats including: - Gemini CLI: "Your quota will reset after 39s." + - Antigravity: "quota will reset after 156h14m36s" - Generic: "quota will reset after 120s", "retry after 60s" Args: @@ -35,21 +83,21 @@ def extract_retry_after_from_body(error_body: Optional[str]) -> Optional[int]: if not error_body: return None - # Pattern to match various "reset after Xs" or "retry after Xs" formats + # Pattern to match various "reset after" formats - capture the full duration string patterns = [ - r"quota will reset after\s*(\d+)s", - r"reset after\s*(\d+)s", - r"retry after\s*(\d+)s", + r"quota will reset after\s*([\dhmso.]+)", # Matches compound: 156h14m36s or 120s + r"reset after\s*([\dhmso.]+)", + r"retry after\s*([\dhmso.]+)", r"try again in\s*(\d+)\s*seconds?", ] for pattern in patterns: match = re.search(pattern, error_body, re.IGNORECASE) if match: - try: - return int(match.group(1)) - except (ValueError, IndexError): - continue + duration_str = match.group(1) + result = _parse_duration_string(duration_str) + if result is not None: + return result return None @@ -311,6 +359,11 @@ def get_retry_after(error: Exception) -> Optional[int]: Extracts the 'retry-after' duration in seconds from an exception message. Handles both integer and string representations of the duration, as well as JSON bodies. Also checks HTTP response headers for httpx.HTTPStatusError instances. + + Supports Antigravity/Google API error formats: + - RetryInfo with retryDelay: "562476.752463453s" + - ErrorInfo metadata with quotaResetDelay: "156h14m36.752463453s" + - Human-readable message: "quota will reset after 156h14m36s" """ # 0. For httpx errors, check response headers first (most reliable) if isinstance(error, httpx.HTTPStatusError): @@ -341,79 +394,81 @@ def get_retry_after(error: Exception) -> Optional[int]: error_str = str(error).lower() - # 1. Try to parse JSON from the error string to find 'retryDelay' + # 1. Try to parse JSON from the error string to find retry info + # Antigravity errors have details array with RetryInfo and/or ErrorInfo try: # It's common for the actual JSON to be embedded in the string representation json_match = re.search(r"(\{.*\})", error_str, re.DOTALL) if json_match: error_json = json.loads(json_match.group(1)) - retry_info = error_json.get("error", {}).get("details", [{}])[0] - if retry_info.get("@type") == "type.googleapis.com/google.rpc.RetryInfo": - delay_str = retry_info.get("retryDelay", {}).get("seconds") - if delay_str: - return int(delay_str) - # Fallback for the other format - delay_str = retry_info.get("retryDelay") - if isinstance(delay_str, str) and delay_str.endswith("s"): - return int(delay_str[:-1]) + details = error_json.get("error", {}).get("details", []) + + # Iterate through ALL details items (not just index 0) + for detail in details: + detail_type = detail.get("@type", "") + + # Check RetryInfo for retryDelay (most authoritative) + if detail_type == "type.googleapis.com/google.rpc.retryinfo": + delay_str = detail.get("retrydelay") + if delay_str: + # Handle both {"seconds": "123"} format and "123.456s" string format + if isinstance(delay_str, dict): + seconds = delay_str.get("seconds") + if seconds: + return int(float(seconds)) + elif isinstance(delay_str, str): + result = _parse_duration_string(delay_str) + if result is not None: + return result + + # Check ErrorInfo metadata for quotaResetDelay (Antigravity-specific) + if detail_type == "type.googleapis.com/google.rpc.errorinfo": + metadata = detail.get("metadata", {}) + quota_reset_delay = metadata.get("quotaresetdelay") + if quota_reset_delay: + result = _parse_duration_string(quota_reset_delay) + if result is not None: + return result except (json.JSONDecodeError, IndexError, KeyError, TypeError): pass # If JSON parsing fails, proceed to regex and attribute checks - # 2. Common regex patterns for 'retry-after' (with duration format support) + # 2. Common regex patterns for 'retry-after' (with compound duration support) patterns = [ r"retry[-_\s]after:?\s*(\d+)", # Matches: retry-after, retry_after, retry after r"retry in\s*(\d+)\s*seconds?", r"wait for\s*(\d+)\s*seconds?", - r'"retryDelay":\s*"(\d+)s"', + r'"retrydelay":\s*"([\d.]+)s?"', # retryDelay in JSON r"x-ratelimit-reset:?\s*(\d+)", - r"quota will reset after\s*(\d+)s", # Gemini CLI rate limit format - r"reset after\s*(\d+)s", # Generic reset after format + # Compound duration patterns (Antigravity format) + r"quota will reset after\s*([\dhms.]+)", # e.g., "156h14m36s" or "120s" + r"reset after\s*([\dhms.]+)", + r'"quotaresetdelay":\s*"([\dhms.]+)"', # quotaResetDelay in JSON ] for pattern in patterns: match = re.search(pattern, error_str) if match: + duration_str = match.group(1) + # Try parsing as compound duration first + result = _parse_duration_string(duration_str) + if result is not None: + return result + # Fallback to simple integer try: - return int(match.group(1)) + return int(duration_str) except (ValueError, IndexError): continue - # 3. Handle duration formats like "60s", "2m", "1h" - duration_match = re.search(r"(\d+)\s*([smh])", error_str) - if duration_match: - try: - value = int(duration_match.group(1)) - unit = duration_match.group(2) - if unit == "s": - return value - elif unit == "m": - return value * 60 - elif unit == "h": - return value * 3600 - except (ValueError, IndexError): - pass - - # 4. Handle cases where the error object itself has the attribute + # 3. Handle cases where the error object itself has the attribute if hasattr(error, "retry_after"): value = getattr(error, "retry_after") if isinstance(value, int): return value if isinstance(value, str): - # Try to parse string formats - if value.isdigit(): - return int(value) - # Handle "60s", "2m" format in attribute - duration_match = re.search(r"(\d+)\s*([smh])", value.lower()) - if duration_match: - val = int(duration_match.group(1)) - unit = duration_match.group(2) - if unit == "s": - return val - elif unit == "m": - return val * 60 - elif unit == "h": - return val * 3600 + result = _parse_duration_string(value) + if result is not None: + return result return None From cde9cb012be4e31ec83fe70c5282aba7ff4255be Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sat, 6 Dec 2025 00:34:04 +0100 Subject: [PATCH 084/221] fix(antigravity-provider): handle multiple consecutive system messages in prompt processing The previous implementation only extracted the first system message from the messages array. This refactor changes the logic to use a while loop that processes all consecutive system messages at the beginning of the array, accumulating their content parts before constructing the system instruction. - Changed from single system message extraction to loop-based consecutive system message handling - Accumulate all system message parts into a single system_parts list - Construct system_instruction only after all system messages are processed - Ensures no system message content is lost when multiple system messages are provided --- .../providers/antigravity_provider.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index b17b21d9..b8226a8a 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -1919,15 +1919,18 @@ def _transform_messages( system_instruction = None gemini_contents = [] - # Extract system prompt - if messages and messages[0].get("role") == "system": + # Extract system prompts (handle multiple consecutive system messages) + system_parts = [] + while messages and messages[0].get("role") == "system": system_content = messages.pop(0).get("content", "") if system_content: - system_parts = self._parse_content_parts( + new_parts = self._parse_content_parts( system_content, _strip_cache_control=True ) - if system_parts: - system_instruction = {"role": "user", "parts": system_parts} + system_parts.extend(new_parts) + + if system_parts: + system_instruction = {"role": "user", "parts": system_parts} # Build tool_call_id → name mapping tool_id_to_name = {} From abdc406e9cdf1593c78afad743aa6a62d87a1a65 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sat, 6 Dec 2025 00:49:39 +0100 Subject: [PATCH 085/221] =?UTF-8?q?fix(error-handler):=20=F0=9F=94=A8=20ex?= =?UTF-8?q?tract=20JSON=20retry=20parsing=20into=20dedicated=20function?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extracted retry delay parsing logic from `get_retry_after()` into a new `_extract_retry_from_json_body()` helper function to improve code organization and maintainability. - New `_extract_retry_from_json_body()` function handles parsing of Antigravity/Google API JSON error responses - Preserves case-sensitive key handling for API responses (RetryInfo, ErrorInfo) - Prioritizes response body JSON parsing over HTTP headers for httpx errors - Maintains backward compatibility with all existing retry delay extraction patterns - Improves code readability by separating JSON parsing concerns from the main retry extraction logic --- src/rotator_library/error_handler.py | 128 ++++++++++++++++++--------- 1 file changed, 84 insertions(+), 44 deletions(-) diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index 038d4f19..3676d44c 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -354,6 +354,66 @@ def __str__(self): return f"ClassifiedError(type={self.error_type}, status={self.status_code}, retry_after={self.retry_after}, original_exc={self.original_exception})" +def _extract_retry_from_json_body(json_text: str) -> Optional[int]: + """ + Extract retry delay from a JSON error response body. + + Handles Antigravity/Google API error formats with details array containing: + - RetryInfo with retryDelay: "562476.752463453s" + - ErrorInfo metadata with quotaResetDelay: "156h14m36.752463453s" + + Args: + json_text: JSON string (original case, not lowercased) + + Returns: + Retry delay in seconds, or None if not found + """ + try: + # Find JSON object in the text + json_match = re.search(r"(\{.*\})", json_text, re.DOTALL) + if not json_match: + return None + + error_json = json.loads(json_match.group(1)) + details = error_json.get("error", {}).get("details", []) + + # Iterate through ALL details items (not just index 0) + for detail in details: + detail_type = detail.get("@type", "") + + # Check RetryInfo for retryDelay (most authoritative) + # Note: Case-sensitive key names as returned by API + if "google.rpc.RetryInfo" in detail_type: + delay_str = detail.get("retryDelay") + if delay_str: + # Handle both {"seconds": "123"} format and "123.456s" string format + if isinstance(delay_str, dict): + seconds = delay_str.get("seconds") + if seconds: + return int(float(seconds)) + elif isinstance(delay_str, str): + result = _parse_duration_string(delay_str) + if result is not None: + return result + + # Check ErrorInfo metadata for quotaResetDelay (Antigravity-specific) + if "google.rpc.ErrorInfo" in detail_type: + metadata = detail.get("metadata", {}) + # Try both camelCase and lowercase variants + quota_reset_delay = metadata.get("quotaResetDelay") or metadata.get( + "quotaresetdelay" + ) + if quota_reset_delay: + result = _parse_duration_string(quota_reset_delay) + if result is not None: + return result + + except (json.JSONDecodeError, IndexError, KeyError, TypeError): + pass + + return None + + def get_retry_after(error: Exception) -> Optional[int]: """ Extracts the 'retry-after' duration in seconds from an exception message. @@ -365,8 +425,20 @@ def get_retry_after(error: Exception) -> Optional[int]: - ErrorInfo metadata with quotaResetDelay: "156h14m36.752463453s" - Human-readable message: "quota will reset after 156h14m36s" """ - # 0. For httpx errors, check response headers first (most reliable) + # 0. For httpx errors, check response body and headers if isinstance(error, httpx.HTTPStatusError): + # First, try to parse the response body JSON (contains retryDelay/quotaResetDelay) + # This is where Antigravity puts the retry information + try: + response_text = error.response.text + if response_text: + result = _extract_retry_from_json_body(response_text) + if result is not None: + return result + except Exception: + pass # Response body may not be available + + # Fallback to HTTP headers headers = error.response.headers # Check standard Retry-After header (case-insensitive) retry_header = headers.get("retry-after") or headers.get("Retry-After") @@ -392,62 +464,30 @@ def get_retry_after(error: Exception) -> Optional[int]: except (ValueError, TypeError): pass - error_str = str(error).lower() - - # 1. Try to parse JSON from the error string to find retry info - # Antigravity errors have details array with RetryInfo and/or ErrorInfo - try: - # It's common for the actual JSON to be embedded in the string representation - json_match = re.search(r"(\{.*\})", error_str, re.DOTALL) - if json_match: - error_json = json.loads(json_match.group(1)) - details = error_json.get("error", {}).get("details", []) - - # Iterate through ALL details items (not just index 0) - for detail in details: - detail_type = detail.get("@type", "") - - # Check RetryInfo for retryDelay (most authoritative) - if detail_type == "type.googleapis.com/google.rpc.retryinfo": - delay_str = detail.get("retrydelay") - if delay_str: - # Handle both {"seconds": "123"} format and "123.456s" string format - if isinstance(delay_str, dict): - seconds = delay_str.get("seconds") - if seconds: - return int(float(seconds)) - elif isinstance(delay_str, str): - result = _parse_duration_string(delay_str) - if result is not None: - return result - - # Check ErrorInfo metadata for quotaResetDelay (Antigravity-specific) - if detail_type == "type.googleapis.com/google.rpc.errorinfo": - metadata = detail.get("metadata", {}) - quota_reset_delay = metadata.get("quotaresetdelay") - if quota_reset_delay: - result = _parse_duration_string(quota_reset_delay) - if result is not None: - return result - - except (json.JSONDecodeError, IndexError, KeyError, TypeError): - pass # If JSON parsing fails, proceed to regex and attribute checks + # 1. Try to parse JSON from the error string representation + # Some exceptions embed JSON in their string representation + error_str = str(error) + result = _extract_retry_from_json_body(error_str) + if result is not None: + return result # 2. Common regex patterns for 'retry-after' (with compound duration support) + # Use lowercase for pattern matching + error_str_lower = error_str.lower() patterns = [ r"retry[-_\s]after:?\s*(\d+)", # Matches: retry-after, retry_after, retry after r"retry in\s*(\d+)\s*seconds?", r"wait for\s*(\d+)\s*seconds?", - r'"retrydelay":\s*"([\d.]+)s?"', # retryDelay in JSON + r'"retrydelay":\s*"([\d.]+)s?"', # retryDelay in JSON (lowercased) r"x-ratelimit-reset:?\s*(\d+)", # Compound duration patterns (Antigravity format) r"quota will reset after\s*([\dhms.]+)", # e.g., "156h14m36s" or "120s" r"reset after\s*([\dhms.]+)", - r'"quotaresetdelay":\s*"([\dhms.]+)"', # quotaResetDelay in JSON + r'"quotaresetdelay":\s*"([\dhms.]+)"', # quotaResetDelay in JSON (lowercased) ] for pattern in patterns: - match = re.search(pattern, error_str) + match = re.search(pattern, error_str_lower) if match: duration_str = match.group(1) # Try parsing as compound duration first From 4dfb828cf02537dfa7a42d148b2e6559e8997406 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sat, 6 Dec 2025 04:11:15 +0100 Subject: [PATCH 086/221] =?UTF-8?q?feat(providers):=20=E2=9C=A8=20implemen?= =?UTF-8?q?t=20credential=20tier=20initialization=20and=20persistence=20sy?= =?UTF-8?q?stem?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces a comprehensive credential tier management system across the library, enabling automatic tier detection, persistence, and intelligent credential prioritization at startup. - Add `initialize_credentials()` method to `ProviderInterface` for startup credential loading - Add `get_credential_tier_name()` method to expose human-readable tier names for logging - Implement tier persistence in credential files via `_proxy_metadata` field - Add lazy-loading fallback for tier data when not in memory cache - Introduce `BackgroundRefresher._initialize_credentials()` to pre-load all provider tiers before refresh loop - Pass `credential_tier_names` map through client to usage_manager for enhanced logging - Update `UsageManager.acquire_key()` to display tier information in acquisition logs - Make `ModelDefinitions` a singleton to prevent duplicate loading across providers - Add comprehensive 3-line startup summary showing provider counts, credentials, and tier breakdown - Implement tier-aware logging in Antigravity and GeminiCli providers with disk persistence - Fix provider instance lookup for OAuth providers by handling `_oauth` suffix correctly This ensures all credential priorities are known before any API calls, preventing unknown credentials from getting priority 999 and improving load balancing from the first request. --- src/rotator_library/background_refresher.py | 109 +- src/rotator_library/client.py | 33 +- src/rotator_library/model_definitions.py | 21 +- .../providers/antigravity_provider.py | 125 ++ .../providers/gemini_cli_provider.py | 1269 +++++++++++------ .../providers/provider_interface.py | 83 +- src/rotator_library/usage_manager.py | 198 ++- 7 files changed, 1330 insertions(+), 508 deletions(-) diff --git a/src/rotator_library/background_refresher.py b/src/rotator_library/background_refresher.py index 4c1fc26f..a6830fa8 100644 --- a/src/rotator_library/background_refresher.py +++ b/src/rotator_library/background_refresher.py @@ -8,28 +8,35 @@ if TYPE_CHECKING: from .client import RotatingClient -lib_logger = logging.getLogger('rotator_library') +lib_logger = logging.getLogger("rotator_library") + class BackgroundRefresher: """ A background task that periodically checks and refreshes OAuth tokens to ensure they remain valid. """ - def __init__(self, client: 'RotatingClient'): + + def __init__(self, client: "RotatingClient"): try: interval_str = os.getenv("OAUTH_REFRESH_INTERVAL", "600") self._interval = int(interval_str) except ValueError: - lib_logger.warning(f"Invalid OAUTH_REFRESH_INTERVAL '{interval_str}'. Falling back to 600s.") + lib_logger.warning( + f"Invalid OAUTH_REFRESH_INTERVAL '{interval_str}'. Falling back to 600s." + ) self._interval = 600 self._client = client self._task: Optional[asyncio.Task] = None + self._initialized = False def start(self): """Starts the background refresh task.""" if self._task is None: self._task = asyncio.create_task(self._run()) - lib_logger.info(f"Background token refresher started. Check interval: {self._interval} seconds.") + lib_logger.info( + f"Background token refresher started. Check interval: {self._interval} seconds." + ) # [NEW] Log if custom interval is set async def stop(self): @@ -42,23 +49,107 @@ async def stop(self): pass lib_logger.info("Background token refresher stopped.") + async def _initialize_credentials(self): + """ + Initialize all providers by loading credentials and persisted tier data. + Called once before the main refresh loop starts. + """ + if self._initialized: + return + + api_summary = {} # provider -> count + oauth_summary = {} # provider -> {"count": N, "tiers": {tier: count}} + + all_credentials = self._client.all_credentials + oauth_providers = self._client.oauth_providers + + for provider, credentials in all_credentials.items(): + if not credentials: + continue + + provider_plugin = self._client._get_provider_instance(provider) + + # Call initialize_credentials if provider supports it + if provider_plugin and hasattr(provider_plugin, "initialize_credentials"): + try: + await provider_plugin.initialize_credentials(credentials) + except Exception as e: + lib_logger.error( + f"Error initializing credentials for provider '{provider}': {e}" + ) + + # Build summary based on provider type + if provider in oauth_providers: + tier_breakdown = {} + if provider_plugin and hasattr( + provider_plugin, "get_credential_tier_name" + ): + for cred in credentials: + tier = provider_plugin.get_credential_tier_name(cred) + if tier: + tier_breakdown[tier] = tier_breakdown.get(tier, 0) + 1 + oauth_summary[provider] = { + "count": len(credentials), + "tiers": tier_breakdown, + } + else: + api_summary[provider] = len(credentials) + + # Log 3-line summary + total_providers = len(api_summary) + len(oauth_summary) + total_credentials = sum(api_summary.values()) + sum( + d["count"] for d in oauth_summary.values() + ) + + if total_providers > 0: + lib_logger.info( + f"Providers initialized: {total_providers} providers, {total_credentials} credentials" + ) + + # API providers line + if api_summary: + api_parts = [f"{p}:{c}" for p, c in sorted(api_summary.items())] + lib_logger.info(f" API: {', '.join(api_parts)}") + + # OAuth providers line with tier breakdown + if oauth_summary: + oauth_parts = [] + for provider, data in sorted(oauth_summary.items()): + if data["tiers"]: + tier_str = ", ".join( + f"{t}:{c}" for t, c in sorted(data["tiers"].items()) + ) + oauth_parts.append(f"{provider}:{data['count']} ({tier_str})") + else: + oauth_parts.append(f"{provider}:{data['count']}") + lib_logger.info(f" OAuth: {', '.join(oauth_parts)}") + + self._initialized = True + async def _run(self): """The main loop for the background task.""" + # Initialize credentials (load persisted tiers) before starting the refresh loop + await self._initialize_credentials() + while True: try: - #lib_logger.info("Running proactive token refresh check...") + # lib_logger.info("Running proactive token refresh check...") oauth_configs = self._client.get_oauth_credentials() for provider, paths in oauth_configs.items(): - provider_plugin = self._client._get_provider_instance(f"{provider}_oauth") - if provider_plugin and hasattr(provider_plugin, 'proactively_refresh'): + provider_plugin = self._client._get_provider_instance(provider) + if provider_plugin and hasattr( + provider_plugin, "proactively_refresh" + ): for path in paths: try: await provider_plugin.proactively_refresh(path) except Exception as e: - lib_logger.error(f"Error during proactive refresh for '{path}': {e}") + lib_logger.error( + f"Error during proactive refresh for '{path}': {e}" + ) await asyncio.sleep(self._interval) except asyncio.CancelledError: break except Exception as e: - lib_logger.error(f"Unexpected error in background refresher loop: {e}") \ No newline at end of file + lib_logger.error(f"Unexpected error in background refresher loop: {e}") diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index cf1bb1cf..befa39ed 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -447,12 +447,23 @@ def _get_provider_instance(self, provider_name: str): Args: provider_name: The name of the provider to get an instance for. + For OAuth providers, this may include "_oauth" suffix + (e.g., "antigravity_oauth"), but credentials are stored + under the base name (e.g., "antigravity"). Returns: Provider instance if credentials exist, None otherwise. """ + # For OAuth providers, credentials are stored under base name (without _oauth suffix) + # e.g., "antigravity_oauth" plugin → credentials under "antigravity" + credential_key = provider_name + if provider_name.endswith("_oauth"): + base_name = provider_name[:-6] # Remove "_oauth" + if base_name in self.oauth_providers: + credential_key = base_name + # Only initialize providers for which we have credentials - if provider_name not in self.all_credentials: + if credential_key not in self.all_credentials: lib_logger.debug( f"Skipping provider '{provider_name}' initialization: no credentials configured" ) @@ -824,13 +835,20 @@ async def _execute_with_retry( f"Request will likely fail." ) - # Build priority map for usage_manager + # Build priority map and tier names map for usage_manager + credential_tier_names = None if provider_plugin and hasattr(provider_plugin, "get_credential_priority"): credential_priorities = {} + credential_tier_names = {} for cred in credentials_for_provider: priority = provider_plugin.get_credential_priority(cred) if priority is not None: credential_priorities[cred] = priority + # Also get tier name for logging + if hasattr(provider_plugin, "get_credential_tier_name"): + tier_name = provider_plugin.get_credential_tier_name(cred) + if tier_name: + credential_tier_names[cred] = tier_name if credential_priorities: lib_logger.debug( @@ -883,6 +901,7 @@ async def _execute_with_retry( deadline=deadline, max_concurrent=max_concurrent, credential_priorities=credential_priorities, + credential_tier_names=credential_tier_names, ) key_acquired = True tried_creds.add(current_cred) @@ -1371,13 +1390,20 @@ async def _streaming_acompletion_with_retry( f"Request will likely fail." ) - # Build priority map for usage_manager + # Build priority map and tier names map for usage_manager + credential_tier_names = None if provider_plugin and hasattr(provider_plugin, "get_credential_priority"): credential_priorities = {} + credential_tier_names = {} for cred in credentials_for_provider: priority = provider_plugin.get_credential_priority(cred) if priority is not None: credential_priorities[cred] = priority + # Also get tier name for logging + if hasattr(provider_plugin, "get_credential_tier_name"): + tier_name = provider_plugin.get_credential_tier_name(cred) + if tier_name: + credential_tier_names[cred] = tier_name if credential_priorities: lib_logger.debug( @@ -1433,6 +1459,7 @@ async def _streaming_acompletion_with_retry( deadline=deadline, max_concurrent=max_concurrent, credential_priorities=credential_priorities, + credential_tier_names=credential_tier_names, ) key_acquired = True tried_creds.add(current_cred) diff --git a/src/rotator_library/model_definitions.py b/src/rotator_library/model_definitions.py index 12219bcd..cb2aabf6 100644 --- a/src/rotator_library/model_definitions.py +++ b/src/rotator_library/model_definitions.py @@ -24,10 +24,23 @@ class ModelDefinitions: - IFLOW_MODELS='{"glm-4.6": {}}' - dict format, uses "glm-4.6" as both name and ID - IFLOW_MODELS='{"custom-name": {"id": "actual-id"}}' - dict format with custom ID - IFLOW_MODELS='{"model": {"id": "id", "options": {"temperature": 0.7}}}' - with options + + This class is a singleton - instantiated once and shared across all providers. """ + _instance: Optional["ModelDefinitions"] = None + _initialized: bool = False + + def __new__(cls, config_path: Optional[str] = None): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + def __init__(self, config_path: Optional[str] = None): - """Initialize model definitions loader.""" + """Initialize model definitions loader (only runs once due to singleton).""" + if ModelDefinitions._initialized: + return + ModelDefinitions._initialized = True self.config_path = config_path self.definitions = {} self._load_definitions() @@ -49,7 +62,11 @@ def _load_definitions(self): # Handle array format: ["model-1", "model-2", "model-3"] elif isinstance(models_json, list): # Convert array to dict format with empty definitions - models_dict = {model_name: {} for model_name in models_json if isinstance(model_name, str)} + models_dict = { + model_name: {} + for model_name in models_json + if isinstance(model_name, str) + } self.definitions[provider_name] = models_dict lib_logger.info( f"Loaded {len(models_dict)} models for provider: {provider_name} (array format)" diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index b8226a8a..7ed85f4b 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -595,6 +595,11 @@ def get_credential_priority(self, credential: str) -> Optional[int]: Priority level (1-10) or None if tier not yet discovered """ tier = self.project_tier_cache.get(credential) + + # Lazy load from file if not in cache + if not tier: + tier = self._load_tier_from_file(credential) + if not tier: return None # Not yet discovered @@ -609,6 +614,60 @@ def get_credential_priority(self, credential: str) -> Optional[int]: # Legacy and unknown get even lower return 10 + def _load_tier_from_file(self, credential_path: str) -> Optional[str]: + """ + Load tier from credential file's _proxy_metadata and cache it. + + This is used as a fallback when the tier isn't in the memory cache, + typically on first access before initialize_credentials() has run. + + Args: + credential_path: Path to the credential file + + Returns: + Tier string if found, None otherwise + """ + # Skip env:// paths (environment-based credentials) + if self._parse_env_credential_path(credential_path) is not None: + return None + + try: + with open(credential_path, "r") as f: + creds = json.load(f) + + metadata = creds.get("_proxy_metadata", {}) + tier = metadata.get("tier") + project_id = metadata.get("project_id") + + if tier: + self.project_tier_cache[credential_path] = tier + lib_logger.debug( + f"Lazy-loaded tier '{tier}' for credential: {Path(credential_path).name}" + ) + + if project_id and credential_path not in self.project_id_cache: + self.project_id_cache[credential_path] = project_id + + return tier + except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: + lib_logger.debug(f"Could not lazy-load tier from {credential_path}: {e}") + return None + + def get_credential_tier_name(self, credential: str) -> Optional[str]: + """ + Returns the human-readable tier name for a credential. + + Args: + credential: The credential path + + Returns: + Tier name string (e.g., "free-tier") or None if unknown + """ + tier = self.project_tier_cache.get(credential) + if not tier: + tier = self._load_tier_from_file(credential) + return tier + def get_model_tier_requirement(self, model: str) -> Optional[int]: """ Returns the minimum priority tier required for a model. @@ -622,6 +681,72 @@ def get_model_tier_requirement(self, model: str) -> Optional[int]: """ return None + async def initialize_credentials(self, credential_paths: List[str]) -> None: + """ + Load persisted tier information from credential files at startup. + + This ensures all credential priorities are known before any API calls, + preventing unknown credentials from getting priority 999. + """ + await self._load_persisted_tiers(credential_paths) + + async def _load_persisted_tiers( + self, credential_paths: List[str] + ) -> Dict[str, str]: + """ + Load persisted tier information from credential files into memory cache. + + Args: + credential_paths: List of credential file paths + + Returns: + Dict mapping credential path to tier name for logging purposes + """ + loaded = {} + for path in credential_paths: + # Skip env:// paths (environment-based credentials) + if self._parse_env_credential_path(path) is not None: + continue + + # Skip if already in cache + if path in self.project_tier_cache: + continue + + try: + with open(path, "r") as f: + creds = json.load(f) + + metadata = creds.get("_proxy_metadata", {}) + tier = metadata.get("tier") + project_id = metadata.get("project_id") + + if tier: + self.project_tier_cache[path] = tier + loaded[path] = tier + lib_logger.debug( + f"Loaded persisted tier '{tier}' for credential: {Path(path).name}" + ) + + if project_id: + self.project_id_cache[path] = project_id + + except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: + lib_logger.debug(f"Could not load persisted tier from {path}: {e}") + + if loaded: + # Log summary at debug level + tier_counts: Dict[str, int] = {} + for tier in loaded.values(): + tier_counts[tier] = tier_counts.get(tier, 0) + 1 + lib_logger.debug( + f"Antigravity: Loaded {len(loaded)} credential tiers from disk: " + + ", ".join( + f"{tier}={count}" for tier, count in sorted(tier_counts.items()) + ) + ) + + return loaded + # ========================================================================= # MODEL UTILITIES # ========================================================================= diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 259fb831..e4109ef9 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -19,13 +19,15 @@ import uuid from datetime import datetime -lib_logger = logging.getLogger('rotator_library') +lib_logger = logging.getLogger("rotator_library") LOGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "logs" GEMINI_CLI_LOGS_DIR = LOGS_DIR / "gemini_cli_logs" + class _GeminiCliFileLogger: """A simple file logger for a single Gemini CLI transaction.""" + def __init__(self, model_name: str, enabled: bool = True): self.enabled = enabled if not self.enabled: @@ -34,8 +36,10 @@ def __init__(self, model_name: str, enabled: bool = True): timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") request_id = str(uuid.uuid4()) # Sanitize model name for directory - safe_model_name = model_name.replace('/', '_').replace(':', '_') - self.log_dir = GEMINI_CLI_LOGS_DIR / f"{timestamp}_{safe_model_name}_{request_id}" + safe_model_name = model_name.replace("/", "_").replace(":", "_") + self.log_dir = ( + GEMINI_CLI_LOGS_DIR / f"{timestamp}_{safe_model_name}_{request_id}" + ) try: self.log_dir.mkdir(parents=True, exist_ok=True) except Exception as e: @@ -44,25 +48,32 @@ def __init__(self, model_name: str, enabled: bool = True): def log_request(self, payload: Dict[str, Any]): """Logs the request payload sent to Gemini.""" - if not self.enabled: return + if not self.enabled: + return try: - with open(self.log_dir / "request_payload.json", "w", encoding="utf-8") as f: + with open( + self.log_dir / "request_payload.json", "w", encoding="utf-8" + ) as f: json.dump(payload, f, indent=2, ensure_ascii=False) except Exception as e: lib_logger.error(f"_GeminiCliFileLogger: Failed to write request: {e}") def log_response_chunk(self, chunk: str): """Logs a raw chunk from the Gemini response stream.""" - if not self.enabled: return + if not self.enabled: + return try: with open(self.log_dir / "response_stream.log", "a", encoding="utf-8") as f: f.write(chunk + "\n") except Exception as e: - lib_logger.error(f"_GeminiCliFileLogger: Failed to write response chunk: {e}") + lib_logger.error( + f"_GeminiCliFileLogger: Failed to write response chunk: {e}" + ) def log_error(self, error_message: str): """Logs an error message.""" - if not self.enabled: return + if not self.enabled: + return try: with open(self.log_dir / "error.log", "a", encoding="utf-8") as f: f.write(f"[{datetime.utcnow().isoformat()}] {error_message}\n") @@ -71,12 +82,16 @@ def log_error(self, error_message: str): def log_final_response(self, response_data: Dict[str, Any]): """Logs the final, reassembled response.""" - if not self.enabled: return + if not self.enabled: + return try: with open(self.log_dir / "final_response.json", "w", encoding="utf-8") as f: json.dump(response_data, f, indent=2, ensure_ascii=False) except Exception as e: - lib_logger.error(f"_GeminiCliFileLogger: Failed to write final response: {e}") + lib_logger.error( + f"_GeminiCliFileLogger: Failed to write final response: {e}" + ) + CODE_ASSIST_ENDPOINT = "https://cloudcode-pa.googleapis.com/v1internal" @@ -84,11 +99,13 @@ def log_final_response(self, response_data: Dict[str, Any]): "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite", - "gemini-3-pro-preview" + "gemini-3-pro-preview", ] # Cache directory for Gemini CLI -CACHE_DIR = Path(__file__).resolve().parent.parent.parent.parent / "cache" / "gemini_cli" +CACHE_DIR = ( + Path(__file__).resolve().parent.parent.parent.parent / "cache" / "gemini_cli" +) GEMINI3_SIGNATURE_CACHE_FILE = CACHE_DIR / "gemini3_signatures.json" # Gemini 3 tool fix system instruction (prevents hallucination) @@ -172,36 +189,49 @@ class GeminiCliProvider(GeminiAuthBase, ProviderInterface): def __init__(self): super().__init__() self.model_definitions = ModelDefinitions() - self.project_id_cache: Dict[str, str] = {} # Cache project ID per credential path - self.project_tier_cache: Dict[str, str] = {} # Cache project tier per credential path - + self.project_id_cache: Dict[ + str, str + ] = {} # Cache project ID per credential path + self.project_tier_cache: Dict[ + str, str + ] = {} # Cache project tier per credential path + # Gemini 3 configuration from environment memory_ttl = _env_int("GEMINI_CLI_SIGNATURE_CACHE_TTL", 3600) disk_ttl = _env_int("GEMINI_CLI_SIGNATURE_DISK_TTL", 86400) - + # Initialize signature cache for Gemini 3 thoughtSignatures self._signature_cache = ProviderCache( - GEMINI3_SIGNATURE_CACHE_FILE, memory_ttl, disk_ttl, - env_prefix="GEMINI_CLI_SIGNATURE" + GEMINI3_SIGNATURE_CACHE_FILE, + memory_ttl, + disk_ttl, + env_prefix="GEMINI_CLI_SIGNATURE", ) - + # Gemini 3 feature flags - self._preserve_signatures_in_client = _env_bool("GEMINI_CLI_PRESERVE_THOUGHT_SIGNATURES", True) - self._enable_signature_cache = _env_bool("GEMINI_CLI_ENABLE_SIGNATURE_CACHE", True) + self._preserve_signatures_in_client = _env_bool( + "GEMINI_CLI_PRESERVE_THOUGHT_SIGNATURES", True + ) + self._enable_signature_cache = _env_bool( + "GEMINI_CLI_ENABLE_SIGNATURE_CACHE", True + ) self._enable_gemini3_tool_fix = _env_bool("GEMINI_CLI_GEMINI3_TOOL_FIX", True) - self._gemini3_enforce_strict_schema = _env_bool("GEMINI_CLI_GEMINI3_STRICT_SCHEMA", True) - + self._gemini3_enforce_strict_schema = _env_bool( + "GEMINI_CLI_GEMINI3_STRICT_SCHEMA", True + ) + # Gemini 3 tool fix configuration - self._gemini3_tool_prefix = os.getenv("GEMINI_CLI_GEMINI3_TOOL_PREFIX", "gemini3_") + self._gemini3_tool_prefix = os.getenv( + "GEMINI_CLI_GEMINI3_TOOL_PREFIX", "gemini3_" + ) self._gemini3_description_prompt = os.getenv( "GEMINI_CLI_GEMINI3_DESCRIPTION_PROMPT", - "\n\n⚠️ STRICT PARAMETERS (use EXACTLY as shown): {params}. Do NOT use parameters from your training data - use ONLY these parameter names." + "\n\n⚠️ STRICT PARAMETERS (use EXACTLY as shown): {params}. Do NOT use parameters from your training data - use ONLY these parameter names.", ) self._gemini3_system_instruction = os.getenv( - "GEMINI_CLI_GEMINI3_SYSTEM_INSTRUCTION", - DEFAULT_GEMINI3_SYSTEM_INSTRUCTION + "GEMINI_CLI_GEMINI3_SYSTEM_INSTRUCTION", DEFAULT_GEMINI3_SYSTEM_INSTRUCTION ) - + lib_logger.debug( f"GeminiCli config: signatures_in_client={self._preserve_signatures_in_client}, " f"cache={self._enable_signature_cache}, gemini3_fix={self._enable_gemini3_tool_fix}, " @@ -211,75 +241,200 @@ def __init__(self): # ========================================================================= # CREDENTIAL PRIORITIZATION # ========================================================================= - + def get_credential_priority(self, credential: str) -> Optional[int]: """ Returns priority based on Gemini tier. Paid tiers: priority 1 (highest) Free/Legacy tiers: priority 2 Unknown: priority 10 (lowest) - + Args: credential: The credential path - + Returns: Priority level (1-10) or None if tier not yet discovered """ tier = self.project_tier_cache.get(credential) + + # Lazy load from file if not in cache + if not tier: + tier = self._load_tier_from_file(credential) + if not tier: return None # Not yet discovered - + # Paid tiers get highest priority - if tier not in ['free-tier', 'legacy-tier', 'unknown']: + if tier not in ["free-tier", "legacy-tier", "unknown"]: return 1 - + # Free tier gets lower priority - if tier == 'free-tier': + if tier == "free-tier": return 2 - + # Legacy and unknown get even lower return 10 - + + def _load_tier_from_file(self, credential_path: str) -> Optional[str]: + """ + Load tier from credential file's _proxy_metadata and cache it. + + This is used as a fallback when the tier isn't in the memory cache, + typically on first access before initialize_credentials() has run. + + Args: + credential_path: Path to the credential file + + Returns: + Tier string if found, None otherwise + """ + # Skip env:// paths (environment-based credentials) + if self._parse_env_credential_path(credential_path) is not None: + return None + + try: + with open(credential_path, "r") as f: + creds = json.load(f) + + metadata = creds.get("_proxy_metadata", {}) + tier = metadata.get("tier") + project_id = metadata.get("project_id") + + if tier: + self.project_tier_cache[credential_path] = tier + lib_logger.debug( + f"Lazy-loaded tier '{tier}' for credential: {Path(credential_path).name}" + ) + + if project_id and credential_path not in self.project_id_cache: + self.project_id_cache[credential_path] = project_id + + return tier + except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: + lib_logger.debug(f"Could not lazy-load tier from {credential_path}: {e}") + return None + + def get_credential_tier_name(self, credential: str) -> Optional[str]: + """ + Returns the human-readable tier name for a credential. + + Args: + credential: The credential path + + Returns: + Tier name string (e.g., "free-tier") or None if unknown + """ + tier = self.project_tier_cache.get(credential) + if not tier: + tier = self._load_tier_from_file(credential) + return tier + def get_model_tier_requirement(self, model: str) -> Optional[int]: """ Returns the minimum priority tier required for a model. Gemini 3 requires paid tier (priority 1). - + Args: model: The model name (with or without provider prefix) - + Returns: Minimum required priority level or None if no restrictions """ - model_name = model.split('/')[-1].replace(':thinking', '') - + model_name = model.split("/")[-1].replace(":thinking", "") + # Gemini 3 requires paid tier if model_name.startswith("gemini-3-"): return 1 # Only priority 1 (paid) credentials - + return None # All other models have no restrictions + async def initialize_credentials(self, credential_paths: List[str]) -> None: + """ + Load persisted tier information from credential files at startup. + This ensures all credential priorities are known before any API calls, + preventing unknown credentials from getting priority 999. + """ + await self._load_persisted_tiers(credential_paths) + + async def _load_persisted_tiers( + self, credential_paths: List[str] + ) -> Dict[str, str]: + """ + Load persisted tier information from credential files into memory cache. + + Args: + credential_paths: List of credential file paths + + Returns: + Dict mapping credential path to tier name for logging purposes + """ + loaded = {} + for path in credential_paths: + # Skip env:// paths (environment-based credentials) + if self._parse_env_credential_path(path) is not None: + continue + + # Skip if already in cache + if path in self.project_tier_cache: + continue + + try: + with open(path, "r") as f: + creds = json.load(f) + + metadata = creds.get("_proxy_metadata", {}) + tier = metadata.get("tier") + project_id = metadata.get("project_id") + + if tier: + self.project_tier_cache[path] = tier + loaded[path] = tier + lib_logger.debug( + f"Loaded persisted tier '{tier}' for credential: {Path(path).name}" + ) + + if project_id: + self.project_id_cache[path] = project_id + + except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: + lib_logger.debug(f"Could not load persisted tier from {path}: {e}") + + if loaded: + # Log summary at debug level + tier_counts: Dict[str, int] = {} + for tier in loaded.values(): + tier_counts[tier] = tier_counts.get(tier, 0) + 1 + lib_logger.debug( + f"GeminiCli: Loaded {len(loaded)} credential tiers from disk: " + + ", ".join( + f"{tier}={count}" for tier, count in sorted(tier_counts.items()) + ) + ) + + return loaded # ========================================================================= # MODEL UTILITIES # ========================================================================= - + def _is_gemini_3(self, model: str) -> bool: """Check if model is Gemini 3 (requires special handling).""" - model_name = model.split('/')[-1].replace(':thinking', '') + model_name = model.split("/")[-1].replace(":thinking", "") return model_name.startswith("gemini-3-") - + def _strip_gemini3_prefix(self, name: str) -> str: """Strip the Gemini 3 namespace prefix from a tool name.""" if name and name.startswith(self._gemini3_tool_prefix): - return name[len(self._gemini3_tool_prefix):] + return name[len(self._gemini3_tool_prefix) :] return name - async def _discover_project_id(self, credential_path: str, access_token: str, litellm_params: Dict[str, Any]) -> str: + async def _discover_project_id( + self, credential_path: str, access_token: str, litellm_params: Dict[str, Any] + ) -> str: """ Discovers the Google Cloud Project ID, with caching and onboarding for new accounts. - + This follows the official Gemini CLI discovery flow: 1. Check in-memory cache 2. Check configured project_id override (litellm_params or env var) @@ -293,7 +448,9 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li - PAID tier: pass cloudaicompanionProject=configured_project_id 6. Fallback to GCP Resource Manager project listing """ - lib_logger.debug(f"Starting project discovery for credential: {credential_path}") + lib_logger.debug( + f"Starting project discovery for credential: {credential_path}" + ) # Check in-memory cache first if credential_path in self.project_id_cache: @@ -305,7 +462,9 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li # This is REQUIRED for paid tier users per the official CLI behavior configured_project_id = litellm_params.get("project_id") if configured_project_id: - lib_logger.debug(f"Found configured project_id override: {configured_project_id}") + lib_logger.debug( + f"Found configured project_id override: {configured_project_id}" + ) # Load credentials from file to check for persisted project_id and tier # Skip for env:// paths (environment-based credentials don't persist to files) @@ -313,35 +472,44 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li if credential_index is None: # Only try to load from file if it's not an env:// path try: - with open(credential_path, 'r') as f: + with open(credential_path, "r") as f: creds = json.load(f) - + metadata = creds.get("_proxy_metadata", {}) persisted_project_id = metadata.get("project_id") persisted_tier = metadata.get("tier") - + if persisted_project_id: - lib_logger.info(f"Loaded persisted project ID from credential file: {persisted_project_id}") + lib_logger.info( + f"Loaded persisted project ID from credential file: {persisted_project_id}" + ) self.project_id_cache[credential_path] = persisted_project_id - + # Also load tier if available if persisted_tier: self.project_tier_cache[credential_path] = persisted_tier lib_logger.debug(f"Loaded persisted tier: {persisted_tier}") - + return persisted_project_id except (FileNotFoundError, json.JSONDecodeError, KeyError) as e: lib_logger.debug(f"Could not load persisted project ID from file: {e}") - lib_logger.debug("No cached or configured project ID found, initiating discovery...") - headers = {'Authorization': f'Bearer {access_token}', 'Content-Type': 'application/json'} + lib_logger.debug( + "No cached or configured project ID found, initiating discovery..." + ) + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + } discovered_project_id = None discovered_tier = None async with httpx.AsyncClient() as client: # 1. Try discovery endpoint with loadCodeAssist - lib_logger.debug("Attempting project discovery via Code Assist loadCodeAssist endpoint...") + lib_logger.debug( + "Attempting project discovery via Code Assist loadCodeAssist endpoint..." + ) try: # Build metadata - include duetProject only if we have a configured project core_client_metadata = { @@ -351,53 +519,65 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li } if configured_project_id: core_client_metadata["duetProject"] = configured_project_id - + # Build load request - pass configured_project_id if available, otherwise None load_request = { "cloudaicompanionProject": configured_project_id, # Can be None "metadata": core_client_metadata, } - - lib_logger.debug(f"Sending loadCodeAssist request with cloudaicompanionProject={configured_project_id}") - response = await client.post(f"{CODE_ASSIST_ENDPOINT}:loadCodeAssist", headers=headers, json=load_request, timeout=20) + + lib_logger.debug( + f"Sending loadCodeAssist request with cloudaicompanionProject={configured_project_id}" + ) + response = await client.post( + f"{CODE_ASSIST_ENDPOINT}:loadCodeAssist", + headers=headers, + json=load_request, + timeout=20, + ) response.raise_for_status() data = response.json() # Log full response for debugging - lib_logger.debug(f"loadCodeAssist full response keys: {list(data.keys())}") + lib_logger.debug( + f"loadCodeAssist full response keys: {list(data.keys())}" + ) # Extract and log ALL tier information for debugging - allowed_tiers = data.get('allowedTiers', []) - current_tier = data.get('currentTier') - + allowed_tiers = data.get("allowedTiers", []) + current_tier = data.get("currentTier") + lib_logger.debug(f"=== Tier Information ===") lib_logger.debug(f"currentTier: {current_tier}") lib_logger.debug(f"allowedTiers count: {len(allowed_tiers)}") for i, tier in enumerate(allowed_tiers): - tier_id = tier.get('id', 'unknown') - is_default = tier.get('isDefault', False) - user_defined = tier.get('userDefinedCloudaicompanionProject', False) - lib_logger.debug(f" Tier {i+1}: id={tier_id}, isDefault={is_default}, userDefinedProject={user_defined}") + tier_id = tier.get("id", "unknown") + is_default = tier.get("isDefault", False) + user_defined = tier.get("userDefinedCloudaicompanionProject", False) + lib_logger.debug( + f" Tier {i + 1}: id={tier_id}, isDefault={is_default}, userDefinedProject={user_defined}" + ) lib_logger.debug(f"========================") # Determine the current tier ID current_tier_id = None if current_tier: - current_tier_id = current_tier.get('id') + current_tier_id = current_tier.get("id") lib_logger.debug(f"User has currentTier: {current_tier_id}") # Check if user is already known to server (has currentTier) if current_tier_id: # User is already onboarded - check for project from server - server_project = data.get('cloudaicompanionProject') - + server_project = data.get("cloudaicompanionProject") + # Check if this tier requires user-defined project (paid tiers) requires_user_project = any( - t.get('id') == current_tier_id and t.get('userDefinedCloudaicompanionProject', False) + t.get("id") == current_tier_id + and t.get("userDefinedCloudaicompanionProject", False) for t in allowed_tiers ) - is_free_tier = current_tier_id == 'free-tier' - + is_free_tier = current_tier_id == "free-tier" + if server_project: # Server returned a project - use it (server wins) # This is the normal case for FREE tier users @@ -407,11 +587,15 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li # No server project but we have configured one - use it # This is the PAID TIER case where server doesn't return a project project_id = configured_project_id - lib_logger.debug(f"No server project, using configured: {project_id}") + lib_logger.debug( + f"No server project, using configured: {project_id}" + ) elif is_free_tier: # Free tier user without server project - this shouldn't happen normally # but let's not fail, just proceed to onboarding - lib_logger.debug("Free tier user with currentTier but no project - will try onboarding") + lib_logger.debug( + "Free tier user with currentTier but no project - will try onboarding" + ) project_id = None elif requires_user_project: # Paid tier requires a project ID to be set @@ -421,7 +605,9 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li ) else: # Unknown tier without project - proceed carefully - lib_logger.warning(f"Tier '{current_tier_id}' has no project and none configured - will try onboarding") + lib_logger.warning( + f"Tier '{current_tier_id}' has no project and none configured - will try onboarding" + ) project_id = None if project_id: @@ -430,54 +616,70 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li discovered_tier = current_tier_id # Log appropriately based on tier - is_paid = current_tier_id and current_tier_id not in ['free-tier', 'legacy-tier', 'unknown'] + is_paid = current_tier_id and current_tier_id not in [ + "free-tier", + "legacy-tier", + "unknown", + ] if is_paid: - lib_logger.info(f"Using Gemini paid tier '{current_tier_id}' with project: {project_id}") + lib_logger.info( + f"Using Gemini paid tier '{current_tier_id}' with project: {project_id}" + ) else: - lib_logger.info(f"Discovered Gemini project ID via loadCodeAssist: {project_id}") + lib_logger.info( + f"Discovered Gemini project ID via loadCodeAssist: {project_id}" + ) self.project_id_cache[credential_path] = project_id discovered_project_id = project_id - + # Persist to credential file - await self._persist_project_metadata(credential_path, project_id, discovered_tier) - + await self._persist_project_metadata( + credential_path, project_id, discovered_tier + ) + return project_id - + # 2. User needs onboarding - no currentTier - lib_logger.info("No existing Gemini session found (no currentTier), attempting to onboard user...") - + lib_logger.info( + "No existing Gemini session found (no currentTier), attempting to onboard user..." + ) + # Determine which tier to onboard with onboard_tier = None for tier in allowed_tiers: - if tier.get('isDefault'): + if tier.get("isDefault"): onboard_tier = tier break - + # Fallback to LEGACY tier if no default (requires user project) if not onboard_tier and allowed_tiers: # Look for legacy-tier as fallback for tier in allowed_tiers: - if tier.get('id') == 'legacy-tier': + if tier.get("id") == "legacy-tier": onboard_tier = tier break # If still no tier, use first available if not onboard_tier: onboard_tier = allowed_tiers[0] - + if not onboard_tier: raise ValueError("No onboarding tiers available from server") - - tier_id = onboard_tier.get('id', 'free-tier') - requires_user_project = onboard_tier.get('userDefinedCloudaicompanionProject', False) - - lib_logger.debug(f"Onboarding with tier: {tier_id}, requiresUserProject: {requires_user_project}") - + + tier_id = onboard_tier.get("id", "free-tier") + requires_user_project = onboard_tier.get( + "userDefinedCloudaicompanionProject", False + ) + + lib_logger.debug( + f"Onboarding with tier: {tier_id}, requiresUserProject: {requires_user_project}" + ) + # Build onboard request based on tier type (following official CLI logic) # FREE tier: cloudaicompanionProject = None (server-managed) # PAID tier: cloudaicompanionProject = configured_project_id (user must provide) - is_free_tier = tier_id == 'free-tier' - + is_free_tier = tier_id == "free-tier" + if is_free_tier: # Free tier uses server-managed project onboard_request = { @@ -485,7 +687,9 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li "cloudaicompanionProject": None, # Server will create/manage "metadata": core_client_metadata, } - lib_logger.debug("Free tier onboarding: using server-managed project") + lib_logger.debug( + "Free tier onboarding: using server-managed project" + ) else: # Paid/legacy tier requires user-provided project if not configured_project_id and requires_user_project: @@ -499,51 +703,85 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li "metadata": { **core_client_metadata, "duetProject": configured_project_id, - } if configured_project_id else core_client_metadata, + } + if configured_project_id + else core_client_metadata, } - lib_logger.debug(f"Paid tier onboarding: using project {configured_project_id}") + lib_logger.debug( + f"Paid tier onboarding: using project {configured_project_id}" + ) lib_logger.debug("Initiating onboardUser request...") - lro_response = await client.post(f"{CODE_ASSIST_ENDPOINT}:onboardUser", headers=headers, json=onboard_request, timeout=30) + lro_response = await client.post( + f"{CODE_ASSIST_ENDPOINT}:onboardUser", + headers=headers, + json=onboard_request, + timeout=30, + ) lro_response.raise_for_status() lro_data = lro_response.json() - lib_logger.debug(f"Initial onboarding response: done={lro_data.get('done')}") + lib_logger.debug( + f"Initial onboarding response: done={lro_data.get('done')}" + ) for i in range(150): # Poll for up to 5 minutes (150 × 2s) - if lro_data.get('done'): - lib_logger.debug(f"Onboarding completed after {i} polling attempts") + if lro_data.get("done"): + lib_logger.debug( + f"Onboarding completed after {i} polling attempts" + ) break await asyncio.sleep(2) if (i + 1) % 15 == 0: # Log every 30 seconds - lib_logger.info(f"Still waiting for onboarding completion... ({(i+1)*2}s elapsed)") - lib_logger.debug(f"Polling onboarding status... (Attempt {i+1}/150)") - lro_response = await client.post(f"{CODE_ASSIST_ENDPOINT}:onboardUser", headers=headers, json=onboard_request, timeout=30) + lib_logger.info( + f"Still waiting for onboarding completion... ({(i + 1) * 2}s elapsed)" + ) + lib_logger.debug( + f"Polling onboarding status... (Attempt {i + 1}/150)" + ) + lro_response = await client.post( + f"{CODE_ASSIST_ENDPOINT}:onboardUser", + headers=headers, + json=onboard_request, + timeout=30, + ) lro_response.raise_for_status() lro_data = lro_response.json() - if not lro_data.get('done'): + if not lro_data.get("done"): lib_logger.error("Onboarding process timed out after 5 minutes") - raise ValueError("Onboarding process timed out after 5 minutes. Please try again or contact support.") + raise ValueError( + "Onboarding process timed out after 5 minutes. Please try again or contact support." + ) # Extract project ID from LRO response # Note: onboardUser returns response.cloudaicompanionProject as an object with .id - lro_response_data = lro_data.get('response', {}) - lro_project_obj = lro_response_data.get('cloudaicompanionProject', {}) - project_id = lro_project_obj.get('id') if isinstance(lro_project_obj, dict) else None - + lro_response_data = lro_data.get("response", {}) + lro_project_obj = lro_response_data.get("cloudaicompanionProject", {}) + project_id = ( + lro_project_obj.get("id") + if isinstance(lro_project_obj, dict) + else None + ) + # Fallback to configured project if LRO didn't return one if not project_id and configured_project_id: project_id = configured_project_id - lib_logger.debug(f"LRO didn't return project, using configured: {project_id}") - + lib_logger.debug( + f"LRO didn't return project, using configured: {project_id}" + ) + if not project_id: - lib_logger.error("Onboarding completed but no project ID in response and none configured") + lib_logger.error( + "Onboarding completed but no project ID in response and none configured" + ) raise ValueError( "Onboarding completed, but no project ID was returned. " "For paid tiers, set GEMINI_CLI_PROJECT_ID environment variable." ) - lib_logger.debug(f"Successfully extracted project ID from onboarding response: {project_id}") + lib_logger.debug( + f"Successfully extracted project ID from onboarding response: {project_id}" + ) # Cache tier info self.project_tier_cache[credential_path] = tier_id @@ -551,18 +789,24 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li lib_logger.debug(f"Cached tier information: {tier_id}") # Log concise message for paid projects - is_paid = tier_id and tier_id not in ['free-tier', 'legacy-tier'] + is_paid = tier_id and tier_id not in ["free-tier", "legacy-tier"] if is_paid: - lib_logger.info(f"Using Gemini paid tier '{tier_id}' with project: {project_id}") + lib_logger.info( + f"Using Gemini paid tier '{tier_id}' with project: {project_id}" + ) else: - lib_logger.info(f"Successfully onboarded user and discovered project ID: {project_id}") + lib_logger.info( + f"Successfully onboarded user and discovered project ID: {project_id}" + ) self.project_id_cache[credential_path] = project_id discovered_project_id = project_id - + # Persist to credential file - await self._persist_project_metadata(credential_path, project_id, discovered_tier) - + await self._persist_project_metadata( + credential_path, project_id, discovered_tier + ) + return project_id except httpx.HTTPStatusError as e: @@ -572,50 +816,86 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li except Exception: pass if e.response.status_code == 403: - lib_logger.error(f"Gemini Code Assist API access denied (403). Response: {error_body}") - lib_logger.error("Possible causes: 1) cloudaicompanion.googleapis.com API not enabled, 2) Wrong project ID for paid tier, 3) Account lacks permissions") + lib_logger.error( + f"Gemini Code Assist API access denied (403). Response: {error_body}" + ) + lib_logger.error( + "Possible causes: 1) cloudaicompanion.googleapis.com API not enabled, 2) Wrong project ID for paid tier, 3) Account lacks permissions" + ) elif e.response.status_code == 404: - lib_logger.warning(f"Gemini Code Assist endpoint not found (404). Falling back to project listing.") + lib_logger.warning( + f"Gemini Code Assist endpoint not found (404). Falling back to project listing." + ) elif e.response.status_code == 412: # Precondition Failed - often means wrong project for free tier onboarding - lib_logger.error(f"Precondition failed (412): {error_body}. This may mean the project ID is incompatible with the selected tier.") + lib_logger.error( + f"Precondition failed (412): {error_body}. This may mean the project ID is incompatible with the selected tier." + ) else: - lib_logger.warning(f"Gemini onboarding/discovery failed with status {e.response.status_code}: {error_body}. Falling back to project listing.") + lib_logger.warning( + f"Gemini onboarding/discovery failed with status {e.response.status_code}: {error_body}. Falling back to project listing." + ) except httpx.RequestError as e: - lib_logger.warning(f"Gemini onboarding/discovery network error: {e}. Falling back to project listing.") + lib_logger.warning( + f"Gemini onboarding/discovery network error: {e}. Falling back to project listing." + ) # 3. Fallback to listing all available GCP projects (last resort) - lib_logger.debug("Attempting to discover project via GCP Resource Manager API...") + lib_logger.debug( + "Attempting to discover project via GCP Resource Manager API..." + ) try: async with httpx.AsyncClient() as client: - lib_logger.debug("Querying Cloud Resource Manager for available projects...") - response = await client.get("https://cloudresourcemanager.googleapis.com/v1/projects", headers=headers, timeout=20) + lib_logger.debug( + "Querying Cloud Resource Manager for available projects..." + ) + response = await client.get( + "https://cloudresourcemanager.googleapis.com/v1/projects", + headers=headers, + timeout=20, + ) response.raise_for_status() - projects = response.json().get('projects', []) + projects = response.json().get("projects", []) lib_logger.debug(f"Found {len(projects)} total projects") - active_projects = [p for p in projects if p.get('lifecycleState') == 'ACTIVE'] + active_projects = [ + p for p in projects if p.get("lifecycleState") == "ACTIVE" + ] lib_logger.debug(f"Found {len(active_projects)} active projects") if not projects: - lib_logger.error("No GCP projects found for this account. Please create a project in Google Cloud Console.") + lib_logger.error( + "No GCP projects found for this account. Please create a project in Google Cloud Console." + ) elif not active_projects: - lib_logger.error("No active GCP projects found. Please activate a project in Google Cloud Console.") + lib_logger.error( + "No active GCP projects found. Please activate a project in Google Cloud Console." + ) else: - project_id = active_projects[0]['projectId'] - lib_logger.info(f"Discovered Gemini project ID from active projects list: {project_id}") - lib_logger.debug(f"Selected first active project: {project_id} (out of {len(active_projects)} active projects)") + project_id = active_projects[0]["projectId"] + lib_logger.info( + f"Discovered Gemini project ID from active projects list: {project_id}" + ) + lib_logger.debug( + f"Selected first active project: {project_id} (out of {len(active_projects)} active projects)" + ) self.project_id_cache[credential_path] = project_id discovered_project_id = project_id - + # [NEW] Persist to credential file (no tier info from resource manager) - await self._persist_project_metadata(credential_path, project_id, None) - + await self._persist_project_metadata( + credential_path, project_id, None + ) + return project_id except httpx.HTTPStatusError as e: if e.response.status_code == 403: - lib_logger.error("Failed to list GCP projects due to a 403 Forbidden error. The Cloud Resource Manager API may not be enabled, or your account lacks the 'resourcemanager.projects.list' permission.") + lib_logger.error( + "Failed to list GCP projects due to a 403 Forbidden error. The Cloud Resource Manager API may not be enabled, or your account lacks the 'resourcemanager.projects.list' permission." + ) else: - lib_logger.error(f"Failed to list GCP projects with status {e.response.status_code}: {e}") + lib_logger.error( + f"Failed to list GCP projects with status {e.response.status_code}: {e}" + ) except httpx.RequestError as e: lib_logger.error(f"Network error while listing GCP projects: {e}") @@ -626,20 +906,24 @@ async def _discover_project_id(self, credential_path: str, access_token: str, li " 3. Account lacks necessary permissions\n" "To manually specify a project, set GEMINI_CLI_PROJECT_ID in your .env file." ) - - async def _persist_project_metadata(self, credential_path: str, project_id: str, tier: Optional[str]): + + async def _persist_project_metadata( + self, credential_path: str, project_id: str, tier: Optional[str] + ): """Persists project ID and tier to the credential file for faster future startups.""" # Skip persistence for env:// paths (environment-based credentials) credential_index = self._parse_env_credential_path(credential_path) if credential_index is not None: - lib_logger.debug(f"Skipping project metadata persistence for env:// credential path: {credential_path}") + lib_logger.debug( + f"Skipping project metadata persistence for env:// credential path: {credential_path}" + ) return - + try: # Load current credentials - with open(credential_path, 'r') as f: + with open(credential_path, "r") as f: creds = json.load(f) - + # Update metadata if "_proxy_metadata" not in creds: creds["_proxy_metadata"] = {} @@ -647,33 +931,36 @@ async def _persist_project_metadata(self, credential_path: str, project_id: str, creds["_proxy_metadata"]["project_id"] = project_id if tier: creds["_proxy_metadata"]["tier"] = tier - + # Save back using the existing save method (handles atomic writes and permissions) await self._save_credentials(credential_path, creds) - - lib_logger.debug(f"Persisted project_id and tier to credential file: {credential_path}") + + lib_logger.debug( + f"Persisted project_id and tier to credential file: {credential_path}" + ) except Exception as e: - lib_logger.warning(f"Failed to persist project metadata to credential file: {e}") + lib_logger.warning( + f"Failed to persist project metadata to credential file: {e}" + ) # Non-fatal - just means slower startup next time - def _check_mixed_tier_warning(self): """Check if mixed free/paid tier credentials are loaded and emit warning.""" if not self.project_tier_cache: return # No tiers loaded yet - + tiers = set(self.project_tier_cache.values()) if len(tiers) <= 1: return # All same tier or only one credential - + # Define paid vs free tiers - free_tiers = {'free-tier', 'legacy-tier', 'unknown'} + free_tiers = {"free-tier", "legacy-tier", "unknown"} paid_tiers = tiers - free_tiers - + # Check if we have both free and paid has_free = bool(tiers & free_tiers) has_paid = bool(paid_tiers) - + if has_free and has_paid: lib_logger.warning( f"Mixed Gemini tier credentials detected! You have both free-tier and paid-tier " @@ -688,12 +975,12 @@ def _cli_preview_fallback_order(self, model: str) -> List[str]: """ Returns a list of model names to try in order for rate limit fallback. First model in list is the original model, subsequent models are fallback options. - + Since all fallbacks have been deprecated, this now only returns the base model. The fallback logic will check if there are actual fallbacks available. """ # Remove provider prefix if present - model_name = model.split('/')[-1].replace(':thinking', '') + model_name = model.split("/")[-1].replace(":thinking", "") # Define fallback chains for models with preview versions # All fallbacks have been deprecated, so only base models are returned @@ -706,10 +993,12 @@ def _cli_preview_fallback_order(self, model: str) -> List[str]: # Return fallback chain if available, otherwise just return the original model return fallback_chains.get(model_name, [model_name]) - def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: + def _transform_messages( + self, messages: List[Dict[str, Any]], model: str = "" + ) -> Tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]: """ Transform OpenAI messages to Gemini CLI format. - + Handles: - System instruction extraction - Multi-part content (text, images) @@ -720,14 +1009,14 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - system_instruction = None gemini_contents = [] is_gemini_3 = self._is_gemini_3(model) - + # Separate system prompt from other messages - if messages and messages[0].get('role') == 'system': - system_prompt_content = messages.pop(0).get('content', '') + if messages and messages[0].get("role") == "system": + system_prompt_content = messages.pop(0).get("content", "") if system_prompt_content: system_instruction = { "role": "user", - "parts": [{"text": system_prompt_content}] + "parts": [{"text": system_prompt_content}], } tool_call_id_to_name = {} @@ -735,18 +1024,22 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - if msg.get("role") == "assistant" and msg.get("tool_calls"): for tool_call in msg["tool_calls"]: if tool_call.get("type") == "function": - tool_call_id_to_name[tool_call["id"]] = tool_call["function"]["name"] + tool_call_id_to_name[tool_call["id"]] = tool_call["function"][ + "name" + ] # Process messages and consolidate consecutive tool responses # Per Gemini docs: parallel function responses must be in a single user message, # not interleaved as separate messages pending_tool_parts = [] # Accumulate tool responses - + for msg in messages: role = msg.get("role") content = msg.get("content") parts = [] - gemini_role = "model" if role == "assistant" else "user" # tool -> user in Gemini + gemini_role = ( + "model" if role == "assistant" else "user" + ) # tool -> user in Gemini # If we have pending tool parts and hit a non-tool message, flush them first if pending_tool_parts and role != "tool": @@ -773,16 +1066,22 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - # Parse: data:image/png;base64,iVBORw0KG... header, data = image_url.split(",", 1) mime_type = header.split(":")[1].split(";")[0] - parts.append({ - "inlineData": { - "mimeType": mime_type, - "data": data + parts.append( + { + "inlineData": { + "mimeType": mime_type, + "data": data, + } } - }) + ) except Exception as e: - lib_logger.warning(f"Failed to parse image data URL: {e}") + lib_logger.warning( + f"Failed to parse image data URL: {e}" + ) else: - lib_logger.warning(f"Non-data-URL images not supported: {image_url[:50]}...") + lib_logger.warning( + f"Non-data-URL images not supported: {image_url[:50]}..." + ) elif role == "assistant": if isinstance(content, str): @@ -794,25 +1093,27 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - for tool_call in msg["tool_calls"]: if tool_call.get("type") == "function": try: - args_dict = json.loads(tool_call["function"]["arguments"]) + args_dict = json.loads( + tool_call["function"]["arguments"] + ) except (json.JSONDecodeError, TypeError): args_dict = {} - + tool_id = tool_call.get("id", "") func_name = tool_call["function"]["name"] - + # Add prefix for Gemini 3 if is_gemini_3 and self._enable_gemini3_tool_fix: func_name = f"{self._gemini3_tool_prefix}{func_name}" - + func_part = { "functionCall": { "name": func_name, "args": args_dict, - "id": tool_id + "id": tool_id, } } - + # Add thoughtSignature for Gemini 3 # Per Gemini docs: Only the FIRST parallel function call gets a signature. # Subsequent parallel calls should NOT have a thoughtSignature field. @@ -820,17 +1121,21 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - sig = tool_call.get("thought_signature") if not sig and tool_id and self._enable_signature_cache: sig = self._signature_cache.retrieve(tool_id) - + if sig: func_part["thoughtSignature"] = sig elif first_func_in_msg: # Only add bypass to the first function call if no sig available - func_part["thoughtSignature"] = "skip_thought_signature_validator" - lib_logger.warning(f"Missing thoughtSignature for first func call {tool_id}, using bypass") + func_part["thoughtSignature"] = ( + "skip_thought_signature_validator" + ) + lib_logger.warning( + f"Missing thoughtSignature for first func call {tool_id}, using bypass" + ) # Subsequent parallel calls: no signature field at all - + first_func_in_msg = False - + parts.append(func_part) elif role == "tool": @@ -840,17 +1145,19 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - # Add prefix for Gemini 3 if is_gemini_3 and self._enable_gemini3_tool_fix: function_name = f"{self._gemini3_tool_prefix}{function_name}" - + # Wrap the tool response in a 'result' object response_content = {"result": content} # Accumulate tool responses - they'll be combined into one user message - pending_tool_parts.append({ - "functionResponse": { - "name": function_name, - "response": response_content, - "id": tool_call_id + pending_tool_parts.append( + { + "functionResponse": { + "name": function_name, + "response": response_content, + "id": tool_call_id, + } } - }) + ) # Don't add parts here - tool responses are handled via pending_tool_parts continue @@ -861,15 +1168,17 @@ def _transform_messages(self, messages: List[Dict[str, Any]], model: str = "") - if pending_tool_parts: gemini_contents.append({"role": "user", "parts": pending_tool_parts}) - if not gemini_contents or gemini_contents[0]['role'] != 'user': + if not gemini_contents or gemini_contents[0]["role"] != "user": gemini_contents.insert(0, {"role": "user", "parts": [{"text": ""}]}) return system_instruction, gemini_contents - def _handle_reasoning_parameters(self, payload: Dict[str, Any], model: str) -> Optional[Dict[str, Any]]: + def _handle_reasoning_parameters( + self, payload: Dict[str, Any], model: str + ) -> Optional[Dict[str, Any]]: """ Map reasoning_effort to thinking configuration. - + - Gemini 2.5: thinkingBudget (integer tokens) - Gemini 3: thinkingLevel (string: "low"/"high") """ @@ -887,13 +1196,13 @@ def _handle_reasoning_parameters(self, payload: Dict[str, Any], model: str) -> O payload.pop("reasoning_effort", None) payload.pop("custom_reasoning_budget", None) return None - + # Gemini 3: String-based thinkingLevel if is_gemini_3: # Clean up the original payload payload.pop("reasoning_effort", None) payload.pop("custom_reasoning_budget", None) - + if reasoning_effort == "low": return {"thinkingLevel": "low", "include_thoughts": True} return {"thinkingLevel": "high", "include_thoughts": True} @@ -918,122 +1227,137 @@ def _handle_reasoning_parameters(self, payload: Dict[str, Any], model: str) -> O budget = budgets.get(reasoning_effort, -1) if reasoning_effort == "disable": budget = 0 - + if not custom_reasoning_budget: budget = budget // 4 # Clean up the original payload payload.pop("reasoning_effort", None) payload.pop("custom_reasoning_budget", None) - + return {"thinkingBudget": budget, "include_thoughts": True} - def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumulator: Optional[Dict[str, Any]] = None): + def _convert_chunk_to_openai( + self, + chunk: Dict[str, Any], + model_id: str, + accumulator: Optional[Dict[str, Any]] = None, + ): """ Convert Gemini response chunk to OpenAI streaming format. - + Args: chunk: Gemini API response chunk model_id: Model name accumulator: Optional dict to accumulate data for post-processing (signatures, etc.) """ - response_data = chunk.get('response', chunk) - candidates = response_data.get('candidates', []) + response_data = chunk.get("response", chunk) + candidates = response_data.get("candidates", []) if not candidates: return candidate = candidates[0] - parts = candidate.get('content', {}).get('parts', []) + parts = candidate.get("content", {}).get("parts", []) is_gemini_3 = self._is_gemini_3(model_id) for part in parts: delta = {} - - has_func = 'functionCall' in part - has_text = 'text' in part - has_sig = bool(part.get('thoughtSignature')) - is_thought = part.get('thought') is True or (isinstance(part.get('thought'), str) and str(part.get('thought')).lower() == 'true') - + + has_func = "functionCall" in part + has_text = "text" in part + has_sig = bool(part.get("thoughtSignature")) + is_thought = part.get("thought") is True or ( + isinstance(part.get("thought"), str) + and str(part.get("thought")).lower() == "true" + ) + # Skip standalone signature parts (no function, no meaningful text) - if has_sig and not has_func and (not has_text or not part.get('text')): + if has_sig and not has_func and (not has_text or not part.get("text")): continue if has_func: - function_call = part['functionCall'] - function_name = function_call.get('name', 'unknown') - + function_call = part["functionCall"] + function_name = function_call.get("name", "unknown") + # Strip Gemini 3 prefix from tool name if is_gemini_3 and self._enable_gemini3_tool_fix: function_name = self._strip_gemini3_prefix(function_name) - + # Use provided ID or generate unique one with nanosecond precision - tool_call_id = function_call.get('id') or f"call_{function_name}_{int(time.time() * 1_000_000_000)}" - + tool_call_id = ( + function_call.get("id") + or f"call_{function_name}_{int(time.time() * 1_000_000_000)}" + ) + # Get current tool index from accumulator (default 0) and increment - current_tool_idx = accumulator.get('tool_idx', 0) if accumulator else 0 - + current_tool_idx = accumulator.get("tool_idx", 0) if accumulator else 0 + tool_call = { "index": current_tool_idx, "id": tool_call_id, "type": "function", "function": { "name": function_name, - "arguments": json.dumps(function_call.get('args', {})) - } + "arguments": json.dumps(function_call.get("args", {})), + }, } - + # Handle thoughtSignature for Gemini 3 # Store signature for each tool call (needed for parallel tool calls) if is_gemini_3 and has_sig: - sig = part['thoughtSignature'] - + sig = part["thoughtSignature"] + if self._enable_signature_cache: self._signature_cache.store(tool_call_id, sig) lib_logger.debug(f"Stored signature for {tool_call_id}") - + if self._preserve_signatures_in_client: tool_call["thought_signature"] = sig - - delta['tool_calls'] = [tool_call] + + delta["tool_calls"] = [tool_call] # Mark that we've sent tool calls and increment tool_idx if accumulator is not None: - accumulator['has_tool_calls'] = True - accumulator['tool_idx'] = current_tool_idx + 1 - + accumulator["has_tool_calls"] = True + accumulator["tool_idx"] = current_tool_idx + 1 + elif has_text: # Use an explicit check for the 'thought' flag, as its type can be inconsistent if is_thought: - delta['reasoning_content'] = part['text'] + delta["reasoning_content"] = part["text"] else: - delta['content'] = part['text'] - + delta["content"] = part["text"] + if not delta: continue # Mark that we have tool calls for accumulator tracking # finish_reason determination is handled by the client - + # Mark stream complete if we have usageMetadata - is_final_chunk = 'usageMetadata' in response_data + is_final_chunk = "usageMetadata" in response_data if is_final_chunk and accumulator is not None: - accumulator['is_complete'] = True + accumulator["is_complete"] = True # Build choice - don't include finish_reason, let client handle it choice = {"index": 0, "delta": delta} - + openai_chunk = { - "choices": [choice], "model": model_id, "object": "chat.completion.chunk", - "id": chunk.get("responseId", f"chatcmpl-geminicli-{time.time()}"), "created": int(time.time()) + "choices": [choice], + "model": model_id, + "object": "chat.completion.chunk", + "id": chunk.get("responseId", f"chatcmpl-geminicli-{time.time()}"), + "created": int(time.time()), } - if 'usageMetadata' in response_data: - usage = response_data['usageMetadata'] + if "usageMetadata" in response_data: + usage = response_data["usageMetadata"] prompt_tokens = usage.get("promptTokenCount", 0) thoughts_tokens = usage.get("thoughtsTokenCount", 0) candidate_tokens = usage.get("candidatesTokenCount", 0) openai_chunk["usage"] = { - "prompt_tokens": prompt_tokens + thoughts_tokens, # Include thoughts in prompt tokens + "prompt_tokens": prompt_tokens + + thoughts_tokens, # Include thoughts in prompt tokens "completion_tokens": candidate_tokens, "total_tokens": usage.get("totalTokenCount", 0), } @@ -1042,14 +1366,18 @@ def _convert_chunk_to_openai(self, chunk: Dict[str, Any], model_id: str, accumul if thoughts_tokens > 0: if "completion_tokens_details" not in openai_chunk["usage"]: openai_chunk["usage"]["completion_tokens_details"] = {} - openai_chunk["usage"]["completion_tokens_details"]["reasoning_tokens"] = thoughts_tokens - + openai_chunk["usage"]["completion_tokens_details"][ + "reasoning_tokens" + ] = thoughts_tokens + yield openai_chunk - def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> litellm.ModelResponse: + def _stream_to_completion_response( + self, chunks: List[litellm.ModelResponse] + ) -> litellm.ModelResponse: """ Manually reassembles streaming chunks into a complete response. - + Key improvements: - Determines finish_reason based on accumulated state - Priority: tool_calls > chunk's finish_reason (length, content_filter, etc.) > stop @@ -1069,7 +1397,7 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> # Process each chunk to aggregate content for chunk in chunks: - if not hasattr(chunk, 'choices') or not chunk.choices: + if not hasattr(chunk, "choices") or not chunk.choices: continue choice = chunk.choices[0] @@ -1092,25 +1420,48 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> for tc_chunk in delta["tool_calls"]: index = tc_chunk.get("index", 0) if index not in aggregated_tool_calls: - aggregated_tool_calls[index] = {"type": "function", "function": {"name": "", "arguments": ""}} + aggregated_tool_calls[index] = { + "type": "function", + "function": {"name": "", "arguments": ""}, + } if "id" in tc_chunk: aggregated_tool_calls[index]["id"] = tc_chunk["id"] if "type" in tc_chunk: aggregated_tool_calls[index]["type"] = tc_chunk["type"] if "function" in tc_chunk: - if "name" in tc_chunk["function"] and tc_chunk["function"]["name"] is not None: - aggregated_tool_calls[index]["function"]["name"] += tc_chunk["function"]["name"] - if "arguments" in tc_chunk["function"] and tc_chunk["function"]["arguments"] is not None: - aggregated_tool_calls[index]["function"]["arguments"] += tc_chunk["function"]["arguments"] + if ( + "name" in tc_chunk["function"] + and tc_chunk["function"]["name"] is not None + ): + aggregated_tool_calls[index]["function"]["name"] += ( + tc_chunk["function"]["name"] + ) + if ( + "arguments" in tc_chunk["function"] + and tc_chunk["function"]["arguments"] is not None + ): + aggregated_tool_calls[index]["function"]["arguments"] += ( + tc_chunk["function"]["arguments"] + ) # Aggregate function calls (legacy format) if "function_call" in delta and delta["function_call"] is not None: if "function_call" not in final_message: final_message["function_call"] = {"name": "", "arguments": ""} - if "name" in delta["function_call"] and delta["function_call"]["name"] is not None: - final_message["function_call"]["name"] += delta["function_call"]["name"] - if "arguments" in delta["function_call"] and delta["function_call"]["arguments"] is not None: - final_message["function_call"]["arguments"] += delta["function_call"]["arguments"] + if ( + "name" in delta["function_call"] + and delta["function_call"]["name"] is not None + ): + final_message["function_call"]["name"] += delta["function_call"][ + "name" + ] + if ( + "arguments" in delta["function_call"] + and delta["function_call"]["arguments"] is not None + ): + final_message["function_call"]["arguments"] += delta[ + "function_call" + ]["arguments"] # Track finish_reason from chunks (respects length, content_filter, etc.) if choice.get("finish_reason"): @@ -1118,7 +1469,7 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> # Handle usage data from the last chunk that has it for chunk in reversed(chunks): - if hasattr(chunk, 'usage') and chunk.usage: + if hasattr(chunk, "usage") and chunk.usage: usage_data = chunk.usage break @@ -1139,12 +1490,12 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> finish_reason = chunk_finish_reason else: finish_reason = "stop" - + # Construct the final response final_choice = { "index": 0, "message": final_message, - "finish_reason": finish_reason + "finish_reason": finish_reason, } # Create the final ModelResponse @@ -1154,7 +1505,7 @@ def _stream_to_completion_response(self, chunks: List[litellm.ModelResponse]) -> "created": first_chunk.created, "model": first_chunk.model, "choices": [final_choice], - "usage": usage_data + "usage": usage_data, } return litellm.ModelResponse(**final_response_data) @@ -1169,63 +1520,72 @@ def _gemini_cli_transform_schema(self, schema: Dict[str, Any]) -> Dict[str, Any] return schema # Handle nullable types - if 'type' in schema and isinstance(schema['type'], list): - types = schema['type'] - if 'null' in types: - schema['nullable'] = True - remaining_types = [t for t in types if t != 'null'] + if "type" in schema and isinstance(schema["type"], list): + types = schema["type"] + if "null" in types: + schema["nullable"] = True + remaining_types = [t for t in types if t != "null"] if len(remaining_types) == 1: - schema['type'] = remaining_types[0] + schema["type"] = remaining_types[0] elif len(remaining_types) > 1: - schema['type'] = remaining_types # Let's see if Gemini supports this + schema["type"] = ( + remaining_types # Let's see if Gemini supports this + ) else: - del schema['type'] + del schema["type"] # Recurse into properties - if 'properties' in schema and isinstance(schema['properties'], dict): - for prop_schema in schema['properties'].values(): + if "properties" in schema and isinstance(schema["properties"], dict): + for prop_schema in schema["properties"].values(): self._gemini_cli_transform_schema(prop_schema) # Recurse into items (for arrays) - if 'items' in schema and isinstance(schema['items'], dict): - self._gemini_cli_transform_schema(schema['items']) + if "items" in schema and isinstance(schema["items"], dict): + self._gemini_cli_transform_schema(schema["items"]) # Clean up unsupported properties schema.pop("strict", None) schema.pop("additionalProperties", None) - + return schema def _enforce_strict_schema(self, schema: Any) -> Any: """ Enforce strict JSON schema for Gemini 3 to prevent hallucinated parameters. - + Adds 'additionalProperties: false' recursively to all object schemas, which tells the model it CANNOT add properties not in the schema. """ if not isinstance(schema, dict): return schema - + result = {} for key, value in schema.items(): if isinstance(value, dict): result[key] = self._enforce_strict_schema(value) elif isinstance(value, list): - result[key] = [self._enforce_strict_schema(item) if isinstance(item, dict) else item for item in value] + result[key] = [ + self._enforce_strict_schema(item) + if isinstance(item, dict) + else item + for item in value + ] else: result[key] = value - + # Add additionalProperties: false to object schemas if result.get("type") == "object" and "properties" in result: result["additionalProperties"] = False - + return result - def _transform_tool_schemas(self, tools: List[Dict[str, Any]], model: str = "") -> List[Dict[str, Any]]: + def _transform_tool_schemas( + self, tools: List[Dict[str, Any]], model: str = "" + ) -> List[Dict[str, Any]]: """ Transforms a list of OpenAI-style tool schemas into the format required by the Gemini CLI API. This uses a custom schema transformer instead of litellm's generic one. - + For Gemini 3 models, also applies: - Namespace prefix to tool names - Parameter signature injection into descriptions @@ -1233,22 +1593,27 @@ def _transform_tool_schemas(self, tools: List[Dict[str, Any]], model: str = "") """ transformed_declarations = [] is_gemini_3 = self._is_gemini_3(model) - + for tool in tools: if tool.get("type") == "function" and "function" in tool: new_function = json.loads(json.dumps(tool["function"])) - + # The Gemini CLI API does not support the 'strict' property. new_function.pop("strict", None) # Gemini CLI expects 'parametersJsonSchema' instead of 'parameters' if "parameters" in new_function: - schema = self._gemini_cli_transform_schema(new_function["parameters"]) + schema = self._gemini_cli_transform_schema( + new_function["parameters"] + ) new_function["parametersJsonSchema"] = schema del new_function["parameters"] elif "parametersJsonSchema" not in new_function: # Set default empty schema if neither exists - new_function["parametersJsonSchema"] = {"type": "object", "properties": {}} + new_function["parametersJsonSchema"] = { + "type": "object", + "properties": {}, + } # Gemini 3 specific transformations if is_gemini_3 and self._enable_gemini3_tool_fix: @@ -1256,64 +1621,73 @@ def _transform_tool_schemas(self, tools: List[Dict[str, Any]], model: str = "") name = new_function.get("name", "") if name: new_function["name"] = f"{self._gemini3_tool_prefix}{name}" - + # Enforce strict schema (additionalProperties: false) - if self._gemini3_enforce_strict_schema and "parametersJsonSchema" in new_function: - new_function["parametersJsonSchema"] = self._enforce_strict_schema(new_function["parametersJsonSchema"]) - + if ( + self._gemini3_enforce_strict_schema + and "parametersJsonSchema" in new_function + ): + new_function["parametersJsonSchema"] = ( + self._enforce_strict_schema( + new_function["parametersJsonSchema"] + ) + ) + # Inject parameter signature into description new_function = self._inject_signature_into_description(new_function) transformed_declarations.append(new_function) - + return transformed_declarations - def _inject_signature_into_description(self, func_decl: Dict[str, Any]) -> Dict[str, Any]: + def _inject_signature_into_description( + self, func_decl: Dict[str, Any] + ) -> Dict[str, Any]: """Inject parameter signatures into tool description for Gemini 3.""" schema = func_decl.get("parametersJsonSchema", {}) if not schema: return func_decl - + required = schema.get("required", []) properties = schema.get("properties", {}) - + if not properties: return func_decl - + param_list = [] for prop_name, prop_data in properties.items(): if not isinstance(prop_data, dict): continue - + type_hint = self._format_type_hint(prop_data) is_required = prop_name in required param_list.append( f"{prop_name} ({type_hint}{', REQUIRED' if is_required else ''})" ) - + if param_list: sig_str = self._gemini3_description_prompt.replace( "{params}", ", ".join(param_list) ) func_decl["description"] = func_decl.get("description", "") + sig_str - + return func_decl def _format_type_hint(self, prop_data: Dict[str, Any], depth: int = 0) -> str: """Format a detailed type hint for a property schema.""" type_hint = prop_data.get("type", "unknown") - + # Handle enum values - show allowed options if "enum" in prop_data: enum_vals = prop_data["enum"] if len(enum_vals) <= 5: return f"string ENUM[{', '.join(repr(v) for v in enum_vals)}]" return f"string ENUM[{len(enum_vals)} options]" - + # Handle const values if "const" in prop_data: return f"string CONST={repr(prop_data['const'])}" - + if type_hint == "array": items = prop_data.get("items", {}) if isinstance(items, dict): @@ -1336,7 +1710,7 @@ def _format_type_hint(self, prop_data: Dict[str, Any], depth: int = 0) -> str: return "ARRAY_OF_OBJECTS" return f"ARRAY_OF_{item_type.upper()}" return "ARRAY" - + if type_hint == "object": nested_props = prop_data.get("properties", {}) nested_req = prop_data.get("required", []) @@ -1348,31 +1722,39 @@ def _format_type_hint(self, prop_data: Dict[str, Any], depth: int = 0) -> str: req = " REQUIRED" if n in nested_req else "" nested_list.append(f"{n}: {t}{req}") return f"object{{{', '.join(nested_list)}}}" - + return type_hint - def _inject_gemini3_system_instruction(self, request_payload: Dict[str, Any]) -> None: + def _inject_gemini3_system_instruction( + self, request_payload: Dict[str, Any] + ) -> None: """Inject Gemini 3 tool fix system instruction if tools are present.""" if not request_payload.get("request", {}).get("tools"): return - + existing_system = request_payload.get("request", {}).get("systemInstruction") - + if existing_system: # Prepend to existing system instruction existing_parts = existing_system.get("parts", []) if existing_parts and existing_parts[0].get("text"): - existing_parts[0]["text"] = self._gemini3_system_instruction + "\n\n" + existing_parts[0]["text"] + existing_parts[0]["text"] = ( + self._gemini3_system_instruction + + "\n\n" + + existing_parts[0]["text"] + ) else: existing_parts.insert(0, {"text": self._gemini3_system_instruction}) else: # Create new system instruction request_payload["request"]["systemInstruction"] = { "role": "user", - "parts": [{"text": self._gemini3_system_instruction}] + "parts": [{"text": self._gemini3_system_instruction}], } - def _translate_tool_choice(self, tool_choice: Union[str, Dict[str, Any]], model: str = "") -> Optional[Dict[str, Any]]: + def _translate_tool_choice( + self, tool_choice: Union[str, Dict[str, Any]], model: str = "" + ) -> Optional[Dict[str, Any]]: """ Translates OpenAI's `tool_choice` to Gemini's `toolConfig`. Handles Gemini 3 namespace prefixes for specific tool selection. @@ -1397,18 +1779,20 @@ def _translate_tool_choice(self, tool_choice: Union[str, Dict[str, Any]], model: # Add Gemini 3 prefix if needed if is_gemini_3 and self._enable_gemini3_tool_fix: function_name = f"{self._gemini3_tool_prefix}{function_name}" - - mode = "ANY" # Force a call, but only to this function + + mode = "ANY" # Force a call, but only to this function config["functionCallingConfig"] = { "mode": mode, - "allowedFunctionNames": [function_name] + "allowedFunctionNames": [function_name], } return config config["functionCallingConfig"] = {"mode": mode} return config - async def acompletion(self, client: httpx.AsyncClient, **kwargs) -> Union[litellm.ModelResponse, AsyncGenerator[litellm.ModelResponse, None]]: + async def acompletion( + self, client: httpx.AsyncClient, **kwargs + ) -> Union[litellm.ModelResponse, AsyncGenerator[litellm.ModelResponse, None]]: model = kwargs["model"] credential_path = kwargs.pop("credential_identifier") enable_request_logging = kwargs.pop("enable_request_logging", False) @@ -1423,28 +1807,37 @@ async def do_call(attempt_model: str, is_fallback: bool = False): # Discover project ID only if not already cached project_id = self.project_id_cache.get(credential_path) if not project_id: - access_token = auth_header['Authorization'].split(' ')[1] - project_id = await self._discover_project_id(credential_path, access_token, kwargs.get("litellm_params", {})) + access_token = auth_header["Authorization"].split(" ")[1] + project_id = await self._discover_project_id( + credential_path, access_token, kwargs.get("litellm_params", {}) + ) # Log paid tier usage visibly on each request credential_tier = self.project_tier_cache.get(credential_path) - if credential_tier and credential_tier not in ['free-tier', 'legacy-tier', 'unknown']: - lib_logger.info(f"[PAID TIER] Using Gemini '{credential_tier}' subscription for this request") + if credential_tier and credential_tier not in [ + "free-tier", + "legacy-tier", + "unknown", + ]: + lib_logger.info( + f"[PAID TIER] Using Gemini '{credential_tier}' subscription for this request" + ) # Handle :thinking suffix - model_name = attempt_model.split('/')[-1].replace(':thinking', '') + model_name = attempt_model.split("/")[-1].replace(":thinking", "") # [NEW] Create a dedicated file logger for this request file_logger = _GeminiCliFileLogger( - model_name=model_name, - enabled=enable_request_logging + model_name=model_name, enabled=enable_request_logging ) - + is_gemini_3 = self._is_gemini_3(model_name) gen_config = { - "maxOutputTokens": kwargs.get("max_tokens", 64000), # Increased default - "temperature": kwargs.get("temperature", 1), # Default to 1 if not provided + "maxOutputTokens": kwargs.get("max_tokens", 64000), # Increased default + "temperature": kwargs.get( + "temperature", 1 + ), # Default to 1 if not provided } if "top_k" in kwargs: gen_config["topK"] = kwargs["top_k"] @@ -1456,7 +1849,9 @@ async def do_call(attempt_model: str, is_fallback: bool = False): if thinking_config: gen_config["thinkingConfig"] = thinking_config - system_instruction, contents = self._transform_messages(kwargs.get("messages", []), model_name) + system_instruction, contents = self._transform_messages( + kwargs.get("messages", []), model_name + ) request_payload = { "model": model_name, "project": project_id, @@ -1470,16 +1865,22 @@ async def do_call(attempt_model: str, is_fallback: bool = False): request_payload["request"]["systemInstruction"] = system_instruction if "tools" in kwargs and kwargs["tools"]: - function_declarations = self._transform_tool_schemas(kwargs["tools"], model_name) + function_declarations = self._transform_tool_schemas( + kwargs["tools"], model_name + ) if function_declarations: - request_payload["request"]["tools"] = [{"functionDeclarations": function_declarations}] + request_payload["request"]["tools"] = [ + {"functionDeclarations": function_declarations} + ] # [NEW] Handle tool_choice translation if "tool_choice" in kwargs and kwargs["tool_choice"]: - tool_config = self._translate_tool_choice(kwargs["tool_choice"], model_name) + tool_config = self._translate_tool_choice( + kwargs["tool_choice"], model_name + ) if tool_config: request_payload["request"]["toolConfig"] = tool_config - + # Inject Gemini 3 system instruction if using tools if is_gemini_3 and self._enable_gemini3_tool_fix: self._inject_gemini3_system_instruction(request_payload) @@ -1491,52 +1892,77 @@ async def do_call(attempt_model: str, is_fallback: bool = False): {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"}, {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"}, {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"}, - {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"}, + { + "category": "HARM_CATEGORY_CIVIC_INTEGRITY", + "threshold": "BLOCK_NONE", + }, ] # Log the final payload for debugging and to the dedicated file - #lib_logger.debug(f"Gemini CLI Request Payload: {json.dumps(request_payload, indent=2)}") + # lib_logger.debug(f"Gemini CLI Request Payload: {json.dumps(request_payload, indent=2)}") file_logger.log_request(request_payload) - + url = f"{CODE_ASSIST_ENDPOINT}:streamGenerateContent" async def stream_handler(): # Track state across chunks for tool indexing - accumulator = {"has_tool_calls": False, "tool_idx": 0, "is_complete": False} - + accumulator = { + "has_tool_calls": False, + "tool_idx": 0, + "is_complete": False, + } + final_headers = auth_header.copy() - final_headers.update({ - "User-Agent": "google-api-nodejs-client/9.15.1", - "X-Goog-Api-Client": "gl-node/22.17.0", - "Client-Metadata": "ideType=IDE_UNSPECIFIED,platform=PLATFORM_UNSPECIFIED,pluginType=GEMINI", - "Accept": "application/json", - }) + final_headers.update( + { + "User-Agent": "google-api-nodejs-client/9.15.1", + "X-Goog-Api-Client": "gl-node/22.17.0", + "Client-Metadata": "ideType=IDE_UNSPECIFIED,platform=PLATFORM_UNSPECIFIED,pluginType=GEMINI", + "Accept": "application/json", + } + ) try: - async with client.stream("POST", url, headers=final_headers, json=request_payload, params={"alt": "sse"}, timeout=600) as response: + async with client.stream( + "POST", + url, + headers=final_headers, + json=request_payload, + params={"alt": "sse"}, + timeout=600, + ) as response: # Read and log error body before raise_for_status for better debugging if response.status_code >= 400: try: error_body = await response.aread() - lib_logger.error(f"Gemini CLI API error {response.status_code}: {error_body.decode()}") - file_logger.log_error(f"API error {response.status_code}: {error_body.decode()}") + lib_logger.error( + f"Gemini CLI API error {response.status_code}: {error_body.decode()}" + ) + file_logger.log_error( + f"API error {response.status_code}: {error_body.decode()}" + ) except Exception: pass - + # This will raise an HTTPStatusError for 4xx/5xx responses response.raise_for_status() async for line in response.aiter_lines(): file_logger.log_response_chunk(line) - if line.startswith('data: '): + if line.startswith("data: "): data_str = line[6:] - if data_str == "[DONE]": break + if data_str == "[DONE]": + break try: chunk = json.loads(data_str) - for openai_chunk in self._convert_chunk_to_openai(chunk, model, accumulator): + for openai_chunk in self._convert_chunk_to_openai( + chunk, model, accumulator + ): yield litellm.ModelResponse(**openai_chunk) except json.JSONDecodeError: - lib_logger.warning(f"Could not decode JSON from Gemini CLI: {line}") - + lib_logger.warning( + f"Could not decode JSON from Gemini CLI: {line}" + ) + # Emit final chunk if stream ended without usageMetadata # Client will determine the correct finish_reason if not accumulator.get("is_complete"): @@ -1545,9 +1971,15 @@ async def stream_handler(): "object": "chat.completion.chunk", "created": int(time.time()), "model": model, - "choices": [{"index": 0, "delta": {}, "finish_reason": None}], + "choices": [ + {"index": 0, "delta": {}, "finish_reason": None} + ], # Include minimal usage to signal this is the final chunk - "usage": {"prompt_tokens": 0, "completion_tokens": 1, "total_tokens": 1} + "usage": { + "prompt_tokens": 0, + "completion_tokens": 1, + "total_tokens": 1, + }, } yield litellm.ModelResponse(**final_chunk) @@ -1558,27 +1990,35 @@ async def stream_handler(): error_body = e.response.text except Exception: pass - + # Only log to file logger (for detailed logging) if error_body: - file_logger.log_error(f"HTTPStatusError {e.response.status_code}: {error_body}") + file_logger.log_error( + f"HTTPStatusError {e.response.status_code}: {error_body}" + ) else: - file_logger.log_error(f"HTTPStatusError {e.response.status_code}: {str(e)}") - + file_logger.log_error( + f"HTTPStatusError {e.response.status_code}: {str(e)}" + ) + if e.response.status_code == 429: # Extract retry-after time from the error body retry_after = extract_retry_after_from_body(error_body) - retry_info = f" (retry after {retry_after}s)" if retry_after else "" + retry_info = ( + f" (retry after {retry_after}s)" if retry_after else "" + ) error_msg = f"Gemini CLI rate limit exceeded{retry_info}" if error_body: error_msg = f"{error_msg} | {error_body}" # Only log at debug level - rotation happens silently - lib_logger.debug(f"Gemini CLI 429 rate limit: retry_after={retry_after}s") + lib_logger.debug( + f"Gemini CLI 429 rate limit: retry_after={retry_after}s" + ) raise RateLimitError( message=error_msg, llm_provider="gemini_cli", model=model, - response=e.response + response=e.response, ) # Re-raise other status errors to be handled by the main acompletion logic raise e @@ -1595,29 +2035,41 @@ async def logging_stream_wrapper(): yield chunk finally: if openai_chunks: - final_response = self._stream_to_completion_response(openai_chunks) + final_response = self._stream_to_completion_response( + openai_chunks + ) file_logger.log_final_response(final_response.dict()) return logging_stream_wrapper() # Check if there are actual fallback models available # If fallback_models is empty or contains only the base model (no actual fallbacks), skip fallback logic - has_fallbacks = len(fallback_models) > 1 and any(model != fallback_models[0] for model in fallback_models[1:]) - + has_fallbacks = len(fallback_models) > 1 and any( + model != fallback_models[0] for model in fallback_models[1:] + ) + lib_logger.debug(f"Fallback models available: {fallback_models}") if not has_fallbacks: - lib_logger.debug("No actual fallback models available, proceeding with single model attempt") - + lib_logger.debug( + "No actual fallback models available, proceeding with single model attempt" + ) + last_error = None for idx, attempt_model in enumerate(fallback_models): is_fallback = idx > 0 if is_fallback: # Silent rotation - only log at debug level - lib_logger.debug(f"Rate limited on previous model, trying fallback: {attempt_model}") + lib_logger.debug( + f"Rate limited on previous model, trying fallback: {attempt_model}" + ) elif has_fallbacks: - lib_logger.debug(f"Attempting primary model: {attempt_model} (with {len(fallback_models)-1} fallback(s) available)") + lib_logger.debug( + f"Attempting primary model: {attempt_model} (with {len(fallback_models) - 1} fallback(s) available)" + ) else: - lib_logger.debug(f"Attempting model: {attempt_model} (no fallbacks available)") + lib_logger.debug( + f"Attempting model: {attempt_model} (no fallbacks available)" + ) try: response_gen = await do_call(attempt_model, is_fallback) @@ -1633,10 +2085,14 @@ async def logging_stream_wrapper(): last_error = e # If this is not the last model in the fallback chain, continue to next model if idx + 1 < len(fallback_models): - lib_logger.debug(f"Rate limit hit on {attempt_model}, trying next fallback...") + lib_logger.debug( + f"Rate limit hit on {attempt_model}, trying next fallback..." + ) continue # If this was the last fallback option, log error and raise - lib_logger.warning(f"Rate limit exhausted on all fallback models (tried {len(fallback_models)} models)") + lib_logger.warning( + f"Rate limit exhausted on all fallback models (tried {len(fallback_models)} models)" + ) raise # Should not reach here, but raise last error if we do @@ -1651,7 +2107,7 @@ async def count_tokens( model: str, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, - litellm_params: Optional[Dict[str, Any]] = None + litellm_params: Optional[Dict[str, Any]] = None, ) -> Dict[str, int]: """ Counts tokens for the given prompt using the Gemini CLI :countTokens endpoint. @@ -1673,11 +2129,13 @@ async def count_tokens( # Discover project ID project_id = self.project_id_cache.get(credential_path) if not project_id: - access_token = auth_header['Authorization'].split(' ')[1] - project_id = await self._discover_project_id(credential_path, access_token, litellm_params or {}) + access_token = auth_header["Authorization"].split(" ")[1] + project_id = await self._discover_project_id( + credential_path, access_token, litellm_params or {} + ) # Handle :thinking suffix - model_name = model.split('/')[-1].replace(':thinking', '') + model_name = model.split("/")[-1].replace(":thinking", "") # Transform messages to Gemini format system_instruction, contents = self._transform_messages(messages) @@ -1695,35 +2153,41 @@ async def count_tokens( if tools: function_declarations = self._transform_tool_schemas(tools) if function_declarations: - request_payload["request"]["tools"] = [{"functionDeclarations": function_declarations}] + request_payload["request"]["tools"] = [ + {"functionDeclarations": function_declarations} + ] # Make the request url = f"{CODE_ASSIST_ENDPOINT}:countTokens" headers = auth_header.copy() - headers.update({ - "User-Agent": "google-api-nodejs-client/9.15.1", - "X-Goog-Api-Client": "gl-node/22.17.0", - "Client-Metadata": "ideType=IDE_UNSPECIFIED,platform=PLATFORM_UNSPECIFIED,pluginType=GEMINI", - "Accept": "application/json", - }) + headers.update( + { + "User-Agent": "google-api-nodejs-client/9.15.1", + "X-Goog-Api-Client": "gl-node/22.17.0", + "Client-Metadata": "ideType=IDE_UNSPECIFIED,platform=PLATFORM_UNSPECIFIED,pluginType=GEMINI", + "Accept": "application/json", + } + ) try: - response = await client.post(url, headers=headers, json=request_payload, timeout=30) + response = await client.post( + url, headers=headers, json=request_payload, timeout=30 + ) response.raise_for_status() data = response.json() # Extract token counts from response - total_tokens = data.get('totalTokens', 0) + total_tokens = data.get("totalTokens", 0) return { - 'prompt_tokens': total_tokens, - 'total_tokens': total_tokens, + "prompt_tokens": total_tokens, + "total_tokens": total_tokens, } except httpx.HTTPStatusError as e: lib_logger.error(f"Failed to count tokens: {e}") # Return 0 on error rather than raising - return {'prompt_tokens': 0, 'total_tokens': 0} + return {"prompt_tokens": 0, "total_tokens": 0} # Use the shared GeminiAuthBase for auth logic async def get_models(self, credential: str, client: httpx.AsyncClient) -> List[str]: @@ -1738,9 +2202,11 @@ async def get_models(self, credential: str, client: httpx.AsyncClient) -> List[s """ # Check for mixed tier credentials and warn if detected self._check_mixed_tier_warning() - + models = [] - env_var_ids = set() # Track IDs from env vars to prevent hardcoded/dynamic duplicates + env_var_ids = ( + set() + ) # Track IDs from env vars to prevent hardcoded/dynamic duplicates def extract_model_id(item) -> str: """Extract model ID from various formats (dict, string with/without provider prefix).""" @@ -1770,7 +2236,9 @@ def extract_model_id(item) -> str: # Track the ID to prevent hardcoded/dynamic duplicates if model_id: env_var_ids.add(model_id) - lib_logger.info(f"Loaded {len(static_models)} static models for gemini_cli from environment variables") + lib_logger.info( + f"Loaded {len(static_models)} static models for gemini_cli from environment variables" + ) # Source 2: Add hardcoded models (only if ID not already in env vars) for model_id in HARDCODED_MODELS: @@ -1782,7 +2250,7 @@ def extract_model_id(item) -> str: try: # Get access token for API calls auth_header = await self.get_auth_header(credential) - access_token = auth_header['Authorization'].split(' ')[1] + access_token = auth_header["Authorization"].split(" ")[1] # Try Vertex AI models endpoint # Note: Gemini may not support a simple /models endpoint like OpenAI @@ -1790,8 +2258,7 @@ def extract_model_id(item) -> str: models_url = f"https://generativelanguage.googleapis.com/v1beta/models" response = await client.get( - models_url, - headers={"Authorization": f"Bearer {access_token}"} + models_url, headers={"Authorization": f"Bearer {access_token}"} ) response.raise_for_status() @@ -1803,17 +2270,23 @@ def extract_model_id(item) -> str: for model in model_list: model_id = extract_model_id(model) # Only include Gemini models that aren't already in env vars - if model_id and model_id not in env_var_ids and model_id.startswith("gemini"): + if ( + model_id + and model_id not in env_var_ids + and model_id.startswith("gemini") + ): models.append(f"gemini_cli/{model_id}") env_var_ids.add(model_id) dynamic_count += 1 if dynamic_count > 0: - lib_logger.debug(f"Discovered {dynamic_count} additional models for gemini_cli from API") + lib_logger.debug( + f"Discovered {dynamic_count} additional models for gemini_cli from API" + ) except Exception as e: # Silently ignore dynamic discovery errors lib_logger.debug(f"Dynamic model discovery failed for gemini_cli: {e}") pass - return models \ No newline at end of file + return models diff --git a/src/rotator_library/providers/provider_interface.py b/src/rotator_library/providers/provider_interface.py index 8a20a64c..996f3a7e 100644 --- a/src/rotator_library/providers/provider_interface.py +++ b/src/rotator_library/providers/provider_interface.py @@ -3,13 +3,15 @@ import httpx import litellm + class ProviderInterface(ABC): """ An interface for API provider-specific functionality, including model discovery and custom API call handling for non-standard providers. """ + skip_cost_calculation: bool = False - + @abstractmethod async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: """ @@ -32,28 +34,38 @@ def has_custom_logic(self) -> bool: """ return False - async def acompletion(self, client: httpx.AsyncClient, **kwargs) -> Union[litellm.ModelResponse, AsyncGenerator[litellm.ModelResponse, None]]: + async def acompletion( + self, client: httpx.AsyncClient, **kwargs + ) -> Union[litellm.ModelResponse, AsyncGenerator[litellm.ModelResponse, None]]: """ Handles the entire completion call for non-standard providers. """ - raise NotImplementedError(f"{self.__class__.__name__} does not implement custom acompletion.") + raise NotImplementedError( + f"{self.__class__.__name__} does not implement custom acompletion." + ) - async def aembedding(self, client: httpx.AsyncClient, **kwargs) -> litellm.EmbeddingResponse: + async def aembedding( + self, client: httpx.AsyncClient, **kwargs + ) -> litellm.EmbeddingResponse: """Handles the entire embedding call for non-standard providers.""" - raise NotImplementedError(f"{self.__class__.__name__} does not implement custom aembedding.") - - def convert_safety_settings(self, settings: Dict[str, str]) -> Optional[List[Dict[str, Any]]]: + raise NotImplementedError( + f"{self.__class__.__name__} does not implement custom aembedding." + ) + + def convert_safety_settings( + self, settings: Dict[str, str] + ) -> Optional[List[Dict[str, Any]]]: """ Converts a generic safety settings dictionary to the provider-specific format. - + Args: settings: A dictionary with generic harm categories and thresholds. - + Returns: A list of provider-specific safety setting objects or None. """ return None - + # [NEW] Add new methods for OAuth providers async def get_auth_header(self, credential_identifier: str) -> Dict[str, str]: """ @@ -67,23 +79,23 @@ async def proactively_refresh(self, credential_path: str): Proactively refreshes a token if it's nearing expiry. """ pass - + # [NEW] Credential Prioritization System def get_credential_priority(self, credential: str) -> Optional[int]: """ Returns the priority level for a credential. Lower numbers = higher priority (1 is highest). Returns None if provider doesn't use priorities. - + This allows providers to auto-detect credential tiers (e.g., paid vs free) and ensure higher-tier credentials are always tried first. - + Args: credential: The credential identifier (API key or path) - + Returns: Priority level (1-10) or None if no priority system - + Example: For Gemini CLI: - Paid tier credentials: priority 1 (highest) @@ -91,24 +103,53 @@ def get_credential_priority(self, credential: str) -> Optional[int]: - Unknown tier: priority 10 (lowest) """ return None - + def get_model_tier_requirement(self, model: str) -> Optional[int]: """ Returns the minimum priority tier required for a model. If a model requires priority 1, only credentials with priority <= 1 can use it. - + This allows providers to restrict certain models to specific credential tiers. For example, Gemini 3 models require paid-tier credentials. - + Args: model: The model name (with or without provider prefix) - + Returns: Minimum required priority level or None if no restrictions - + Example: For Gemini CLI: - gemini-3-*: requires priority 1 (paid tier only) - gemini-2.5-*: no restriction (None) """ - return None \ No newline at end of file + return None + + async def initialize_credentials(self, credential_paths: List[str]) -> None: + """ + Called at startup to initialize provider with all available credentials. + + Providers can override this to load cached tier data, discover priorities, + or perform any other initialization needed before the first API request. + + This is called once during startup by the BackgroundRefresher before + the main refresh loop begins. + + Args: + credential_paths: List of credential file paths for this provider + """ + pass + + def get_credential_tier_name(self, credential: str) -> Optional[str]: + """ + Returns the human-readable tier name for a credential. + + This is used for logging purposes to show which plan tier a credential belongs to. + + Args: + credential: The credential identifier (API key or path) + + Returns: + Tier name string (e.g., "free-tier", "paid-tier") or None if unknown + """ + return None diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index c72d9769..577bf4aa 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -22,24 +22,24 @@ class UsageManager: """ Manages usage statistics and cooldowns for API keys with asyncio-safe locking, asynchronous file I/O, lazy-loading mechanism, and weighted random credential rotation. - + The credential rotation strategy can be configured via the `rotation_tolerance` parameter: - + - **tolerance = 0.0**: Deterministic least-used selection. The credential with the lowest usage count is always selected. This provides predictable, perfectly balanced load distribution but may be vulnerable to fingerprinting. - + - **tolerance = 2.0 - 4.0 (default, recommended)**: Balanced weighted randomness. Credentials are selected randomly with weights biased toward less-used ones. Credentials within 2 uses of the maximum can still be selected with reasonable probability. This provides security through unpredictability while maintaining good load balance. - + - **tolerance = 5.0+**: High randomness. Even heavily-used credentials have significant selection probability. Useful for stress testing or maximum unpredictability, but may result in less balanced load distribution. - + The weight formula is: `weight = (max_usage - credential_usage) + tolerance + 1` - + This ensures lower-usage credentials are preferred while tolerance controls how much randomness is introduced into the selection process. """ @@ -52,7 +52,7 @@ def __init__( ): """ Initialize the UsageManager. - + Args: file_path: Path to the usage data JSON file daily_reset_time_utc: Time in UTC when daily stats should reset (HH:MM format) @@ -139,7 +139,9 @@ async def _reset_daily_stats_if_needed(self): last_reset_dt is None or last_reset_dt < reset_threshold_today <= now_utc ): - lib_logger.debug(f"Performing daily reset for key {mask_credential(key)}") + lib_logger.debug( + f"Performing daily reset for key {mask_credential(key)}" + ) needs_saving = True # Reset cooldowns @@ -194,24 +196,20 @@ def _initialize_key_states(self, keys: List[str]): "models_in_use": {}, # Dict[model_name, concurrent_count] } - def _select_weighted_random( - self, - candidates: List[tuple], - tolerance: float - ) -> str: + def _select_weighted_random(self, candidates: List[tuple], tolerance: float) -> str: """ Selects a credential using weighted random selection based on usage counts. - + Args: candidates: List of (credential_id, usage_count) tuples tolerance: Tolerance value for weight calculation - + Returns: Selected credential ID - + Formula: weight = (max_usage - credential_usage) + tolerance + 1 - + This formula ensures: - Lower usage = higher weight = higher selection probability - Tolerance adds variability: higher tolerance means more randomness @@ -219,63 +217,66 @@ def _select_weighted_random( """ if not candidates: raise ValueError("Cannot select from empty candidate list") - + if len(candidates) == 1: return candidates[0][0] - + # Extract usage counts usage_counts = [usage for _, usage in candidates] max_usage = max(usage_counts) - + # Calculate weights using the formula: (max - current) + tolerance + 1 weights = [] for credential, usage in candidates: weight = (max_usage - usage) + tolerance + 1 weights.append(weight) - + # Log weight distribution for debugging if lib_logger.isEnabledFor(logging.DEBUG): total_weight = sum(weights) weight_info = ", ".join( - f"{mask_credential(cred)}: w={w:.1f} ({w/total_weight*100:.1f}%)" + f"{mask_credential(cred)}: w={w:.1f} ({w / total_weight * 100:.1f}%)" for (cred, _), w in zip(candidates, weights) ) - #lib_logger.debug(f"Weighted selection candidates: {weight_info}") - + # lib_logger.debug(f"Weighted selection candidates: {weight_info}") + # Random selection with weights selected_credential = random.choices( - [cred for cred, _ in candidates], - weights=weights, - k=1 + [cred for cred, _ in candidates], weights=weights, k=1 )[0] - + return selected_credential async def acquire_key( - self, available_keys: List[str], model: str, deadline: float, + self, + available_keys: List[str], + model: str, + deadline: float, max_concurrent: int = 1, - credential_priorities: Optional[Dict[str, int]] = None + credential_priorities: Optional[Dict[str, int]] = None, + credential_tier_names: Optional[Dict[str, str]] = None, ) -> str: """ Acquires the best available key using a tiered, model-aware locking strategy, respecting a global deadline and credential priorities. - + Priority Logic: - Groups credentials by priority level (1=highest, 2=lower, etc.) - Always tries highest priority (lowest number) first - Within same priority, sorts by usage count (load balancing) - Only moves to next priority if all higher-priority keys exhausted/busy - + Args: available_keys: List of credential identifiers to choose from model: Model name being requested deadline: Timestamp after which to stop trying max_concurrent: Maximum concurrent requests allowed per credential credential_priorities: Optional dict mapping credentials to priority levels (1=highest) - + credential_tier_names: Optional dict mapping credentials to tier names (for logging) + Returns: Selected credential identifier - + Raises: NoAvailableKeysError: If no key could be acquired within the deadline """ @@ -294,16 +295,16 @@ async def acquire_key( async with self._data_lock: for key in available_keys: key_data = self._usage_data.get(key, {}) - + # Skip keys on cooldown if (key_data.get("key_cooldown_until") or 0) > now or ( key_data.get("model_cooldowns", {}).get(model) or 0 ) > now: continue - + # Get priority for this key (default to 999 if not specified) priority = credential_priorities.get(key, 999) - + # Get usage count for load balancing within priority groups usage_count = ( key_data.get("daily", {}) @@ -311,58 +312,75 @@ async def acquire_key( .get(model, {}) .get("success_count", 0) ) - + # Group by priority if priority not in priority_groups: priority_groups[priority] = [] priority_groups[priority].append((key, usage_count)) - + # Try priority groups in order (1, 2, 3, ...) sorted_priorities = sorted(priority_groups.keys()) - + for priority_level in sorted_priorities: keys_in_priority = priority_groups[priority_level] - + # Within each priority group, use existing tier1/tier2 logic tier1_keys, tier2_keys = [], [] for key, usage_count in keys_in_priority: key_state = self.key_states[key] - + # Tier 1: Completely idle keys (preferred) if not key_state["models_in_use"]: tier1_keys.append((key, usage_count)) # Tier 2: Keys that can accept more concurrent requests elif key_state["models_in_use"].get(model, 0) < max_concurrent: tier2_keys.append((key, usage_count)) - + # Apply weighted random selection or deterministic sorting - selection_method = "weighted-random" if self.rotation_tolerance > 0 else "least-used" - + selection_method = ( + "weighted-random" + if self.rotation_tolerance > 0 + else "least-used" + ) + if self.rotation_tolerance > 0: # Weighted random selection within each tier if tier1_keys: - selected_key = self._select_weighted_random(tier1_keys, self.rotation_tolerance) - tier1_keys = [(k, u) for k, u in tier1_keys if k == selected_key] + selected_key = self._select_weighted_random( + tier1_keys, self.rotation_tolerance + ) + tier1_keys = [ + (k, u) for k, u in tier1_keys if k == selected_key + ] if tier2_keys: - selected_key = self._select_weighted_random(tier2_keys, self.rotation_tolerance) - tier2_keys = [(k, u) for k, u in tier2_keys if k == selected_key] + selected_key = self._select_weighted_random( + tier2_keys, self.rotation_tolerance + ) + tier2_keys = [ + (k, u) for k, u in tier2_keys if k == selected_key + ] else: # Deterministic: sort by usage within each tier tier1_keys.sort(key=lambda x: x[1]) tier2_keys.sort(key=lambda x: x[1]) - + # Try to acquire from Tier 1 first for key, usage in tier1_keys: state = self.key_states[key] async with state["lock"]: if not state["models_in_use"]: state["models_in_use"][model] = 1 + tier_name = ( + credential_tier_names.get(key, "unknown") + if credential_tier_names + else "unknown" + ) lib_logger.info( - f"Acquired Priority-{priority_level} Tier-1 key {mask_credential(key)} for model {model} " - f"(selection: {selection_method}, usage: {usage})" + f"Acquired key {mask_credential(key)} for model {model} " + f"(tier: {tier_name}, priority: {priority_level}, selection: {selection_method}, usage: {usage})" ) return key - + # Then try Tier 2 for key, usage in tier2_keys: state = self.key_states[key] @@ -370,35 +388,40 @@ async def acquire_key( current_count = state["models_in_use"].get(model, 0) if current_count < max_concurrent: state["models_in_use"][model] = current_count + 1 + tier_name = ( + credential_tier_names.get(key, "unknown") + if credential_tier_names + else "unknown" + ) lib_logger.info( - f"Acquired Priority-{priority_level} Tier-2 key {mask_credential(key)} for model {model} " - f"(selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" + f"Acquired key {mask_credential(key)} for model {model} " + f"(tier: {tier_name}, priority: {priority_level}, selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" ) return key - + # If we get here, all priority groups were exhausted but keys might become available # Collect all keys across all priorities for waiting all_potential_keys = [] for keys_list in priority_groups.values(): all_potential_keys.extend(keys_list) - + if not all_potential_keys: lib_logger.warning( "No keys are eligible (all on cooldown or filtered out). Waiting before re-evaluating." ) await asyncio.sleep(1) continue - + # Wait for the highest priority key with lowest usage best_priority = min(priority_groups.keys()) best_priority_keys = priority_groups[best_priority] best_wait_key = min(best_priority_keys, key=lambda x: x[1])[0] wait_condition = self.key_states[best_wait_key]["condition"] - + lib_logger.info( f"All Priority-{best_priority} keys are busy. Waiting for highest priority credential to become available..." ) - + else: # Original logic when no priorities specified tier1_keys, tier2_keys = [], [] @@ -430,16 +453,26 @@ async def acquire_key( tier2_keys.append((key, usage_count)) # Apply weighted random selection or deterministic sorting - selection_method = "weighted-random" if self.rotation_tolerance > 0 else "least-used" - + selection_method = ( + "weighted-random" if self.rotation_tolerance > 0 else "least-used" + ) + if self.rotation_tolerance > 0: # Weighted random selection within each tier if tier1_keys: - selected_key = self._select_weighted_random(tier1_keys, self.rotation_tolerance) - tier1_keys = [(k, u) for k, u in tier1_keys if k == selected_key] + selected_key = self._select_weighted_random( + tier1_keys, self.rotation_tolerance + ) + tier1_keys = [ + (k, u) for k, u in tier1_keys if k == selected_key + ] if tier2_keys: - selected_key = self._select_weighted_random(tier2_keys, self.rotation_tolerance) - tier2_keys = [(k, u) for k, u in tier2_keys if k == selected_key] + selected_key = self._select_weighted_random( + tier2_keys, self.rotation_tolerance + ) + tier2_keys = [ + (k, u) for k, u in tier2_keys if k == selected_key + ] else: # Deterministic: sort by usage within each tier tier1_keys.sort(key=lambda x: x[1]) @@ -451,9 +484,15 @@ async def acquire_key( async with state["lock"]: if not state["models_in_use"]: state["models_in_use"][model] = 1 + tier_name = ( + credential_tier_names.get(key) + if credential_tier_names + else None + ) + tier_info = f"tier: {tier_name}, " if tier_name else "" lib_logger.info( - f"Acquired Tier 1 key {mask_credential(key)} for model {model} " - f"(selection: {selection_method}, usage: {usage})" + f"Acquired key {mask_credential(key)} for model {model} " + f"({tier_info}selection: {selection_method}, usage: {usage})" ) return key @@ -464,9 +503,15 @@ async def acquire_key( current_count = state["models_in_use"].get(model, 0) if current_count < max_concurrent: state["models_in_use"][model] = current_count + 1 + tier_name = ( + credential_tier_names.get(key) + if credential_tier_names + else None + ) + tier_info = f"tier: {tier_name}, " if tier_name else "" lib_logger.info( - f"Acquired Tier 2 key {mask_credential(key)} for model {model} " - f"(selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" + f"Acquired key {mask_credential(key)} for model {model} " + f"({tier_info}selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" ) return key @@ -506,8 +551,6 @@ async def acquire_key( f"Could not acquire a key for model {model} within the global time budget." ) - - async def release_key(self, key: str, model: str): """Releases a key's lock for a specific model and notifies waiting tasks.""" if key not in self.key_states: @@ -640,8 +683,11 @@ async def record_success( await self._save_usage() async def record_failure( - self, key: str, model: str, classified_error: ClassifiedError, - increment_consecutive_failures: bool = True + self, + key: str, + model: str, + classified_error: ClassifiedError, + increment_consecutive_failures: bool = True, ): """Records a failure and applies cooldowns based on an escalating backoff strategy. @@ -705,7 +751,9 @@ async def record_failure( # If cooldown wasn't set by specific error type, use escalating backoff if cooldown_seconds is None: backoff_tiers = {1: 10, 2: 30, 3: 60, 4: 120} - cooldown_seconds = backoff_tiers.get(count, 7200) # Default to 2 hours for "spent" keys + cooldown_seconds = backoff_tiers.get( + count, 7200 + ) # Default to 2 hours for "spent" keys lib_logger.warning( f"Failure #{count} for key {mask_credential(key)} with model {model}. " f"Error type: {classified_error.error_type}" From bd84d38c96b435187e230b7724a4a98481836ea2 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sat, 6 Dec 2025 05:19:16 +0100 Subject: [PATCH 087/221] =?UTF-8?q?feat(rotation):=20=E2=9C=A8=20add=20seq?= =?UTF-8?q?uential=20rotation=20mode=20with=20provider-specific=20quota=20?= =?UTF-8?q?parsing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a new credential rotation mode system that allows providers to choose between "balanced" (distribute load evenly) and "sequential" (use until exhausted) strategies. Sequential mode is particularly beneficial for providers with cache-preserving features like Antigravity's thinking signature caches. Key changes: - Added ROTATION_MODE_{PROVIDER} environment variable support with comprehensive documentation in .env.example - Implemented provider-specific quota error parsing for Antigravity and Gemini CLI providers, extracting retry_after from Google RPC error format (handles compound durations like "143h4m52.73s") - Extended ProviderInterface with rotation mode configuration and parse_quota_error() method - Updated UsageManager to support sequential credential selection that preserves sticky credential usage until quota exhaustion - Enhanced error_handler.py classify_error() to attempt provider-specific parsing before falling back to generic classification - Added rotation mode management UI in settings_tool.py with visual indicators for configured vs default modes - Preserved long-term cooldowns during daily reset to prevent premature quota retry - Updated all classify_error() call sites to pass provider parameter for context-aware parsing Provider defaults: - Antigravity: sequential (preserves thinking caches, handles weekly quota reset) - Gemini CLI: balanced (short cooldowns in seconds/minutes) - All others: balanced (standard per-minute rate limits) The sequential mode ensures the same credential is reused until it hits a cooldown (429 error), at which point the system switches to the next available credential. This maximizes cache hit rates for providers that maintain request context across API calls. --- .env.example | 26 + src/proxy_app/settings_tool.py | 934 +++++++++++++----- src/rotator_library/client.py | 48 +- src/rotator_library/error_handler.py | 62 +- .../providers/antigravity_provider.py | 141 +++ .../providers/gemini_cli_provider.py | 25 + .../providers/provider_interface.py | 72 ++ src/rotator_library/usage_manager.py | 216 +++- 8 files changed, 1217 insertions(+), 307 deletions(-) diff --git a/.env.example b/.env.example index e856b21e..9ce21139 100644 --- a/.env.example +++ b/.env.example @@ -159,6 +159,32 @@ MAX_CONCURRENT_REQUESTS_PER_KEY_GEMINI=1 MAX_CONCURRENT_REQUESTS_PER_KEY_ANTHROPIC=1 MAX_CONCURRENT_REQUESTS_PER_KEY_IFLOW=1 +# --- Credential Rotation Mode --- +# Controls how credentials are rotated when multiple are available for a provider. +# This affects how the proxy selects the next credential to use for requests. +# +# Available modes: +# balanced - (Default) Rotate credentials evenly across requests to distribute load. +# Best for API keys with per-minute rate limits. +# sequential - Use one credential until it's exhausted (429 error), then switch to next. +# Best for credentials with daily/weekly quotas (e.g., free tier accounts). +# When a credential hits quota, it's put on cooldown based on the reset time +# parsed from the provider's error response. +# +# Format: ROTATION_MODE_= +# +# Provider Defaults: +# - antigravity: sequential (free tier accounts with daily quotas) +# - All others: balanced +# +# Example: +# ROTATION_MODE_GEMINI=sequential # Use Gemini keys until quota exhausted +# ROTATION_MODE_OPENAI=balanced # Distribute load across OpenAI keys (default) +# ROTATION_MODE_ANTIGRAVITY=balanced # Override Antigravity's sequential default +# +# ROTATION_MODE_GEMINI=balanced +# ROTATION_MODE_ANTIGRAVITY=sequential + # ------------------------------------------------------------------------------ # | [ADVANCED] Proxy Configuration | # ------------------------------------------------------------------------------ diff --git a/src/proxy_app/settings_tool.py b/src/proxy_app/settings_tool.py index 59d91d5e..66b81e2e 100644 --- a/src/proxy_app/settings_tool.py +++ b/src/proxy_app/settings_tool.py @@ -17,37 +17,38 @@ def clear_screen(): """ - Cross-platform terminal clear that works robustly on both + Cross-platform terminal clear that works robustly on both classic Windows conhost and modern terminals (Windows Terminal, Linux, Mac). - + Uses native OS commands instead of ANSI escape sequences: - Windows (conhost & Windows Terminal): cls - Unix-like systems (Linux, Mac): clear """ - os.system('cls' if os.name == 'nt' else 'clear') + os.system("cls" if os.name == "nt" else "clear") class AdvancedSettings: """Manages pending changes to .env""" - + def __init__(self): self.env_file = Path.cwd() / ".env" self.pending_changes = {} # key -> value (None means delete) self.load_current_settings() - + def load_current_settings(self): """Load current .env values into env vars""" from dotenv import load_dotenv + load_dotenv(override=True) - + def set(self, key: str, value: str): """Stage a change""" self.pending_changes[key] = value - + def remove(self, key: str): """Stage a removal""" self.pending_changes[key] = None - + def save(self): """Write pending changes to .env""" for key, value in self.pending_changes.items(): @@ -57,14 +58,14 @@ def save(self): else: # Set key set_key(str(self.env_file), key, value) - + self.pending_changes.clear() self.load_current_settings() - + def discard(self): """Discard pending changes""" self.pending_changes.clear() - + def has_pending(self) -> bool: """Check if there are pending changes""" return bool(self.pending_changes) @@ -72,14 +73,14 @@ def has_pending(self) -> bool: class CustomProviderManager: """Manages custom provider API bases""" - + def __init__(self, settings: AdvancedSettings): self.settings = settings - + def get_current_providers(self) -> Dict[str, str]: """Get currently configured custom providers""" from proxy_app.provider_urls import PROVIDER_URL_MAP - + providers = {} for key, value in os.environ.items(): if key.endswith("_API_BASE"): @@ -88,16 +89,16 @@ def get_current_providers(self) -> Dict[str, str]: if provider not in PROVIDER_URL_MAP: providers[provider] = value return providers - + def add_provider(self, name: str, api_base: str): """Add PROVIDER_API_BASE""" key = f"{name.upper()}_API_BASE" self.settings.set(key, api_base) - + def edit_provider(self, name: str, api_base: str): """Edit PROVIDER_API_BASE""" self.add_provider(name, api_base) - + def remove_provider(self, name: str): """Remove PROVIDER_API_BASE""" key = f"{name.upper()}_API_BASE" @@ -106,10 +107,10 @@ def remove_provider(self, name: str): class ModelDefinitionManager: """Manages PROVIDER_MODELS""" - + def __init__(self, settings: AdvancedSettings): self.settings = settings - + def get_current_provider_models(self, provider: str) -> Optional[Dict]: """Get currently configured models for a provider""" key = f"{provider.upper()}_MODELS" @@ -120,7 +121,7 @@ def get_current_provider_models(self, provider: str) -> Optional[Dict]: except (json.JSONDecodeError, ValueError): return None return None - + def get_all_providers_with_models(self) -> Dict[str, int]: """Get all providers with model definitions""" providers = {} @@ -136,13 +137,13 @@ def get_all_providers_with_models(self) -> Dict[str, int]: except (json.JSONDecodeError, ValueError): pass return providers - + def set_models(self, provider: str, models: Dict[str, Dict[str, Any]]): """Set PROVIDER_MODELS""" key = f"{provider.upper()}_MODELS" value = json.dumps(models) self.settings.set(key, value) - + def remove_models(self, provider: str): """Remove PROVIDER_MODELS""" key = f"{provider.upper()}_MODELS" @@ -151,10 +152,10 @@ def remove_models(self, provider: str): class ConcurrencyManager: """Manages MAX_CONCURRENT_REQUESTS_PER_KEY_PROVIDER""" - + def __init__(self, settings: AdvancedSettings): self.settings = settings - + def get_current_limits(self) -> Dict[str, int]: """Get currently configured concurrency limits""" limits = {} @@ -166,18 +167,73 @@ def get_current_limits(self) -> Dict[str, int]: except (json.JSONDecodeError, ValueError): pass return limits - + def set_limit(self, provider: str, limit: int): """Set concurrency limit""" key = f"MAX_CONCURRENT_REQUESTS_PER_KEY_{provider.upper()}" self.settings.set(key, str(limit)) - + def remove_limit(self, provider: str): """Remove concurrency limit (reset to default)""" key = f"MAX_CONCURRENT_REQUESTS_PER_KEY_{provider.upper()}" self.settings.remove(key) +class RotationModeManager: + """Manages ROTATION_MODE_PROVIDER settings for sequential/balanced credential rotation""" + + VALID_MODES = ["balanced", "sequential"] + + def __init__(self, settings: AdvancedSettings): + self.settings = settings + + def get_current_modes(self) -> Dict[str, str]: + """Get currently configured rotation modes""" + modes = {} + for key, value in os.environ.items(): + if key.startswith("ROTATION_MODE_"): + provider = key.replace("ROTATION_MODE_", "").lower() + if value.lower() in self.VALID_MODES: + modes[provider] = value.lower() + return modes + + def get_default_mode(self, provider: str) -> str: + """Get the default rotation mode for a provider""" + # Import here to avoid circular imports + try: + from rotator_library.providers.provider_interface import ( + LLMProviderInterface, + ) + + return LLMProviderInterface.get_rotation_mode(provider) + except ImportError: + # Fallback defaults if import fails + if provider.lower() == "antigravity": + return "sequential" + return "balanced" + + def get_effective_mode(self, provider: str) -> str: + """Get the effective rotation mode (configured or default)""" + configured = self.get_current_modes().get(provider.lower()) + if configured: + return configured + return self.get_default_mode(provider) + + def set_mode(self, provider: str, mode: str): + """Set rotation mode for a provider""" + if mode.lower() not in self.VALID_MODES: + raise ValueError( + f"Invalid rotation mode: {mode}. Must be one of {self.VALID_MODES}" + ) + key = f"ROTATION_MODE_{provider.upper()}" + self.settings.set(key, mode.lower()) + + def remove_mode(self, provider: str): + """Remove rotation mode (reset to provider default)""" + key = f"ROTATION_MODE_{provider.upper()}" + self.settings.remove(key) + + # ============================================================================= # PROVIDER-SPECIFIC SETTINGS DEFINITIONS # ============================================================================= @@ -294,24 +350,26 @@ def remove_limit(self, provider: str): class ProviderSettingsManager: """Manages provider-specific configuration settings""" - + def __init__(self, settings: AdvancedSettings): self.settings = settings - + def get_available_providers(self) -> List[str]: """Get list of providers with specific settings available""" return list(PROVIDER_SETTINGS_MAP.keys()) - - def get_provider_settings_definitions(self, provider: str) -> Dict[str, Dict[str, Any]]: + + def get_provider_settings_definitions( + self, provider: str + ) -> Dict[str, Dict[str, Any]]: """Get settings definitions for a provider""" return PROVIDER_SETTINGS_MAP.get(provider, {}) - + def get_current_value(self, key: str, definition: Dict[str, Any]) -> Any: """Get current value of a setting from environment""" env_value = os.getenv(key) if env_value is None: return definition.get("default") - + setting_type = definition.get("type", "str") try: if setting_type == "bool": @@ -322,7 +380,7 @@ def get_current_value(self, key: str, definition: Dict[str, Any]) -> Any: return env_value except (ValueError, AttributeError): return definition.get("default") - + def get_all_current_values(self, provider: str) -> Dict[str, Any]: """Get all current values for a provider""" definitions = self.get_provider_settings_definitions(provider) @@ -330,7 +388,7 @@ def get_all_current_values(self, provider: str) -> Dict[str, Any]: for key, definition in definitions.items(): values[key] = self.get_current_value(key, definition) return values - + def set_value(self, key: str, value: Any, definition: Dict[str, Any]): """Set a setting value, converting to string for .env storage""" setting_type = definition.get("type", "str") @@ -339,11 +397,11 @@ def set_value(self, key: str, value: Any, definition: Dict[str, Any]): else: str_value = str(value) self.settings.set(key, str_value) - + def reset_to_default(self, key: str): """Remove a setting to reset it to default""" self.settings.remove(key) - + def get_modified_settings(self, provider: str) -> Dict[str, Any]: """Get settings that differ from defaults""" definitions = self.get_provider_settings_definitions(provider) @@ -358,80 +416,96 @@ def get_modified_settings(self, provider: str) -> Dict[str, Any]: class SettingsTool: """Main settings tool TUI""" - + def __init__(self): self.console = Console() self.settings = AdvancedSettings() self.provider_mgr = CustomProviderManager(self.settings) self.model_mgr = ModelDefinitionManager(self.settings) self.concurrency_mgr = ConcurrencyManager(self.settings) + self.rotation_mgr = RotationModeManager(self.settings) self.provider_settings_mgr = ProviderSettingsManager(self.settings) self.running = True - + def get_available_providers(self) -> List[str]: """Get list of providers that have credentials configured""" env_file = Path.cwd() / ".env" providers = set() - + # Scan for providers with API keys from local .env if env_file.exists(): try: - with open(env_file, 'r', encoding='utf-8') as f: + with open(env_file, "r", encoding="utf-8") as f: for line in f: line = line.strip() - if "_API_KEY" in line and "PROXY_API_KEY" not in line and "=" in line: + if ( + "_API_KEY" in line + and "PROXY_API_KEY" not in line + and "=" in line + ): provider = line.split("_API_KEY")[0].strip().lower() providers.add(provider) except (IOError, OSError): pass - + # Also check for OAuth providers from files oauth_dir = Path("oauth_credentials") if oauth_dir.exists(): for file in oauth_dir.glob("*_oauth_*.json"): provider = file.name.split("_oauth_")[0] providers.add(provider) - + return sorted(list(providers)) def run(self): """Main loop""" while self.running: self.show_main_menu() - + def show_main_menu(self): """Display settings categories""" clear_screen() - - self.console.print(Panel.fit( - "[bold cyan]🔧 Advanced Settings Configuration[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + "[bold cyan]🔧 Advanced Settings Configuration[/bold cyan]", + border_style="cyan", + ) + ) + self.console.print() self.console.print("[bold]⚙️ Configuration Categories[/bold]") self.console.print() self.console.print(" 1. 🌐 Custom Provider API Bases") self.console.print(" 2. 📦 Provider Model Definitions") self.console.print(" 3. ⚡ Concurrency Limits") - self.console.print(" 4. 🔬 Provider-Specific Settings") - self.console.print(" 5. 💾 Save & Exit") - self.console.print(" 6. 🚫 Exit Without Saving") - + self.console.print(" 4. 🔄 Rotation Modes") + self.console.print(" 5. 🔬 Provider-Specific Settings") + self.console.print(" 6. 💾 Save & Exit") + self.console.print(" 7. 🚫 Exit Without Saving") + self.console.print() self.console.print("━" * 70) - + if self.settings.has_pending(): - self.console.print("[yellow]ℹ️ Changes are pending until you select \"Save & Exit\"[/yellow]") + self.console.print( + '[yellow]ℹ️ Changes are pending until you select "Save & Exit"[/yellow]' + ) else: self.console.print("[dim]ℹ️ No pending changes[/dim]") - + self.console.print() - self.console.print("[dim]⚠️ Model filters not supported - edit .env for IGNORE_MODELS_* / WHITELIST_MODELS_*[/dim]") + self.console.print( + "[dim]⚠️ Model filters not supported - edit .env for IGNORE_MODELS_* / WHITELIST_MODELS_*[/dim]" + ) self.console.print() - - choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5", "6"], show_choices=False) - + + choice = Prompt.ask( + "Select option", + choices=["1", "2", "3", "4", "5", "6", "7"], + show_choices=False, + ) + if choice == "1": self.manage_custom_providers() elif choice == "2": @@ -439,34 +513,38 @@ def show_main_menu(self): elif choice == "3": self.manage_concurrency_limits() elif choice == "4": - self.manage_provider_settings() + self.manage_rotation_modes() elif choice == "5": - self.save_and_exit() + self.manage_provider_settings() elif choice == "6": + self.save_and_exit() + elif choice == "7": self.exit_without_saving() - + def manage_custom_providers(self): """Manage custom provider API bases""" while True: clear_screen() - + providers = self.provider_mgr.get_current_providers() - - self.console.print(Panel.fit( - "[bold cyan]🌐 Custom Provider API Bases[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + "[bold cyan]🌐 Custom Provider API Bases[/bold cyan]", + border_style="cyan", + ) + ) + self.console.print() self.console.print("[bold]📋 Configured Custom Providers[/bold]") self.console.print("━" * 70) - + if providers: for name, base in providers.items(): self.console.print(f" • {name:15} {base}") else: self.console.print(" [dim]No custom providers configured[/dim]") - + self.console.print() self.console.print("━" * 70) self.console.print() @@ -476,94 +554,116 @@ def manage_custom_providers(self): self.console.print(" 2. ✏️ Edit Existing Provider") self.console.print(" 3. 🗑️ Remove Provider") self.console.print(" 4. ↩️ Back to Settings Menu") - + self.console.print() self.console.print("━" * 70) self.console.print() - - choice = Prompt.ask("Select option", choices=["1", "2", "3", "4"], show_choices=False) - + + choice = Prompt.ask( + "Select option", choices=["1", "2", "3", "4"], show_choices=False + ) + if choice == "1": name = Prompt.ask("Provider name (e.g., 'opencode')").strip().lower() if name: api_base = Prompt.ask("API Base URL").strip() if api_base: self.provider_mgr.add_provider(name, api_base) - self.console.print(f"\n[green]✅ Custom provider '{name}' configured![/green]") - self.console.print(f" To use: set {name.upper()}_API_KEY in credentials") + self.console.print( + f"\n[green]✅ Custom provider '{name}' configured![/green]" + ) + self.console.print( + f" To use: set {name.upper()}_API_KEY in credentials" + ) input("\nPress Enter to continue...") - + elif choice == "2": if not providers: self.console.print("\n[yellow]No providers to edit[/yellow]") input("\nPress Enter to continue...") continue - + # Show numbered list self.console.print("\n[bold]Select provider to edit:[/bold]") providers_list = list(providers.keys()) for idx, prov in enumerate(providers_list, 1): self.console.print(f" {idx}. {prov}") - - choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers_list) + 1)]) + + choice_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(providers_list) + 1)], + ) name = providers_list[choice_idx - 1] current_base = providers.get(name, "") - + self.console.print(f"\nCurrent API Base: {current_base}") - new_base = Prompt.ask("New API Base [press Enter to keep current]", default=current_base).strip() - + new_base = Prompt.ask( + "New API Base [press Enter to keep current]", default=current_base + ).strip() + if new_base and new_base != current_base: self.provider_mgr.edit_provider(name, new_base) - self.console.print(f"\n[green]✅ Custom provider '{name}' updated![/green]") + self.console.print( + f"\n[green]✅ Custom provider '{name}' updated![/green]" + ) else: self.console.print("\n[yellow]No changes made[/yellow]") input("\nPress Enter to continue...") - + elif choice == "3": if not providers: self.console.print("\n[yellow]No providers to remove[/yellow]") input("\nPress Enter to continue...") continue - + # Show numbered list self.console.print("\n[bold]Select provider to remove:[/bold]") providers_list = list(providers.keys()) for idx, prov in enumerate(providers_list, 1): self.console.print(f" {idx}. {prov}") - - choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers_list) + 1)]) + + choice_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(providers_list) + 1)], + ) name = providers_list[choice_idx - 1] - + if Confirm.ask(f"Remove '{name}'?"): self.provider_mgr.remove_provider(name) - self.console.print(f"\n[green]✅ Provider '{name}' removed![/green]") + self.console.print( + f"\n[green]✅ Provider '{name}' removed![/green]" + ) input("\nPress Enter to continue...") - + elif choice == "4": break - + def manage_model_definitions(self): """Manage provider model definitions""" while True: clear_screen() - + all_providers = self.model_mgr.get_all_providers_with_models() - - self.console.print(Panel.fit( - "[bold cyan]📦 Provider Model Definitions[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + "[bold cyan]📦 Provider Model Definitions[/bold cyan]", + border_style="cyan", + ) + ) + self.console.print() self.console.print("[bold]📋 Configured Provider Models[/bold]") self.console.print("━" * 70) - + if all_providers: for provider, count in all_providers.items(): - self.console.print(f" • {provider:15} {count} model{'s' if count > 1 else ''}") + self.console.print( + f" • {provider:15} {count} model{'s' if count > 1 else ''}" + ) else: self.console.print(" [dim]No model definitions configured[/dim]") - + self.console.print() self.console.print("━" * 70) self.console.print() @@ -574,13 +674,15 @@ def manage_model_definitions(self): self.console.print(" 3. 👁️ View Provider Models") self.console.print(" 4. 🗑️ Remove Provider Models") self.console.print(" 5. ↩️ Back to Settings Menu") - + self.console.print() self.console.print("━" * 70) self.console.print() - - choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5"], show_choices=False) - + + choice = Prompt.ask( + "Select option", choices=["1", "2", "3", "4", "5"], show_choices=False + ) + if choice == "1": self.add_model_definitions() elif choice == "2": @@ -600,57 +702,71 @@ def manage_model_definitions(self): self.console.print("\n[yellow]No providers to remove[/yellow]") input("\nPress Enter to continue...") continue - + # Show numbered list - self.console.print("\n[bold]Select provider to remove models from:[/bold]") + self.console.print( + "\n[bold]Select provider to remove models from:[/bold]" + ) providers_list = list(all_providers.keys()) for idx, prov in enumerate(providers_list, 1): self.console.print(f" {idx}. {prov}") - - choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers_list) + 1)]) + + choice_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(providers_list) + 1)], + ) provider = providers_list[choice_idx - 1] - + if Confirm.ask(f"Remove all model definitions for '{provider}'?"): self.model_mgr.remove_models(provider) - self.console.print(f"\n[green]✅ Model definitions removed for '{provider}'![/green]") + self.console.print( + f"\n[green]✅ Model definitions removed for '{provider}'![/green]" + ) input("\nPress Enter to continue...") elif choice == "5": break - + def add_model_definitions(self): """Add model definitions for a provider""" # Get available providers from credentials available_providers = self.get_available_providers() - + if not available_providers: - self.console.print("\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]") + self.console.print( + "\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]" + ) input("\nPress Enter to continue...") return - + # Show provider selection menu self.console.print("\n[bold]Select provider:[/bold]") for idx, prov in enumerate(available_providers, 1): self.console.print(f" {idx}. {prov}") - self.console.print(f" {len(available_providers) + 1}. Enter custom provider name") - - choice = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(available_providers) + 2)]) - + self.console.print( + f" {len(available_providers) + 1}. Enter custom provider name" + ) + + choice = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(available_providers) + 2)], + ) + if choice == len(available_providers) + 1: provider = Prompt.ask("Provider name").strip().lower() else: provider = available_providers[choice - 1] - + if not provider: return - + self.console.print("\nHow would you like to define models?") self.console.print(" 1. Simple list (names only)") self.console.print(" 2. Advanced (names with IDs and options)") - + mode = Prompt.ask("Select mode", choices=["1", "2"], show_choices=False) - + models = {} - + if mode == "1": # Simple mode while True: @@ -667,13 +783,19 @@ def add_model_definitions(self): break if name: model_def = {} - model_id = Prompt.ask(f"Model ID [press Enter to use '{name}']", default=name).strip() + model_id = Prompt.ask( + f"Model ID [press Enter to use '{name}']", default=name + ).strip() if model_id and model_id != name: model_def["id"] = model_id - + # Optional: model options - if Confirm.ask("Add model options (e.g., temperature limits)?", default=False): - self.console.print("\nEnter options as key=value pairs (one per line, 'done' to finish):") + if Confirm.ask( + "Add model options (e.g., temperature limits)?", default=False + ): + self.console.print( + "\nEnter options as key=value pairs (one per line, 'done' to finish):" + ) options = {} while True: opt = Prompt.ask("Option").strip() @@ -690,121 +812,143 @@ def add_model_definitions(self): options[key.strip()] = value if options: model_def["options"] = options - + models[name] = model_def - + if models: self.model_mgr.set_models(provider, models) - self.console.print(f"\n[green]✅ Model definitions saved for '{provider}'![/green]") + self.console.print( + f"\n[green]✅ Model definitions saved for '{provider}'![/green]" + ) else: self.console.print("\n[yellow]No models added[/yellow]") - + input("\nPress Enter to continue...") - + def edit_model_definitions(self, providers: List[str]): """Edit existing model definitions""" # Show numbered list self.console.print("\n[bold]Select provider to edit:[/bold]") for idx, prov in enumerate(providers, 1): self.console.print(f" {idx}. {prov}") - - choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers) + 1)]) + + choice_idx = IntPrompt.ask( + "Select option", choices=[str(i) for i in range(1, len(providers) + 1)] + ) provider = providers[choice_idx - 1] - + current_models = self.model_mgr.get_current_provider_models(provider) if not current_models: self.console.print(f"\n[yellow]No models found for '{provider}'[/yellow]") input("\nPress Enter to continue...") return - + # Convert to dict if list if isinstance(current_models, list): current_models = {m: {} for m in current_models} - + while True: clear_screen() self.console.print(f"[bold]Editing models for: {provider}[/bold]\n") self.console.print("Current models:") for i, (name, definition) in enumerate(current_models.items(), 1): - model_id = definition.get("id", name) if isinstance(definition, dict) else name + model_id = ( + definition.get("id", name) if isinstance(definition, dict) else name + ) self.console.print(f" {i}. {name} (ID: {model_id})") - + self.console.print("\nOptions:") self.console.print(" 1. Add new model") self.console.print(" 2. Edit existing model") self.console.print(" 3. Remove model") self.console.print(" 4. Done") - - choice = Prompt.ask("\nSelect option", choices=["1", "2", "3", "4"], show_choices=False) - + + choice = Prompt.ask( + "\nSelect option", choices=["1", "2", "3", "4"], show_choices=False + ) + if choice == "1": name = Prompt.ask("New model name").strip() if name and name not in current_models: model_id = Prompt.ask("Model ID", default=name).strip() current_models[name] = {"id": model_id} if model_id != name else {} - + elif choice == "2": # Show numbered list models_list = list(current_models.keys()) self.console.print("\n[bold]Select model to edit:[/bold]") for idx, model_name in enumerate(models_list, 1): self.console.print(f" {idx}. {model_name}") - - model_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(models_list) + 1)]) + + model_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(models_list) + 1)], + ) name = models_list[model_idx - 1] - + current_def = current_models[name] - current_id = current_def.get("id", name) if isinstance(current_def, dict) else name - + current_id = ( + current_def.get("id", name) + if isinstance(current_def, dict) + else name + ) + new_id = Prompt.ask("Model ID", default=current_id).strip() current_models[name] = {"id": new_id} if new_id != name else {} - + elif choice == "3": # Show numbered list models_list = list(current_models.keys()) self.console.print("\n[bold]Select model to remove:[/bold]") for idx, model_name in enumerate(models_list, 1): self.console.print(f" {idx}. {model_name}") - - model_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(models_list) + 1)]) + + model_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(models_list) + 1)], + ) name = models_list[model_idx - 1] - + if Confirm.ask(f"Remove '{name}'?"): del current_models[name] - + elif choice == "4": break - + if current_models: self.model_mgr.set_models(provider, current_models) self.console.print(f"\n[green]✅ Models updated for '{provider}'![/green]") else: - self.console.print("\n[yellow]No models left - removing definition[/yellow]") + self.console.print( + "\n[yellow]No models left - removing definition[/yellow]" + ) self.model_mgr.remove_models(provider) - + input("\nPress Enter to continue...") - + def view_model_definitions(self, providers: List[str]): """View model definitions for a provider""" # Show numbered list self.console.print("\n[bold]Select provider to view:[/bold]") for idx, prov in enumerate(providers, 1): self.console.print(f" {idx}. {prov}") - - choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(providers) + 1)]) + + choice_idx = IntPrompt.ask( + "Select option", choices=[str(i) for i in range(1, len(providers) + 1)] + ) provider = providers[choice_idx - 1] - + models = self.model_mgr.get_current_provider_models(provider) if not models: self.console.print(f"\n[yellow]No models found for '{provider}'[/yellow]") input("\nPress Enter to continue...") return - + clear_screen() self.console.print(f"[bold]Provider: {provider}[/bold]\n") self.console.print("[bold]📦 Configured Models:[/bold]") self.console.print("━" * 50) - + # Handle both dict and list formats if isinstance(models, dict): for name, definition in models.items(): @@ -822,74 +966,88 @@ def view_model_definitions(self, providers: List[str]): for name in models: self.console.print(f" Name: {name}") self.console.print() - + input("Press Enter to return...") - + def manage_provider_settings(self): """Manage provider-specific settings (Antigravity, Gemini CLI)""" while True: clear_screen() - + available_providers = self.provider_settings_mgr.get_available_providers() - - self.console.print(Panel.fit( - "[bold cyan]🔬 Provider-Specific Settings[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + "[bold cyan]🔬 Provider-Specific Settings[/bold cyan]", + border_style="cyan", + ) + ) + self.console.print() - self.console.print("[bold]📋 Available Providers with Custom Settings[/bold]") + self.console.print( + "[bold]📋 Available Providers with Custom Settings[/bold]" + ) self.console.print("━" * 70) - + for provider in available_providers: modified = self.provider_settings_mgr.get_modified_settings(provider) - status = f"[yellow]{len(modified)} modified[/yellow]" if modified else "[dim]defaults[/dim]" + status = ( + f"[yellow]{len(modified)} modified[/yellow]" + if modified + else "[dim]defaults[/dim]" + ) display_name = provider.replace("_", " ").title() self.console.print(f" • {display_name:20} {status}") - + self.console.print() self.console.print("━" * 70) self.console.print() self.console.print("[bold]⚙️ Select Provider to Configure[/bold]") self.console.print() - + for idx, provider in enumerate(available_providers, 1): display_name = provider.replace("_", " ").title() self.console.print(f" {idx}. {display_name}") - self.console.print(f" {len(available_providers) + 1}. ↩️ Back to Settings Menu") - + self.console.print( + f" {len(available_providers) + 1}. ↩️ Back to Settings Menu" + ) + self.console.print() self.console.print("━" * 70) self.console.print() - + choices = [str(i) for i in range(1, len(available_providers) + 2)] choice = Prompt.ask("Select option", choices=choices, show_choices=False) choice_idx = int(choice) - + if choice_idx == len(available_providers) + 1: break - + provider = available_providers[choice_idx - 1] self._manage_single_provider_settings(provider) - + def _manage_single_provider_settings(self, provider: str): """Manage settings for a single provider""" while True: clear_screen() - + display_name = provider.replace("_", " ").title() - definitions = self.provider_settings_mgr.get_provider_settings_definitions(provider) + definitions = self.provider_settings_mgr.get_provider_settings_definitions( + provider + ) current_values = self.provider_settings_mgr.get_all_current_values(provider) - - self.console.print(Panel.fit( - f"[bold cyan]🔬 {display_name} Settings[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + f"[bold cyan]🔬 {display_name} Settings[/bold cyan]", + border_style="cyan", + ) + ) + self.console.print() self.console.print("[bold]📋 Current Settings[/bold]") self.console.print("━" * 70) - + # Display all settings with current values settings_list = list(definitions.keys()) for idx, key in enumerate(settings_list, 1): @@ -898,25 +1056,35 @@ def _manage_single_provider_settings(self, provider: str): default = definition.get("default") setting_type = definition.get("type", "str") description = definition.get("description", "") - + # Format value display if setting_type == "bool": - value_display = "[green]✓ Enabled[/green]" if current else "[red]✗ Disabled[/red]" + value_display = ( + "[green]✓ Enabled[/green]" + if current + else "[red]✗ Disabled[/red]" + ) elif setting_type == "int": value_display = f"[cyan]{current}[/cyan]" else: - value_display = f"[cyan]{current or '(not set)'}[/cyan]" if current else "[dim](not set)[/dim]" - + value_display = ( + f"[cyan]{current or '(not set)'}[/cyan]" + if current + else "[dim](not set)[/dim]" + ) + # Check if modified from default modified = current != default mod_marker = "[yellow]*[/yellow]" if modified else " " - + # Short key name for display (strip provider prefix) short_key = key.replace(f"{provider.upper()}_", "") - - self.console.print(f" {mod_marker}{idx:2}. {short_key:35} {value_display}") + + self.console.print( + f" {mod_marker}{idx:2}. {short_key:35} {value_display}" + ) self.console.print(f" [dim]{description}[/dim]") - + self.console.print() self.console.print("━" * 70) self.console.print("[dim]* = modified from default[/dim]") @@ -927,13 +1095,17 @@ def _manage_single_provider_settings(self, provider: str): self.console.print(" R. 🔄 Reset Setting to Default") self.console.print(" A. 🔄 Reset All to Defaults") self.console.print(" B. ↩️ Back to Provider Selection") - + self.console.print() self.console.print("━" * 70) self.console.print() - - choice = Prompt.ask("Select action", choices=["e", "r", "a", "b", "E", "R", "A", "B"], show_choices=False).lower() - + + choice = Prompt.ask( + "Select action", + choices=["e", "r", "a", "b", "E", "R", "A", "B"], + show_choices=False, + ).lower() + if choice == "b": break elif choice == "e": @@ -942,26 +1114,31 @@ def _manage_single_provider_settings(self, provider: str): self._reset_provider_setting(provider, settings_list, definitions) elif choice == "a": self._reset_all_provider_settings(provider, settings_list) - - def _edit_provider_setting(self, provider: str, settings_list: List[str], definitions: Dict[str, Dict[str, Any]]): + + def _edit_provider_setting( + self, + provider: str, + settings_list: List[str], + definitions: Dict[str, Dict[str, Any]], + ): """Edit a single provider setting""" self.console.print("\n[bold]Select setting number to edit:[/bold]") - + choices = [str(i) for i in range(1, len(settings_list) + 1)] choice = IntPrompt.ask("Setting number", choices=choices) key = settings_list[choice - 1] definition = definitions[key] - + current = self.provider_settings_mgr.get_current_value(key, definition) default = definition.get("default") setting_type = definition.get("type", "str") short_key = key.replace(f"{provider.upper()}_", "") - + self.console.print(f"\n[bold]Editing: {short_key}[/bold]") self.console.print(f"Current value: [cyan]{current}[/cyan]") self.console.print(f"Default value: [dim]{default}[/dim]") self.console.print(f"Type: {setting_type}") - + if setting_type == "bool": new_value = Confirm.ask("\nEnable this setting?", default=current) self.provider_settings_mgr.set_value(key, new_value, definition) @@ -972,71 +1149,252 @@ def _edit_provider_setting(self, provider: str, settings_list: List[str], defini self.provider_settings_mgr.set_value(key, new_value, definition) self.console.print(f"\n[green]✅ {short_key} set to {new_value}![/green]") else: - new_value = Prompt.ask("\nNew value", default=str(current) if current else "").strip() + new_value = Prompt.ask( + "\nNew value", default=str(current) if current else "" + ).strip() if new_value: self.provider_settings_mgr.set_value(key, new_value, definition) self.console.print(f"\n[green]✅ {short_key} updated![/green]") else: self.console.print("\n[yellow]No changes made[/yellow]") - + input("\nPress Enter to continue...") - - def _reset_provider_setting(self, provider: str, settings_list: List[str], definitions: Dict[str, Dict[str, Any]]): + + def _reset_provider_setting( + self, + provider: str, + settings_list: List[str], + definitions: Dict[str, Dict[str, Any]], + ): """Reset a single provider setting to default""" self.console.print("\n[bold]Select setting number to reset:[/bold]") - + choices = [str(i) for i in range(1, len(settings_list) + 1)] choice = IntPrompt.ask("Setting number", choices=choices) key = settings_list[choice - 1] definition = definitions[key] - + default = definition.get("default") short_key = key.replace(f"{provider.upper()}_", "") - + if Confirm.ask(f"\nReset {short_key} to default ({default})?"): self.provider_settings_mgr.reset_to_default(key) self.console.print(f"\n[green]✅ {short_key} reset to default![/green]") else: self.console.print("\n[yellow]No changes made[/yellow]") - + input("\nPress Enter to continue...") - + def _reset_all_provider_settings(self, provider: str, settings_list: List[str]): """Reset all provider settings to defaults""" display_name = provider.replace("_", " ").title() - - if Confirm.ask(f"\n[bold red]Reset ALL {display_name} settings to defaults?[/bold red]"): + + if Confirm.ask( + f"\n[bold red]Reset ALL {display_name} settings to defaults?[/bold red]" + ): for key in settings_list: self.provider_settings_mgr.reset_to_default(key) - self.console.print(f"\n[green]✅ All {display_name} settings reset to defaults![/green]") + self.console.print( + f"\n[green]✅ All {display_name} settings reset to defaults![/green]" + ) else: self.console.print("\n[yellow]No changes made[/yellow]") - + input("\nPress Enter to continue...") - + + def manage_rotation_modes(self): + """Manage credential rotation modes (sequential vs balanced)""" + while True: + clear_screen() + + modes = self.rotation_mgr.get_current_modes() + available_providers = self.get_available_providers() + + self.console.print( + Panel.fit( + "[bold cyan]🔄 Credential Rotation Mode Configuration[/bold cyan]", + border_style="cyan", + ) + ) + + self.console.print() + self.console.print("[bold]📋 Rotation Modes Explained[/bold]") + self.console.print("━" * 70) + self.console.print( + " [cyan]balanced[/cyan] - Rotate credentials evenly across requests (default)" + ) + self.console.print( + " [cyan]sequential[/cyan] - Use one credential until exhausted (429), then switch" + ) + self.console.print() + self.console.print("[bold]📋 Current Rotation Mode Settings[/bold]") + self.console.print("━" * 70) + + if modes: + for provider, mode in modes.items(): + default_mode = self.rotation_mgr.get_default_mode(provider) + is_custom = mode != default_mode + marker = "[yellow]*[/yellow]" if is_custom else " " + mode_display = ( + f"[green]{mode}[/green]" + if mode == "sequential" + else f"[blue]{mode}[/blue]" + ) + self.console.print(f" {marker}• {provider:20} {mode_display}") + + # Show providers with default modes + providers_with_defaults = [p for p in available_providers if p not in modes] + if providers_with_defaults: + self.console.print() + self.console.print("[dim]Providers using default modes:[/dim]") + for provider in providers_with_defaults: + default_mode = self.rotation_mgr.get_default_mode(provider) + mode_display = ( + f"[green]{default_mode}[/green]" + if default_mode == "sequential" + else f"[blue]{default_mode}[/blue]" + ) + self.console.print( + f" • {provider:20} {mode_display} [dim](default)[/dim]" + ) + + self.console.print() + self.console.print("━" * 70) + self.console.print( + "[dim]* = custom setting (differs from provider default)[/dim]" + ) + self.console.print() + self.console.print("[bold]⚙️ Actions[/bold]") + self.console.print() + self.console.print(" 1. ➕ Set Rotation Mode for Provider") + self.console.print(" 2. 🗑️ Reset to Provider Default") + self.console.print(" 3. ↩️ Back to Settings Menu") + + self.console.print() + self.console.print("━" * 70) + self.console.print() + + choice = Prompt.ask( + "Select option", choices=["1", "2", "3"], show_choices=False + ) + + if choice == "1": + if not available_providers: + self.console.print( + "\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]" + ) + input("\nPress Enter to continue...") + continue + + # Show provider selection menu + self.console.print("\n[bold]Select provider:[/bold]") + for idx, prov in enumerate(available_providers, 1): + current_mode = self.rotation_mgr.get_effective_mode(prov) + mode_display = ( + f"[green]{current_mode}[/green]" + if current_mode == "sequential" + else f"[blue]{current_mode}[/blue]" + ) + self.console.print(f" {idx}. {prov} ({mode_display})") + self.console.print( + f" {len(available_providers) + 1}. Enter custom provider name" + ) + + choice_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(available_providers) + 2)], + ) + + if choice_idx == len(available_providers) + 1: + provider = Prompt.ask("Provider name").strip().lower() + else: + provider = available_providers[choice_idx - 1] + + if provider: + current_mode = self.rotation_mgr.get_effective_mode(provider) + self.console.print( + f"\nCurrent mode for {provider}: [cyan]{current_mode}[/cyan]" + ) + self.console.print("\nSelect new rotation mode:") + self.console.print( + " 1. [blue]balanced[/blue] - Rotate credentials evenly" + ) + self.console.print( + " 2. [green]sequential[/green] - Use until exhausted" + ) + + mode_choice = Prompt.ask( + "Select mode", choices=["1", "2"], show_choices=False + ) + new_mode = "balanced" if mode_choice == "1" else "sequential" + + self.rotation_mgr.set_mode(provider, new_mode) + self.console.print( + f"\n[green]✅ Rotation mode for '{provider}' set to {new_mode}![/green]" + ) + input("\nPress Enter to continue...") + + elif choice == "2": + if not modes: + self.console.print( + "\n[yellow]No custom rotation modes to reset[/yellow]" + ) + input("\nPress Enter to continue...") + continue + + # Show numbered list + self.console.print( + "\n[bold]Select provider to reset to default:[/bold]" + ) + modes_list = list(modes.keys()) + for idx, prov in enumerate(modes_list, 1): + default_mode = self.rotation_mgr.get_default_mode(prov) + self.console.print( + f" {idx}. {prov} (will reset to: {default_mode})" + ) + + choice_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(modes_list) + 1)], + ) + provider = modes_list[choice_idx - 1] + default_mode = self.rotation_mgr.get_default_mode(provider) + + if Confirm.ask(f"Reset '{provider}' to default mode ({default_mode})?"): + self.rotation_mgr.remove_mode(provider) + self.console.print( + f"\n[green]✅ Rotation mode for '{provider}' reset to default ({default_mode})![/green]" + ) + input("\nPress Enter to continue...") + + elif choice == "3": + break + def manage_concurrency_limits(self): """Manage concurrency limits""" while True: clear_screen() - + limits = self.concurrency_mgr.get_current_limits() - - self.console.print(Panel.fit( - "[bold cyan]⚡ Concurrency Limits Configuration[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + "[bold cyan]⚡ Concurrency Limits Configuration[/bold cyan]", + border_style="cyan", + ) + ) + self.console.print() self.console.print("[bold]📋 Current Concurrency Settings[/bold]") self.console.print("━" * 70) - + if limits: for provider, limit in limits.items(): self.console.print(f" • {provider:15} {limit} requests/key") self.console.print(f" • Default: 1 request/key (all others)") else: self.console.print(" • Default: 1 request/key (all providers)") - + self.console.print() self.console.print("━" * 70) self.console.print() @@ -1046,96 +1404,128 @@ def manage_concurrency_limits(self): self.console.print(" 2. ✏️ Edit Existing Limit") self.console.print(" 3. 🗑️ Remove Limit (reset to default)") self.console.print(" 4. ↩️ Back to Settings Menu") - + self.console.print() self.console.print("━" * 70) self.console.print() - - choice = Prompt.ask("Select option", choices=["1", "2", "3", "4"], show_choices=False) - + + choice = Prompt.ask( + "Select option", choices=["1", "2", "3", "4"], show_choices=False + ) + if choice == "1": # Get available providers available_providers = self.get_available_providers() - + if not available_providers: - self.console.print("\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]") + self.console.print( + "\n[yellow]No providers with credentials found. Please add credentials first.[/yellow]" + ) input("\nPress Enter to continue...") continue - + # Show provider selection menu self.console.print("\n[bold]Select provider:[/bold]") for idx, prov in enumerate(available_providers, 1): self.console.print(f" {idx}. {prov}") - self.console.print(f" {len(available_providers) + 1}. Enter custom provider name") - - choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(available_providers) + 2)]) - + self.console.print( + f" {len(available_providers) + 1}. Enter custom provider name" + ) + + choice_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(available_providers) + 2)], + ) + if choice_idx == len(available_providers) + 1: provider = Prompt.ask("Provider name").strip().lower() else: provider = available_providers[choice_idx - 1] - + if provider: - limit = IntPrompt.ask("Max concurrent requests per key (1-100)", default=1) + limit = IntPrompt.ask( + "Max concurrent requests per key (1-100)", default=1 + ) if 1 <= limit <= 100: self.concurrency_mgr.set_limit(provider, limit) - self.console.print(f"\n[green]✅ Concurrency limit set for '{provider}': {limit} requests/key[/green]") + self.console.print( + f"\n[green]✅ Concurrency limit set for '{provider}': {limit} requests/key[/green]" + ) else: - self.console.print("\n[red]❌ Limit must be between 1-100[/red]") + self.console.print( + "\n[red]❌ Limit must be between 1-100[/red]" + ) input("\nPress Enter to continue...") - + elif choice == "2": if not limits: self.console.print("\n[yellow]No limits to edit[/yellow]") input("\nPress Enter to continue...") continue - + # Show numbered list self.console.print("\n[bold]Select provider to edit:[/bold]") limits_list = list(limits.keys()) for idx, prov in enumerate(limits_list, 1): self.console.print(f" {idx}. {prov}") - - choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(limits_list) + 1)]) + + choice_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(limits_list) + 1)], + ) provider = limits_list[choice_idx - 1] current_limit = limits.get(provider, 1) - + self.console.print(f"\nCurrent limit: {current_limit} requests/key") - new_limit = IntPrompt.ask("New limit (1-100) [press Enter to keep current]", default=current_limit) - + new_limit = IntPrompt.ask( + "New limit (1-100) [press Enter to keep current]", + default=current_limit, + ) + if 1 <= new_limit <= 100: if new_limit != current_limit: self.concurrency_mgr.set_limit(provider, new_limit) - self.console.print(f"\n[green]✅ Concurrency limit updated for '{provider}': {new_limit} requests/key[/green]") + self.console.print( + f"\n[green]✅ Concurrency limit updated for '{provider}': {new_limit} requests/key[/green]" + ) else: self.console.print("\n[yellow]No changes made[/yellow]") else: self.console.print("\n[red]Limit must be between 1-100[/red]") input("\nPress Enter to continue...") - + elif choice == "3": if not limits: self.console.print("\n[yellow]No limits to remove[/yellow]") input("\nPress Enter to continue...") continue - + # Show numbered list - self.console.print("\n[bold]Select provider to remove limit from:[/bold]") + self.console.print( + "\n[bold]Select provider to remove limit from:[/bold]" + ) limits_list = list(limits.keys()) for idx, prov in enumerate(limits_list, 1): self.console.print(f" {idx}. {prov}") - - choice_idx = IntPrompt.ask("Select option", choices=[str(i) for i in range(1, len(limits_list) + 1)]) + + choice_idx = IntPrompt.ask( + "Select option", + choices=[str(i) for i in range(1, len(limits_list) + 1)], + ) provider = limits_list[choice_idx - 1] - - if Confirm.ask(f"Remove concurrency limit for '{provider}' (reset to default 1)?"): + + if Confirm.ask( + f"Remove concurrency limit for '{provider}' (reset to default 1)?" + ): self.concurrency_mgr.remove_limit(provider) - self.console.print(f"\n[green]✅ Limit removed for '{provider}' - using default (1 request/key)[/green]") + self.console.print( + f"\n[green]✅ Limit removed for '{provider}' - using default (1 request/key)[/green]" + ) input("\nPress Enter to continue...") - + elif choice == "4": break - + def save_and_exit(self): """Save pending changes and exit""" if self.settings.has_pending(): @@ -1150,9 +1540,9 @@ def save_and_exit(self): else: self.console.print("\n[dim]No changes to save[/dim]") input("\nPress Enter to return to launcher...") - + self.running = False - + def exit_without_saving(self): """Exit without saving""" if self.settings.has_pending(): diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index befa39ed..179cd09b 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -139,8 +139,28 @@ def __init__( self.max_retries = max_retries self.global_timeout = global_timeout self.abort_on_callback_error = abort_on_callback_error + + # Build provider rotation modes map + # Each provider can specify its preferred rotation mode ("balanced" or "sequential") + provider_rotation_modes = {} + for provider in self.all_credentials.keys(): + provider_class = self._provider_plugins.get(provider) + if provider_class and hasattr(provider_class, "get_rotation_mode"): + # Use class method to get rotation mode (checks env var + class default) + mode = provider_class.get_rotation_mode(provider) + else: + # Fallback: check environment variable directly + env_key = f"ROTATION_MODE_{provider.upper()}" + mode = os.getenv(env_key, "balanced") + + provider_rotation_modes[provider] = mode + if mode != "balanced": + lib_logger.info(f"Provider '{provider}' using rotation mode: {mode}") + self.usage_manager = UsageManager( - file_path=usage_file_path, rotation_tolerance=rotation_tolerance + file_path=usage_file_path, + rotation_tolerance=rotation_tolerance, + provider_rotation_modes=provider_rotation_modes, ) self._model_list_cache = {} self._provider_plugins = PROVIDER_PLUGINS @@ -1070,7 +1090,7 @@ async def _execute_with_retry( if request else {}, ) - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) # Extract a clean error message for the user-facing log error_message = str(e).split("\n")[0] @@ -1114,7 +1134,7 @@ async def _execute_with_retry( if request else {}, ) - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) error_message = str(e).split("\n")[0] # Provider-level error: don't increment consecutive failures @@ -1170,7 +1190,7 @@ async def _execute_with_retry( else {}, ) - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) error_message = str(e).split("\n")[0] lib_logger.warning( @@ -1239,7 +1259,7 @@ async def _execute_with_retry( ) raise last_exception - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) error_message = str(e).split("\n")[0] lib_logger.warning( @@ -1566,7 +1586,9 @@ async def _streaming_acompletion_with_retry( last_exception = e # If the exception is our custom wrapper, unwrap the original error original_exc = getattr(e, "data", e) - classified_error = classify_error(original_exc) + classified_error = classify_error( + original_exc, provider=provider + ) error_message = str(original_exc).split("\n")[0] log_failure( @@ -1623,7 +1645,7 @@ async def _streaming_acompletion_with_retry( if request else {}, ) - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) error_message = str(e).split("\n")[0] # Provider-level error: don't increment consecutive failures @@ -1673,7 +1695,7 @@ async def _streaming_acompletion_with_retry( if request else {}, ) - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) error_message = str(e).split("\n")[0] # Record in accumulator @@ -1812,7 +1834,9 @@ async def _streaming_acompletion_with_retry( cleaned_str = None # The actual exception might be wrapped in our StreamedAPIError. original_exc = getattr(e, "data", e) - classified_error = classify_error(original_exc) + classified_error = classify_error( + original_exc, provider=provider + ) # Check if this error should trigger rotation if not should_rotate_on_error(classified_error): @@ -1939,7 +1963,7 @@ async def _streaming_acompletion_with_retry( if request else {}, ) - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) error_message_text = str(e).split("\n")[0] # Record error in accumulator (server errors are transient, not abnormal) @@ -1990,7 +2014,7 @@ async def _streaming_acompletion_with_retry( if request else {}, ) - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) error_message_text = str(e).split("\n")[0] # Record error in accumulator @@ -2232,7 +2256,7 @@ async def get_available_models(self, provider: str) -> List[str]: self._model_list_cache[provider] = final_models return final_models except Exception as e: - classified_error = classify_error(e) + classified_error = classify_error(e, provider=provider) cred_display = mask_credential(credential) lib_logger.debug( f"Failed to get models for provider {provider} with credential {cred_display}: {classified_error.error_type}. Trying next credential." diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index 3676d44c..51692c49 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -1,6 +1,7 @@ import re import json import os +import logging from typing import Optional, Dict, Any import httpx @@ -17,6 +18,8 @@ ContextWindowExceededError, ) +lib_logger = logging.getLogger("rotator_library") + def _parse_duration_string(duration_str: str) -> Optional[int]: """ @@ -513,11 +516,15 @@ def get_retry_after(error: Exception) -> Optional[int]: return None -def classify_error(e: Exception) -> ClassifiedError: +def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedError: """ Classifies an exception into a structured ClassifiedError object. Now handles both litellm and httpx exceptions. + If provider is specified and has a parse_quota_error() method, + attempts provider-specific error parsing first before falling back + to generic classification. + Error types and their typical handling: - rate_limit (429): Rotate key, may retry with backoff - server_error (5xx): Retry with backoff, then rotate @@ -528,7 +535,60 @@ def classify_error(e: Exception) -> ClassifiedError: - context_window_exceeded: Don't retry - request too large - api_connection: Retry with backoff, then rotate - unknown: Rotate key (safer to try another) + + Args: + e: The exception to classify + provider: Optional provider name for provider-specific error parsing + + Returns: + ClassifiedError with error_type, status_code, retry_after, etc. """ + # Try provider-specific parsing first for 429/rate limit errors + if provider: + try: + from .providers import PROVIDER_PLUGINS + + provider_class = PROVIDER_PLUGINS.get(provider) + + if provider_class and hasattr(provider_class, "parse_quota_error"): + # Get error body if available + error_body = None + if hasattr(e, "response") and hasattr(e.response, "text"): + try: + error_body = e.response.text + except Exception: + pass + elif hasattr(e, "body"): + error_body = str(e.body) + + quota_info = provider_class.parse_quota_error(e, error_body) + + if quota_info and quota_info.get("retry_after"): + retry_after = quota_info["retry_after"] + reason = quota_info.get("reason", "QUOTA_EXHAUSTED") + reset_ts = quota_info.get("reset_timestamp") + + # Log the parsed result with human-readable duration + hours = retry_after / 3600 + lib_logger.info( + f"Provider '{provider}' parsed quota error: " + f"retry_after={retry_after}s ({hours:.1f}h), reason={reason}" + + (f", resets at {reset_ts}" if reset_ts else "") + ) + + return ClassifiedError( + error_type="quota_exceeded", + original_exception=e, + status_code=429, + retry_after=retry_after, + ) + except Exception as parse_error: + lib_logger.debug( + f"Provider-specific error parsing failed for '{provider}': {parse_error}" + ) + # Fall through to generic classification + + # Generic classification logic status_code = getattr(e, "status_code", None) if isinstance(e, httpx.HTTPStatusError): # [NEW] Handle httpx errors first diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 7ed85f4b..bdd319b5 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -494,6 +494,147 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface): skip_cost_calculation = True + # Sequential mode by default - preserves thinking signature caches between requests + default_rotation_mode: str = "sequential" + + @staticmethod + def parse_quota_error( + error: Exception, error_body: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """ + Parse Antigravity/Google RPC quota errors. + + Handles the Google Cloud API error format with ErrorInfo and RetryInfo details. + + Example error format: + { + "error": { + "code": 429, + "details": [ + { + "@type": "type.googleapis.com/google.rpc.ErrorInfo", + "reason": "QUOTA_EXHAUSTED", + "metadata": { + "quotaResetDelay": "143h4m52.730699158s", + "quotaResetTimeStamp": "2025-12-11T22:53:16Z" + } + }, + { + "@type": "type.googleapis.com/google.rpc.RetryInfo", + "retryDelay": "515092.730699158s" + } + ] + } + } + + Args: + error: The caught exception + error_body: Optional raw response body string + + Returns: + None if not a parseable quota error, otherwise: + { + "retry_after": int, + "reason": str, + "reset_timestamp": str | None, + } + """ + import re as regex_module + + def parse_duration(duration_str: str) -> Optional[int]: + """Parse duration strings like '143h4m52.73s' or '515092.73s' to seconds.""" + if not duration_str: + return None + + # Handle pure seconds format: "515092.730699158s" + pure_seconds_match = regex_module.match(r"^([\d.]+)s$", duration_str) + if pure_seconds_match: + return int(float(pure_seconds_match.group(1))) + + # Handle compound format: "143h4m52.730699158s" + total_seconds = 0 + patterns = [ + (r"(\d+)h", 3600), # hours + (r"(\d+)m", 60), # minutes + (r"([\d.]+)s", 1), # seconds + ] + for pattern, multiplier in patterns: + match = regex_module.search(pattern, duration_str) + if match: + total_seconds += float(match.group(1)) * multiplier + + return int(total_seconds) if total_seconds > 0 else None + + # Get error body from exception if not provided + body = error_body + if not body: + # Try to extract from various exception attributes + if hasattr(error, "response") and hasattr(error.response, "text"): + body = error.response.text + elif hasattr(error, "body"): + body = str(error.body) + elif hasattr(error, "message"): + body = str(error.message) + else: + body = str(error) + + # Try to find JSON in the body + try: + # Handle cases where JSON is embedded in a larger string + json_match = regex_module.search(r"\{[\s\S]*\}", body) + if not json_match: + return None + + data = json.loads(json_match.group(0)) + except (json.JSONDecodeError, AttributeError, TypeError): + return None + + # Navigate to error.details + error_obj = data.get("error", data) + details = error_obj.get("details", []) + + if not details: + return None + + result = { + "retry_after": None, + "reason": None, + "reset_timestamp": None, + } + + for detail in details: + detail_type = detail.get("@type", "") + + # Parse RetryInfo - most authoritative source for retry delay + if "RetryInfo" in detail_type: + retry_delay = detail.get("retryDelay") + if retry_delay: + parsed = parse_duration(retry_delay) + if parsed: + result["retry_after"] = parsed + + # Parse ErrorInfo - contains reason and quota reset metadata + elif "ErrorInfo" in detail_type: + result["reason"] = detail.get("reason") + metadata = detail.get("metadata", {}) + + # Get quotaResetDelay as fallback if RetryInfo not present + if not result["retry_after"]: + quota_delay = metadata.get("quotaResetDelay") + if quota_delay: + parsed = parse_duration(quota_delay) + if parsed: + result["retry_after"] = parsed + + # Capture reset timestamp for logging + result["reset_timestamp"] = metadata.get("quotaResetTimeStamp") + + # Return None if we couldn't extract retry_after + if not result["retry_after"]: + return None + + return result + def __init__(self): super().__init__() self.model_definitions = ModelDefinitions() diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index e4109ef9..745f934d 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -186,6 +186,31 @@ def _env_int(key: str, default: int) -> int: class GeminiCliProvider(GeminiAuthBase, ProviderInterface): skip_cost_calculation = True + # Balanced by default - Gemini CLI has short cooldowns (seconds, not hours) + default_rotation_mode: str = "balanced" + + @staticmethod + def parse_quota_error( + error: Exception, error_body: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """ + Parse Gemini CLI quota errors. + + Uses the same Google RPC format as Antigravity but typically has + much shorter cooldown durations (seconds to minutes, not hours). + + Args: + error: The caught exception + error_body: Optional raw response body string + + Returns: + Same format as AntigravityProvider.parse_quota_error() + """ + # Reuse the same parsing logic as Antigravity since both use Google RPC format + from .antigravity_provider import AntigravityProvider + + return AntigravityProvider.parse_quota_error(error, error_body) + def __init__(self): super().__init__() self.model_definitions = ModelDefinitions() diff --git a/src/rotator_library/providers/provider_interface.py b/src/rotator_library/providers/provider_interface.py index 996f3a7e..f0f2d695 100644 --- a/src/rotator_library/providers/provider_interface.py +++ b/src/rotator_library/providers/provider_interface.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod from typing import List, Dict, Any, Optional, AsyncGenerator, Union +import os import httpx import litellm @@ -12,6 +13,11 @@ class ProviderInterface(ABC): skip_cost_calculation: bool = False + # Default rotation mode for this provider ("balanced" or "sequential") + # - "balanced": Rotate credentials to distribute load evenly + # - "sequential": Use one credential until exhausted, then switch to next + default_rotation_mode: str = "balanced" + @abstractmethod async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: """ @@ -153,3 +159,69 @@ def get_credential_tier_name(self, credential: str) -> Optional[str]: Tier name string (e.g., "free-tier", "paid-tier") or None if unknown """ return None + + # ========================================================================= + # Sequential Rotation Support + # ========================================================================= + + @classmethod + def get_rotation_mode(cls, provider_name: str) -> str: + """ + Get the rotation mode for this provider. + + Checks ROTATION_MODE_{PROVIDER} environment variable first, + then falls back to the class's default_rotation_mode. + + Args: + provider_name: The provider name (e.g., "antigravity", "gemini_cli") + + Returns: + "balanced" or "sequential" + """ + env_key = f"ROTATION_MODE_{provider_name.upper()}" + return os.getenv(env_key, cls.default_rotation_mode) + + @staticmethod + def parse_quota_error( + error: Exception, error_body: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """ + Parse a quota/rate-limit error and extract structured information. + + Providers should override this method to handle their specific error formats. + This allows the error_handler to use provider-specific parsing when available, + falling back to generic parsing otherwise. + + Args: + error: The caught exception + error_body: Optional raw response body string + + Returns: + None if not a parseable quota error, otherwise: + { + "retry_after": int, # seconds until quota resets + "reason": str, # e.g., "QUOTA_EXHAUSTED", "RATE_LIMITED" + "reset_timestamp": str | None, # ISO timestamp if available + } + """ + return None # Default: no provider-specific parsing + + # TODO: Implement provider-specific quota reset schedules + # Different providers have different quota reset periods: + # - Most providers: Daily reset at a specific time + # - Antigravity free tier: Weekly reset + # - Antigravity paid tier: 5-hour rolling window + # + # Future implementation should add: + # @classmethod + # def get_quota_reset_behavior(cls) -> Dict[str, Any]: + # """ + # Get provider-specific quota reset behavior. + # Returns: + # { + # "type": "daily" | "weekly" | "rolling", + # "reset_time_utc": "03:00", # For daily/weekly + # "rolling_hours": 5, # For rolling + # } + # """ + # return {"type": "daily", "reset_time_utc": "03:00"} diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 577bf4aa..108c1b47 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -5,7 +5,7 @@ import asyncio import random from datetime import date, datetime, timezone, time as dt_time -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, Optional, Set, Tuple import aiofiles import litellm @@ -42,6 +42,10 @@ class UsageManager: This ensures lower-usage credentials are preferred while tolerance controls how much randomness is introduced into the selection process. + + Additionally, providers can specify a rotation mode: + - "balanced" (default): Rotate credentials to distribute load evenly + - "sequential": Use one credential until exhausted (preserves caching) """ def __init__( @@ -49,6 +53,7 @@ def __init__( file_path: str = "key_usage.json", daily_reset_time_utc: Optional[str] = "03:00", rotation_tolerance: float = 0.0, + provider_rotation_modes: Optional[Dict[str, str]] = None, ): """ Initialize the UsageManager. @@ -60,9 +65,13 @@ def __init__( - 0.0: Deterministic, least-used credential always selected - tolerance = 2.0 - 4.0 (default, recommended): Balanced randomness, can pick credentials within 2 uses of max - 5.0+: High randomness, more unpredictable selection patterns + provider_rotation_modes: Dict mapping provider names to rotation modes. + - "balanced": Rotate credentials to distribute load evenly (default) + - "sequential": Use one credential until exhausted (preserves caching) """ self.file_path = file_path self.rotation_tolerance = rotation_tolerance + self.provider_rotation_modes = provider_rotation_modes or {} self.key_states: Dict[str, Dict[str, Any]] = {} self._data_lock = asyncio.Lock() @@ -81,6 +90,72 @@ def __init__( else: self.daily_reset_time_utc = None + def _get_rotation_mode(self, provider: str) -> str: + """ + Get the rotation mode for a provider. + + Args: + provider: Provider name (e.g., "antigravity", "gemini_cli") + + Returns: + "balanced" or "sequential" + """ + return self.provider_rotation_modes.get(provider, "balanced") + + def _select_sequential( + self, + candidates: List[Tuple[str, int]], + credential_priorities: Optional[Dict[str, int]] = None, + ) -> str: + """ + Select credential in strict sequential order for cache-preserving rotation. + + This method ensures the same credential is reused until it hits a cooldown, + which preserves provider-side caching (e.g., thinking signature caches). + + Selection logic: + 1. Sort by priority (lowest number = highest priority) + 2. Within same priority, sort by last_used_ts (most recent first = sticky) + 3. Return the first candidate + + Args: + candidates: List of (credential_id, usage_count) tuples + credential_priorities: Optional dict mapping credentials to priority levels + + Returns: + Selected credential ID + """ + if not candidates: + raise ValueError("Cannot select from empty candidate list") + + if len(candidates) == 1: + return candidates[0][0] + + def sort_key(item: Tuple[str, int]) -> Tuple[int, float]: + cred, _ = item + # Priority: lower is better (1 = highest priority) + priority = ( + credential_priorities.get(cred, 999) if credential_priorities else 999 + ) + # Last used: higher (more recent) is better for stickiness + last_used = ( + self._usage_data.get(cred, {}).get("last_used_ts", 0) + if self._usage_data + else 0 + ) + # Negative last_used so most recent sorts first + return (priority, -last_used) + + sorted_candidates = sorted(candidates, key=sort_key) + selected = sorted_candidates[0][0] + + lib_logger.debug( + f"Sequential selection: chose {mask_credential(selected)} " + f"(priority={credential_priorities.get(selected, 999) if credential_priorities else 'N/A'})" + ) + + return selected + async def _lazy_init(self): """Initializes the usage data by loading it from the file asynchronously.""" async with self._init_lock: @@ -144,14 +219,63 @@ async def _reset_daily_stats_if_needed(self): ) needs_saving = True - # Reset cooldowns - data["model_cooldowns"] = {} - data["key_cooldown_until"] = None + # Reset cooldowns - BUT preserve unexpired long-term cooldowns + # This is important for quota errors with long cooldowns (e.g., 143 hours) + now_ts = time.time() + if "model_cooldowns" in data: + active_cooldowns = { + model: end_time + for model, end_time in data["model_cooldowns"].items() + if end_time > now_ts + } + if active_cooldowns: + # Calculate how long the longest cooldown has remaining + max_remaining = max( + end_time - now_ts + for end_time in active_cooldowns.values() + ) + hours_remaining = max_remaining / 3600 + lib_logger.info( + f"Preserving {len(active_cooldowns)} active cooldown(s) " + f"for key {mask_credential(key)} during daily reset " + f"(longest: {hours_remaining:.1f}h remaining)" + ) + data["model_cooldowns"] = active_cooldowns + else: + data["model_cooldowns"] = {} + + # Clear key-level cooldown only if expired + if data.get("key_cooldown_until"): + if data["key_cooldown_until"] <= now_ts: + data["key_cooldown_until"] = None + else: + hours_remaining = ( + data["key_cooldown_until"] - now_ts + ) / 3600 + lib_logger.info( + f"Preserving key-level cooldown for {mask_credential(key)} " + f"during daily reset ({hours_remaining:.1f}h remaining)" + ) + else: + data["key_cooldown_until"] = None # Reset consecutive failures if "failures" in data: data["failures"] = {} + # TODO: Implement provider-specific reset schedules + # Different providers have different quota reset periods: + # - Most providers: Daily reset at daily_reset_time_utc + # - Antigravity free tier: Weekly reset + # - Antigravity paid tier: 5-hour rolling window + # + # Future implementation should: + # 1. Group credentials by provider (extracted from key path or metadata) + # 2. Check each provider's get_quota_reset_behavior() + # 3. Apply provider-specific reset logic instead of universal daily reset + # + # For now, we preserve unexpired cooldowns which handles long cooldowns correctly. + # Archive global stats from the previous day's 'daily' daily_data = data.get("daily", {}) if daily_data: @@ -336,15 +460,30 @@ async def acquire_key( elif key_state["models_in_use"].get(model, 0) < max_concurrent: tier2_keys.append((key, usage_count)) - # Apply weighted random selection or deterministic sorting - selection_method = ( - "weighted-random" - if self.rotation_tolerance > 0 - else "least-used" - ) + # Determine selection method based on provider's rotation mode + provider = model.split("/")[0] if "/" in model else "" + rotation_mode = self._get_rotation_mode(provider) - if self.rotation_tolerance > 0: - # Weighted random selection within each tier + if rotation_mode == "sequential": + # Sequential mode: stick with same credential until exhausted + selection_method = "sequential" + if tier1_keys: + selected_key = self._select_sequential( + tier1_keys, credential_priorities + ) + tier1_keys = [ + (k, u) for k, u in tier1_keys if k == selected_key + ] + if tier2_keys: + selected_key = self._select_sequential( + tier2_keys, credential_priorities + ) + tier2_keys = [ + (k, u) for k, u in tier2_keys if k == selected_key + ] + elif self.rotation_tolerance > 0: + # Balanced mode with weighted randomness + selection_method = "weighted-random" if tier1_keys: selected_key = self._select_weighted_random( tier1_keys, self.rotation_tolerance @@ -361,6 +500,7 @@ async def acquire_key( ] else: # Deterministic: sort by usage within each tier + selection_method = "least-used" tier1_keys.sort(key=lambda x: x[1]) tier2_keys.sort(key=lambda x: x[1]) @@ -452,13 +592,30 @@ async def acquire_key( elif key_state["models_in_use"].get(model, 0) < max_concurrent: tier2_keys.append((key, usage_count)) - # Apply weighted random selection or deterministic sorting - selection_method = ( - "weighted-random" if self.rotation_tolerance > 0 else "least-used" - ) + # Determine selection method based on provider's rotation mode + provider = model.split("/")[0] if "/" in model else "" + rotation_mode = self._get_rotation_mode(provider) - if self.rotation_tolerance > 0: - # Weighted random selection within each tier + if rotation_mode == "sequential": + # Sequential mode: stick with same credential until exhausted + selection_method = "sequential" + if tier1_keys: + selected_key = self._select_sequential( + tier1_keys, credential_priorities + ) + tier1_keys = [ + (k, u) for k, u in tier1_keys if k == selected_key + ] + if tier2_keys: + selected_key = self._select_sequential( + tier2_keys, credential_priorities + ) + tier2_keys = [ + (k, u) for k, u in tier2_keys if k == selected_key + ] + elif self.rotation_tolerance > 0: + # Balanced mode with weighted randomness + selection_method = "weighted-random" if tier1_keys: selected_key = self._select_weighted_random( tier1_keys, self.rotation_tolerance @@ -475,6 +632,7 @@ async def acquire_key( ] else: # Deterministic: sort by usage within each tier + selection_method = "least-used" tier1_keys.sort(key=lambda x: x[1]) tier2_keys.sort(key=lambda x: x[1]) @@ -726,10 +884,24 @@ async def record_failure( if classified_error.error_type in ["rate_limit", "quota_exceeded"]: # Rate limit / Quota errors: use retry_after if available, otherwise default to 60s cooldown_seconds = classified_error.retry_after or 60 - lib_logger.info( - f"Rate limit error on key {mask_credential(key)} for model {model}. " - f"Using {'provided' if classified_error.retry_after else 'default'} retry_after: {cooldown_seconds}s" - ) + if classified_error.retry_after: + # Log with human-readable duration for provider-parsed cooldowns + hours = cooldown_seconds / 3600 + if hours >= 1: + lib_logger.info( + f"Quota/rate limit on key {mask_credential(key)} for model {model}. " + f"Applying provider-specified cooldown: {cooldown_seconds}s ({hours:.1f}h)" + ) + else: + lib_logger.info( + f"Rate limit on key {mask_credential(key)} for model {model}. " + f"Applying provider-specified cooldown: {cooldown_seconds}s" + ) + else: + lib_logger.info( + f"Rate limit on key {mask_credential(key)} for model {model}. " + f"Using default cooldown: {cooldown_seconds}s" + ) elif classified_error.error_type == "authentication": # Apply a 5-minute key-level lockout for auth errors key_data["key_cooldown_until"] = time.time() + 300 From 98f6823355fe3b71dda0387eb5ab66cb6e4b3fa0 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sat, 6 Dec 2025 05:53:50 +0100 Subject: [PATCH 088/221] =?UTF-8?q?feat(usage):=20=E2=9C=A8=20add=20provid?= =?UTF-8?q?er-specific=20rolling=20window=20usage=20tracking?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement flexible per-provider usage reset configurations to support different quota windows (5h rolling for Antigravity paid tiers, 7-day for free tier) instead of universal daily resets. - Add `get_usage_reset_config()` and `get_default_usage_field_name()` methods to ProviderInterface for provider-specific configuration - Implement Antigravity-specific reset config returning different windows based on credential tier (5h for paid, 7-day for free) - Refactor UsageManager to support custom usage field names ("5h_window", "weekly") instead of hardcoded "daily" - Add window start timestamp tracking that begins on first request and resets after window expiration - Extract reset logic into separate methods (`_check_window_reset`, `_check_daily_reset`) for cleaner separation - Add credential-to-provider mapping via regex pattern matching for OAuth credential paths - Archive expired window stats to "global" field matching existing daily reset behavior - Preserve unexpired cooldowns during all reset types to maintain long-term quota error handling - Pass provider_plugins to UsageManager initialization for access to provider configuration This enables accurate quota tracking for providers with non-daily reset schedules while maintaining backward compatibility with existing daily reset behavior for providers without custom configuration. --- src/rotator_library/client.py | 1 + .../providers/antigravity_provider.py | 54 ++ .../providers/provider_interface.py | 82 ++- src/rotator_library/usage_manager.py | 563 +++++++++++++----- 4 files changed, 533 insertions(+), 167 deletions(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 179cd09b..9e1a3042 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -161,6 +161,7 @@ def __init__( file_path=usage_file_path, rotation_tolerance=rotation_tolerance, provider_rotation_modes=provider_rotation_modes, + provider_plugins=PROVIDER_PLUGINS, ) self._model_list_cache = {} self._provider_plugins = PROVIDER_PLUGINS diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index bdd319b5..599c4040 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -822,6 +822,60 @@ def get_model_tier_requirement(self, model: str) -> Optional[int]: """ return None + def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: + """ + Get Antigravity-specific usage tracking configuration based on credential tier. + + Antigravity has different quota reset windows by tier: + - Paid tiers (priority 1): 5-hour rolling window + - Free tier (priority 2): 7-day rolling window + - Unknown/legacy: 7-day rolling window (conservative default) + + Args: + credential: The credential path + + Returns: + Usage reset configuration dict + """ + tier = self.project_tier_cache.get(credential) + if not tier: + tier = self._load_tier_from_file(credential) + + # Paid tiers: 5-hour window + if tier and tier not in ["free-tier", "legacy-tier", "unknown"]: + return { + "window_seconds": 5 * 60 * 60, # 18000 seconds = 5 hours + "field_name": "5h_window", + "priority": 1, + "description": "5-hour rolling window (paid tier)", + } + + # Free tier: 7-day window + if tier == "free-tier": + return { + "window_seconds": 7 * 24 * 60 * 60, # 604800 seconds = 7 days + "field_name": "weekly", + "priority": 2, + "description": "7-day rolling window (free tier)", + } + + # Unknown/legacy: use 7-day window as conservative default + return { + "window_seconds": 7 * 24 * 60 * 60, # 604800 seconds = 7 days + "field_name": "weekly", + "priority": 10, + "description": "7-day rolling window (unknown tier - conservative default)", + } + + def get_default_usage_field_name(self) -> str: + """ + Get the default usage tracking field name for Antigravity. + + Returns: + "weekly" as the conservative default for unknown credentials + """ + return "weekly" + async def initialize_credentials(self, credential_paths: List[str]) -> None: """ Load persisted tier information from credential files at startup. diff --git a/src/rotator_library/providers/provider_interface.py b/src/rotator_library/providers/provider_interface.py index f0f2d695..e12cbabc 100644 --- a/src/rotator_library/providers/provider_interface.py +++ b/src/rotator_library/providers/provider_interface.py @@ -206,22 +206,66 @@ def parse_quota_error( """ return None # Default: no provider-specific parsing - # TODO: Implement provider-specific quota reset schedules - # Different providers have different quota reset periods: - # - Most providers: Daily reset at a specific time - # - Antigravity free tier: Weekly reset - # - Antigravity paid tier: 5-hour rolling window - # - # Future implementation should add: - # @classmethod - # def get_quota_reset_behavior(cls) -> Dict[str, Any]: - # """ - # Get provider-specific quota reset behavior. - # Returns: - # { - # "type": "daily" | "weekly" | "rolling", - # "reset_time_utc": "03:00", # For daily/weekly - # "rolling_hours": 5, # For rolling - # } - # """ - # return {"type": "daily", "reset_time_utc": "03:00"} + # ========================================================================= + # Per-Provider Usage Tracking Configuration + # ========================================================================= + + def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: + """ + Get provider-specific usage tracking configuration for a credential. + + This allows providers to define custom usage reset windows based on + credential tier (e.g., paid vs free accounts with different quota periods). + + The UsageManager will use this configuration to: + 1. Track usage in a custom-named field (instead of default "daily") + 2. Reset usage based on a rolling window from first request + 3. Archive stats to "global" when the window expires + + Args: + credential: The credential identifier (API key or path) + + Returns: + None to use default daily reset, otherwise a dict with: + { + "window_seconds": int, # Duration in seconds (e.g., 18000 for 5h) + "field_name": str, # Custom field name (e.g., "5h_window", "weekly") + "priority": int, # Priority level this config applies to (for docs) + "description": str, # Human-readable description (for logging) + } + + Examples: + Antigravity paid tier: + { + "window_seconds": 18000, # 5 hours + "field_name": "5h_window", + "priority": 1, + "description": "5-hour rolling window (paid tier)" + } + + Antigravity free tier: + { + "window_seconds": 604800, # 7 days + "field_name": "weekly", + "priority": 2, + "description": "7-day rolling window (free tier)" + } + + Note: + - window_seconds: Time from first request until stats reset + - When window expires, stats move to "global" (same as daily reset) + - First request after window expiry starts a new window + """ + return None # Default: use daily reset at daily_reset_time_utc + + def get_default_usage_field_name(self) -> str: + """ + Get the default usage tracking field name for this provider. + + Providers can override this to use a custom field name for usage tracking + when no credential-specific config is available. + + Returns: + Field name string (default: "daily") + """ + return "daily" diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 108c1b47..1ae93277 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -54,6 +54,7 @@ def __init__( daily_reset_time_utc: Optional[str] = "03:00", rotation_tolerance: float = 0.0, provider_rotation_modes: Optional[Dict[str, str]] = None, + provider_plugins: Optional[Dict[str, Any]] = None, ): """ Initialize the UsageManager. @@ -68,10 +69,13 @@ def __init__( provider_rotation_modes: Dict mapping provider names to rotation modes. - "balanced": Rotate credentials to distribute load evenly (default) - "sequential": Use one credential until exhausted (preserves caching) + provider_plugins: Dict mapping provider names to provider plugin instances. + Used for per-provider usage reset configuration (window durations, field names). """ self.file_path = file_path self.rotation_tolerance = rotation_tolerance self.provider_rotation_modes = provider_rotation_modes or {} + self.provider_plugins = provider_plugins or PROVIDER_PLUGINS self.key_states: Dict[str, Dict[str, Any]] = {} self._data_lock = asyncio.Lock() @@ -102,6 +106,112 @@ def _get_rotation_mode(self, provider: str) -> str: """ return self.provider_rotation_modes.get(provider, "balanced") + def _get_provider_from_credential(self, credential: str) -> Optional[str]: + """ + Extract provider name from credential path or identifier. + + Supports multiple credential formats: + - OAuth: "oauth_creds/antigravity_oauth_15.json" -> "antigravity" + - OAuth: "C:\\...\\oauth_creds\\gemini_cli_oauth_1.json" -> "gemini_cli" + - API key style: stored with provider prefix metadata + + Args: + credential: The credential identifier (path or key) + + Returns: + Provider name string or None if cannot be determined + """ + import re + + # Normalize path separators + normalized = credential.replace("\\", "/") + + # Pattern: {provider}_oauth_{number}.json + match = re.search(r"/([a-z_]+)_oauth_\d+\.json$", normalized, re.IGNORECASE) + if match: + return match.group(1).lower() + + # Pattern: oauth_creds/{provider}_... + match = re.search(r"oauth_creds/([a-z_]+)_", normalized, re.IGNORECASE) + if match: + return match.group(1).lower() + + return None + + def _get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: + """ + Get the usage reset configuration for a credential from its provider plugin. + + Args: + credential: The credential identifier + + Returns: + Configuration dict with window_seconds, field_name, etc. + or None to use default daily reset. + """ + provider = self._get_provider_from_credential(credential) + if not provider: + return None + + plugin = self.provider_plugins.get(provider) + if not plugin: + return None + + if hasattr(plugin, "get_usage_reset_config"): + return plugin.get_usage_reset_config(credential) + + return None + + def _get_usage_field_name(self, credential: str) -> str: + """ + Get the usage tracking field name for a credential. + + Returns the provider-specific field name if configured, + otherwise falls back to "daily". + + Args: + credential: The credential identifier + + Returns: + Field name string (e.g., "5h_window", "weekly", "daily") + """ + config = self._get_usage_reset_config(credential) + if config and "field_name" in config: + return config["field_name"] + + # Check provider default + provider = self._get_provider_from_credential(credential) + if provider: + plugin = self.provider_plugins.get(provider) + if plugin and hasattr(plugin, "get_default_usage_field_name"): + return plugin.get_default_usage_field_name() + + return "daily" + + def _get_usage_count(self, key: str, model: str) -> int: + """ + Get the current usage count for a model from the appropriate usage field. + + Args: + key: Credential identifier + model: Model name + + Returns: + Usage count (success_count) for the model in the current window/daily period + """ + if self._usage_data is None: + return 0 + + key_data = self._usage_data.get(key, {}) + usage_field = self._get_usage_field_name(key) + + return ( + key_data.get(usage_field, {}) + .get("models", {}) + .get(model, {}) + .get("success_count", 0) + ) + def _select_sequential( self, candidates: List[Tuple[str, int]], @@ -186,129 +296,233 @@ async def _save_usage(self): await f.write(json.dumps(self._usage_data, indent=2)) async def _reset_daily_stats_if_needed(self): - """Checks if daily stats need to be reset for any key.""" - if self._usage_data is None or not self.daily_reset_time_utc: + """ + Checks if usage stats need to be reset for any key. + + Supports two reset modes: + 1. Provider-specific rolling windows (e.g., 5h for Antigravity paid, 7d for free) + 2. Legacy daily reset at daily_reset_time_utc for providers without custom config + """ + if self._usage_data is None: return now_utc = datetime.now(timezone.utc) + now_ts = time.time() today_str = now_utc.date().isoformat() needs_saving = False for key, data in self._usage_data.items(): - last_reset_str = data.get("last_daily_reset", "") - - if last_reset_str != today_str: - last_reset_dt = None - if last_reset_str: - # Ensure the parsed datetime is timezone-aware (UTC) - last_reset_dt = datetime.fromisoformat(last_reset_str).replace( - tzinfo=timezone.utc - ) + # Check for provider-specific reset configuration + reset_config = self._get_usage_reset_config(key) - # Determine the reset threshold for today - reset_threshold_today = datetime.combine( - now_utc.date(), self.daily_reset_time_utc + if reset_config: + # Provider-specific rolling window reset + needs_saving |= await self._check_window_reset( + key, data, reset_config, now_ts + ) + elif self.daily_reset_time_utc: + # Legacy daily reset for providers without custom config + needs_saving |= await self._check_daily_reset( + key, data, now_utc, today_str, now_ts ) - if ( - last_reset_dt is None - or last_reset_dt < reset_threshold_today <= now_utc - ): - lib_logger.debug( - f"Performing daily reset for key {mask_credential(key)}" - ) - needs_saving = True - - # Reset cooldowns - BUT preserve unexpired long-term cooldowns - # This is important for quota errors with long cooldowns (e.g., 143 hours) - now_ts = time.time() - if "model_cooldowns" in data: - active_cooldowns = { - model: end_time - for model, end_time in data["model_cooldowns"].items() - if end_time > now_ts - } - if active_cooldowns: - # Calculate how long the longest cooldown has remaining - max_remaining = max( - end_time - now_ts - for end_time in active_cooldowns.values() - ) - hours_remaining = max_remaining / 3600 - lib_logger.info( - f"Preserving {len(active_cooldowns)} active cooldown(s) " - f"for key {mask_credential(key)} during daily reset " - f"(longest: {hours_remaining:.1f}h remaining)" - ) - data["model_cooldowns"] = active_cooldowns - else: - data["model_cooldowns"] = {} + if needs_saving: + await self._save_usage() - # Clear key-level cooldown only if expired - if data.get("key_cooldown_until"): - if data["key_cooldown_until"] <= now_ts: - data["key_cooldown_until"] = None - else: - hours_remaining = ( - data["key_cooldown_until"] - now_ts - ) / 3600 - lib_logger.info( - f"Preserving key-level cooldown for {mask_credential(key)} " - f"during daily reset ({hours_remaining:.1f}h remaining)" - ) - else: - data["key_cooldown_until"] = None - - # Reset consecutive failures - if "failures" in data: - data["failures"] = {} - - # TODO: Implement provider-specific reset schedules - # Different providers have different quota reset periods: - # - Most providers: Daily reset at daily_reset_time_utc - # - Antigravity free tier: Weekly reset - # - Antigravity paid tier: 5-hour rolling window - # - # Future implementation should: - # 1. Group credentials by provider (extracted from key path or metadata) - # 2. Check each provider's get_quota_reset_behavior() - # 3. Apply provider-specific reset logic instead of universal daily reset - # - # For now, we preserve unexpired cooldowns which handles long cooldowns correctly. - - # Archive global stats from the previous day's 'daily' - daily_data = data.get("daily", {}) - if daily_data: - global_data = data.setdefault("global", {"models": {}}) - for model, stats in daily_data.get("models", {}).items(): - global_model_stats = global_data["models"].setdefault( - model, - { - "success_count": 0, - "prompt_tokens": 0, - "completion_tokens": 0, - "approx_cost": 0.0, - }, - ) - global_model_stats["success_count"] += stats.get( - "success_count", 0 - ) - global_model_stats["prompt_tokens"] += stats.get( - "prompt_tokens", 0 - ) - global_model_stats["completion_tokens"] += stats.get( - "completion_tokens", 0 - ) - global_model_stats["approx_cost"] += stats.get( - "approx_cost", 0.0 - ) + async def _check_window_reset( + self, + key: str, + data: Dict[str, Any], + reset_config: Dict[str, Any], + now_ts: float, + ) -> bool: + """ + Check and perform rolling window reset for a credential. - # Reset daily stats - data["daily"] = {"date": today_str, "models": {}} - data["last_daily_reset"] = today_str + Args: + key: Credential identifier + data: Usage data for this credential + reset_config: Provider's reset configuration + now_ts: Current timestamp - if needs_saving: - await self._save_usage() + Returns: + True if data was modified and needs saving + """ + window_seconds = reset_config.get("window_seconds", 86400) # Default 24h + field_name = reset_config.get("field_name", "window") + description = reset_config.get("description", "rolling window") + + # Get current window data + window_data = data.get(field_name, {}) + window_start = window_data.get("start_ts") + + # No window started yet - nothing to reset + if window_start is None: + return False + + # Check if window has expired + window_end = window_start + window_seconds + if now_ts < window_end: + # Window still active + return False + + # Window expired - perform reset + hours_elapsed = (now_ts - window_start) / 3600 + lib_logger.info( + f"Resetting {field_name} for {mask_credential(key)} - " + f"{description} expired after {hours_elapsed:.1f}h" + ) + + # Archive to global + self._archive_to_global(data, window_data) + + # Preserve unexpired cooldowns + self._preserve_unexpired_cooldowns(key, data, now_ts) + + # Reset window stats (but don't start new window until first request) + data[field_name] = {"start_ts": None, "models": {}} + + # Reset consecutive failures + if "failures" in data: + data["failures"] = {} + + return True + + async def _check_daily_reset( + self, + key: str, + data: Dict[str, Any], + now_utc: datetime, + today_str: str, + now_ts: float, + ) -> bool: + """ + Check and perform legacy daily reset for a credential. + + Args: + key: Credential identifier + data: Usage data for this credential + now_utc: Current datetime in UTC + today_str: Today's date as ISO string + now_ts: Current timestamp + + Returns: + True if data was modified and needs saving + """ + last_reset_str = data.get("last_daily_reset", "") + + if last_reset_str == today_str: + return False + + last_reset_dt = None + if last_reset_str: + try: + last_reset_dt = datetime.fromisoformat(last_reset_str).replace( + tzinfo=timezone.utc + ) + except ValueError: + pass + + # Determine the reset threshold for today + reset_threshold_today = datetime.combine( + now_utc.date(), self.daily_reset_time_utc + ) + + if not ( + last_reset_dt is None or last_reset_dt < reset_threshold_today <= now_utc + ): + return False + + lib_logger.debug(f"Performing daily reset for key {mask_credential(key)}") + + # Preserve unexpired cooldowns + self._preserve_unexpired_cooldowns(key, data, now_ts) + + # Reset consecutive failures + if "failures" in data: + data["failures"] = {} + + # Archive daily stats to global + daily_data = data.get("daily", {}) + if daily_data: + self._archive_to_global(data, daily_data) + + # Reset daily stats + data["daily"] = {"date": today_str, "models": {}} + data["last_daily_reset"] = today_str + + return True + + def _archive_to_global( + self, data: Dict[str, Any], source_data: Dict[str, Any] + ) -> None: + """ + Archive usage stats from a source field (daily/window) to global. + + Args: + data: The credential's usage data + source_data: The source field data to archive (has "models" key) + """ + global_data = data.setdefault("global", {"models": {}}) + for model, stats in source_data.get("models", {}).items(): + global_model_stats = global_data["models"].setdefault( + model, + { + "success_count": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "approx_cost": 0.0, + }, + ) + global_model_stats["success_count"] += stats.get("success_count", 0) + global_model_stats["prompt_tokens"] += stats.get("prompt_tokens", 0) + global_model_stats["completion_tokens"] += stats.get("completion_tokens", 0) + global_model_stats["approx_cost"] += stats.get("approx_cost", 0.0) + + def _preserve_unexpired_cooldowns( + self, key: str, data: Dict[str, Any], now_ts: float + ) -> None: + """ + Preserve unexpired cooldowns during reset (important for long quota cooldowns). + + Args: + key: Credential identifier (for logging) + data: The credential's usage data + now_ts: Current timestamp + """ + # Preserve unexpired model cooldowns + if "model_cooldowns" in data: + active_cooldowns = { + model: end_time + for model, end_time in data["model_cooldowns"].items() + if end_time > now_ts + } + if active_cooldowns: + max_remaining = max( + end_time - now_ts for end_time in active_cooldowns.values() + ) + hours_remaining = max_remaining / 3600 + lib_logger.info( + f"Preserving {len(active_cooldowns)} active cooldown(s) " + f"for key {mask_credential(key)} during reset " + f"(longest: {hours_remaining:.1f}h remaining)" + ) + data["model_cooldowns"] = active_cooldowns + else: + data["model_cooldowns"] = {} + + # Preserve unexpired key-level cooldown + if data.get("key_cooldown_until"): + if data["key_cooldown_until"] <= now_ts: + data["key_cooldown_until"] = None + else: + hours_remaining = (data["key_cooldown_until"] - now_ts) / 3600 + lib_logger.info( + f"Preserving key-level cooldown for {mask_credential(key)} " + f"during reset ({hours_remaining:.1f}h remaining)" + ) + else: + data["key_cooldown_until"] = None def _initialize_key_states(self, keys: List[str]): """Initializes state tracking for all provided keys if not already present.""" @@ -430,12 +644,7 @@ async def acquire_key( priority = credential_priorities.get(key, 999) # Get usage count for load balancing within priority groups - usage_count = ( - key_data.get("daily", {}) - .get("models", {}) - .get(model, {}) - .get("success_count", 0) - ) + usage_count = self._get_usage_count(key, model) # Group by priority if priority not in priority_groups: @@ -577,12 +786,7 @@ async def acquire_key( continue # Prioritize keys based on their current usage to ensure load balancing. - usage_count = ( - key_data.get("daily", {}) - .get("models", {}) - .get(model, {}) - .get("success_count", 0) - ) + usage_count = self._get_usage_count(key, model) key_state = self.key_states[key] # Tier 1: Completely idle keys (preferred). @@ -743,22 +947,50 @@ async def record_success( """ Records a successful API call, resetting failure counters. It safely handles cases where token usage data is not available. + + Uses provider-specific field names for usage tracking (e.g., "5h_window", "weekly") + and sets window start timestamp on first request. """ await self._lazy_init() async with self._data_lock: + now_ts = time.time() today_utc_str = datetime.now(timezone.utc).date().isoformat() - key_data = self._usage_data.setdefault( - key, - { - "daily": {"date": today_utc_str, "models": {}}, - "global": {"models": {}}, - "model_cooldowns": {}, - "failures": {}, - }, - ) + + # Determine the usage field name for this credential + usage_field = self._get_usage_field_name(key) + reset_config = self._get_usage_reset_config(key) + uses_window = reset_config is not None + + # Initialize key data with appropriate structure + if uses_window: + # Provider-specific rolling window + key_data = self._usage_data.setdefault( + key, + { + usage_field: {"start_ts": None, "models": {}}, + "global": {"models": {}}, + "model_cooldowns": {}, + "failures": {}, + }, + ) + # Ensure the usage field exists (for migration from old format) + if usage_field not in key_data: + key_data[usage_field] = {"start_ts": None, "models": {}} + else: + # Legacy daily reset + key_data = self._usage_data.setdefault( + key, + { + "daily": {"date": today_utc_str, "models": {}}, + "global": {"models": {}}, + "model_cooldowns": {}, + "failures": {}, + }, + ) + usage_field = "daily" # If the key is new, ensure its reset date is initialized to prevent an immediate reset. - if "last_daily_reset" not in key_data: + if not uses_window and "last_daily_reset" not in key_data: key_data["last_daily_reset"] = today_utc_str # Always record a success and reset failures @@ -767,7 +999,24 @@ async def record_success( if model in key_data.get("model_cooldowns", {}): del key_data["model_cooldowns"][model] - daily_model_data = key_data["daily"]["models"].setdefault( + # Get or create the usage field data + usage_data = key_data.setdefault(usage_field, {"models": {}}) + + # For window-based tracking, set start_ts on first request + if uses_window: + if usage_data.get("start_ts") is None: + usage_data["start_ts"] = now_ts + window_hours = reset_config.get("window_seconds", 0) / 3600 + description = reset_config.get("description", "rolling window") + lib_logger.info( + f"Starting new {window_hours:.1f}h window for {mask_credential(key)} - {description}" + ) + + # Ensure models dict exists + if "models" not in usage_data: + usage_data["models"] = {} + + model_data = usage_data["models"].setdefault( model, { "success_count": 0, @@ -776,7 +1025,7 @@ async def record_success( "approx_cost": 0.0, }, ) - daily_model_data["success_count"] += 1 + model_data["success_count"] += 1 # Safely attempt to record token and cost usage if ( @@ -785,8 +1034,8 @@ async def record_success( and completion_response.usage ): usage = completion_response.usage - daily_model_data["prompt_tokens"] += usage.prompt_tokens - daily_model_data["completion_tokens"] += getattr( + model_data["prompt_tokens"] += usage.prompt_tokens + model_data["completion_tokens"] += getattr( usage, "completion_tokens", 0 ) # Not present in embedding responses lib_logger.info( @@ -794,7 +1043,7 @@ async def record_success( ) try: provider_name = model.split("/")[0] - provider_plugin = PROVIDER_PLUGINS.get(provider_name) + provider_plugin = self.provider_plugins.get(provider_name) # Check class attribute directly - no need to instantiate if provider_plugin and getattr( @@ -821,7 +1070,7 @@ async def record_success( ) if cost is not None: - daily_model_data["approx_cost"] += cost + model_data["approx_cost"] += cost except Exception as e: lib_logger.warning( f"Could not calculate cost for model {model}: {e}" @@ -836,7 +1085,7 @@ async def record_success( f"No usage data found in completion response for model {model}. Recording success without token count." ) - key_data["last_used_ts"] = time.time() + key_data["last_used_ts"] = now_ts await self._save_usage() @@ -859,15 +1108,33 @@ async def record_failure( await self._lazy_init() async with self._data_lock: today_utc_str = datetime.now(timezone.utc).date().isoformat() - key_data = self._usage_data.setdefault( - key, - { - "daily": {"date": today_utc_str, "models": {}}, - "global": {"models": {}}, - "model_cooldowns": {}, - "failures": {}, - }, - ) + + # Determine the usage field name for this credential + usage_field = self._get_usage_field_name(key) + reset_config = self._get_usage_reset_config(key) + uses_window = reset_config is not None + + # Initialize key data with appropriate structure + if uses_window: + key_data = self._usage_data.setdefault( + key, + { + usage_field: {"start_ts": None, "models": {}}, + "global": {"models": {}}, + "model_cooldowns": {}, + "failures": {}, + }, + ) + else: + key_data = self._usage_data.setdefault( + key, + { + "daily": {"date": today_utc_str, "models": {}}, + "global": {"models": {}}, + "model_cooldowns": {}, + "failures": {}, + }, + ) # Provider-level errors (transient issues) should not count against the key provider_level_errors = {"server_error", "api_connection"} From 0ca165129f842dc861b569f017eabe562c6d7ac5 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sat, 6 Dec 2025 06:18:39 +0100 Subject: [PATCH 089/221] =?UTF-8?q?feat(usage):=20=E2=9C=A8=20implement=20?= =?UTF-8?q?per-model=20quota=20tracking=20with=20authoritative=20reset=20t?= =?UTF-8?q?imestamps?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces granular per-model quota tracking that supports provider-specific reset timestamps from quota exhausted errors. Key changes: - Add `quota_reset_timestamp` field to `ClassifiedError` to capture authoritative Unix timestamp from provider's quota exhausted responses - Implement per-model usage tracking mode where each model maintains its own window with `window_start_ts` and `quota_reset_ts` - Add quota group support for models that share quota limits (e.g., Claude Sonnet and Opus on Antigravity) - Parse Antigravity's `quotaResetTimeStamp` ISO format to Unix timestamp for precise reset timing - Update reset logic to prioritize authoritative `quota_reset_ts` over fallback window calculations - Distinguish between quota exhausted (sets authoritative reset time) and rate limit (transient cooldown only) - Migrate Antigravity provider to per-model tracking with 5-hour windows for paid tier and 7-day windows for free tier The per-model mode enables more accurate quota tracking by using exact reset times from provider error responses rather than estimated windows, preventing premature resets and improving credential utilization. BREAKING CHANGE: Provider implementations using custom `get_usage_reset_config()` must now return a `mode` field ("per_model" or "credential") instead of `field_name`. The usage data structure has changed from `key_data["field_name"]["models"]` to `key_data["models"]` for per-model tracking. Existing usage data will be preserved but new tracking will use the updated structure. --- src/rotator_library/error_handler.py | 16 +- .../providers/antigravity_provider.py | 96 +++- .../providers/provider_interface.py | 76 ++- src/rotator_library/usage_manager.py | 534 ++++++++++++++---- 4 files changed, 574 insertions(+), 148 deletions(-) diff --git a/src/rotator_library/error_handler.py b/src/rotator_library/error_handler.py index 51692c49..3b9ae81f 100644 --- a/src/rotator_library/error_handler.py +++ b/src/rotator_library/error_handler.py @@ -347,14 +347,26 @@ def __init__( original_exception: Exception, status_code: Optional[int] = None, retry_after: Optional[int] = None, + quota_reset_timestamp: Optional[float] = None, ): self.error_type = error_type self.original_exception = original_exception self.status_code = status_code self.retry_after = retry_after + # Unix timestamp when quota resets (from quota_exhausted errors) + # This is the authoritative reset time parsed from provider's error response + self.quota_reset_timestamp = quota_reset_timestamp def __str__(self): - return f"ClassifiedError(type={self.error_type}, status={self.status_code}, retry_after={self.retry_after}, original_exc={self.original_exception})" + parts = [ + f"type={self.error_type}", + f"status={self.status_code}", + f"retry_after={self.retry_after}", + ] + if self.quota_reset_timestamp: + parts.append(f"quota_reset_ts={self.quota_reset_timestamp}") + parts.append(f"original_exc={self.original_exception}") + return f"ClassifiedError({', '.join(parts)})" def _extract_retry_from_json_body(json_text: str) -> Optional[int]: @@ -567,6 +579,7 @@ def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedEr retry_after = quota_info["retry_after"] reason = quota_info.get("reason", "QUOTA_EXHAUSTED") reset_ts = quota_info.get("reset_timestamp") + quota_reset_timestamp = quota_info.get("quota_reset_timestamp") # Log the parsed result with human-readable duration hours = retry_after / 3600 @@ -581,6 +594,7 @@ def classify_error(e: Exception, provider: Optional[str] = None) -> ClassifiedEr original_exception=e, status_code=429, retry_after=retry_after, + quota_reset_timestamp=quota_reset_timestamp, ) except Exception as parse_error: lib_logger.debug( diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 599c4040..88e5a1d1 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -600,6 +600,7 @@ def parse_duration(duration_str: str) -> Optional[int]: "retry_after": None, "reason": None, "reset_timestamp": None, + "quota_reset_timestamp": None, # Unix timestamp for quota reset } for detail in details: @@ -626,8 +627,22 @@ def parse_duration(duration_str: str) -> Optional[int]: if parsed: result["retry_after"] = parsed - # Capture reset timestamp for logging - result["reset_timestamp"] = metadata.get("quotaResetTimeStamp") + # Capture reset timestamp for logging and authoritative reset time + reset_ts_str = metadata.get("quotaResetTimeStamp") + result["reset_timestamp"] = reset_ts_str + + # Parse ISO timestamp to Unix timestamp for usage tracking + if reset_ts_str: + try: + # Handle ISO format: "2025-12-11T22:53:16Z" + reset_dt = datetime.fromisoformat( + reset_ts_str.replace("Z", "+00:00") + ) + result["quota_reset_timestamp"] = reset_dt.timestamp() + except (ValueError, AttributeError) as e: + lib_logger.warning( + f"Failed to parse quota reset timestamp '{reset_ts_str}': {e}" + ) # Return None if we couldn't extract retry_after if not result["retry_after"]: @@ -826,45 +841,48 @@ def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: """ Get Antigravity-specific usage tracking configuration based on credential tier. - Antigravity has different quota reset windows by tier: - - Paid tiers (priority 1): 5-hour rolling window - - Free tier (priority 2): 7-day rolling window - - Unknown/legacy: 7-day rolling window (conservative default) + Antigravity uses per-model windows with different durations by tier: + - Paid tiers (priority 1): 5-hour per-model window + - Free tier (priority 2): 7-day per-model window + - Unknown/legacy: 7-day per-model window (conservative default) + + When a model hits a quota_exhausted 429 error with exact reset timestamp, + that timestamp becomes the authoritative reset time for the model (and its group). Args: credential: The credential path Returns: - Usage reset configuration dict + Usage reset configuration dict with mode="per_model" """ tier = self.project_tier_cache.get(credential) if not tier: tier = self._load_tier_from_file(credential) - # Paid tiers: 5-hour window + # Paid tiers: 5-hour per-model window if tier and tier not in ["free-tier", "legacy-tier", "unknown"]: return { "window_seconds": 5 * 60 * 60, # 18000 seconds = 5 hours - "field_name": "5h_window", + "mode": "per_model", "priority": 1, - "description": "5-hour rolling window (paid tier)", + "description": "5-hour per-model window (paid tier)", } - # Free tier: 7-day window + # Free tier: 7-day per-model window if tier == "free-tier": return { "window_seconds": 7 * 24 * 60 * 60, # 604800 seconds = 7 days - "field_name": "weekly", + "mode": "per_model", "priority": 2, - "description": "7-day rolling window (free tier)", + "description": "7-day per-model window (free tier)", } - # Unknown/legacy: use 7-day window as conservative default + # Unknown/legacy: use 7-day per-model window as conservative default return { "window_seconds": 7 * 24 * 60 * 60, # 604800 seconds = 7 days - "field_name": "weekly", + "mode": "per_model", "priority": 10, - "description": "7-day rolling window (unknown tier - conservative default)", + "description": "7-day per-model window (unknown tier - conservative default)", } def get_default_usage_field_name(self) -> str: @@ -872,9 +890,51 @@ def get_default_usage_field_name(self) -> str: Get the default usage tracking field name for Antigravity. Returns: - "weekly" as the conservative default for unknown credentials + "models" for per-model tracking + """ + return "models" + + # ========================================================================= + # Model Quota Grouping + # ========================================================================= + + # Models that share quota timing - when one hits quota, all get same reset time + QUOTA_GROUPS = { + # Future: add claude/gemini groups if they share quota + } + + def get_model_quota_group(self, model: str) -> Optional[str]: + """ + Returns the quota group name for a model. + + Claude models (sonnet and opus) share quota on Antigravity. + When one hits quota exhausted, all models in the group get the same reset time. + + Args: + model: Model name (with or without "antigravity/" prefix) + + Returns: + Group name ("claude") or None if not grouped + """ + # Remove provider prefix if present + clean_model = model.replace("antigravity/", "") + + for group_name, models in self.QUOTA_GROUPS.items(): + if clean_model in models: + return group_name + return None + + def get_models_in_quota_group(self, group: str) -> List[str]: + """ + Returns all model names in a quota group. + + Args: + group: Group name (e.g., "claude") + + Returns: + List of model names (without provider prefix) """ - return "weekly" + return self.QUOTA_GROUPS.get(group, []) async def initialize_credentials(self, credential_paths: List[str]) -> None: """ diff --git a/src/rotator_library/providers/provider_interface.py b/src/rotator_library/providers/provider_interface.py index e12cbabc..1cc8879e 100644 --- a/src/rotator_library/providers/provider_interface.py +++ b/src/rotator_library/providers/provider_interface.py @@ -202,6 +202,7 @@ def parse_quota_error( "retry_after": int, # seconds until quota resets "reason": str, # e.g., "QUOTA_EXHAUSTED", "RATE_LIMITED" "reset_timestamp": str | None, # ISO timestamp if available + "quota_reset_timestamp": float | None, # Unix timestamp for quota reset } """ return None # Default: no provider-specific parsing @@ -218,9 +219,9 @@ def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: credential tier (e.g., paid vs free accounts with different quota periods). The UsageManager will use this configuration to: - 1. Track usage in a custom-named field (instead of default "daily") - 2. Reset usage based on a rolling window from first request - 3. Archive stats to "global" when the window expires + 1. Track usage per-model or per-credential based on mode + 2. Reset usage based on a rolling window OR quota exhausted timestamp + 3. Archive stats to "global" when the window/quota expires Args: credential: The credential identifier (API key or path) @@ -229,32 +230,35 @@ def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: None to use default daily reset, otherwise a dict with: { "window_seconds": int, # Duration in seconds (e.g., 18000 for 5h) - "field_name": str, # Custom field name (e.g., "5h_window", "weekly") - "priority": int, # Priority level this config applies to (for docs) + "mode": str, # "credential" or "per_model" + "priority": int, # Priority level this config applies to "description": str, # Human-readable description (for logging) } + Modes: + - "credential": One window per credential. Window starts from first + request of ANY model. All models reset together when window expires. + - "per_model": Separate window per model (or model group). Window starts + from first request of THAT model. Models reset independently unless + grouped. If a quota_exhausted error provides exact reset time, that + becomes the authoritative reset time for the model. + Examples: - Antigravity paid tier: + Antigravity paid tier (per-model): { "window_seconds": 18000, # 5 hours - "field_name": "5h_window", + "mode": "per_model", "priority": 1, - "description": "5-hour rolling window (paid tier)" + "description": "5-hour per-model window (paid tier)" } - Antigravity free tier: + Default provider (credential-level): { - "window_seconds": 604800, # 7 days - "field_name": "weekly", - "priority": 2, - "description": "7-day rolling window (free tier)" + "window_seconds": 86400, # 24 hours + "mode": "credential", + "priority": 1, + "description": "24-hour credential window" } - - Note: - - window_seconds: Time from first request until stats reset - - When window expires, stats move to "global" (same as daily reset) - - First request after window expiry starts a new window """ return None # Default: use daily reset at daily_reset_time_utc @@ -269,3 +273,39 @@ def get_default_usage_field_name(self) -> str: Field name string (default: "daily") """ return "daily" + + # ========================================================================= + # Model Quota Grouping + # ========================================================================= + + def get_model_quota_group(self, model: str) -> Optional[str]: + """ + Returns the quota group name for a model, or None if not grouped. + + Models in the same quota group share cooldown timing - when one model + hits a quota exhausted error, all models in the group get the same + reset timestamp. They also reset (archive stats) together. + + This is useful for providers where multiple model variants share the + same underlying quota (e.g., Claude Sonnet and Opus on Antigravity). + + Args: + model: Model name (with or without provider prefix) + + Returns: + Group name string (e.g., "claude") or None if model is not grouped + """ + return None + + def get_models_in_quota_group(self, group: str) -> List[str]: + """ + Returns all model names that belong to a quota group. + + Args: + group: Group name (e.g., "claude") + + Returns: + List of model names (WITHOUT provider prefix) in the group. + Empty list if group doesn't exist. + """ + return [] diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 1ae93277..7e0fef4b 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -162,6 +162,69 @@ def _get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: return None + def _get_reset_mode(self, credential: str) -> str: + """ + Get the reset mode for a credential: 'credential' or 'per_model'. + + Args: + credential: The credential identifier + + Returns: + "per_model" or "credential" (default) + """ + config = self._get_usage_reset_config(credential) + return config.get("mode", "credential") if config else "credential" + + def _get_model_quota_group(self, credential: str, model: str) -> Optional[str]: + """ + Get the quota group for a model, if the provider defines one. + + Args: + credential: The credential identifier + model: Model name (with or without provider prefix) + + Returns: + Group name (e.g., "claude") or None if not grouped + """ + provider = self._get_provider_from_credential(credential) + if not provider: + return None + + plugin = self.provider_plugins.get(provider) + if not plugin: + return None + + if hasattr(plugin, "get_model_quota_group"): + return plugin.get_model_quota_group(model) + + return None + + def _get_grouped_models(self, credential: str, group: str) -> List[str]: + """ + Get all model names in a quota group (with provider prefix). + + Args: + credential: The credential identifier + group: Group name (e.g., "claude") + + Returns: + List of full model names (e.g., ["antigravity/claude-opus-4-5", ...]) + """ + provider = self._get_provider_from_credential(credential) + if not provider: + return [] + + plugin = self.provider_plugins.get(provider) + if not plugin: + return [] + + if hasattr(plugin, "get_models_in_quota_group"): + models = plugin.get_models_in_quota_group(group) + # Add provider prefix + return [f"{provider}/{m}" for m in models] + + return [] + def _get_usage_field_name(self, credential: str) -> str: """ Get the usage tracking field name for a credential. @@ -190,27 +253,36 @@ def _get_usage_field_name(self, credential: str) -> str: def _get_usage_count(self, key: str, model: str) -> int: """ - Get the current usage count for a model from the appropriate usage field. + Get the current usage count for a model from the appropriate usage structure. + + Supports both: + - New per-model structure: {"models": {"model_name": {"success_count": N, ...}}} + - Legacy structure: {"daily": {"models": {"model_name": {"success_count": N, ...}}}} Args: key: Credential identifier model: Model name Returns: - Usage count (success_count) for the model in the current window/daily period + Usage count (success_count) for the model in the current window/period """ if self._usage_data is None: return 0 key_data = self._usage_data.get(key, {}) - usage_field = self._get_usage_field_name(key) + reset_mode = self._get_reset_mode(key) - return ( - key_data.get(usage_field, {}) - .get("models", {}) - .get(model, {}) - .get("success_count", 0) - ) + if reset_mode == "per_model": + # New per-model structure: key_data["models"][model]["success_count"] + return key_data.get("models", {}).get(model, {}).get("success_count", 0) + else: + # Legacy structure: key_data["daily"]["models"][model]["success_count"] + return ( + key_data.get("daily", {}) + .get("models", {}) + .get(model, {}) + .get("success_count", 0) + ) def _select_sequential( self, @@ -299,9 +371,10 @@ async def _reset_daily_stats_if_needed(self): """ Checks if usage stats need to be reset for any key. - Supports two reset modes: - 1. Provider-specific rolling windows (e.g., 5h for Antigravity paid, 7d for free) - 2. Legacy daily reset at daily_reset_time_utc for providers without custom config + Supports three reset modes: + 1. per_model: Each model has its own window, resets based on quota_reset_ts or fallback window + 2. credential: One window per credential (legacy with custom window duration) + 3. daily: Legacy daily reset at daily_reset_time_utc """ if self._usage_data is None: return @@ -312,16 +385,23 @@ async def _reset_daily_stats_if_needed(self): needs_saving = False for key, data in self._usage_data.items(): - # Check for provider-specific reset configuration reset_config = self._get_usage_reset_config(key) if reset_config: - # Provider-specific rolling window reset - needs_saving |= await self._check_window_reset( - key, data, reset_config, now_ts - ) + reset_mode = reset_config.get("mode", "credential") + + if reset_mode == "per_model": + # Per-model window reset + needs_saving |= await self._check_per_model_resets( + key, data, reset_config, now_ts + ) + else: + # Credential-level window reset (legacy) + needs_saving |= await self._check_window_reset( + key, data, reset_config, now_ts + ) elif self.daily_reset_time_utc: - # Legacy daily reset for providers without custom config + # Legacy daily reset needs_saving |= await self._check_daily_reset( key, data, now_utc, today_str, now_ts ) @@ -329,6 +409,170 @@ async def _reset_daily_stats_if_needed(self): if needs_saving: await self._save_usage() + async def _check_per_model_resets( + self, + key: str, + data: Dict[str, Any], + reset_config: Dict[str, Any], + now_ts: float, + ) -> bool: + """ + Check and perform per-model resets for a credential. + + Each model resets independently based on: + 1. quota_reset_ts (authoritative, from quota exhausted error) if set + 2. window_start_ts + window_seconds (fallback) otherwise + + Grouped models reset together - all models in a group must be ready. + + Args: + key: Credential identifier + data: Usage data for this credential + reset_config: Provider's reset configuration + now_ts: Current timestamp + + Returns: + True if data was modified and needs saving + """ + window_seconds = reset_config.get("window_seconds", 86400) + models_data = data.get("models", {}) + + if not models_data: + return False + + modified = False + processed_groups = set() + + for model, model_data in list(models_data.items()): + # Check if this model is in a quota group + group = self._get_model_quota_group(key, model) + + if group: + if group in processed_groups: + continue # Already handled this group + + # Check if entire group should reset + if self._should_group_reset( + key, group, models_data, window_seconds, now_ts + ): + # Archive and reset all models in group + grouped_models = self._get_grouped_models(key, group) + archived_count = 0 + + for grouped_model in grouped_models: + if grouped_model in models_data: + gm_data = models_data[grouped_model] + self._archive_model_to_global(data, grouped_model, gm_data) + self._reset_model_data(gm_data) + archived_count += 1 + + if archived_count > 0: + lib_logger.info( + f"Reset model group '{group}' ({archived_count} models) for {mask_credential(key)}" + ) + modified = True + + processed_groups.add(group) + + else: + # Ungrouped model - check individually + if self._should_model_reset(model_data, window_seconds, now_ts): + self._archive_model_to_global(data, model, model_data) + self._reset_model_data(model_data) + lib_logger.info(f"Reset model {model} for {mask_credential(key)}") + modified = True + + # Preserve unexpired cooldowns + if modified: + self._preserve_unexpired_cooldowns(key, data, now_ts) + if "failures" in data: + data["failures"] = {} + + return modified + + def _should_model_reset( + self, model_data: Dict[str, Any], window_seconds: int, now_ts: float + ) -> bool: + """ + Check if a single model should reset. + + Returns True if: + - quota_reset_ts is set AND now >= quota_reset_ts, OR + - quota_reset_ts is NOT set AND now >= window_start_ts + window_seconds + """ + quota_reset = model_data.get("quota_reset_ts") + window_start = model_data.get("window_start_ts") + + if quota_reset: + return now_ts >= quota_reset + elif window_start: + return now_ts >= window_start + window_seconds + return False + + def _should_group_reset( + self, + key: str, + group: str, + models_data: Dict[str, Dict], + window_seconds: int, + now_ts: float, + ) -> bool: + """ + Check if all models in a group should reset. + + All models in the group must be ready to reset. + If any model has an active cooldown/window, the whole group waits. + """ + grouped_models = self._get_grouped_models(key, group) + + # Track if any model in group has data + any_has_data = False + + for grouped_model in grouped_models: + model_data = models_data.get(grouped_model, {}) + + if not model_data or ( + model_data.get("window_start_ts") is None + and model_data.get("success_count", 0) == 0 + ): + continue # No stats for this model yet + + any_has_data = True + + if not self._should_model_reset(model_data, window_seconds, now_ts): + return False # At least one model not ready + + return any_has_data + + def _archive_model_to_global( + self, data: Dict[str, Any], model: str, model_data: Dict[str, Any] + ) -> None: + """Archive a single model's stats to global.""" + global_data = data.setdefault("global", {"models": {}}) + global_model = global_data["models"].setdefault( + model, + { + "success_count": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "approx_cost": 0.0, + }, + ) + + global_model["success_count"] += model_data.get("success_count", 0) + global_model["prompt_tokens"] += model_data.get("prompt_tokens", 0) + global_model["completion_tokens"] += model_data.get("completion_tokens", 0) + global_model["approx_cost"] += model_data.get("approx_cost", 0.0) + + def _reset_model_data(self, model_data: Dict[str, Any]) -> None: + """Reset a model's window and stats.""" + model_data["window_start_ts"] = None + model_data["quota_reset_ts"] = None + model_data["success_count"] = 0 + model_data["prompt_tokens"] = 0 + model_data["completion_tokens"] = 0 + model_data["approx_cost"] = 0.0 + async def _check_window_reset( self, key: str, @@ -948,36 +1192,67 @@ async def record_success( Records a successful API call, resetting failure counters. It safely handles cases where token usage data is not available. - Uses provider-specific field names for usage tracking (e.g., "5h_window", "weekly") - and sets window start timestamp on first request. + Supports two modes based on provider configuration: + - per_model: Each model has its own window_start_ts and stats in key_data["models"] + - credential: Legacy mode with key_data["daily"]["models"] """ await self._lazy_init() async with self._data_lock: now_ts = time.time() today_utc_str = datetime.now(timezone.utc).date().isoformat() - # Determine the usage field name for this credential - usage_field = self._get_usage_field_name(key) reset_config = self._get_usage_reset_config(key) - uses_window = reset_config is not None + reset_mode = ( + reset_config.get("mode", "credential") if reset_config else "credential" + ) - # Initialize key data with appropriate structure - if uses_window: - # Provider-specific rolling window + if reset_mode == "per_model": + # New per-model structure key_data = self._usage_data.setdefault( key, { - usage_field: {"start_ts": None, "models": {}}, + "models": {}, "global": {"models": {}}, "model_cooldowns": {}, "failures": {}, }, ) - # Ensure the usage field exists (for migration from old format) - if usage_field not in key_data: - key_data[usage_field] = {"start_ts": None, "models": {}} + + # Ensure models dict exists + if "models" not in key_data: + key_data["models"] = {} + + # Get or create per-model data with window tracking + model_data = key_data["models"].setdefault( + model, + { + "window_start_ts": None, + "quota_reset_ts": None, + "success_count": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "approx_cost": 0.0, + }, + ) + + # Start window on first request for this model + if model_data.get("window_start_ts") is None: + model_data["window_start_ts"] = now_ts + window_hours = ( + reset_config.get("window_seconds", 0) / 3600 + if reset_config + else 0 + ) + lib_logger.info( + f"Started {window_hours:.1f}h window for model {model} on {mask_credential(key)}" + ) + + # Record stats + model_data["success_count"] += 1 + usage_data_ref = model_data # For token/cost recording below + else: - # Legacy daily reset + # Legacy credential-level structure key_data = self._usage_data.setdefault( key, { @@ -987,57 +1262,41 @@ async def record_success( "failures": {}, }, ) - usage_field = "daily" - # If the key is new, ensure its reset date is initialized to prevent an immediate reset. - if not uses_window and "last_daily_reset" not in key_data: - key_data["last_daily_reset"] = today_utc_str + if "last_daily_reset" not in key_data: + key_data["last_daily_reset"] = today_utc_str - # Always record a success and reset failures + # Get or create model data in daily structure + usage_data_ref = key_data["daily"]["models"].setdefault( + model, + { + "success_count": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "approx_cost": 0.0, + }, + ) + usage_data_ref["success_count"] += 1 + + # Reset failures for this model model_failures = key_data.setdefault("failures", {}).setdefault(model, {}) model_failures["consecutive_failures"] = 0 + + # Clear transient cooldown on success (but NOT quota_reset_ts) if model in key_data.get("model_cooldowns", {}): del key_data["model_cooldowns"][model] - # Get or create the usage field data - usage_data = key_data.setdefault(usage_field, {"models": {}}) - - # For window-based tracking, set start_ts on first request - if uses_window: - if usage_data.get("start_ts") is None: - usage_data["start_ts"] = now_ts - window_hours = reset_config.get("window_seconds", 0) / 3600 - description = reset_config.get("description", "rolling window") - lib_logger.info( - f"Starting new {window_hours:.1f}h window for {mask_credential(key)} - {description}" - ) - - # Ensure models dict exists - if "models" not in usage_data: - usage_data["models"] = {} - - model_data = usage_data["models"].setdefault( - model, - { - "success_count": 0, - "prompt_tokens": 0, - "completion_tokens": 0, - "approx_cost": 0.0, - }, - ) - model_data["success_count"] += 1 - - # Safely attempt to record token and cost usage + # Record token and cost usage if ( completion_response and hasattr(completion_response, "usage") and completion_response.usage ): usage = completion_response.usage - model_data["prompt_tokens"] += usage.prompt_tokens - model_data["completion_tokens"] += getattr( + usage_data_ref["prompt_tokens"] += usage.prompt_tokens + usage_data_ref["completion_tokens"] += getattr( usage, "completion_tokens", 0 - ) # Not present in embedding responses + ) lib_logger.info( f"Recorded usage from response object for key {mask_credential(key)}" ) @@ -1045,7 +1304,6 @@ async def record_success( provider_name = model.split("/")[0] provider_plugin = self.provider_plugins.get(provider_name) - # Check class attribute directly - no need to instantiate if provider_plugin and getattr( provider_plugin, "skip_cost_calculation", False ): @@ -1053,9 +1311,7 @@ async def record_success( f"Skipping cost calculation for provider '{provider_name}' (custom provider)." ) else: - # Differentiate cost calculation based on response type if isinstance(completion_response, litellm.EmbeddingResponse): - # Manually calculate cost for embeddings model_info = litellm.get_model_info(model) input_cost = model_info.get("input_cost_per_token") if input_cost: @@ -1070,7 +1326,7 @@ async def record_success( ) if cost is not None: - model_data["approx_cost"] += cost + usage_data_ref["approx_cost"] += cost except Exception as e: lib_logger.warning( f"Could not calculate cost for model {model}: {e}" @@ -1078,8 +1334,7 @@ async def record_success( elif isinstance(completion_response, asyncio.Future) or hasattr( completion_response, "__aiter__" ): - # This is an unconsumed stream object. Do not log a warning, as usage will be recorded from the chunks. - pass + pass # Stream - usage recorded from chunks else: lib_logger.warning( f"No usage data found in completion response for model {model}. Recording success without token count." @@ -1096,7 +1351,13 @@ async def record_failure( classified_error: ClassifiedError, increment_consecutive_failures: bool = True, ): - """Records a failure and applies cooldowns based on an escalating backoff strategy. + """Records a failure and applies cooldowns based on error type. + + Distinguishes between: + - quota_exceeded: Long cooldown with exact reset time (from quota_reset_timestamp) + Sets quota_reset_ts on model (and group) - this becomes authoritative stats reset time + - rate_limit: Short transient cooldown (just wait and retry) + Only sets model_cooldowns - does NOT affect stats reset timing Args: key: The API key or credential identifier @@ -1107,19 +1368,20 @@ async def record_failure( """ await self._lazy_init() async with self._data_lock: + now_ts = time.time() today_utc_str = datetime.now(timezone.utc).date().isoformat() - # Determine the usage field name for this credential - usage_field = self._get_usage_field_name(key) reset_config = self._get_usage_reset_config(key) - uses_window = reset_config is not None + reset_mode = ( + reset_config.get("mode", "credential") if reset_config else "credential" + ) # Initialize key data with appropriate structure - if uses_window: + if reset_mode == "per_model": key_data = self._usage_data.setdefault( key, { - usage_field: {"start_ts": None, "models": {}}, + "models": {}, "global": {"models": {}}, "model_cooldowns": {}, "failures": {}, @@ -1147,36 +1409,94 @@ async def record_failure( # Calculate cooldown duration based on error type cooldown_seconds = None + model_cooldowns = key_data.setdefault("model_cooldowns", {}) - if classified_error.error_type in ["rate_limit", "quota_exceeded"]: - # Rate limit / Quota errors: use retry_after if available, otherwise default to 60s + if classified_error.error_type == "quota_exceeded": + # Quota exhausted - use authoritative reset timestamp if available + quota_reset_ts = classified_error.quota_reset_timestamp cooldown_seconds = classified_error.retry_after or 60 - if classified_error.retry_after: - # Log with human-readable duration for provider-parsed cooldowns - hours = cooldown_seconds / 3600 - if hours >= 1: + + if quota_reset_ts and reset_mode == "per_model": + # Set quota_reset_ts on model - this becomes authoritative stats reset time + models_data = key_data.setdefault("models", {}) + model_data = models_data.setdefault( + model, + { + "window_start_ts": None, + "quota_reset_ts": None, + "success_count": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "approx_cost": 0.0, + }, + ) + model_data["quota_reset_ts"] = quota_reset_ts + + # Apply to all models in the same quota group + group = self._get_model_quota_group(key, model) + if group: + grouped_models = self._get_grouped_models(key, group) + for grouped_model in grouped_models: + group_model_data = models_data.setdefault( + grouped_model, + { + "window_start_ts": None, + "quota_reset_ts": None, + "success_count": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "approx_cost": 0.0, + }, + ) + group_model_data["quota_reset_ts"] = quota_reset_ts + # Also set transient cooldown for selection logic + model_cooldowns[grouped_model] = quota_reset_ts + + reset_dt = datetime.fromtimestamp( + quota_reset_ts, tz=timezone.utc + ) lib_logger.info( - f"Quota/rate limit on key {mask_credential(key)} for model {model}. " - f"Applying provider-specified cooldown: {cooldown_seconds}s ({hours:.1f}h)" + f"Quota exhausted for group '{group}' ({len(grouped_models)} models) " + f"on {mask_credential(key)}. Resets at {reset_dt.isoformat()}" ) else: + reset_dt = datetime.fromtimestamp( + quota_reset_ts, tz=timezone.utc + ) + hours = (quota_reset_ts - now_ts) / 3600 lib_logger.info( - f"Rate limit on key {mask_credential(key)} for model {model}. " - f"Applying provider-specified cooldown: {cooldown_seconds}s" + f"Quota exhausted for model {model} on {mask_credential(key)}. " + f"Resets at {reset_dt.isoformat()} ({hours:.1f}h)" ) + + # Set transient cooldown for selection logic + model_cooldowns[model] = quota_reset_ts else: + # No authoritative timestamp or legacy mode - just use retry_after + model_cooldowns[model] = now_ts + cooldown_seconds + hours = cooldown_seconds / 3600 lib_logger.info( - f"Rate limit on key {mask_credential(key)} for model {model}. " - f"Using default cooldown: {cooldown_seconds}s" + f"Quota exhausted on {mask_credential(key)} for model {model}. " + f"Cooldown: {cooldown_seconds}s ({hours:.1f}h)" ) + + elif classified_error.error_type == "rate_limit": + # Transient rate limit - just set short cooldown (does NOT set quota_reset_ts) + cooldown_seconds = classified_error.retry_after or 60 + model_cooldowns[model] = now_ts + cooldown_seconds + lib_logger.info( + f"Rate limit on {mask_credential(key)} for model {model}. " + f"Transient cooldown: {cooldown_seconds}s" + ) + elif classified_error.error_type == "authentication": # Apply a 5-minute key-level lockout for auth errors - key_data["key_cooldown_until"] = time.time() + 300 + key_data["key_cooldown_until"] = now_ts + 300 + cooldown_seconds = 300 + model_cooldowns[model] = now_ts + cooldown_seconds lib_logger.warning( f"Authentication error on key {mask_credential(key)}. Applying 5-minute key-level lockout." ) - # Auth errors still use escalating backoff for the specific model - cooldown_seconds = 300 # 5 minutes for model cooldown # If we should increment failures, calculate escalating backoff if should_increment: @@ -1190,35 +1510,27 @@ async def record_failure( # If cooldown wasn't set by specific error type, use escalating backoff if cooldown_seconds is None: backoff_tiers = {1: 10, 2: 30, 3: 60, 4: 120} - cooldown_seconds = backoff_tiers.get( - count, 7200 - ) # Default to 2 hours for "spent" keys + cooldown_seconds = backoff_tiers.get(count, 7200) + model_cooldowns[model] = now_ts + cooldown_seconds lib_logger.warning( f"Failure #{count} for key {mask_credential(key)} with model {model}. " - f"Error type: {classified_error.error_type}" + f"Error type: {classified_error.error_type}, cooldown: {cooldown_seconds}s" ) else: # Provider-level errors: apply short cooldown but don't count against key if cooldown_seconds is None: - cooldown_seconds = 30 # 30s cooldown for provider issues + cooldown_seconds = 30 + model_cooldowns[model] = now_ts + cooldown_seconds lib_logger.info( - f"Provider-level error ({classified_error.error_type}) for key {mask_credential(key)} with model {model}. " - f"NOT incrementing consecutive failures. Applying {cooldown_seconds}s cooldown." + f"Provider-level error ({classified_error.error_type}) for key {mask_credential(key)} " + f"with model {model}. NOT incrementing failures. Cooldown: {cooldown_seconds}s" ) - # Apply the cooldown - model_cooldowns = key_data.setdefault("model_cooldowns", {}) - model_cooldowns[model] = time.time() + cooldown_seconds - lib_logger.warning( - f"Cooldown applied for key {mask_credential(key)} with model {model}: {cooldown_seconds}s. " - f"Error type: {classified_error.error_type}" - ) - # Check for key-level lockout condition await self._check_key_lockout(key, key_data) key_data["last_failure"] = { - "timestamp": time.time(), + "timestamp": now_ts, "model": model, "error": str(classified_error.original_exception), } From 4bc76131c13b44b0735b462cfa8f7433ed66a7ba Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sat, 6 Dec 2025 07:03:30 +0100 Subject: [PATCH 090/221] =?UTF-8?q?refactor(client):=20=F0=9F=94=A8=20init?= =?UTF-8?q?ialize=20provider=20plugins=20before=20rotation=20mode=20detect?= =?UTF-8?q?ion?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move provider plugin initialization earlier in the constructor to ensure they are available when building the provider rotation modes map. This prevents potential issues where rotation mode detection logic might need access to provider instances before they were initialized. --- src/rotator_library/client.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 9e1a3042..4ca9d8cf 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -140,6 +140,10 @@ def __init__( self.global_timeout = global_timeout self.abort_on_callback_error = abort_on_callback_error + # Initialize provider plugins early so they can be used for rotation mode detection + self._provider_plugins = PROVIDER_PLUGINS + self._provider_instances = {} + # Build provider rotation modes map # Each provider can specify its preferred rotation mode ("balanced" or "sequential") provider_rotation_modes = {} @@ -164,8 +168,6 @@ def __init__( provider_plugins=PROVIDER_PLUGINS, ) self._model_list_cache = {} - self._provider_plugins = PROVIDER_PLUGINS - self._provider_instances = {} self.http_client = httpx.AsyncClient() self.all_providers = AllProviders() self.cooldown_manager = CooldownManager() From fd014827166b264a7bb45f2482ec772ba268ae74 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sat, 6 Dec 2025 07:04:55 +0100 Subject: [PATCH 091/221] =?UTF-8?q?refactor(usage):=20=F0=9F=94=A8=20cache?= =?UTF-8?q?=20provider=20plugin=20instances=20to=20reduce=20redundant=20in?= =?UTF-8?q?stantiation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduced a provider instance cache (`_provider_instances`) to store and reuse provider plugin instances across multiple method calls. - Added `_get_provider_instance()` helper method to centralize provider plugin instantiation logic with caching support - Refactored `_get_usage_reset_config()`, `_get_model_quota_group()`, `_get_models_in_quota_group()`, `_get_usage_field_name()`, and cost calculation logic to use the cached provider instances - Eliminated redundant provider plugin instantiation that occurred on every method call - Simplified error handling by consolidating null checks in the helper method This change improves performance by avoiding repeated instantiation of the same provider plugin objects and reduces code duplication across provider plugin access patterns. --- src/rotator_library/usage_manager.py | 75 ++++++++++++++++------------ 1 file changed, 44 insertions(+), 31 deletions(-) diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 7e0fef4b..39c8db6f 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -76,6 +76,7 @@ def __init__( self.rotation_tolerance = rotation_tolerance self.provider_rotation_modes = provider_rotation_modes or {} self.provider_plugins = provider_plugins or PROVIDER_PLUGINS + self._provider_instances: Dict[str, Any] = {} # Cache for provider instances self.key_states: Dict[str, Dict[str, Any]] = {} self._data_lock = asyncio.Lock() @@ -138,6 +139,33 @@ def _get_provider_from_credential(self, credential: str) -> Optional[str]: return None + def _get_provider_instance(self, provider: str) -> Optional[Any]: + """ + Get or create a provider plugin instance. + + Args: + provider: The provider name + + Returns: + Provider plugin instance or None + """ + if not provider: + return None + + plugin_class = self.provider_plugins.get(provider) + if not plugin_class: + return None + + # Get or create provider instance from cache + if provider not in self._provider_instances: + # Instantiate the plugin if it's a class, or use it directly if already an instance + if isinstance(plugin_class, type): + self._provider_instances[provider] = plugin_class() + else: + self._provider_instances[provider] = plugin_class + + return self._provider_instances[provider] + def _get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: """ Get the usage reset configuration for a credential from its provider plugin. @@ -150,15 +178,10 @@ def _get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: or None to use default daily reset. """ provider = self._get_provider_from_credential(credential) - if not provider: - return None + plugin_instance = self._get_provider_instance(provider) - plugin = self.provider_plugins.get(provider) - if not plugin: - return None - - if hasattr(plugin, "get_usage_reset_config"): - return plugin.get_usage_reset_config(credential) + if plugin_instance and hasattr(plugin_instance, "get_usage_reset_config"): + return plugin_instance.get_usage_reset_config(credential) return None @@ -187,15 +210,10 @@ def _get_model_quota_group(self, credential: str, model: str) -> Optional[str]: Group name (e.g., "claude") or None if not grouped """ provider = self._get_provider_from_credential(credential) - if not provider: - return None - - plugin = self.provider_plugins.get(provider) - if not plugin: - return None + plugin_instance = self._get_provider_instance(provider) - if hasattr(plugin, "get_model_quota_group"): - return plugin.get_model_quota_group(model) + if plugin_instance and hasattr(plugin_instance, "get_model_quota_group"): + return plugin_instance.get_model_quota_group(model) return None @@ -211,15 +229,10 @@ def _get_grouped_models(self, credential: str, group: str) -> List[str]: List of full model names (e.g., ["antigravity/claude-opus-4-5", ...]) """ provider = self._get_provider_from_credential(credential) - if not provider: - return [] - - plugin = self.provider_plugins.get(provider) - if not plugin: - return [] + plugin_instance = self._get_provider_instance(provider) - if hasattr(plugin, "get_models_in_quota_group"): - models = plugin.get_models_in_quota_group(group) + if plugin_instance and hasattr(plugin_instance, "get_models_in_quota_group"): + models = plugin_instance.get_models_in_quota_group(group) # Add provider prefix return [f"{provider}/{m}" for m in models] @@ -244,10 +257,10 @@ def _get_usage_field_name(self, credential: str) -> str: # Check provider default provider = self._get_provider_from_credential(credential) - if provider: - plugin = self.provider_plugins.get(provider) - if plugin and hasattr(plugin, "get_default_usage_field_name"): - return plugin.get_default_usage_field_name() + plugin_instance = self._get_provider_instance(provider) + + if plugin_instance and hasattr(plugin_instance, "get_default_usage_field_name"): + return plugin_instance.get_default_usage_field_name() return "daily" @@ -1302,10 +1315,10 @@ async def record_success( ) try: provider_name = model.split("/")[0] - provider_plugin = self.provider_plugins.get(provider_name) + provider_instance = self._get_provider_instance(provider_name) - if provider_plugin and getattr( - provider_plugin, "skip_cost_calculation", False + if provider_instance and getattr( + provider_instance, "skip_cost_calculation", False ): lib_logger.debug( f"Skipping cost calculation for provider '{provider_name}' (custom provider)." From 31c3d361ac17c3ea1604d84be8be1e355745a5c3 Mon Sep 17 00:00:00 2001 From: MasuRii Date: Sat, 6 Dec 2025 22:18:14 +0800 Subject: [PATCH 092/221] feat: add runtime resilience for file deletion survival Implement graceful degradation patterns that allow the proxy to continue running even if core files are deleted during runtime. Changes only take effect on restart, enabling safe development while the proxy is serving. ## Changes by Component ### Usage Manager (usage_manager.py) - Wrap `_save_usage()` in try/except with directory auto-recreation - Enhanced `_load_usage()` with explicit error handling - In-memory state continues working if file operations fail ### Failure Logger (failure_logger.py) - Add module-level `_file_handler` and `_fallback_mode` state - Create `_create_file_handler()` with directory auto-recreation - Create `_ensure_handler_valid()` for handler recovery - Use NullHandler as fallback when file logging fails ### Detailed Logger (detailed_logger.py) - Add class-level `_disk_available` and `_console_fallback_warned` flags - Add instance-level `_in_memory_logs` list for fallback storage - Skip disk writes gracefully when filesystem unavailable ### Google OAuth Base (google_oauth_base.py) - Update memory cache FIRST before disk write (memory-first pattern) - Use cached tokens as fallback when refresh/save fails - Log warnings but don't crash on persistence failures ### Provider Cache (provider_cache.py) - Add `_disk_available` health flag and `disk_errors` counter - Track disk health status in get_stats() - Gracefully degrade to memory-only caching on disk failures ### Documentation (DOCUMENTATION.md) - Add Section 5: Runtime Resilience with resilience hierarchy - Document "Develop While Running" workflow - Explain graceful degradation and data loss scenarios --- DOCUMENTATION.md | 31 ++++ src/proxy_app/detailed_logger.py | 31 +++- src/rotator_library/failure_logger.py | 101 +++++++++---- .../providers/google_oauth_base.py | 141 +++++++++++------- .../providers/provider_cache.py | 37 ++++- src/rotator_library/usage_manager.py | 45 +++++- 6 files changed, 294 insertions(+), 92 deletions(-) diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index cf985326..de340c15 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -697,4 +697,35 @@ To facilitate robust debugging, the proxy includes a comprehensive transaction l This level of detail allows developers to trace exactly why a request failed or why a specific key was rotated. +--- + +## 5. Runtime Resilience + +The proxy is engineered to maintain high availability even in the face of runtime filesystem disruptions. This "Runtime Resilience" capability ensures that the service continues to process API requests even if core data directories (like `logs/`, `oauth_creds/`) or files are accidentally deleted or become unwritable while the application is running. + +### 5.1. Resilience Hierarchy + +The system follows a strict hierarchy of survival: + +1. **Core API Handling (Level 1)**: The Python runtime keeps all necessary code in memory (`sys.modules`). Deleting source code files while the proxy is running will **not** crash active requests. +2. **Credential Management (Level 2)**: OAuth tokens are aggressively cached in memory. If credential files are deleted, the proxy continues using the cached tokens. If a token needs refresh and the file cannot be written, the new token is updated in memory only. +3. **Usage Tracking (Level 3)**: Usage statistics (`key_usage.json`) are maintained in memory. If the file is deleted, the system tracks usage internally. It attempts to recreate the file/directory on the next save interval. If save fails, data is effectively "memory-only" until the next successful write. +4. **Logging (Level 4)**: Logging is treated as non-critical. If the `logs/` directory is removed, the system attempts to recreate it. If creation fails (e.g., permission error), logging degrades gracefully (stops or falls back to console) without interrupting the request flow. + +### 5.2. "Develop While Running" + +This architecture supports a robust development workflow: + +* **Log Cleanup**: You can safely run `rm -rf logs/` while the proxy is serving traffic. The system will simply recreate the directory structure on the next request. +* **Config Reset**: Deleting `key_usage.json` resets the persistence layer, but the running instance preserves its current in-memory counts to ensure load balancing consistency. +* **File Recovery**: If you delete a critical file, the system attempts **Directory Auto-Recreation** before every write operation. + +### 5.3. Graceful Degradation & Data Loss + +While functionality is preserved, persistence may be compromised during filesystem failures: + +* **Logs**: If disk writes fail, detailed request logs may be lost (unless console fallback is active). +* **Usage Stats**: If `key_usage.json` cannot be written, usage data since the last successful save will be lost upon application restart. +* **Credentials**: Refreshed tokens held only in memory will require re-authentication after a restart if they cannot be persisted to disk. + diff --git a/src/proxy_app/detailed_logger.py b/src/proxy_app/detailed_logger.py index 4ebaf7e9..107a05cf 100644 --- a/src/proxy_app/detailed_logger.py +++ b/src/proxy_app/detailed_logger.py @@ -13,6 +13,10 @@ class DetailedLogger: """ Logs comprehensive details of each API transaction to a unique, timestamped directory. """ + # Class-level fallback flags for resilience + _disk_available = True + _console_fallback_warned = False + def __init__(self): """ Initializes the logger for a single request, creating a unique directory to store all related log files. @@ -21,16 +25,33 @@ def __init__(self): self.request_id = str(uuid.uuid4()) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") self.log_dir = DETAILED_LOGS_DIR / f"{timestamp}_{self.request_id}" - self.log_dir.mkdir(parents=True, exist_ok=True) self.streaming = False + self._in_memory_logs = [] # Fallback storage + + # Attempt directory creation with resilience + try: + self.log_dir.mkdir(parents=True, exist_ok=True) + DetailedLogger._disk_available = True + except (OSError, PermissionError) as e: + DetailedLogger._disk_available = False + if not DetailedLogger._console_fallback_warned: + logging.warning(f"Detailed logging disabled - cannot create log directory: {e}") + DetailedLogger._console_fallback_warned = True def _write_json(self, filename: str, data: Dict[str, Any]): """Helper to write data to a JSON file in the log directory.""" + if not DetailedLogger._disk_available: + self._in_memory_logs.append({"file": filename, "data": data}) + return + try: + # Attempt directory recreation if needed + self.log_dir.mkdir(parents=True, exist_ok=True) with open(self.log_dir / filename, "w", encoding="utf-8") as f: json.dump(data, f, indent=4, ensure_ascii=False) - except Exception as e: + except (OSError, PermissionError, IOError) as e: logging.error(f"[{self.request_id}] Failed to write to {filename}: {e}") + self._in_memory_logs.append({"file": filename, "data": data}) def log_request(self, headers: Dict[str, Any], body: Dict[str, Any]): """Logs the initial request details.""" @@ -45,14 +66,18 @@ def log_request(self, headers: Dict[str, Any], body: Dict[str, Any]): def log_stream_chunk(self, chunk: Dict[str, Any]): """Logs an individual chunk from a streaming response to a JSON Lines file.""" + if not DetailedLogger._disk_available: + return # Skip chunk logging when disk unavailable + try: + self.log_dir.mkdir(parents=True, exist_ok=True) log_entry = { "timestamp_utc": datetime.utcnow().isoformat(), "chunk": chunk } with open(self.log_dir / "streaming_chunks.jsonl", "a", encoding="utf-8") as f: f.write(json.dumps(log_entry, ensure_ascii=False) + "\n") - except Exception as e: + except (OSError, PermissionError, IOError) as e: logging.error(f"[{self.request_id}] Failed to write stream chunk: {e}") def log_final_response(self, status_code: int, headers: Optional[Dict[str, Any]], body: Dict[str, Any]): diff --git a/src/rotator_library/failure_logger.py b/src/rotator_library/failure_logger.py index 8f1848ae..9379d34e 100644 --- a/src/rotator_library/failure_logger.py +++ b/src/rotator_library/failure_logger.py @@ -5,41 +5,76 @@ from datetime import datetime from .error_handler import mask_credential +# Module-level state for resilience +_file_handler = None +_fallback_mode = False -def setup_failure_logger(): - """Sets up a dedicated JSON logger for writing detailed failure logs to a file.""" - log_dir = "logs" - if not os.path.exists(log_dir): - os.makedirs(log_dir) - # Create a logger specifically for failures. - # This logger will NOT propagate to the root logger. - logger = logging.getLogger("failure_logger") - logger.setLevel(logging.INFO) - logger.propagate = False +# Custom JSON formatter for structured logs (defined at module level for reuse) +class JsonFormatter(logging.Formatter): + def format(self, record): + # The message is already a dict, so we just format it as a JSON string + return json.dumps(record.msg) - # Use a rotating file handler - handler = RotatingFileHandler( - os.path.join(log_dir, "failures.log"), - maxBytes=5 * 1024 * 1024, # 5 MB - backupCount=2, - ) - # Custom JSON formatter for structured logs - class JsonFormatter(logging.Formatter): - def format(self, record): - # The message is already a dict, so we just format it as a JSON string - return json.dumps(record.msg) +def _create_file_handler(): + """Create file handler with directory auto-recreation.""" + global _file_handler, _fallback_mode + log_dir = "logs" + + try: + if not os.path.exists(log_dir): + os.makedirs(log_dir, exist_ok=True) + + handler = RotatingFileHandler( + os.path.join(log_dir, "failures.log"), + maxBytes=5 * 1024 * 1024, # 5 MB + backupCount=2, + ) + + handler.setFormatter(JsonFormatter()) + _file_handler = handler + _fallback_mode = False + return handler + except (OSError, PermissionError, IOError) as e: + logging.warning(f"Cannot create failure log file handler: {e}") + _fallback_mode = True + return None - handler.setFormatter(JsonFormatter()) - # Add handler only if it hasn't been added before - if not logger.handlers: +def setup_failure_logger(): + """Sets up a dedicated JSON logger for writing detailed failure logs.""" + logger = logging.getLogger("failure_logger") + logger.setLevel(logging.INFO) + logger.propagate = False + + # Remove existing handlers to prevent duplicates + logger.handlers.clear() + + # Try to add file handler + handler = _create_file_handler() + if handler: logger.addHandler(handler) - + + # Always add a NullHandler as fallback to prevent "no handlers" warning + if not logger.handlers: + logger.addHandler(logging.NullHandler()) + return logger +def _ensure_handler_valid(): + """Check if file handler is still valid, recreate if needed.""" + global _file_handler, _fallback_mode + + if _file_handler is None or _fallback_mode: + handler = _create_file_handler() + if handler: + failure_logger = logging.getLogger("failure_logger") + failure_logger.handlers.clear() + failure_logger.addHandler(handler) + + # Initialize the dedicated logger for detailed failure logs failure_logger = setup_failure_logger() @@ -145,11 +180,23 @@ def log_failure( "request_headers": request_headers, "error_chain": error_chain if len(error_chain) > 1 else None, } - failure_logger.error(detailed_log_data) - + # 2. Log a concise summary to the main library logger, which will propagate summary_message = ( f"API call failed for model {model} with key {mask_credential(api_key)}. " f"Error: {type(error).__name__}. See failures.log for details." ) + + # Attempt to ensure handler is valid before logging + _ensure_handler_valid() + + # Wrap the actual log call with resilience + try: + failure_logger.error(detailed_log_data) + except (OSError, IOError) as e: + # File logging failed - log to console instead + logging.error(f"Failed to write to failures.log: {e}") + logging.error(f"Failure summary: {summary_message}") + + # Console log always succeeds main_lib_logger.error(summary_message) diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index 0b34153b..a5ca9f4f 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -260,64 +260,76 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: ) async def _save_credentials(self, path: str, creds: Dict[str, Any]): + """Save credentials with in-memory fallback if disk unavailable. + + [RUNTIME RESILIENCE] Always updates the in-memory cache first (memory is reliable), + then attempts disk persistence. If disk write fails, logs a warning but does NOT + raise an exception - the in-memory state continues to work. + """ + # [IN-MEMORY FIRST] Always update cache first (reliable) + self._credentials_cache[path] = creds + # Don't save to file if credentials were loaded from environment if creds.get("_proxy_metadata", {}).get("loaded_from_env"): lib_logger.debug("Credentials loaded from env, skipping file save") - # Still update cache for in-memory consistency - self._credentials_cache[path] = creds return - # [ATOMIC WRITE] Use tempfile + move pattern to ensure atomic writes - # This prevents credential corruption if the process is interrupted during write - parent_dir = os.path.dirname(os.path.abspath(path)) - os.makedirs(parent_dir, exist_ok=True) - - tmp_fd = None - tmp_path = None try: - # Create temp file in same directory as target (ensures same filesystem) - tmp_fd, tmp_path = tempfile.mkstemp( - dir=parent_dir, prefix=".tmp_", suffix=".json", text=True - ) - - # Write JSON to temp file - with os.fdopen(tmp_fd, "w") as f: - json.dump(creds, f, indent=2) - tmp_fd = None # fdopen closes the fd + # [ATOMIC WRITE] Use tempfile + move pattern to ensure atomic writes + # This prevents credential corruption if the process is interrupted during write + parent_dir = os.path.dirname(os.path.abspath(path)) + os.makedirs(parent_dir, exist_ok=True) - # Set secure permissions (0600 = owner read/write only) + tmp_fd = None + tmp_path = None try: - os.chmod(tmp_path, 0o600) - except (OSError, AttributeError): - # Windows may not support chmod, ignore - pass - - # Atomic move (overwrites target if it exists) - shutil.move(tmp_path, path) - tmp_path = None # Successfully moved + # Create temp file in same directory as target (ensures same filesystem) + tmp_fd, tmp_path = tempfile.mkstemp( + dir=parent_dir, prefix=".tmp_", suffix=".json", text=True + ) - # Update cache AFTER successful file write (prevents cache/file inconsistency) - self._credentials_cache[path] = creds - lib_logger.debug( - f"Saved updated {self.ENV_PREFIX} OAuth credentials to '{path}' (atomic write)." - ) + # Write JSON to temp file + with os.fdopen(tmp_fd, "w") as f: + json.dump(creds, f, indent=2) + tmp_fd = None # fdopen closes the fd - except Exception as e: - lib_logger.error( - f"Failed to save updated {self.ENV_PREFIX} OAuth credentials to '{path}': {e}" - ) - # Clean up temp file if it still exists - if tmp_fd is not None: + # Set secure permissions (0600 = owner read/write only) try: - os.close(tmp_fd) - except: + os.chmod(tmp_path, 0o600) + except (OSError, AttributeError): + # Windows may not support chmod, ignore pass - if tmp_path and os.path.exists(tmp_path): - try: - os.unlink(tmp_path) - except: - pass - raise + + # Atomic move (overwrites target if it exists) + shutil.move(tmp_path, path) + tmp_path = None # Successfully moved + + lib_logger.debug( + f"Saved updated {self.ENV_PREFIX} OAuth credentials to '{path}' (atomic write)." + ) + + except Exception as e: + # Clean up temp file if it still exists + if tmp_fd is not None: + try: + os.close(tmp_fd) + except: + pass + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except: + pass + raise + + except (OSError, PermissionError, IOError) as e: + # [FAIL SILENTLY, LOG LOUDLY] Log the error but don't crash + # The in-memory cache was already updated, so we can continue operating + lib_logger.warning( + f"Failed to save credentials to {path}: {e}. " + "Credentials cached in memory only (will be lost on restart)." + ) + # Don't raise - we already updated the memory cache def _is_token_expired(self, creds: Dict[str, Any]) -> bool: expiry = creds.get("token_expiry") # gcloud format @@ -841,10 +853,39 @@ async def handle_callback(reader, writer): ) async def get_auth_header(self, credential_path: str) -> Dict[str, str]: - creds = await self._load_credentials(credential_path) - if self._is_token_expired(creds): - creds = await self._refresh_token(credential_path, creds) - return {"Authorization": f"Bearer {creds['access_token']}"} + """Get auth header with graceful degradation if refresh fails. + + [RUNTIME RESILIENCE] If credential file is deleted or refresh fails, + attempts to use cached credentials. This allows the proxy to continue + operating with potentially stale tokens rather than crashing. + """ + try: + creds = await self._load_credentials(credential_path) + if self._is_token_expired(creds): + try: + creds = await self._refresh_token(credential_path, creds) + except Exception as e: + # [CACHED TOKEN FALLBACK] Check if we have a cached token that might still work + cached = self._credentials_cache.get(credential_path) + if cached and cached.get("access_token"): + lib_logger.warning( + f"Token refresh failed for {Path(credential_path).name}: {e}. " + "Using cached token (may be expired)." + ) + creds = cached + else: + raise + return {"Authorization": f"Bearer {creds['access_token']}"} + except Exception as e: + # [FINAL FALLBACK] Check if any cached credential exists as last resort + cached = self._credentials_cache.get(credential_path) + if cached and cached.get("access_token"): + lib_logger.error( + f"Credential load failed for {credential_path}: {e}. " + "Using stale cached token as last resort." + ) + return {"Authorization": f"Bearer {cached['access_token']}"} + raise async def get_user_info( self, creds_or_path: Union[Dict[str, Any], str] diff --git a/src/rotator_library/providers/provider_cache.py b/src/rotator_library/providers/provider_cache.py index b6bb2db6..1e7f85e6 100644 --- a/src/rotator_library/providers/provider_cache.py +++ b/src/rotator_library/providers/provider_cache.py @@ -104,7 +104,10 @@ def __init__( self._running = False # Statistics - self._stats = {"memory_hits": 0, "disk_hits": 0, "misses": 0, "writes": 0} + self._stats = {"memory_hits": 0, "disk_hits": 0, "misses": 0, "writes": 0, "disk_errors": 0} + + # [RUNTIME RESILIENCE] Track disk health for monitoring + self._disk_available = True # Metadata about this cache instance self._cache_name = cache_file.stem if cache_file else "unnamed" @@ -171,13 +174,27 @@ async def _load_from_disk(self) -> None: # ========================================================================= async def _save_to_disk(self) -> None: - """Persist cache to disk using atomic write.""" + """Persist cache to disk using atomic write with health tracking. + + [RUNTIME RESILIENCE] Tracks disk health and records errors. If disk + operations fail, the memory cache continues to work. Health status + is available via get_stats() for monitoring. + """ if not self._enable_disk: return try: async with self._disk_lock: - self._cache_file.parent.mkdir(parents=True, exist_ok=True) + # [DIRECTORY AUTO-RECREATION] Attempt to create directory + try: + self._cache_file.parent.mkdir(parents=True, exist_ok=True) + except (OSError, PermissionError) as e: + self._stats["disk_errors"] += 1 + self._disk_available = False + lib_logger.warning( + f"ProviderCache[{self._cache_name}]: Cannot create cache directory: {e}" + ) + return cache_data = { "version": "1.0", @@ -210,6 +227,8 @@ async def _save_to_disk(self) -> None: shutil.move(tmp_path, self._cache_file) self._stats["writes"] += 1 + # [RUNTIME RESILIENCE] Mark disk as healthy on success + self._disk_available = True lib_logger.debug( f"ProviderCache[{self._cache_name}]: Saved {len(self._cache)} entries" ) @@ -218,6 +237,9 @@ async def _save_to_disk(self) -> None: os.unlink(tmp_path) raise except Exception as e: + # [RUNTIME RESILIENCE] Track disk errors for monitoring + self._stats["disk_errors"] += 1 + self._disk_available = False lib_logger.error(f"ProviderCache[{self._cache_name}]: Disk save failed: {e}") # ========================================================================= @@ -416,12 +438,17 @@ def contains(self, key: str) -> bool: return False def get_stats(self) -> Dict[str, Any]: - """Get cache statistics.""" + """Get cache statistics including disk health. + + [RUNTIME RESILIENCE] Includes disk_available flag for monitoring + the health of disk persistence. + """ return { **self._stats, "memory_entries": len(self._cache), "dirty": self._dirty, - "disk_enabled": self._enable_disk + "disk_enabled": self._enable_disk, + "disk_available": self._disk_available # [RUNTIME RESILIENCE] Health indicator } async def clear(self) -> None: diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 577bf4aa..1defd7ae 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -90,25 +90,56 @@ async def _lazy_init(self): self._initialized.set() async def _load_usage(self): - """Loads usage data from the JSON file asynchronously.""" + """Loads usage data from the JSON file asynchronously with enhanced resilience. + + [RUNTIME RESILIENCE] Handles various file system errors gracefully, + including race conditions where file is deleted between exists check and open. + """ async with self._data_lock: if not os.path.exists(self.file_path): self._usage_data = {} return + try: async with aiofiles.open(self.file_path, "r") as f: content = await f.read() - self._usage_data = json.loads(content) - except (json.JSONDecodeError, IOError, FileNotFoundError): + self._usage_data = json.loads(content) if content.strip() else {} + except FileNotFoundError: + # [RACE CONDITION HANDLING] File deleted between exists check and open + self._usage_data = {} + except json.JSONDecodeError as e: + lib_logger.warning(f"Corrupted usage file {self.file_path}: {e}. Starting fresh.") + self._usage_data = {} + except (OSError, PermissionError, IOError) as e: + lib_logger.warning(f"Cannot read usage file {self.file_path}: {e}. Using empty state.") self._usage_data = {} async def _save_usage(self): - """Saves the current usage data to the JSON file asynchronously.""" + """Saves the current usage data to the JSON file asynchronously with resilience. + + [RUNTIME RESILIENCE] Wraps file operations in try/except to prevent crashes + if the file or directory is deleted during runtime. The in-memory state + continues to work even if disk persistence fails. + """ if self._usage_data is None: return - async with self._data_lock: - async with aiofiles.open(self.file_path, "w") as f: - await f.write(json.dumps(self._usage_data, indent=2)) + + try: + async with self._data_lock: + # [DIRECTORY AUTO-RECREATION] Ensure directory exists before write + file_dir = os.path.dirname(os.path.abspath(self.file_path)) + if file_dir and not os.path.exists(file_dir): + os.makedirs(file_dir, exist_ok=True) + + async with aiofiles.open(self.file_path, "w") as f: + await f.write(json.dumps(self._usage_data, indent=2)) + except (OSError, PermissionError, IOError) as e: + # [FAIL SILENTLY, LOG LOUDLY] Log the error but don't crash + # In-memory state is preserved and will continue to work + lib_logger.warning( + f"Failed to save usage data to {self.file_path}: {e}. " + "Data will be retained in memory but may be lost on restart." + ) async def _reset_daily_stats_if_needed(self): """Checks if daily stats need to be reset for any key.""" From 3c52746ba68d4614f1cab3cd4f3891742630a50e Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Sun, 7 Dec 2025 06:32:24 +0100 Subject: [PATCH 093/221] =?UTF-8?q?refactor(providers):=20=F0=9F=94=A8=20c?= =?UTF-8?q?entralize=20tier=20and=20quota=20configuration=20in=20ProviderI?= =?UTF-8?q?nterface?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidate provider-specific tier prioritization, usage reset configuration, and quota group logic into the base ProviderInterface class to eliminate code duplication and establish a single source of truth. - Introduce UsageResetConfigDef dataclass for declarative usage configuration - Add tier_priorities, usage_reset_configs, and model_quota_groups as class attributes - Implement centralized _resolve_tier_priority() and _build_usage_reset_config() methods - Move get_credential_priority() and get_usage_reset_config() logic to base class - Add environment variable override support for quota groups (QUOTA_GROUPS_{PROVIDER}_{GROUP}) - Remove duplicate priority/usage logic from AntigravityProvider and GeminiCliProvider - Update .env.example with comprehensive documentation for quota group configuration This refactoring allows providers to define their tier system, usage windows, and quota groups purely through class attributes, while the base class handles all resolution logic. Providers now only need to override get_credential_tier_name() for tier lookup. --- .env.example | 20 ++ .../providers/antigravity_provider.py | 184 +++--------- .../providers/gemini_cli_provider.py | 70 ++--- .../providers/provider_interface.py | 270 +++++++++++++++--- 4 files changed, 336 insertions(+), 208 deletions(-) diff --git a/.env.example b/.env.example index 9ce21139..ad9895f7 100644 --- a/.env.example +++ b/.env.example @@ -185,6 +185,26 @@ MAX_CONCURRENT_REQUESTS_PER_KEY_IFLOW=1 # ROTATION_MODE_GEMINI=balanced # ROTATION_MODE_ANTIGRAVITY=sequential +# --- Model Quota Groups --- +# Models that share quota/cooldown timing. When one model in a group hits +# quota exhausted (429), all models in the group receive the same cooldown timestamp. +# They also reset (archive stats) together when the quota period expires. +# +# This is useful for providers where multiple model variants share the same +# underlying quota (e.g., Claude Sonnet and Opus on Antigravity). +# +# Format: QUOTA_GROUPS__="model1,model2,model3" +# +# To DISABLE a default group, set it to empty string: +# QUOTA_GROUPS_ANTIGRAVITY_CLAUDE="" +# +# Default groups: +# ANTIGRAVITY.CLAUDE: claude-sonnet-4-5,claude-opus-4-5 +# +# Examples: +# QUOTA_GROUPS_ANTIGRAVITY_CLAUDE="claude-sonnet-4-5,claude-opus-4-5" +# QUOTA_GROUPS_ANTIGRAVITY_GEMINI="gemini-3-pro-preview,gemini-3-pro-image-preview" + # ------------------------------------------------------------------------------ # | [ADVANCED] Proxy Configuration | # ------------------------------------------------------------------------------ diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 88e5a1d1..377e7d9d 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -34,7 +34,7 @@ import httpx import litellm -from .provider_interface import ProviderInterface +from .provider_interface import ProviderInterface, UsageResetConfigDef, QuotaGroupMap from .antigravity_auth_base import AntigravityAuthBase from .provider_cache import ProviderCache from ..model_definitions import ModelDefinitions @@ -497,6 +497,52 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface): # Sequential mode by default - preserves thinking signature caches between requests default_rotation_mode: str = "sequential" + # ========================================================================= + # TIER & USAGE CONFIGURATION + # ========================================================================= + + # Provider name for env var lookups (QUOTA_GROUPS_ANTIGRAVITY_*) + provider_env_name: str = "antigravity" + + # Tier name -> priority mapping (Single Source of Truth) + # Lower numbers = higher priority + tier_priorities = { + # Priority 1: Highest paid tier (Google AI Ultra - name unconfirmed) + # "google-ai-ultra": 1, # Uncomment when tier name is confirmed + # Priority 2: Standard paid tier + "standard-tier": 2, + # Priority 3: Free tier + "free-tier": 3, + # Priority 10: Legacy/Unknown (lowest) + "legacy-tier": 10, + "unknown": 10, + } + + # Default priority for tiers not in the mapping + default_tier_priority: int = 10 + + # Usage reset configs keyed by priority sets + # Priorities 1-2 (paid tiers) get 5h window, others get 7d window + usage_reset_configs = { + frozenset({1, 2}): UsageResetConfigDef( + window_seconds=5 * 60 * 60, # 5 hours + mode="per_model", + description="5-hour per-model window (paid tier)", + field_name="models", + ), + "default": UsageResetConfigDef( + window_seconds=7 * 24 * 60 * 60, # 7 days + mode="per_model", + description="7-day per-model window (free/unknown tier)", + field_name="models", + ), + } + + # Model quota groups (can be overridden via QUOTA_GROUPS_ANTIGRAVITY_CLAUDE) + model_quota_groups: QuotaGroupMap = { + # "claude": ["claude-sonnet-4-5", "claude-opus-4-5"], + } + @staticmethod def parse_quota_error( error: Exception, error_body: Optional[str] = None @@ -733,43 +779,6 @@ def _log_config(self) -> None: f"claude_fix={self._enable_claude_tool_fix}, thinking_sanitization={self._enable_thinking_sanitization}" ) - # ========================================================================= - # CREDENTIAL PRIORITIZATION - # ========================================================================= - - def get_credential_priority(self, credential: str) -> Optional[int]: - """ - Returns priority based on Antigravity tier. - Paid tiers: priority 1 (highest) - Free tier: priority 2 - Legacy/Unknown: priority 10 (lowest) - - Args: - credential: The credential path - - Returns: - Priority level (1-10) or None if tier not yet discovered - """ - tier = self.project_tier_cache.get(credential) - - # Lazy load from file if not in cache - if not tier: - tier = self._load_tier_from_file(credential) - - if not tier: - return None # Not yet discovered - - # Paid tiers get highest priority - if tier not in ["free-tier", "legacy-tier", "unknown"]: - return 1 - - # Free tier gets lower priority - if tier == "free-tier": - return 2 - - # Legacy and unknown get even lower - return 10 - def _load_tier_from_file(self, credential_path: str) -> Optional[str]: """ Load tier from credential file's _proxy_metadata and cache it. @@ -837,105 +846,6 @@ def get_model_tier_requirement(self, model: str) -> Optional[int]: """ return None - def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: - """ - Get Antigravity-specific usage tracking configuration based on credential tier. - - Antigravity uses per-model windows with different durations by tier: - - Paid tiers (priority 1): 5-hour per-model window - - Free tier (priority 2): 7-day per-model window - - Unknown/legacy: 7-day per-model window (conservative default) - - When a model hits a quota_exhausted 429 error with exact reset timestamp, - that timestamp becomes the authoritative reset time for the model (and its group). - - Args: - credential: The credential path - - Returns: - Usage reset configuration dict with mode="per_model" - """ - tier = self.project_tier_cache.get(credential) - if not tier: - tier = self._load_tier_from_file(credential) - - # Paid tiers: 5-hour per-model window - if tier and tier not in ["free-tier", "legacy-tier", "unknown"]: - return { - "window_seconds": 5 * 60 * 60, # 18000 seconds = 5 hours - "mode": "per_model", - "priority": 1, - "description": "5-hour per-model window (paid tier)", - } - - # Free tier: 7-day per-model window - if tier == "free-tier": - return { - "window_seconds": 7 * 24 * 60 * 60, # 604800 seconds = 7 days - "mode": "per_model", - "priority": 2, - "description": "7-day per-model window (free tier)", - } - - # Unknown/legacy: use 7-day per-model window as conservative default - return { - "window_seconds": 7 * 24 * 60 * 60, # 604800 seconds = 7 days - "mode": "per_model", - "priority": 10, - "description": "7-day per-model window (unknown tier - conservative default)", - } - - def get_default_usage_field_name(self) -> str: - """ - Get the default usage tracking field name for Antigravity. - - Returns: - "models" for per-model tracking - """ - return "models" - - # ========================================================================= - # Model Quota Grouping - # ========================================================================= - - # Models that share quota timing - when one hits quota, all get same reset time - QUOTA_GROUPS = { - # Future: add claude/gemini groups if they share quota - } - - def get_model_quota_group(self, model: str) -> Optional[str]: - """ - Returns the quota group name for a model. - - Claude models (sonnet and opus) share quota on Antigravity. - When one hits quota exhausted, all models in the group get the same reset time. - - Args: - model: Model name (with or without "antigravity/" prefix) - - Returns: - Group name ("claude") or None if not grouped - """ - # Remove provider prefix if present - clean_model = model.replace("antigravity/", "") - - for group_name, models in self.QUOTA_GROUPS.items(): - if clean_model in models: - return group_name - return None - - def get_models_in_quota_group(self, group: str) -> List[str]: - """ - Returns all model names in a quota group. - - Args: - group: Group name (e.g., "claude") - - Returns: - List of model names (without provider prefix) - """ - return self.QUOTA_GROUPS.get(group, []) - async def initialize_credentials(self, credential_paths: List[str]) -> None: """ Load persisted tier information from credential files at startup. diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 745f934d..9965e449 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -189,6 +189,36 @@ class GeminiCliProvider(GeminiAuthBase, ProviderInterface): # Balanced by default - Gemini CLI has short cooldowns (seconds, not hours) default_rotation_mode: str = "balanced" + # ========================================================================= + # TIER CONFIGURATION + # ========================================================================= + + # Provider name for env var lookups (QUOTA_GROUPS_GEMINI_CLI_*) + provider_env_name: str = "gemini_cli" + + # Tier name -> priority mapping (Single Source of Truth) + # Same tier names as Antigravity (coincidentally), but defined separately + tier_priorities = { + # Priority 1: Highest paid tier (Google AI Ultra - name unconfirmed) + # "google-ai-ultra": 1, # Uncomment when tier name is confirmed + # Priority 2: Standard paid tier + "standard-tier": 2, + # Priority 3: Free tier + "free-tier": 3, + # Priority 10: Legacy/Unknown (lowest) + "legacy-tier": 10, + "unknown": 10, + } + + # Default priority for tiers not in the mapping + default_tier_priority: int = 10 + + # Gemini CLI uses default daily reset - no custom usage_reset_configs + # (Empty dict means inherited get_usage_reset_config returns None) + + # No quota groups defined for Gemini CLI + # (Models don't share quotas) + @staticmethod def parse_quota_error( error: Exception, error_body: Optional[str] = None @@ -264,41 +294,13 @@ def __init__(self): ) # ========================================================================= - # CREDENTIAL PRIORITIZATION + # CREDENTIAL TIER LOOKUP (Provider-specific - uses cache) + # ========================================================================= + # + # NOTE: get_credential_priority() is now inherited from ProviderInterface. + # It uses get_credential_tier_name() to get the tier and resolve priority + # from the tier_priorities class attribute. # ========================================================================= - - def get_credential_priority(self, credential: str) -> Optional[int]: - """ - Returns priority based on Gemini tier. - Paid tiers: priority 1 (highest) - Free/Legacy tiers: priority 2 - Unknown: priority 10 (lowest) - - Args: - credential: The credential path - - Returns: - Priority level (1-10) or None if tier not yet discovered - """ - tier = self.project_tier_cache.get(credential) - - # Lazy load from file if not in cache - if not tier: - tier = self._load_tier_from_file(credential) - - if not tier: - return None # Not yet discovered - - # Paid tiers get highest priority - if tier not in ["free-tier", "legacy-tier", "unknown"]: - return 1 - - # Free tier gets lower priority - if tier == "free-tier": - return 2 - - # Legacy and unknown get even lower - return 10 def _load_tier_from_file(self, credential_path: str) -> Optional[str]: """ diff --git a/src/rotator_library/providers/provider_interface.py b/src/rotator_library/providers/provider_interface.py index 1cc8879e..4fde24ec 100644 --- a/src/rotator_library/providers/provider_interface.py +++ b/src/rotator_library/providers/provider_interface.py @@ -1,10 +1,46 @@ from abc import ABC, abstractmethod -from typing import List, Dict, Any, Optional, AsyncGenerator, Union +from dataclasses import dataclass +from typing import List, Dict, Any, Optional, AsyncGenerator, Union, FrozenSet import os import httpx import litellm +# ============================================================================= +# TIER & USAGE CONFIGURATION TYPES +# ============================================================================= + + +@dataclass(frozen=True) +class UsageResetConfigDef: + """ + Definition for usage reset configuration per tier type. + + Providers define these as class attributes to specify how usage stats + should reset based on credential tier (paid vs free). + + Attributes: + window_seconds: Duration of the usage tracking window in seconds. + mode: Either "credential" (one window per credential) or "per_model" + (separate window per model or model group). + description: Human-readable description for logging. + field_name: The key used in usage data JSON structure. + Typically "models" for per_model mode, "daily" for credential mode. + """ + + window_seconds: int + mode: str # "credential" or "per_model" + description: str + field_name: str = "daily" # Default for backwards compatibility + + +# Type aliases for provider configuration +TierPriorityMap = Dict[str, int] # tier_name -> priority +UsageConfigKey = Union[FrozenSet[int], str] # frozenset of priorities OR "default" +UsageConfigMap = Dict[UsageConfigKey, UsageResetConfigDef] # priority_set -> config +QuotaGroupMap = Dict[str, List[str]] # group_name -> [models] + + class ProviderInterface(ABC): """ An interface for API provider-specific functionality, including model @@ -18,6 +54,40 @@ class ProviderInterface(ABC): # - "sequential": Use one credential until exhausted, then switch to next default_rotation_mode: str = "balanced" + # ========================================================================= + # TIER CONFIGURATION - Override in subclass + # ========================================================================= + + # Provider name for env var lookups (e.g., "antigravity", "gemini_cli") + # Used for: QUOTA_GROUPS_{provider_env_name}_{GROUP} + provider_env_name: str = "" + + # Tier name -> priority mapping (Single Source of Truth) + # Lower numbers = higher priority (1 is highest) + # Multiple tiers can map to the same priority + # Unknown tiers fall back to default_tier_priority + tier_priorities: TierPriorityMap = {} + + # Default priority for tiers not in tier_priorities mapping + default_tier_priority: int = 10 + + # ========================================================================= + # USAGE RESET CONFIGURATION - Override in subclass + # ========================================================================= + + # Usage reset configurations keyed by priority sets + # Keys: frozenset of priority values (e.g., frozenset({1, 2})) OR "default" + # The "default" key is used for any priority not matched by a frozenset + usage_reset_configs: UsageConfigMap = {} + + # ========================================================================= + # MODEL QUOTA GROUPS - Override in subclass + # ========================================================================= + + # Models that share quota/cooldown timing + # Can be overridden via env: QUOTA_GROUPS_{PROVIDER}_{GROUP}="model1,model2" + model_quota_groups: QuotaGroupMap = {} + @abstractmethod async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: """ @@ -87,28 +157,50 @@ async def proactively_refresh(self, credential_path: str): pass # [NEW] Credential Prioritization System + + # ========================================================================= + # TIER RESOLUTION LOGIC (Centralized) + # ========================================================================= + + def _resolve_tier_priority(self, tier_name: Optional[str]) -> int: + """ + Resolve priority for a tier name using provider's tier_priorities mapping. + + Args: + tier_name: The tier name string (e.g., "free-tier", "standard-tier") + + Returns: + Priority level from tier_priorities, or default_tier_priority if + tier_name is None or not found in the mapping. + """ + if tier_name is None: + return self.default_tier_priority + return self.tier_priorities.get(tier_name, self.default_tier_priority) + def get_credential_priority(self, credential: str) -> Optional[int]: """ Returns the priority level for a credential. Lower numbers = higher priority (1 is highest). - Returns None if provider doesn't use priorities. + Returns None if tier not yet discovered. + + Uses the provider's tier_priorities mapping to resolve priority from + tier name. Unknown tiers fall back to default_tier_priority. - This allows providers to auto-detect credential tiers (e.g., paid vs free) - and ensure higher-tier credentials are always tried first. + Subclasses should: + 1. Define tier_priorities dict with all known tier names + 2. Override get_credential_tier_name() for tier lookup + Do NOT override this method. Args: credential: The credential identifier (API key or path) Returns: - Priority level (1-10) or None if no priority system - - Example: - For Gemini CLI: - - Paid tier credentials: priority 1 (highest) - - Free tier credentials: priority 2 - - Unknown tier: priority 10 (lowest) + Priority level (1-10) or None if tier not yet discovered """ - return None + tier = self.get_credential_tier_name(credential) + if tier is None: + return None # Tier not yet discovered + return self._resolve_tier_priority(tier) def get_model_tier_requirement(self, model: str) -> Optional[int]: """ @@ -211,12 +303,76 @@ def parse_quota_error( # Per-Provider Usage Tracking Configuration # ========================================================================= + # ========================================================================= + # USAGE RESET CONFIG LOGIC (Centralized) + # ========================================================================= + + def _find_usage_config_for_priority( + self, priority: int + ) -> Optional[UsageResetConfigDef]: + """ + Find usage config that applies to a priority value. + + Checks frozenset keys first (priority must be in the set), + then falls back to "default" key if no match found. + + Args: + priority: The credential priority level + + Returns: + UsageResetConfigDef if found, None otherwise + """ + # First, check frozenset keys for explicit priority match + for key, config in self.usage_reset_configs.items(): + if isinstance(key, frozenset) and priority in key: + return config + + # Fall back to "default" key + return self.usage_reset_configs.get("default") + + def _build_usage_reset_config( + self, tier_name: Optional[str] + ) -> Optional[Dict[str, Any]]: + """ + Build usage reset configuration dict for a tier. + + Resolves tier to priority, then finds matching usage config. + Returns None if provider doesn't define usage_reset_configs. + + Args: + tier_name: The tier name string + + Returns: + Usage config dict with window_seconds, mode, priority, description, + field_name, or None if no config applies + """ + if not self.usage_reset_configs: + return None + + priority = self._resolve_tier_priority(tier_name) + config = self._find_usage_config_for_priority(priority) + + if config is None: + return None + + return { + "window_seconds": config.window_seconds, + "mode": config.mode, + "priority": priority, + "description": config.description, + "field_name": config.field_name, + } + def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: """ Get provider-specific usage tracking configuration for a credential. - This allows providers to define custom usage reset windows based on - credential tier (e.g., paid vs free accounts with different quota periods). + Uses the provider's usage_reset_configs class attribute to build + the configuration dict. Priority is auto-derived from tier. + + Subclasses should define usage_reset_configs as a class attribute + instead of overriding this method. Only override get_credential_tier_name() + to provide the tier lookup mechanism. The UsageManager will use this configuration to: 1. Track usage per-model or per-credential based on mode @@ -231,7 +387,7 @@ def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: { "window_seconds": int, # Duration in seconds (e.g., 18000 for 5h) "mode": str, # "credential" or "per_model" - "priority": int, # Priority level this config applies to + "priority": int, # Priority level (auto-derived from tier) "description": str, # Human-readable description (for logging) } @@ -242,25 +398,9 @@ def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]: from first request of THAT model. Models reset independently unless grouped. If a quota_exhausted error provides exact reset time, that becomes the authoritative reset time for the model. - - Examples: - Antigravity paid tier (per-model): - { - "window_seconds": 18000, # 5 hours - "mode": "per_model", - "priority": 1, - "description": "5-hour per-model window (paid tier)" - } - - Default provider (credential-level): - { - "window_seconds": 86400, # 24 hours - "mode": "credential", - "priority": 1, - "description": "24-hour credential window" - } """ - return None # Default: use daily reset at daily_reset_time_utc + tier = self.get_credential_tier_name(credential) + return self._build_usage_reset_config(tier) def get_default_usage_field_name(self) -> str: """ @@ -278,16 +418,68 @@ def get_default_usage_field_name(self) -> str: # Model Quota Grouping # ========================================================================= + # ========================================================================= + # QUOTA GROUPS LOGIC (Centralized) + # ========================================================================= + + def _get_effective_quota_groups(self) -> QuotaGroupMap: + """ + Get quota groups with .env overrides applied. + + Env format: QUOTA_GROUPS_{PROVIDER}_{GROUP}="model1,model2" + Set empty string to disable a default group. + """ + if not self.provider_env_name or not self.model_quota_groups: + return self.model_quota_groups + + result: QuotaGroupMap = {} + + for group_name, default_models in self.model_quota_groups.items(): + env_key = ( + f"QUOTA_GROUPS_{self.provider_env_name.upper()}_{group_name.upper()}" + ) + env_value = os.getenv(env_key) + + if env_value is not None: + # Env override present + if env_value.strip(): + # Parse comma-separated models + result[group_name] = [ + m.strip() for m in env_value.split(",") if m.strip() + ] + # Empty string = group disabled, don't add to result + else: + # Use default + result[group_name] = list(default_models) + + return result + + def _find_model_quota_group(self, model: str) -> Optional[str]: + """Find which quota group a model belongs to.""" + groups = self._get_effective_quota_groups() + for group_name, models in groups.items(): + if model in models: + return group_name + return None + + def _get_quota_group_models(self, group: str) -> List[str]: + """Get all models in a quota group.""" + groups = self._get_effective_quota_groups() + return groups.get(group, []) + def get_model_quota_group(self, model: str) -> Optional[str]: """ Returns the quota group name for a model, or None if not grouped. + Uses the provider's model_quota_groups class attribute with .env overrides + via QUOTA_GROUPS_{PROVIDER}_{GROUP}="model1,model2". + Models in the same quota group share cooldown timing - when one model hits a quota exhausted error, all models in the group get the same reset timestamp. They also reset (archive stats) together. - This is useful for providers where multiple model variants share the - same underlying quota (e.g., Claude Sonnet and Opus on Antigravity). + Subclasses should define model_quota_groups as a class attribute + instead of overriding this method. Args: model: Model name (with or without provider prefix) @@ -295,12 +487,16 @@ def get_model_quota_group(self, model: str) -> Optional[str]: Returns: Group name string (e.g., "claude") or None if model is not grouped """ - return None + # Strip provider prefix if present + clean_model = model.split("/")[-1] if "/" in model else model + return self._find_model_quota_group(clean_model) def get_models_in_quota_group(self, group: str) -> List[str]: """ Returns all model names that belong to a quota group. + Uses the provider's model_quota_groups class attribute with .env overrides. + Args: group: Group name (e.g., "claude") @@ -308,4 +504,4 @@ def get_models_in_quota_group(self, group: str) -> List[str]: List of model names (WITHOUT provider prefix) in the group. Empty list if group doesn't exist. """ - return [] + return self._get_quota_group_models(group) From 5e42536dc5b67ed5e06a095ae06da5ae93b9c4d1 Mon Sep 17 00:00:00 2001 From: MasuRii Date: Mon, 8 Dec 2025 02:44:59 +0800 Subject: [PATCH 094/221] fix(resilience): complete circuit breaker patterns per PR review Address bot review feedback on PR #32: - Add _disk_available flag update in _write_json exception handler - Add _disk_available flag update in log_stream_chunk (critical for streams) - Document intentional no-memory-fallback design for streams - Add _fallback_mode update in failure_logger exception handler - Add complete circuit breaker pattern to usage_manager --- src/proxy_app/detailed_logger.py | 5 ++++- src/rotator_library/failure_logger.py | 2 ++ src/rotator_library/usage_manager.py | 11 +++++++++++ 3 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/proxy_app/detailed_logger.py b/src/proxy_app/detailed_logger.py index 107a05cf..0d0dd9a9 100644 --- a/src/proxy_app/detailed_logger.py +++ b/src/proxy_app/detailed_logger.py @@ -50,6 +50,7 @@ def _write_json(self, filename: str, data: Dict[str, Any]): with open(self.log_dir / filename, "w", encoding="utf-8") as f: json.dump(data, f, indent=4, ensure_ascii=False) except (OSError, PermissionError, IOError) as e: + DetailedLogger._disk_available = False logging.error(f"[{self.request_id}] Failed to write to {filename}: {e}") self._in_memory_logs.append({"file": filename, "data": data}) @@ -66,8 +67,9 @@ def log_request(self, headers: Dict[str, Any], body: Dict[str, Any]): def log_stream_chunk(self, chunk: Dict[str, Any]): """Logs an individual chunk from a streaming response to a JSON Lines file.""" + # Intentionally skip memory fallback for streams to prevent OOM - unlike _write_json, we don't buffer stream chunks in memory if not DetailedLogger._disk_available: - return # Skip chunk logging when disk unavailable + return try: self.log_dir.mkdir(parents=True, exist_ok=True) @@ -78,6 +80,7 @@ def log_stream_chunk(self, chunk: Dict[str, Any]): with open(self.log_dir / "streaming_chunks.jsonl", "a", encoding="utf-8") as f: f.write(json.dumps(log_entry, ensure_ascii=False) + "\n") except (OSError, PermissionError, IOError) as e: + DetailedLogger._disk_available = False logging.error(f"[{self.request_id}] Failed to write stream chunk: {e}") def log_final_response(self, status_code: int, headers: Optional[Dict[str, Any]], body: Dict[str, Any]): diff --git a/src/rotator_library/failure_logger.py b/src/rotator_library/failure_logger.py index 9379d34e..a3e07d33 100644 --- a/src/rotator_library/failure_logger.py +++ b/src/rotator_library/failure_logger.py @@ -194,6 +194,8 @@ def log_failure( try: failure_logger.error(detailed_log_data) except (OSError, IOError) as e: + global _fallback_mode + _fallback_mode = True # File logging failed - log to console instead logging.error(f"Failed to write to failures.log: {e}") logging.error(f"Failure summary: {summary_message}") diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 1defd7ae..d6398f32 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -72,6 +72,9 @@ def __init__( self._timeout_lock = asyncio.Lock() self._claimed_on_timeout: Set[str] = set() + + # Circuit breaker for disk write failures + self._disk_available = True if daily_reset_time_utc: hour, minute = map(int, daily_reset_time_utc.split(":")) @@ -113,6 +116,9 @@ async def _load_usage(self): except (OSError, PermissionError, IOError) as e: lib_logger.warning(f"Cannot read usage file {self.file_path}: {e}. Using empty state.") self._usage_data = {} + else: + # [CIRCUIT BREAKER RESET] Successfully loaded, re-enable disk writes + self._disk_available = True async def _save_usage(self): """Saves the current usage data to the JSON file asynchronously with resilience. @@ -123,6 +129,9 @@ async def _save_usage(self): """ if self._usage_data is None: return + + if not self._disk_available: + return # Skip disk write when unavailable try: async with self._data_lock: @@ -134,6 +143,8 @@ async def _save_usage(self): async with aiofiles.open(self.file_path, "w") as f: await f.write(json.dumps(self._usage_data, indent=2)) except (OSError, PermissionError, IOError) as e: + # [CIRCUIT BREAKER] Disable disk writes to prevent repeated failures + self._disk_available = False # [FAIL SILENTLY, LOG LOUDLY] Log the error but don't crash # In-memory state is preserved and will continue to work lib_logger.warning( From 67e70d91d41dbeb79694f768755f9d4573822944 Mon Sep 17 00:00:00 2001 From: MasuRii Date: Mon, 8 Dec 2025 03:09:57 +0800 Subject: [PATCH 095/221] fix(google-oauth): prevent credentials from becoming permanently stuck Fixed a bug where OAuth credentials would become permanently unavailable after token refresh due to improper cleanup of _unavailable_credentials. Changes: - Added cleanup to finally block (always executes) - Added cleanup before timeout exit path - Added cleanup to CancelledError handler - Changed _unavailable_credentials from set to Dict with 5-min TTL for automatic stale entry cleanup as defense in depth This resolves the 'No keys are eligible' loop that required restart. --- .../providers/google_oauth_base.py | 91 ++++++++++++++++--- 1 file changed, 80 insertions(+), 11 deletions(-) diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index 0b34153b..96684ef4 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -85,9 +85,12 @@ def __init__(self): # [QUEUE SYSTEM] Sequential refresh processing self._refresh_queue: asyncio.Queue = asyncio.Queue() self._queued_credentials: set = set() # Track credentials already in queue - self._unavailable_credentials: set = ( - set() - ) # Mark credentials unavailable during re-auth + # [FIX 4] Changed from set to dict mapping credential path to timestamp + # This enables TTL-based stale entry cleanup as defense in depth + self._unavailable_credentials: Dict[str, float] = ( + {} + ) # Maps credential path -> timestamp when marked unavailable + self._unavailable_ttl_seconds: int = 300 # 5 minutes TTL for stale entries self._queue_tracking_lock = asyncio.Lock() # Protects queue sets self._queue_processor_task: Optional[asyncio.Task] = ( None # Background worker task @@ -526,8 +529,33 @@ async def _get_lock(self, path: str) -> asyncio.Lock: return self._refresh_locks[path] def is_credential_available(self, path: str) -> bool: - """Check if a credential is available for rotation (not queued/refreshing).""" - return path not in self._unavailable_credentials + """Check if a credential is available for rotation (not queued/refreshing). + + [FIX 4] Now includes TTL-based stale entry cleanup as defense in depth. + If a credential has been unavailable for longer than _unavailable_ttl_seconds, + it is automatically cleaned up and considered available. + """ + if path not in self._unavailable_credentials: + return True + + # [FIX 4] Check if the entry is stale (TTL expired) + marked_time = self._unavailable_credentials.get(path) + if marked_time is not None: + now = time.time() + if now - marked_time > self._unavailable_ttl_seconds: + # Entry is stale - clean it up and return available + lib_logger.warning( + f"Credential '{Path(path).name}' was stuck in unavailable state for " + f"{int(now - marked_time)}s (TTL: {self._unavailable_ttl_seconds}s). " + f"Auto-cleaning stale entry." + ) + # Note: This is a sync method, so we can't use async lock here. + # However, discard from dict is thread-safe for single operations. + # The _queue_tracking_lock protects concurrent modifications in async context. + self._unavailable_credentials.pop(path, None) + return True + + return False async def _ensure_queue_processor_running(self): """Lazily starts the queue processor if not already running.""" @@ -563,7 +591,12 @@ async def _queue_refresh( async with self._queue_tracking_lock: if path not in self._queued_credentials: self._queued_credentials.add(path) - self._unavailable_credentials.add(path) # Mark as unavailable + # [FIX 4] Store timestamp when marking unavailable (for TTL cleanup) + self._unavailable_credentials[path] = time.time() + lib_logger.debug( + f"Marked '{Path(path).name}' as unavailable. " + f"Total unavailable: {len(self._unavailable_credentials)}" + ) await self._refresh_queue.put((path, force, needs_reauth)) await self._ensure_queue_processor_running() @@ -578,7 +611,16 @@ async def _process_refresh_queue(self): self._refresh_queue.get(), timeout=60.0 ) except asyncio.TimeoutError: - # No items for 60s, exit to save resources + # [FIX 2] Clean up any stale unavailable entries before exiting + # If we're idle for 60s, no refreshes are in progress + async with self._queue_tracking_lock: + if self._unavailable_credentials: + stale_count = len(self._unavailable_credentials) + lib_logger.warning( + f"Queue processor idle timeout. Cleaning {stale_count} " + f"stale unavailable credentials: {list(self._unavailable_credentials.keys())}" + ) + self._unavailable_credentials.clear() self._queue_processor_task = None return @@ -590,7 +632,11 @@ async def _process_refresh_queue(self): if creds and not self._is_token_expired(creds): # No longer expired, mark as available async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Credential '{Path(path).name}' no longer expired, marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) continue # Perform refresh @@ -600,21 +646,44 @@ async def _process_refresh_queue(self): # SUCCESS: Mark as available again async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) finally: - # Remove from queued set + # [FIX 1] Remove from BOTH queued set AND unavailable credentials + # This ensures cleanup happens in ALL exit paths (success, exception, etc.) async with self._queue_tracking_lock: self._queued_credentials.discard(path) + # [FIX 1] Always clean up unavailable credentials in finally block + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Finally cleanup for '{Path(path).name}'. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) self._refresh_queue.task_done() except asyncio.CancelledError: + # [FIX 3] Clean up the current credential before breaking + if path: + async with self._queue_tracking_lock: + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"CancelledError cleanup for '{Path(path).name}'. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) break except Exception as e: lib_logger.error(f"Error in queue processor: {e}") # Even on error, mark as available (backoff will prevent immediate retry) if path: async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Error cleanup for '{Path(path).name}': {e}. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) async def initialize_token( self, creds_or_path: Union[Dict[str, Any], str] From 4cdd2618be00ef0db8ba20d31495b9822f0b2e84 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 00:31:51 +0100 Subject: [PATCH 096/221] =?UTF-8?q?feat(usage):=20=E2=9C=A8=20add=20human-?= =?UTF-8?q?readable=20timestamp=20fields=20to=20usage=20data=20for=20debug?= =?UTF-8?q?ging?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces helper methods to automatically generate and persist human-readable timestamp fields alongside Unix timestamps in the usage tracking data. - Add `_format_timestamp_local()` method to convert Unix timestamps to local time strings with timezone offset - Add `_add_readable_timestamps()` method to enrich usage data with 'window_started' and 'quota_resets' fields - Integrate timestamp formatting into the save flow, automatically updating readable fields before persisting to disk - Set `quota_reset_ts` when initializing new model windows based on provider's window configuration The readable timestamps improve observability and debugging by making it easier to understand when quota windows started and when they will reset, without requiring manual timestamp conversion. --- src/rotator_library/usage_manager.py | 77 ++++++++++++++++++++++++++-- 1 file changed, 73 insertions(+), 4 deletions(-) diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index 39c8db6f..c05a31a9 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -297,6 +297,69 @@ def _get_usage_count(self, key: str, model: str) -> int: .get("success_count", 0) ) + # ========================================================================= + # TIMESTAMP FORMATTING HELPERS + # ========================================================================= + + def _format_timestamp_local(self, ts: Optional[float]) -> Optional[str]: + """ + Format Unix timestamp as local time string with timezone offset. + + Args: + ts: Unix timestamp or None + + Returns: + Formatted string like "2025-12-07 14:30:17 +0100" or None + """ + if ts is None: + return None + try: + dt = datetime.fromtimestamp(ts).astimezone() # Local timezone + # Use UTC offset for conciseness (works on all platforms) + return dt.strftime("%Y-%m-%d %H:%M:%S %z") + except (OSError, ValueError, OverflowError): + return None + + def _add_readable_timestamps(self, data: Dict) -> Dict: + """ + Add human-readable timestamp fields to usage data before saving. + + Adds 'window_started' and 'quota_resets' fields derived from + Unix timestamps for easier debugging and monitoring. + + Args: + data: The usage data dict to enhance + + Returns: + The same dict with readable timestamp fields added + """ + for key, key_data in data.items(): + # Handle per-model structure + models = key_data.get("models", {}) + for model_name, model_stats in models.items(): + if not isinstance(model_stats, dict): + continue + + # Add readable window start time + window_start = model_stats.get("window_start_ts") + if window_start: + model_stats["window_started"] = self._format_timestamp_local( + window_start + ) + elif "window_started" in model_stats: + del model_stats["window_started"] + + # Add readable reset time + quota_reset = model_stats.get("quota_reset_ts") + if quota_reset: + model_stats["quota_resets"] = self._format_timestamp_local( + quota_reset + ) + elif "quota_resets" in model_stats: + del model_stats["quota_resets"] + + return data + def _select_sequential( self, candidates: List[Tuple[str, int]], @@ -377,6 +440,8 @@ async def _save_usage(self): if self._usage_data is None: return async with self._data_lock: + # Add human-readable timestamp fields before saving + self._add_readable_timestamps(self._usage_data) async with aiofiles.open(self.file_path, "w") as f: await f.write(json.dumps(self._usage_data, indent=2)) @@ -1251,11 +1316,15 @@ async def record_success( # Start window on first request for this model if model_data.get("window_start_ts") is None: model_data["window_start_ts"] = now_ts - window_hours = ( - reset_config.get("window_seconds", 0) / 3600 - if reset_config - else 0 + + # Set expected quota reset time from provider config + window_seconds = ( + reset_config.get("window_seconds", 0) if reset_config else 0 ) + if window_seconds > 0: + model_data["quota_reset_ts"] = now_ts + window_seconds + + window_hours = window_seconds / 3600 if window_seconds else 0 lib_logger.info( f"Started {window_hours:.1f}h window for model {model} on {mask_credential(key)}" ) From 136eb6cf5fab3fa2876f6eb8a14ce9ae40d928f7 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 00:56:28 +0100 Subject: [PATCH 097/221] fix - addressing review findings. --- src/proxy_app/settings_tool.py | 4 ++-- src/rotator_library/providers/antigravity_provider.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/proxy_app/settings_tool.py b/src/proxy_app/settings_tool.py index 66b81e2e..7a07b07e 100644 --- a/src/proxy_app/settings_tool.py +++ b/src/proxy_app/settings_tool.py @@ -202,10 +202,10 @@ def get_default_mode(self, provider: str) -> str: # Import here to avoid circular imports try: from rotator_library.providers.provider_interface import ( - LLMProviderInterface, + ProviderInterface, ) - return LLMProviderInterface.get_rotation_mode(provider) + return ProviderInterface.get_rotation_mode(provider) except ImportError: # Fallback defaults if import fails if provider.lower() == "antigravity": diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 377e7d9d..ab3c92f7 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -540,7 +540,7 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface): # Model quota groups (can be overridden via QUOTA_GROUPS_ANTIGRAVITY_CLAUDE) model_quota_groups: QuotaGroupMap = { - # "claude": ["claude-sonnet-4-5", "claude-opus-4-5"], + # "claude": ["claude-sonnet-4-5", "claude-opus-4-5"], - commented out for later use if needed } @staticmethod From aefb70669f12137544b1ca5353996101e35f6a71 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 02:17:54 +0100 Subject: [PATCH 098/221] =?UTF-8?q?feat(concurrency):=20=E2=9C=A8=20add=20?= =?UTF-8?q?priority-based=20concurrency=20multipliers=20for=20credential?= =?UTF-8?q?=20tiers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces a flexible priority-based concurrency multiplier system that allows higher-priority credentials (e.g., paid tiers) to handle more concurrent requests than lower-priority credentials, regardless of rotation mode. Key changes: - Added `default_priority_multipliers` and `default_sequential_fallback_multiplier` to `ProviderInterface` for provider-level configuration - Implemented multiplier lookup with mode-specific overrides via environment variables (format: `CONCURRENCY_MULTIPLIER__PRIORITY_[_]=`) - Modified `UsageManager` to calculate effective concurrency limits by applying multipliers to base `MAX_CONCURRENT_REQUESTS_PER_KEY` values - Added `PriorityMultiplierManager` to `settings_tool.py` for runtime configuration and display of multipliers - Configured default multipliers for Antigravity (P1: 5x, P2: 3x, sequential fallback: 2x) and Gemini CLI (P1: 5x, P2: 3x) - Introduced `model_usage_weights` to account for models with different quota consumption rates (e.g., Opus counts 2x vs Sonnet) - Implemented `_get_grouped_usage_count()` for weighted usage calculation across quota groups - Refactored `_sort_sequential()` to return sorted lists instead of single selection, allowing multipliers to enable multiple concurrent requests in sequential mode - Enhanced logging to display effective concurrency limits and priority tiers during credential acquisition - Added comprehensive documentation in `.env.example` explaining the multiplier system and configuration options The multiplier system preserves existing rotation behavior while allowing paid credentials to maximize throughput. In sequential mode, multipliers enable controlled concurrency while maintaining cache-preserving stickiness. In balanced mode, multipliers provide fair load distribution with tier-appropriate capacity. --- .env.example | 31 +++ src/proxy_app/settings_tool.py | 256 +++++++++++++++++- src/rotator_library/client.py | 84 ++++++ .../providers/antigravity_provider.py | 21 +- .../providers/gemini_cli_provider.py | 10 + .../providers/provider_interface.py | 41 +++ src/rotator_library/usage_manager.py | 246 +++++++++++++---- 7 files changed, 628 insertions(+), 61 deletions(-) diff --git a/.env.example b/.env.example index ad9895f7..c5bce0bb 100644 --- a/.env.example +++ b/.env.example @@ -185,6 +185,37 @@ MAX_CONCURRENT_REQUESTS_PER_KEY_IFLOW=1 # ROTATION_MODE_GEMINI=balanced # ROTATION_MODE_ANTIGRAVITY=sequential +# --- Priority-Based Concurrency Multipliers --- +# Credentials can be assigned to priority tiers (1=highest, 2, 3, etc.). +# Each tier can have a concurrency multiplier that increases the effective +# concurrent request limit for credentials in that tier. +# +# How it works: +# effective_concurrent_limit = MAX_CONCURRENT_REQUESTS_PER_KEY * tier_multiplier +# +# This allows paid/premium credentials to handle more concurrent requests than +# free tier credentials, regardless of rotation mode. +# +# Provider Defaults (built into provider classes): +# Antigravity: +# Priority 1: 5x (paid ultra tier) +# Priority 2: 3x (standard paid tier) +# Priority 3+: 2x (sequential mode) or 1x (balanced mode) +# Gemini CLI: +# Priority 1: 5x +# Priority 2: 3x +# Others: 1x (all modes) +# +# Format: CONCURRENCY_MULTIPLIER__PRIORITY_= +# +# Mode-specific overrides (optional): +# Format: CONCURRENCY_MULTIPLIER__PRIORITY__= +# +# Examples: +# CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_1=10 # Override P1 to 10x +# CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_3=1 # Override P3 to 1x +# CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_2_BALANCED=1 # P2 = 1x in balanced mode only + # --- Model Quota Groups --- # Models that share quota/cooldown timing. When one model in a group hits # quota exhausted (429), all models in the group receive the same cooldown timestamp. diff --git a/src/proxy_app/settings_tool.py b/src/proxy_app/settings_tool.py index 7a07b07e..fe51cdf0 100644 --- a/src/proxy_app/settings_tool.py +++ b/src/proxy_app/settings_tool.py @@ -234,6 +234,94 @@ def remove_mode(self, provider: str): self.settings.remove(key) +class PriorityMultiplierManager: + """Manages CONCURRENCY_MULTIPLIER__PRIORITY_ settings""" + + def __init__(self, settings: AdvancedSettings): + self.settings = settings + + def get_provider_defaults(self, provider: str) -> Dict[int, int]: + """Get default priority multipliers from provider class""" + try: + from rotator_library.providers import PROVIDER_PLUGINS + + provider_class = PROVIDER_PLUGINS.get(provider.lower()) + if provider_class and hasattr( + provider_class, "default_priority_multipliers" + ): + return dict(provider_class.default_priority_multipliers) + except ImportError: + pass + return {} + + def get_sequential_fallback(self, provider: str) -> int: + """Get sequential fallback multiplier from provider class""" + try: + from rotator_library.providers import PROVIDER_PLUGINS + + provider_class = PROVIDER_PLUGINS.get(provider.lower()) + if provider_class and hasattr( + provider_class, "default_sequential_fallback_multiplier" + ): + return provider_class.default_sequential_fallback_multiplier + except ImportError: + pass + return 1 + + def get_current_multipliers(self) -> Dict[str, Dict[int, int]]: + """Get currently configured priority multipliers from env vars""" + multipliers: Dict[str, Dict[int, int]] = {} + for key, value in os.environ.items(): + if key.startswith("CONCURRENCY_MULTIPLIER_") and "_PRIORITY_" in key: + try: + # Parse: CONCURRENCY_MULTIPLIER__PRIORITY_ + parts = key.split("_PRIORITY_") + provider = parts[0].replace("CONCURRENCY_MULTIPLIER_", "").lower() + remainder = parts[1] + + # Check if mode-specific (has _SEQUENTIAL or _BALANCED suffix) + if "_" in remainder: + continue # Skip mode-specific for now (show in separate view) + + priority = int(remainder) + multiplier = int(value) + + if provider not in multipliers: + multipliers[provider] = {} + multipliers[provider][priority] = multiplier + except (ValueError, IndexError): + pass + return multipliers + + def get_effective_multiplier(self, provider: str, priority: int) -> int: + """Get effective multiplier (configured, provider default, or 1)""" + # Check env var override + current = self.get_current_multipliers() + if provider.lower() in current: + if priority in current[provider.lower()]: + return current[provider.lower()][priority] + + # Check provider defaults + defaults = self.get_provider_defaults(provider) + if priority in defaults: + return defaults[priority] + + # Return 1 (no multiplier) + return 1 + + def set_multiplier(self, provider: str, priority: int, multiplier: int): + """Set priority multiplier for a provider""" + if multiplier < 1: + raise ValueError("Multiplier must be >= 1") + key = f"CONCURRENCY_MULTIPLIER_{provider.upper()}_PRIORITY_{priority}" + self.settings.set(key, str(multiplier)) + + def remove_multiplier(self, provider: str, priority: int): + """Remove multiplier (reset to provider default)""" + key = f"CONCURRENCY_MULTIPLIER_{provider.upper()}_PRIORITY_{priority}" + self.settings.remove(key) + + # ============================================================================= # PROVIDER-SPECIFIC SETTINGS DEFINITIONS # ============================================================================= @@ -424,6 +512,7 @@ def __init__(self): self.model_mgr = ModelDefinitionManager(self.settings) self.concurrency_mgr = ConcurrencyManager(self.settings) self.rotation_mgr = RotationModeManager(self.settings) + self.priority_multiplier_mgr = PriorityMultiplierManager(self.settings) self.provider_settings_mgr = ProviderSettingsManager(self.settings) self.running = True @@ -1268,14 +1357,15 @@ def manage_rotation_modes(self): self.console.print() self.console.print(" 1. ➕ Set Rotation Mode for Provider") self.console.print(" 2. 🗑️ Reset to Provider Default") - self.console.print(" 3. ↩️ Back to Settings Menu") + self.console.print(" 3. ⚡ Configure Priority Concurrency Multipliers") + self.console.print(" 4. ↩️ Back to Settings Menu") self.console.print() self.console.print("━" * 70) self.console.print() choice = Prompt.ask( - "Select option", choices=["1", "2", "3"], show_choices=False + "Select option", choices=["1", "2", "3", "4"], show_choices=False ) if choice == "1": @@ -1368,8 +1458,170 @@ def manage_rotation_modes(self): input("\nPress Enter to continue...") elif choice == "3": + self.manage_priority_multipliers() + + elif choice == "4": break + def manage_priority_multipliers(self): + """Manage priority-based concurrency multipliers per provider""" + clear_screen() + + current_multipliers = self.priority_multiplier_mgr.get_current_multipliers() + available_providers = self.get_available_providers() + + self.console.print( + Panel.fit( + "[bold cyan]⚡ Priority Concurrency Multipliers[/bold cyan]", + border_style="cyan", + ) + ) + + self.console.print() + self.console.print("[bold]📋 Current Priority Multiplier Settings[/bold]") + self.console.print("━" * 70) + + # Show all providers with their priority multipliers + has_settings = False + for provider in available_providers: + defaults = self.priority_multiplier_mgr.get_provider_defaults(provider) + overrides = current_multipliers.get(provider, {}) + seq_fallback = self.priority_multiplier_mgr.get_sequential_fallback( + provider + ) + rotation_mode = self.rotation_mgr.get_effective_mode(provider) + + if defaults or overrides or seq_fallback != 1: + has_settings = True + self.console.print( + f"\n [bold]{provider}[/bold] ({rotation_mode} mode)" + ) + + # Combine and display priorities + all_priorities = set(defaults.keys()) | set(overrides.keys()) + for priority in sorted(all_priorities): + default_val = defaults.get(priority, 1) + override_val = overrides.get(priority) + + if override_val is not None: + self.console.print( + f" Priority {priority}: [cyan]{override_val}x[/cyan] (override, default: {default_val}x)" + ) + else: + self.console.print( + f" Priority {priority}: {default_val}x [dim](default)[/dim]" + ) + + # Show sequential fallback if applicable + if rotation_mode == "sequential" and seq_fallback != 1: + self.console.print( + f" Others (seq): {seq_fallback}x [dim](fallback)[/dim]" + ) + + if not has_settings: + self.console.print(" [dim]No priority multipliers configured[/dim]") + + self.console.print() + self.console.print("[bold]ℹ️ About Priority Multipliers:[/bold]") + self.console.print( + " Higher priority tiers (lower numbers) can have higher multipliers." + ) + self.console.print(" Example: Priority 1 = 5x, Priority 2 = 3x, Others = 1x") + self.console.print() + self.console.print("━" * 70) + self.console.print() + self.console.print(" 1. ✏️ Set Priority Multiplier") + self.console.print(" 2. 🔄 Reset to Provider Default") + self.console.print(" 3. ↩️ Back") + + choice = Prompt.ask( + "Select option", choices=["1", "2", "3"], show_choices=False + ) + + if choice == "1": + if not available_providers: + self.console.print("\n[yellow]No providers available[/yellow]") + input("\nPress Enter to continue...") + return + + # Select provider + self.console.print("\n[bold]Select provider:[/bold]") + for idx, prov in enumerate(available_providers, 1): + self.console.print(f" {idx}. {prov}") + + prov_idx = IntPrompt.ask( + "Provider", + choices=[str(i) for i in range(1, len(available_providers) + 1)], + ) + provider = available_providers[prov_idx - 1] + + # Get priority level + priority = IntPrompt.ask("Priority level (e.g., 1, 2, 3)") + + # Get current value + current = self.priority_multiplier_mgr.get_effective_multiplier( + provider, priority + ) + self.console.print( + f"\nCurrent multiplier for priority {priority}: {current}x" + ) + + multiplier = IntPrompt.ask("New multiplier (1-10)", default=current) + if 1 <= multiplier <= 10: + self.priority_multiplier_mgr.set_multiplier( + provider, priority, multiplier + ) + self.console.print( + f"\n[green]✅ Priority {priority} multiplier for '{provider}' set to {multiplier}x[/green]" + ) + else: + self.console.print( + "\n[yellow]Multiplier must be between 1 and 10[/yellow]" + ) + input("\nPress Enter to continue...") + + elif choice == "2": + # Find providers with overrides + providers_with_overrides = [ + p for p in available_providers if p in current_multipliers + ] + if not providers_with_overrides: + self.console.print("\n[yellow]No custom multipliers to reset[/yellow]") + input("\nPress Enter to continue...") + return + + self.console.print("\n[bold]Select provider to reset:[/bold]") + for idx, prov in enumerate(providers_with_overrides, 1): + self.console.print(f" {idx}. {prov}") + + prov_idx = IntPrompt.ask( + "Provider", + choices=[str(i) for i in range(1, len(providers_with_overrides) + 1)], + ) + provider = providers_with_overrides[prov_idx - 1] + + # Get priority to reset + overrides = current_multipliers.get(provider, {}) + if len(overrides) == 1: + priority = list(overrides.keys())[0] + else: + self.console.print(f"\nOverrides for {provider}: {overrides}") + priority = IntPrompt.ask("Priority level to reset") + + if priority in overrides: + self.priority_multiplier_mgr.remove_multiplier(provider, priority) + default = self.priority_multiplier_mgr.get_effective_multiplier( + provider, priority + ) + self.console.print( + f"\n[green]✅ Reset priority {priority} for '{provider}' to default ({default}x)[/green]" + ) + else: + self.console.print( + f"\n[yellow]No override for priority {priority}[/yellow]" + ) + input("\nPress Enter to continue...") + def manage_concurrency_limits(self): """Manage concurrency limits""" while True: diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 4ca9d8cf..6a3b8907 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -161,11 +161,95 @@ def __init__( if mode != "balanced": lib_logger.info(f"Provider '{provider}' using rotation mode: {mode}") + # Build priority-based concurrency multiplier maps + # These are universal multipliers based on credential tier/priority + priority_multipliers: Dict[str, Dict[int, int]] = {} + priority_multipliers_by_mode: Dict[str, Dict[str, Dict[int, int]]] = {} + sequential_fallback_multipliers: Dict[str, int] = {} + + for provider in self.all_credentials.keys(): + provider_class = self._provider_plugins.get(provider) + + # Start with provider class defaults + if provider_class: + # Get default priority multipliers from provider class + if hasattr(provider_class, "default_priority_multipliers"): + default_multipliers = provider_class.default_priority_multipliers + if default_multipliers: + priority_multipliers[provider] = dict(default_multipliers) + + # Get sequential fallback from provider class + if hasattr(provider_class, "default_sequential_fallback_multiplier"): + fallback = provider_class.default_sequential_fallback_multiplier + if fallback != 1: # Only store if different from global default + sequential_fallback_multipliers[provider] = fallback + + # Override with environment variables + # Format: CONCURRENCY_MULTIPLIER__PRIORITY_= + # Format: CONCURRENCY_MULTIPLIER__PRIORITY__= + for key, value in os.environ.items(): + prefix = f"CONCURRENCY_MULTIPLIER_{provider.upper()}_PRIORITY_" + if key.startswith(prefix): + remainder = key[len(prefix) :] + try: + multiplier = int(value) + if multiplier < 1: + lib_logger.warning(f"Invalid {key}: {value}. Must be >= 1.") + continue + + # Check if mode-specific (e.g., _PRIORITY_1_SEQUENTIAL) + if "_" in remainder: + parts = remainder.rsplit("_", 1) + priority = int(parts[0]) + mode = parts[1].lower() + if mode in ("sequential", "balanced"): + # Mode-specific override + if provider not in priority_multipliers_by_mode: + priority_multipliers_by_mode[provider] = {} + if mode not in priority_multipliers_by_mode[provider]: + priority_multipliers_by_mode[provider][mode] = {} + priority_multipliers_by_mode[provider][mode][ + priority + ] = multiplier + lib_logger.info( + f"Provider '{provider}' priority {priority} ({mode} mode) multiplier: {multiplier}x" + ) + else: + # Assume it's part of the priority number (unlikely but handle gracefully) + lib_logger.warning(f"Unknown mode in {key}: {mode}") + else: + # Universal priority multiplier + priority = int(remainder) + if provider not in priority_multipliers: + priority_multipliers[provider] = {} + priority_multipliers[provider][priority] = multiplier + lib_logger.info( + f"Provider '{provider}' priority {priority} multiplier: {multiplier}x" + ) + except ValueError: + lib_logger.warning( + f"Invalid {key}: {value}. Could not parse priority or multiplier." + ) + + # Log configured multipliers + for provider, multipliers in priority_multipliers.items(): + if multipliers: + lib_logger.info( + f"Provider '{provider}' priority multipliers: {multipliers}" + ) + for provider, fallback in sequential_fallback_multipliers.items(): + lib_logger.info( + f"Provider '{provider}' sequential fallback multiplier: {fallback}x" + ) + self.usage_manager = UsageManager( file_path=usage_file_path, rotation_tolerance=rotation_tolerance, provider_rotation_modes=provider_rotation_modes, provider_plugins=PROVIDER_PLUGINS, + priority_multipliers=priority_multipliers, + priority_multipliers_by_mode=priority_multipliers_by_mode, + sequential_fallback_multipliers=sequential_fallback_multipliers, ) self._model_list_cache = {} self.http_client = httpx.AsyncClient() diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index ab3c92f7..a29a63ab 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -539,10 +539,29 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface): } # Model quota groups (can be overridden via QUOTA_GROUPS_ANTIGRAVITY_CLAUDE) + # Models in the same group share quota - when one is exhausted, all are model_quota_groups: QuotaGroupMap = { - # "claude": ["claude-sonnet-4-5", "claude-opus-4-5"], - commented out for later use if needed + #"claude": ["claude-sonnet-4-5", "claude-opus-4-5"], - commented out for later use if needed } + # Model usage weights for grouped usage calculation + # Opus consumes more quota per request, so its usage counts 2x when + # comparing credentials for selection + model_usage_weights = { + "claude-opus-4-5": 2, + } + + # Priority-based concurrency multipliers + # Higher priority credentials (lower number) get higher multipliers + # Priority 1 (paid ultra): 5x concurrent requests + # Priority 2 (standard paid): 3x concurrent requests + # Others: Use sequential fallback (2x) or balanced default (1x) + default_priority_multipliers = {1: 5, 2: 3} + + # For sequential mode, lower priority tiers still get 2x to maintain stickiness + # For balanced mode, this doesn't apply (falls back to 1x) + default_sequential_fallback_multiplier = 2 + @staticmethod def parse_quota_error( error: Exception, error_body: Optional[str] = None diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 9965e449..52f15d68 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -219,6 +219,16 @@ class GeminiCliProvider(GeminiAuthBase, ProviderInterface): # No quota groups defined for Gemini CLI # (Models don't share quotas) + # Priority-based concurrency multipliers + # Same structure as Antigravity (by coincidence, tiers share naming) + # Priority 1 (paid ultra): 5x concurrent requests + # Priority 2 (standard paid): 3x concurrent requests + # Others: 1x (no sequential fallback, uses global default) + default_priority_multipliers = {1: 5, 2: 3} + + # No sequential fallback for Gemini CLI (uses balanced mode default) + # default_sequential_fallback_multiplier = 1 (inherited from ProviderInterface) + @staticmethod def parse_quota_error( error: Exception, error_body: Optional[str] = None diff --git a/src/rotator_library/providers/provider_interface.py b/src/rotator_library/providers/provider_interface.py index 4fde24ec..08c1e228 100644 --- a/src/rotator_library/providers/provider_interface.py +++ b/src/rotator_library/providers/provider_interface.py @@ -88,6 +88,30 @@ class ProviderInterface(ABC): # Can be overridden via env: QUOTA_GROUPS_{PROVIDER}_{GROUP}="model1,model2" model_quota_groups: QuotaGroupMap = {} + # Model usage weights for grouped usage calculation + # When calculating combined usage for quota groups, each model's usage + # is multiplied by its weight. This accounts for models that consume + # more quota per request (e.g., Opus uses more than Sonnet). + # Models not in the map default to weight 1. + # Example: {"claude-opus-4-5": 2} means Opus usage counts 2x + model_usage_weights: Dict[str, int] = {} + + # ========================================================================= + # PRIORITY CONCURRENCY MULTIPLIERS - Override in subclass + # ========================================================================= + + # Priority-based concurrency multipliers (universal, applies to all modes) + # Maps priority level -> multiplier + # Higher priority credentials (lower number) can have higher multipliers + # to allow more concurrent requests + # Example: {1: 5, 2: 3} means Priority 1 gets 5x, Priority 2 gets 3x + default_priority_multipliers: Dict[int, int] = {} + + # Fallback multiplier for sequential mode when priority not in default_priority_multipliers + # This is used for lower-priority tiers in sequential mode to maintain some stickiness + # Default: 1 (no multiplier effect) + default_sequential_fallback_multiplier: int = 1 + @abstractmethod async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]: """ @@ -505,3 +529,20 @@ def get_models_in_quota_group(self, group: str) -> List[str]: Empty list if group doesn't exist. """ return self._get_quota_group_models(group) + + def get_model_usage_weight(self, model: str) -> int: + """ + Returns the usage weight for a model when calculating grouped usage. + + Models with higher weights contribute more to the combined group usage. + This accounts for models that consume more quota per request. + + Args: + model: Model name (with or without provider prefix) + + Returns: + Weight multiplier (default 1 if not configured) + """ + # Strip provider prefix if present + clean_model = model.split("/")[-1] if "/" in model else model + return self.model_usage_weights.get(clean_model, 1) diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index c05a31a9..4cee8f14 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -55,6 +55,11 @@ def __init__( rotation_tolerance: float = 0.0, provider_rotation_modes: Optional[Dict[str, str]] = None, provider_plugins: Optional[Dict[str, Any]] = None, + priority_multipliers: Optional[Dict[str, Dict[int, int]]] = None, + priority_multipliers_by_mode: Optional[ + Dict[str, Dict[str, Dict[int, int]]] + ] = None, + sequential_fallback_multipliers: Optional[Dict[str, int]] = None, ): """ Initialize the UsageManager. @@ -71,11 +76,22 @@ def __init__( - "sequential": Use one credential until exhausted (preserves caching) provider_plugins: Dict mapping provider names to provider plugin instances. Used for per-provider usage reset configuration (window durations, field names). + priority_multipliers: Dict mapping provider -> priority -> multiplier. + Universal multipliers that apply regardless of rotation mode. + Example: {"antigravity": {1: 5, 2: 3}} + priority_multipliers_by_mode: Dict mapping provider -> mode -> priority -> multiplier. + Mode-specific overrides. Example: {"antigravity": {"balanced": {3: 1}}} + sequential_fallback_multipliers: Dict mapping provider -> fallback multiplier. + Used in sequential mode when priority not in priority_multipliers. + Example: {"antigravity": 2} """ self.file_path = file_path self.rotation_tolerance = rotation_tolerance self.provider_rotation_modes = provider_rotation_modes or {} self.provider_plugins = provider_plugins or PROVIDER_PLUGINS + self.priority_multipliers = priority_multipliers or {} + self.priority_multipliers_by_mode = priority_multipliers_by_mode or {} + self.sequential_fallback_multipliers = sequential_fallback_multipliers or {} self._provider_instances: Dict[str, Any] = {} # Cache for provider instances self.key_states: Dict[str, Dict[str, Any]] = {} @@ -107,6 +123,48 @@ def _get_rotation_mode(self, provider: str) -> str: """ return self.provider_rotation_modes.get(provider, "balanced") + def _get_priority_multiplier( + self, provider: str, priority: int, rotation_mode: str + ) -> int: + """ + Get the concurrency multiplier for a provider/priority/mode combination. + + Lookup order: + 1. Mode-specific tier override: priority_multipliers_by_mode[provider][mode][priority] + 2. Universal tier multiplier: priority_multipliers[provider][priority] + 3. Sequential fallback (if mode is sequential): sequential_fallback_multipliers[provider] + 4. Global default: 1 (no multiplier effect) + + Args: + provider: Provider name (e.g., "antigravity") + priority: Priority level (1 = highest priority) + rotation_mode: Current rotation mode ("sequential" or "balanced") + + Returns: + Multiplier value + """ + provider_lower = provider.lower() + + # 1. Check mode-specific override + if provider_lower in self.priority_multipliers_by_mode: + mode_multipliers = self.priority_multipliers_by_mode[provider_lower] + if rotation_mode in mode_multipliers: + if priority in mode_multipliers[rotation_mode]: + return mode_multipliers[rotation_mode][priority] + + # 2. Check universal tier multiplier + if provider_lower in self.priority_multipliers: + if priority in self.priority_multipliers[provider_lower]: + return self.priority_multipliers[provider_lower][priority] + + # 3. Sequential fallback (only for sequential mode) + if rotation_mode == "sequential": + if provider_lower in self.sequential_fallback_multipliers: + return self.sequential_fallback_multipliers[provider_lower] + + # 4. Global default + return 1 + def _get_provider_from_credential(self, credential: str) -> Optional[str]: """ Extract provider name from credential path or identifier. @@ -238,6 +296,60 @@ def _get_grouped_models(self, credential: str, group: str) -> List[str]: return [] + def _get_model_usage_weight(self, credential: str, model: str) -> int: + """ + Get the usage weight for a model when calculating grouped usage. + + Args: + credential: The credential identifier + model: Model name (with or without provider prefix) + + Returns: + Weight multiplier (default 1 if not configured) + """ + provider = self._get_provider_from_credential(credential) + plugin_instance = self._get_provider_instance(provider) + + if plugin_instance and hasattr(plugin_instance, "get_model_usage_weight"): + return plugin_instance.get_model_usage_weight(model) + + return 1 + + def _get_grouped_usage_count(self, key: str, model: str) -> int: + """ + Get usage count for credential selection, considering quota groups. + + If the model belongs to a quota group, returns the weighted combined usage + across all models in the group. Otherwise returns individual model usage. + + Weights are applied per-model to account for models that consume more quota + per request (e.g., Opus might count 2x compared to Sonnet). + + Args: + key: Credential identifier + model: Model name (with provider prefix, e.g., "antigravity/claude-sonnet-4-5") + + Returns: + Weighted combined usage if grouped, otherwise individual model usage + """ + # Check if model is in a quota group + group = self._get_model_quota_group(key, model) + + if group: + # Get all models in the group + grouped_models = self._get_grouped_models(key, group) + + # Sum weighted usage across all models in the group + total_weighted_usage = 0 + for grouped_model in grouped_models: + usage = self._get_usage_count(key, grouped_model) + weight = self._get_model_usage_weight(key, grouped_model) + total_weighted_usage += usage * weight + return total_weighted_usage + + # Not grouped - return individual model usage (no weight applied) + return self._get_usage_count(key, model) + def _get_usage_field_name(self, credential: str) -> str: """ Get the usage tracking field name for a credential. @@ -360,59 +472,64 @@ def _add_readable_timestamps(self, data: Dict) -> Dict: return data - def _select_sequential( + def _sort_sequential( self, candidates: List[Tuple[str, int]], credential_priorities: Optional[Dict[str, int]] = None, - ) -> str: + ) -> List[Tuple[str, int]]: """ - Select credential in strict sequential order for cache-preserving rotation. + Sort credentials for sequential mode with position retention. - This method ensures the same credential is reused until it hits a cooldown, - which preserves provider-side caching (e.g., thinking signature caches). + Credentials maintain their position based on established usage patterns, + ensuring that actively-used credentials remain primary until exhausted. - Selection logic: - 1. Sort by priority (lowest number = highest priority) - 2. Within same priority, sort by last_used_ts (most recent first = sticky) - 3. Return the first candidate + Sorting order (within each sort key, lower value = higher priority): + 1. Priority tier (lower number = higher priority) + 2. Usage count (higher = more established in rotation, maintains position) + 3. Last used timestamp (higher = more recent, tiebreaker for stickiness) + 4. Credential ID (alphabetical, stable ordering) Args: candidates: List of (credential_id, usage_count) tuples credential_priorities: Optional dict mapping credentials to priority levels Returns: - Selected credential ID + Sorted list of candidates (same format as input) """ if not candidates: - raise ValueError("Cannot select from empty candidate list") + return [] if len(candidates) == 1: - return candidates[0][0] + return candidates - def sort_key(item: Tuple[str, int]) -> Tuple[int, float]: - cred, _ = item - # Priority: lower is better (1 = highest priority) + def sort_key(item: Tuple[str, int]) -> Tuple[int, int, float, str]: + cred, usage_count = item priority = ( credential_priorities.get(cred, 999) if credential_priorities else 999 ) - # Last used: higher (more recent) is better for stickiness last_used = ( self._usage_data.get(cred, {}).get("last_used_ts", 0) if self._usage_data else 0 ) - # Negative last_used so most recent sorts first - return (priority, -last_used) + return ( + priority, # ASC: lower priority number = higher priority + -usage_count, # DESC: higher usage = more established + -last_used, # DESC: more recent = preferred for ties + cred, # ASC: stable alphabetical ordering + ) sorted_candidates = sorted(candidates, key=sort_key) - selected = sorted_candidates[0][0] - lib_logger.debug( - f"Sequential selection: chose {mask_credential(selected)} " - f"(priority={credential_priorities.get(selected, 999) if credential_priorities else 'N/A'})" - ) + # Debug logging - show top 3 credentials in ordering + if lib_logger.isEnabledFor(logging.DEBUG): + order_info = [ + f"{mask_credential(c)}(p={credential_priorities.get(c, 999) if credential_priorities else 'N/A'}, u={u})" + for c, u in sorted_candidates[:3] + ] + lib_logger.debug(f"Sequential ordering: {' → '.join(order_info)}") - return selected + return sorted_candidates async def _lazy_init(self): """Initializes the usage data by loading it from the file asynchronously.""" @@ -966,7 +1083,8 @@ async def acquire_key( priority = credential_priorities.get(key, 999) # Get usage count for load balancing within priority groups - usage_count = self._get_usage_count(key, model) + # Uses grouped usage if model is in a quota group + usage_count = self._get_grouped_usage_count(key, model) # Group by priority if priority not in priority_groups: @@ -979,6 +1097,16 @@ async def acquire_key( for priority_level in sorted_priorities: keys_in_priority = priority_groups[priority_level] + # Determine selection method based on provider's rotation mode + provider = model.split("/")[0] if "/" in model else "" + rotation_mode = self._get_rotation_mode(provider) + + # Calculate effective concurrency based on priority tier + multiplier = self._get_priority_multiplier( + provider, priority_level, rotation_mode + ) + effective_max_concurrent = max_concurrent * multiplier + # Within each priority group, use existing tier1/tier2 logic tier1_keys, tier2_keys = [], [] for key, usage_count in keys_in_priority: @@ -988,30 +1116,24 @@ async def acquire_key( if not key_state["models_in_use"]: tier1_keys.append((key, usage_count)) # Tier 2: Keys that can accept more concurrent requests - elif key_state["models_in_use"].get(model, 0) < max_concurrent: + elif ( + key_state["models_in_use"].get(model, 0) + < effective_max_concurrent + ): tier2_keys.append((key, usage_count)) - # Determine selection method based on provider's rotation mode - provider = model.split("/")[0] if "/" in model else "" - rotation_mode = self._get_rotation_mode(provider) - if rotation_mode == "sequential": - # Sequential mode: stick with same credential until exhausted + # Sequential mode: sort credentials by priority, usage, recency + # Keep all candidates in sorted order (no filtering to single key) selection_method = "sequential" if tier1_keys: - selected_key = self._select_sequential( + tier1_keys = self._sort_sequential( tier1_keys, credential_priorities ) - tier1_keys = [ - (k, u) for k, u in tier1_keys if k == selected_key - ] if tier2_keys: - selected_key = self._select_sequential( + tier2_keys = self._sort_sequential( tier2_keys, credential_priorities ) - tier2_keys = [ - (k, u) for k, u in tier2_keys if k == selected_key - ] elif self.rotation_tolerance > 0: # Balanced mode with weighted randomness selection_method = "weighted-random" @@ -1057,7 +1179,7 @@ async def acquire_key( state = self.key_states[key] async with state["lock"]: current_count = state["models_in_use"].get(model, 0) - if current_count < max_concurrent: + if current_count < effective_max_concurrent: state["models_in_use"][model] = current_count + 1 tier_name = ( credential_tier_names.get(key, "unknown") @@ -1066,7 +1188,7 @@ async def acquire_key( ) lib_logger.info( f"Acquired key {mask_credential(key)} for model {model} " - f"(tier: {tier_name}, priority: {priority_level}, selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" + f"(tier: {tier_name}, priority: {priority_level}, selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{effective_max_concurrent}, usage: {usage})" ) return key @@ -1095,6 +1217,19 @@ async def acquire_key( else: # Original logic when no priorities specified + + # Determine selection method based on provider's rotation mode + provider = model.split("/")[0] if "/" in model else "" + rotation_mode = self._get_rotation_mode(provider) + + # Calculate effective concurrency for default priority (999) + # When no priorities are specified, all credentials get default priority + default_priority = 999 + multiplier = self._get_priority_multiplier( + provider, default_priority, rotation_mode + ) + effective_max_concurrent = max_concurrent * multiplier + tier1_keys, tier2_keys = [], [] # First, filter the list of available keys to exclude any on cooldown. @@ -1108,37 +1243,32 @@ async def acquire_key( continue # Prioritize keys based on their current usage to ensure load balancing. - usage_count = self._get_usage_count(key, model) + # Uses grouped usage if model is in a quota group + usage_count = self._get_grouped_usage_count(key, model) key_state = self.key_states[key] # Tier 1: Completely idle keys (preferred). if not key_state["models_in_use"]: tier1_keys.append((key, usage_count)) # Tier 2: Keys that can accept more concurrent requests for this model. - elif key_state["models_in_use"].get(model, 0) < max_concurrent: + elif ( + key_state["models_in_use"].get(model, 0) + < effective_max_concurrent + ): tier2_keys.append((key, usage_count)) - # Determine selection method based on provider's rotation mode - provider = model.split("/")[0] if "/" in model else "" - rotation_mode = self._get_rotation_mode(provider) - if rotation_mode == "sequential": - # Sequential mode: stick with same credential until exhausted + # Sequential mode: sort credentials by priority, usage, recency + # Keep all candidates in sorted order (no filtering to single key) selection_method = "sequential" if tier1_keys: - selected_key = self._select_sequential( + tier1_keys = self._sort_sequential( tier1_keys, credential_priorities ) - tier1_keys = [ - (k, u) for k, u in tier1_keys if k == selected_key - ] if tier2_keys: - selected_key = self._select_sequential( + tier2_keys = self._sort_sequential( tier2_keys, credential_priorities ) - tier2_keys = [ - (k, u) for k, u in tier2_keys if k == selected_key - ] elif self.rotation_tolerance > 0: # Balanced mode with weighted randomness selection_method = "weighted-random" @@ -1185,7 +1315,7 @@ async def acquire_key( state = self.key_states[key] async with state["lock"]: current_count = state["models_in_use"].get(model, 0) - if current_count < max_concurrent: + if current_count < effective_max_concurrent: state["models_in_use"][model] = current_count + 1 tier_name = ( credential_tier_names.get(key) @@ -1195,7 +1325,7 @@ async def acquire_key( tier_info = f"tier: {tier_name}, " if tier_name else "" lib_logger.info( f"Acquired key {mask_credential(key)} for model {model} " - f"({tier_info}selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{max_concurrent}, usage: {usage})" + f"({tier_info}selection: {selection_method}, concurrent: {state['models_in_use'][model]}/{effective_max_concurrent}, usage: {usage})" ) return key From 672c6bd67817aafad1849d7db6588ab45c5a904c Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 02:41:31 +0100 Subject: [PATCH 099/221] =?UTF-8?q?docs:=20=F0=9F=93=9A=20document=20seque?= =?UTF-8?q?ntial=20rotation=20and=20per-model=20quota=20tracking=20feature?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit comprehensively documents the new credential rotation and quota management system introduced in PR #31. Documentation updates include: - **Rotation Modes**: Detailed explanation of balanced vs sequential rotation strategies, with configuration examples and use cases for each mode - **Per-Model Quota Tracking**: Complete documentation of the granular per-model usage tracking system with authoritative quota reset timestamps - **Provider-Specific Quota Parsing**: Documentation of the `parse_quota_error()` extension point with Google RPC format examples - **Model Quota Groups**: Explanation of shared quota limits across model groups with configuration syntax - **Priority-Based Concurrency**: Documentation of tier-based concurrency multipliers with mode-specific override capabilities - **Reset Window Configuration**: Details on flexible rolling windows (5-hour, 7-day, etc.) replacing hardcoded daily resets - **Usage Flow**: Step-by-step explanation of the complete request lifecycle from credential selection through quota enforcement README updates include: - Feature highlights for all new capabilities in the features section - Configuration examples for rotation modes, concurrency multipliers, and quota groups - TUI enhancements for managing new settings - Provider-specific defaults and behaviors for Antigravity and Gemini CLI --- DOCUMENTATION.md | 248 ++++++++++++++++++++++++++++++++++++++++++++++- README.md | 49 ++++++++++ 2 files changed, 292 insertions(+), 5 deletions(-) diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index cf985326..1e96809d 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -96,22 +96,30 @@ The `_safe_streaming_wrapper` is a critical component for stability. It: ### 2.2. `usage_manager.py` - Stateful Concurrency & Usage Management -This class is the stateful core of the library, managing concurrency, usage tracking, and cooldowns. +This class is the stateful core of the library, managing concurrency, usage tracking, cooldowns, and quota resets. #### Key Concepts * **Async-Native & Lazy-Loaded**: Fully asynchronous, using `aiofiles` for non-blocking file I/O. Usage data is loaded only when needed. * **Fine-Grained Locking**: Each API key has its own `asyncio.Lock` and `asyncio.Condition`. This allows for highly granular control. +* **Multiple Reset Modes**: Supports three reset strategies: + - **per_model**: Each model has independent usage window with authoritative `quota_reset_ts` (from provider errors) + - **credential**: One window per credential with custom duration (e.g., 5 hours, 7 days) + - **daily**: Legacy daily reset at `daily_reset_time_utc` +* **Model Quota Groups**: Models can be grouped to share quota limits. When one model in a group hits quota, all receive the same reset timestamp. #### Tiered Key Acquisition Strategy The `acquire_key` method uses a sophisticated strategy to balance load: 1. **Filtering**: Keys currently on cooldown (global or model-specific) are excluded. -2. **Tiering**: Valid keys are split into two tiers: +2. **Rotation Mode**: Determines credential selection strategy: + * **Balanced Mode** (default): Credentials sorted by usage count - least-used first for even distribution + * **Sequential Mode**: Credentials sorted by usage count descending - most-used first to maintain sticky behavior until exhausted +3. **Tiering**: Valid keys are split into two tiers: * **Tier 1 (Ideal)**: Keys that are completely idle (0 concurrent requests). * **Tier 2 (Acceptable)**: Keys that are busy but still under their configured `MAX_CONCURRENT_REQUESTS_PER_KEY_` limit for the requested model. This allows a single key to be used multiple times for the same model, maximizing throughput. -3. **Selection Strategy** (configurable via `rotation_tolerance`): +4. **Selection Strategy** (configurable via `rotation_tolerance`): * **Deterministic (tolerance=0.0)**: Within each tier, keys are sorted by daily usage count and the least-used key is always selected. This provides perfect load balance but predictable patterns. * **Weighted Random (tolerance>0, default)**: Keys are selected randomly with weights biased toward less-used ones: - Formula: `weight = (max_usage - credential_usage) + tolerance + 1` @@ -119,14 +127,19 @@ The `acquire_key` method uses a sophisticated strategy to balance load: - `tolerance=5.0+`: High randomness - even heavily-used credentials have significant probability - **Security Benefit**: Unpredictable selection patterns make rate limit detection and fingerprinting harder - **Load Balance**: Lower-usage credentials still preferred, maintaining reasonable distribution -4. **Concurrency Limits**: Checks against `max_concurrent` limits to prevent overloading a single key. -5. **Priority Groups**: When credential prioritization is enabled, higher-tier credentials (lower priority numbers) are tried first before moving to lower tiers. +5. **Concurrency Limits**: Checks against `max_concurrent` limits (with priority multipliers applied) to prevent overloading a single key. +6. **Priority Groups**: When credential prioritization is enabled, higher-tier credentials (lower priority numbers) are tried first before moving to lower tiers. #### Failure Handling & Cooldowns * **Escalating Backoff**: When a failure occurs, the key gets a temporary cooldown for that specific model. Consecutive failures increase this time (10s -> 30s -> 60s -> 120s). * **Key-Level Lockouts**: If a key accumulates failures across multiple distinct models (3+), it is assumed to be dead/revoked and placed on a global 5-minute lockout. * **Authentication Errors**: Immediate 5-minute global lockout. +* **Quota Exhausted Errors**: When a provider returns a quota exhausted error with an authoritative reset timestamp: + - The `quota_reset_ts` is extracted from the error response (via provider's `parse_quota_error()` method) + - Applied to the affected model (and all models in its quota group if defined) + - Cooldown preserved even during daily/window resets until the actual quota reset time + - Logs show the exact reset time in local timezone with ISO format ### 2.3. `batch_manager.py` - Efficient Request Aggregation @@ -406,6 +419,10 @@ The most sophisticated provider implementation, supporting Google's internal Ant - **Thought Signature Caching**: Server-side caching of encrypted signatures for multi-turn Gemini 3 conversations - **Model-Specific Logic**: Automatic configuration based on model type (Gemini 3, Claude Sonnet, Claude Opus) - **Credential Prioritization**: Automatic tier detection with paid credentials prioritized over free (paid tier resets every 5 hours, free tier resets weekly) +- **Sequential Rotation Mode**: Default rotation mode is sequential (use credentials until exhausted) to maximize thought signature cache hits +- **Per-Model Quota Tracking**: Each model tracks independent usage windows with authoritative reset timestamps from quota errors +- **Quota Groups**: Claude models (Sonnet 4.5 + Opus 4.5) can be grouped to share quota limits (disabled by default, configurable via `QUOTA_GROUPS_ANTIGRAVITY_CLAUDE`) +- **Priority Multipliers**: Paid tier credentials get higher concurrency limits (Priority 1: 5x, Priority 2: 3x, Priority 3+: 2x in sequential mode) #### Model Support @@ -585,6 +602,221 @@ cache/ --- +### 2.13. Sequential Rotation & Per-Model Quota Tracking + +A comprehensive credential rotation and quota management system introduced in PR #31. + +#### Rotation Modes + +Two rotation strategies are available per provider: + +**Balanced Mode (Default)**: +- Distributes load evenly across all credentials +- Least-used credentials selected first +- Best for providers with per-minute rate limits +- Prevents any single credential from being overused + +**Sequential Mode**: +- Uses one credential until it's exhausted (429 quota error) +- Switches to next credential only after current one fails +- Most-used credentials selected first (sticky behavior) +- Best for providers with daily/weekly quotas +- Maximizes cache hit rates (e.g., Antigravity thought signatures) +- Default for Antigravity provider + +**Configuration**: +```env +# Set per provider +ROTATION_MODE_GEMINI=sequential +ROTATION_MODE_OPENAI=balanced +ROTATION_MODE_ANTIGRAVITY=balanced # Override default +``` + +#### Per-Model Quota Tracking + +Instead of tracking usage at the credential level, the system now supports granular per-model tracking: + +**Data Structure** (when `mode="per_model"`): +```json +{ + "credential_id": { + "models": { + "gemini-2.5-pro": { + "window_start_ts": 1733678400.0, + "quota_reset_ts": 1733696400.0, + "success_count": 15, + "prompt_tokens": 5000, + "completion_tokens": 1000, + "approx_cost": 0.05, + "window_started": "2025-12-08 14:00:00 +0100", + "quota_resets": "2025-12-08 19:00:00 +0100" + } + }, + "global": {...}, + "model_cooldowns": {...} + } +} +``` + +**Key Features**: +- Each model tracks its own usage window independently +- `window_start_ts`: When the current quota period started +- `quota_reset_ts`: Authoritative reset time from provider error response +- Human-readable timestamps added for debugging +- Supports custom window durations (5h, 7d, etc.) + +#### Provider-Specific Quota Parsing + +Providers can implement `parse_quota_error()` to extract precise reset times from error responses: + +```python +@staticmethod +def parse_quota_error(error, error_body) -> Optional[Dict]: + """Extract quota reset timestamp from provider error. + + Returns: + { + 'quota_reset_timestamp': 1733696400.0, # Unix timestamp + 'retry_after': 18000 # Seconds until reset + } + """ +``` + +**Google RPC Format** (Antigravity, Gemini CLI): +- Parses `RetryInfo` and `ErrorInfo` from error details +- Handles duration strings: `"143h4m52.73s"` or `"515092.73s"` +- Extracts `quotaResetTimeStamp` and converts to Unix timestamp +- Falls back to `quotaResetDelay` if timestamp not available + +**Example Error Response**: +```json +{ + "error": { + "code": 429, + "message": "Quota exceeded", + "details": [{ + "@type": "type.googleapis.com/google.rpc.RetryInfo", + "retryDelay": "143h4m52.73s" + }, { + "@type": "type.googleapis.com/google.rpc.ErrorInfo", + "metadata": { + "quotaResetTimeStamp": "2025-12-08T19:00:00Z" + } + }] + } +} +``` + +#### Model Quota Groups + +Models that share the same quota limits can be grouped: + +**Configuration**: +```env +# Models in a group share quota/cooldown timing +QUOTA_GROUPS_ANTIGRAVITY_CLAUDE="claude-sonnet-4-5,claude-opus-4-5" + +# To disable a default group: +QUOTA_GROUPS_ANTIGRAVITY_CLAUDE="" +``` + +**Behavior**: +- When one model hits quota, all models in the group receive the same `quota_reset_ts` +- Combined weighted usage for credential selection (e.g., Opus counts 2x vs Sonnet) +- Group resets only when ALL models' quotas have reset +- Preserves unexpired cooldowns during other resets + +**Provider Implementation**: +```python +class AntigravityProvider(ProviderInterface): + model_quota_groups = { + "claude": ["claude-sonnet-4-5", "claude-opus-4-5"] + } + + model_usage_weights = { + "claude-opus-4-5": 2 # Opus counts 2x vs Sonnet + } +``` + +#### Priority-Based Concurrency Multipliers + +Credentials can be assigned to priority tiers with configurable concurrency limits: + +**Configuration**: +```env +# Universal multipliers (all modes) +CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_1=10 +CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_2=3 + +# Mode-specific overrides +CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_2_BALANCED=1 # Lower in balanced mode +``` + +**How it works**: +```python +effective_concurrent_limit = MAX_CONCURRENT_REQUESTS_PER_KEY * tier_multiplier +``` + +**Provider Defaults** (Antigravity): +- Priority 1 (paid ultra): 5x multiplier +- Priority 2 (standard paid): 3x multiplier +- Priority 3+ (free): 2x (sequential mode) or 1x (balanced mode) + +**Benefits**: +- Paid credentials handle more load without manual configuration +- Different concurrency for different rotation modes +- Automatic tier detection based on credential properties + +#### Reset Window Configuration + +Providers can specify custom reset windows per priority tier: + +```python +class AntigravityProvider(ProviderInterface): + usage_reset_configs = { + frozenset([1, 2]): UsageResetConfigDef( + mode="per_model", + window_hours=5, # 5-hour rolling window for paid tiers + field_name="5h_window" + ), + frozenset([3, 4, 5]): UsageResetConfigDef( + mode="per_model", + window_hours=168, # 7-day window for free tier + field_name="7d_window" + ) + } +``` + +**Supported Modes**: +- `per_model`: Independent window per model with authoritative reset times +- `credential`: Single window per credential (legacy) +- `daily`: Daily reset at configured UTC hour (legacy) + +#### Usage Flow + +1. **Request arrives** for model X with credential Y +2. **Check rotation mode**: Sequential or balanced? +3. **Select credential**: + - Filter by priority tier requirements + - Apply concurrency multiplier for effective limit + - Sort by rotation mode strategy +4. **Check quota**: + - Load model's usage data + - Check if within window (window_start_ts to quota_reset_ts) + - Check model quota groups for combined usage +5. **Execute request** +6. **On success**: Increment model usage count +7. **On quota error**: + - Parse error for `quota_reset_ts` + - Apply to model (and quota group) + - Credential remains on cooldown until reset time +8. **On window expiration**: + - Archive model data to global stats + - Start fresh window with new `window_start_ts` + - Preserve unexpired quota cooldowns + +--- + ### 2.12. Google OAuth Base (`providers/google_oauth_base.py`) A refactored, reusable OAuth2 base class that eliminates code duplication across Google-based providers. @@ -637,6 +869,12 @@ The library handles provider idiosyncrasies through specialized "Provider" class The `GeminiCliProvider` is the most complex implementation, mimicking the Google Cloud Code extension. +**New in PR #31**: +- **Quota Parsing**: Implements `parse_quota_error()` using Google RPC format parser +- **Tier Configuration**: Defines `tier_priorities` and `usage_reset_configs` for automatic priority resolution +- **Balanced Rotation**: Defaults to balanced mode (unlike Antigravity which uses sequential) +- **Priority Multipliers**: Same as Antigravity (P1: 5x, P2: 3x, others: 1x) + #### Authentication (`gemini_auth_base.py`) * **Device Flow**: Uses a standard OAuth 2.0 flow. The `credential_tool` spins up a local web server (`localhost:8085`) to capture the callback from Google's auth page. diff --git a/README.md b/README.md index 9c3e3809..e746d422 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,12 @@ This project provides a powerful solution for developers building complex applic - Automatic thinking block sanitization for Claude models (with recovery strategies) - Note: Claude thinking mode requires careful conversation state management (see [Antigravity documentation](DOCUMENTATION.md#antigravity-claude-extended-thinking-sanitization) for details) - **🆕 Credential Prioritization**: Automatic tier detection and priority-based credential selection ensures paid-tier credentials are used for premium models that require them. +- **🆕 Sequential Rotation Mode**: Choose between balanced (distribute load evenly) or sequential (use until exhausted) credential rotation strategies. Sequential mode maximizes cache hit rates for providers like Antigravity. +- **🆕 Per-Model Quota Tracking**: Granular per-model usage tracking with authoritative quota reset timestamps from provider error responses. Each model maintains its own window with `window_start_ts` and `quota_reset_ts`. +- **🆕 Model Quota Groups**: Group models that share quota limits (e.g., Claude Sonnet and Opus). When one model in a group hits quota, all receive the same cooldown timestamp. +- **🆕 Priority-Based Concurrency**: Assign credentials to priority tiers (1=highest) with configurable concurrency multipliers. Paid-tier credentials can handle more concurrent requests than free-tier ones. +- **🆕 Provider-Specific Quota Parsing**: Extended provider interface with `parse_quota_error()` method to extract precise retry-after times from provider-specific error formats (e.g., Google RPC format). +- **🆕 Flexible Rolling Windows**: Support for provider-specific quota reset configurations (5-hour, 7-day, etc.) replacing hardcoded daily resets. - **🆕 Weighted Random Rotation**: Configurable credential rotation strategy - choose between deterministic (perfect balance) or weighted random (unpredictable, harder to fingerprint) selection. - **🆕 Enhanced Gemini CLI**: Improved project discovery, paid vs free tier detection, and Gemini 3 support with thoughtSignature caching. - **🆕 Temperature Override**: Global temperature=0 override option to prevent tool hallucination issues with low-temperature settings. @@ -129,6 +135,8 @@ The proxy now includes a powerful **interactive Text User Interface (TUI)** that - Configure custom OpenAI-compatible providers - Define provider models (simple or advanced JSON format) - Set concurrency limits per provider + - Configure rotation modes (balanced vs sequential) + - Manage priority-based concurrency multipliers - Interactive numbered menus for easy selection - Pending changes system with save/discard options @@ -545,6 +553,47 @@ ANTIGRAVITY_GEMINI3_TOOL_FIX=true # Prevent tool hallucination ``` +#### Credential Rotation Modes + +- **`ROTATION_MODE_`**: Controls how credentials are rotated when multiple are available. Default: `balanced` (except Antigravity which defaults to `sequential`). + - `balanced`: Rotate credentials evenly across requests to distribute load. Best for per-minute rate limits. + - `sequential`: Use one credential until exhausted (429 error), then switch to next. Best for daily/weekly quotas. + ```env + ROTATION_MODE_GEMINI=sequential # Use Gemini keys until quota exhausted + ROTATION_MODE_OPENAI=balanced # Distribute load across OpenAI keys (default) + ROTATION_MODE_ANTIGRAVITY=balanced # Override Antigravity's sequential default + ``` + +#### Priority-Based Concurrency Multipliers + +- **`CONCURRENCY_MULTIPLIER__PRIORITY_`**: Assign concurrency multipliers to priority tiers. Higher-tier credentials handle more concurrent requests. + ```env + # Universal multipliers (apply to all rotation modes) + CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_1=10 # 10x for paid ultra tier + CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_3=1 # 1x for lower tiers + + # Mode-specific overrides + CONCURRENCY_MULTIPLIER_ANTIGRAVITY_PRIORITY_2_BALANCED=1 # P2 = 1x in balanced mode only + ``` + + **Provider Defaults** (built into provider classes): + - **Antigravity**: Priority 1: 5x, Priority 2: 3x, Priority 3+: 2x (sequential) or 1x (balanced) + - **Gemini CLI**: Priority 1: 5x, Priority 2: 3x, Others: 1x + +#### Model Quota Groups + +- **`QUOTA_GROUPS__`**: Define models that share quota/cooldown timing. When one model hits quota, all in the group receive the same cooldown timestamp. + ```env + QUOTA_GROUPS_ANTIGRAVITY_CLAUDE="claude-sonnet-4-5,claude-opus-4-5" + QUOTA_GROUPS_ANTIGRAVITY_GEMINI="gemini-3-pro-preview,gemini-3-pro-image-preview" + + # To disable a default group: + QUOTA_GROUPS_ANTIGRAVITY_CLAUDE="" + ``` + + **Default Groups**: + - **Antigravity**: Claude group (Sonnet 4.5 + Opus 4.5) with Opus counting 2x vs Sonnet + #### Concurrency Control - **`MAX_CONCURRENT_REQUESTS_PER_KEY_`**: Set the maximum number of simultaneous requests allowed per API key for a specific provider. Default is `1` (no concurrency). Useful for high-throughput providers. From d655ada64961b2faead22b681ba6dd0e33326be7 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 03:17:42 +0100 Subject: [PATCH 100/221] =?UTF-8?q?fix(settings):=20=F0=9F=94=A8=20improve?= =?UTF-8?q?=20provider=20detection=20and=20configuration=20loading?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Refactor `get_default_mode` to use centralized PROVIDER_PLUGINS registry instead of ProviderInterface method, accessing default_rotation_mode directly from provider classes - Add comment filtering logic when parsing .env files to skip empty lines and comments starting with '#' - Update OAuth credentials directory path from 'oauth_credentials' to 'oauth_creds' for consistency --- src/proxy_app/settings_tool.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/proxy_app/settings_tool.py b/src/proxy_app/settings_tool.py index fe51cdf0..ddc0dae1 100644 --- a/src/proxy_app/settings_tool.py +++ b/src/proxy_app/settings_tool.py @@ -199,13 +199,13 @@ def get_current_modes(self) -> Dict[str, str]: def get_default_mode(self, provider: str) -> str: """Get the default rotation mode for a provider""" - # Import here to avoid circular imports try: - from rotator_library.providers.provider_interface import ( - ProviderInterface, - ) + from rotator_library.providers import PROVIDER_PLUGINS - return ProviderInterface.get_rotation_mode(provider) + provider_class = PROVIDER_PLUGINS.get(provider.lower()) + if provider_class and hasattr(provider_class, "default_rotation_mode"): + return provider_class.default_rotation_mode + return "balanced" except ImportError: # Fallback defaults if import fails if provider.lower() == "antigravity": @@ -527,6 +527,9 @@ def get_available_providers(self) -> List[str]: with open(env_file, "r", encoding="utf-8") as f: for line in f: line = line.strip() + # Skip comments and empty lines + if not line or line.startswith("#"): + continue if ( "_API_KEY" in line and "PROXY_API_KEY" not in line @@ -538,7 +541,7 @@ def get_available_providers(self) -> List[str]: pass # Also check for OAuth providers from files - oauth_dir = Path("oauth_credentials") + oauth_dir = Path("oauth_creds") if oauth_dir.exists(): for file in oauth_dir.glob("*_oauth_*.json"): provider = file.name.split("_oauth_")[0] From c5716c1fe58efdd916da57df906fa897739426c5 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 04:39:44 +0100 Subject: [PATCH 101/221] =?UTF-8?q?feat(tui):=20=F0=9F=94=A8=20add=20warni?= =?UTF-8?q?ngs=20for=20changing=20the=20proxy=20settings,=20and=20add=20re?= =?UTF-8?q?set=20to=20default=20button?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/proxy_app/launcher_tui.py | 642 +++++++++++++++++++++++----------- src/proxy_app/main.py | 522 ++++++++++++++++++--------- 2 files changed, 793 insertions(+), 371 deletions(-) diff --git a/src/proxy_app/launcher_tui.py b/src/proxy_app/launcher_tui.py index 2db109f9..954083dc 100644 --- a/src/proxy_app/launcher_tui.py +++ b/src/proxy_app/launcher_tui.py @@ -18,32 +18,33 @@ def clear_screen(): """ - Cross-platform terminal clear that works robustly on both + Cross-platform terminal clear that works robustly on both classic Windows conhost and modern terminals (Windows Terminal, Linux, Mac). - + Uses native OS commands instead of ANSI escape sequences: - Windows (conhost & Windows Terminal): cls - Unix-like systems (Linux, Mac): clear """ - os.system('cls' if os.name == 'nt' else 'clear') + os.system("cls" if os.name == "nt" else "clear") + class LauncherConfig: """Manages launcher_config.json (host, port, logging only)""" - + def __init__(self, config_path: Path = Path("launcher_config.json")): self.config_path = config_path self.defaults = { "host": "127.0.0.1", "port": 8000, - "enable_request_logging": False + "enable_request_logging": False, } self.config = self.load() - + def load(self) -> dict: """Load config from file or create with defaults.""" if self.config_path.exists(): try: - with open(self.config_path, 'r') as f: + with open(self.config_path, "r") as f: config = json.load(f) # Merge with defaults for any missing keys for key, value in self.defaults.items(): @@ -53,22 +54,23 @@ def load(self) -> dict: except (json.JSONDecodeError, IOError): return self.defaults.copy() return self.defaults.copy() - + def save(self): """Save current config to file.""" import datetime + self.config["last_updated"] = datetime.datetime.now().isoformat() try: - with open(self.config_path, 'w') as f: + with open(self.config_path, "w") as f: json.dump(self.config, f, indent=2) except IOError as e: console.print(f"[red]Error saving config: {e}[/red]") - + def update(self, **kwargs): """Update config values.""" self.config.update(kwargs) self.save() - + @staticmethod def update_proxy_api_key(new_key: str): """Update PROXY_API_KEY in .env only""" @@ -79,7 +81,7 @@ def update_proxy_api_key(new_key: str): class SettingsDetector: """Detects settings from .env for display""" - + @staticmethod def _load_local_env() -> dict: """Load environment variables from local .env file only""" @@ -88,13 +90,13 @@ def _load_local_env() -> dict: if not env_file.exists(): return env_dict try: - with open(env_file, 'r', encoding='utf-8') as f: + with open(env_file, "r", encoding="utf-8") as f: for line in f: line = line.strip() - if not line or line.startswith('#'): + if not line or line.startswith("#"): continue - if '=' in line: - key, _, value = line.partition('=') + if "=" in line: + key, _, value = line.partition("=") key, value = key.strip(), value.strip() if value and value[0] in ('"', "'") and value[-1] == value[0]: value = value[1:-1] @@ -112,16 +114,16 @@ def get_all_settings() -> dict: "model_definitions": SettingsDetector.detect_model_definitions(), "concurrency_limits": SettingsDetector.detect_concurrency_limits(), "model_filters": SettingsDetector.detect_model_filters(), - "provider_settings": SettingsDetector.detect_provider_settings() + "provider_settings": SettingsDetector.detect_provider_settings(), } - + @staticmethod def detect_credentials() -> dict: """Detect API keys and OAuth credentials""" from pathlib import Path - + providers = {} - + # Scan for API keys env_vars = SettingsDetector._load_local_env() for key, value in env_vars.items(): @@ -130,7 +132,7 @@ def detect_credentials() -> dict: if provider not in providers: providers[provider] = {"api_keys": 0, "oauth": 0, "custom": False} providers[provider]["api_keys"] += 1 - + # Scan for OAuth credentials oauth_dir = Path("oauth_credentials") if oauth_dir.exists(): @@ -139,19 +141,19 @@ def detect_credentials() -> dict: if provider not in providers: providers[provider] = {"api_keys": 0, "oauth": 0, "custom": False} providers[provider]["oauth"] += 1 - + # Mark custom providers (have API_BASE set) for provider in providers: if os.getenv(f"{provider.upper()}_API_BASE"): providers[provider]["custom"] = True - + return providers - + @staticmethod def detect_custom_api_bases() -> dict: """Detect custom API base URLs (not in hardcoded map)""" from proxy_app.provider_urls import PROVIDER_URL_MAP - + bases = {} env_vars = SettingsDetector._load_local_env() for key, value in env_vars.items(): @@ -161,7 +163,7 @@ def detect_custom_api_bases() -> dict: if provider not in PROVIDER_URL_MAP: bases[provider] = value return bases - + @staticmethod def detect_model_definitions() -> dict: """Detect provider model definitions""" @@ -179,7 +181,7 @@ def detect_model_definitions() -> dict: except (json.JSONDecodeError, ValueError): pass return models - + @staticmethod def detect_concurrency_limits() -> dict: """Detect max concurrent requests per key""" @@ -193,7 +195,7 @@ def detect_concurrency_limits() -> dict: except (json.JSONDecodeError, ValueError): pass return limits - + @staticmethod def detect_model_filters() -> dict: """Detect active model filters (basic info only: defined or not)""" @@ -210,7 +212,7 @@ def detect_model_filters() -> dict: else: filters[provider]["has_whitelist"] = True return filters - + @staticmethod def detect_provider_settings() -> dict: """Detect provider-specific settings (Antigravity, Gemini CLI)""" @@ -219,10 +221,10 @@ def detect_provider_settings() -> dict: except ImportError: # Fallback for direct execution or testing from .settings_tool import PROVIDER_SETTINGS_MAP - + provider_settings = {} env_vars = SettingsDetector._load_local_env() - + for provider, definitions in PROVIDER_SETTINGS_MAP.items(): modified_count = 0 for key, definition in definitions.items(): @@ -231,7 +233,7 @@ def detect_provider_settings() -> dict: # Check if value differs from default default = definition.get("default") setting_type = definition.get("type", "str") - + try: if setting_type == "bool": current = env_value.lower() in ("true", "1", "yes") @@ -239,21 +241,21 @@ def detect_provider_settings() -> dict: current = int(env_value) else: current = env_value - + if current != default: modified_count += 1 except (ValueError, AttributeError): pass - + if modified_count > 0: provider_settings[provider] = modified_count - + return provider_settings class LauncherTUI: """Main launcher interface""" - + def __init__(self): self.console = Console() self.config = LauncherConfig() @@ -261,90 +263,100 @@ def __init__(self): self.env_file = Path.cwd() / ".env" # Load .env file to ensure environment variables are available load_dotenv(dotenv_path=self.env_file, override=True) - + def needs_onboarding(self) -> bool: """Check if onboarding is needed""" return not self.env_file.exists() or not os.getenv("PROXY_API_KEY") - + def run(self): """Main TUI loop""" while self.running: self.show_main_menu() - + def show_main_menu(self): """Display main menu and handle selection""" clear_screen() - + # Detect all settings settings = SettingsDetector.get_all_settings() credentials = settings["credentials"] custom_bases = settings["custom_bases"] - + # Check if setup is needed show_warning = self.needs_onboarding() - + # Build title with GitHub link - self.console.print(Panel.fit( - "[bold cyan]🚀 LLM API Key Proxy - Interactive Launcher[/bold cyan]", - border_style="cyan" - )) - self.console.print("[dim]GitHub: [blue underline]https://github.com/Mirrowel/LLM-API-Key-Proxy[/blue underline][/dim]") - + self.console.print( + Panel.fit( + "[bold cyan]🚀 LLM API Key Proxy - Interactive Launcher[/bold cyan]", + border_style="cyan", + ) + ) + self.console.print( + "[dim]GitHub: [blue underline]https://github.com/Mirrowel/LLM-API-Key-Proxy[/blue underline][/dim]" + ) + # Show warning if .env file doesn't exist if show_warning: self.console.print() - self.console.print(Panel( - Text.from_markup( - "⚠️ [bold yellow]INITIAL SETUP REQUIRED[/bold yellow]\n\n" - "The proxy needs initial configuration:\n" - " ❌ No .env file found\n\n" - "Why this matters:\n" - " • The .env file stores your credentials and settings\n" - " • PROXY_API_KEY protects your proxy from unauthorized access\n" - " • Provider API keys enable LLM access\n\n" - "What to do:\n" - " 1. Select option \"3. Manage Credentials\" to launch the credential tool\n" - " 2. The tool will create .env and set up PROXY_API_KEY automatically\n" - " 3. You can add provider credentials (API keys or OAuth)\n\n" - "⚠️ Note: The credential tool adds PROXY_API_KEY by default.\n" - " You can remove it later if you want an unsecured proxy." - ), - border_style="yellow", - expand=False - )) + self.console.print( + Panel( + Text.from_markup( + "⚠️ [bold yellow]INITIAL SETUP REQUIRED[/bold yellow]\n\n" + "The proxy needs initial configuration:\n" + " ❌ No .env file found\n\n" + "Why this matters:\n" + " • The .env file stores your credentials and settings\n" + " • PROXY_API_KEY protects your proxy from unauthorized access\n" + " • Provider API keys enable LLM access\n\n" + "What to do:\n" + ' 1. Select option "3. Manage Credentials" to launch the credential tool\n' + " 2. The tool will create .env and set up PROXY_API_KEY automatically\n" + " 3. You can add provider credentials (API keys or OAuth)\n\n" + "⚠️ Note: The credential tool adds PROXY_API_KEY by default.\n" + " You can remove it later if you want an unsecured proxy." + ), + border_style="yellow", + expand=False, + ) + ) # Show security warning if PROXY_API_KEY is missing (but .env exists) elif not os.getenv("PROXY_API_KEY"): self.console.print() - self.console.print(Panel( - Text.from_markup( - "⚠️ [bold red]SECURITY WARNING: PROXY_API_KEY Not Set[/bold red]\n\n" - "Your proxy is currently UNSECURED!\n" - "Anyone can access it without authentication.\n\n" - "This is a serious security risk if your proxy is accessible\n" - "from the internet or untrusted networks.\n\n" - "👉 [bold]Recommended:[/bold] Set PROXY_API_KEY in .env file\n" - " Use option \"2. Configure Proxy Settings\" → \"3. Set Proxy API Key\"\n" - " or option \"3. Manage Credentials\"" - ), - border_style="red", - expand=False - )) - + self.console.print( + Panel( + Text.from_markup( + "⚠️ [bold red]SECURITY WARNING: PROXY_API_KEY Not Set[/bold red]\n\n" + "Your proxy is currently UNSECURED!\n" + "Anyone can access it without authentication.\n\n" + "This is a serious security risk if your proxy is accessible\n" + "from the internet or untrusted networks.\n\n" + "👉 [bold]Recommended:[/bold] Set PROXY_API_KEY in .env file\n" + ' Use option "2. Configure Proxy Settings" → "3. Set Proxy API Key"\n' + ' or option "3. Manage Credentials"' + ), + border_style="red", + expand=False, + ) + ) + # Show config self.console.print() self.console.print("[bold]📋 Proxy Configuration[/bold]") self.console.print("━" * 70) self.console.print(f" Host: {self.config.config['host']}") self.console.print(f" Port: {self.config.config['port']}") - self.console.print(f" Request Logging: {'✅ Enabled' if self.config.config['enable_request_logging'] else '❌ Disabled'}") - + self.console.print( + f" Request Logging: {'✅ Enabled' if self.config.config['enable_request_logging'] else '❌ Disabled'}" + ) + # Show actual API key value - proxy_key = os.getenv('PROXY_API_KEY') + proxy_key = os.getenv("PROXY_API_KEY") if proxy_key: self.console.print(f" Proxy API Key: {proxy_key}") else: self.console.print(" Proxy API Key: [red]Not Set (INSECURE!)[/red]") - + # Show status summary self.console.print() self.console.print("[bold]📊 Status Summary[/bold]") @@ -352,12 +364,19 @@ def show_main_menu(self): provider_count = len(credentials) custom_count = len(custom_bases) provider_settings = settings.get("provider_settings", {}) - has_advanced = bool(settings["model_definitions"] or settings["concurrency_limits"] or settings["model_filters"] or provider_settings) - + has_advanced = bool( + settings["model_definitions"] + or settings["concurrency_limits"] + or settings["model_filters"] + or provider_settings + ) + self.console.print(f" Providers: {provider_count} configured") self.console.print(f" Custom Providers: {custom_count} configured") - self.console.print(f" Advanced Settings: {'Active (view in menu 4)' if has_advanced else 'None'}") - + self.console.print( + f" Advanced Settings: {'Active (view in menu 4)' if has_advanced else 'None'}" + ) + # Show menu self.console.print() self.console.print("━" * 70) @@ -367,23 +386,29 @@ def show_main_menu(self): if show_warning: self.console.print(" 1. ▶️ Run Proxy Server") self.console.print(" 2. ⚙️ Configure Proxy Settings") - self.console.print(" 3. 🔑 Manage Credentials ⬅️ [bold yellow]Start here![/bold yellow]") + self.console.print( + " 3. 🔑 Manage Credentials ⬅️ [bold yellow]Start here![/bold yellow]" + ) else: self.console.print(" 1. ▶️ Run Proxy Server") self.console.print(" 2. ⚙️ Configure Proxy Settings") self.console.print(" 3. 🔑 Manage Credentials") - + self.console.print(" 4. 📊 View Provider & Advanced Settings") self.console.print(" 5. 🔄 Reload Configuration") self.console.print(" 6. ℹ️ About") self.console.print(" 7. 🚪 Exit") - + self.console.print() self.console.print("━" * 70) self.console.print() - - choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5", "6", "7"], show_choices=False) - + + choice = Prompt.ask( + "Select option", + choices=["1", "2", "3", "4", "5", "6", "7"], + show_choices=False, + ) + if choice == "1": self.run_proxy() elif choice == "2": @@ -393,7 +418,7 @@ def show_main_menu(self): elif choice == "4": self.show_provider_settings_menu() elif choice == "5": - load_dotenv(dotenv_path=Path.cwd() / ".env",override=True) + load_dotenv(dotenv_path=Path.cwd() / ".env", override=True) self.config = LauncherConfig() # Reload config self.console.print("\n[green]✅ Configuration reloaded![/green]") elif choice == "6": @@ -401,25 +426,64 @@ def show_main_menu(self): elif choice == "7": self.running = False sys.exit(0) - + + def confirm_setting_change(self, setting_name: str, warning_lines: list) -> bool: + """ + Display a warning and require Y/N (case-sensitive) confirmation. + Re-prompts until user enters exactly 'Y' or 'N'. + Returns True only if user enters 'Y'. + """ + clear_screen() + self.console.print() + self.console.print( + Panel( + Text.from_markup( + f"[bold yellow]⚠️ WARNING: You are about to change the {setting_name}[/bold yellow]\n\n" + + "\n".join(warning_lines) + + "\n\n[bold]If you are not sure about changing this - don't.[/bold]" + ), + border_style="yellow", + expand=False, + ) + ) + + while True: + response = Prompt.ask( + "Enter [bold]Y[/bold] to confirm, [bold]N[/bold] to cancel (case-sensitive)" + ) + if response == "Y": + return True + elif response == "N": + self.console.print("\n[dim]Operation cancelled.[/dim]") + return False + else: + self.console.print( + "[red]Please enter exactly 'Y' or 'N' (case-sensitive)[/red]" + ) + def show_config_menu(self): """Display configuration sub-menu""" while True: clear_screen() - - self.console.print(Panel.fit( - "[bold cyan]⚙️ Proxy Configuration[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + "[bold cyan]⚙️ Proxy Configuration[/bold cyan]", border_style="cyan" + ) + ) + self.console.print() self.console.print("[bold]📋 Current Settings[/bold]") self.console.print("━" * 70) self.console.print(f" Host: {self.config.config['host']}") self.console.print(f" Port: {self.config.config['port']}") - self.console.print(f" Request Logging: {'✅ Enabled' if self.config.config['enable_request_logging'] else '❌ Disabled'}") - self.console.print(f" Proxy API Key: {'✅ Set' if os.getenv('PROXY_API_KEY') else '❌ Not Set'}") - + self.console.print( + f" Request Logging: {'✅ Enabled' if self.config.config['enable_request_logging'] else '❌ Disabled'}" + ) + self.console.print( + f" Proxy API Key: {'✅ Set' if os.getenv('PROXY_API_KEY') else '❌ Not Set'}" + ) + self.console.print() self.console.print("━" * 70) self.console.print() @@ -429,45 +493,172 @@ def show_config_menu(self): self.console.print(" 2. 🔌 Set Port") self.console.print(" 3. 🔑 Set Proxy API Key") self.console.print(" 4. 📝 Toggle Request Logging") - self.console.print(" 5. ↩️ Back to Main Menu") - + self.console.print(" 5. 🔄 Reset to Default Settings") + self.console.print(" 6. ↩️ Back to Main Menu") + self.console.print() self.console.print("━" * 70) self.console.print() - - choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5"], show_choices=False) - + + choice = Prompt.ask( + "Select option", + choices=["1", "2", "3", "4", "5", "6"], + show_choices=False, + ) + if choice == "1": - new_host = Prompt.ask("Enter new host IP", default=self.config.config["host"]) + # Show warning and require confirmation + confirmed = self.confirm_setting_change( + "Host IP", + [ + "Changing the host IP affects which network interfaces the proxy listens on:", + " • [cyan]127.0.0.1[/cyan] = Local access only (recommended for development)", + " • [cyan]0.0.0.0[/cyan] = Accessible from all network interfaces", + "", + "Applications configured to connect to the old host may fail to connect.", + ], + ) + if not confirmed: + continue + + new_host = Prompt.ask( + "Enter new host IP", default=self.config.config["host"] + ) self.config.update(host=new_host) self.console.print(f"\n[green]✅ Host updated to: {new_host}[/green]") elif choice == "2": - new_port = IntPrompt.ask("Enter new port", default=self.config.config["port"]) + # Show warning and require confirmation + confirmed = self.confirm_setting_change( + "Port", + [ + "Changing the port will affect all applications currently configured", + "to connect to your proxy on the existing port.", + "", + "Applications using the old port will fail to connect.", + ], + ) + if not confirmed: + continue + + new_port = IntPrompt.ask( + "Enter new port", default=self.config.config["port"] + ) if 1 <= new_port <= 65535: self.config.update(port=new_port) - self.console.print(f"\n[green]✅ Port updated to: {new_port}[/green]") + self.console.print( + f"\n[green]✅ Port updated to: {new_port}[/green]" + ) else: self.console.print("\n[red]❌ Port must be between 1-65535[/red]") elif choice == "3": + # Show warning and require confirmation + confirmed = self.confirm_setting_change( + "Proxy API Key", + [ + "This is the authentication key that applications use to access your proxy.", + "", + "[bold red]⚠️ Changing this will BREAK all applications currently configured", + " with the existing API key![/bold red]", + "", + "[bold cyan]💡 If you want to add provider API keys (OpenAI, Gemini, etc.),", + ' go to "3. 🔑 Manage Credentials" in the main menu instead.[/bold cyan]', + ], + ) + if not confirmed: + continue + current = os.getenv("PROXY_API_KEY", "") - new_key = Prompt.ask("Enter new Proxy API Key", default=current) - if new_key and new_key != current: + new_key = Prompt.ask( + "Enter new Proxy API Key (leave empty to disable authentication)", + default=current, + ) + + if new_key != current: + # If setting to empty, show additional warning + if not new_key: + self.console.print( + "\n[bold red]⚠️ Authentication will be DISABLED - anyone can access your proxy![/bold red]" + ) + Prompt.ask("Press Enter to continue", default="") + LauncherConfig.update_proxy_api_key(new_key) - self.console.print("\n[green]✅ Proxy API Key updated successfully![/green]") - self.console.print(" Updated in .env file") + + if new_key: + self.console.print( + "\n[green]✅ Proxy API Key updated successfully![/green]" + ) + self.console.print(" Updated in .env file") + else: + self.console.print( + "\n[yellow]⚠️ Proxy API Key cleared - authentication disabled![/yellow]" + ) + self.console.print(" Updated in .env file") else: self.console.print("\n[yellow]No changes made[/yellow]") elif choice == "4": current = self.config.config["enable_request_logging"] self.config.update(enable_request_logging=not current) - self.console.print(f"\n[green]✅ Request Logging {'enabled' if not current else 'disabled'}![/green]") + self.console.print( + f"\n[green]✅ Request Logging {'enabled' if not current else 'disabled'}![/green]" + ) elif choice == "5": + # Reset to Default Settings + # Define defaults + default_host = "127.0.0.1" + default_port = 8000 + default_logging = False + default_api_key = "VerysecretKey" + + # Get current values + current_host = self.config.config["host"] + current_port = self.config.config["port"] + current_logging = self.config.config["enable_request_logging"] + current_api_key = os.getenv("PROXY_API_KEY", "") + + # Build comparison table + warning_lines = [ + "This will reset ALL proxy settings to their defaults:", + "", + "[bold] Setting Current Value → Default Value[/bold]", + " " + "─" * 62, + f" Host IP {current_host:20} → {default_host}", + f" Port {str(current_port):20} → {default_port}", + f" Request Logging {'Enabled':20} → Disabled" + if current_logging + else f" Request Logging {'Disabled':20} → Disabled", + f" Proxy API Key {current_api_key[:20]:20} → {default_api_key}", + "", + "[bold red]⚠️ This may break applications configured with current settings![/bold red]", + ] + + confirmed = self.confirm_setting_change( + "Settings (Reset to Defaults)", warning_lines + ) + if not confirmed: + continue + + # Apply defaults + self.config.update( + host=default_host, + port=default_port, + enable_request_logging=default_logging, + ) + LauncherConfig.update_proxy_api_key(default_api_key) + + self.console.print( + "\n[green]✅ All settings have been reset to defaults![/green]" + ) + self.console.print(f" Host: {default_host}") + self.console.print(f" Port: {default_port}") + self.console.print(f" Request Logging: Disabled") + self.console.print(f" Proxy API Key: {default_api_key}") + elif choice == "6": break - + def show_provider_settings_menu(self): """Display provider/advanced settings (read-only + launch tool)""" clear_screen() - + settings = SettingsDetector.get_all_settings() credentials = settings["credentials"] custom_bases = settings["custom_bases"] @@ -475,12 +666,14 @@ def show_provider_settings_menu(self): concurrency = settings["concurrency_limits"] filters = settings["model_filters"] provider_settings = settings.get("provider_settings", {}) - - self.console.print(Panel.fit( - "[bold cyan]📊 Provider & Advanced Settings[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + "[bold cyan]📊 Provider & Advanced Settings[/bold cyan]", + border_style="cyan", + ) + ) + # Configured Providers self.console.print() self.console.print("[bold]📊 Configured Providers[/bold]") @@ -490,18 +683,22 @@ def show_provider_settings_menu(self): provider_name = provider.title() parts = [] if info["api_keys"] > 0: - parts.append(f"{info['api_keys']} API key{'s' if info['api_keys'] > 1 else ''}") + parts.append( + f"{info['api_keys']} API key{'s' if info['api_keys'] > 1 else ''}" + ) if info["oauth"] > 0: - parts.append(f"{info['oauth']} OAuth credential{'s' if info['oauth'] > 1 else ''}") - + parts.append( + f"{info['oauth']} OAuth credential{'s' if info['oauth'] > 1 else ''}" + ) + display = " + ".join(parts) if info["custom"]: display += " (Custom)" - + self.console.print(f" ✅ {provider_name:20} {display}") else: self.console.print(" [dim]No providers configured[/dim]") - + # Custom API Bases if custom_bases: self.console.print() @@ -509,15 +706,17 @@ def show_provider_settings_menu(self): self.console.print("━" * 70) for provider, base in custom_bases.items(): self.console.print(f" • {provider:15} {base}") - + # Model Definitions if model_defs: self.console.print() self.console.print("[bold]📦 Provider Model Definitions[/bold]") self.console.print("━" * 70) for provider, count in model_defs.items(): - self.console.print(f" • {provider:15} {count} model{'s' if count > 1 else ''} configured") - + self.console.print( + f" • {provider:15} {count} model{'s' if count > 1 else ''} configured" + ) + # Concurrency Limits if concurrency: self.console.print() @@ -526,7 +725,7 @@ def show_provider_settings_menu(self): for provider, limit in concurrency.items(): self.console.print(f" • {provider:15} {limit} requests/key") self.console.print(" • Default: 1 request/key (all others)") - + # Model Filters (basic info only) if filters: self.console.print() @@ -540,7 +739,7 @@ def show_provider_settings_menu(self): status_parts.append("Ignore list") status = " + ".join(status_parts) if status_parts else "None" self.console.print(f" • {provider:15} ✅ {status}") - + # Provider-Specific Settings self.console.print() self.console.print("[bold]🔬 Provider-Specific Settings[/bold]") @@ -553,158 +752,207 @@ def show_provider_settings_menu(self): display_name = provider.replace("_", " ").title() modified = provider_settings.get(provider, 0) if modified > 0: - self.console.print(f" • {display_name:20} [yellow]{modified} setting{'s' if modified > 1 else ''} modified[/yellow]") + self.console.print( + f" • {display_name:20} [yellow]{modified} setting{'s' if modified > 1 else ''} modified[/yellow]" + ) else: self.console.print(f" • {display_name:20} [dim]using defaults[/dim]") - + # Actions self.console.print() self.console.print("━" * 70) self.console.print() self.console.print("[bold]💡 Actions[/bold]") self.console.print() - self.console.print(" 1. 🔧 Launch Settings Tool (configure advanced settings)") + self.console.print( + " 1. 🔧 Launch Settings Tool (configure advanced settings)" + ) self.console.print(" 2. ↩️ Back to Main Menu") - + self.console.print() self.console.print("━" * 70) - self.console.print("[dim]ℹ️ Advanced settings are stored in .env file.\n Use the Settings Tool to configure them interactively.[/dim]") + self.console.print( + "[dim]ℹ️ Advanced settings are stored in .env file.\n Use the Settings Tool to configure them interactively.[/dim]" + ) self.console.print() - self.console.print("[dim]⚠️ Note: Settings Tool supports only common configuration types.\n For complex settings, edit .env directly.[/dim]") + self.console.print( + "[dim]⚠️ Note: Settings Tool supports only common configuration types.\n For complex settings, edit .env directly.[/dim]" + ) self.console.print() - + choice = Prompt.ask("Select option", choices=["1", "2"], show_choices=False) - + if choice == "1": self.launch_settings_tool() # choice == "2" returns to main menu - + def launch_credential_tool(self): """Launch credential management tool""" import time - + # CRITICAL: Show full loading UI to replace the 6-7 second blank wait clear_screen() - + _start_time = time.time() - + # Show the same header as standalone mode self.console.print("━" * 70) self.console.print("Interactive Credential Setup Tool") self.console.print("GitHub: https://github.com/Mirrowel/LLM-API-Key-Proxy") self.console.print("━" * 70) self.console.print("Loading credential management components...") - + # Now import with spinner (this is where the 6-7 second delay happens) with self.console.status("Initializing credential tool...", spinner="dots"): - from rotator_library.credential_tool import run_credential_tool, _ensure_providers_loaded + from rotator_library.credential_tool import ( + run_credential_tool, + _ensure_providers_loaded, + ) + _, PROVIDER_PLUGINS = _ensure_providers_loaded() self.console.print("✓ Credential tool initialized") _elapsed = time.time() - _start_time - self.console.print(f"✓ Tool ready in {_elapsed:.2f}s ({len(PROVIDER_PLUGINS)} providers available)") - + self.console.print( + f"✓ Tool ready in {_elapsed:.2f}s ({len(PROVIDER_PLUGINS)} providers available)" + ) + # Small delay to let user see the ready message time.sleep(0.5) - + # Run the tool with from_launcher=True to skip duplicate loading screen run_credential_tool(from_launcher=True) # Reload environment after credential tool load_dotenv(dotenv_path=Path.cwd() / ".env", override=True) - + def launch_settings_tool(self): """Launch settings configuration tool""" from proxy_app.settings_tool import run_settings_tool + run_settings_tool() # Reload environment after settings tool load_dotenv(dotenv_path=Path.cwd() / ".env", override=True) - + def show_about(self): """Display About page with project information""" clear_screen() - - self.console.print(Panel.fit( - "[bold cyan]ℹ️ About LLM API Key Proxy[/bold cyan]", - border_style="cyan" - )) - + + self.console.print( + Panel.fit( + "[bold cyan]ℹ️ About LLM API Key Proxy[/bold cyan]", border_style="cyan" + ) + ) + self.console.print() self.console.print("[bold]📦 Project Information[/bold]") self.console.print("━" * 70) self.console.print(" [bold cyan]LLM API Key Proxy[/bold cyan]") - self.console.print(" A lightweight, high-performance proxy server for managing") + self.console.print( + " A lightweight, high-performance proxy server for managing" + ) self.console.print(" LLM API keys with automatic rotation and OAuth support") self.console.print() - self.console.print(" [dim]GitHub:[/dim] [blue underline]https://github.com/Mirrowel/LLM-API-Key-Proxy[/blue underline]") - + self.console.print( + " [dim]GitHub:[/dim] [blue underline]https://github.com/Mirrowel/LLM-API-Key-Proxy[/blue underline]" + ) + self.console.print() self.console.print("[bold]✨ Key Features[/bold]") self.console.print("━" * 70) - self.console.print(" • [green]Smart Key Rotation[/green] - Automatic rotation across multiple API keys") - self.console.print(" • [green]OAuth Support[/green] - Automated OAuth flows for supported providers") - self.console.print(" • [green]Multiple Providers[/green] - Support for 10+ LLM providers") - self.console.print(" • [green]Custom Providers[/green] - Easy integration of custom OpenAI-compatible APIs") - self.console.print(" • [green]Advanced Filtering[/green] - Model whitelists and ignore lists per provider") - self.console.print(" • [green]Concurrency Control[/green] - Per-key rate limiting and request management") - self.console.print(" • [green]Cost Tracking[/green] - Track usage and costs across all providers") - self.console.print(" • [green]Interactive TUI[/green] - Beautiful terminal interface for easy configuration") - + self.console.print( + " • [green]Smart Key Rotation[/green] - Automatic rotation across multiple API keys" + ) + self.console.print( + " • [green]OAuth Support[/green] - Automated OAuth flows for supported providers" + ) + self.console.print( + " • [green]Multiple Providers[/green] - Support for 10+ LLM providers" + ) + self.console.print( + " • [green]Custom Providers[/green] - Easy integration of custom OpenAI-compatible APIs" + ) + self.console.print( + " • [green]Advanced Filtering[/green] - Model whitelists and ignore lists per provider" + ) + self.console.print( + " • [green]Concurrency Control[/green] - Per-key rate limiting and request management" + ) + self.console.print( + " • [green]Cost Tracking[/green] - Track usage and costs across all providers" + ) + self.console.print( + " • [green]Interactive TUI[/green] - Beautiful terminal interface for easy configuration" + ) + self.console.print() self.console.print("[bold]📝 License & Credits[/bold]") self.console.print("━" * 70) self.console.print(" Made with ❤️ by the community") self.console.print(" Open source - contributions welcome!") - + self.console.print() self.console.print("━" * 70) self.console.print() - + Prompt.ask("Press Enter to return to main menu", default="") - + def run_proxy(self): """Prepare and launch proxy in same window""" # Check if forced onboarding needed if self.needs_onboarding(): clear_screen() - self.console.print(Panel( - Text.from_markup( - "⚠️ [bold yellow]Setup Required[/bold yellow]\n\n" - "Cannot start without .env.\n" - "Launching credential tool..." - ), - border_style="yellow" - )) - + self.console.print( + Panel( + Text.from_markup( + "⚠️ [bold yellow]Setup Required[/bold yellow]\n\n" + "Cannot start without .env.\n" + "Launching credential tool..." + ), + border_style="yellow", + ) + ) + # Force credential tool - from rotator_library.credential_tool import ensure_env_defaults, run_credential_tool + from rotator_library.credential_tool import ( + ensure_env_defaults, + run_credential_tool, + ) + ensure_env_defaults() load_dotenv(dotenv_path=Path.cwd() / ".env", override=True) run_credential_tool() load_dotenv(dotenv_path=Path.cwd() / ".env", override=True) - + # Check again after credential tool if not os.getenv("PROXY_API_KEY"): - self.console.print("\n[red]❌ PROXY_API_KEY still not set. Cannot start proxy.[/red]") + self.console.print( + "\n[red]❌ PROXY_API_KEY still not set. Cannot start proxy.[/red]" + ) return - + # Clear console and modify sys.argv clear_screen() - self.console.print(f"\n[bold green]🚀 Starting proxy on {self.config.config['host']}:{self.config.config['port']}...[/bold green]\n") - + self.console.print( + f"\n[bold green]🚀 Starting proxy on {self.config.config['host']}:{self.config.config['port']}...[/bold green]\n" + ) + # Clear console again to remove the starting message before main.py shows loading details import time + time.sleep(0.5) # Brief pause so user sees the message clear_screen() - + # Reconstruct sys.argv for main.py sys.argv = [ "main.py", - "--host", self.config.config["host"], - "--port", str(self.config.config["port"]) + "--host", + self.config.config["host"], + "--port", + str(self.config.config["port"]), ] if self.config.config["enable_request_logging"]: sys.argv.append("--enable-request-logging") - + # Exit TUI - main.py will continue execution self.running = False diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 258a69f3..6b2c75d2 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -10,10 +10,18 @@ # --- Argument Parsing (BEFORE heavy imports) --- parser = argparse.ArgumentParser(description="API Key Proxy Server") -parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind the server to.") +parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Host to bind the server to." +) parser.add_argument("--port", type=int, default=8000, help="Port to run the server on.") -parser.add_argument("--enable-request-logging", action="store_true", help="Enable request logging.") -parser.add_argument("--add-credential", action="store_true", help="Launch the interactive tool to add a new OAuth credential.") +parser.add_argument( + "--enable-request-logging", action="store_true", help="Enable request logging." +) +parser.add_argument( + "--add-credential", + action="store_true", + help="Launch the interactive tool to add a new OAuth credential.", +) args, _ = parser.parse_known_args() # Add the 'src' directory to the Python path @@ -23,6 +31,7 @@ if len(sys.argv) == 1: # TUI MODE - Load ONLY what's needed for the launcher (fast path!) from proxy_app.launcher_tui import run_launcher_tui + run_launcher_tui() # Launcher modifies sys.argv and returns, or exits if user chose Exit # If we get here, user chose "Run Proxy" and sys.argv is modified @@ -32,6 +41,7 @@ # Check if credential tool mode (also doesn't need heavy proxy imports) if args.add_credential: from rotator_library.credential_tool import run_credential_tool + run_credential_tool() sys.exit(0) @@ -74,6 +84,7 @@ # Phase 2: Load Rich for loading spinner (lightweight) from rich.console import Console + _console = Console() # Phase 3: Heavy dependencies with granular loading messages @@ -92,7 +103,7 @@ import json from typing import AsyncGenerator, Any, List, Optional, Union from pydantic import BaseModel, Field - + # --- Early Log Level Configuration --- logging.getLogger("LiteLLM").setLevel(logging.WARNING) @@ -100,7 +111,7 @@ with _console.status("[dim]Loading LiteLLM library...", spinner="dots"): import litellm -# Phase 4: Application imports with granular loading messages +# Phase 4: Application imports with granular loading messages print(" → Initializing proxy core...") with _console.status("[dim]Initializing proxy core...", spinner="dots"): from rotator_library import RotatingClient @@ -115,12 +126,15 @@ # Provider lazy loading happens during import, so time it here _provider_start = time.time() with _console.status("[dim]Discovering provider plugins...", spinner="dots"): - from rotator_library import PROVIDER_PLUGINS # This triggers lazy load via __getattr__ + from rotator_library import ( + PROVIDER_PLUGINS, + ) # This triggers lazy load via __getattr__ _provider_time = time.time() - _provider_start # Get count after import (without timing to avoid double-counting) _plugin_count = len(PROVIDER_PLUGINS) + # --- Pydantic Models --- class EmbeddingRequest(BaseModel): model: str @@ -129,15 +143,19 @@ class EmbeddingRequest(BaseModel): dimensions: Optional[int] = None user: Optional[str] = None + class ModelCard(BaseModel): """Basic model card for minimal response.""" + id: str object: str = "model" created: int = Field(default_factory=lambda: int(time.time())) owned_by: str = "Mirro-Proxy" + class ModelCapabilities(BaseModel): """Model capability flags.""" + tool_choice: bool = False function_calling: bool = False reasoning: bool = False @@ -146,8 +164,10 @@ class ModelCapabilities(BaseModel): prompt_caching: bool = False assistant_prefill: bool = False + class EnrichedModelCard(BaseModel): """Extended model card with pricing and capabilities.""" + id: str object: str = "model" created: int = Field(default_factory=lambda: int(time.time())) @@ -169,28 +189,36 @@ class EnrichedModelCard(BaseModel): # Debug info (optional) _sources: Optional[List[str]] = None _match_type: Optional[str] = None - + class Config: extra = "allow" # Allow extra fields from the service + class ModelList(BaseModel): """List of models response.""" + object: str = "list" data: List[ModelCard] + class EnrichedModelList(BaseModel): """List of enriched models with pricing and capabilities.""" + object: str = "list" data: List[EnrichedModelCard] + # Calculate total loading time _elapsed = time.time() - _start_time -print(f"✓ Server ready in {_elapsed:.2f}s ({_plugin_count} providers discovered in {_provider_time:.2f}s)") +print( + f"✓ Server ready in {_elapsed:.2f}s ({_plugin_count} providers discovered in {_provider_time:.2f}s)" +) # Clear screen and reprint header for clean startup view # This pushes loading messages up (still in scroll history) but shows a clean final screen import os as _os_module -_os_module.system('cls' if _os_module.name == 'nt' else 'clear') + +_os_module.system("cls" if _os_module.name == "nt" else "clear") # Reprint header print("━" * 70) @@ -198,7 +226,9 @@ class EnrichedModelList(BaseModel): print(f"Proxy API Key: {key_display}") print(f"GitHub: https://github.com/Mirrowel/LLM-API-Key-Proxy") print("━" * 70) -print(f"✓ Server ready in {_elapsed:.2f}s ({_plugin_count} providers discovered in {_provider_time:.2f}s)") +print( + f"✓ Server ready in {_elapsed:.2f}s ({_plugin_count} providers discovered in {_provider_time:.2f}s)" +) # Note: Debug logging will be added after logging configuration below @@ -211,52 +241,64 @@ class EnrichedModelList(BaseModel): console_handler = colorlog.StreamHandler(sys.stdout) console_handler.setLevel(logging.INFO) formatter = colorlog.ColoredFormatter( - '%(log_color)s%(message)s', + "%(log_color)s%(message)s", log_colors={ - 'DEBUG': 'cyan', - 'INFO': 'green', - 'WARNING': 'yellow', - 'ERROR': 'red', - 'CRITICAL': 'red,bg_white', - } + "DEBUG": "cyan", + "INFO": "green", + "WARNING": "yellow", + "ERROR": "red", + "CRITICAL": "red,bg_white", + }, ) console_handler.setFormatter(formatter) # Configure a file handler for INFO-level logs and higher info_file_handler = logging.FileHandler(LOG_DIR / "proxy.log", encoding="utf-8") info_file_handler.setLevel(logging.INFO) -info_file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) +info_file_handler.setFormatter( + logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") +) # Configure a dedicated file handler for all DEBUG-level logs debug_file_handler = logging.FileHandler(LOG_DIR / "proxy_debug.log", encoding="utf-8") debug_file_handler.setLevel(logging.DEBUG) -debug_file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) +debug_file_handler.setFormatter( + logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") +) + # Create a filter to ensure the debug handler ONLY gets DEBUG messages from the rotator_library class RotatorDebugFilter(logging.Filter): def filter(self, record): - return record.levelno == logging.DEBUG and record.name.startswith('rotator_library') + return record.levelno == logging.DEBUG and record.name.startswith( + "rotator_library" + ) + + debug_file_handler.addFilter(RotatorDebugFilter()) # Configure a console handler with color console_handler = colorlog.StreamHandler(sys.stdout) console_handler.setLevel(logging.INFO) formatter = colorlog.ColoredFormatter( - '%(log_color)s%(message)s', + "%(log_color)s%(message)s", log_colors={ - 'DEBUG': 'cyan', - 'INFO': 'green', - 'WARNING': 'yellow', - 'ERROR': 'red', - 'CRITICAL': 'red,bg_white', - } + "DEBUG": "cyan", + "INFO": "green", + "WARNING": "yellow", + "ERROR": "red", + "CRITICAL": "red,bg_white", + }, ) console_handler.setFormatter(formatter) + # Add a filter to prevent any LiteLLM logs from cluttering the console class NoLiteLLMLogFilter(logging.Filter): def filter(self, record): - return not record.name.startswith('LiteLLM') + return not record.name.startswith("LiteLLM") + + console_handler.addFilter(NoLiteLLMLogFilter()) # Get the root logger and set it to DEBUG to capture all messages @@ -306,18 +348,26 @@ def filter(self, record): for key, value in os.environ.items(): if key.startswith("IGNORE_MODELS_"): provider = key.replace("IGNORE_MODELS_", "").lower() - models_to_ignore = [model.strip() for model in value.split(',') if model.strip()] + models_to_ignore = [ + model.strip() for model in value.split(",") if model.strip() + ] ignore_models[provider] = models_to_ignore - logging.debug(f"Loaded ignore list for provider '{provider}': {models_to_ignore}") + logging.debug( + f"Loaded ignore list for provider '{provider}': {models_to_ignore}" + ) # Load model whitelist from environment variables whitelist_models = {} for key, value in os.environ.items(): if key.startswith("WHITELIST_MODELS_"): provider = key.replace("WHITELIST_MODELS_", "").lower() - models_to_whitelist = [model.strip() for model in value.split(',') if model.strip()] + models_to_whitelist = [ + model.strip() for model in value.split(",") if model.strip() + ] whitelist_models[provider] = models_to_whitelist - logging.debug(f"Loaded whitelist for provider '{provider}': {models_to_whitelist}") + logging.debug( + f"Loaded whitelist for provider '{provider}': {models_to_whitelist}" + ) # Load max concurrent requests per key from environment variables max_concurrent_requests_per_key = {} @@ -327,12 +377,19 @@ def filter(self, record): try: max_concurrent = int(value) if max_concurrent < 1: - logging.warning(f"Invalid max_concurrent value for provider '{provider}': {value}. Must be >= 1. Using default (1).") + logging.warning( + f"Invalid max_concurrent value for provider '{provider}': {value}. Must be >= 1. Using default (1)." + ) max_concurrent = 1 max_concurrent_requests_per_key[provider] = max_concurrent - logging.debug(f"Loaded max concurrent requests for provider '{provider}': {max_concurrent}") + logging.debug( + f"Loaded max concurrent requests for provider '{provider}': {max_concurrent}" + ) except ValueError: - logging.warning(f"Invalid max_concurrent value for provider '{provider}': {value}. Using default (1).") + logging.warning( + f"Invalid max_concurrent value for provider '{provider}': {value}. Using default (1)." + ) + # --- Lifespan Management --- @asynccontextmanager @@ -349,11 +406,11 @@ async def lifespan(app: FastAPI): if not skip_oauth_init and oauth_credentials: logging.info("Starting OAuth credential validation and deduplication...") processed_emails = {} # email -> {provider: path} - credentials_to_initialize = {} # provider -> [paths] + credentials_to_initialize = {} # provider -> [paths] final_oauth_credentials = {} # --- Pass 1: Pre-initialization Scan & Deduplication --- - #logging.info("Pass 1: Scanning for existing metadata to find duplicates...") + # logging.info("Pass 1: Scanning for existing metadata to find duplicates...") for provider, paths in oauth_credentials.items(): if provider not in credentials_to_initialize: credentials_to_initialize[provider] = [] @@ -362,9 +419,9 @@ async def lifespan(app: FastAPI): if path.startswith("env://"): credentials_to_initialize[provider].append(path) continue - + try: - with open(path, 'r') as f: + with open(path, "r") as f: data = json.load(f) metadata = data.get("_proxy_metadata", {}) email = metadata.get("email") @@ -372,28 +429,32 @@ async def lifespan(app: FastAPI): if email: if email not in processed_emails: processed_emails[email] = {} - + if provider in processed_emails[email]: original_path = processed_emails[email][provider] - logging.warning(f"Duplicate for '{email}' on '{provider}' found in pre-scan: '{Path(path).name}'. Original: '{Path(original_path).name}'. Skipping.") + logging.warning( + f"Duplicate for '{email}' on '{provider}' found in pre-scan: '{Path(path).name}'. Original: '{Path(original_path).name}'. Skipping." + ) continue else: processed_emails[email][provider] = path - + credentials_to_initialize[provider].append(path) except (FileNotFoundError, json.JSONDecodeError) as e: - logging.warning(f"Could not pre-read metadata from '{path}': {e}. Will process during initialization.") + logging.warning( + f"Could not pre-read metadata from '{path}': {e}. Will process during initialization." + ) credentials_to_initialize[provider].append(path) - + # --- Pass 2: Parallel Initialization of Filtered Credentials --- - #logging.info("Pass 2: Initializing unique credentials and performing final check...") + # logging.info("Pass 2: Initializing unique credentials and performing final check...") async def process_credential(provider: str, path: str, provider_instance): """Process a single credential: initialize and fetch user info.""" try: await provider_instance.initialize_token(path) - if not hasattr(provider_instance, 'get_user_info'): + if not hasattr(provider_instance, "get_user_info"): return (provider, path, None, None) user_info = await provider_instance.get_user_info(path) @@ -401,7 +462,9 @@ async def process_credential(provider: str, path: str, provider_instance): return (provider, path, email, None) except Exception as e: - logging.error(f"Failed to process OAuth token for {provider} at '{path}': {e}") + logging.error( + f"Failed to process OAuth token for {provider} at '{path}': {e}" + ) return (provider, path, None, e) # Collect all tasks for parallel execution @@ -413,9 +476,9 @@ async def process_credential(provider: str, path: str, provider_instance): provider_plugin_class = PROVIDER_PLUGINS.get(provider) if not provider_plugin_class: continue - + provider_instance = provider_plugin_class() - + for path in paths: tasks.append(process_credential(provider, path, provider_instance)) @@ -430,7 +493,7 @@ async def process_credential(provider: str, path: str, provider_instance): continue provider, path, email, error = result - + # Skip if there was an error if error: continue @@ -444,7 +507,9 @@ async def process_credential(provider: str, path: str, provider_instance): # Handle empty email if not email: - logging.warning(f"Could not retrieve email for '{path}'. Treating as unique.") + logging.warning( + f"Could not retrieve email for '{path}'. Treating as unique." + ) if provider not in final_oauth_credentials: final_oauth_credentials[provider] = [] final_oauth_credentials[provider].append(path) @@ -453,10 +518,15 @@ async def process_credential(provider: str, path: str, provider_instance): # Deduplication check if email not in processed_emails: processed_emails[email] = {} - - if provider in processed_emails[email] and processed_emails[email][provider] != path: + + if ( + provider in processed_emails[email] + and processed_emails[email][provider] != path + ): original_path = processed_emails[email][provider] - logging.warning(f"Duplicate for '{email}' on '{provider}' found post-init: '{Path(path).name}'. Original: '{Path(original_path).name}'. Skipping.") + logging.warning( + f"Duplicate for '{email}' on '{provider}' found post-init: '{Path(path).name}'. Original: '{Path(original_path).name}'. Skipping." + ) continue else: processed_emails[email][provider] = path @@ -467,7 +537,7 @@ async def process_credential(provider: str, path: str, provider_instance): # Update metadata (skip for env-based credentials - they don't have files) if not path.startswith("env://"): try: - with open(path, 'r+') as f: + with open(path, "r+") as f: data = json.load(f) metadata = data.get("_proxy_metadata", {}) metadata["email"] = email @@ -490,33 +560,47 @@ async def process_credential(provider: str, path: str, provider_instance): # The client now uses the root logger configuration client = RotatingClient( api_keys=api_keys, - oauth_credentials=oauth_credentials, # Pass OAuth config + oauth_credentials=oauth_credentials, # Pass OAuth config configure_logging=True, litellm_provider_params=litellm_provider_params, ignore_models=ignore_models, whitelist_models=whitelist_models, enable_request_logging=ENABLE_REQUEST_LOGGING, - max_concurrent_requests_per_key=max_concurrent_requests_per_key + max_concurrent_requests_per_key=max_concurrent_requests_per_key, ) - + # Log loaded credentials summary (compact, always visible for deployment verification) - _api_summary = ', '.join([f"{p}:{len(c)}" for p, c in api_keys.items()]) if api_keys else "none" - _oauth_summary = ', '.join([f"{p}:{len(c)}" for p, c in oauth_credentials.items()]) if oauth_credentials else "none" - _total_summary = ', '.join([f"{p}:{len(c)}" for p, c in client.all_credentials.items()]) - print(f"🔑 Credentials loaded: {_total_summary} (API: {_api_summary} | OAuth: {_oauth_summary})") - client.background_refresher.start() # Start the background task + _api_summary = ( + ", ".join([f"{p}:{len(c)}" for p, c in api_keys.items()]) + if api_keys + else "none" + ) + _oauth_summary = ( + ", ".join([f"{p}:{len(c)}" for p, c in oauth_credentials.items()]) + if oauth_credentials + else "none" + ) + _total_summary = ", ".join( + [f"{p}:{len(c)}" for p, c in client.all_credentials.items()] + ) + print( + f"🔑 Credentials loaded: {_total_summary} (API: {_api_summary} | OAuth: {_oauth_summary})" + ) + client.background_refresher.start() # Start the background task app.state.rotating_client = client - + # Warn if no provider credentials are configured if not client.all_credentials: logging.warning("=" * 70) logging.warning("⚠️ NO PROVIDER CREDENTIALS CONFIGURED") logging.warning("The proxy is running but cannot serve any LLM requests.") - logging.warning("Launch the credential tool to add API keys or OAuth credentials.") + logging.warning( + "Launch the credential tool to add API keys or OAuth credentials." + ) logging.warning(" • Executable: Run with --add-credential flag") logging.warning(" • Source: python src/proxy_app/main.py --add-credential") logging.warning("=" * 70) - + os.environ["LITELLM_LOG"] = "ERROR" litellm.set_verbose = False litellm.drop_params = True @@ -527,29 +611,30 @@ async def process_credential(provider: str, path: str, provider_instance): else: app.state.embedding_batcher = None logging.info("RotatingClient initialized (EmbeddingBatcher disabled).") - + # Start model info service in background (fetches pricing/capabilities data) # This runs asynchronously and doesn't block proxy startup model_info_service = await init_model_info_service() app.state.model_info_service = model_info_service logging.info("Model info service started (fetching pricing data in background).") - + yield - - await client.background_refresher.stop() # Stop the background task on shutdown + + await client.background_refresher.stop() # Stop the background task on shutdown if app.state.embedding_batcher: await app.state.embedding_batcher.stop() await client.close() - + # Stop model info service - if hasattr(app.state, 'model_info_service') and app.state.model_info_service: + if hasattr(app.state, "model_info_service") and app.state.model_info_service: await app.state.model_info_service.stop() - + if app.state.embedding_batcher: logging.info("RotatingClient and EmbeddingBatcher closed.") else: logging.info("RotatingClient closed.") + # --- FastAPI App Setup --- app = FastAPI(lifespan=lifespan) @@ -563,25 +648,32 @@ async def process_credential(provider: str, path: str, provider_instance): ) api_key_header = APIKeyHeader(name="Authorization", auto_error=False) + def get_rotating_client(request: Request) -> RotatingClient: """Dependency to get the rotating client instance from the app state.""" return request.app.state.rotating_client + def get_embedding_batcher(request: Request) -> EmbeddingBatcher: """Dependency to get the embedding batcher instance from the app state.""" return request.app.state.embedding_batcher + async def verify_api_key(auth: str = Depends(api_key_header)): """Dependency to verify the proxy API key.""" + # If PROXY_API_KEY is not set or empty, skip verification (open access) + if not PROXY_API_KEY: + return auth if not auth or auth != f"Bearer {PROXY_API_KEY}": raise HTTPException(status_code=401, detail="Invalid or missing API Key") return auth + async def streaming_response_wrapper( request: Request, request_data: dict, response_stream: AsyncGenerator[str, None], - logger: Optional[DetailedLogger] = None + logger: Optional[DetailedLogger] = None, ) -> AsyncGenerator[str, None]: """ Wraps a streaming response to log the full response after completion @@ -589,7 +681,7 @@ async def streaming_response_wrapper( """ response_chunks = [] full_response = {} - + try: async for chunk_str in response_stream: if await request.is_disconnected(): @@ -597,7 +689,7 @@ async def streaming_response_wrapper( break yield chunk_str if chunk_str.strip() and chunk_str.startswith("data:"): - content = chunk_str[len("data:"):].strip() + content = chunk_str[len("data:") :].strip() if content != "[DONE]": try: chunk_data = json.loads(content) @@ -613,15 +705,17 @@ async def streaming_response_wrapper( "error": { "message": f"An unexpected error occurred during the stream: {str(e)}", "type": "proxy_internal_error", - "code": 500 + "code": 500, } } yield f"data: {json.dumps(error_payload)}\n\n" yield "data: [DONE]\n\n" # Also log this as a failed request if logger: - logger.log_final_response(status_code=500, headers=None, body={"error": str(e)}) - return # Stop further processing + logger.log_final_response( + status_code=500, headers=None, body={"error": str(e)} + ) + return # Stop further processing finally: if response_chunks: # --- Aggregation Logic --- @@ -645,36 +739,56 @@ async def streaming_response_wrapper( final_message["content"] = "" if value: final_message["content"] += value - + elif key == "tool_calls": for tc_chunk in value: index = tc_chunk["index"] if index not in aggregated_tool_calls: - aggregated_tool_calls[index] = {"type": "function", "function": {"name": "", "arguments": ""}} + aggregated_tool_calls[index] = { + "type": "function", + "function": {"name": "", "arguments": ""}, + } # Ensure 'function' key exists for this index before accessing its sub-keys if "function" not in aggregated_tool_calls[index]: - aggregated_tool_calls[index]["function"] = {"name": "", "arguments": ""} + aggregated_tool_calls[index]["function"] = { + "name": "", + "arguments": "", + } if tc_chunk.get("id"): aggregated_tool_calls[index]["id"] = tc_chunk["id"] if "function" in tc_chunk: if "name" in tc_chunk["function"]: if tc_chunk["function"]["name"] is not None: - aggregated_tool_calls[index]["function"]["name"] += tc_chunk["function"]["name"] + aggregated_tool_calls[index]["function"][ + "name" + ] += tc_chunk["function"]["name"] if "arguments" in tc_chunk["function"]: - if tc_chunk["function"]["arguments"] is not None: - aggregated_tool_calls[index]["function"]["arguments"] += tc_chunk["function"]["arguments"] - + if ( + tc_chunk["function"]["arguments"] + is not None + ): + aggregated_tool_calls[index]["function"][ + "arguments" + ] += tc_chunk["function"]["arguments"] + elif key == "function_call": if "function_call" not in final_message: - final_message["function_call"] = {"name": "", "arguments": ""} + final_message["function_call"] = { + "name": "", + "arguments": "", + } if "name" in value: if value["name"] is not None: - final_message["function_call"]["name"] += value["name"] + final_message["function_call"]["name"] += value[ + "name" + ] if "arguments" in value: if value["arguments"] is not None: - final_message["function_call"]["arguments"] += value["arguments"] - - else: # Generic key handling for other data like 'reasoning' + final_message["function_call"]["arguments"] += ( + value["arguments"] + ) + + else: # Generic key handling for other data like 'reasoning' # FIX: Role should always replace, never concatenate if key == "role": final_message[key] = value @@ -707,7 +821,7 @@ async def streaming_response_wrapper( final_choice = { "index": 0, "message": final_message, - "finish_reason": finish_reason + "finish_reason": finish_reason, } full_response = { @@ -716,21 +830,22 @@ async def streaming_response_wrapper( "created": first_chunk.get("created"), "model": first_chunk.get("model"), "choices": [final_choice], - "usage": usage_data + "usage": usage_data, } if logger: logger.log_final_response( status_code=200, headers=None, # Headers are not available at this stage - body=full_response + body=full_response, ) + @app.post("/v1/chat/completions") async def chat_completions( request: Request, client: RotatingClient = Depends(get_rotating_client), - _ = Depends(verify_api_key) + _=Depends(verify_api_key), ): """ OpenAI-compatible endpoint powered by the RotatingClient. @@ -749,16 +864,24 @@ async def chat_completions( # instead of actual schemas, which can cause tool hallucination # Modes: "remove" = delete temperature key, "set" = change to 1.0, "false" = disabled override_temp_zero = os.getenv("OVERRIDE_TEMPERATURE_ZERO", "false").lower() - - if override_temp_zero in ("remove", "set", "true", "1", "yes") and "temperature" in request_data and request_data["temperature"] == 0: + + if ( + override_temp_zero in ("remove", "set", "true", "1", "yes") + and "temperature" in request_data + and request_data["temperature"] == 0 + ): if override_temp_zero == "remove": # Remove temperature key entirely del request_data["temperature"] - logging.debug("OVERRIDE_TEMPERATURE_ZERO=remove: Removed temperature=0 from request") + logging.debug( + "OVERRIDE_TEMPERATURE_ZERO=remove: Removed temperature=0 from request" + ) else: # Set to 1.0 (for "set", "true", "1", "yes") request_data["temperature"] = 1.0 - logging.debug("OVERRIDE_TEMPERATURE_ZERO=set: Converting temperature=0 to temperature=1.0") + logging.debug( + "OVERRIDE_TEMPERATURE_ZERO=set: Converting temperature=0 to temperature=1.0" + ) # If logging is enabled, perform all logging operations using the parsed data. if logger: @@ -766,9 +889,17 @@ async def chat_completions( # Extract and log specific reasoning parameters for monitoring. model = request_data.get("model") - generation_cfg = request_data.get("generationConfig", {}) or request_data.get("generation_config", {}) or {} - reasoning_effort = request_data.get("reasoning_effort") or generation_cfg.get("reasoning_effort") - custom_reasoning_budget = request_data.get("custom_reasoning_budget") or generation_cfg.get("custom_reasoning_budget", False) + generation_cfg = ( + request_data.get("generationConfig", {}) + or request_data.get("generation_config", {}) + or {} + ) + reasoning_effort = request_data.get("reasoning_effort") or generation_cfg.get( + "reasoning_effort" + ) + custom_reasoning_budget = request_data.get( + "custom_reasoning_budget" + ) or generation_cfg.get("custom_reasoning_budget", False) logging.getLogger("rotator_library").debug( f"Handling reasoning parameters: model={model}, reasoning_effort={reasoning_effort}, custom_reasoning_budget={custom_reasoning_budget}" @@ -779,31 +910,41 @@ async def chat_completions( url=str(request.url), headers=dict(request.headers), client_info=(request.client.host, request.client.port), - request_data=request_data + request_data=request_data, ) is_streaming = request_data.get("stream", False) if is_streaming: response_generator = client.acompletion(request=request, **request_data) return StreamingResponse( - streaming_response_wrapper(request, request_data, response_generator, logger), - media_type="text/event-stream" + streaming_response_wrapper( + request, request_data, response_generator, logger + ), + media_type="text/event-stream", ) else: response = await client.acompletion(request=request, **request_data) if logger: # Assuming response has status_code and headers attributes # This might need adjustment based on the actual response object - response_headers = response.headers if hasattr(response, 'headers') else None - status_code = response.status_code if hasattr(response, 'status_code') else 200 + response_headers = ( + response.headers if hasattr(response, "headers") else None + ) + status_code = ( + response.status_code if hasattr(response, "status_code") else 200 + ) logger.log_final_response( status_code=status_code, headers=response_headers, - body=response.model_dump() + body=response.model_dump(), ) return response - except (litellm.InvalidRequestError, ValueError, litellm.ContextWindowExceededError) as e: + except ( + litellm.InvalidRequestError, + ValueError, + litellm.ContextWindowExceededError, + ) as e: raise HTTPException(status_code=400, detail=f"Invalid Request: {str(e)}") except litellm.AuthenticationError as e: raise HTTPException(status_code=401, detail=f"Authentication Error: {str(e)}") @@ -824,16 +965,19 @@ async def chat_completions( except json.JSONDecodeError: request_data = {"error": "Could not parse request body"} if logger: - logger.log_final_response(status_code=500, headers=None, body={"error": str(e)}) + logger.log_final_response( + status_code=500, headers=None, body={"error": str(e)} + ) raise HTTPException(status_code=500, detail=str(e)) + @app.post("/v1/embeddings") async def embeddings( request: Request, body: EmbeddingRequest, client: RotatingClient = Depends(get_rotating_client), batcher: Optional[EmbeddingBatcher] = Depends(get_embedding_batcher), - _ = Depends(verify_api_key) + _=Depends(verify_api_key), ): """ OpenAI-compatible endpoint for creating embeddings. @@ -847,7 +991,7 @@ async def embeddings( url=str(request.url), headers=dict(request.headers), client_info=(request.client.host, request.client.port), - request_data=request_data + request_data=request_data, ) if USE_EMBEDDING_BATCHER and batcher: # --- Server-Side Batching Logic --- @@ -861,7 +1005,7 @@ async def embeddings( individual_request = request_data.copy() individual_request["input"] = single_input tasks.append(batcher.add_request(individual_request)) - + results = await asyncio.gather(*tasks) all_data = [] @@ -877,16 +1021,19 @@ async def embeddings( "object": "list", "model": results[0]["model"], "data": all_data, - "usage": { "prompt_tokens": total_prompt_tokens, "total_tokens": total_tokens }, + "usage": { + "prompt_tokens": total_prompt_tokens, + "total_tokens": total_tokens, + }, } response = litellm.EmbeddingResponse(**final_response_data) - + else: # --- Direct Pass-Through Logic --- request_data = body.model_dump(exclude_none=True) if isinstance(request_data.get("input"), str): request_data["input"] = [request_data["input"]] - + response = await client.aembedding(request=request, **request_data) return response @@ -894,7 +1041,11 @@ async def embeddings( except HTTPException as e: # Re-raise HTTPException to ensure it's not caught by the generic Exception handler raise e - except (litellm.InvalidRequestError, ValueError, litellm.ContextWindowExceededError) as e: + except ( + litellm.InvalidRequestError, + ValueError, + litellm.ContextWindowExceededError, + ) as e: raise HTTPException(status_code=400, detail=f"Invalid Request: {str(e)}") except litellm.AuthenticationError as e: raise HTTPException(status_code=401, detail=f"Authentication Error: {str(e)}") @@ -910,10 +1061,12 @@ async def embeddings( logging.error(f"Embedding request failed: {e}") raise HTTPException(status_code=500, detail=str(e)) + @app.get("/") def read_root(): return {"Status": "API Key Proxy is running"} + @app.get("/v1/models") async def list_models( request: Request, @@ -923,22 +1076,30 @@ async def list_models( ): """ Returns a list of available models in the OpenAI-compatible format. - + Query Parameters: enriched: If True (default), returns detailed model info with pricing and capabilities. If False, returns minimal OpenAI-compatible response. """ model_ids = await client.get_all_available_models(grouped=False) - - if enriched and hasattr(request.app.state, 'model_info_service'): + + if enriched and hasattr(request.app.state, "model_info_service"): model_info_service = request.app.state.model_info_service if model_info_service.is_ready: # Return enriched model data enriched_data = model_info_service.enrich_model_list(model_ids) return {"object": "list", "data": enriched_data} - + # Fallback to basic model cards - model_cards = [{"id": model_id, "object": "model", "created": int(time.time()), "owned_by": "Mirro-Proxy"} for model_id in model_ids] + model_cards = [ + { + "id": model_id, + "object": "model", + "created": int(time.time()), + "owned_by": "Mirro-Proxy", + } + for model_id in model_ids + ] return {"object": "list", "data": model_cards} @@ -950,17 +1111,17 @@ async def get_model( ): """ Returns detailed information about a specific model. - + Path Parameters: model_id: The model ID (e.g., "anthropic/claude-3-opus", "openrouter/openai/gpt-4") """ - if hasattr(request.app.state, 'model_info_service'): + if hasattr(request.app.state, "model_info_service"): model_info_service = request.app.state.model_info_service if model_info_service.is_ready: info = model_info_service.get_model_info(model_id) if info: return info.to_dict() - + # Return basic info if service not ready or model not found return { "id": model_id, @@ -978,7 +1139,7 @@ async def model_info_stats( """ Returns statistics about the model info service (for monitoring/debugging). """ - if hasattr(request.app.state, 'model_info_service'): + if hasattr(request.app.state, "model_info_service"): return request.app.state.model_info_service.get_stats() return {"error": "Model info service not initialized"} @@ -990,11 +1151,12 @@ async def list_providers(_=Depends(verify_api_key)): """ return list(PROVIDER_PLUGINS.keys()) + @app.post("/v1/token-count") async def token_count( - request: Request, + request: Request, client: RotatingClient = Depends(get_rotating_client), - _=Depends(verify_api_key) + _=Depends(verify_api_key), ): """ Calculates the token count for a given list of messages and a model. @@ -1005,7 +1167,9 @@ async def token_count( messages = data.get("messages") if not model or not messages: - raise HTTPException(status_code=400, detail="'model' and 'messages' are required.") + raise HTTPException( + status_code=400, detail="'model' and 'messages' are required." + ) count = client.token_count(**data) return {"token_count": count} @@ -1016,13 +1180,10 @@ async def token_count( @app.post("/v1/cost-estimate") -async def cost_estimate( - request: Request, - _=Depends(verify_api_key) -): +async def cost_estimate(request: Request, _=Depends(verify_api_key)): """ Estimates the cost for a request based on token counts and model pricing. - + Request body: { "model": "anthropic/claude-3-opus", @@ -1031,7 +1192,7 @@ async def cost_estimate( "cache_read_tokens": 0, # optional "cache_creation_tokens": 0 # optional } - + Returns: { "model": "anthropic/claude-3-opus", @@ -1051,25 +1212,28 @@ async def cost_estimate( completion_tokens = data.get("completion_tokens", 0) cache_read_tokens = data.get("cache_read_tokens", 0) cache_creation_tokens = data.get("cache_creation_tokens", 0) - + if not model: raise HTTPException(status_code=400, detail="'model' is required.") - + result = { "model": model, "cost": None, "currency": "USD", "pricing": {}, - "source": None + "source": None, } - + # Try model info service first - if hasattr(request.app.state, 'model_info_service'): + if hasattr(request.app.state, "model_info_service"): model_info_service = request.app.state.model_info_service if model_info_service.is_ready: cost = model_info_service.calculate_cost( - model, prompt_tokens, completion_tokens, - cache_read_tokens, cache_creation_tokens + model, + prompt_tokens, + completion_tokens, + cache_read_tokens, + cache_creation_tokens, ) if cost is not None: cost_info = model_info_service.get_cost_info(model) @@ -1077,31 +1241,32 @@ async def cost_estimate( result["pricing"] = cost_info or {} result["source"] = "model_info_service" return result - + # Fallback to litellm try: import litellm + # Create a mock response for cost calculation model_info = litellm.get_model_info(model) input_cost = model_info.get("input_cost_per_token", 0) output_cost = model_info.get("output_cost_per_token", 0) - + if input_cost or output_cost: cost = (prompt_tokens * input_cost) + (completion_tokens * output_cost) result["cost"] = cost result["pricing"] = { "input_cost_per_token": input_cost, - "output_cost_per_token": output_cost + "output_cost_per_token": output_cost, } result["source"] = "litellm_fallback" return result except Exception: pass - + result["source"] = "unknown" result["error"] = "Pricing data not available for this model" return result - + except HTTPException: raise except Exception as e: @@ -1112,17 +1277,18 @@ async def cost_estimate( if __name__ == "__main__": # Define ENV_FILE for onboarding checks ENV_FILE = Path.cwd() / ".env" - + # Check if launcher TUI should be shown (no arguments provided) if len(sys.argv) == 1: # No arguments - show launcher TUI (lazy import) from proxy_app.launcher_tui import run_launcher_tui + run_launcher_tui() # Launcher modifies sys.argv and returns, or exits if user chose Exit # If we get here, user chose "Run Proxy" and sys.argv is modified # Re-parse arguments with modified sys.argv args = parser.parse_args() - + def needs_onboarding() -> bool: """ Check if the proxy needs onboarding (first-time setup). @@ -1132,40 +1298,49 @@ def needs_onboarding() -> bool: # PROXY_API_KEY is optional (will show warning if not set) if not ENV_FILE.is_file(): return True - + return False def show_onboarding_message(): """Display clear explanatory message for why onboarding is needed.""" - os.system('cls' if os.name == 'nt' else 'clear') # Clear terminal for clean presentation - console.print(Panel.fit( - "[bold cyan]🚀 LLM API Key Proxy - First Time Setup[/bold cyan]", - border_style="cyan" - )) + os.system( + "cls" if os.name == "nt" else "clear" + ) # Clear terminal for clean presentation + console.print( + Panel.fit( + "[bold cyan]🚀 LLM API Key Proxy - First Time Setup[/bold cyan]", + border_style="cyan", + ) + ) console.print("[bold yellow]⚠️ Configuration Required[/bold yellow]\n") - + console.print("The proxy needs initial configuration:") console.print(" [red]❌ No .env file found[/red]") - + console.print("\n[bold]Why this matters:[/bold]") console.print(" • The .env file stores your credentials and settings") console.print(" • PROXY_API_KEY protects your proxy from unauthorized access") console.print(" • Provider API keys enable LLM access") - + console.print("\n[bold]What happens next:[/bold]") console.print(" 1. We'll create a .env file with PROXY_API_KEY") console.print(" 2. You can add LLM provider credentials (API keys or OAuth)") console.print(" 3. The proxy will then start normally") - - console.print("\n[bold yellow]⚠️ Note:[/bold yellow] The credential tool adds PROXY_API_KEY by default.") + + console.print( + "\n[bold yellow]⚠️ Note:[/bold yellow] The credential tool adds PROXY_API_KEY by default." + ) console.print(" You can remove it later if you want an unsecured proxy.\n") - - console.input("[bold green]Press Enter to launch the credential setup tool...[/bold green]") + + console.input( + "[bold green]Press Enter to launch the credential setup tool...[/bold green]" + ) # Check if user explicitly wants to add credentials if args.add_credential: # Import and call ensure_env_defaults to create .env and PROXY_API_KEY if needed from rotator_library.credential_tool import ensure_env_defaults + ensure_env_defaults() # Reload environment variables after ensure_env_defaults creates/updates .env load_dotenv(override=True) @@ -1176,36 +1351,35 @@ def show_onboarding_message(): # Import console from rich for better messaging from rich.console import Console from rich.panel import Panel + console = Console() - + # Show clear explanatory message show_onboarding_message() - + # Launch credential tool automatically from rotator_library.credential_tool import ensure_env_defaults + ensure_env_defaults() load_dotenv(override=True) run_credential_tool() - + # After credential tool exits, reload and re-check load_dotenv(override=True) # Re-read PROXY_API_KEY from environment PROXY_API_KEY = os.getenv("PROXY_API_KEY") - + # Verify onboarding is complete if needs_onboarding(): console.print("\n[bold red]❌ Configuration incomplete.[/bold red]") - console.print("The proxy still cannot start. Please ensure PROXY_API_KEY is set in .env\n") + console.print( + "The proxy still cannot start. Please ensure PROXY_API_KEY is set in .env\n" + ) sys.exit(1) else: console.print("\n[bold green]✅ Configuration complete![/bold green]") console.print("\nStarting proxy server...\n") - - # Validate PROXY_API_KEY before starting the server - if not PROXY_API_KEY: - raise ValueError("PROXY_API_KEY environment variable not set. Please run with --add-credential to set up your environment.") - - import uvicorn - uvicorn.run(app, host=args.host, port=args.port) + import uvicorn + uvicorn.run(app, host=args.host, port=args.port) From a725feba53b661b2b203c00a25992530e6e4c25a Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 04:44:42 +0100 Subject: [PATCH 102/221] =?UTF-8?q?refactor(client):=20=F0=9F=94=A8=20add?= =?UTF-8?q?=20comprehensive=20error=20handling=20and=20retry=20logic=20for?= =?UTF-8?q?=20custom=20provider=20non-streaming=20calls?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This change brings the non-streaming custom provider call path in line with the streaming path's robust error handling strategy. - Implements a retry loop with attempt tracking and logging for custom provider calls - Adds pre-request callback execution with configurable error handling - Integrates error classification and rotation logic for rate limits, HTTP errors, and server errors - Records errors in the accumulator for client-level reporting and visibility - Implements exponential backoff with jitter for transient server errors - Adds cooldown management for rate-limited providers - Respects time budget constraints when calculating retry wait times - Properly manages credential state (success/failure recording and key release) - Distinguishes between recoverable errors (which trigger rotation) and non-recoverable errors (which fail immediately) The retry loop handles three categories of exceptions: 1. Rate limits and HTTP status errors: trigger immediate rotation after recording 2. Connection and server errors: retry with backoff, rotate only after max retries 3. General exceptions: classify and rotate if recoverable, fail if not --- src/rotator_library/client.py | 190 +++++++++++++++++++++++++++++++--- 1 file changed, 178 insertions(+), 12 deletions(-) diff --git a/src/rotator_library/client.py b/src/rotator_library/client.py index 6a3b8907..a220020e 100644 --- a/src/rotator_library/client.py +++ b/src/rotator_library/client.py @@ -1065,19 +1065,185 @@ async def _execute_with_retry( is_budget_enabled ) - # The plugin handles the entire call, including retries on 401, etc. - # The main retry loop here is for key rotation on other errors. - response = await provider_plugin.acompletion( - self.http_client, **litellm_kwargs - ) + # Retry loop for custom providers - mirrors streaming path error handling + for attempt in range(self.max_retries): + try: + lib_logger.info( + f"Attempting call with credential {mask_credential(current_cred)} (Attempt {attempt + 1}/{self.max_retries})" + ) - # For non-streaming, success is immediate, and this function only handles non-streaming. - await self.usage_manager.record_success( - current_cred, model, response - ) - await self.usage_manager.release_key(current_cred, model) - key_acquired = False - return response + if pre_request_callback: + try: + await pre_request_callback(request, litellm_kwargs) + except Exception as e: + if self.abort_on_callback_error: + raise PreRequestCallbackError( + f"Pre-request callback failed: {e}" + ) from e + else: + lib_logger.warning( + f"Pre-request callback failed but abort_on_callback_error is False. Proceeding with request. Error: {e}" + ) + + response = await provider_plugin.acompletion( + self.http_client, **litellm_kwargs + ) + + # For non-streaming, success is immediate + await self.usage_manager.record_success( + current_cred, model, response + ) + await self.usage_manager.release_key(current_cred, model) + key_acquired = False + return response + + except ( + litellm.RateLimitError, + httpx.HTTPStatusError, + ) as e: + last_exception = e + classified_error = classify_error(e, provider=provider) + error_message = str(e).split("\n")[0] + + log_failure( + api_key=current_cred, + model=model, + attempt=attempt + 1, + error=e, + request_headers=dict(request.headers) + if request + else {}, + ) + + # Record in accumulator for client reporting + error_accumulator.record_error( + current_cred, classified_error, error_message + ) + + # Check if this error should trigger rotation + if not should_rotate_on_error(classified_error): + lib_logger.error( + f"Non-recoverable error ({classified_error.error_type}) during custom provider call. Failing." + ) + raise last_exception + + # Handle rate limits with cooldown (exclude quota_exceeded) + if classified_error.error_type == "rate_limit": + cooldown_duration = classified_error.retry_after or 60 + await self.cooldown_manager.start_cooldown( + provider, cooldown_duration + ) + + await self.usage_manager.record_failure( + current_cred, model, classified_error + ) + lib_logger.warning( + f"Cred {mask_credential(current_cred)} {classified_error.error_type} (HTTP {classified_error.status_code}). Rotating." + ) + break # Rotate to next credential + + except ( + APIConnectionError, + litellm.InternalServerError, + litellm.ServiceUnavailableError, + ) as e: + last_exception = e + log_failure( + api_key=current_cred, + model=model, + attempt=attempt + 1, + error=e, + request_headers=dict(request.headers) + if request + else {}, + ) + classified_error = classify_error(e, provider=provider) + error_message = str(e).split("\n")[0] + + # Provider-level error: don't increment consecutive failures + await self.usage_manager.record_failure( + current_cred, + model, + classified_error, + increment_consecutive_failures=False, + ) + + if attempt >= self.max_retries - 1: + error_accumulator.record_error( + current_cred, classified_error, error_message + ) + lib_logger.warning( + f"Cred {mask_credential(current_cred)} failed after max retries. Rotating." + ) + break + + wait_time = classified_error.retry_after or ( + 2**attempt + ) + random.uniform(0, 1) + remaining_budget = deadline - time.time() + if wait_time > remaining_budget: + error_accumulator.record_error( + current_cred, classified_error, error_message + ) + lib_logger.warning( + f"Retry wait ({wait_time:.2f}s) exceeds budget. Rotating." + ) + break + + lib_logger.warning( + f"Cred {mask_credential(current_cred)} server error. Retrying in {wait_time:.2f}s." + ) + await asyncio.sleep(wait_time) + continue + + except Exception as e: + last_exception = e + log_failure( + api_key=current_cred, + model=model, + attempt=attempt + 1, + error=e, + request_headers=dict(request.headers) + if request + else {}, + ) + classified_error = classify_error(e, provider=provider) + error_message = str(e).split("\n")[0] + + # Record in accumulator + error_accumulator.record_error( + current_cred, classified_error, error_message + ) + + lib_logger.warning( + f"Cred {mask_credential(current_cred)} {classified_error.error_type} (HTTP {classified_error.status_code})." + ) + + # Check if this error should trigger rotation + if not should_rotate_on_error(classified_error): + lib_logger.error( + f"Non-recoverable error ({classified_error.error_type}). Failing." + ) + raise last_exception + + # Handle rate limits with cooldown (exclude quota_exceeded) + if ( + classified_error.status_code == 429 + and classified_error.error_type != "quota_exceeded" + ) or classified_error.error_type == "rate_limit": + cooldown_duration = classified_error.retry_after or 60 + await self.cooldown_manager.start_cooldown( + provider, cooldown_duration + ) + + await self.usage_manager.record_failure( + current_cred, model, classified_error + ) + break # Rotate to next credential + + # If the inner loop breaks, it means the key failed and we need to rotate. + # Continue to the next iteration of the outer while loop to pick a new key. + continue else: # This is the standard API Key / litellm-handled provider logic is_oauth = provider in self.oauth_providers From 640efbfedece68315031d94c6649a715a2310f19 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 04:50:48 +0100 Subject: [PATCH 103/221] fix(providers): disable endpoint in antigravity provider --- src/rotator_library/providers/antigravity_provider.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index a29a63ab..42109f52 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -50,7 +50,7 @@ # Priority: daily (sandbox) → autopush (sandbox) → production BASE_URLS = [ "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal", - "https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal", + #"https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal", "https://cloudcode-pa.googleapis.com/v1internal", # Production fallback ] From 73a2395fc7f8c8e35063cfa542e65c6b18b88c94 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 04:55:52 +0100 Subject: [PATCH 104/221] =?UTF-8?q?refactor(providers):=20=F0=9F=94=A8=20i?= =?UTF-8?q?mprove=20error=20handling=20and=20reduce=20debug=20logging=20in?= =?UTF-8?q?=20antigravity=20provider?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add specific handling for 429 HTTP status errors to prevent unnecessary fallback URL retries, as quota exhaustion is credential-bound - Separate HTTP errors from network errors in exception handling for more intelligent retry logic - Comment out verbose debug logging for function grouping operations to reduce noise - Fix code style formatting for commented URLs and quota group configuration - Enable claude model quota group for production use The changes improve the provider's resilience by distinguishing between errors that benefit from URL fallback (network issues, server errors) versus those that don't (rate limits). Debug log reduction improves terminal readability while maintaining error logging in failures.log. --- .../providers/antigravity_provider.py | 57 +++++++++++++------ 1 file changed, 40 insertions(+), 17 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 42109f52..ebf950ee 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -50,7 +50,7 @@ # Priority: daily (sandbox) → autopush (sandbox) → production BASE_URLS = [ "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal", - #"https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal", + # "https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal", "https://cloudcode-pa.googleapis.com/v1internal", # Production fallback ] @@ -541,7 +541,7 @@ class AntigravityProvider(AntigravityAuthBase, ProviderInterface): # Model quota groups (can be overridden via QUOTA_GROUPS_ANTIGRAVITY_CLAUDE) # Models in the same group share quota - when one is exhausted, all are model_quota_groups: QuotaGroupMap = { - #"claude": ["claude-sonnet-4-5", "claude-opus-4-5"], - commented out for later use if needed + "claude": ["claude-sonnet-4-5", "claude-opus-4-5"], } # Model usage weights for grouped usage calculation @@ -2559,9 +2559,9 @@ def _fix_tool_response_grouping( f"Ignoring duplicate - this may indicate malformed conversation history." ) continue - lib_logger.debug( - f"[Grouping] Collected response for ID: {resp_id}" - ) + #lib_logger.debug( + # f"[Grouping] Collected response for ID: {resp_id}" + #) collected_responses[resp_id] = resp # Try to satisfy pending groups (newest first) @@ -2576,10 +2576,10 @@ def _fix_tool_response_grouping( collected_responses.pop(gid) for gid in group_ids ] new_contents.append({"parts": group_responses, "role": "user"}) - lib_logger.debug( - f"[Grouping] Satisfied group with {len(group_responses)} responses: " - f"ids={group_ids}" - ) + #lib_logger.debug( + # f"[Grouping] Satisfied group with {len(group_responses)} responses: " + # f"ids={group_ids}" + #) pending_groups.pop(i) break continue @@ -2599,10 +2599,10 @@ def _fix_tool_response_grouping( ] if call_ids: - lib_logger.debug( - f"[Grouping] Created pending group expecting {len(call_ids)} responses: " - f"ids={call_ids}, names={func_names}" - ) + #lib_logger.debug( + # f"[Grouping] Created pending group expecting {len(call_ids)} responses: " + # f"ids={call_ids}, names={func_names}" + #) pending_groups.append( { "ids": call_ids, @@ -3634,7 +3634,28 @@ async def acompletion( return await self._handle_non_streaming( client, url, headers, payload, model, file_logger ) + except httpx.HTTPStatusError as e: + # 429 = Rate limit/quota exhausted - tied to credential, not URL + # Do NOT retry on different URL, just raise immediately + if e.response.status_code == 429: + lib_logger.debug(f"429 quota error - not retrying on fallback URL: {e}") + raise + + # For other HTTP errors (403, 500, etc.), try fallback URL + if self._try_next_base_url(): + lib_logger.warning(f"Retrying with fallback URL: {e}") + url = f"{self._get_base_url()}{endpoint}" + if stream: + return self._handle_streaming( + client, url, headers, payload, model, file_logger + ) + else: + return await self._handle_non_streaming( + client, url, headers, payload, model, file_logger + ) + raise except Exception as e: + # Non-HTTP errors (network issues, timeouts, etc.) - try fallback URL if self._try_next_base_url(): lib_logger.warning(f"Retrying with fallback URL: {e}") url = f"{self._get_base_url()}{endpoint}" @@ -3718,11 +3739,13 @@ async def _handle_streaming( "POST", url, headers=headers, json=payload, timeout=600.0 ) as response: if response.status_code >= 400: + # Read error body for raise_for_status to include in exception + # Terminal logging commented out - errors are logged in failures.log try: - error_body = await response.aread() - lib_logger.error( - f"API error {response.status_code}: {error_body.decode()}" - ) + await response.aread() + # lib_logger.error( + # f"API error {response.status_code}: {error_body.decode()}" + # ) except Exception: pass From 219a7a9dfb56633812f294d3d6e5a9a3d7206c24 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 06:57:42 +0100 Subject: [PATCH 105/221] =?UTF-8?q?feat(auth):=20=E2=9C=A8=20implement=20g?= =?UTF-8?q?lobal=20reauth=20coordinator=20to=20serialize=20interactive=20O?= =?UTF-8?q?Auth=20flows?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces a centralized ReauthCoordinator singleton that ensures only one interactive OAuth flow (across all providers: Google, iFlow, Qwen) executes at a time. This prevents port conflicts, reduces user confusion, and improves credential state management reliability. Key changes: - Add new `ReauthCoordinator` class with global semaphore-based serialization - Extract interactive OAuth logic into separate `_perform_interactive_oauth()` methods for each provider - Update `initialize_token()` methods to delegate to coordinator instead of running OAuth inline - Change `_unavailable_credentials` from set to dict with timestamps for TTL-based stale entry cleanup - Add comprehensive logging and statistics tracking for reauth operations - Update all providers (GoogleOAuthBase, IFlowAuthBase, QwenAuthBase) to use the coordinator - Add 300-second timeout for interactive flows with automatic cleanup on timeout/cancellation - Implement defense-in-depth with TTL-based cleanup (5 minutes) to prevent credentials from becoming permanently stuck The coordinator provides: - Queue management for pending reauth requests - Status tracking and observability (success/failure/timeout counts) - Graceful handling of timeouts, cancellations, and errors - Consistent cleanup in all exit paths (success, exception, timeout) Refs PR#34 --- .../providers/google_oauth_base.py | 402 ++++++++------- .../providers/iflow_auth_base.py | 348 ++++++++----- .../providers/qwen_auth_base.py | 476 +++++++++++------- src/rotator_library/utils/__init__.py | 3 +- .../utils/reauth_coordinator.py | 235 +++++++++ 5 files changed, 954 insertions(+), 510 deletions(-) create mode 100644 src/rotator_library/utils/reauth_coordinator.py diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index 96684ef4..68979cdf 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -19,6 +19,7 @@ from rich.markup import escape as rich_escape from ..utils.headless_detection import is_headless_environment +from ..utils.reauth_coordinator import get_reauth_coordinator lib_logger = logging.getLogger("rotator_library") @@ -85,11 +86,11 @@ def __init__(self): # [QUEUE SYSTEM] Sequential refresh processing self._refresh_queue: asyncio.Queue = asyncio.Queue() self._queued_credentials: set = set() # Track credentials already in queue - # [FIX 4] Changed from set to dict mapping credential path to timestamp + # [FIX PR#34] Changed from set to dict mapping credential path to timestamp # This enables TTL-based stale entry cleanup as defense in depth - self._unavailable_credentials: Dict[str, float] = ( - {} - ) # Maps credential path -> timestamp when marked unavailable + self._unavailable_credentials: Dict[ + str, float + ] = {} # Maps credential path -> timestamp when marked unavailable self._unavailable_ttl_seconds: int = 300 # 5 minutes TTL for stale entries self._queue_tracking_lock = asyncio.Lock() # Protects queue sets self._queue_processor_task: Optional[asyncio.Task] = ( @@ -530,15 +531,15 @@ async def _get_lock(self, path: str) -> asyncio.Lock: def is_credential_available(self, path: str) -> bool: """Check if a credential is available for rotation (not queued/refreshing). - - [FIX 4] Now includes TTL-based stale entry cleanup as defense in depth. + + [FIX PR#34] Now includes TTL-based stale entry cleanup as defense in depth. If a credential has been unavailable for longer than _unavailable_ttl_seconds, it is automatically cleaned up and considered available. """ if path not in self._unavailable_credentials: return True - - # [FIX 4] Check if the entry is stale (TTL expired) + + # [FIX PR#34] Check if the entry is stale (TTL expired) marked_time = self._unavailable_credentials.get(path) if marked_time is not None: now = time.time() @@ -550,11 +551,11 @@ def is_credential_available(self, path: str) -> bool: f"Auto-cleaning stale entry." ) # Note: This is a sync method, so we can't use async lock here. - # However, discard from dict is thread-safe for single operations. + # However, pop from dict is thread-safe for single operations. # The _queue_tracking_lock protects concurrent modifications in async context. self._unavailable_credentials.pop(path, None) return True - + return False async def _ensure_queue_processor_running(self): @@ -591,7 +592,7 @@ async def _queue_refresh( async with self._queue_tracking_lock: if path not in self._queued_credentials: self._queued_credentials.add(path) - # [FIX 4] Store timestamp when marking unavailable (for TTL cleanup) + # [FIX PR#34] Store timestamp when marking unavailable (for TTL cleanup) self._unavailable_credentials[path] = time.time() lib_logger.debug( f"Marked '{Path(path).name}' as unavailable. " @@ -611,7 +612,7 @@ async def _process_refresh_queue(self): self._refresh_queue.get(), timeout=60.0 ) except asyncio.TimeoutError: - # [FIX 2] Clean up any stale unavailable entries before exiting + # [FIX PR#34] Clean up any stale unavailable entries before exiting # If we're idle for 60s, no refreshes are in progress async with self._queue_tracking_lock: if self._unavailable_credentials: @@ -653,11 +654,11 @@ async def _process_refresh_queue(self): ) finally: - # [FIX 1] Remove from BOTH queued set AND unavailable credentials + # [FIX PR#34] Remove from BOTH queued set AND unavailable credentials # This ensures cleanup happens in ALL exit paths (success, exception, etc.) async with self._queue_tracking_lock: self._queued_credentials.discard(path) - # [FIX 1] Always clean up unavailable credentials in finally block + # [FIX PR#34] Always clean up unavailable credentials in finally block self._unavailable_credentials.pop(path, None) lib_logger.debug( f"Finally cleanup for '{Path(path).name}'. " @@ -665,7 +666,7 @@ async def _process_refresh_queue(self): ) self._refresh_queue.task_done() except asyncio.CancelledError: - # [FIX 3] Clean up the current credential before breaking + # [FIX PR#34] Clean up the current credential before breaking if path: async with self._queue_tracking_lock: self._unavailable_credentials.pop(path, None) @@ -685,9 +686,196 @@ async def _process_refresh_queue(self): f"Remaining unavailable: {len(self._unavailable_credentials)}" ) + async def _perform_interactive_oauth( + self, path: str, creds: Dict[str, Any], display_name: str + ) -> Dict[str, Any]: + """ + Perform interactive OAuth flow (browser-based authentication). + + This method is called via the global ReauthCoordinator to ensure + only one interactive OAuth flow runs at a time across all providers. + + Args: + path: Credential file path + creds: Current credentials dict (will be updated) + display_name: Display name for logging/UI + + Returns: + Updated credentials dict with new tokens + """ + # [HEADLESS DETECTION] Check if running in headless environment + is_headless = is_headless_environment() + + auth_code_future = asyncio.get_event_loop().create_future() + server = None + + async def handle_callback(reader, writer): + try: + request_line_bytes = await reader.readline() + if not request_line_bytes: + return + path_str = request_line_bytes.decode("utf-8").strip().split(" ")[1] + while await reader.readline() != b"\r\n": + pass + from urllib.parse import urlparse, parse_qs + + query_params = parse_qs(urlparse(path_str).query) + writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n") + if "code" in query_params: + if not auth_code_future.done(): + auth_code_future.set_result(query_params["code"][0]) + writer.write( + b"

Authentication successful!

You can close this window.

" + ) + else: + error = query_params.get("error", ["Unknown error"])[0] + if not auth_code_future.done(): + auth_code_future.set_exception( + Exception(f"OAuth failed: {error}") + ) + writer.write( + f"

Authentication Failed

Error: {error}. Please try again.

".encode() + ) + await writer.drain() + except Exception as e: + lib_logger.error(f"Error in OAuth callback handler: {e}") + finally: + writer.close() + + try: + server = await asyncio.start_server( + handle_callback, "127.0.0.1", self.CALLBACK_PORT + ) + from urllib.parse import urlencode + + auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode( + { + "client_id": self.CLIENT_ID, + "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", + "scope": " ".join(self.OAUTH_SCOPES), + "access_type": "offline", + "response_type": "code", + "prompt": "consent", + } + ) + + # [HEADLESS SUPPORT] Display appropriate instructions + if is_headless: + auth_panel_text = Text.from_markup( + "Running in headless environment (no GUI detected).\n" + "Please open the URL below in a browser on another machine to authorize:\n" + ) + else: + auth_panel_text = Text.from_markup( + "1. Your browser will now open to log in and authorize the application.\n" + "2. If it doesn't open automatically, please open the URL below manually." + ) + + console.print( + Panel( + auth_panel_text, + title=f"{self.ENV_PREFIX} OAuth Setup for [bold yellow]{display_name}[/bold yellow]", + style="bold blue", + ) + ) + # [URL DISPLAY] Print URL with proper escaping to prevent Rich markup issues. + # IMPORTANT: OAuth URLs contain special characters (=, &, etc.) that Rich might + # interpret as markup in some terminal configurations. We escape the URL to + # ensure it displays correctly. + # + # KNOWN ISSUE: If Rich rendering fails entirely (e.g., terminal doesn't support + # ANSI codes, or output is piped), the escaped URL should still be valid. + # However, if the terminal strips or mangles the output, users should copy + # the URL directly from logs or use --verbose to see the raw URL. + # + # The [link=...] markup creates a clickable hyperlink in supported terminals + # (iTerm2, Windows Terminal, etc.), but the displayed text is the escaped URL + # which can be safely copied even if the hyperlink doesn't work. + escaped_url = rich_escape(auth_url) + console.print(f"[bold]URL:[/bold] [link={auth_url}]{escaped_url}[/link]\n") + + # [HEADLESS SUPPORT] Only attempt browser open if NOT headless + if not is_headless: + try: + webbrowser.open(auth_url) + lib_logger.info("Browser opened successfully for OAuth flow") + except Exception as e: + lib_logger.warning( + f"Failed to open browser automatically: {e}. Please open the URL manually." + ) + + with console.status( + f"[bold green]Waiting for you to complete authentication in the browser...[/bold green]", + spinner="dots", + ): + # Note: The 300s timeout here is handled by the ReauthCoordinator + # We use a slightly longer internal timeout to let the coordinator handle it + auth_code = await asyncio.wait_for(auth_code_future, timeout=310) + except asyncio.TimeoutError: + raise Exception("OAuth flow timed out. Please try again.") + finally: + if server: + server.close() + await server.wait_closed() + + lib_logger.info(f"Attempting to exchange authorization code for tokens...") + async with httpx.AsyncClient() as client: + response = await client.post( + self.TOKEN_URI, + data={ + "code": auth_code.strip(), + "client_id": self.CLIENT_ID, + "client_secret": self.CLIENT_SECRET, + "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", + "grant_type": "authorization_code", + }, + ) + response.raise_for_status() + token_data = response.json() + # Start with the full token data from the exchange + new_creds = token_data.copy() + + # Convert 'expires_in' to 'expiry_date' in milliseconds + new_creds["expiry_date"] = ( + time.time() + new_creds.pop("expires_in") + ) * 1000 + + # Ensure client_id and client_secret are present + new_creds["client_id"] = self.CLIENT_ID + new_creds["client_secret"] = self.CLIENT_SECRET + + new_creds["token_uri"] = self.TOKEN_URI + new_creds["universe_domain"] = "googleapis.com" + + # Fetch user info and add metadata + user_info_response = await client.get( + self.USER_INFO_URI, + headers={"Authorization": f"Bearer {new_creds['access_token']}"}, + ) + user_info_response.raise_for_status() + user_info = user_info_response.json() + new_creds["_proxy_metadata"] = { + "email": user_info.get("email"), + "last_check_timestamp": time.time(), + } + + if path: + await self._save_credentials(path, new_creds) + lib_logger.info( + f"{self.ENV_PREFIX} OAuth initialized successfully for '{display_name}'." + ) + return new_creds + async def initialize_token( self, creds_or_path: Union[Dict[str, Any], str] ) -> Dict[str, Any]: + """ + Initialize OAuth token, triggering interactive OAuth flow if needed. + + If interactive OAuth is required (expired refresh token, missing credentials, etc.), + the flow is coordinated globally via ReauthCoordinator to ensure only one + interactive OAuth flow runs at a time across all providers. + """ path = creds_or_path if isinstance(creds_or_path, str) else None # Get display name from metadata if available, otherwise derive from path @@ -724,181 +912,23 @@ async def initialize_token( f"{self.ENV_PREFIX} OAuth token for '{display_name}' needs setup: {reason}." ) - # [HEADLESS DETECTION] Check if running in headless environment - is_headless = is_headless_environment() - - auth_code_future = asyncio.get_event_loop().create_future() - server = None - - async def handle_callback(reader, writer): - try: - request_line_bytes = await reader.readline() - if not request_line_bytes: - return - path_str = ( - request_line_bytes.decode("utf-8").strip().split(" ")[1] - ) - while await reader.readline() != b"\r\n": - pass - from urllib.parse import urlparse, parse_qs - - query_params = parse_qs(urlparse(path_str).query) - writer.write( - b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n" - ) - if "code" in query_params: - if not auth_code_future.done(): - auth_code_future.set_result(query_params["code"][0]) - writer.write( - b"

Authentication successful!

You can close this window.

" - ) - else: - error = query_params.get("error", ["Unknown error"])[0] - if not auth_code_future.done(): - auth_code_future.set_exception( - Exception(f"OAuth failed: {error}") - ) - writer.write( - f"

Authentication Failed

Error: {error}. Please try again.

".encode() - ) - await writer.drain() - except Exception as e: - lib_logger.error(f"Error in OAuth callback handler: {e}") - finally: - writer.close() + # [GLOBAL REAUTH COORDINATION] Use the global coordinator to ensure + # only one interactive OAuth flow runs at a time across all providers + coordinator = get_reauth_coordinator() - try: - server = await asyncio.start_server( - handle_callback, "127.0.0.1", self.CALLBACK_PORT - ) - from urllib.parse import urlencode - - auth_url = ( - "https://accounts.google.com/o/oauth2/v2/auth?" - + urlencode( - { - "client_id": self.CLIENT_ID, - "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", - "scope": " ".join(self.OAUTH_SCOPES), - "access_type": "offline", - "response_type": "code", - "prompt": "consent", - } - ) + # Define the interactive OAuth function to be executed by coordinator + async def _do_interactive_oauth(): + return await self._perform_interactive_oauth( + path, creds, display_name ) - # [HEADLESS SUPPORT] Display appropriate instructions - if is_headless: - auth_panel_text = Text.from_markup( - "Running in headless environment (no GUI detected).\n" - "Please open the URL below in a browser on another machine to authorize:\n" - ) - else: - auth_panel_text = Text.from_markup( - "1. Your browser will now open to log in and authorize the application.\n" - "2. If it doesn't open automatically, please open the URL below manually." - ) - - console.print( - Panel( - auth_panel_text, - title=f"{self.ENV_PREFIX} OAuth Setup for [bold yellow]{display_name}[/bold yellow]", - style="bold blue", - ) - ) - # [URL DISPLAY] Print URL with proper escaping to prevent Rich markup issues. - # IMPORTANT: OAuth URLs contain special characters (=, &, etc.) that Rich might - # interpret as markup in some terminal configurations. We escape the URL to - # ensure it displays correctly. - # - # KNOWN ISSUE: If Rich rendering fails entirely (e.g., terminal doesn't support - # ANSI codes, or output is piped), the escaped URL should still be valid. - # However, if the terminal strips or mangles the output, users should copy - # the URL directly from logs or use --verbose to see the raw URL. - # - # The [link=...] markup creates a clickable hyperlink in supported terminals - # (iTerm2, Windows Terminal, etc.), but the displayed text is the escaped URL - # which can be safely copied even if the hyperlink doesn't work. - escaped_url = rich_escape(auth_url) - console.print( - f"[bold]URL:[/bold] [link={auth_url}]{escaped_url}[/link]\n" - ) - - # [HEADLESS SUPPORT] Only attempt browser open if NOT headless - if not is_headless: - try: - webbrowser.open(auth_url) - lib_logger.info( - "Browser opened successfully for OAuth flow" - ) - except Exception as e: - lib_logger.warning( - f"Failed to open browser automatically: {e}. Please open the URL manually." - ) - - with console.status( - f"[bold green]Waiting for you to complete authentication in the browser...[/bold green]", - spinner="dots", - ): - auth_code = await asyncio.wait_for( - auth_code_future, timeout=300 - ) - except asyncio.TimeoutError: - raise Exception("OAuth flow timed out. Please try again.") - finally: - if server: - server.close() - await server.wait_closed() - - lib_logger.info( - f"Attempting to exchange authorization code for tokens..." + # Execute via global coordinator (ensures only one at a time) + return await coordinator.execute_reauth( + credential_path=path or display_name, + provider_name=self.ENV_PREFIX, + reauth_func=_do_interactive_oauth, + timeout=300.0, # 5 minute timeout for user to complete OAuth ) - async with httpx.AsyncClient() as client: - response = await client.post( - self.TOKEN_URI, - data={ - "code": auth_code.strip(), - "client_id": self.CLIENT_ID, - "client_secret": self.CLIENT_SECRET, - "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", - "grant_type": "authorization_code", - }, - ) - response.raise_for_status() - token_data = response.json() - # Start with the full token data from the exchange - creds = token_data.copy() - - # Convert 'expires_in' to 'expiry_date' in milliseconds - creds["expiry_date"] = ( - time.time() + creds.pop("expires_in") - ) * 1000 - - # Ensure client_id and client_secret are present - creds["client_id"] = self.CLIENT_ID - creds["client_secret"] = self.CLIENT_SECRET - - creds["token_uri"] = self.TOKEN_URI - creds["universe_domain"] = "googleapis.com" - - # Fetch user info and add metadata - user_info_response = await client.get( - self.USER_INFO_URI, - headers={"Authorization": f"Bearer {creds['access_token']}"}, - ) - user_info_response.raise_for_status() - user_info = user_info_response.json() - creds["_proxy_metadata"] = { - "email": user_info.get("email"), - "last_check_timestamp": time.time(), - } - - if path: - await self._save_credentials(path, creds) - lib_logger.info( - f"{self.ENV_PREFIX} OAuth initialized successfully for '{display_name}'." - ) - return creds lib_logger.info( f"{self.ENV_PREFIX} OAuth token at '{display_name}' is valid." diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index 021c3100..4d20f14c 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -23,6 +23,7 @@ from rich.text import Text from rich.markup import escape as rich_escape from ..utils.headless_detection import is_headless_environment +from ..utils.reauth_coordinator import get_reauth_coordinator lib_logger = logging.getLogger("rotator_library") @@ -173,9 +174,12 @@ def __init__(self): # [QUEUE SYSTEM] Sequential refresh processing self._refresh_queue: asyncio.Queue = asyncio.Queue() self._queued_credentials: set = set() # Track credentials already in queue - self._unavailable_credentials: set = ( - set() - ) # Mark credentials unavailable during re-auth + # [FIX PR#34] Changed from set to dict mapping credential path to timestamp + # This enables TTL-based stale entry cleanup as defense in depth + self._unavailable_credentials: Dict[ + str, float + ] = {} # Maps credential path -> timestamp when marked unavailable + self._unavailable_ttl_seconds: int = 300 # 5 minutes TTL for stale entries self._queue_tracking_lock = asyncio.Lock() # Protects queue sets self._queue_processor_task: Optional[asyncio.Task] = ( None # Background worker task @@ -768,8 +772,30 @@ async def _get_lock(self, path: str) -> asyncio.Lock: return self._refresh_locks[path] def is_credential_available(self, path: str) -> bool: - """Check if a credential is available for rotation (not queued/refreshing).""" - return path not in self._unavailable_credentials + """Check if a credential is available for rotation (not queued/refreshing). + + [FIX PR#34] Now includes TTL-based stale entry cleanup as defense in depth. + If a credential has been unavailable for longer than _unavailable_ttl_seconds, + it is automatically cleaned up and considered available. + """ + if path not in self._unavailable_credentials: + return True + + # [FIX PR#34] Check if the entry is stale (TTL expired) + marked_time = self._unavailable_credentials.get(path) + if marked_time is not None: + now = time.time() + if now - marked_time > self._unavailable_ttl_seconds: + # Entry is stale - clean it up and return available + lib_logger.warning( + f"Credential '{Path(path).name}' was stuck in unavailable state for " + f"{int(now - marked_time)}s (TTL: {self._unavailable_ttl_seconds}s). " + f"Auto-cleaning stale entry." + ) + self._unavailable_credentials.pop(path, None) + return True + + return False async def _ensure_queue_processor_running(self): """Lazily starts the queue processor if not already running.""" @@ -805,7 +831,12 @@ async def _queue_refresh( async with self._queue_tracking_lock: if path not in self._queued_credentials: self._queued_credentials.add(path) - self._unavailable_credentials.add(path) # Mark as unavailable + # [FIX PR#34] Store timestamp when marking unavailable (for TTL cleanup) + self._unavailable_credentials[path] = time.time() + lib_logger.debug( + f"Marked '{Path(path).name}' as unavailable. " + f"Total unavailable: {len(self._unavailable_credentials)}" + ) await self._refresh_queue.put((path, force, needs_reauth)) await self._ensure_queue_processor_running() @@ -820,7 +851,16 @@ async def _process_refresh_queue(self): self._refresh_queue.get(), timeout=60.0 ) except asyncio.TimeoutError: - # No items for 60s, exit to save resources + # [FIX PR#34] Clean up any stale unavailable entries before exiting + # If we're idle for 60s, no refreshes are in progress + async with self._queue_tracking_lock: + if self._unavailable_credentials: + stale_count = len(self._unavailable_credentials) + lib_logger.warning( + f"Queue processor idle timeout. Cleaning {stale_count} " + f"stale unavailable credentials: {list(self._unavailable_credentials.keys())}" + ) + self._unavailable_credentials.clear() self._queue_processor_task = None return @@ -832,7 +872,11 @@ async def _process_refresh_queue(self): if creds and not self._is_token_expired(creds): # No longer expired, mark as available async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Credential '{Path(path).name}' no longer expired, marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) continue # Perform refresh @@ -842,28 +886,174 @@ async def _process_refresh_queue(self): # SUCCESS: Mark as available again async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) finally: - # Remove from queued set + # [FIX PR#34] Remove from BOTH queued set AND unavailable credentials + # This ensures cleanup happens in ALL exit paths (success, exception, etc.) async with self._queue_tracking_lock: self._queued_credentials.discard(path) + # [FIX PR#34] Always clean up unavailable credentials in finally block + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Finally cleanup for '{Path(path).name}'. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) self._refresh_queue.task_done() except asyncio.CancelledError: + # [FIX PR#34] Clean up the current credential before breaking + if path: + async with self._queue_tracking_lock: + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"CancelledError cleanup for '{Path(path).name}'. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) break except Exception as e: lib_logger.error(f"Error in queue processor: {e}") # Even on error, mark as available (backoff will prevent immediate retry) if path: async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Error cleanup for '{Path(path).name}': {e}. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) + + async def _perform_interactive_oauth( + self, path: str, creds: Dict[str, Any], display_name: str + ) -> Dict[str, Any]: + """ + Perform interactive OAuth authorization code flow (browser-based authentication). + + This method is called via the global ReauthCoordinator to ensure + only one interactive OAuth flow runs at a time across all providers. + + Args: + path: Credential file path + creds: Current credentials dict (will be updated) + display_name: Display name for logging/UI + + Returns: + Updated credentials dict with new tokens + """ + # [HEADLESS DETECTION] Check if running in headless environment + is_headless = is_headless_environment() + + # Generate random state for CSRF protection + state = secrets.token_urlsafe(32) + + # Build authorization URL + redirect_uri = f"http://localhost:{CALLBACK_PORT}/oauth2callback" + auth_params = { + "loginMethod": "phone", + "type": "phone", + "redirect": redirect_uri, + "state": state, + "client_id": IFLOW_CLIENT_ID, + } + auth_url = f"{IFLOW_OAUTH_AUTHORIZE_ENDPOINT}?{urlencode(auth_params)}" + + # Start OAuth callback server + callback_server = OAuthCallbackServer(port=CALLBACK_PORT) + try: + await callback_server.start(expected_state=state) + + # [HEADLESS SUPPORT] Display appropriate instructions + if is_headless: + auth_panel_text = Text.from_markup( + "Running in headless environment (no GUI detected).\n" + "Please open the URL below in a browser on another machine to authorize:\n" + "1. Visit the URL below to sign in with your phone number.\n" + "2. [bold]Authorize the application[/bold] to access your account.\n" + "3. You will be automatically redirected after authorization." + ) + else: + auth_panel_text = Text.from_markup( + "1. Visit the URL below to sign in with your phone number.\n" + "2. [bold]Authorize the application[/bold] to access your account.\n" + "3. You will be automatically redirected after authorization." + ) + + console.print( + Panel( + auth_panel_text, + title=f"iFlow OAuth Setup for [bold yellow]{display_name}[/bold yellow]", + style="bold blue", + ) + ) + escaped_url = rich_escape(auth_url) + console.print(f"[bold]URL:[/bold] [link={auth_url}]{escaped_url}[/link]\n") + + # [HEADLESS SUPPORT] Only attempt browser open if NOT headless + if not is_headless: + try: + webbrowser.open(auth_url) + lib_logger.info("Browser opened successfully for iFlow OAuth flow") + except Exception as e: + lib_logger.warning( + f"Failed to open browser automatically: {e}. Please open the URL manually." + ) + + # Wait for callback + with console.status( + "[bold green]Waiting for authorization in the browser...[/bold green]", + spinner="dots", + ): + # Note: The 300s timeout here is handled by the ReauthCoordinator + # We use a slightly longer internal timeout to let the coordinator handle it + code = await callback_server.wait_for_callback(timeout=310.0) + + lib_logger.info("Received authorization code, exchanging for tokens...") + + # Exchange code for tokens and API key + token_data = await self._exchange_code_for_tokens(code, redirect_uri) + + # Update credentials + creds.update( + { + "access_token": token_data["access_token"], + "refresh_token": token_data["refresh_token"], + "api_key": token_data["api_key"], + "email": token_data["email"], + "expiry_date": token_data["expiry_date"], + "token_type": token_data["token_type"], + "scope": token_data["scope"], + } + ) + + # Create metadata object + if not creds.get("_proxy_metadata"): + creds["_proxy_metadata"] = { + "email": token_data["email"], + "last_check_timestamp": time.time(), + } + + if path: + await self._save_credentials(path, creds) + + lib_logger.info( + f"iFlow OAuth initialized successfully for '{display_name}'." + ) + return creds + + finally: + await callback_server.stop() async def initialize_token( self, creds_or_path: Union[Dict[str, Any], str] ) -> Dict[str, Any]: """ - Initiates OAuth authorization code flow if tokens are missing or invalid. - Uses local callback server to receive authorization code. + Initialize OAuth token, triggering interactive authorization flow if needed. + + If interactive OAuth is required (expired refresh token, missing credentials, etc.), + the flow is coordinated globally via ReauthCoordinator to ensure only one + interactive OAuth flow runs at a time across all providers. """ path = creds_or_path if isinstance(creds_or_path, str) else None @@ -903,127 +1093,23 @@ async def initialize_token( f"iFlow OAuth token for '{display_name}' needs setup: {reason}." ) - # [HEADLESS DETECTION] Check if running in headless environment - is_headless = is_headless_environment() - - # Generate random state for CSRF protection - state = secrets.token_urlsafe(32) - - # Build authorization URL - redirect_uri = f"http://localhost:{CALLBACK_PORT}/oauth2callback" - auth_params = { - "loginMethod": "phone", - "type": "phone", - "redirect": redirect_uri, - "state": state, - "client_id": IFLOW_CLIENT_ID, - } - auth_url = f"{IFLOW_OAUTH_AUTHORIZE_ENDPOINT}?{urlencode(auth_params)}" - - # Start OAuth callback server - callback_server = OAuthCallbackServer(port=CALLBACK_PORT) - try: - await callback_server.start(expected_state=state) - - # [HEADLESS SUPPORT] Display appropriate instructions - if is_headless: - auth_panel_text = Text.from_markup( - "Running in headless environment (no GUI detected).\n" - "Please open the URL below in a browser on another machine to authorize:\n" - "1. Visit the URL below to sign in with your phone number.\n" - "2. [bold]Authorize the application[/bold] to access your account.\n" - "3. You will be automatically redirected after authorization." - ) - else: - auth_panel_text = Text.from_markup( - "1. Visit the URL below to sign in with your phone number.\n" - "2. [bold]Authorize the application[/bold] to access your account.\n" - "3. You will be automatically redirected after authorization." - ) - - console.print( - Panel( - auth_panel_text, - title=f"iFlow OAuth Setup for [bold yellow]{display_name}[/bold yellow]", - style="bold blue", - ) - ) - # [URL DISPLAY] Print URL with proper escaping to prevent Rich markup issues. - # IMPORTANT: OAuth URLs contain special characters (=, &, etc.) that Rich might - # interpret as markup in some terminal configurations. We escape the URL to - # ensure it displays correctly. - # - # KNOWN ISSUE: If Rich rendering fails entirely (e.g., terminal doesn't support - # ANSI codes, or output is piped), the escaped URL should still be valid. - # However, if the terminal strips or mangles the output, users should copy - # the URL directly from logs or use --verbose to see the raw URL. - # - # The [link=...] markup creates a clickable hyperlink in supported terminals - # (iTerm2, Windows Terminal, etc.), but the displayed text is the escaped URL - # which can be safely copied even if the hyperlink doesn't work. - escaped_url = rich_escape(auth_url) - console.print( - f"[bold]URL:[/bold] [link={auth_url}]{escaped_url}[/link]\n" - ) + # [GLOBAL REAUTH COORDINATION] Use the global coordinator to ensure + # only one interactive OAuth flow runs at a time across all providers + coordinator = get_reauth_coordinator() - # [HEADLESS SUPPORT] Only attempt browser open if NOT headless - if not is_headless: - try: - webbrowser.open(auth_url) - lib_logger.info( - "Browser opened successfully for iFlow OAuth flow" - ) - except Exception as e: - lib_logger.warning( - f"Failed to open browser automatically: {e}. Please open the URL manually." - ) - - # Wait for callback - with console.status( - "[bold green]Waiting for authorization in the browser...[/bold green]", - spinner="dots", - ): - code = await callback_server.wait_for_callback(timeout=300.0) - - lib_logger.info( - "Received authorization code, exchanging for tokens..." + # Define the interactive OAuth function to be executed by coordinator + async def _do_interactive_oauth(): + return await self._perform_interactive_oauth( + path, creds, display_name ) - # Exchange code for tokens and API key - token_data = await self._exchange_code_for_tokens( - code, redirect_uri - ) - - # Update credentials - creds.update( - { - "access_token": token_data["access_token"], - "refresh_token": token_data["refresh_token"], - "api_key": token_data["api_key"], - "email": token_data["email"], - "expiry_date": token_data["expiry_date"], - "token_type": token_data["token_type"], - "scope": token_data["scope"], - } - ) - - # Create metadata object - if not creds.get("_proxy_metadata"): - creds["_proxy_metadata"] = { - "email": token_data["email"], - "last_check_timestamp": time.time(), - } - - if path: - await self._save_credentials(path, creds) - - lib_logger.info( - f"iFlow OAuth initialized successfully for '{display_name}'." - ) - return creds - - finally: - await callback_server.stop() + # Execute via global coordinator (ensures only one at a time) + return await coordinator.execute_reauth( + credential_path=path or display_name, + provider_name="IFLOW", + reauth_func=_do_interactive_oauth, + timeout=300.0, # 5 minute timeout for user to complete OAuth + ) lib_logger.info(f"iFlow OAuth token at '{display_name}' is valid.") return creds diff --git a/src/rotator_library/providers/qwen_auth_base.py b/src/rotator_library/providers/qwen_auth_base.py index 66e1d685..090c1716 100644 --- a/src/rotator_library/providers/qwen_auth_base.py +++ b/src/rotator_library/providers/qwen_auth_base.py @@ -22,6 +22,7 @@ from rich.markup import escape as rich_escape from ..utils.headless_detection import is_headless_environment +from ..utils.reauth_coordinator import get_reauth_coordinator lib_logger = logging.getLogger("rotator_library") @@ -53,9 +54,12 @@ def __init__(self): # [QUEUE SYSTEM] Sequential refresh processing self._refresh_queue: asyncio.Queue = asyncio.Queue() self._queued_credentials: set = set() # Track credentials already in queue - self._unavailable_credentials: set = ( - set() - ) # Mark credentials unavailable during re-auth + # [FIX PR#34] Changed from set to dict mapping credential path to timestamp + # This enables TTL-based stale entry cleanup as defense in depth + self._unavailable_credentials: Dict[ + str, float + ] = {} # Maps credential path -> timestamp when marked unavailable + self._unavailable_ttl_seconds: int = 300 # 5 minutes TTL for stale entries self._queue_tracking_lock = asyncio.Lock() # Protects queue sets self._queue_processor_task: Optional[asyncio.Task] = ( None # Background worker task @@ -494,8 +498,30 @@ async def _get_lock(self, path: str) -> asyncio.Lock: return self._refresh_locks[path] def is_credential_available(self, path: str) -> bool: - """Check if a credential is available for rotation (not queued/refreshing).""" - return path not in self._unavailable_credentials + """Check if a credential is available for rotation (not queued/refreshing). + + [FIX PR#34] Now includes TTL-based stale entry cleanup as defense in depth. + If a credential has been unavailable for longer than _unavailable_ttl_seconds, + it is automatically cleaned up and considered available. + """ + if path not in self._unavailable_credentials: + return True + + # [FIX PR#34] Check if the entry is stale (TTL expired) + marked_time = self._unavailable_credentials.get(path) + if marked_time is not None: + now = time.time() + if now - marked_time > self._unavailable_ttl_seconds: + # Entry is stale - clean it up and return available + lib_logger.warning( + f"Credential '{Path(path).name}' was stuck in unavailable state for " + f"{int(now - marked_time)}s (TTL: {self._unavailable_ttl_seconds}s). " + f"Auto-cleaning stale entry." + ) + self._unavailable_credentials.pop(path, None) + return True + + return False async def _ensure_queue_processor_running(self): """Lazily starts the queue processor if not already running.""" @@ -531,7 +557,12 @@ async def _queue_refresh( async with self._queue_tracking_lock: if path not in self._queued_credentials: self._queued_credentials.add(path) - self._unavailable_credentials.add(path) # Mark as unavailable + # [FIX PR#34] Store timestamp when marking unavailable (for TTL cleanup) + self._unavailable_credentials[path] = time.time() + lib_logger.debug( + f"Marked '{Path(path).name}' as unavailable. " + f"Total unavailable: {len(self._unavailable_credentials)}" + ) await self._refresh_queue.put((path, force, needs_reauth)) await self._ensure_queue_processor_running() @@ -546,7 +577,16 @@ async def _process_refresh_queue(self): self._refresh_queue.get(), timeout=60.0 ) except asyncio.TimeoutError: - # No items for 60s, exit to save resources + # [FIX PR#34] Clean up any stale unavailable entries before exiting + # If we're idle for 60s, no refreshes are in progress + async with self._queue_tracking_lock: + if self._unavailable_credentials: + stale_count = len(self._unavailable_credentials) + lib_logger.warning( + f"Queue processor idle timeout. Cleaning {stale_count} " + f"stale unavailable credentials: {list(self._unavailable_credentials.keys())}" + ) + self._unavailable_credentials.clear() self._queue_processor_task = None return @@ -558,7 +598,11 @@ async def _process_refresh_queue(self): if creds and not self._is_token_expired(creds): # No longer expired, mark as available async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Credential '{Path(path).name}' no longer expired, marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) continue # Perform refresh @@ -568,26 +612,240 @@ async def _process_refresh_queue(self): # SUCCESS: Mark as available again async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) finally: - # Remove from queued set + # [FIX PR#34] Remove from BOTH queued set AND unavailable credentials + # This ensures cleanup happens in ALL exit paths (success, exception, etc.) async with self._queue_tracking_lock: self._queued_credentials.discard(path) + # [FIX PR#34] Always clean up unavailable credentials in finally block + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Finally cleanup for '{Path(path).name}'. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) self._refresh_queue.task_done() except asyncio.CancelledError: + # [FIX PR#34] Clean up the current credential before breaking + if path: + async with self._queue_tracking_lock: + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"CancelledError cleanup for '{Path(path).name}'. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) break except Exception as e: lib_logger.error(f"Error in queue processor: {e}") # Even on error, mark as available (backoff will prevent immediate retry) if path: async with self._queue_tracking_lock: - self._unavailable_credentials.discard(path) + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Error cleanup for '{Path(path).name}': {e}. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) + + async def _perform_interactive_oauth( + self, path: str, creds: Dict[str, Any], display_name: str + ) -> Dict[str, Any]: + """ + Perform interactive OAuth device flow (browser-based authentication). + + This method is called via the global ReauthCoordinator to ensure + only one interactive OAuth flow runs at a time across all providers. + + Args: + path: Credential file path + creds: Current credentials dict (will be updated) + display_name: Display name for logging/UI + + Returns: + Updated credentials dict with new tokens + """ + # [HEADLESS DETECTION] Check if running in headless environment + is_headless = is_headless_environment() + + code_verifier = ( + base64.urlsafe_b64encode(secrets.token_bytes(32)) + .decode("utf-8") + .rstrip("=") + ) + code_challenge = ( + base64.urlsafe_b64encode( + hashlib.sha256(code_verifier.encode("utf-8")).digest() + ) + .decode("utf-8") + .rstrip("=") + ) + + headers = { + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", + "Content-Type": "application/x-www-form-urlencoded", + "Accept": "application/json", + } + async with httpx.AsyncClient() as client: + request_data = { + "client_id": CLIENT_ID, + "scope": SCOPE, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + lib_logger.debug(f"Qwen device code request data: {request_data}") + try: + dev_response = await client.post( + "https://chat.qwen.ai/api/v1/oauth2/device/code", + headers=headers, + data=request_data, + ) + dev_response.raise_for_status() + dev_data = dev_response.json() + lib_logger.debug(f"Qwen device auth response: {dev_data}") + except httpx.HTTPStatusError as e: + lib_logger.error( + f"Qwen device code request failed with status {e.response.status_code}: {e.response.text}" + ) + raise e + + # [HEADLESS SUPPORT] Display appropriate instructions + if is_headless: + auth_panel_text = Text.from_markup( + "Running in headless environment (no GUI detected).\n" + "Please open the URL below in a browser on another machine to authorize:\n" + "1. Visit the URL below to sign in.\n" + "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\n" + "3. You will be prompted to enter your identifier after authorization." + ) + else: + auth_panel_text = Text.from_markup( + "1. Visit the URL below to sign in.\n" + "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\n" + "3. You will be prompted to enter your identifier after authorization." + ) + + console.print( + Panel( + auth_panel_text, + title=f"Qwen OAuth Setup for [bold yellow]{display_name}[/bold yellow]", + style="bold blue", + ) + ) + verification_url = dev_data["verification_uri_complete"] + escaped_url = rich_escape(verification_url) + console.print( + f"[bold]URL:[/bold] [link={verification_url}]{escaped_url}[/link]\n" + ) + + # [HEADLESS SUPPORT] Only attempt browser open if NOT headless + if not is_headless: + try: + webbrowser.open(dev_data["verification_uri_complete"]) + lib_logger.info("Browser opened successfully for Qwen OAuth flow") + except Exception as e: + lib_logger.warning( + f"Failed to open browser automatically: {e}. Please open the URL manually." + ) + + token_data = None + start_time = time.time() + interval = dev_data.get("interval", 5) + + with console.status( + "[bold green]Polling for token, please complete authentication in the browser...[/bold green]", + spinner="dots", + ) as status: + while time.time() - start_time < dev_data["expires_in"]: + poll_response = await client.post( + TOKEN_ENDPOINT, + headers=headers, + data={ + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + "device_code": dev_data["device_code"], + "client_id": CLIENT_ID, + "code_verifier": code_verifier, + }, + ) + if poll_response.status_code == 200: + token_data = poll_response.json() + lib_logger.info("Successfully received token.") + break + elif poll_response.status_code == 400: + poll_data = poll_response.json() + error_type = poll_data.get("error") + if error_type == "authorization_pending": + lib_logger.debug( + f"Polling status: {error_type}, waiting {interval}s" + ) + elif error_type == "slow_down": + interval = int(interval * 1.5) + if interval > 10: + interval = 10 + lib_logger.debug( + f"Polling status: {error_type}, waiting {interval}s" + ) + else: + raise ValueError( + f"Token polling failed: {poll_data.get('error_description', error_type)}" + ) + else: + poll_response.raise_for_status() + + await asyncio.sleep(interval) + + if not token_data: + raise TimeoutError("Qwen device flow timed out.") + + creds.update( + { + "access_token": token_data["access_token"], + "refresh_token": token_data.get("refresh_token"), + "expiry_date": (time.time() + token_data["expires_in"]) * 1000, + "resource_url": token_data.get("resource_url"), + } + ) + + # Prompt for user identifier and create metadata object if needed + if not creds.get("_proxy_metadata", {}).get("email"): + try: + prompt_text = Text.from_markup( + f"\\n[bold]Please enter your email or a unique identifier for [yellow]'{display_name}'[/yellow][/bold]" + ) + email = Prompt.ask(prompt_text) + creds["_proxy_metadata"] = { + "email": email.strip(), + "last_check_timestamp": time.time(), + } + except (EOFError, KeyboardInterrupt): + console.print( + "\\n[bold yellow]No identifier provided. Deduplication will not be possible.[/bold yellow]" + ) + creds["_proxy_metadata"] = { + "email": None, + "last_check_timestamp": time.time(), + } + + if path: + await self._save_credentials(path, creds) + lib_logger.info( + f"Qwen OAuth initialized successfully for '{display_name}'." + ) + return creds async def initialize_token( self, creds_or_path: Union[Dict[str, Any], str] ) -> Dict[str, Any]: - """Initiates device flow if tokens are missing or invalid.""" + """ + Initialize OAuth token, triggering interactive device flow if needed. + + If interactive OAuth is required (expired refresh token, missing credentials, etc.), + the flow is coordinated globally via ReauthCoordinator to ensure only one + interactive OAuth flow runs at a time across all providers. + """ path = creds_or_path if isinstance(creds_or_path, str) else None # Get display name from metadata if available, otherwise derive from path @@ -623,189 +881,23 @@ async def initialize_token( f"Qwen OAuth token for '{display_name}' needs setup: {reason}." ) - # [HEADLESS DETECTION] Check if running in headless environment - is_headless = is_headless_environment() - - code_verifier = ( - base64.urlsafe_b64encode(secrets.token_bytes(32)) - .decode("utf-8") - .rstrip("=") - ) - code_challenge = ( - base64.urlsafe_b64encode( - hashlib.sha256(code_verifier.encode("utf-8")).digest() - ) - .decode("utf-8") - .rstrip("=") - ) - - headers = { - "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - "Content-Type": "application/x-www-form-urlencoded", - "Accept": "application/json", - } - async with httpx.AsyncClient() as client: - request_data = { - "client_id": CLIENT_ID, - "scope": SCOPE, - "code_challenge": code_challenge, - "code_challenge_method": "S256", - } - lib_logger.debug(f"Qwen device code request data: {request_data}") - try: - dev_response = await client.post( - "https://chat.qwen.ai/api/v1/oauth2/device/code", - headers=headers, - data=request_data, - ) - dev_response.raise_for_status() - dev_data = dev_response.json() - lib_logger.debug(f"Qwen device auth response: {dev_data}") - except httpx.HTTPStatusError as e: - lib_logger.error( - f"Qwen device code request failed with status {e.response.status_code}: {e.response.text}" - ) - raise e - - # [HEADLESS SUPPORT] Display appropriate instructions - if is_headless: - auth_panel_text = Text.from_markup( - "Running in headless environment (no GUI detected).\n" - "Please open the URL below in a browser on another machine to authorize:\n" - "1. Visit the URL below to sign in.\n" - "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\n" - "3. You will be prompted to enter your identifier after authorization." - ) - else: - auth_panel_text = Text.from_markup( - "1. Visit the URL below to sign in.\n" - "2. [bold]Copy your email[/bold] or another unique identifier and authorize the application.\n" - "3. You will be prompted to enter your identifier after authorization." - ) - - console.print( - Panel( - auth_panel_text, - title=f"Qwen OAuth Setup for [bold yellow]{display_name}[/bold yellow]", - style="bold blue", - ) - ) - # [URL DISPLAY] Print URL with proper escaping to prevent Rich markup issues. - # IMPORTANT: OAuth URLs contain special characters (=, &, etc.) that Rich might - # interpret as markup in some terminal configurations. We escape the URL to - # ensure it displays correctly. - # - # KNOWN ISSUE: If Rich rendering fails entirely (e.g., terminal doesn't support - # ANSI codes, or output is piped), the escaped URL should still be valid. - # However, if the terminal strips or mangles the output, users should copy - # the URL directly from logs or use --verbose to see the raw URL. - # - # The [link=...] markup creates a clickable hyperlink in supported terminals - # (iTerm2, Windows Terminal, etc.), but the displayed text is the escaped URL - # which can be safely copied even if the hyperlink doesn't work. - verification_url = dev_data["verification_uri_complete"] - escaped_url = rich_escape(verification_url) - console.print( - f"[bold]URL:[/bold] [link={verification_url}]{escaped_url}[/link]\n" - ) + # [GLOBAL REAUTH COORDINATION] Use the global coordinator to ensure + # only one interactive OAuth flow runs at a time across all providers + coordinator = get_reauth_coordinator() - # [HEADLESS SUPPORT] Only attempt browser open if NOT headless - if not is_headless: - try: - webbrowser.open(dev_data["verification_uri_complete"]) - lib_logger.info( - "Browser opened successfully for Qwen OAuth flow" - ) - except Exception as e: - lib_logger.warning( - f"Failed to open browser automatically: {e}. Please open the URL manually." - ) - - token_data = None - start_time = time.time() - interval = dev_data.get("interval", 5) - - with console.status( - "[bold green]Polling for token, please complete authentication in the browser...[/bold green]", - spinner="dots", - ) as status: - while time.time() - start_time < dev_data["expires_in"]: - poll_response = await client.post( - TOKEN_ENDPOINT, - headers=headers, - data={ - "grant_type": "urn:ietf:params:oauth:grant-type:device_code", - "device_code": dev_data["device_code"], - "client_id": CLIENT_ID, - "code_verifier": code_verifier, - }, - ) - if poll_response.status_code == 200: - token_data = poll_response.json() - lib_logger.info("Successfully received token.") - break - elif poll_response.status_code == 400: - poll_data = poll_response.json() - error_type = poll_data.get("error") - if error_type == "authorization_pending": - lib_logger.debug( - f"Polling status: {error_type}, waiting {interval}s" - ) - elif error_type == "slow_down": - interval = int(interval * 1.5) - if interval > 10: - interval = 10 - lib_logger.debug( - f"Polling status: {error_type}, waiting {interval}s" - ) - else: - raise ValueError( - f"Token polling failed: {poll_data.get('error_description', error_type)}" - ) - else: - poll_response.raise_for_status() - - await asyncio.sleep(interval) - - if not token_data: - raise TimeoutError("Qwen device flow timed out.") - - creds.update( - { - "access_token": token_data["access_token"], - "refresh_token": token_data.get("refresh_token"), - "expiry_date": (time.time() + token_data["expires_in"]) - * 1000, - "resource_url": token_data.get("resource_url"), - } + # Define the interactive OAuth function to be executed by coordinator + async def _do_interactive_oauth(): + return await self._perform_interactive_oauth( + path, creds, display_name ) - # Prompt for user identifier and create metadata object if needed - if not creds.get("_proxy_metadata", {}).get("email"): - try: - prompt_text = Text.from_markup( - f"\\n[bold]Please enter your email or a unique identifier for [yellow]'{display_name}'[/yellow][/bold]" - ) - email = Prompt.ask(prompt_text) - creds["_proxy_metadata"] = { - "email": email.strip(), - "last_check_timestamp": time.time(), - } - except (EOFError, KeyboardInterrupt): - console.print( - "\\n[bold yellow]No identifier provided. Deduplication will not be possible.[/bold yellow]" - ) - creds["_proxy_metadata"] = { - "email": None, - "last_check_timestamp": time.time(), - } - - if path: - await self._save_credentials(path, creds) - lib_logger.info( - f"Qwen OAuth initialized successfully for '{display_name}'." - ) - return creds + # Execute via global coordinator (ensures only one at a time) + return await coordinator.execute_reauth( + credential_path=path or display_name, + provider_name="QWEN_CODE", + reauth_func=_do_interactive_oauth, + timeout=300.0, # 5 minute timeout for user to complete OAuth + ) lib_logger.info(f"Qwen OAuth token at '{display_name}' is valid.") return creds diff --git a/src/rotator_library/utils/__init__.py b/src/rotator_library/utils/__init__.py index 83a86429..86a48dee 100644 --- a/src/rotator_library/utils/__init__.py +++ b/src/rotator_library/utils/__init__.py @@ -1,5 +1,6 @@ # src/rotator_library/utils/__init__.py from .headless_detection import is_headless_environment +from .reauth_coordinator import get_reauth_coordinator, ReauthCoordinator -__all__ = ['is_headless_environment'] +__all__ = ["is_headless_environment", "get_reauth_coordinator", "ReauthCoordinator"] diff --git a/src/rotator_library/utils/reauth_coordinator.py b/src/rotator_library/utils/reauth_coordinator.py new file mode 100644 index 00000000..dec3fa3e --- /dev/null +++ b/src/rotator_library/utils/reauth_coordinator.py @@ -0,0 +1,235 @@ +# src/rotator_library/utils/reauth_coordinator.py + +""" +Global Re-authentication Coordinator + +Ensures only ONE interactive OAuth flow runs at a time across ALL providers. +This prevents port conflicts and user confusion when multiple credentials +need re-authentication simultaneously. + +When a credential needs interactive re-auth (expired refresh token, revoked, etc.), +it queues a request here. The coordinator ensures only one re-auth happens at a time, +regardless of which provider the credential belongs to. +""" + +import asyncio +import logging +import time +from typing import Callable, Optional, Dict, Any, Awaitable +from pathlib import Path + +lib_logger = logging.getLogger("rotator_library") + + +class ReauthCoordinator: + """ + Singleton coordinator for global re-authentication serialization. + + When a credential needs interactive re-auth (expired refresh token, revoked, etc.), + it queues a request here. The coordinator ensures only one re-auth happens at a time. + + This is critical because: + 1. Different providers may use the same callback ports + 2. User can only complete one OAuth flow at a time + 3. Prevents race conditions in credential state management + """ + + _instance: Optional["ReauthCoordinator"] = None + + def __new__(cls): + # Singleton pattern - only one coordinator exists + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + if self._initialized: + return + + # Global semaphore - only 1 re-auth at a time + self._reauth_semaphore: asyncio.Semaphore = asyncio.Semaphore(1) + + # Tracking for observability + self._pending_reauths: Dict[str, float] = {} # credential -> queue_time + self._current_reauth: Optional[str] = None + self._current_provider: Optional[str] = None + self._reauth_start_time: Optional[float] = None + + # Lock for tracking dict modifications + self._tracking_lock: asyncio.Lock = asyncio.Lock() + + # Statistics + self._total_reauths: int = 0 + self._successful_reauths: int = 0 + self._failed_reauths: int = 0 + self._timeout_reauths: int = 0 + + self._initialized = True + lib_logger.info("Global ReauthCoordinator initialized") + + def _get_display_name(self, credential_path: str) -> str: + """Get a display-friendly name for a credential path.""" + if credential_path.startswith("env://"): + return credential_path + return Path(credential_path).name + + async def execute_reauth( + self, + credential_path: str, + provider_name: str, + reauth_func: Callable[[], Awaitable[Dict[str, Any]]], + timeout: float = 300.0, # 5 minutes default timeout + ) -> Dict[str, Any]: + """ + Execute a re-authentication function with global serialization. + + Only one re-auth can run at a time across all providers. + Other requests wait in queue. + + Args: + credential_path: Path/identifier of the credential needing re-auth + provider_name: Name of the provider (for logging) + reauth_func: Async function that performs the actual re-auth + timeout: Maximum time to wait for re-auth to complete + + Returns: + The result from reauth_func (new credentials dict) + + Raises: + TimeoutError: If re-auth doesn't complete within timeout + Exception: Any exception from reauth_func is re-raised + """ + display_name = self._get_display_name(credential_path) + + # Track that this credential is waiting + async with self._tracking_lock: + self._pending_reauths[credential_path] = time.time() + pending_count = len(self._pending_reauths) + + # Log queue status + if self._current_reauth: + current_display = self._get_display_name(self._current_reauth) + lib_logger.info( + f"[ReauthCoordinator] Credential '{display_name}' ({provider_name}) queued for re-auth. " + f"Position in queue: {pending_count}. " + f"Currently processing: '{current_display}' ({self._current_provider})" + ) + else: + lib_logger.info( + f"[ReauthCoordinator] Credential '{display_name}' ({provider_name}) requesting re-auth." + ) + + try: + # Acquire global semaphore - blocks until our turn + async with self._reauth_semaphore: + # Calculate how long we waited in queue + async with self._tracking_lock: + queue_time = self._pending_reauths.pop(credential_path, time.time()) + wait_duration = time.time() - queue_time + self._current_reauth = credential_path + self._current_provider = provider_name + self._reauth_start_time = time.time() + self._total_reauths += 1 + + if wait_duration > 1.0: + lib_logger.info( + f"[ReauthCoordinator] Starting re-auth for '{display_name}' ({provider_name}) " + f"after waiting {wait_duration:.1f}s in queue" + ) + else: + lib_logger.info( + f"[ReauthCoordinator] Starting re-auth for '{display_name}' ({provider_name})" + ) + + try: + # Execute the actual re-auth with timeout + result = await asyncio.wait_for(reauth_func(), timeout=timeout) + + async with self._tracking_lock: + self._successful_reauths += 1 + duration = time.time() - self._reauth_start_time + + lib_logger.info( + f"[ReauthCoordinator] Re-auth SUCCESS for '{display_name}' ({provider_name}) " + f"in {duration:.1f}s" + ) + return result + + except asyncio.TimeoutError: + async with self._tracking_lock: + self._failed_reauths += 1 + self._timeout_reauths += 1 + lib_logger.error( + f"[ReauthCoordinator] Re-auth TIMEOUT for '{display_name}' ({provider_name}) " + f"after {timeout}s. User did not complete OAuth flow in time." + ) + raise TimeoutError( + f"Re-authentication timed out after {timeout}s. " + f"Please try again and complete the OAuth flow within the time limit." + ) + + except Exception as e: + async with self._tracking_lock: + self._failed_reauths += 1 + lib_logger.error( + f"[ReauthCoordinator] Re-auth FAILED for '{display_name}' ({provider_name}): {e}" + ) + raise + + finally: + async with self._tracking_lock: + self._current_reauth = None + self._current_provider = None + self._reauth_start_time = None + + # Log if there are still pending reauths + if self._pending_reauths: + lib_logger.info( + f"[ReauthCoordinator] {len(self._pending_reauths)} credential(s) " + f"still waiting for re-auth" + ) + + finally: + # Ensure we're removed from pending even if something goes wrong + async with self._tracking_lock: + self._pending_reauths.pop(credential_path, None) + + def is_reauth_in_progress(self) -> bool: + """Check if a re-auth is currently in progress.""" + return self._current_reauth is not None + + def get_pending_count(self) -> int: + """Get number of credentials waiting for re-auth.""" + return len(self._pending_reauths) + + def get_status(self) -> Dict[str, Any]: + """Get current coordinator status for debugging/monitoring.""" + return { + "current_reauth": self._current_reauth, + "current_provider": self._current_provider, + "reauth_in_progress": self._current_reauth is not None, + "reauth_duration": (time.time() - self._reauth_start_time) + if self._reauth_start_time + else None, + "pending_count": len(self._pending_reauths), + "pending_credentials": list(self._pending_reauths.keys()), + "stats": { + "total": self._total_reauths, + "successful": self._successful_reauths, + "failed": self._failed_reauths, + "timeouts": self._timeout_reauths, + }, + } + + +# Global singleton instance +_coordinator: Optional[ReauthCoordinator] = None + + +def get_reauth_coordinator() -> ReauthCoordinator: + """Get the global ReauthCoordinator instance.""" + global _coordinator + if _coordinator is None: + _coordinator = ReauthCoordinator() + return _coordinator From 1456ae3fb6ee347fdc02c7ae7da3a6355e26939d Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 07:29:37 +0100 Subject: [PATCH 106/221] =?UTF-8?q?fix(auth):=20=F0=9F=90=9B=20improve=20c?= =?UTF-8?q?redential=20refresh=20detection=20and=20prevent=20queue=20proce?= =?UTF-8?q?ssor=20stuck=20state?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enhanced the proactive token refresh logic in both IFlowAuthBase and QwenAuthBase to more robustly detect OAuth credentials versus direct API keys: - Changed from checking file existence/env:// prefix to attempting credential load in try/except block - Added comprehensive debug logging throughout the refresh flow to track credential lifecycle - Fixed BUG#6 where queued credentials were not cleared on queue processor timeout, potentially causing stuck state - Now clears both unavailable_credentials and queued_credentials when processor times out The previous approach of checking `is_env_path` and `os.path.isfile()` could incorrectly classify credentials. The new approach leverages the existing `_load_credentials()` exception handling to make a definitive determination. --- .../providers/iflow_auth_base.py | 33 +++++++++++++++---- .../providers/qwen_auth_base.py | 33 +++++++++++++++---- 2 files changed, 52 insertions(+), 14 deletions(-) diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index 4d20f14c..ccdae302 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -749,15 +749,28 @@ async def proactively_refresh(self, credential_identifier: str): Proactively refreshes tokens if they're close to expiry. Only applies to OAuth credentials (file paths or env:// paths). Direct API keys are skipped. """ - # Check if it's an env:// virtual path (OAuth credentials from environment) - is_env_path = credential_identifier.startswith("env://") + lib_logger.debug(f"proactively_refresh called for: {credential_identifier}") - # Only refresh if it's an OAuth credential (file path or env:// path) - if not is_env_path and not os.path.isfile(credential_identifier): - return # Direct API key, no refresh needed + # Try to load credentials - this will fail for direct API keys + # and succeed for OAuth credentials (file paths or env:// paths) + try: + creds = await self._load_credentials(credential_identifier) + except IOError as e: + # Not a valid credential path (likely a direct API key string) + lib_logger.debug( + f"Skipping refresh for '{credential_identifier}' - not an OAuth credential: {e}" + ) + return - creds = await self._load_credentials(credential_identifier) - if self._is_token_expired(creds): + is_expired = self._is_token_expired(creds) + lib_logger.debug( + f"Token expired check for '{Path(credential_identifier).name}': {is_expired}" + ) + + if is_expired: + lib_logger.debug( + f"Queueing refresh for '{Path(credential_identifier).name}'" + ) # Queue for refresh with needs_reauth=False (automated refresh) await self._queue_refresh( credential_identifier, force=False, needs_reauth=False @@ -861,6 +874,12 @@ async def _process_refresh_queue(self): f"stale unavailable credentials: {list(self._unavailable_credentials.keys())}" ) self._unavailable_credentials.clear() + # [FIX BUG#6] Also clear queued credentials to prevent stuck state + if self._queued_credentials: + lib_logger.debug( + f"Clearing {len(self._queued_credentials)} queued credentials on timeout" + ) + self._queued_credentials.clear() self._queue_processor_task = None return diff --git a/src/rotator_library/providers/qwen_auth_base.py b/src/rotator_library/providers/qwen_auth_base.py index 090c1716..7065bbe6 100644 --- a/src/rotator_library/providers/qwen_auth_base.py +++ b/src/rotator_library/providers/qwen_auth_base.py @@ -476,15 +476,28 @@ async def proactively_refresh(self, credential_identifier: str): Proactively refreshes tokens if they're close to expiry. Only applies to OAuth credentials (file paths or env:// paths). Direct API keys are skipped. """ - # Check if it's an env:// virtual path (OAuth credentials from environment) - is_env_path = credential_identifier.startswith("env://") + lib_logger.debug(f"proactively_refresh called for: {credential_identifier}") - # Only refresh if it's an OAuth credential (file path or env:// path) - if not is_env_path and not os.path.isfile(credential_identifier): - return # Direct API key, no refresh needed + # Try to load credentials - this will fail for direct API keys + # and succeed for OAuth credentials (file paths or env:// paths) + try: + creds = await self._load_credentials(credential_identifier) + except IOError as e: + # Not a valid credential path (likely a direct API key string) + lib_logger.debug( + f"Skipping refresh for '{credential_identifier}' - not an OAuth credential: {e}" + ) + return - creds = await self._load_credentials(credential_identifier) - if self._is_token_expired(creds): + is_expired = self._is_token_expired(creds) + lib_logger.debug( + f"Token expired check for '{Path(credential_identifier).name}': {is_expired}" + ) + + if is_expired: + lib_logger.debug( + f"Queueing refresh for '{Path(credential_identifier).name}'" + ) # Queue for refresh with needs_reauth=False (automated refresh) await self._queue_refresh( credential_identifier, force=False, needs_reauth=False @@ -587,6 +600,12 @@ async def _process_refresh_queue(self): f"stale unavailable credentials: {list(self._unavailable_credentials.keys())}" ) self._unavailable_credentials.clear() + # [FIX BUG#6] Also clear queued credentials to prevent stuck state + if self._queued_credentials: + lib_logger.debug( + f"Clearing {len(self._queued_credentials)} queued credentials on timeout" + ) + self._queued_credentials.clear() self._queue_processor_task = None return From d76b29a2fce8cb79066baeb2d173119cf1f7fe6e Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 07:58:05 +0100 Subject: [PATCH 107/221] =?UTF-8?q?refactor(auth):=20=F0=9F=94=A8=20reloca?= =?UTF-8?q?te=20attribute=20declarations=20in=20BackgroundRefresher?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Positioned the instance variable initializations earlier in the constructor to establish clear variable declarations at the outset of the method. This adjustment enhances code predictability by ensuring all attributes are defined before any complex logic execution. --- src/rotator_library/background_refresher.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/rotator_library/background_refresher.py b/src/rotator_library/background_refresher.py index a6830fa8..8c371388 100644 --- a/src/rotator_library/background_refresher.py +++ b/src/rotator_library/background_refresher.py @@ -18,6 +18,9 @@ class BackgroundRefresher: """ def __init__(self, client: "RotatingClient"): + self._client = client + self._task: Optional[asyncio.Task] = None + self._initialized = False try: interval_str = os.getenv("OAUTH_REFRESH_INTERVAL", "600") self._interval = int(interval_str) @@ -26,9 +29,6 @@ def __init__(self, client: "RotatingClient"): f"Invalid OAUTH_REFRESH_INTERVAL '{interval_str}'. Falling back to 600s." ) self._interval = 600 - self._client = client - self._task: Optional[asyncio.Task] = None - self._initialized = False def start(self): """Starts the background refresh task.""" From 4ecfabac17718def0998e0271b1bd449c90e8b67 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 08:03:15 +0100 Subject: [PATCH 108/221] =?UTF-8?q?refactor(proxy):=20=F0=9F=94=A8=20remov?= =?UTF-8?q?e=20debug=20print=20statement=20for=20credentials?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The debug print statement that was logging credential summaries to the console has been commented out. This removes unnecessary console output in the proxy application while keeping the credential loading logic intact. --- src/proxy_app/main.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/proxy_app/main.py b/src/proxy_app/main.py index 258a69f3..167bd985 100644 --- a/src/proxy_app/main.py +++ b/src/proxy_app/main.py @@ -500,10 +500,10 @@ async def process_credential(provider: str, path: str, provider_instance): ) # Log loaded credentials summary (compact, always visible for deployment verification) - _api_summary = ', '.join([f"{p}:{len(c)}" for p, c in api_keys.items()]) if api_keys else "none" - _oauth_summary = ', '.join([f"{p}:{len(c)}" for p, c in oauth_credentials.items()]) if oauth_credentials else "none" - _total_summary = ', '.join([f"{p}:{len(c)}" for p, c in client.all_credentials.items()]) - print(f"🔑 Credentials loaded: {_total_summary} (API: {_api_summary} | OAuth: {_oauth_summary})") + #_api_summary = ', '.join([f"{p}:{len(c)}" for p, c in api_keys.items()]) if api_keys else "none" + #_oauth_summary = ', '.join([f"{p}:{len(c)}" for p, c in oauth_credentials.items()]) if oauth_credentials else "none" + #_total_summary = ', '.join([f"{p}:{len(c)}" for p, c in client.all_credentials.items()]) + #print(f"🔑 Credentials loaded: {_total_summary} (API: {_api_summary} | OAuth: {_oauth_summary})") client.background_refresher.start() # Start the background task app.state.rotating_client = client From 0af8a39f85ce8a793ce12e8da76177fe0c6f65b6 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 08:06:14 +0100 Subject: [PATCH 109/221] Fix to satisfy pylint --- src/rotator_library/utils/reauth_coordinator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/rotator_library/utils/reauth_coordinator.py b/src/rotator_library/utils/reauth_coordinator.py index dec3fa3e..7d5f3cd0 100644 --- a/src/rotator_library/utils/reauth_coordinator.py +++ b/src/rotator_library/utils/reauth_coordinator.py @@ -35,6 +35,7 @@ class ReauthCoordinator: """ _instance: Optional["ReauthCoordinator"] = None + _initialized: bool = False # Class-level declaration for Pylint def __new__(cls): # Singleton pattern - only one coordinator exists From 7f148b3ce45e83c2b6d2efab093fa6ddfce8b3e5 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 09:29:09 +0100 Subject: [PATCH 110/221] =?UTF-8?q?feat(io):=20=E2=9C=A8=20add=20fault-tol?= =?UTF-8?q?erant=20file=20operations=20with=20automatic=20recovery?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enhances application reliability by introducing a comprehensive I/O abstraction layer that eliminates crashes from filesystem issues. The system distinguishes between critical state files (credentials, usage data) that require memory buffering with retry logic, and disposable logs that can be safely dropped on failure. Key improvements: - New `ResilientStateWriter` class maintains in-memory state for critical files with background retry mechanism on disk failure - Introduced `safe_write_json`, `safe_log_write`, and `safe_mkdir` utility functions for one-shot operations with graceful degradation - Logging subsystems (`DetailedLogger`, `failure_logger`) now drop data on disk failure to prevent memory exhaustion during streaming - Authentication providers (`GoogleOAuthBase`, `IFlowAuthBase`, `QwenAuthBase`) preserve credentials in memory when filesystem becomes unavailable - `UsageManager` delegates persistence to `ResilientStateWriter` for automatic recovery from transient failures - `ProviderCache` disk operations now fail silently while maintaining in-memory functionality - Replaced scattered tempfile/atomic write patterns with centralized implementation featuring consistent error handling - All directory creation operations now proceed gracefully if parent paths are inaccessible - Thread-safe writer implementation supports concurrent usage from async contexts BREAKING CHANGE: `ProviderCache._save_to_disk()` no longer raises exceptions on filesystem errors. Consumers relying on exception handling for disk write failures must now check the `disk_available` field in `get_stats()` return value for monitoring disk health. --- src/proxy_app/detailed_logger.py | 104 +++--- src/rotator_library/failure_logger.py | 80 ++--- .../providers/google_oauth_base.py | 84 +---- .../providers/iflow_auth_base.py | 66 +--- .../providers/provider_cache.py | 301 ++++++++-------- .../providers/qwen_auth_base.py | 64 +--- src/rotator_library/usage_manager.py | 51 +-- src/rotator_library/utils/__init__.py | 16 +- src/rotator_library/utils/resilient_io.py | 339 ++++++++++++++++++ 9 files changed, 618 insertions(+), 487 deletions(-) create mode 100644 src/rotator_library/utils/resilient_io.py diff --git a/src/proxy_app/detailed_logger.py b/src/proxy_app/detailed_logger.py index 0d0dd9a9..9afceef0 100644 --- a/src/proxy_app/detailed_logger.py +++ b/src/proxy_app/detailed_logger.py @@ -3,20 +3,27 @@ import uuid from datetime import datetime from pathlib import Path -from typing import Any, Dict, Optional, List +from typing import Any, Dict, Optional import logging +from rotator_library.utils.resilient_io import ( + safe_write_json, + safe_log_write, + safe_mkdir, +) + LOGS_DIR = Path(__file__).resolve().parent.parent.parent / "logs" DETAILED_LOGS_DIR = LOGS_DIR / "detailed_logs" + class DetailedLogger: """ Logs comprehensive details of each API transaction to a unique, timestamped directory. + + Uses fire-and-forget logging - if disk writes fail, logs are dropped (not buffered) + to prevent memory issues, especially with streaming responses. """ - # Class-level fallback flags for resilience - _disk_available = True - _console_fallback_warned = False - + def __init__(self): """ Initializes the logger for a single request, creating a unique directory to store all related log files. @@ -26,33 +33,24 @@ def __init__(self): timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") self.log_dir = DETAILED_LOGS_DIR / f"{timestamp}_{self.request_id}" self.streaming = False - self._in_memory_logs = [] # Fallback storage - - # Attempt directory creation with resilience - try: - self.log_dir.mkdir(parents=True, exist_ok=True) - DetailedLogger._disk_available = True - except (OSError, PermissionError) as e: - DetailedLogger._disk_available = False - if not DetailedLogger._console_fallback_warned: - logging.warning(f"Detailed logging disabled - cannot create log directory: {e}") - DetailedLogger._console_fallback_warned = True + self._dir_available = safe_mkdir(self.log_dir, logging) def _write_json(self, filename: str, data: Dict[str, Any]): """Helper to write data to a JSON file in the log directory.""" - if not DetailedLogger._disk_available: - self._in_memory_logs.append({"file": filename, "data": data}) - return - - try: - # Attempt directory recreation if needed - self.log_dir.mkdir(parents=True, exist_ok=True) - with open(self.log_dir / filename, "w", encoding="utf-8") as f: - json.dump(data, f, indent=4, ensure_ascii=False) - except (OSError, PermissionError, IOError) as e: - DetailedLogger._disk_available = False - logging.error(f"[{self.request_id}] Failed to write to {filename}: {e}") - self._in_memory_logs.append({"file": filename, "data": data}) + if not self._dir_available: + # Try to create directory again in case it was recreated + self._dir_available = safe_mkdir(self.log_dir, logging) + if not self._dir_available: + return + + safe_write_json( + self.log_dir / filename, + data, + logging, + atomic=False, + indent=4, + ensure_ascii=False, + ) def log_request(self, headers: Dict[str, Any], body: Dict[str, Any]): """Logs the initial request details.""" @@ -61,29 +59,22 @@ def log_request(self, headers: Dict[str, Any], body: Dict[str, Any]): "request_id": self.request_id, "timestamp_utc": datetime.utcnow().isoformat(), "headers": dict(headers), - "body": body + "body": body, } self._write_json("request.json", request_data) def log_stream_chunk(self, chunk: Dict[str, Any]): """Logs an individual chunk from a streaming response to a JSON Lines file.""" - # Intentionally skip memory fallback for streams to prevent OOM - unlike _write_json, we don't buffer stream chunks in memory - if not DetailedLogger._disk_available: + if not self._dir_available: return - - try: - self.log_dir.mkdir(parents=True, exist_ok=True) - log_entry = { - "timestamp_utc": datetime.utcnow().isoformat(), - "chunk": chunk - } - with open(self.log_dir / "streaming_chunks.jsonl", "a", encoding="utf-8") as f: - f.write(json.dumps(log_entry, ensure_ascii=False) + "\n") - except (OSError, PermissionError, IOError) as e: - DetailedLogger._disk_available = False - logging.error(f"[{self.request_id}] Failed to write stream chunk: {e}") - - def log_final_response(self, status_code: int, headers: Optional[Dict[str, Any]], body: Dict[str, Any]): + + log_entry = {"timestamp_utc": datetime.utcnow().isoformat(), "chunk": chunk} + content = json.dumps(log_entry, ensure_ascii=False) + "\n" + safe_log_write(self.log_dir / "streaming_chunks.jsonl", content, logging) + + def log_final_response( + self, status_code: int, headers: Optional[Dict[str, Any]], body: Dict[str, Any] + ): """Logs the complete final response, either from a non-streaming call or after reassembling a stream.""" end_time = time.time() duration_ms = (end_time - self.start_time) * 1000 @@ -94,7 +85,7 @@ def log_final_response(self, status_code: int, headers: Optional[Dict[str, Any]] "status_code": status_code, "duration_ms": round(duration_ms), "headers": dict(headers) if headers else None, - "body": body + "body": body, } self._write_json("final_response.json", response_data) self._log_metadata(response_data) @@ -103,10 +94,10 @@ def _extract_reasoning(self, response_body: Dict[str, Any]) -> Optional[str]: """Recursively searches for and extracts 'reasoning' fields from the response body.""" if not isinstance(response_body, dict): return None - + if "reasoning" in response_body: return response_body["reasoning"] - + if "choices" in response_body and response_body["choices"]: message = response_body["choices"][0].get("message", {}) if "reasoning" in message: @@ -121,8 +112,13 @@ def _log_metadata(self, response_data: Dict[str, Any]): usage = response_data.get("body", {}).get("usage") or {} model = response_data.get("body", {}).get("model", "N/A") finish_reason = "N/A" - if "choices" in response_data.get("body", {}) and response_data["body"]["choices"]: - finish_reason = response_data["body"]["choices"][0].get("finish_reason", "N/A") + if ( + "choices" in response_data.get("body", {}) + and response_data["body"]["choices"] + ): + finish_reason = response_data["body"]["choices"][0].get( + "finish_reason", "N/A" + ) metadata = { "request_id": self.request_id, @@ -138,12 +134,12 @@ def _log_metadata(self, response_data: Dict[str, Any]): }, "finish_reason": finish_reason, "reasoning_found": False, - "reasoning_content": None + "reasoning_content": None, } reasoning = self._extract_reasoning(response_data.get("body", {})) if reasoning: metadata["reasoning_found"] = True metadata["reasoning_content"] = reasoning - - self._write_json("metadata.json", metadata) \ No newline at end of file + + self._write_json("metadata.json", metadata) diff --git a/src/rotator_library/failure_logger.py b/src/rotator_library/failure_logger.py index a3e07d33..3fbda577 100644 --- a/src/rotator_library/failure_logger.py +++ b/src/rotator_library/failure_logger.py @@ -5,74 +5,42 @@ from datetime import datetime from .error_handler import mask_credential -# Module-level state for resilience -_file_handler = None -_fallback_mode = False - -# Custom JSON formatter for structured logs (defined at module level for reuse) class JsonFormatter(logging.Formatter): + """Custom JSON formatter for structured logs.""" + def format(self, record): # The message is already a dict, so we just format it as a JSON string return json.dumps(record.msg) -def _create_file_handler(): - """Create file handler with directory auto-recreation.""" - global _file_handler, _fallback_mode +def setup_failure_logger(): + """Sets up a dedicated JSON logger for writing detailed failure logs to a file.""" log_dir = "logs" - + logger = logging.getLogger("failure_logger") + logger.setLevel(logging.INFO) + logger.propagate = False + + # Clear existing handlers to prevent duplicates on re-setup + logger.handlers.clear() + try: if not os.path.exists(log_dir): os.makedirs(log_dir, exist_ok=True) - + handler = RotatingFileHandler( os.path.join(log_dir, "failures.log"), maxBytes=5 * 1024 * 1024, # 5 MB backupCount=2, ) - handler.setFormatter(JsonFormatter()) - _file_handler = handler - _fallback_mode = False - return handler + logger.addHandler(handler) except (OSError, PermissionError, IOError) as e: logging.warning(f"Cannot create failure log file handler: {e}") - _fallback_mode = True - return None - - -def setup_failure_logger(): - """Sets up a dedicated JSON logger for writing detailed failure logs.""" - logger = logging.getLogger("failure_logger") - logger.setLevel(logging.INFO) - logger.propagate = False - - # Remove existing handlers to prevent duplicates - logger.handlers.clear() - - # Try to add file handler - handler = _create_file_handler() - if handler: - logger.addHandler(handler) - - # Always add a NullHandler as fallback to prevent "no handlers" warning - if not logger.handlers: + # Add NullHandler to prevent "no handlers" warning logger.addHandler(logging.NullHandler()) - - return logger - -def _ensure_handler_valid(): - """Check if file handler is still valid, recreate if needed.""" - global _file_handler, _fallback_mode - - if _file_handler is None or _fallback_mode: - handler = _create_file_handler() - if handler: - failure_logger = logging.getLogger("failure_logger") - failure_logger.handlers.clear() - failure_logger.addHandler(handler) + return logger # Initialize the dedicated logger for detailed failure logs @@ -180,25 +148,19 @@ def log_failure( "request_headers": request_headers, "error_chain": error_chain if len(error_chain) > 1 else None, } - + # 2. Log a concise summary to the main library logger, which will propagate summary_message = ( f"API call failed for model {model} with key {mask_credential(api_key)}. " f"Error: {type(error).__name__}. See failures.log for details." ) - - # Attempt to ensure handler is valid before logging - _ensure_handler_valid() - - # Wrap the actual log call with resilience + + # Log to failure logger with resilience - if it fails, just continue try: failure_logger.error(detailed_log_data) except (OSError, IOError) as e: - global _fallback_mode - _fallback_mode = True - # File logging failed - log to console instead - logging.error(f"Failed to write to failures.log: {e}") - logging.error(f"Failure summary: {summary_message}") - + # Log file write failed - log to console instead + logging.warning(f"Failed to write to failures.log: {e}") + # Console log always succeeds main_lib_logger.error(summary_message) diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index 9120a44c..5f8a09b3 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -9,8 +9,6 @@ import logging from pathlib import Path from typing import Dict, Any -import tempfile -import shutil import httpx from rich.console import Console @@ -20,6 +18,7 @@ from ..utils.headless_detection import is_headless_environment from ..utils.reauth_coordinator import get_reauth_coordinator +from ..utils.resilient_io import safe_write_json lib_logger = logging.getLogger("rotator_library") @@ -264,13 +263,8 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: ) async def _save_credentials(self, path: str, creds: Dict[str, Any]): - """Save credentials with in-memory fallback if disk unavailable. - - [RUNTIME RESILIENCE] Always updates the in-memory cache first (memory is reliable), - then attempts disk persistence. If disk write fails, logs a warning but does NOT - raise an exception - the in-memory state continues to work. - """ - # [IN-MEMORY FIRST] Always update cache first (reliable) + """Save credentials with in-memory fallback if disk unavailable.""" + # Always update cache first (memory is reliable) self._credentials_cache[path] = creds # Don't save to file if credentials were loaded from environment @@ -278,62 +272,15 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): lib_logger.debug("Credentials loaded from env, skipping file save") return - try: - # [ATOMIC WRITE] Use tempfile + move pattern to ensure atomic writes - # This prevents credential corruption if the process is interrupted during write - parent_dir = os.path.dirname(os.path.abspath(path)) - os.makedirs(parent_dir, exist_ok=True) - - tmp_fd = None - tmp_path = None - try: - # Create temp file in same directory as target (ensures same filesystem) - tmp_fd, tmp_path = tempfile.mkstemp( - dir=parent_dir, prefix=".tmp_", suffix=".json", text=True - ) - - # Write JSON to temp file - with os.fdopen(tmp_fd, "w") as f: - json.dump(creds, f, indent=2) - tmp_fd = None # fdopen closes the fd - - # Set secure permissions (0600 = owner read/write only) - try: - os.chmod(tmp_path, 0o600) - except (OSError, AttributeError): - # Windows may not support chmod, ignore - pass - - # Atomic move (overwrites target if it exists) - shutil.move(tmp_path, path) - tmp_path = None # Successfully moved - - lib_logger.debug( - f"Saved updated {self.ENV_PREFIX} OAuth credentials to '{path}' (atomic write)." - ) - - except Exception as e: - # Clean up temp file if it still exists - if tmp_fd is not None: - try: - os.close(tmp_fd) - except: - pass - if tmp_path and os.path.exists(tmp_path): - try: - os.unlink(tmp_path) - except: - pass - raise - - except (OSError, PermissionError, IOError) as e: - # [FAIL SILENTLY, LOG LOUDLY] Log the error but don't crash - # The in-memory cache was already updated, so we can continue operating + # Attempt disk write - if it fails, we still have the cache + if safe_write_json(path, creds, lib_logger, secure_permissions=True): + lib_logger.debug( + f"Saved updated {self.ENV_PREFIX} OAuth credentials to '{path}'." + ) + else: lib_logger.warning( - f"Failed to save credentials to {path}: {e}. " - "Credentials cached in memory only (will be lost on restart)." + f"Credentials for {self.ENV_PREFIX} cached in memory only (will be lost on restart)." ) - # Don't raise - we already updated the memory cache def _is_token_expired(self, creds: Dict[str, Any]) -> bool: expiry = creds.get("token_expiry") # gcloud format @@ -952,19 +899,14 @@ async def _do_interactive_oauth(): ) async def get_auth_header(self, credential_path: str) -> Dict[str, str]: - """Get auth header with graceful degradation if refresh fails. - - [RUNTIME RESILIENCE] If credential file is deleted or refresh fails, - attempts to use cached credentials. This allows the proxy to continue - operating with potentially stale tokens rather than crashing. - """ + """Get auth header with graceful degradation if refresh fails.""" try: creds = await self._load_credentials(credential_path) if self._is_token_expired(creds): try: creds = await self._refresh_token(credential_path, creds) except Exception as e: - # [CACHED TOKEN FALLBACK] Check if we have a cached token that might still work + # Check if we have a cached token that might still work cached = self._credentials_cache.get(credential_path) if cached and cached.get("access_token"): lib_logger.warning( @@ -976,7 +918,7 @@ async def get_auth_header(self, credential_path: str) -> Dict[str, str]: raise return {"Authorization": f"Bearer {creds['access_token']}"} except Exception as e: - # [FINAL FALLBACK] Check if any cached credential exists as last resort + # Check if any cached credential exists as last resort cached = self._credentials_cache.get(credential_path) if cached and cached.get("access_token"): lib_logger.error( diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index ccdae302..a2096df3 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -12,8 +12,6 @@ from pathlib import Path from typing import Dict, Any, Tuple, Union, Optional from urllib.parse import urlencode, parse_qs, urlparse -import tempfile -import shutil import httpx from aiohttp import web @@ -24,6 +22,7 @@ from rich.markup import escape as rich_escape from ..utils.headless_detection import is_headless_environment from ..utils.reauth_coordinator import get_reauth_coordinator +from ..utils.resilient_io import safe_write_json lib_logger = logging.getLogger("rotator_library") @@ -316,65 +315,22 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: return await self._read_creds_from_file(path) async def _save_credentials(self, path: str, creds: Dict[str, Any]): - """Saves credentials to cache and file using atomic writes.""" + """Save credentials with in-memory fallback if disk unavailable.""" + # Always update cache first (memory is reliable) + self._credentials_cache[path] = creds + # Don't save to file if credentials were loaded from environment if creds.get("_proxy_metadata", {}).get("loaded_from_env"): lib_logger.debug("Credentials loaded from env, skipping file save") - # Still update cache for in-memory consistency - self._credentials_cache[path] = creds return - # [ATOMIC WRITE] Use tempfile + move pattern to ensure atomic writes - # This prevents credential corruption if the process is interrupted during write - parent_dir = os.path.dirname(os.path.abspath(path)) - os.makedirs(parent_dir, exist_ok=True) - - tmp_fd = None - tmp_path = None - try: - # Create temp file in same directory as target (ensures same filesystem) - tmp_fd, tmp_path = tempfile.mkstemp( - dir=parent_dir, prefix=".tmp_", suffix=".json", text=True - ) - - # Write JSON to temp file - with os.fdopen(tmp_fd, "w") as f: - json.dump(creds, f, indent=2) - tmp_fd = None # fdopen closes the fd - - # Set secure permissions (0600 = owner read/write only) - try: - os.chmod(tmp_path, 0o600) - except (OSError, AttributeError): - # Windows may not support chmod, ignore - pass - - # Atomic move (overwrites target if it exists) - shutil.move(tmp_path, path) - tmp_path = None # Successfully moved - - # Update cache AFTER successful file write - self._credentials_cache[path] = creds - lib_logger.debug( - f"Saved updated iFlow OAuth credentials to '{path}' (atomic write)." - ) - - except Exception as e: - lib_logger.error( - f"Failed to save updated iFlow OAuth credentials to '{path}': {e}" + # Attempt disk write - if it fails, we still have the cache + if safe_write_json(path, creds, lib_logger, secure_permissions=True): + lib_logger.debug(f"Saved updated iFlow OAuth credentials to '{path}'.") + else: + lib_logger.warning( + "iFlow credentials cached in memory only (will be lost on restart)." ) - # Clean up temp file if it still exists - if tmp_fd is not None: - try: - os.close(tmp_fd) - except: - pass - if tmp_path and os.path.exists(tmp_path): - try: - os.unlink(tmp_path) - except: - pass - raise def _is_token_expired(self, creds: Dict[str, Any]) -> bool: """Checks if the token is expired (with buffer for proactive refresh).""" diff --git a/src/rotator_library/providers/provider_cache.py b/src/rotator_library/providers/provider_cache.py index 1e7f85e6..8b0f835b 100644 --- a/src/rotator_library/providers/provider_cache.py +++ b/src/rotator_library/providers/provider_cache.py @@ -20,19 +20,20 @@ import json import logging import os -import shutil -import tempfile import time from pathlib import Path from typing import Any, Dict, Optional, Tuple -lib_logger = logging.getLogger('rotator_library') +from ..utils.resilient_io import safe_write_json + +lib_logger = logging.getLogger("rotator_library") # ============================================================================= # UTILITY FUNCTIONS # ============================================================================= + def _env_bool(key: str, default: bool = False) -> bool: """Get boolean from environment variable.""" return os.getenv(key, str(default).lower()).lower() in ("true", "1", "yes") @@ -47,18 +48,19 @@ def _env_int(key: str, default: int) -> int: # PROVIDER CACHE CLASS # ============================================================================= + class ProviderCache: """ Server-side cache for provider conversation state preservation. - + A generic, modular cache supporting any key-value data that providers need to persist across requests. Features: - + - Dual-TTL system: configurable memory TTL, longer disk TTL - Async disk persistence with batched writes - Background cleanup task for expired entries - Statistics tracking (hits, misses, writes) - + Args: cache_file: Path to disk cache file memory_ttl_seconds: In-memory entry lifetime (default: 1 hour) @@ -67,13 +69,13 @@ class ProviderCache: write_interval: Seconds between background disk writes (default: 60) cleanup_interval: Seconds between expired entry cleanup (default: 30 min) env_prefix: Environment variable prefix for configuration overrides - + Environment Variables (with default prefix "PROVIDER_CACHE"): {PREFIX}_ENABLE: Enable/disable disk persistence {PREFIX}_WRITE_INTERVAL: Background write interval in seconds {PREFIX}_CLEANUP_INTERVAL: Cleanup interval in seconds """ - + def __init__( self, cache_file: Path, @@ -82,7 +84,7 @@ def __init__( enable_disk: Optional[bool] = None, write_interval: Optional[int] = None, cleanup_interval: Optional[int] = None, - env_prefix: str = "PROVIDER_CACHE" + env_prefix: str = "PROVIDER_CACHE", ): # In-memory cache: {cache_key: (data, timestamp)} self._cache: Dict[str, Tuple[str, float]] = {} @@ -90,28 +92,42 @@ def __init__( self._disk_ttl = disk_ttl_seconds self._lock = asyncio.Lock() self._disk_lock = asyncio.Lock() - + # Disk persistence configuration self._cache_file = cache_file - self._enable_disk = enable_disk if enable_disk is not None else _env_bool(f"{env_prefix}_ENABLE", True) + self._enable_disk = ( + enable_disk + if enable_disk is not None + else _env_bool(f"{env_prefix}_ENABLE", True) + ) self._dirty = False - self._write_interval = write_interval or _env_int(f"{env_prefix}_WRITE_INTERVAL", 60) - self._cleanup_interval = cleanup_interval or _env_int(f"{env_prefix}_CLEANUP_INTERVAL", 1800) - + self._write_interval = write_interval or _env_int( + f"{env_prefix}_WRITE_INTERVAL", 60 + ) + self._cleanup_interval = cleanup_interval or _env_int( + f"{env_prefix}_CLEANUP_INTERVAL", 1800 + ) + # Background tasks self._writer_task: Optional[asyncio.Task] = None self._cleanup_task: Optional[asyncio.Task] = None self._running = False - + # Statistics - self._stats = {"memory_hits": 0, "disk_hits": 0, "misses": 0, "writes": 0, "disk_errors": 0} - - # [RUNTIME RESILIENCE] Track disk health for monitoring + self._stats = { + "memory_hits": 0, + "disk_hits": 0, + "misses": 0, + "writes": 0, + "disk_errors": 0, + } + + # Track disk health for monitoring self._disk_available = True - + # Metadata about this cache instance self._cache_name = cache_file.stem if cache_file else "unnamed" - + if self._enable_disk: lib_logger.debug( f"ProviderCache[{self._cache_name}]: Disk enabled " @@ -120,142 +136,114 @@ def __init__( asyncio.create_task(self._async_init()) else: lib_logger.debug(f"ProviderCache[{self._cache_name}]: Memory-only mode") - + # ========================================================================= # INITIALIZATION # ========================================================================= - + async def _async_init(self) -> None: """Async initialization: load from disk and start background tasks.""" try: await self._load_from_disk() await self._start_background_tasks() except Exception as e: - lib_logger.error(f"ProviderCache[{self._cache_name}] async init failed: {e}") - + lib_logger.error( + f"ProviderCache[{self._cache_name}] async init failed: {e}" + ) + async def _load_from_disk(self) -> None: """Load cache from disk file with TTL validation.""" if not self._enable_disk or not self._cache_file.exists(): return - + try: async with self._disk_lock: - with open(self._cache_file, 'r', encoding='utf-8') as f: + with open(self._cache_file, "r", encoding="utf-8") as f: data = json.load(f) - + if data.get("version") != "1.0": - lib_logger.warning(f"ProviderCache[{self._cache_name}]: Version mismatch, starting fresh") + lib_logger.warning( + f"ProviderCache[{self._cache_name}]: Version mismatch, starting fresh" + ) return - + now = time.time() entries = data.get("entries", {}) loaded = expired = 0 - + for cache_key, entry in entries.items(): age = now - entry.get("timestamp", 0) if age <= self._disk_ttl: - value = entry.get("value", entry.get("signature", "")) # Support both formats + value = entry.get( + "value", entry.get("signature", "") + ) # Support both formats if value: self._cache[cache_key] = (value, entry["timestamp"]) loaded += 1 else: expired += 1 - + lib_logger.debug( f"ProviderCache[{self._cache_name}]: Loaded {loaded} entries ({expired} expired)" ) except json.JSONDecodeError as e: - lib_logger.warning(f"ProviderCache[{self._cache_name}]: File corrupted: {e}") + lib_logger.warning( + f"ProviderCache[{self._cache_name}]: File corrupted: {e}" + ) except Exception as e: lib_logger.error(f"ProviderCache[{self._cache_name}]: Load failed: {e}") - + # ========================================================================= # DISK PERSISTENCE # ========================================================================= - + async def _save_to_disk(self) -> None: - """Persist cache to disk using atomic write with health tracking. - - [RUNTIME RESILIENCE] Tracks disk health and records errors. If disk - operations fail, the memory cache continues to work. Health status - is available via get_stats() for monitoring. - """ + """Persist cache to disk using atomic write with health tracking.""" if not self._enable_disk: return - - try: - async with self._disk_lock: - # [DIRECTORY AUTO-RECREATION] Attempt to create directory - try: - self._cache_file.parent.mkdir(parents=True, exist_ok=True) - except (OSError, PermissionError) as e: - self._stats["disk_errors"] += 1 - self._disk_available = False - lib_logger.warning( - f"ProviderCache[{self._cache_name}]: Cannot create cache directory: {e}" - ) - return - - cache_data = { - "version": "1.0", - "memory_ttl_seconds": self._memory_ttl, - "disk_ttl_seconds": self._disk_ttl, - "entries": { - key: {"value": val, "timestamp": ts} - for key, (val, ts) in self._cache.items() - }, - "statistics": { - "total_entries": len(self._cache), - "last_write": time.time(), - **self._stats - } - } - - # Atomic write using temp file - parent_dir = self._cache_file.parent - tmp_fd, tmp_path = tempfile.mkstemp(dir=parent_dir, prefix='.tmp_', suffix='.json') - - try: - with os.fdopen(tmp_fd, 'w', encoding='utf-8') as f: - json.dump(cache_data, f, indent=2) - - # Set restrictive permissions (if supported) - try: - os.chmod(tmp_path, 0o600) - except (OSError, AttributeError): - pass - - shutil.move(tmp_path, self._cache_file) - self._stats["writes"] += 1 - # [RUNTIME RESILIENCE] Mark disk as healthy on success - self._disk_available = True - lib_logger.debug( - f"ProviderCache[{self._cache_name}]: Saved {len(self._cache)} entries" - ) - except Exception: - if tmp_path and os.path.exists(tmp_path): - os.unlink(tmp_path) - raise - except Exception as e: - # [RUNTIME RESILIENCE] Track disk errors for monitoring - self._stats["disk_errors"] += 1 - self._disk_available = False - lib_logger.error(f"ProviderCache[{self._cache_name}]: Disk save failed: {e}") - + + async with self._disk_lock: + cache_data = { + "version": "1.0", + "memory_ttl_seconds": self._memory_ttl, + "disk_ttl_seconds": self._disk_ttl, + "entries": { + key: {"value": val, "timestamp": ts} + for key, (val, ts) in self._cache.items() + }, + "statistics": { + "total_entries": len(self._cache), + "last_write": time.time(), + **self._stats, + }, + } + + if safe_write_json( + self._cache_file, cache_data, lib_logger, secure_permissions=True + ): + self._stats["writes"] += 1 + self._disk_available = True + lib_logger.debug( + f"ProviderCache[{self._cache_name}]: Saved {len(self._cache)} entries" + ) + else: + self._stats["disk_errors"] += 1 + self._disk_available = False + # ========================================================================= # BACKGROUND TASKS # ========================================================================= - + async def _start_background_tasks(self) -> None: """Start background writer and cleanup tasks.""" if not self._enable_disk or self._running: return - + self._running = True self._writer_task = asyncio.create_task(self._writer_loop()) self._cleanup_task = asyncio.create_task(self._cleanup_loop()) lib_logger.debug(f"ProviderCache[{self._cache_name}]: Started background tasks") - + async def _writer_loop(self) -> None: """Background task: periodically flush dirty cache to disk.""" try: @@ -266,10 +254,12 @@ async def _writer_loop(self) -> None: await self._save_to_disk() self._dirty = False except Exception as e: - lib_logger.error(f"ProviderCache[{self._cache_name}]: Writer error: {e}") + lib_logger.error( + f"ProviderCache[{self._cache_name}]: Writer error: {e}" + ) except asyncio.CancelledError: pass - + async def _cleanup_loop(self) -> None: """Background task: periodically clean up expired entries.""" try: @@ -278,12 +268,14 @@ async def _cleanup_loop(self) -> None: await self._cleanup_expired() except asyncio.CancelledError: pass - + async def _cleanup_expired(self) -> None: """Remove expired entries from memory cache.""" async with self._lock: now = time.time() - expired = [k for k, (_, ts) in self._cache.items() if now - ts > self._memory_ttl] + expired = [ + k for k, (_, ts) in self._cache.items() if now - ts > self._memory_ttl + ] for k in expired: del self._cache[k] if expired: @@ -291,42 +283,42 @@ async def _cleanup_expired(self) -> None: lib_logger.debug( f"ProviderCache[{self._cache_name}]: Cleaned {len(expired)} expired entries" ) - + # ========================================================================= # CORE OPERATIONS # ========================================================================= - + def store(self, key: str, value: str) -> None: """ Store a value synchronously (schedules async storage). - + Args: key: Cache key value: Value to store (typically JSON-serialized data) """ asyncio.create_task(self._async_store(key, value)) - + async def _async_store(self, key: str, value: str) -> None: """Async implementation of store.""" async with self._lock: self._cache[key] = (value, time.time()) self._dirty = True - + async def store_async(self, key: str, value: str) -> None: """ Store a value asynchronously (awaitable). - + Use this when you need to ensure the value is stored before continuing. """ await self._async_store(key, value) - + def retrieve(self, key: str) -> Optional[str]: """ Retrieve a value by key (synchronous, with optional async disk fallback). - + Args: key: Cache key - + Returns: Cached value if found and not expired, None otherwise """ @@ -338,17 +330,17 @@ def retrieve(self, key: str) -> Optional[str]: else: del self._cache[key] self._dirty = True - + self._stats["misses"] += 1 if self._enable_disk: # Schedule async disk lookup for next time asyncio.create_task(self._check_disk_fallback(key)) return None - + async def retrieve_async(self, key: str) -> Optional[str]: """ Retrieve a value asynchronously (checks disk if not in memory). - + Use this when you can await and need guaranteed disk fallback. """ # Check memory first @@ -362,24 +354,24 @@ async def retrieve_async(self, key: str) -> Optional[str]: if key in self._cache: del self._cache[key] self._dirty = True - + # Check disk if self._enable_disk: return await self._disk_retrieve(key) - + self._stats["misses"] += 1 return None - + async def _check_disk_fallback(self, key: str) -> None: """Check disk for key and load into memory if found (background).""" try: if not self._cache_file.exists(): return - + async with self._disk_lock: - with open(self._cache_file, 'r', encoding='utf-8') as f: + with open(self._cache_file, "r", encoding="utf-8") as f: data = json.load(f) - + entries = data.get("entries", {}) if key in entries: entry = entries[key] @@ -394,19 +386,21 @@ async def _check_disk_fallback(self, key: str) -> None: f"ProviderCache[{self._cache_name}]: Loaded {key} from disk" ) except Exception as e: - lib_logger.debug(f"ProviderCache[{self._cache_name}]: Disk fallback failed: {e}") - + lib_logger.debug( + f"ProviderCache[{self._cache_name}]: Disk fallback failed: {e}" + ) + async def _disk_retrieve(self, key: str) -> Optional[str]: """Direct disk retrieval with loading into memory.""" try: if not self._cache_file.exists(): self._stats["misses"] += 1 return None - + async with self._disk_lock: - with open(self._cache_file, 'r', encoding='utf-8') as f: + with open(self._cache_file, "r", encoding="utf-8") as f: data = json.load(f) - + entries = data.get("entries", {}) if key in entries: entry = entries[key] @@ -418,39 +412,37 @@ async def _disk_retrieve(self, key: str) -> Optional[str]: self._cache[key] = (value, ts) self._stats["disk_hits"] += 1 return value - + self._stats["misses"] += 1 return None except Exception as e: - lib_logger.debug(f"ProviderCache[{self._cache_name}]: Disk retrieve failed: {e}") + lib_logger.debug( + f"ProviderCache[{self._cache_name}]: Disk retrieve failed: {e}" + ) self._stats["misses"] += 1 return None - + # ========================================================================= # UTILITY METHODS # ========================================================================= - + def contains(self, key: str) -> bool: """Check if key exists in memory cache (without updating stats).""" if key in self._cache: _, timestamp = self._cache[key] return time.time() - timestamp <= self._memory_ttl return False - + def get_stats(self) -> Dict[str, Any]: - """Get cache statistics including disk health. - - [RUNTIME RESILIENCE] Includes disk_available flag for monitoring - the health of disk persistence. - """ + """Get cache statistics including disk health.""" return { **self._stats, "memory_entries": len(self._cache), "dirty": self._dirty, "disk_enabled": self._enable_disk, - "disk_available": self._disk_available # [RUNTIME RESILIENCE] Health indicator + "disk_available": self._disk_available, } - + async def clear(self) -> None: """Clear all cached data.""" async with self._lock: @@ -458,12 +450,12 @@ async def clear(self) -> None: self._dirty = True if self._enable_disk: await self._save_to_disk() - + async def shutdown(self) -> None: """Graceful shutdown: flush pending writes and stop background tasks.""" lib_logger.info(f"ProviderCache[{self._cache_name}]: Shutting down...") self._running = False - + # Cancel background tasks for task in (self._writer_task, self._cleanup_task): if task: @@ -472,11 +464,11 @@ async def shutdown(self) -> None: await task except asyncio.CancelledError: pass - + # Final save if self._dirty and self._enable_disk: await self._save_to_disk() - + lib_logger.info( f"ProviderCache[{self._cache_name}]: Shutdown complete " f"(stats: mem_hits={self._stats['memory_hits']}, " @@ -488,38 +480,39 @@ async def shutdown(self) -> None: # CONVENIENCE FACTORY # ============================================================================= + def create_provider_cache( name: str, cache_dir: Optional[Path] = None, memory_ttl_seconds: int = 3600, disk_ttl_seconds: int = 86400, - env_prefix: Optional[str] = None + env_prefix: Optional[str] = None, ) -> ProviderCache: """ Factory function to create a provider cache with sensible defaults. - + Args: name: Cache name (used as filename and for logging) cache_dir: Directory for cache file (default: project_root/cache/provider_name) memory_ttl_seconds: In-memory TTL disk_ttl_seconds: Disk TTL env_prefix: Environment variable prefix (default: derived from name) - + Returns: Configured ProviderCache instance """ if cache_dir is None: cache_dir = Path(__file__).resolve().parent.parent.parent.parent / "cache" - + cache_file = cache_dir / f"{name}.json" - + if env_prefix is None: # Convert name to env prefix: "gemini3_signatures" -> "GEMINI3_SIGNATURES_CACHE" env_prefix = f"{name.upper().replace('-', '_')}_CACHE" - + return ProviderCache( cache_file=cache_file, memory_ttl_seconds=memory_ttl_seconds, disk_ttl_seconds=disk_ttl_seconds, - env_prefix=env_prefix + env_prefix=env_prefix, ) diff --git a/src/rotator_library/providers/qwen_auth_base.py b/src/rotator_library/providers/qwen_auth_base.py index 7065bbe6..b95416a5 100644 --- a/src/rotator_library/providers/qwen_auth_base.py +++ b/src/rotator_library/providers/qwen_auth_base.py @@ -11,8 +11,6 @@ import os from pathlib import Path from typing import Dict, Any, Tuple, Union, Optional -import tempfile -import shutil import httpx from rich.console import Console @@ -23,6 +21,7 @@ from ..utils.headless_detection import is_headless_environment from ..utils.reauth_coordinator import get_reauth_coordinator +from ..utils.resilient_io import safe_write_json lib_logger = logging.getLogger("rotator_library") @@ -201,63 +200,22 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: return await self._read_creds_from_file(path) async def _save_credentials(self, path: str, creds: Dict[str, Any]): + """Save credentials with in-memory fallback if disk unavailable.""" + # Always update cache first (memory is reliable) + self._credentials_cache[path] = creds + # Don't save to file if credentials were loaded from environment if creds.get("_proxy_metadata", {}).get("loaded_from_env"): lib_logger.debug("Credentials loaded from env, skipping file save") - # Still update cache for in-memory consistency - self._credentials_cache[path] = creds return - # [ATOMIC WRITE] Use tempfile + move pattern to ensure atomic writes - parent_dir = os.path.dirname(os.path.abspath(path)) - os.makedirs(parent_dir, exist_ok=True) - - tmp_fd = None - tmp_path = None - try: - # Create temp file in same directory as target (ensures same filesystem) - tmp_fd, tmp_path = tempfile.mkstemp( - dir=parent_dir, prefix=".tmp_", suffix=".json", text=True - ) - - # Write JSON to temp file - with os.fdopen(tmp_fd, "w") as f: - json.dump(creds, f, indent=2) - tmp_fd = None # fdopen closes the fd - - # Set secure permissions (0600 = owner read/write only) - try: - os.chmod(tmp_path, 0o600) - except (OSError, AttributeError): - # Windows may not support chmod, ignore - pass - - # Atomic move (overwrites target if it exists) - shutil.move(tmp_path, path) - tmp_path = None # Successfully moved - - # Update cache AFTER successful file write - self._credentials_cache[path] = creds - lib_logger.debug( - f"Saved updated Qwen OAuth credentials to '{path}' (atomic write)." - ) - - except Exception as e: - lib_logger.error( - f"Failed to save updated Qwen OAuth credentials to '{path}': {e}" + # Attempt disk write - if it fails, we still have the cache + if safe_write_json(path, creds, lib_logger, secure_permissions=True): + lib_logger.debug(f"Saved updated Qwen OAuth credentials to '{path}'.") + else: + lib_logger.warning( + "Qwen credentials cached in memory only (will be lost on restart)." ) - # Clean up temp file if it still exists - if tmp_fd is not None: - try: - os.close(tmp_fd) - except: - pass - if tmp_path and os.path.exists(tmp_path): - try: - os.unlink(tmp_path) - except: - pass - raise def _is_token_expired(self, creds: Dict[str, Any]) -> bool: expiry_timestamp = creds.get("expiry_date", 0) / 1000 diff --git a/src/rotator_library/usage_manager.py b/src/rotator_library/usage_manager.py index ac8ca739..613b4c33 100644 --- a/src/rotator_library/usage_manager.py +++ b/src/rotator_library/usage_manager.py @@ -11,6 +11,7 @@ from .error_handler import ClassifiedError, NoAvailableKeysError, mask_credential from .providers import PROVIDER_PLUGINS +from .utils.resilient_io import ResilientStateWriter lib_logger = logging.getLogger("rotator_library") lib_logger.propagate = False @@ -103,8 +104,8 @@ def __init__( self._timeout_lock = asyncio.Lock() self._claimed_on_timeout: Set[str] = set() - # Circuit breaker for disk write failures - self._disk_available = True + # Resilient writer for usage data persistence + self._state_writer = ResilientStateWriter(file_path, lib_logger) if daily_reset_time_utc: hour, minute = map(int, daily_reset_time_utc.split(":")) @@ -543,11 +544,7 @@ async def _lazy_init(self): self._initialized.set() async def _load_usage(self): - """Loads usage data from the JSON file asynchronously with enhanced resilience. - - [RUNTIME RESILIENCE] Handles various file system errors gracefully, - including race conditions where file is deleted between exists check and open. - """ + """Loads usage data from the JSON file asynchronously with resilience.""" async with self._data_lock: if not os.path.exists(self.file_path): self._usage_data = {} @@ -558,7 +555,7 @@ async def _load_usage(self): content = await f.read() self._usage_data = json.loads(content) if content.strip() else {} except FileNotFoundError: - # [RACE CONDITION HANDLING] File deleted between exists check and open + # File deleted between exists check and open self._usage_data = {} except json.JSONDecodeError as e: lib_logger.warning( @@ -570,43 +567,17 @@ async def _load_usage(self): f"Cannot read usage file {self.file_path}: {e}. Using empty state." ) self._usage_data = {} - else: - # [CIRCUIT BREAKER RESET] Successfully loaded, re-enable disk writes - self._disk_available = True async def _save_usage(self): - """Saves the current usage data to the JSON file asynchronously with resilience. - - [RUNTIME RESILIENCE] Wraps file operations in try/except to prevent crashes - if the file or directory is deleted during runtime. The in-memory state - continues to work even if disk persistence fails. - """ + """Saves the current usage data using the resilient state writer.""" if self._usage_data is None: return - if not self._disk_available: - return # Skip disk write when unavailable - - try: - async with self._data_lock: - # [DIRECTORY AUTO-RECREATION] Ensure directory exists before write - file_dir = os.path.dirname(os.path.abspath(self.file_path)) - if file_dir and not os.path.exists(file_dir): - os.makedirs(file_dir, exist_ok=True) - - # Add human-readable timestamp fields before saving - self._add_readable_timestamps(self._usage_data) - async with aiofiles.open(self.file_path, "w") as f: - await f.write(json.dumps(self._usage_data, indent=2)) - except (OSError, PermissionError, IOError) as e: - # [CIRCUIT BREAKER] Disable disk writes to prevent repeated failures - self._disk_available = False - # [FAIL SILENTLY, LOG LOUDLY] Log the error but don't crash - # In-memory state is preserved and will continue to work - lib_logger.warning( - f"Failed to save usage data to {self.file_path}: {e}. " - "Data will be retained in memory but may be lost on restart." - ) + async with self._data_lock: + # Add human-readable timestamp fields before saving + self._add_readable_timestamps(self._usage_data) + # Hand off to resilient writer - handles retries and disk failures + self._state_writer.write(self._usage_data) async def _reset_daily_stats_if_needed(self): """ diff --git a/src/rotator_library/utils/__init__.py b/src/rotator_library/utils/__init__.py index 86a48dee..22d1ea78 100644 --- a/src/rotator_library/utils/__init__.py +++ b/src/rotator_library/utils/__init__.py @@ -2,5 +2,19 @@ from .headless_detection import is_headless_environment from .reauth_coordinator import get_reauth_coordinator, ReauthCoordinator +from .resilient_io import ( + ResilientStateWriter, + safe_write_json, + safe_log_write, + safe_mkdir, +) -__all__ = ["is_headless_environment", "get_reauth_coordinator", "ReauthCoordinator"] +__all__ = [ + "is_headless_environment", + "get_reauth_coordinator", + "ReauthCoordinator", + "ResilientStateWriter", + "safe_write_json", + "safe_log_write", + "safe_mkdir", +] diff --git a/src/rotator_library/utils/resilient_io.py b/src/rotator_library/utils/resilient_io.py new file mode 100644 index 00000000..47aa4ca4 --- /dev/null +++ b/src/rotator_library/utils/resilient_io.py @@ -0,0 +1,339 @@ +# src/rotator_library/utils/resilient_io.py +""" +Resilient I/O utilities for handling file operations gracefully. + +Provides two main patterns: +1. ResilientStateWriter - For stateful files (usage.json, credentials, cache) + that should be buffered in memory and retried on disk failure. +2. safe_log_write / safe_write_json - For logs that can be dropped on failure. +""" + +import json +import os +import shutil +import tempfile +import threading +import time +import logging +from pathlib import Path +from typing import Any, Callable, Dict, Optional, Union + + +class ResilientStateWriter: + """ + Manages resilient writes for stateful files (usage stats, credentials, cache). + + Design: + - Caller hands off data via write() - always succeeds (memory update) + - Attempts disk write immediately + - If disk fails, retries periodically in background + - On recovery, writes full current state (not just new data) + + Thread-safe for use in async contexts with sync file I/O. + + Usage: + writer = ResilientStateWriter("data.json", logger) + writer.write({"key": "value"}) # Always succeeds + # ... later ... + if not writer.is_healthy: + logger.warning("Disk writes failing, data in memory only") + """ + + def __init__( + self, + path: Union[str, Path], + logger: logging.Logger, + retry_interval: float = 30.0, + serializer: Optional[Callable[[Any], str]] = None, + ): + """ + Initialize the resilient writer. + + Args: + path: File path to write to + logger: Logger for warnings/errors + retry_interval: Seconds between retry attempts when disk is unhealthy + serializer: Custom serializer function (defaults to JSON with indent=2) + """ + self.path = Path(path) + self.logger = logger + self.retry_interval = retry_interval + self._serializer = serializer or (lambda d: json.dumps(d, indent=2)) + + self._current_state: Optional[Any] = None + self._disk_healthy = True + self._last_attempt: float = 0 + self._last_success: Optional[float] = None + self._failure_count = 0 + self._lock = threading.Lock() + + def write(self, data: Any) -> bool: + """ + Update state and attempt disk write. + + Always updates in-memory state (guaranteed to succeed). + Attempts disk write - if it fails, schedules for retry. + + Args: + data: Data to persist (must be serializable) + + Returns: + True if disk write succeeded, False if failed (data still in memory) + """ + with self._lock: + self._current_state = data + return self._try_disk_write() + + def retry_if_needed(self) -> bool: + """ + Retry disk write if unhealthy and retry interval has passed. + + Call this periodically (e.g., on each save attempt) to recover + from transient disk failures. + + Returns: + True if healthy (no retry needed or retry succeeded) + """ + with self._lock: + if self._disk_healthy: + return True + + if self._current_state is None: + return True + + now = time.time() + if now - self._last_attempt < self.retry_interval: + return False + + return self._try_disk_write() + + def _try_disk_write(self) -> bool: + """ + Attempt atomic write to disk. Updates health status. + + Uses tempfile + move pattern for atomic writes on POSIX systems. + On Windows, uses direct write (still safe for our use case). + """ + if self._current_state is None: + return True + + self._last_attempt = time.time() + + try: + # Ensure directory exists + self.path.parent.mkdir(parents=True, exist_ok=True) + + # Serialize data + content = self._serializer(self._current_state) + + # Atomic write: write to temp file, then move + tmp_fd = None + tmp_path = None + try: + tmp_fd, tmp_path = tempfile.mkstemp( + dir=self.path.parent, prefix=".tmp_", suffix=".json", text=True + ) + + with os.fdopen(tmp_fd, "w", encoding="utf-8") as f: + f.write(content) + tmp_fd = None # fdopen closes the fd + + # Atomic move + shutil.move(tmp_path, self.path) + tmp_path = None + + finally: + # Cleanup on failure + if tmp_fd is not None: + try: + os.close(tmp_fd) + except OSError: + pass + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except OSError: + pass + + # Success - update health + self._disk_healthy = True + self._last_success = time.time() + self._failure_count = 0 + return True + + except (OSError, PermissionError, IOError) as e: + self._disk_healthy = False + self._failure_count += 1 + + # Log warning (rate-limited to avoid flooding) + if self._failure_count == 1 or self._failure_count % 10 == 0: + self.logger.warning( + f"Failed to write {self.path.name}: {e}. " + f"Data retained in memory (failure #{self._failure_count})." + ) + return False + + @property + def is_healthy(self) -> bool: + """Check if disk writes are currently working.""" + return self._disk_healthy + + @property + def current_state(self) -> Optional[Any]: + """Get the current in-memory state (for inspection/debugging).""" + return self._current_state + + def get_health_info(self) -> Dict[str, Any]: + """ + Get detailed health information for monitoring. + + Returns dict with: + - healthy: bool + - failure_count: int + - last_success: Optional[float] (timestamp) + - last_attempt: float (timestamp) + - path: str + """ + return { + "healthy": self._disk_healthy, + "failure_count": self._failure_count, + "last_success": self._last_success, + "last_attempt": self._last_attempt, + "path": str(self.path), + } + + +def safe_write_json( + path: Union[str, Path], + data: Dict[str, Any], + logger: logging.Logger, + atomic: bool = True, + indent: int = 2, + ensure_ascii: bool = True, + secure_permissions: bool = False, +) -> bool: + """ + Write JSON data to file with error handling. No buffering or retry. + + Suitable for one-off writes where failure is acceptable (e.g., logs). + Creates parent directories if needed. + + Args: + path: File path to write to + data: JSON-serializable data + logger: Logger for warnings + atomic: Use atomic write pattern (tempfile + move) + indent: JSON indentation level (default: 2) + ensure_ascii: Escape non-ASCII characters (default: True) + secure_permissions: Set file permissions to 0o600 (default: False) + + Returns: + True on success, False on failure (never raises) + """ + path = Path(path) + + try: + path.parent.mkdir(parents=True, exist_ok=True) + content = json.dumps(data, indent=indent, ensure_ascii=ensure_ascii) + + if atomic: + tmp_fd = None + tmp_path = None + try: + tmp_fd, tmp_path = tempfile.mkstemp( + dir=path.parent, prefix=".tmp_", suffix=".json", text=True + ) + with os.fdopen(tmp_fd, "w", encoding="utf-8") as f: + f.write(content) + tmp_fd = None + + # Set secure permissions if requested (before move for security) + if secure_permissions: + try: + os.chmod(tmp_path, 0o600) + except (OSError, AttributeError): + # Windows may not support chmod, ignore + pass + + shutil.move(tmp_path, path) + tmp_path = None + finally: + if tmp_fd is not None: + try: + os.close(tmp_fd) + except OSError: + pass + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except OSError: + pass + else: + with open(path, "w", encoding="utf-8") as f: + f.write(content) + + # Set secure permissions if requested + if secure_permissions: + try: + os.chmod(path, 0o600) + except (OSError, AttributeError): + pass + + return True + + except (OSError, PermissionError, IOError, TypeError, ValueError) as e: + logger.warning(f"Failed to write JSON to {path}: {e}") + return False + + +def safe_log_write( + path: Union[str, Path], + content: str, + logger: logging.Logger, + mode: str = "a", +) -> bool: + """ + Write content to log file with error handling. No buffering or retry. + + Suitable for log files where occasional loss is acceptable. + Creates parent directories if needed. + + Args: + path: File path to write to + content: String content to write + logger: Logger for warnings + mode: File mode ('a' for append, 'w' for overwrite) + + Returns: + True on success, False on failure (never raises) + """ + path = Path(path) + + try: + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, mode, encoding="utf-8") as f: + f.write(content) + return True + + except (OSError, PermissionError, IOError) as e: + logger.warning(f"Failed to write log to {path}: {e}") + return False + + +def safe_mkdir(path: Union[str, Path], logger: logging.Logger) -> bool: + """ + Create directory with error handling. + + Args: + path: Directory path to create + logger: Logger for warnings + + Returns: + True on success (or already exists), False on failure + """ + try: + Path(path).mkdir(parents=True, exist_ok=True) + return True + except (OSError, PermissionError) as e: + logger.warning(f"Failed to create directory {path}: {e}") + return False From ea1e9f13f99a68b31c5500196fd3c4183ae0db16 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 10:02:00 +0100 Subject: [PATCH 111/221] =?UTF-8?q?feat(io):=20=E2=9C=A8=20add=20shutdown?= =?UTF-8?q?=20flush=20mechanism=20for=20buffered=20writes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces a global buffered write registry with automatic shutdown flush, ensuring critical data (auth tokens, usage stats) is saved even when disk writes fail temporarily. - Add `BufferedWriteRegistry` singleton for centralized buffered write management - Implement periodic retry (30s interval) and atexit shutdown flush for pending writes - Enable `buffer_on_failure` parameter in `safe_write_json()` for credential files - Integrate buffering with `ResilientStateWriter` for automatic registry registration - Update OAuth providers (Google, Qwen, iFlow) to use buffered credential writes - Change provider cache `_save_to_disk()` to return success status for better tracking - Reduce log noise by changing missing thoughtSignature warnings to debug level - Export `BufferedWriteRegistry` from utils module for monitoring access The new architecture ensures data is never lost on graceful shutdown (Ctrl+C), with console output showing flush progress and results. All buffered writes are retried in a background thread and guaranteed a final save attempt on application exit. --- DOCUMENTATION.md | 170 ++++++++- .../providers/antigravity_provider.py | 2 +- .../providers/gemini_cli_provider.py | 2 +- .../providers/google_oauth_base.py | 7 +- .../providers/iflow_auth_base.py | 7 +- .../providers/provider_cache.py | 18 +- .../providers/qwen_auth_base.py | 7 +- src/rotator_library/utils/__init__.py | 2 + src/rotator_library/utils/resilient_io.py | 348 +++++++++++++++++- 9 files changed, 525 insertions(+), 38 deletions(-) diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index f8060c32..30020176 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -939,31 +939,173 @@ This level of detail allows developers to trace exactly why a request failed or ## 5. Runtime Resilience -The proxy is engineered to maintain high availability even in the face of runtime filesystem disruptions. This "Runtime Resilience" capability ensures that the service continues to process API requests even if core data directories (like `logs/`, `oauth_creds/`) or files are accidentally deleted or become unwritable while the application is running. +The proxy is engineered to maintain high availability even in the face of runtime filesystem disruptions. This "Runtime Resilience" capability ensures that the service continues to process API requests even if data files or directories are deleted while the application is running. -### 5.1. Resilience Hierarchy +### 5.1. Centralized Resilient I/O (`resilient_io.py`) + +All file operations are centralized in a single utility module that provides consistent error handling, graceful degradation, and automatic retry with shutdown flush: + +#### `BufferedWriteRegistry` (Singleton) + +Global registry for buffered writes with periodic retry and shutdown flush. Ensures critical data is saved even if disk writes fail temporarily: + +- **Per-file buffering**: Each file path has its own pending write (latest data always wins) +- **Periodic retries**: Background thread retries failed writes every 30 seconds +- **Shutdown flush**: `atexit` hook ensures final write attempt on app exit (Ctrl+C) +- **Thread-safe**: Safe for concurrent access from multiple threads + +```python +# Get the singleton instance +registry = BufferedWriteRegistry.get_instance() + +# Check pending writes (for monitoring) +pending_count = registry.get_pending_count() +pending_files = registry.get_pending_paths() + +# Manual flush (optional - atexit handles this automatically) +results = registry.flush_all() # Returns {path: success_bool} + +# Manual shutdown (if needed before atexit) +results = registry.shutdown() +``` + +#### `ResilientStateWriter` + +For stateful files that must persist (usage stats): +- **Memory-first**: Always updates in-memory state before attempting disk write +- **Atomic writes**: Uses tempfile + move pattern to prevent corruption +- **Automatic retry with backoff**: If disk fails, waits `retry_interval` seconds before trying again +- **Shutdown integration**: Registers with `BufferedWriteRegistry` on failure for final flush +- **Health monitoring**: Exposes `is_healthy` property for monitoring + +```python +writer = ResilientStateWriter("data.json", logger, retry_interval=30.0) +writer.write({"key": "value"}) # Always succeeds (memory update) +if not writer.is_healthy: + logger.warning("Disk writes failing, data in memory only") +# On next write() call after retry_interval, disk write is attempted again +# On app exit (Ctrl+C), BufferedWriteRegistry attempts final save +``` + +#### `safe_write_json()` + +For JSON writes with configurable options (credentials, cache): + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `path` | required | File path to write to | +| `data` | required | JSON-serializable data | +| `logger` | required | Logger for warnings | +| `atomic` | `True` | Use atomic write pattern (tempfile + move) | +| `indent` | `2` | JSON indentation level | +| `ensure_ascii` | `True` | Escape non-ASCII characters | +| `secure_permissions` | `False` | Set file permissions to 0o600 | +| `buffer_on_failure` | `False` | Register with BufferedWriteRegistry on failure | + +When `buffer_on_failure=True`: +- Failed writes are registered with `BufferedWriteRegistry` +- Data is retried every 30 seconds in background +- On app exit, final write attempt is made automatically +- Success unregisters the pending write + +```python +# For critical data (auth tokens) - use buffer_on_failure +safe_write_json(path, creds, logger, secure_permissions=True, buffer_on_failure=True) + +# For non-critical data (logs) - no buffering needed +safe_write_json(path, data, logger) +``` + +#### `safe_log_write()` + +For log files where occasional loss is acceptable: +- Fire-and-forget pattern +- Creates parent directories if needed +- Returns `True`/`False`, never raises +- **No buffering** - logs are dropped on failure + +#### `safe_mkdir()` + +For directory creation with error handling. + +### 5.2. Resilience Hierarchy The system follows a strict hierarchy of survival: -1. **Core API Handling (Level 1)**: The Python runtime keeps all necessary code in memory (`sys.modules`). Deleting source code files while the proxy is running will **not** crash active requests. -2. **Credential Management (Level 2)**: OAuth tokens are aggressively cached in memory. If credential files are deleted, the proxy continues using the cached tokens. If a token needs refresh and the file cannot be written, the new token is updated in memory only. -3. **Usage Tracking (Level 3)**: Usage statistics (`key_usage.json`) are maintained in memory. If the file is deleted, the system tracks usage internally. It attempts to recreate the file/directory on the next save interval. If save fails, data is effectively "memory-only" until the next successful write. -4. **Logging (Level 4)**: Logging is treated as non-critical. If the `logs/` directory is removed, the system attempts to recreate it. If creation fails (e.g., permission error), logging degrades gracefully (stops or falls back to console) without interrupting the request flow. +1. **Core API Handling (Level 1)**: The Python runtime keeps all necessary code in memory. Deleting source code files while the proxy is running will **not** crash active requests. + +2. **Credential Management (Level 2)**: OAuth tokens are cached in memory first. If credential files are deleted, the proxy continues using cached tokens. If a token refresh succeeds but the file cannot be written, the new token is buffered for retry and saved on shutdown. + +3. **Usage Tracking (Level 3)**: Usage statistics (`key_usage.json`) are maintained in memory via `ResilientStateWriter`. If the file is deleted, the system tracks usage internally and attempts to recreate the file on the next save interval. Pending writes are flushed on shutdown. + +4. **Provider Cache (Level 4)**: The provider cache tracks disk health and continues operating in memory-only mode if disk writes fail. Has its own shutdown mechanism. + +5. **Logging (Level 5)**: Logging is treated as non-critical. If the `logs/` directory is removed, the system attempts to recreate it. If creation fails, logging degrades gracefully without interrupting the request flow. **No buffering or retry**. + +### 5.3. Component Integration -### 5.2. "Develop While Running" +| Component | Utility Used | Behavior on Disk Failure | Shutdown Flush | +|-----------|--------------|--------------------------|----------------| +| `UsageManager` | `ResilientStateWriter` | Continues in memory, retries after 30s | Yes (via registry) | +| `GoogleOAuthBase` | `safe_write_json(buffer_on_failure=True)` | Memory cache preserved, buffered for retry | Yes (via registry) | +| `QwenAuthBase` | `safe_write_json(buffer_on_failure=True)` | Memory cache preserved, buffered for retry | Yes (via registry) | +| `IFlowAuthBase` | `safe_write_json(buffer_on_failure=True)` | Memory cache preserved, buffered for retry | Yes (via registry) | +| `ProviderCache` | `safe_write_json` + own shutdown | Retries via own background loop | Yes (own mechanism) | +| `DetailedLogger` | `safe_write_json` | Logs dropped, no crash | No | +| `failure_logger` | Python `logging.RotatingFileHandler` | Falls back to NullHandler | No | + +### 5.4. Shutdown Behavior + +When the application exits (including Ctrl+C): + +1. **atexit handler fires**: `BufferedWriteRegistry._atexit_handler()` is called +2. **Pending writes counted**: Registry checks how many files have pending writes +3. **Flush attempted**: Each pending file gets a final write attempt +4. **Results logged**: + - Success: `"Shutdown flush: all N write(s) succeeded"` + - Partial: `"Shutdown flush: X succeeded, Y failed"` with failed file names + +**Console output example:** +``` +INFO:rotator_library.resilient_io:Flushing 2 pending write(s) on shutdown... +INFO:rotator_library.resilient_io:Shutdown flush: all 2 write(s) succeeded +``` + +### 5.5. "Develop While Running" This architecture supports a robust development workflow: -* **Log Cleanup**: You can safely run `rm -rf logs/` while the proxy is serving traffic. The system will simply recreate the directory structure on the next request. -* **Config Reset**: Deleting `key_usage.json` resets the persistence layer, but the running instance preserves its current in-memory counts to ensure load balancing consistency. -* **File Recovery**: If you delete a critical file, the system attempts **Directory Auto-Recreation** before every write operation. +- **Log Cleanup**: You can safely run `rm -rf logs/` while the proxy is serving traffic. The system will recreate the directory structure on the next request. +- **Config Reset**: Deleting `key_usage.json` resets the persistence layer, but the running instance preserves its current in-memory counts for load balancing consistency. +- **File Recovery**: If you delete a critical file, the system attempts directory auto-recreation before every write operation. +- **Safe Exit**: Ctrl+C triggers graceful shutdown with final data flush attempt. -### 5.3. Graceful Degradation & Data Loss +### 5.6. Graceful Degradation & Data Loss While functionality is preserved, persistence may be compromised during filesystem failures: -* **Logs**: If disk writes fail, detailed request logs may be lost (unless console fallback is active). -* **Usage Stats**: If `key_usage.json` cannot be written, usage data since the last successful save will be lost upon application restart. -* **Credentials**: Refreshed tokens held only in memory will require re-authentication after a restart if they cannot be persisted to disk. +- **Logs**: If disk writes fail, detailed request logs may be lost (no buffering). +- **Usage Stats**: Buffered in memory and flushed on shutdown. Data loss only if shutdown flush also fails. +- **Credentials**: Buffered in memory and flushed on shutdown. Re-authentication only needed if shutdown flush fails. +- **Cache**: Provider cache entries may need to be regenerated after restart if its own shutdown mechanism fails. + +### 5.7. Monitoring Disk Health + +Components expose health information for monitoring: +```python +# BufferedWriteRegistry +registry = BufferedWriteRegistry.get_instance() +pending = registry.get_pending_count() # Number of files with pending writes +files = registry.get_pending_paths() # List of pending file names + +# UsageManager +writer = usage_manager._state_writer +health = writer.get_health_info() +# Returns: {"healthy": True, "failure_count": 0, "last_success": 1234567890.0, ...} + +# ProviderCache +stats = cache.get_stats() +# Includes: {"disk_available": True, "disk_errors": 0, ...} +``` diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index ebf950ee..3a803fdf 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -2424,7 +2424,7 @@ def _transform_assistant_message( elif first_func_in_msg: # Only add bypass to the first function call if no sig available func_part["thoughtSignature"] = "skip_thought_signature_validator" - lib_logger.warning( + lib_logger.debug( f"Missing thoughtSignature for first func call {tool_id}, using bypass" ) # Subsequent parallel calls: no signature field at all diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 52f15d68..64791b29 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -1166,7 +1166,7 @@ def _transform_messages( func_part["thoughtSignature"] = ( "skip_thought_signature_validator" ) - lib_logger.warning( + lib_logger.debug( f"Missing thoughtSignature for first func call {tool_id}, using bypass" ) # Subsequent parallel calls: no signature field at all diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index 5f8a09b3..ba99b96d 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -273,13 +273,16 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): return # Attempt disk write - if it fails, we still have the cache - if safe_write_json(path, creds, lib_logger, secure_permissions=True): + # buffer_on_failure ensures data is retried periodically and saved on shutdown + if safe_write_json( + path, creds, lib_logger, secure_permissions=True, buffer_on_failure=True + ): lib_logger.debug( f"Saved updated {self.ENV_PREFIX} OAuth credentials to '{path}'." ) else: lib_logger.warning( - f"Credentials for {self.ENV_PREFIX} cached in memory only (will be lost on restart)." + f"Credentials for {self.ENV_PREFIX} cached in memory only (buffered for retry)." ) def _is_token_expired(self, creds: Dict[str, Any]) -> bool: diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index a2096df3..29258138 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -325,11 +325,14 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): return # Attempt disk write - if it fails, we still have the cache - if safe_write_json(path, creds, lib_logger, secure_permissions=True): + # buffer_on_failure ensures data is retried periodically and saved on shutdown + if safe_write_json( + path, creds, lib_logger, secure_permissions=True, buffer_on_failure=True + ): lib_logger.debug(f"Saved updated iFlow OAuth credentials to '{path}'.") else: lib_logger.warning( - "iFlow credentials cached in memory only (will be lost on restart)." + "iFlow credentials cached in memory only (buffered for retry)." ) def _is_token_expired(self, creds: Dict[str, Any]) -> bool: diff --git a/src/rotator_library/providers/provider_cache.py b/src/rotator_library/providers/provider_cache.py index 8b0f835b..1fc94374 100644 --- a/src/rotator_library/providers/provider_cache.py +++ b/src/rotator_library/providers/provider_cache.py @@ -197,10 +197,14 @@ async def _load_from_disk(self) -> None: # DISK PERSISTENCE # ========================================================================= - async def _save_to_disk(self) -> None: - """Persist cache to disk using atomic write with health tracking.""" + async def _save_to_disk(self) -> bool: + """Persist cache to disk using atomic write with health tracking. + + Returns: + True if write succeeded, False otherwise. + """ if not self._enable_disk: - return + return True # Not an error if disk is disabled async with self._disk_lock: cache_data = { @@ -226,9 +230,11 @@ async def _save_to_disk(self) -> None: lib_logger.debug( f"ProviderCache[{self._cache_name}]: Saved {len(self._cache)} entries" ) + return True else: self._stats["disk_errors"] += 1 self._disk_available = False + return False # ========================================================================= # BACKGROUND TASKS @@ -251,8 +257,10 @@ async def _writer_loop(self) -> None: await asyncio.sleep(self._write_interval) if self._dirty: try: - await self._save_to_disk() - self._dirty = False + success = await self._save_to_disk() + if success: + self._dirty = False + # If save failed, _dirty remains True so we retry next interval except Exception as e: lib_logger.error( f"ProviderCache[{self._cache_name}]: Writer error: {e}" diff --git a/src/rotator_library/providers/qwen_auth_base.py b/src/rotator_library/providers/qwen_auth_base.py index b95416a5..df07b776 100644 --- a/src/rotator_library/providers/qwen_auth_base.py +++ b/src/rotator_library/providers/qwen_auth_base.py @@ -210,11 +210,14 @@ async def _save_credentials(self, path: str, creds: Dict[str, Any]): return # Attempt disk write - if it fails, we still have the cache - if safe_write_json(path, creds, lib_logger, secure_permissions=True): + # buffer_on_failure ensures data is retried periodically and saved on shutdown + if safe_write_json( + path, creds, lib_logger, secure_permissions=True, buffer_on_failure=True + ): lib_logger.debug(f"Saved updated Qwen OAuth credentials to '{path}'.") else: lib_logger.warning( - "Qwen credentials cached in memory only (will be lost on restart)." + "Qwen credentials cached in memory only (buffered for retry)." ) def _is_token_expired(self, creds: Dict[str, Any]) -> bool: diff --git a/src/rotator_library/utils/__init__.py b/src/rotator_library/utils/__init__.py index 22d1ea78..fa3bb12c 100644 --- a/src/rotator_library/utils/__init__.py +++ b/src/rotator_library/utils/__init__.py @@ -3,6 +3,7 @@ from .headless_detection import is_headless_environment from .reauth_coordinator import get_reauth_coordinator, ReauthCoordinator from .resilient_io import ( + BufferedWriteRegistry, ResilientStateWriter, safe_write_json, safe_log_write, @@ -13,6 +14,7 @@ "is_headless_environment", "get_reauth_coordinator", "ReauthCoordinator", + "BufferedWriteRegistry", "ResilientStateWriter", "safe_write_json", "safe_log_write", diff --git a/src/rotator_library/utils/resilient_io.py b/src/rotator_library/utils/resilient_io.py index 47aa4ca4..a9c623a7 100644 --- a/src/rotator_library/utils/resilient_io.py +++ b/src/rotator_library/utils/resilient_io.py @@ -2,12 +2,17 @@ """ Resilient I/O utilities for handling file operations gracefully. -Provides two main patterns: -1. ResilientStateWriter - For stateful files (usage.json, credentials, cache) - that should be buffered in memory and retried on disk failure. -2. safe_log_write / safe_write_json - For logs that can be dropped on failure. +Provides three main patterns: +1. BufferedWriteRegistry - Global singleton for buffered writes with periodic + retry and shutdown flush. Ensures data is saved on app exit (Ctrl+C). +2. ResilientStateWriter - For stateful files (usage.json) that should be + buffered in memory and retried on disk failure. +3. safe_write_json (with buffer_on_failure) - For critical files (auth tokens) + that should be buffered and retried if write fails. +4. safe_log_write - For logs that can be dropped on failure. """ +import atexit import json import os import shutil @@ -16,7 +21,284 @@ import time import logging from pathlib import Path -from typing import Any, Callable, Dict, Optional, Union +from typing import Any, Callable, Dict, Optional, Tuple, Union + + +# ============================================================================= +# BUFFERED WRITE REGISTRY (SINGLETON) +# ============================================================================= + + +class BufferedWriteRegistry: + """ + Global singleton registry for buffered writes with periodic retry and shutdown flush. + + This ensures that critical data (auth tokens, usage stats) is saved even if + disk writes fail temporarily. On app exit (including Ctrl+C), all pending + writes are flushed. + + Features: + - Per-file buffering: each file path has its own pending write + - Periodic retries: background thread retries failed writes every N seconds + - Shutdown flush: atexit hook ensures final write attempt on app exit + - Thread-safe: safe for concurrent access from multiple threads + + Usage: + # Get the singleton instance + registry = BufferedWriteRegistry.get_instance() + + # Register a pending write (usually called by safe_write_json on failure) + registry.register_pending(path, data, serializer_fn, options) + + # Manual flush (optional - atexit handles this automatically) + results = registry.flush_all() + """ + + _instance: Optional["BufferedWriteRegistry"] = None + _instance_lock = threading.Lock() + + def __init__(self, retry_interval: float = 30.0): + """ + Initialize the registry. Use get_instance() instead of direct construction. + + Args: + retry_interval: Seconds between retry attempts (default: 30) + """ + self._pending: Dict[str, Tuple[Any, Callable[[Any], str], Dict[str, Any]]] = {} + self._retry_interval = retry_interval + self._lock = threading.Lock() + self._running = False + self._retry_thread: Optional[threading.Thread] = None + self._logger = logging.getLogger("rotator_library.resilient_io") + + # Start background retry thread + self._start_retry_thread() + + # Register atexit handler for shutdown flush + atexit.register(self._atexit_handler) + + @classmethod + def get_instance(cls, retry_interval: float = 30.0) -> "BufferedWriteRegistry": + """ + Get or create the singleton instance. + + Args: + retry_interval: Seconds between retry attempts (only used on first call) + + Returns: + The singleton BufferedWriteRegistry instance + """ + if cls._instance is None: + with cls._instance_lock: + if cls._instance is None: + cls._instance = cls(retry_interval) + return cls._instance + + def _start_retry_thread(self) -> None: + """Start the background retry thread.""" + if self._running: + return + + self._running = True + self._retry_thread = threading.Thread( + target=self._retry_loop, + name="BufferedWriteRegistry-Retry", + daemon=True, # Daemon so it doesn't block app exit + ) + self._retry_thread.start() + + def _retry_loop(self) -> None: + """Background thread: periodically retry pending writes.""" + while self._running: + time.sleep(self._retry_interval) + if not self._running: + break + self._retry_pending() + + def _retry_pending(self) -> None: + """Attempt to write all pending files.""" + with self._lock: + if not self._pending: + return + + # Copy paths to avoid modifying dict during iteration + paths = list(self._pending.keys()) + + for path_str in paths: + self._try_write(path_str, remove_on_success=True) + + def register_pending( + self, + path: Union[str, Path], + data: Any, + serializer: Callable[[Any], str], + options: Optional[Dict[str, Any]] = None, + ) -> None: + """ + Register a pending write for later retry. + + If a write is already pending for this path, it is replaced with the new data + (we always want to write the latest state). + + Args: + path: File path to write to + data: Data to serialize and write + serializer: Function to serialize data to string + options: Additional options (e.g., secure_permissions) + """ + path_str = str(Path(path).resolve()) + with self._lock: + self._pending[path_str] = (data, serializer, options or {}) + self._logger.debug(f"Registered pending write for {Path(path).name}") + + def unregister(self, path: Union[str, Path]) -> None: + """ + Remove a pending write (called when write succeeds elsewhere). + + Args: + path: File path to remove from pending + """ + path_str = str(Path(path).resolve()) + with self._lock: + self._pending.pop(path_str, None) + + def _try_write(self, path_str: str, remove_on_success: bool = True) -> bool: + """ + Attempt to write a pending file. + + Args: + path_str: Resolved path string + remove_on_success: Remove from pending if successful + + Returns: + True if write succeeded, False otherwise + """ + with self._lock: + if path_str not in self._pending: + return True + data, serializer, options = self._pending[path_str] + + path = Path(path_str) + try: + # Ensure directory exists + path.parent.mkdir(parents=True, exist_ok=True) + + # Serialize data + content = serializer(data) + + # Atomic write + tmp_fd = None + tmp_path = None + try: + tmp_fd, tmp_path = tempfile.mkstemp( + dir=path.parent, prefix=".tmp_", suffix=".json", text=True + ) + with os.fdopen(tmp_fd, "w", encoding="utf-8") as f: + f.write(content) + tmp_fd = None + + # Set secure permissions if requested + if options.get("secure_permissions"): + try: + os.chmod(tmp_path, 0o600) + except (OSError, AttributeError): + pass + + shutil.move(tmp_path, path) + tmp_path = None + + finally: + if tmp_fd is not None: + try: + os.close(tmp_fd) + except OSError: + pass + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except OSError: + pass + + # Success - remove from pending + if remove_on_success: + with self._lock: + self._pending.pop(path_str, None) + + self._logger.debug(f"Retry succeeded for {path.name}") + return True + + except (OSError, PermissionError, IOError) as e: + self._logger.debug(f"Retry failed for {path.name}: {e}") + return False + + def flush_all(self) -> Dict[str, bool]: + """ + Attempt to write all pending files immediately. + + Returns: + Dict mapping file paths to success status + """ + with self._lock: + paths = list(self._pending.keys()) + + results = {} + for path_str in paths: + results[path_str] = self._try_write(path_str, remove_on_success=True) + + return results + + def _atexit_handler(self) -> None: + """Called on app exit to flush pending writes.""" + self._running = False + + with self._lock: + pending_count = len(self._pending) + + if pending_count == 0: + return + + self._logger.info(f"Flushing {pending_count} pending write(s) on shutdown...") + results = self.flush_all() + + succeeded = sum(1 for v in results.values() if v) + failed = pending_count - succeeded + + if failed > 0: + self._logger.warning( + f"Shutdown flush: {succeeded} succeeded, {failed} failed" + ) + for path_str, success in results.items(): + if not success: + self._logger.warning(f" Failed to save: {Path(path_str).name}") + else: + self._logger.info(f"Shutdown flush: all {succeeded} write(s) succeeded") + + def get_pending_count(self) -> int: + """Get the number of pending writes.""" + with self._lock: + return len(self._pending) + + def get_pending_paths(self) -> list: + """Get list of paths with pending writes (for monitoring).""" + with self._lock: + return [Path(p).name for p in self._pending.keys()] + + def shutdown(self) -> Dict[str, bool]: + """ + Manually trigger shutdown: stop retry thread and flush all pending writes. + + Returns: + Dict mapping file paths to success status + """ + self._running = False + if self._retry_thread and self._retry_thread.is_alive(): + self._retry_thread.join(timeout=1.0) + return self.flush_all() + + +# ============================================================================= +# RESILIENT STATE WRITER +# ============================================================================= class ResilientStateWriter: @@ -72,7 +354,8 @@ def write(self, data: Any) -> bool: Update state and attempt disk write. Always updates in-memory state (guaranteed to succeed). - Attempts disk write - if it fails, schedules for retry. + Attempts disk write - if disk is unhealthy, respects retry_interval + before attempting again to avoid flooding with failed writes. Args: data: Data to persist (must be serializable) @@ -82,6 +365,14 @@ def write(self, data: Any) -> bool: """ with self._lock: self._current_state = data + + # If disk is unhealthy, only retry after retry_interval has passed + if not self._disk_healthy: + now = time.time() + if now - self._last_attempt < self.retry_interval: + # Too soon to retry, data is safe in memory + return False + return self._try_disk_write() def retry_if_needed(self) -> bool: @@ -113,6 +404,8 @@ def _try_disk_write(self) -> bool: Uses tempfile + move pattern for atomic writes on POSIX systems. On Windows, uses direct write (still safe for our use case). + + Also registers/unregisters with BufferedWriteRegistry for shutdown flush. """ if self._current_state is None: return True @@ -155,16 +448,26 @@ def _try_disk_write(self) -> bool: except OSError: pass - # Success - update health + # Success - update health and unregister from shutdown flush self._disk_healthy = True self._last_success = time.time() self._failure_count = 0 + BufferedWriteRegistry.get_instance().unregister(self.path) return True except (OSError, PermissionError, IOError) as e: self._disk_healthy = False self._failure_count += 1 + # Register with BufferedWriteRegistry for shutdown flush + registry = BufferedWriteRegistry.get_instance() + registry.register_pending( + self.path, + self._current_state, + self._serializer, + {}, # No special options for ResilientStateWriter + ) + # Log warning (rate-limited to avoid flooding) if self._failure_count == 1 or self._failure_count % 10 == 0: self.logger.warning( @@ -211,12 +514,14 @@ def safe_write_json( indent: int = 2, ensure_ascii: bool = True, secure_permissions: bool = False, + buffer_on_failure: bool = False, ) -> bool: """ - Write JSON data to file with error handling. No buffering or retry. + Write JSON data to file with error handling and optional buffering. - Suitable for one-off writes where failure is acceptable (e.g., logs). - Creates parent directories if needed. + When buffer_on_failure is True, failed writes are registered with the + BufferedWriteRegistry for periodic retry and shutdown flush. This ensures + critical data (like auth tokens) is eventually saved. Args: path: File path to write to @@ -226,15 +531,20 @@ def safe_write_json( indent: JSON indentation level (default: 2) ensure_ascii: Escape non-ASCII characters (default: True) secure_permissions: Set file permissions to 0o600 (default: False) + buffer_on_failure: Register with BufferedWriteRegistry on failure (default: False) Returns: True on success, False on failure (never raises) """ path = Path(path) + # Create serializer function that matches the requested formatting + def serializer(d: Any) -> str: + return json.dumps(d, indent=indent, ensure_ascii=ensure_ascii) + try: path.parent.mkdir(parents=True, exist_ok=True) - content = json.dumps(data, indent=indent, ensure_ascii=ensure_ascii) + content = serializer(data) if atomic: tmp_fd = None @@ -279,10 +589,26 @@ def safe_write_json( except (OSError, AttributeError): pass + # Success - remove from pending if it was there + if buffer_on_failure: + BufferedWriteRegistry.get_instance().unregister(path) + return True except (OSError, PermissionError, IOError, TypeError, ValueError) as e: logger.warning(f"Failed to write JSON to {path}: {e}") + + # Register for retry if buffering is enabled + if buffer_on_failure: + registry = BufferedWriteRegistry.get_instance() + registry.register_pending( + path, + data, + serializer, + {"secure_permissions": secure_permissions}, + ) + logger.debug(f"Buffered {path.name} for retry on next interval or shutdown") + return False From 2ef272f3cd98eb2329f524ac5df115753e71a889 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 10:15:51 +0100 Subject: [PATCH 112/221] =?UTF-8?q?fix(auth):=20=F0=9F=94=A8=20prioritize?= =?UTF-8?q?=20file-based=20credential=20loading=20over=20environment=20var?= =?UTF-8?q?iables?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Change credential loading strategy across all auth providers to prefer file-based credentials when an explicit file path is provided, falling back to legacy environment variables only when the file is not found. - Modified `_load_credentials()` in GoogleOAuthBase, IFlowAuthBase, and QwenAuthBase to attempt file loading first - Environment variable fallback now only triggers on FileNotFoundError, improving error clarity - Removed redundant exception handling in GoogleOAuthBase (duplicate catch blocks) - Fixed potential deadlock in credential refresh queue by removing nested lock acquisition - _refresh_token() already handles its own locking, so removed outer lock to prevent deadlock - Improved logging to indicate when fallback to environment variables occurs - Maintains backwards compatibility for existing deployments using environment variables This change addresses two issues: 1. Ensures explicit file paths are respected as the primary credential source 2. Prevents deadlock scenario where refresh queue would acquire lock before calling _refresh_token() which also acquires the same lock --- .../providers/google_oauth_base.py | 67 +++++++++---------- .../providers/iflow_auth_base.py | 64 +++++++++--------- .../providers/qwen_auth_base.py | 66 +++++++++--------- 3 files changed, 99 insertions(+), 98 deletions(-) diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index ba99b96d..1a3f5d92 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -227,17 +227,7 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: f"Environment variables for {self.ENV_PREFIX} credential index {credential_index} not found" ) - # For file paths, first try loading from legacy env vars (for backwards compatibility) - env_creds = self._load_from_env() - if env_creds: - lib_logger.info( - f"Using {self.ENV_PREFIX} credentials from environment variables" - ) - # Cache env-based credentials using the path as key - self._credentials_cache[path] = env_creds - return env_creds - - # Fall back to file-based loading + # Try file-based loading first (preferred for explicit file paths) try: lib_logger.debug( f"Loading {self.ENV_PREFIX} credentials from file: {path}" @@ -250,6 +240,15 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: self._credentials_cache[path] = creds return creds except FileNotFoundError: + # File not found - fall back to legacy env vars for backwards compatibility + # This handles the case where only env vars are set and file paths are placeholders + env_creds = self._load_from_env() + if env_creds: + lib_logger.info( + f"File '{path}' not found, using {self.ENV_PREFIX} credentials from environment variables" + ) + self._credentials_cache[path] = env_creds + return env_creds raise IOError( f"{self.ENV_PREFIX} OAuth credential file not found at '{path}'" ) @@ -257,10 +256,6 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: raise IOError( f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}" ) - except Exception as e: - raise IOError( - f"Failed to load {self.ENV_PREFIX} OAuth credentials from '{path}': {e}" - ) async def _save_credentials(self, path: str, creds: Dict[str, Any]): """Save credentials with in-memory fallback if disk unavailable.""" @@ -588,32 +583,32 @@ async def _process_refresh_queue(self): return try: - # Perform the actual refresh (still using per-credential lock) - async with await self._get_lock(path): - # Re-check if still expired (may have changed since queueing) - creds = self._credentials_cache.get(path) - if creds and not self._is_token_expired(creds): - # No longer expired, mark as available - async with self._queue_tracking_lock: - self._unavailable_credentials.pop(path, None) - lib_logger.debug( - f"Credential '{Path(path).name}' no longer expired, marked available. " - f"Remaining unavailable: {len(self._unavailable_credentials)}" - ) - continue - - # Perform refresh - if not creds: - creds = await self._load_credentials(path) - await self._refresh_token(path, creds, force=force) - - # SUCCESS: Mark as available again + # Quick check if still expired (optimization to avoid unnecessary refresh) + # Note: _refresh_token() will do its own locking and expiry check + creds = self._credentials_cache.get(path) + if creds and not self._is_token_expired(creds): + # No longer expired, mark as available async with self._queue_tracking_lock: self._unavailable_credentials.pop(path, None) lib_logger.debug( - f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Credential '{Path(path).name}' no longer expired, marked available. " f"Remaining unavailable: {len(self._unavailable_credentials)}" ) + continue + + # Perform refresh - _refresh_token handles its own locking + # DO NOT acquire lock here as _refresh_token also acquires it (would deadlock) + if not creds: + creds = await self._load_credentials(path) + await self._refresh_token(path, creds, force=force) + + # SUCCESS: Mark as available again + async with self._queue_tracking_lock: + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) finally: # [FIX PR#34] Remove from BOTH queued set AND unavailable credentials diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index 29258138..8854c493 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -304,15 +304,19 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: f"Environment variables for iFlow credential index {credential_index} not found" ) - # For file paths, try loading from legacy env vars first - env_creds = self._load_from_env() - if env_creds: - lib_logger.info("Using iFlow credentials from environment variables") - self._credentials_cache[path] = env_creds - return env_creds - - # Fall back to file-based loading - return await self._read_creds_from_file(path) + # Try file-based loading first (preferred for explicit file paths) + try: + return await self._read_creds_from_file(path) + except IOError: + # File not found - fall back to legacy env vars for backwards compatibility + env_creds = self._load_from_env() + if env_creds: + lib_logger.info( + f"File '{path}' not found, using iFlow credentials from environment variables" + ) + self._credentials_cache[path] = env_creds + return env_creds + raise # Re-raise the original file not found error async def _save_credentials(self, path: str, creds: Dict[str, Any]): """Save credentials with in-memory fallback if disk unavailable.""" @@ -843,32 +847,32 @@ async def _process_refresh_queue(self): return try: - # Perform the actual refresh (still using per-credential lock) - async with await self._get_lock(path): - # Re-check if still expired (may have changed since queueing) - creds = self._credentials_cache.get(path) - if creds and not self._is_token_expired(creds): - # No longer expired, mark as available - async with self._queue_tracking_lock: - self._unavailable_credentials.pop(path, None) - lib_logger.debug( - f"Credential '{Path(path).name}' no longer expired, marked available. " - f"Remaining unavailable: {len(self._unavailable_credentials)}" - ) - continue - - # Perform refresh - if not creds: - creds = await self._load_credentials(path) - await self._refresh_token(path, force=force) - - # SUCCESS: Mark as available again + # Quick check if still expired (optimization to avoid unnecessary refresh) + # Note: _refresh_token() will do its own locking and expiry check + creds = self._credentials_cache.get(path) + if creds and not self._is_token_expired(creds): + # No longer expired, mark as available async with self._queue_tracking_lock: self._unavailable_credentials.pop(path, None) lib_logger.debug( - f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Credential '{Path(path).name}' no longer expired, marked available. " f"Remaining unavailable: {len(self._unavailable_credentials)}" ) + continue + + # Perform refresh - _refresh_token handles its own locking + # DO NOT acquire lock here as _refresh_token also acquires it (would deadlock) + if not creds: + creds = await self._load_credentials(path) + await self._refresh_token(path, force=force) + + # SUCCESS: Mark as available again + async with self._queue_tracking_lock: + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) finally: # [FIX PR#34] Remove from BOTH queued set AND unavailable credentials diff --git a/src/rotator_library/providers/qwen_auth_base.py b/src/rotator_library/providers/qwen_auth_base.py index df07b776..28657f74 100644 --- a/src/rotator_library/providers/qwen_auth_base.py +++ b/src/rotator_library/providers/qwen_auth_base.py @@ -187,17 +187,19 @@ async def _load_credentials(self, path: str) -> Dict[str, Any]: f"Environment variables for Qwen Code credential index {credential_index} not found" ) - # For file paths, try loading from legacy env vars first - env_creds = self._load_from_env() - if env_creds: - lib_logger.info( - "Using Qwen Code credentials from environment variables" - ) - self._credentials_cache[path] = env_creds - return env_creds - - # Fall back to file-based loading - return await self._read_creds_from_file(path) + # Try file-based loading first (preferred for explicit file paths) + try: + return await self._read_creds_from_file(path) + except IOError: + # File not found - fall back to legacy env vars for backwards compatibility + env_creds = self._load_from_env() + if env_creds: + lib_logger.info( + f"File '{path}' not found, using Qwen Code credentials from environment variables" + ) + self._credentials_cache[path] = env_creds + return env_creds + raise # Re-raise the original file not found error async def _save_credentials(self, path: str, creds: Dict[str, Any]): """Save credentials with in-memory fallback if disk unavailable.""" @@ -571,32 +573,32 @@ async def _process_refresh_queue(self): return try: - # Perform the actual refresh (still using per-credential lock) - async with await self._get_lock(path): - # Re-check if still expired (may have changed since queueing) - creds = self._credentials_cache.get(path) - if creds and not self._is_token_expired(creds): - # No longer expired, mark as available - async with self._queue_tracking_lock: - self._unavailable_credentials.pop(path, None) - lib_logger.debug( - f"Credential '{Path(path).name}' no longer expired, marked available. " - f"Remaining unavailable: {len(self._unavailable_credentials)}" - ) - continue - - # Perform refresh - if not creds: - creds = await self._load_credentials(path) - await self._refresh_token(path, force=force) - - # SUCCESS: Mark as available again + # Quick check if still expired (optimization to avoid unnecessary refresh) + # Note: _refresh_token() will do its own locking and expiry check + creds = self._credentials_cache.get(path) + if creds and not self._is_token_expired(creds): + # No longer expired, mark as available async with self._queue_tracking_lock: self._unavailable_credentials.pop(path, None) lib_logger.debug( - f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Credential '{Path(path).name}' no longer expired, marked available. " f"Remaining unavailable: {len(self._unavailable_credentials)}" ) + continue + + # Perform refresh - _refresh_token handles its own locking + # DO NOT acquire lock here as _refresh_token also acquires it (would deadlock) + if not creds: + creds = await self._load_credentials(path) + await self._refresh_token(path, force=force) + + # SUCCESS: Mark as available again + async with self._queue_tracking_lock: + self._unavailable_credentials.pop(path, None) + lib_logger.debug( + f"Refresh SUCCESS for '{Path(path).name}', marked available. " + f"Remaining unavailable: {len(self._unavailable_credentials)}" + ) finally: # [FIX PR#34] Remove from BOTH queued set AND unavailable credentials From 683c1c110208458911180afd534afa5ac66cea85 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 10:28:49 +0100 Subject: [PATCH 113/221] =?UTF-8?q?refactor(providers):=20=F0=9F=94=A8=20i?= =?UTF-8?q?mprove=20error=20handling=20and=20logging=20specificity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Change antigravity cache logging from info to debug level to reduce noise - Replace Gemini CLI's delegated error parsing with native implementation - Add comprehensive duration parsing for multiple time formats (2s, 156h14m36s, 515092.73s) - Extract retry timing from human-readable error messages instead of relying on structured metadata - Improve error body extraction with multiple fallback strategies The Gemini CLI provider now handles its own quota error parsing rather than delegating to Antigravity, since the two providers use fundamentally different error formats: Gemini embeds reset times in human-readable messages while Antigravity uses structured RetryInfo/quotaResetDelay metadata. --- .../providers/antigravity_provider.py | 2 +- .../providers/gemini_cli_provider.py | 145 +++++++++++++++++- 2 files changed, 138 insertions(+), 9 deletions(-) diff --git a/src/rotator_library/providers/antigravity_provider.py b/src/rotator_library/providers/antigravity_provider.py index 3a803fdf..2a29509b 100644 --- a/src/rotator_library/providers/antigravity_provider.py +++ b/src/rotator_library/providers/antigravity_provider.py @@ -3405,7 +3405,7 @@ def _cache_thinking( } self._thinking_cache.store(cache_key, json.dumps(data)) - lib_logger.info(f"Cached thinking: {cache_key[:50]}...") + lib_logger.debug(f"Cached thinking: {cache_key[:50]}...") # ========================================================================= # PROVIDER INTERFACE IMPLEMENTATION diff --git a/src/rotator_library/providers/gemini_cli_provider.py b/src/rotator_library/providers/gemini_cli_provider.py index 64791b29..1d4588ea 100644 --- a/src/rotator_library/providers/gemini_cli_provider.py +++ b/src/rotator_library/providers/gemini_cli_provider.py @@ -234,22 +234,151 @@ def parse_quota_error( error: Exception, error_body: Optional[str] = None ) -> Optional[Dict[str, Any]]: """ - Parse Gemini CLI quota errors. - - Uses the same Google RPC format as Antigravity but typically has - much shorter cooldown durations (seconds to minutes, not hours). + Parse Gemini CLI rate limit/quota errors. + + Handles the Gemini CLI error format which embeds reset time in the message: + "You have exhausted your capacity on this model. Your quota will reset after 2s." + + Unlike Antigravity which uses structured RetryInfo/quotaResetDelay metadata, + Gemini CLI embeds the reset time in a human-readable message. + + Example error format: + { + "error": { + "code": 429, + "message": "You have exhausted your capacity on this model. Your quota will reset after 2s.", + "status": "RESOURCE_EXHAUSTED", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.ErrorInfo", + "reason": "RATE_LIMIT_EXCEEDED", + "domain": "cloudcode-pa.googleapis.com", + "metadata": { "uiMessage": "true", "model": "gemini-3-pro-preview" } + } + ] + } + } Args: error: The caught exception error_body: Optional raw response body string Returns: - Same format as AntigravityProvider.parse_quota_error() + None if not a parseable quota error, otherwise: + { + "retry_after": int, + "reason": str | None, + "reset_timestamp": str | None, + "quota_reset_timestamp": float | None, + } """ - # Reuse the same parsing logic as Antigravity since both use Google RPC format - from .antigravity_provider import AntigravityProvider + import re as regex_module + + # Get error body from exception if not provided + body = error_body + if not body: + if hasattr(error, "response") and hasattr(error.response, "text"): + try: + body = error.response.text + except Exception: + pass + if not body and hasattr(error, "body"): + body = str(error.body) + if not body and hasattr(error, "message"): + body = str(error.message) + if not body: + body = str(error) + + if not body: + return None + + result = { + "retry_after": None, + "reason": None, + "reset_timestamp": None, + "quota_reset_timestamp": None, + } + + # 1. Try to extract retry time from human-readable message + # Pattern: "Your quota will reset after 2s." or "quota will reset after 156h14m36s" + retry_after = extract_retry_after_from_body(body) + if retry_after: + result["retry_after"] = retry_after + + # 2. Try to parse JSON to get structured details (reason, any RetryInfo fallback) + try: + json_match = regex_module.search(r"\{[\s\S]*\}", body) + if json_match: + data = json.loads(json_match.group(0)) + error_obj = data.get("error", data) + details = error_obj.get("details", []) + + for detail in details: + detail_type = detail.get("@type", "") + + # Extract reason from ErrorInfo + if "ErrorInfo" in detail_type: + if not result["reason"]: + result["reason"] = detail.get("reason") + # Check metadata for any additional timing info + metadata = detail.get("metadata", {}) + quota_delay = metadata.get("quotaResetDelay") + if quota_delay and not result["retry_after"]: + parsed = GeminiCliProvider._parse_duration(quota_delay) + if parsed: + result["retry_after"] = parsed + + # Check for RetryInfo (fallback, in case format changes) + if "RetryInfo" in detail_type and not result["retry_after"]: + retry_delay = detail.get("retryDelay") + if retry_delay: + parsed = GeminiCliProvider._parse_duration(retry_delay) + if parsed: + result["retry_after"] = parsed + + except (json.JSONDecodeError, AttributeError, TypeError): + pass + + # Return None if we couldn't extract retry_after + if not result["retry_after"]: + return None + + return result + + @staticmethod + def _parse_duration(duration_str: str) -> Optional[int]: + """ + Parse duration strings like '2s', '156h14m36.73s', '515092.73s' to seconds. + + Args: + duration_str: Duration string to parse + + Returns: + Total seconds as integer, or None if parsing fails + """ + import re as regex_module + + if not duration_str: + return None - return AntigravityProvider.parse_quota_error(error, error_body) + # Handle pure seconds format: "515092.730699158s" or "2s" + pure_seconds_match = regex_module.match(r"^([\d.]+)s$", duration_str) + if pure_seconds_match: + return int(float(pure_seconds_match.group(1))) + + # Handle compound format: "143h4m52.730699158s" + total_seconds = 0 + patterns = [ + (r"(\d+)h", 3600), # hours + (r"(\d+)m", 60), # minutes + (r"([\d.]+)s", 1), # seconds + ] + for pattern, multiplier in patterns: + match = regex_module.search(pattern, duration_str) + if match: + total_seconds += float(match.group(1)) * multiplier + + return int(total_seconds) if total_seconds > 0 else None def __init__(self): super().__init__() From 92211ea358b21773e988517b3112189cc808f90b Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 11:05:24 +0100 Subject: [PATCH 114/221] =?UTF-8?q?feat(auth):=20=E2=9C=A8=20add=20configu?= =?UTF-8?q?rable=20OAuth=20callback=20ports=20for=20all=20providers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce environment variable configuration for OAuth callback server ports across Gemini CLI, Antigravity, and iFlow providers to prevent port conflicts and support multi-instance deployments. - Add `*_OAUTH_PORT` environment variables (GEMINI_CLI_OAUTH_PORT, ANTIGRAVITY_OAUTH_PORT, IFLOW_OAUTH_PORT) - Implement dynamic port resolution with fallback to hardcoded defaults - Add comprehensive documentation section explaining port configuration methods - Integrate port settings into Settings Tool UI for easy configuration - Update provider implementations to use configurable ports via property/function - Optimize launcher TUI startup by deferring heavy provider imports to Settings Tool - Add validation and warning logging for invalid port values Configuration can be managed via TUI settings menu or `.env` file. Port changes take effect on next authentication attempt without affecting existing tokens. --- DOCUMENTATION.md | 42 ++++++++++- src/proxy_app/launcher_tui.py | 69 ++++++++++++------- src/proxy_app/settings_tool.py | 43 ++++++++++++ .../providers/google_oauth_base.py | 25 ++++++- .../providers/iflow_auth_base.py | 24 ++++++- 5 files changed, 169 insertions(+), 34 deletions(-) diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index 30020176..5d43b610 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -856,6 +856,42 @@ class AntigravityAuthBase(GoogleOAuthBase): - Headless environment detection - Sequential refresh queue processing +#### OAuth Callback Port Configuration + +Each OAuth provider uses a local callback server during authentication. The callback port can be customized via environment variables to avoid conflicts with other services. + +**Default Ports:** + +| Provider | Default Port | Environment Variable | +|----------|-------------|---------------------| +| Gemini CLI | 8085 | `GEMINI_CLI_OAUTH_PORT` | +| Antigravity | 51121 | `ANTIGRAVITY_OAUTH_PORT` | +| iFlow | 11451 | `IFLOW_OAUTH_PORT` | + +**Configuration Methods:** + +1. **Via TUI Settings Menu:** + - Main Menu → `4. View Provider & Advanced Settings` → `1. Launch Settings Tool` + - Select the provider (Gemini CLI, Antigravity, or iFlow) + - Modify the `*_OAUTH_PORT` setting + - Use "Reset to Default" to restore the original port + +2. **Via `.env` file:** + ```env + # Custom OAuth callback ports (optional) + GEMINI_CLI_OAUTH_PORT=8085 + ANTIGRAVITY_OAUTH_PORT=51121 + IFLOW_OAUTH_PORT=11451 + ``` + +**When to Change Ports:** + +- If the default port conflicts with another service on your system +- If running multiple proxy instances on the same machine +- If firewall rules require specific port ranges + +**Note:** Port changes take effect on the next OAuth authentication attempt. Existing tokens are not affected. + --- @@ -877,8 +913,8 @@ The `GeminiCliProvider` is the most complex implementation, mimicking the Google #### Authentication (`gemini_auth_base.py`) - * **Device Flow**: Uses a standard OAuth 2.0 flow. The `credential_tool` spins up a local web server (`localhost:8085`) to capture the callback from Google's auth page. -* **Token Lifecycle**: + * **Device Flow**: Uses a standard OAuth 2.0 flow. The `credential_tool` spins up a local web server (default: `localhost:8085`, configurable via `GEMINI_CLI_OAUTH_PORT`) to capture the callback from Google's auth page. + * **Token Lifecycle**: * **Proactive Refresh**: Tokens are refreshed 5 minutes before expiry. * **Atomic Writes**: Credential files are updated using a temp-file-and-move strategy to prevent corruption during writes. * **Revocation Handling**: If a `400` or `401` occurs during refresh, the token is marked as revoked, preventing infinite retry loops. @@ -907,7 +943,7 @@ The provider employs a sophisticated, cached discovery mechanism to find a valid ### 3.3. iFlow (`iflow_provider.py`) * **Hybrid Auth**: Uses a custom OAuth flow (Authorization Code) to obtain an `access_token`. However, the *actual* API calls use a separate `apiKey` that is retrieved from the user's profile (`/api/oauth/getUserInfo`) using the access token. -* **Callback Server**: The auth flow spins up a local server on port `11451` to capture the redirect. +* **Callback Server**: The auth flow spins up a local server (default: port `11451`, configurable via `IFLOW_OAUTH_PORT`) to capture the redirect. * **Token Management**: Automatically refreshes the OAuth token and re-fetches the API key if needed. * **Schema Cleaning**: Similar to Qwen, it aggressively sanitizes tool schemas to prevent 400 errors. * **Dedicated Logging**: Implements `_IFlowFileLogger` to capture raw chunks for debugging proprietary API behaviors. diff --git a/src/proxy_app/launcher_tui.py b/src/proxy_app/launcher_tui.py index 954083dc..52940048 100644 --- a/src/proxy_app/launcher_tui.py +++ b/src/proxy_app/launcher_tui.py @@ -107,7 +107,7 @@ def _load_local_env() -> dict: @staticmethod def get_all_settings() -> dict: - """Returns comprehensive settings overview""" + """Returns comprehensive settings overview (includes provider_settings which triggers heavy imports)""" return { "credentials": SettingsDetector.detect_credentials(), "custom_bases": SettingsDetector.detect_custom_api_bases(), @@ -117,6 +117,17 @@ def get_all_settings() -> dict: "provider_settings": SettingsDetector.detect_provider_settings(), } + @staticmethod + def get_basic_settings() -> dict: + """Returns basic settings overview without provider_settings (avoids heavy imports)""" + return { + "credentials": SettingsDetector.detect_credentials(), + "custom_bases": SettingsDetector.detect_custom_api_bases(), + "model_definitions": SettingsDetector.detect_model_definitions(), + "concurrency_limits": SettingsDetector.detect_concurrency_limits(), + "model_filters": SettingsDetector.detect_model_filters(), + } + @staticmethod def detect_credentials() -> dict: """Detect API keys and OAuth credentials""" @@ -277,8 +288,8 @@ def show_main_menu(self): """Display main menu and handle selection""" clear_screen() - # Detect all settings - settings = SettingsDetector.get_all_settings() + # Detect basic settings (excludes provider_settings to avoid heavy imports) + settings = SettingsDetector.get_basic_settings() credentials = settings["credentials"] custom_bases = settings["custom_bases"] @@ -363,18 +374,17 @@ def show_main_menu(self): self.console.print("━" * 70) provider_count = len(credentials) custom_count = len(custom_bases) - provider_settings = settings.get("provider_settings", {}) + + self.console.print(f" Providers: {provider_count} configured") + self.console.print(f" Custom Providers: {custom_count} configured") + # Note: provider_settings detection is deferred to avoid heavy imports on startup has_advanced = bool( settings["model_definitions"] or settings["concurrency_limits"] or settings["model_filters"] - or provider_settings ) - - self.console.print(f" Providers: {provider_count} configured") - self.console.print(f" Custom Providers: {custom_count} configured") self.console.print( - f" Advanced Settings: {'Active (view in menu 4)' if has_advanced else 'None'}" + f" Advanced Settings: {'Active (view in menu 4)' if has_advanced else 'None (view menu 4 for details)'}" ) # Show menu @@ -659,13 +669,14 @@ def show_provider_settings_menu(self): """Display provider/advanced settings (read-only + launch tool)""" clear_screen() - settings = SettingsDetector.get_all_settings() + # Use basic settings to avoid heavy imports - provider_settings deferred to Settings Tool + settings = SettingsDetector.get_basic_settings() + credentials = settings["credentials"] custom_bases = settings["custom_bases"] model_defs = settings["model_definitions"] concurrency = settings["concurrency_limits"] filters = settings["model_filters"] - provider_settings = settings.get("provider_settings", {}) self.console.print( Panel.fit( @@ -740,23 +751,13 @@ def show_provider_settings_menu(self): status = " + ".join(status_parts) if status_parts else "None" self.console.print(f" • {provider:15} ✅ {status}") - # Provider-Specific Settings + # Provider-Specific Settings (deferred to Settings Tool to avoid heavy imports) self.console.print() self.console.print("[bold]🔬 Provider-Specific Settings[/bold]") self.console.print("━" * 70) - try: - from proxy_app.settings_tool import PROVIDER_SETTINGS_MAP - except ImportError: - from .settings_tool import PROVIDER_SETTINGS_MAP - for provider in PROVIDER_SETTINGS_MAP.keys(): - display_name = provider.replace("_", " ").title() - modified = provider_settings.get(provider, 0) - if modified > 0: - self.console.print( - f" • {display_name:20} [yellow]{modified} setting{'s' if modified > 1 else ''} modified[/yellow]" - ) - else: - self.console.print(f" • {display_name:20} [dim]using defaults[/dim]") + self.console.print( + " [dim]Launch Settings Tool to view/configure provider-specific settings[/dim]" + ) # Actions self.console.print() @@ -827,7 +828,23 @@ def launch_credential_tool(self): def launch_settings_tool(self): """Launch settings configuration tool""" - from proxy_app.settings_tool import run_settings_tool + import time + + clear_screen() + + self.console.print("━" * 70) + self.console.print("Advanced Settings Configuration Tool") + self.console.print("━" * 70) + + _start_time = time.time() + + with self.console.status("Initializing settings tool...", spinner="dots"): + from proxy_app.settings_tool import run_settings_tool + + _elapsed = time.time() - _start_time + self.console.print(f"✓ Settings tool ready in {_elapsed:.2f}s") + + time.sleep(0.3) run_settings_tool() # Reload environment after settings tool diff --git a/src/proxy_app/settings_tool.py b/src/proxy_app/settings_tool.py index ddc0dae1..69e0b851 100644 --- a/src/proxy_app/settings_tool.py +++ b/src/proxy_app/settings_tool.py @@ -14,6 +14,29 @@ console = Console() +# Import default OAuth port values from provider modules +# These serve as the source of truth for default port values +try: + from rotator_library.providers.gemini_auth_base import GeminiAuthBase + + GEMINI_CLI_DEFAULT_OAUTH_PORT = GeminiAuthBase.CALLBACK_PORT +except ImportError: + GEMINI_CLI_DEFAULT_OAUTH_PORT = 8085 + +try: + from rotator_library.providers.antigravity_auth_base import AntigravityAuthBase + + ANTIGRAVITY_DEFAULT_OAUTH_PORT = AntigravityAuthBase.CALLBACK_PORT +except ImportError: + ANTIGRAVITY_DEFAULT_OAUTH_PORT = 51121 + +try: + from rotator_library.providers.iflow_auth_base import ( + CALLBACK_PORT as IFLOW_DEFAULT_OAUTH_PORT, + ) +except ImportError: + IFLOW_DEFAULT_OAUTH_PORT = 11451 + def clear_screen(): """ @@ -383,6 +406,11 @@ def remove_multiplier(self, provider: str, priority: int): "default": "\n\nSTRICT PARAMETERS: {params}.", "description": "Template for Claude strict parameter hints in tool descriptions", }, + "ANTIGRAVITY_OAUTH_PORT": { + "type": "int", + "default": ANTIGRAVITY_DEFAULT_OAUTH_PORT, + "description": "Local port for OAuth callback server during authentication", + }, } # Gemini CLI provider environment variables @@ -427,12 +455,27 @@ def remove_multiplier(self, provider: str, priority: int): "default": "", "description": "GCP Project ID for paid tier users (required for paid tiers)", }, + "GEMINI_CLI_OAUTH_PORT": { + "type": "int", + "default": GEMINI_CLI_DEFAULT_OAUTH_PORT, + "description": "Local port for OAuth callback server during authentication", + }, +} + +# iFlow provider environment variables +IFLOW_SETTINGS = { + "IFLOW_OAUTH_PORT": { + "type": "int", + "default": IFLOW_DEFAULT_OAUTH_PORT, + "description": "Local port for OAuth callback server during authentication", + }, } # Map provider names to their settings definitions PROVIDER_SETTINGS_MAP = { "antigravity": ANTIGRAVITY_SETTINGS, "gemini_cli": GEMINI_CLI_SETTINGS, + "iflow": IFLOW_SETTINGS, } diff --git a/src/rotator_library/providers/google_oauth_base.py b/src/rotator_library/providers/google_oauth_base.py index 1a3f5d92..f618ac22 100644 --- a/src/rotator_library/providers/google_oauth_base.py +++ b/src/rotator_library/providers/google_oauth_base.py @@ -54,6 +54,25 @@ class GoogleOAuthBase: CALLBACK_PATH: str = "/oauth2callback" REFRESH_EXPIRY_BUFFER_SECONDS: int = 30 * 60 # 30 minutes + @property + def callback_port(self) -> int: + """ + Get the OAuth callback port, checking environment variable first. + + Reads from {ENV_PREFIX}_OAUTH_PORT environment variable, falling back + to the class's CALLBACK_PORT default if not set. + """ + env_var = f"{self.ENV_PREFIX}_OAUTH_PORT" + env_value = os.getenv(env_var) + if env_value: + try: + return int(env_value) + except ValueError: + lib_logger.warning( + f"Invalid {env_var} value: {env_value}, using default {self.CALLBACK_PORT}" + ) + return self.CALLBACK_PORT + def __init__(self): # Validate that subclass has set required attributes if self.CLIENT_ID is None: @@ -701,14 +720,14 @@ async def handle_callback(reader, writer): try: server = await asyncio.start_server( - handle_callback, "127.0.0.1", self.CALLBACK_PORT + handle_callback, "127.0.0.1", self.callback_port ) from urllib.parse import urlencode auth_url = "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode( { "client_id": self.CLIENT_ID, - "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", + "redirect_uri": f"http://localhost:{self.callback_port}{self.CALLBACK_PATH}", "scope": " ".join(self.OAUTH_SCOPES), "access_type": "offline", "response_type": "code", @@ -783,7 +802,7 @@ async def handle_callback(reader, writer): "code": auth_code.strip(), "client_id": self.CLIENT_ID, "client_secret": self.CLIENT_SECRET, - "redirect_uri": f"http://localhost:{self.CALLBACK_PORT}{self.CALLBACK_PATH}", + "redirect_uri": f"http://localhost:{self.callback_port}{self.CALLBACK_PATH}", "grant_type": "authorization_code", }, ) diff --git a/src/rotator_library/providers/iflow_auth_base.py b/src/rotator_library/providers/iflow_auth_base.py index 8854c493..589b4338 100644 --- a/src/rotator_library/providers/iflow_auth_base.py +++ b/src/rotator_library/providers/iflow_auth_base.py @@ -39,6 +39,25 @@ # Local callback server port CALLBACK_PORT = 11451 + +def get_callback_port() -> int: + """ + Get the OAuth callback port, checking environment variable first. + + Reads from IFLOW_OAUTH_PORT environment variable, falling back + to the default CALLBACK_PORT if not set. + """ + env_value = os.getenv("IFLOW_OAUTH_PORT") + if env_value: + try: + return int(env_value) + except ValueError: + logging.getLogger("rotator_library").warning( + f"Invalid IFLOW_OAUTH_PORT value: {env_value}, using default {CALLBACK_PORT}" + ) + return CALLBACK_PORT + + # Refresh tokens 24 hours before expiry REFRESH_EXPIRY_BUFFER_SECONDS = 24 * 60 * 60 @@ -931,7 +950,8 @@ async def _perform_interactive_oauth( state = secrets.token_urlsafe(32) # Build authorization URL - redirect_uri = f"http://localhost:{CALLBACK_PORT}/oauth2callback" + callback_port = get_callback_port() + redirect_uri = f"http://localhost:{callback_port}/oauth2callback" auth_params = { "loginMethod": "phone", "type": "phone", @@ -942,7 +962,7 @@ async def _perform_interactive_oauth( auth_url = f"{IFLOW_OAUTH_AUTHORIZE_ENDPOINT}?{urlencode(auth_params)}" # Start OAuth callback server - callback_server = OAuthCallbackServer(port=CALLBACK_PORT) + callback_server = OAuthCallbackServer(port=callback_port) try: await callback_server.start(expected_state=state) From 846ba251b519c6436649a80ec3b08ff1843e4ab9 Mon Sep 17 00:00:00 2001 From: Mirrowel <28632877+Mirrowel@users.noreply.github.com> Date: Mon, 8 Dec 2025 20:36:31 +0100 Subject: [PATCH 115/221] =?UTF-8?q?feat(ui):=20=E2=9C=A8=20add=20GUI=20for?= =?UTF-8?q?=20visual=20model=20filter=20configuration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a comprehensive CustomTkinter-based GUI application for managing model ignore/whitelist rules per provider, accessible from the settings tool. - Created model_filter_gui.py with full-featured visual editor (2600+ lines) - Implemented dual synchronized model lists showing unfiltered and filtered states - Added color-coded rule chips with visual association to affected models - Real-time pattern preview as users type filter rules - Interactive click/right-click functionality for model-rule relationships - Context menus for quick actions (add to ignore/whitelist, copy names) - Comprehensive help documentation with keyboard shortcuts - Unsaved changes detection with save/discard/cancel workflow - Background prefetching of models for all providers to improve responsiveness - Integration with settings tool as menu option #6 The GUI provides pattern matching with exact match, prefix wildcard (*), and match-all support. Whitelist rules take priority over ignore rules. All changes are persisted to .env file using IGNORE_MODELS_* and WHITELIST_MODELS_* variables. --- requirements.txt | 3 + src/proxy_app/model_filter_gui.py | 2601 +++++++++++++++++++++++++++++ src/proxy_app/settings_tool.py | 37 +- 3 files changed, 2633 insertions(+), 8 deletions(-) create mode 100644 src/proxy_app/model_filter_gui.py diff --git a/requirements.txt b/requirements.txt index edb2bcea..64f6aca7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,3 +19,6 @@ aiohttp colorlog rich + +# GUI for model filter configuration +customtkinter diff --git a/src/proxy_app/model_filter_gui.py b/src/proxy_app/model_filter_gui.py new file mode 100644 index 00000000..45d57b66 --- /dev/null +++ b/src/proxy_app/model_filter_gui.py @@ -0,0 +1,2601 @@ +""" +Model Filter GUI - Visual editor for model ignore/whitelist rules. + +A CustomTkinter application that provides a friendly interface for managing +which models are available per provider through ignore lists and whitelists. + +Features: +- Two synchronized model lists showing all fetched models and their filtered status +- Color-coded rules with visual association to affected models +- Real-time filtering preview as you type patterns +- Click interactions to highlight rule-model relationships +- Right-click context menus for quick actions +- Comprehensive help documentation +""" + +import customtkinter as ctk +from tkinter import Menu +import asyncio +import threading +import os +import re +from pathlib import Path +from dataclasses import dataclass, field +from typing import List, Dict, Tuple, Optional, Callable, Set +from dotenv import load_dotenv, set_key, unset_key + + +# ════════════════════════════════════════════════════════════════════════════════ +# CONSTANTS & CONFIGURATION +# ════════════════════════════════════════════════════════════════════════════════ + +# Window settings +WINDOW_TITLE = "Model Filter Configuration" +WINDOW_DEFAULT_SIZE = "1000x750" +WINDOW_MIN_WIDTH = 850 +WINDOW_MIN_HEIGHT = 600 + +# Color scheme (dark mode) +BG_PRIMARY = "#1a1a2e" # Main background +BG_SECONDARY = "#16213e" # Card/panel background +BG_TERTIARY = "#0f0f1a" # Input fields, lists +BG_HOVER = "#1f2b47" # Hover state +BORDER_COLOR = "#2a2a4a" # Subtle borders +TEXT_PRIMARY = "#e8e8e8" # Main text +TEXT_SECONDARY = "#a0a0a0" # Muted text +TEXT_MUTED = "#666680" # Very muted text +ACCENT_BLUE = "#4a9eff" # Primary accent +ACCENT_GREEN = "#2ecc71" # Success/normal +ACCENT_RED = "#e74c3c" # Danger/ignore +ACCENT_YELLOW = "#f1c40f" # Warning + +# Status colors +NORMAL_COLOR = "#2ecc71" # Green - models not affected by any rule +HIGHLIGHT_BG = "#2a3a5a" # Background for highlighted items + +# Ignore rules - warm color progression (reds/oranges) +IGNORE_COLORS = [ + "#e74c3c", # Bright red + "#c0392b", # Dark red + "#e67e22", # Orange + "#d35400", # Dark orange + "#f39c12", # Gold + "#e91e63", # Pink + "#ff5722", # Deep orange + "#f44336", # Material red + "#ff6b6b", # Coral + "#ff8a65", # Light deep orange +] + +# Whitelist rules - cool color progression (blues/teals) +WHITELIST_COLORS = [ + "#3498db", # Blue + "#2980b9", # Dark blue + "#1abc9c", # Teal + "#16a085", # Dark teal + "#9b59b6", # Purple + "#8e44ad", # Dark purple + "#00bcd4", # Cyan + "#2196f3", # Material blue + "#64b5f6", # Light blue + "#4dd0e1", # Light cyan +] + +# Font configuration +FONT_FAMILY = "Segoe UI" +FONT_SIZE_SMALL = 11 +FONT_SIZE_NORMAL = 12 +FONT_SIZE_LARGE = 14 +FONT_SIZE_TITLE = 16 +FONT_SIZE_HEADER = 20 + + +# ════════════════════════════════════════════════════════════════════════════════ +# DATA CLASSES +# ════════════════════════════════════════════════════════════════════════════════ + + +@dataclass +class FilterRule: + """Represents a single filter rule (ignore or whitelist pattern).""" + + pattern: str + color: str + rule_type: str # 'ignore' or 'whitelist' + affected_count: int = 0 + affected_models: List[str] = field(default_factory=list) + + def __hash__(self): + return hash((self.pattern, self.rule_type)) + + def __eq__(self, other): + if not isinstance(other, FilterRule): + return False + return self.pattern == other.pattern and self.rule_type == other.rule_type + + +@dataclass +class ModelStatus: + """Status information for a single model.""" + + model_id: str + status: str # 'normal', 'ignored', 'whitelisted' + color: str + affecting_rule: Optional[FilterRule] = None + + @property + def display_name(self) -> str: + """Get the model name without provider prefix for display.""" + if "/" in self.model_id: + return self.model_id.split("/", 1)[1] + return self.model_id + + @property + def provider(self) -> str: + """Extract provider from model ID.""" + if "/" in self.model_id: + return self.model_id.split("/")[0] + return "" + + +# ════════════════════════════════════════════════════════════════════════════════ +# FILTER ENGINE +# ════════════════════════════════════════════════════════════════════════════════ + + +class FilterEngine: + """ + Core filtering logic with rule management. + + Handles pattern matching, rule storage, and status calculation. + Tracks changes for save/discard functionality. + """ + + def __init__(self): + self.ignore_rules: List[FilterRule] = [] + self.whitelist_rules: List[FilterRule] = [] + self._ignore_color_index = 0 + self._whitelist_color_index = 0 + self._original_ignore_patterns: Set[str] = set() + self._original_whitelist_patterns: Set[str] = set() + self._current_provider: Optional[str] = None + + def reset(self): + """Clear all rules and reset state.""" + self.ignore_rules.clear() + self.whitelist_rules.clear() + self._ignore_color_index = 0 + self._whitelist_color_index = 0 + self._original_ignore_patterns.clear() + self._original_whitelist_patterns.clear() + + def _get_next_ignore_color(self) -> str: + """Get next color for ignore rules (cycles through palette).""" + color = IGNORE_COLORS[self._ignore_color_index % len(IGNORE_COLORS)] + self._ignore_color_index += 1 + return color + + def _get_next_whitelist_color(self) -> str: + """Get next color for whitelist rules (cycles through palette).""" + color = WHITELIST_COLORS[self._whitelist_color_index % len(WHITELIST_COLORS)] + self._whitelist_color_index += 1 + return color + + def add_ignore_rule(self, pattern: str) -> Optional[FilterRule]: + """Add a new ignore rule. Returns the rule if added, None if duplicate.""" + pattern = pattern.strip() + if not pattern: + return None + + # Check for duplicates + for rule in self.ignore_rules: + if rule.pattern == pattern: + return None + + rule = FilterRule( + pattern=pattern, color=self._get_next_ignore_color(), rule_type="ignore" + ) + self.ignore_rules.append(rule) + return rule + + def add_whitelist_rule(self, pattern: str) -> Optional[FilterRule]: + """Add a new whitelist rule. Returns the rule if added, None if duplicate.""" + pattern = pattern.strip() + if not pattern: + return None + + # Check for duplicates + for rule in self.whitelist_rules: + if rule.pattern == pattern: + return None + + rule = FilterRule( + pattern=pattern, + color=self._get_next_whitelist_color(), + rule_type="whitelist", + ) + self.whitelist_rules.append(rule) + return rule + + def remove_ignore_rule(self, pattern: str) -> bool: + """Remove an ignore rule by pattern. Returns True if removed.""" + for i, rule in enumerate(self.ignore_rules): + if rule.pattern == pattern: + self.ignore_rules.pop(i) + return True + return False + + def remove_whitelist_rule(self, pattern: str) -> bool: + """Remove a whitelist rule by pattern. Returns True if removed.""" + for i, rule in enumerate(self.whitelist_rules): + if rule.pattern == pattern: + self.whitelist_rules.pop(i) + return True + return False + + def _pattern_matches(self, model_id: str, pattern: str) -> bool: + """ + Check if a pattern matches a model ID. + + Supports: + - Exact match: "gpt-4" matches only "gpt-4" + - Prefix wildcard: "gpt-4*" matches "gpt-4", "gpt-4-turbo", etc. + - Match all: "*" matches everything + """ + # Extract model name without provider prefix + if "/" in model_id: + provider_model_name = model_id.split("/", 1)[1] + else: + provider_model_name = model_id + + if pattern == "*": + return True + elif pattern.endswith("*"): + prefix = pattern[:-1] + return provider_model_name.startswith(prefix) or model_id.startswith(prefix) + else: + # Exact match against full ID or provider model name + return model_id == pattern or provider_model_name == pattern + + def get_model_status(self, model_id: str) -> ModelStatus: + """ + Determine the status of a model based on current rules. + + Priority: Whitelist > Ignore > Normal + """ + # Check whitelist first (takes priority) + for rule in self.whitelist_rules: + if self._pattern_matches(model_id, rule.pattern): + return ModelStatus( + model_id=model_id, + status="whitelisted", + color=rule.color, + affecting_rule=rule, + ) + + # Then check ignore + for rule in self.ignore_rules: + if self._pattern_matches(model_id, rule.pattern): + return ModelStatus( + model_id=model_id, + status="ignored", + color=rule.color, + affecting_rule=rule, + ) + + # Default: normal + return ModelStatus( + model_id=model_id, status="normal", color=NORMAL_COLOR, affecting_rule=None + ) + + def get_all_statuses(self, models: List[str]) -> List[ModelStatus]: + """Get status for all models.""" + return [self.get_model_status(m) for m in models] + + def update_affected_counts(self, models: List[str]): + """Update the affected_count and affected_models for all rules.""" + # Reset counts + for rule in self.ignore_rules + self.whitelist_rules: + rule.affected_count = 0 + rule.affected_models = [] + + # Count affected models + for model_id in models: + status = self.get_model_status(model_id) + if status.affecting_rule: + status.affecting_rule.affected_count += 1 + status.affecting_rule.affected_models.append(model_id) + + def get_available_count(self, models: List[str]) -> Tuple[int, int]: + """Returns (available_count, total_count).""" + available = 0 + for model_id in models: + status = self.get_model_status(model_id) + if status.status != "ignored": + available += 1 + return available, len(models) + + def preview_pattern( + self, pattern: str, rule_type: str, models: List[str] + ) -> List[str]: + """ + Preview which models would be affected by a pattern without adding it. + Returns list of affected model IDs. + """ + affected = [] + pattern = pattern.strip() + if not pattern: + return affected + + for model_id in models: + if self._pattern_matches(model_id, pattern): + affected.append(model_id) + + return affected + + def load_from_env(self, provider: str): + """Load ignore/whitelist rules for a provider from environment.""" + self.reset() + self._current_provider = provider + load_dotenv(override=True) + + # Load ignore list + ignore_key = f"IGNORE_MODELS_{provider.upper()}" + ignore_value = os.getenv(ignore_key, "") + if ignore_value: + patterns = [p.strip() for p in ignore_value.split(",") if p.strip()] + for pattern in patterns: + self.add_ignore_rule(pattern) + self._original_ignore_patterns = set(patterns) + + # Load whitelist + whitelist_key = f"WHITELIST_MODELS_{provider.upper()}" + whitelist_value = os.getenv(whitelist_key, "") + if whitelist_value: + patterns = [p.strip() for p in whitelist_value.split(",") if p.strip()] + for pattern in patterns: + self.add_whitelist_rule(pattern) + self._original_whitelist_patterns = set(patterns) + + def save_to_env(self, provider: str) -> bool: + """ + Save current rules to .env file. + Returns True if successful. + """ + env_path = Path.cwd() / ".env" + + try: + ignore_key = f"IGNORE_MODELS_{provider.upper()}" + whitelist_key = f"WHITELIST_MODELS_{provider.upper()}" + + # Save ignore patterns + ignore_patterns = [rule.pattern for rule in self.ignore_rules] + if ignore_patterns: + set_key(str(env_path), ignore_key, ",".join(ignore_patterns)) + else: + # Remove the key if no patterns + unset_key(str(env_path), ignore_key) + + # Save whitelist patterns + whitelist_patterns = [rule.pattern for rule in self.whitelist_rules] + if whitelist_patterns: + set_key(str(env_path), whitelist_key, ",".join(whitelist_patterns)) + else: + unset_key(str(env_path), whitelist_key) + + # Update original state + self._original_ignore_patterns = set(ignore_patterns) + self._original_whitelist_patterns = set(whitelist_patterns) + + return True + except Exception as e: + print(f"Error saving to .env: {e}") + return False + + def has_unsaved_changes(self) -> bool: + """Check if current rules differ from saved state.""" + current_ignore = set(rule.pattern for rule in self.ignore_rules) + current_whitelist = set(rule.pattern for rule in self.whitelist_rules) + + return ( + current_ignore != self._original_ignore_patterns + or current_whitelist != self._original_whitelist_patterns + ) + + def discard_changes(self): + """Reload rules from environment, discarding unsaved changes.""" + if self._current_provider: + self.load_from_env(self._current_provider) + + +# ════════════════════════════════════════════════════════════════════════════════ +# MODEL FETCHER +# ════════════════════════════════════════════════════════════════════════════════ + +# Global cache for fetched models (persists across provider switches) +_model_cache: Dict[str, List[str]] = {} + + +class ModelFetcher: + """ + Handles async model fetching from providers. + + Runs fetching in a background thread to avoid blocking the GUI. + Includes caching to avoid refetching on every provider switch. + """ + + @staticmethod + def get_cached_models(provider: str) -> Optional[List[str]]: + """Get cached models for a provider, if available.""" + return _model_cache.get(provider) + + @staticmethod + def clear_cache(provider: Optional[str] = None): + """Clear model cache. If provider specified, only clear that provider.""" + if provider: + _model_cache.pop(provider, None) + else: + _model_cache.clear() + + @staticmethod + def get_available_providers() -> List[str]: + """Get list of providers that have credentials configured.""" + providers = set() + load_dotenv(override=True) + + # Scan environment for API keys (handles numbered keys like GEMINI_API_KEY_1) + for key in os.environ: + if "_API_KEY" in key and "PROXY_API_KEY" not in key: + # Extract provider: NVIDIA_NIM_API_KEY_1 -> nvidia_nim + provider = key.split("_API_KEY")[0].lower() + providers.add(provider) + + # Check for OAuth providers + oauth_dir = Path("oauth_creds") + if oauth_dir.exists(): + for file in oauth_dir.glob("*_oauth_*.json"): + provider = file.name.split("_oauth_")[0] + providers.add(provider) + + return sorted(list(providers)) + + @staticmethod + def _find_credential(provider: str) -> Optional[str]: + """Find a credential for a provider (handles numbered keys).""" + load_dotenv(override=True) + provider_upper = provider.upper() + + # Try exact match first (e.g., GEMINI_API_KEY) + exact_key = f"{provider_upper}_API_KEY" + if os.getenv(exact_key): + return os.getenv(exact_key) + + # Look for numbered keys (e.g., GEMINI_API_KEY_1, NVIDIA_NIM_API_KEY_1) + for key, value in os.environ.items(): + if key.startswith(f"{provider_upper}_API_KEY") and value: + return value + + # Check for OAuth credentials + oauth_dir = Path("oauth_creds") + if oauth_dir.exists(): + oauth_files = list(oauth_dir.glob(f"{provider}_oauth_*.json")) + if oauth_files: + return str(oauth_files[0]) + + return None + + @staticmethod + async def _fetch_models_async(provider: str) -> Tuple[List[str], Optional[str]]: + """ + Async implementation of model fetching. + Returns: (models_list, error_message_or_none) + """ + try: + import httpx + from rotator_library.providers import PROVIDER_PLUGINS + + # Get credential + credential = ModelFetcher._find_credential(provider) + if not credential: + return [], f"No credentials found for '{provider}'" + + # Get provider class + provider_class = PROVIDER_PLUGINS.get(provider.lower()) + if not provider_class: + return [], f"Unknown provider: '{provider}'" + + # Fetch models + async with httpx.AsyncClient(timeout=30.0) as client: + instance = provider_class() + models = await instance.get_models(credential, client) + return models, None + + except ImportError as e: + return [], f"Import error: {e}" + except Exception as e: + return [], f"Failed to fetch: {str(e)}" + + @staticmethod + def fetch_models( + provider: str, + on_success: Callable[[List[str]], None], + on_error: Callable[[str], None], + on_start: Optional[Callable[[], None]] = None, + force_refresh: bool = False, + ): + """ + Fetch models in a background thread. + + Args: + provider: Provider name (e.g., 'openai', 'gemini') + on_success: Callback with list of model IDs + on_error: Callback with error message + on_start: Optional callback when fetching starts + force_refresh: If True, bypass cache and fetch fresh + """ + # Check cache first (unless force refresh) + if not force_refresh: + cached = ModelFetcher.get_cached_models(provider) + if cached is not None: + on_success(cached) + return + + def run_fetch(): + if on_start: + on_start() + + try: + # Run async fetch in new event loop + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + models, error = loop.run_until_complete( + ModelFetcher._fetch_models_async(provider) + ) + # Clean up any pending tasks to avoid warnings + pending = asyncio.all_tasks(loop) + for task in pending: + task.cancel() + if pending: + loop.run_until_complete( + asyncio.gather(*pending, return_exceptions=True) + ) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + if error: + on_error(error) + else: + # Cache the results + _model_cache[provider] = models + on_success(models) + + except Exception as e: + on_error(str(e)) + + thread = threading.Thread(target=run_fetch, daemon=True) + thread.start() + + +# ════════════════════════════════════════════════════════════════════════════════ +# HELP WINDOW +# ════════════════════════════════════════════════════════════════════════════════ + + +class HelpWindow(ctk.CTkToplevel): + """ + Modal help popup with comprehensive filtering documentation. + """ + + def __init__(self, parent): + super().__init__(parent) + + self.title("Help - Model Filtering") + self.geometry("700x600") + self.minsize(600, 500) + + # Make modal + self.transient(parent) + self.grab_set() + + # Configure appearance + self.configure(fg_color=BG_PRIMARY) + + # Build content + self._create_content() + + # Center on parent + self.update_idletasks() + x = parent.winfo_x() + (parent.winfo_width() - self.winfo_width()) // 2 + y = parent.winfo_y() + (parent.winfo_height() - self.winfo_height()) // 2 + self.geometry(f"+{x}+{y}") + + # Focus + self.focus_force() + + # Bind escape to close + self.bind("", lambda e: self.destroy()) + + def _create_content(self): + """Build the help content.""" + # Main scrollable frame + main_frame = ctk.CTkScrollableFrame( + self, + fg_color=BG_PRIMARY, + scrollbar_fg_color=BG_SECONDARY, + scrollbar_button_color=BORDER_COLOR, + ) + main_frame.pack(fill="both", expand=True, padx=20, pady=20) + + # Title + title = ctk.CTkLabel( + main_frame, + text="📖 Model Filtering Guide", + font=(FONT_FAMILY, FONT_SIZE_HEADER, "bold"), + text_color=TEXT_PRIMARY, + ) + title.pack(anchor="w", pady=(0, 20)) + + # Sections + sections = [ + ( + "🎯 Overview", + """ +Model filtering allows you to control which models are available through your proxy for each provider. + +• Use the IGNORE list to block specific models +• Use the WHITELIST to ensure specific models are always available +• Whitelist ALWAYS takes priority over Ignore""", + ), + ( + "⚖️ Filtering Priority", + """ +When a model is checked, the following order is used: + +1. WHITELIST CHECK + If the model matches any whitelist pattern → AVAILABLE + (Whitelist overrides everything else) + +2. IGNORE CHECK + If the model matches any ignore pattern → BLOCKED + +3. DEFAULT + If no patterns match → AVAILABLE""", + ), + ( + "✏️ Pattern Syntax", + """ +Three types of patterns are supported: + +EXACT MATCH + Pattern: gpt-4 + Matches: only "gpt-4", nothing else + +PREFIX WILDCARD + Pattern: gpt-4* + Matches: "gpt-4", "gpt-4-turbo", "gpt-4-preview", etc. + +MATCH ALL + Pattern: * + Matches: every model for this provider""", + ), + ( + "💡 Common Patterns", + """ +BLOCK ALL, ALLOW SPECIFIC: + Ignore: * + Whitelist: gpt-4o, gpt-4o-mini + Result: Only gpt-4o and gpt-4o-mini available + +BLOCK PREVIEW MODELS: + Ignore: *-preview, *-preview* + Result: All preview variants blocked + +BLOCK SPECIFIC SERIES: + Ignore: o1*, dall-e* + Result: All o1 and DALL-E models blocked + +ALLOW ONLY LATEST: + Ignore: * + Whitelist: *-latest + Result: Only models ending in "-latest" available""", + ), + ( + "🖱️ Interface Guide", + """ +PROVIDER DROPDOWN + Select which provider to configure + +MODEL LISTS + • Left list: All fetched models (unfiltered) + • Right list: Same models with colored status + • Green = Available (normal) + • Red/Orange tones = Blocked (ignored) + • Blue/Teal tones = Whitelisted + +SEARCH BOX + Filter both lists to find specific models quickly + +CLICKING MODELS + • Left-click: Highlight the rule affecting this model + • Right-click: Context menu with quick actions + +CLICKING RULES + • Highlights all models affected by that rule + • Shows which models will be blocked/allowed + +RULE INPUT + • Enter patterns separated by commas + • Press Add or Enter to create rules + • Preview updates in real-time as you type + +DELETE RULES + • Click the × button on any rule to remove it""", + ), + ( + "⌨️ Keyboard Shortcuts", + """ +Ctrl+S Save changes +Ctrl+R Refresh models from provider +Ctrl+F Focus search box +F1 Open this help window +Escape Clear search / Close dialogs""", + ), + ( + "💾 Saving Changes", + """ +Changes are saved to your .env file in this format: + + IGNORE_MODELS_OPENAI=pattern1,pattern2* + WHITELIST_MODELS_OPENAI=specific-model + +Click "Save" to persist changes, or "Discard" to revert. +Closing the window with unsaved changes will prompt you.""", + ), + ] + + for title_text, content in sections: + self._add_section(main_frame, title_text, content) + + # Close button + close_btn = ctk.CTkButton( + main_frame, + text="Got it!", + font=(FONT_FAMILY, FONT_SIZE_NORMAL, "bold"), + fg_color=ACCENT_BLUE, + hover_color="#3a8aee", + height=40, + width=120, + command=self.destroy, + ) + close_btn.pack(pady=20) + + def _add_section(self, parent, title: str, content: str): + """Add a help section.""" + # Section title + title_label = ctk.CTkLabel( + parent, + text=title, + font=(FONT_FAMILY, FONT_SIZE_LARGE, "bold"), + text_color=ACCENT_BLUE, + ) + title_label.pack(anchor="w", pady=(15, 5)) + + # Separator + sep = ctk.CTkFrame(parent, height=1, fg_color=BORDER_COLOR) + sep.pack(fill="x", pady=(0, 10)) + + # Content + content_label = ctk.CTkLabel( + parent, + text=content.strip(), + font=(FONT_FAMILY, FONT_SIZE_NORMAL), + text_color=TEXT_SECONDARY, + justify="left", + anchor="w", + ) + content_label.pack(anchor="w", fill="x") + + +# ════════════════════════════════════════════════════════════════════════════════ +# CUSTOM DIALOG +# ════════════════════════════════════════════════════════════════════════════════ + + +class UnsavedChangesDialog(ctk.CTkToplevel): + """Modal dialog for unsaved changes confirmation.""" + + def __init__(self, parent): + super().__init__(parent) + + self.result: Optional[str] = None # 'save', 'discard', 'cancel' + + self.title("Unsaved Changes") + self.geometry("400x180") + self.resizable(False, False) + + # Make modal + self.transient(parent) + self.grab_set() + + # Configure appearance + self.configure(fg_color=BG_PRIMARY) + + # Build content + self._create_content() + + # Center on parent + self.update_idletasks() + x = parent.winfo_x() + (parent.winfo_width() - self.winfo_width()) // 2 + y = parent.winfo_y() + (parent.winfo_height() - self.winfo_height()) // 2 + self.geometry(f"+{x}+{y}") + + # Focus + self.focus_force() + + # Bind escape to cancel + self.bind("", lambda e: self._on_cancel()) + + # Handle window close + self.protocol("WM_DELETE_WINDOW", self._on_cancel) + + def _create_content(self): + """Build dialog content.""" + # Icon and message + msg_frame = ctk.CTkFrame(self, fg_color="transparent") + msg_frame.pack(fill="x", padx=30, pady=(25, 15)) + + icon = ctk.CTkLabel( + msg_frame, text="⚠️", font=(FONT_FAMILY, 32), text_color=ACCENT_YELLOW + ) + icon.pack(side="left", padx=(0, 15)) + + text_frame = ctk.CTkFrame(msg_frame, fg_color="transparent") + text_frame.pack(side="left", fill="x", expand=True) + + title = ctk.CTkLabel( + text_frame, + text="Unsaved Changes", + font=(FONT_FAMILY, FONT_SIZE_LARGE, "bold"), + text_color=TEXT_PRIMARY, + anchor="w", + ) + title.pack(anchor="w") + + subtitle = ctk.CTkLabel( + text_frame, + text="You have unsaved filter changes.\nWhat would you like to do?", + font=(FONT_FAMILY, FONT_SIZE_NORMAL), + text_color=TEXT_SECONDARY, + anchor="w", + justify="left", + ) + subtitle.pack(anchor="w") + + # Buttons + btn_frame = ctk.CTkFrame(self, fg_color="transparent") + btn_frame.pack(fill="x", padx=30, pady=(10, 25)) + + cancel_btn = ctk.CTkButton( + btn_frame, + text="Cancel", + font=(FONT_FAMILY, FONT_SIZE_NORMAL), + fg_color=BG_SECONDARY, + hover_color=BG_HOVER, + border_width=1, + border_color=BORDER_COLOR, + width=100, + command=self._on_cancel, + ) + cancel_btn.pack(side="right", padx=(10, 0)) + + discard_btn = ctk.CTkButton( + btn_frame, + text="Discard", + font=(FONT_FAMILY, FONT_SIZE_NORMAL), + fg_color=ACCENT_RED, + hover_color="#c0392b", + width=100, + command=self._on_discard, + ) + discard_btn.pack(side="right", padx=(10, 0)) + + save_btn = ctk.CTkButton( + btn_frame, + text="Save", + font=(FONT_FAMILY, FONT_SIZE_NORMAL), + fg_color=ACCENT_GREEN, + hover_color="#27ae60", + width=100, + command=self._on_save, + ) + save_btn.pack(side="right") + + def _on_save(self): + self.result = "save" + self.destroy() + + def _on_discard(self): + self.result = "discard" + self.destroy() + + def _on_cancel(self): + self.result = "cancel" + self.destroy() + + def show(self) -> Optional[str]: + """Show dialog and return result.""" + self.wait_window() + return self.result + + +# ════════════════════════════════════════════════════════════════════════════════ +# TOOLTIP +# ════════════════════════════════════════════════════════════════════════════════ + + +class ToolTip: + """Simple tooltip implementation for CustomTkinter widgets.""" + + def __init__(self, widget, text: str, delay: int = 500): + self.widget = widget + self.text = text + self.delay = delay + self.tooltip_window = None + self.after_id = None + + widget.bind("", self._schedule_show) + widget.bind("", self._hide) + widget.bind(" + + + + +
+ +
+
+

+ ⚙ API Configuration + Checking... +

+ ▼ +
+
+
+ + +
+
+ + +
+
+ +
+
+
+ + +
+
+ +
Configure API settings above to connect...
+
+
+
+ + + +