code-puppy 0.0.214__py3-none-any.whl → 0.0.366__py3-none-any.whl

code_puppy/model_factory.py

@@ -4,33 +4,134 @@ import os
 import pathlib
 from typing import Any, Dict
 
-import httpx
 from anthropic import AsyncAnthropic
 from openai import AsyncAzureOpenAI
-from pydantic_ai.models.anthropic import AnthropicModel
-from pydantic_ai.models.google import GoogleModel
-from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
+from pydantic_ai.models.openai import (
+    OpenAIChatModel,
+    OpenAIChatModelSettings,
+    OpenAIResponsesModel,
+)
+from pydantic_ai.profiles import ModelProfile
 from pydantic_ai.providers.anthropic import AnthropicProvider
 from pydantic_ai.providers.cerebras import CerebrasProvider
-from pydantic_ai.providers.google import GoogleProvider
 from pydantic_ai.providers.openai import OpenAIProvider
 from pydantic_ai.providers.openrouter import OpenRouterProvider
+from pydantic_ai.settings import ModelSettings
 
+from code_puppy.gemini_model import GeminiModel
 from code_puppy.messaging import emit_warning
 
 from . import callbacks
-from .config import EXTRA_MODELS_FILE
-from .http_utils import create_async_client
+from .claude_cache_client import ClaudeCacheAsyncClient, patch_anthropic_client_messages
+from .config import EXTRA_MODELS_FILE, get_value
+from .http_utils import create_async_client, get_cert_bundle_path, get_http2
 from .round_robin_model import RoundRobinModel
 
-# Environment variables used in this module:
-# - GEMINI_API_KEY: API key for Google's Gemini models. Required when using Gemini models.
-# - OPENAI_API_KEY: API key for OpenAI models. Required when using OpenAI models or custom_openai endpoints.
-# - TOGETHER_AI_KEY: API key for Together AI models. Required when using Together AI models.
-#
-# When using custom endpoints (type: "custom_openai" in models.json):
-# - Environment variables can be referenced in header values by prefixing with $ in models.json.
-#   Example: "X-Api-Key": "$OPENAI_API_KEY" will use the value from os.environ.get("OPENAI_API_KEY")
+logger = logging.getLogger(__name__)
+
+
+def get_api_key(env_var_name: str) -> str | None:
+    """Get an API key from config first, then fall back to environment variable.
+
+    This allows users to set API keys via `/set KIMI_API_KEY=xxx` in addition to
+    setting them as environment variables.
+
+    Args:
+        env_var_name: The name of the environment variable (e.g., "OPENAI_API_KEY")
+
+    Returns:
+        The API key value, or None if not found in either config or environment.
+    """
+    # First check config (case-insensitive key lookup)
+    config_value = get_value(env_var_name.lower())
+    if config_value:
+        return config_value
+
+    # Fall back to environment variable
+    return os.environ.get(env_var_name)
+
+
+def make_model_settings(
+    model_name: str, max_tokens: int | None = None
+) -> ModelSettings:
+    """Create appropriate ModelSettings for a given model.
+
+    This handles model-specific settings:
+    - GPT-5 models: reasoning_effort and verbosity (non-codex only)
+    - Claude/Anthropic models: extended_thinking and budget_tokens
+    - Automatic max_tokens calculation based on model context length
+
+    Args:
+        model_name: The name of the model to create settings for.
+        max_tokens: Optional max tokens limit. If None, automatically calculated
+            as: max(2048, min(15% of context_length, 65536))
+
+    Returns:
+        Appropriate ModelSettings subclass instance for the model.
+    """
+    from code_puppy.config import (
+        get_effective_model_settings,
+        get_openai_reasoning_effort,
+        get_openai_verbosity,
+    )
+
+    model_settings_dict: dict = {}
+
+    # Calculate max_tokens if not explicitly provided
+    if max_tokens is None:
+        # Load model config to get context length
+        try:
+            models_config = ModelFactory.load_config()
+            model_config = models_config.get(model_name, {})
+            context_length = model_config.get("context_length", 128000)
+        except Exception:
+            # Fallback if config loading fails (e.g., in CI environments)
+            context_length = 128000
+        # min 2048, 15% of context, max 65536
+        max_tokens = max(2048, min(int(0.15 * context_length), 65536))
+
+    model_settings_dict["max_tokens"] = max_tokens
+    effective_settings = get_effective_model_settings(model_name)
+    model_settings_dict.update(effective_settings)
+
+    # Default to clear_thinking=False for GLM-4.7 models (preserved thinking)
+    if "glm-4.7" in model_name.lower():
+        clear_thinking = effective_settings.get("clear_thinking", False)
+        model_settings_dict["thinking"] = {
+            "type": "enabled",
+            "clear_thinking": clear_thinking,
+        }
+
+    model_settings: ModelSettings = ModelSettings(**model_settings_dict)
+
+    if "gpt-5" in model_name:
+        model_settings_dict["openai_reasoning_effort"] = get_openai_reasoning_effort()
+        # Verbosity only applies to non-codex GPT-5 models (codex only supports "medium")
+        if "codex" not in model_name:
+            verbosity = get_openai_verbosity()
+            model_settings_dict["extra_body"] = {"verbosity": verbosity}
+        model_settings = OpenAIChatModelSettings(**model_settings_dict)
+    elif model_name.startswith("claude-") or model_name.startswith("anthropic-"):
+        # Handle Anthropic extended thinking settings
+        # Remove top_p as Anthropic doesn't support it with extended thinking
+        model_settings_dict.pop("top_p", None)
+
+        # Claude extended thinking requires temperature=1.0 (API restriction)
+        # Default to 1.0 if not explicitly set by user
+        if model_settings_dict.get("temperature") is None:
+            model_settings_dict["temperature"] = 1.0
+
+        extended_thinking = effective_settings.get("extended_thinking", True)
+        budget_tokens = effective_settings.get("budget_tokens", 10000)
+        if extended_thinking and budget_tokens:
+            model_settings_dict["anthropic_thinking"] = {
+                "type": "enabled",
+                "budget_tokens": budget_tokens,
+            }
+        model_settings = AnthropicModelSettings(**model_settings_dict)
+
+    return model_settings
 
 
 class ZaiChatModel(OpenAIChatModel):
@@ -52,10 +153,10 @@ def get_custom_config(model_config):
     for key, value in custom_config.get("headers", {}).items():
         if value.startswith("$"):
             env_var_name = value[1:]
-            resolved_value = os.environ.get(env_var_name)
+            resolved_value = get_api_key(env_var_name)
             if resolved_value is None:
                 emit_warning(
-                    f"Environment variable '{env_var_name}' is not set for custom endpoint header '{key}'. Proceeding with empty value."
+                    f"'{env_var_name}' is not set (check config or environment) for custom endpoint header '{key}'. Proceeding with empty value."
                 )
                 resolved_value = ""
             value = resolved_value
@@ -65,10 +166,10 @@ def get_custom_config(model_config):
             for token in tokens:
                 if token.startswith("$"):
                     env_var = token[1:]
-                    resolved_value = os.environ.get(env_var)
+                    resolved_value = get_api_key(env_var)
                     if resolved_value is None:
                         emit_warning(
-                            f"Environment variable '{env_var}' is not set for custom endpoint header '{key}'. Proceeding with empty value."
+                            f"'{env_var}' is not set (check config or environment) for custom endpoint header '{key}'. Proceeding with empty value."
                         )
                         resolved_values.append("")
                     else:
@@ -81,10 +182,10 @@ def get_custom_config(model_config):
     if "api_key" in custom_config:
         if custom_config["api_key"].startswith("$"):
             env_var_name = custom_config["api_key"][1:]
-            api_key = os.environ.get(env_var_name)
+            api_key = get_api_key(env_var_name)
             if api_key is None:
                 emit_warning(
-                    f"Environment variable '{env_var_name}' is not set for custom endpoint API key; proceeding without API key."
+                    f"API key '{env_var_name}' is not set (checked config and environment); proceeding without API key."
                 )
         else:
             api_key = custom_config["api_key"]
@@ -117,26 +218,63 @@ class ModelFactory:
             with open(MODELS_FILE, "r") as f:
                 config = json.load(f)
 
-        if pathlib.Path(EXTRA_MODELS_FILE).exists():
+        # Import OAuth model file paths from main config
+        from code_puppy.config import (
+            ANTIGRAVITY_MODELS_FILE,
+            CHATGPT_MODELS_FILE,
+            CLAUDE_MODELS_FILE,
+            GEMINI_MODELS_FILE,
+        )
+
+        # Build list of extra model sources
+        extra_sources: list[tuple[pathlib.Path, str, bool]] = [
+            (pathlib.Path(EXTRA_MODELS_FILE), "extra models", False),
+            (pathlib.Path(CHATGPT_MODELS_FILE), "ChatGPT OAuth models", False),
+            (pathlib.Path(CLAUDE_MODELS_FILE), "Claude Code OAuth models", True),
+            (pathlib.Path(GEMINI_MODELS_FILE), "Gemini OAuth models", False),
+            (pathlib.Path(ANTIGRAVITY_MODELS_FILE), "Antigravity OAuth models", False),
+        ]
+
+        for source_path, label, use_filtered in extra_sources:
+            if not source_path.exists():
+                continue
             try:
-                with open(EXTRA_MODELS_FILE, "r") as f:
-                    extra_config = json.load(f)
-                    config.update(extra_config)
-            except json.JSONDecodeError as e:
+                # Use filtered loading for Claude Code OAuth models to show only latest versions
+                if use_filtered:
+                    try:
+                        from code_puppy.plugins.claude_code_oauth.utils import (
+                            load_claude_models_filtered,
+                        )
+
+                        extra_config = load_claude_models_filtered()
+                    except ImportError:
+                        # Plugin not available, fall back to standard JSON loading
+                        logging.getLogger(__name__).debug(
+                            f"claude_code_oauth plugin not available, loading {label} as plain JSON"
+                        )
+                        with open(source_path, "r") as f:
+                            extra_config = json.load(f)
+                else:
+                    with open(source_path, "r") as f:
+                        extra_config = json.load(f)
+                config.update(extra_config)
+            except json.JSONDecodeError as exc:
                 logging.getLogger(__name__).warning(
-                    f"Failed to load extra models config from {EXTRA_MODELS_FILE}: Invalid JSON - {e}\n"
-                    f"Please check your extra_models.json file for syntax errors."
+                    f"Failed to load {label} config from {source_path}: Invalid JSON - {exc}"
                 )
-            except Exception as e:
+            except Exception as exc:
                 logging.getLogger(__name__).warning(
-                    f"Failed to load extra models config from {EXTRA_MODELS_FILE}: {e}\n"
-                    f"The extra models configuration will be ignored."
+                    f"Failed to load {label} config from {source_path}: {exc}"
                 )
         return config
 
     @staticmethod
     def get_model(model_name: str, config: Dict[str, Any]) -> Any:
-        """Returns a configured model instance based on the provided name and config."""
+        """Returns a configured model instance based on the provided name and config.
+
+        API key validation happens naturally within each model type's initialization,
+        which emits warnings and returns None if keys are missing.
+        """
         model_config = config.get(model_name)
         if not model_config:
             raise ValueError(f"Model '{model_name}' not found in configuration.")
@@ -144,41 +282,189 @@ class ModelFactory:
         model_type = model_config.get("type")
 
         if model_type == "gemini":
-            provider = GoogleProvider(api_key=os.environ.get("GEMINI_API_KEY", ""))
+            api_key = get_api_key("GEMINI_API_KEY")
+            if not api_key:
+                emit_warning(
+                    f"GEMINI_API_KEY is not set (check config or environment); skipping Gemini model '{model_config.get('name')}'."
+                )
+                return None
 
-            model = GoogleModel(model_name=model_config["name"], provider=provider)
-            setattr(model, "provider", provider)
+            model = GeminiModel(model_name=model_config["name"], api_key=api_key)
             return model
 
         elif model_type == "openai":
-            provider = OpenAIProvider(api_key=os.environ.get("OPENAI_API_KEY", ""))
+            api_key = get_api_key("OPENAI_API_KEY")
+            if not api_key:
+                emit_warning(
+                    f"OPENAI_API_KEY is not set (check config or environment); skipping OpenAI model '{model_config.get('name')}'."
+                )
+                return None
 
+            provider = OpenAIProvider(api_key=api_key)
             model = OpenAIChatModel(model_name=model_config["name"], provider=provider)
+            if "codex" in model_name:
+                model = OpenAIResponsesModel(
+                    model_name=model_config["name"], provider=provider
+                )
             setattr(model, "provider", provider)
             return model
 
         elif model_type == "anthropic":
-            api_key = os.environ.get("ANTHROPIC_API_KEY", None)
+            api_key = get_api_key("ANTHROPIC_API_KEY")
             if not api_key:
                 emit_warning(
-                    f"ANTHROPIC_API_KEY is not set; skipping Anthropic model '{model_config.get('name')}'."
+                    f"ANTHROPIC_API_KEY is not set (check config or environment); skipping Anthropic model '{model_config.get('name')}'."
                 )
                 return None
-            anthropic_client = AsyncAnthropic(api_key=api_key)
+
+            # Use the same caching client as claude_code models
+            verify = get_cert_bundle_path()
+            http2_enabled = get_http2()
+
+            client = ClaudeCacheAsyncClient(
+                verify=verify,
+                timeout=180,
+                http2=http2_enabled,
+            )
+
+            # Check if interleaved thinking is enabled for this model
+            # Only applies to Claude 4 models (Opus 4.5, Opus 4.1, Opus 4, Sonnet 4)
+            from code_puppy.config import get_effective_model_settings
+
+            effective_settings = get_effective_model_settings(model_name)
+            interleaved_thinking = effective_settings.get("interleaved_thinking", False)
+
+            default_headers = {}
+            if interleaved_thinking:
+                default_headers["anthropic-beta"] = "interleaved-thinking-2025-05-14"
+
+            anthropic_client = AsyncAnthropic(
+                api_key=api_key,
+                http_client=client,
+                default_headers=default_headers if default_headers else None,
+            )
+
+            # Ensure cache_control is injected at the Anthropic SDK layer
+            patch_anthropic_client_messages(anthropic_client)
+
             provider = AnthropicProvider(anthropic_client=anthropic_client)
             return AnthropicModel(model_name=model_config["name"], provider=provider)
 
         elif model_type == "custom_anthropic":
             url, headers, verify, api_key = get_custom_config(model_config)
-            client = create_async_client(headers=headers, verify=verify)
+            if not api_key:
+                emit_warning(
+                    f"API key is not set for custom Anthropic endpoint; skipping model '{model_config.get('name')}'."
+                )
+                return None
+
+            # Use the same caching client as claude_code models
+            if verify is None:
+                verify = get_cert_bundle_path()
+
+            http2_enabled = get_http2()
+
+            client = ClaudeCacheAsyncClient(
+                headers=headers,
+                verify=verify,
+                timeout=180,
+                http2=http2_enabled,
+            )
+
+            # Check if interleaved thinking is enabled for this model
+            from code_puppy.config import get_effective_model_settings
+
+            effective_settings = get_effective_model_settings(model_name)
+            interleaved_thinking = effective_settings.get("interleaved_thinking", False)
+
+            default_headers = {}
+            if interleaved_thinking:
+                default_headers["anthropic-beta"] = "interleaved-thinking-2025-05-14"
+
             anthropic_client = AsyncAnthropic(
                 base_url=url,
                 http_client=client,
                 api_key=api_key,
+                default_headers=default_headers if default_headers else None,
             )
+
+            # Ensure cache_control is injected at the Anthropic SDK layer
+            patch_anthropic_client_messages(anthropic_client)
+
             provider = AnthropicProvider(anthropic_client=anthropic_client)
             return AnthropicModel(model_name=model_config["name"], provider=provider)
+        elif model_type == "claude_code":
+            url, headers, verify, api_key = get_custom_config(model_config)
+            if model_config.get("oauth_source") == "claude-code-plugin":
+                try:
+                    from code_puppy.plugins.claude_code_oauth.utils import (
+                        get_valid_access_token,
+                    )
+
+                    refreshed_token = get_valid_access_token()
+                    if refreshed_token:
+                        api_key = refreshed_token
+                        custom_endpoint = model_config.get("custom_endpoint")
+                        if isinstance(custom_endpoint, dict):
+                            custom_endpoint["api_key"] = refreshed_token
+                except ImportError:
+                    pass
+            if not api_key:
+                emit_warning(
+                    f"API key is not set for Claude Code endpoint; skipping model '{model_config.get('name')}'."
+                )
+                return None
+
+            # Check if interleaved thinking is enabled (defaults to True for OAuth models)
+            from code_puppy.config import get_effective_model_settings
 
+            effective_settings = get_effective_model_settings(model_name)
+            interleaved_thinking = effective_settings.get("interleaved_thinking", True)
+
+            # Handle anthropic-beta header based on interleaved_thinking setting
+            if "anthropic-beta" in headers:
+                beta_parts = [p.strip() for p in headers["anthropic-beta"].split(",")]
+                if interleaved_thinking:
+                    # Ensure interleaved-thinking is in the header
+                    if "interleaved-thinking-2025-05-14" not in beta_parts:
+                        beta_parts.append("interleaved-thinking-2025-05-14")
+                else:
+                    # Remove interleaved-thinking from the header
+                    beta_parts = [
+                        p for p in beta_parts if "interleaved-thinking" not in p
+                    ]
+                headers["anthropic-beta"] = ",".join(beta_parts) if beta_parts else None
+                if headers.get("anthropic-beta") is None:
+                    del headers["anthropic-beta"]
+            elif interleaved_thinking:
+                # No existing beta header, add one for interleaved thinking
+                headers["anthropic-beta"] = "interleaved-thinking-2025-05-14"
+
+            # Use a dedicated client wrapper that injects cache_control on /v1/messages
+            if verify is None:
+                verify = get_cert_bundle_path()
+
+            http2_enabled = get_http2()
+
+            client = ClaudeCacheAsyncClient(
+                headers=headers,
+                verify=verify,
+                timeout=180,
+                http2=http2_enabled,
+            )
+
+            anthropic_client = AsyncAnthropic(
+                base_url=url,
+                http_client=client,
+                auth_token=api_key,
+            )
+            # Ensure cache_control is injected at the Anthropic SDK layer too
+            # so we don't depend solely on httpx internals.
+            patch_anthropic_client_messages(anthropic_client)
+            anthropic_client.api_key = None
+            anthropic_client.auth_token = api_key
+            provider = AnthropicProvider(anthropic_client=anthropic_client)
+            return AnthropicModel(model_name=model_config["name"], provider=provider)
         elif model_type == "azure_openai":
             azure_endpoint_config = model_config.get("azure_endpoint")
             if not azure_endpoint_config:
@@ -187,10 +473,10 @@ class ModelFactory:
                 )
             azure_endpoint = azure_endpoint_config
             if azure_endpoint_config.startswith("$"):
-                azure_endpoint = os.environ.get(azure_endpoint_config[1:])
+                azure_endpoint = get_api_key(azure_endpoint_config[1:])
             if not azure_endpoint:
                 emit_warning(
-                    f"Azure OpenAI endpoint environment variable '{azure_endpoint_config[1:] if azure_endpoint_config.startswith('$') else azure_endpoint_config}' not found or is empty; skipping model '{model_config.get('name')}'."
+                    f"Azure OpenAI endpoint '{azure_endpoint_config[1:] if azure_endpoint_config.startswith('$') else azure_endpoint_config}' not found (check config or environment); skipping model '{model_config.get('name')}'."
                 )
                 return None
 
@@ -201,10 +487,10 @@ class ModelFactory:
                 )
             api_version = api_version_config
             if api_version_config.startswith("$"):
-                api_version = os.environ.get(api_version_config[1:])
+                api_version = get_api_key(api_version_config[1:])
             if not api_version:
                 emit_warning(
-                    f"Azure OpenAI API version environment variable '{api_version_config[1:] if api_version_config.startswith('$') else api_version_config}' not found or is empty; skipping model '{model_config.get('name')}'."
+                    f"Azure OpenAI API version '{api_version_config[1:] if api_version_config.startswith('$') else api_version_config}' not found (check config or environment); skipping model '{model_config.get('name')}'."
                 )
                 return None
 
@@ -215,10 +501,10 @@ class ModelFactory:
                 )
             api_key = api_key_config
             if api_key_config.startswith("$"):
-                api_key = os.environ.get(api_key_config[1:])
+                api_key = get_api_key(api_key_config[1:])
             if not api_key:
                 emit_warning(
-                    f"Azure OpenAI API key environment variable '{api_key_config[1:] if api_key_config.startswith('$') else api_key_config}' not found or is empty; skipping model '{model_config.get('name')}'."
+                    f"Azure OpenAI API key '{api_key_config[1:] if api_key_config.startswith('$') else api_key_config}' not found (check config or environment); skipping model '{model_config.get('name')}'."
                 )
                 return None
 
@@ -246,71 +532,193 @@ class ModelFactory:
             if api_key:
                 provider_args["api_key"] = api_key
             provider = OpenAIProvider(**provider_args)
-
             model = OpenAIChatModel(model_name=model_config["name"], provider=provider)
+            if model_name == "chatgpt-gpt-5-codex":
+                model = OpenAIResponsesModel(model_config["name"], provider=provider)
             setattr(model, "provider", provider)
             return model
         elif model_type == "zai_coding":
-            api_key = os.getenv("ZAI_API_KEY")
+            api_key = get_api_key("ZAI_API_KEY")
             if not api_key:
                 emit_warning(
-                    f"ZAI_API_KEY is not set; skipping ZAI coding model '{model_config.get('name')}'."
+                    f"ZAI_API_KEY is not set (check config or environment); skipping ZAI coding model '{model_config.get('name')}'."
                 )
                 return None
+            provider = OpenAIProvider(
+                api_key=api_key,
+                base_url="https://api.z.ai/api/coding/paas/v4",
+            )
             zai_model = ZaiChatModel(
                 model_name=model_config["name"],
-                provider=OpenAIProvider(
-                    api_key=api_key,
-                    base_url="https://api.z.ai/api/coding/paas/v4",
-                ),
+                provider=provider,
             )
+            setattr(zai_model, "provider", provider)
             return zai_model
         elif model_type == "zai_api":
-            api_key = os.getenv("ZAI_API_KEY")
+            api_key = get_api_key("ZAI_API_KEY")
             if not api_key:
                 emit_warning(
-                    f"ZAI_API_KEY is not set; skipping ZAI API model '{model_config.get('name')}'."
+                    f"ZAI_API_KEY is not set (check config or environment); skipping ZAI API model '{model_config.get('name')}'."
                 )
                 return None
+            provider = OpenAIProvider(
+                api_key=api_key,
+                base_url="https://api.z.ai/api/paas/v4/",
+            )
             zai_model = ZaiChatModel(
                 model_name=model_config["name"],
-                provider=OpenAIProvider(
-                    api_key=api_key,
-                    base_url="https://api.z.ai/api/paas/v4/",
-                ),
+                provider=provider,
             )
+            setattr(zai_model, "provider", provider)
             return zai_model
         elif model_type == "custom_gemini":
             url, headers, verify, api_key = get_custom_config(model_config)
-            os.environ["GEMINI_API_KEY"] = api_key
+            if not api_key:
+                emit_warning(
+                    f"API key is not set for custom Gemini endpoint; skipping model '{model_config.get('name')}'."
+                )
+                return None
 
-            class CustomGoogleGLAProvider(GoogleProvider):
-                def __init__(self, *args, **kwargs):
-                    super().__init__(*args, **kwargs)
+            # Check if this is an Antigravity model
+            if model_config.get("antigravity"):
+                try:
+                    from code_puppy.plugins.antigravity_oauth.token import (
+                        is_token_expired,
+                        refresh_access_token,
+                    )
+                    from code_puppy.plugins.antigravity_oauth.transport import (
+                        create_antigravity_client,
+                    )
+                    from code_puppy.plugins.antigravity_oauth.utils import (
+                        load_stored_tokens,
+                        save_tokens,
+                    )
+
+                    # Try to import custom model for thinking signatures
+                    try:
+                        from code_puppy.plugins.antigravity_oauth.antigravity_model import (
+                            AntigravityModel,
+                        )
+                    except ImportError:
+                        AntigravityModel = None
 
-                @property
-                def base_url(self):
-                    return url
+                    # Get fresh access token (refresh if needed)
+                    tokens = load_stored_tokens()
+                    if not tokens:
+                        emit_warning(
+                            "Antigravity tokens not found; run /antigravity-auth first."
+                        )
+                        return None
+
+                    access_token = tokens.get("access_token", "")
+                    refresh_token = tokens.get("refresh_token", "")
+                    expires_at = tokens.get("expires_at")
+
+                    # Refresh if expired or about to expire (initial check)
+                    if is_token_expired(expires_at):
+                        new_tokens = refresh_access_token(refresh_token)
+                        if new_tokens:
+                            access_token = new_tokens.access_token
+                            refresh_token = new_tokens.refresh_token
+                            expires_at = new_tokens.expires_at
+                            tokens["access_token"] = new_tokens.access_token
+                            tokens["refresh_token"] = new_tokens.refresh_token
+                            tokens["expires_at"] = new_tokens.expires_at
+                            save_tokens(tokens)
+                        else:
+                            emit_warning(
+                                "Failed to refresh Antigravity token; run /antigravity-auth again."
+                            )
+                            return None
+
+                    # Callback to persist tokens when proactively refreshed during session
+                    def on_token_refreshed(new_tokens):
+                        """Persist new tokens when proactively refreshed."""
+                        try:
+                            updated_tokens = load_stored_tokens() or {}
+                            updated_tokens["access_token"] = new_tokens.access_token
+                            updated_tokens["refresh_token"] = new_tokens.refresh_token
+                            updated_tokens["expires_at"] = new_tokens.expires_at
+                            save_tokens(updated_tokens)
+                            logger.debug(
+                                "Persisted proactively refreshed Antigravity tokens"
+                            )
+                        except Exception as e:
+                            logger.warning("Failed to persist refreshed tokens: %s", e)
+
+                    project_id = tokens.get(
+                        "project_id", model_config.get("project_id", "")
+                    )
+                    client = create_antigravity_client(
+                        access_token=access_token,
+                        project_id=project_id,
+                        model_name=model_config["name"],
+                        base_url=url,
+                        headers=headers,
+                        refresh_token=refresh_token,
+                        expires_at=expires_at,
+                        on_token_refreshed=on_token_refreshed,
+                    )
+
+                    # Use custom model with direct httpx client
+                    if AntigravityModel:
+                        model = AntigravityModel(
+                            model_name=model_config["name"],
+                            api_key=api_key
+                            or "",  # Antigravity uses OAuth, key may be empty
+                            base_url=url,
+                            http_client=client,
+                        )
+                    else:
+                        model = GeminiModel(
+                            model_name=model_config["name"],
+                            api_key=api_key or "",
+                            base_url=url,
+                            http_client=client,
+                        )
 
-                @property
-                def client(self) -> httpx.AsyncClient:
-                    _client = create_async_client(headers=headers, verify=verify)
-                    _client.base_url = self.base_url
-                    return _client
+                    return model
+
+                except ImportError:
+                    emit_warning(
+                        f"Antigravity transport not available; skipping model '{model_config.get('name')}'."
+                    )
+                    return None
+            else:
+                client = create_async_client(headers=headers, verify=verify)
 
-            google_gla = CustomGoogleGLAProvider(api_key=api_key)
-            model = GoogleModel(model_name=model_config["name"], provider=google_gla)
+            model = GeminiModel(
+                model_name=model_config["name"],
+                api_key=api_key,
+                base_url=url,
+                http_client=client,
+            )
             return model
         elif model_type == "cerebras":
+
+            class ZaiCerebrasProvider(CerebrasProvider):
+                def model_profile(self, model_name: str) -> ModelProfile | None:
+                    profile = super().model_profile(model_name)
+                    if model_name.startswith("zai"):
+                        from pydantic_ai.profiles.qwen import qwen_model_profile
+
+                        profile = profile.update(qwen_model_profile("qwen-3-coder"))
+                    return profile
+
             url, headers, verify, api_key = get_custom_config(model_config)
+            if not api_key:
+                emit_warning(
+                    f"API key is not set for Cerebras endpoint; skipping model '{model_config.get('name')}'."
+                )
+                return None
+            # Add Cerebras 3rd party integration header
+            headers["X-Cerebras-3rd-Party-Integration"] = "code-puppy"
             client = create_async_client(headers=headers, verify=verify)
             provider_args = dict(
                 api_key=api_key,
                 http_client=client,
             )
-            if api_key:
-                provider_args["api_key"] = api_key
-            provider = CerebrasProvider(**provider_args)
+            provider = ZaiCerebrasProvider(**provider_args)
 
             model = OpenAIChatModel(model_name=model_config["name"], provider=provider)
             setattr(model, "provider", provider)
@@ -325,17 +733,23 @@ class ModelFactory:
                 if api_key_config.startswith("$"):
                     # It's an environment variable reference
                     env_var_name = api_key_config[1:]  # Remove the $ prefix
-                    api_key = os.environ.get(env_var_name)
+                    api_key = get_api_key(env_var_name)
                     if api_key is None:
                         emit_warning(
-                            f"OpenRouter API key environment variable '{env_var_name}' not found or is empty; proceeding without API key."
+                            f"OpenRouter API key '{env_var_name}' not found (check config or environment); skipping model '{model_config.get('name')}'."
                         )
-                    else:
-                        # It's a raw API key value
-                        api_key = api_key_config
+                        return None
+                else:
+                    # It's a raw API key value
+                    api_key = api_key_config
             else:
-                # No API key in config, try to get it from the default environment variable
-                api_key = os.environ.get("OPENROUTER_API_KEY")
+                # No API key in config, try to get it from config or the default environment variable
+                api_key = get_api_key("OPENROUTER_API_KEY")
+                if api_key is None:
+                    emit_warning(
+                        f"OPENROUTER_API_KEY is not set (check config or environment); skipping OpenRouter model '{model_config.get('name')}'."
+                    )
+                    return None
 
             provider = OpenRouterProvider(api_key=api_key)
 
@@ -343,6 +757,143 @@ class ModelFactory:
             setattr(model, "provider", provider)
             return model
 
+        elif model_type == "gemini_oauth":
+            # Gemini OAuth models use the Code Assist API (cloudcode-pa.googleapis.com)
+            # This is a different API than the standard Generative Language API
+            try:
+                # Try user plugin first, then built-in plugin
+                try:
+                    from gemini_oauth.config import GEMINI_OAUTH_CONFIG
+                    from gemini_oauth.utils import (
+                        get_project_id,
+                        get_valid_access_token,
+                    )
+                except ImportError:
+                    from code_puppy.plugins.gemini_oauth.config import (
+                        GEMINI_OAUTH_CONFIG,
+                    )
+                    from code_puppy.plugins.gemini_oauth.utils import (
+                        get_project_id,
+                        get_valid_access_token,
+                    )
+            except ImportError as exc:
+                emit_warning(
+                    f"Gemini OAuth plugin not available; skipping model '{model_config.get('name')}'. "
+                    f"Error: {exc}"
+                )
+                return None
+
+            # Get a valid access token (refreshing if needed)
+            access_token = get_valid_access_token()
+            if not access_token:
+                emit_warning(
+                    f"Failed to get valid Gemini OAuth token; skipping model '{model_config.get('name')}'. "
+                    "Run /gemini-auth to re-authenticate."
+                )
+                return None
+
+            # Get project ID from stored tokens
+            project_id = get_project_id()
+            if not project_id:
+                emit_warning(
+                    f"No Code Assist project ID found; skipping model '{model_config.get('name')}'. "
+                    "Run /gemini-auth to re-authenticate."
+                )
+                return None
+
+            # Import the Code Assist model wrapper
+            from code_puppy.gemini_code_assist import GeminiCodeAssistModel
+
+            # Create the Code Assist model
+            model = GeminiCodeAssistModel(
+                model_name=model_config["name"],
+                access_token=access_token,
+                project_id=project_id,
+                api_base_url=GEMINI_OAUTH_CONFIG["api_base_url"],
+                api_version=GEMINI_OAUTH_CONFIG["api_version"],
+            )
+            return model
+
+        elif model_type == "chatgpt_oauth":
+            # ChatGPT OAuth models use the Codex API at chatgpt.com
+            try:
+                try:
+                    from chatgpt_oauth.config import CHATGPT_OAUTH_CONFIG
+                    from chatgpt_oauth.utils import (
+                        get_valid_access_token,
+                        load_stored_tokens,
+                    )
+                except ImportError:
+                    from code_puppy.plugins.chatgpt_oauth.config import (
+                        CHATGPT_OAUTH_CONFIG,
+                    )
+                    from code_puppy.plugins.chatgpt_oauth.utils import (
+                        get_valid_access_token,
+                        load_stored_tokens,
+                    )
+            except ImportError as exc:
+                emit_warning(
+                    f"ChatGPT OAuth plugin not available; skipping model '{model_config.get('name')}'. "
+                    f"Error: {exc}"
+                )
+                return None
+
+            # Get a valid access token (refreshing if needed)
+            access_token = get_valid_access_token()
+            if not access_token:
+                emit_warning(
+                    f"Failed to get valid ChatGPT OAuth token; skipping model '{model_config.get('name')}'. "
+                    "Run /chatgpt-auth to authenticate."
+                )
+                return None
+
+            # Get account_id from stored tokens (required for ChatGPT-Account-Id header)
+            tokens = load_stored_tokens()
+            account_id = tokens.get("account_id", "") if tokens else ""
+            if not account_id:
+                emit_warning(
+                    f"No account_id found in ChatGPT OAuth tokens; skipping model '{model_config.get('name')}'. "
+                    "Run /chatgpt-auth to re-authenticate."
+                )
+                return None
+
+            # Build headers for ChatGPT Codex API
+            originator = CHATGPT_OAUTH_CONFIG.get("originator", "codex_cli_rs")
+            client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.72.0")
+
+            headers = {
+                "ChatGPT-Account-Id": account_id,
+                "originator": originator,
+                "User-Agent": f"{originator}/{client_version}",
+            }
+            # Merge with any headers from model config
+            config_headers = model_config.get("custom_endpoint", {}).get("headers", {})
+            headers.update(config_headers)
+
+            # Get base URL - Codex API uses chatgpt.com, not api.openai.com
+            base_url = model_config.get("custom_endpoint", {}).get(
+                "url", CHATGPT_OAUTH_CONFIG["api_base_url"]
+            )
+
+            # Create HTTP client with Codex interceptor for store=false injection
+            from code_puppy.chatgpt_codex_client import create_codex_async_client
+
+            verify = get_cert_bundle_path()
+            client = create_codex_async_client(headers=headers, verify=verify)
+
+            provider = OpenAIProvider(
+                api_key=access_token,
+                base_url=base_url,
+                http_client=client,
+            )
+
+            # ChatGPT Codex API only supports Responses format
+            model = OpenAIResponsesModel(
+                model_name=model_config["name"], provider=provider
+            )
+            setattr(model, "provider", provider)
+            return model
+
         elif model_type == "round_robin":
             # Get the list of model names to use in the round-robin
             model_names = model_config.get("models")