code-puppy 0.0.302__py3-none-any.whl → 0.0.323__py3-none-any.whl

code_puppy/plugins/chatgpt_oauth/oauth_flow.py

@@ -19,7 +19,6 @@ from .config import CHATGPT_OAUTH_CONFIG
 from .utils import (
     add_models_to_extra_config,
     assign_redirect_uri,
-    fetch_chatgpt_models,
     load_stored_tokens,
     parse_jwt_claims,
     prepare_oauth_context,
@@ -318,12 +317,12 @@ def run_oauth_flow() -> None:
         )
 
     if api_key:
-        emit_info("Fetching available ChatGPT models…")
-        models = fetch_chatgpt_models(api_key)
+        emit_info("Registering ChatGPT Codex models…")
+        from .utils import DEFAULT_CODEX_MODELS
+
+        models = DEFAULT_CODEX_MODELS
         if models:
-            if add_models_to_extra_config(models, api_key):
+            if add_models_to_extra_config(models):
                 emit_success(
                     "ChatGPT models registered. Use the `chatgpt-` prefix in /model."
                 )
-        else:
-            emit_warning("API key obtained, but model list could not be fetched.")

code_puppy/plugins/chatgpt_oauth/register_callbacks.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import os
 from typing import List, Optional, Tuple
 
+from code_puppy.callbacks import register_callback
 from code_puppy.messaging import emit_info, emit_success, emit_warning
 
 from .config import CHATGPT_OAUTH_CONFIG, get_token_storage_path
@@ -87,6 +88,5 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]:
     return None
 
 
-# Temporarily disabled - chatgpt-oauth plugin not working yet
-# register_callback("custom_command_help", _custom_help)
-# register_callback("custom_command", _handle_custom_command)
+register_callback("custom_command_help", _custom_help)
+register_callback("custom_command", _handle_custom_command)

code_puppy/plugins/chatgpt_oauth/test_plugin.py

@@ -235,25 +235,41 @@ def test_exchange_code_for_tokens(mock_post):
 
 @patch("code_puppy.plugins.chatgpt_oauth.utils.requests.get")
 def test_fetch_chatgpt_models(mock_get):
-    """Test fetching models from OpenAI API."""
+    """Test fetching models from ChatGPT Codex API."""
     mock_response = MagicMock()
     mock_response.status_code = 200
+    # New response format uses "models" key with "slug" field
     mock_response.json.return_value = {
-        "data": [
-            {"id": "gpt-4o"},
-            {"id": "gpt-3.5-turbo"},
-            {"id": "whisper-1"},  # Should be filtered out
-            {"id": "o1-preview"},
+        "models": [
+            {"slug": "gpt-4o"},
+            {"slug": "gpt-3.5-turbo"},
+            {"slug": "o1-preview"},
+            {"slug": "codex-mini"},
         ]
     }
     mock_get.return_value = mock_response
 
-    models = utils.fetch_chatgpt_models("test_api_key")
+    models = utils.fetch_chatgpt_models("test_access_token", "test_account_id")
     assert models is not None
     assert "gpt-4o" in models
     assert "gpt-3.5-turbo" in models
     assert "o1-preview" in models
-    assert "whisper-1" not in models  # Should be filtered
+    assert "codex-mini" in models
+
+
+@patch("code_puppy.plugins.chatgpt_oauth.utils.requests.get")
+def test_fetch_chatgpt_models_fallback(mock_get):
+    """Test that fetch_chatgpt_models returns default list on API failure."""
+    mock_response = MagicMock()
+    mock_response.status_code = 404
+    mock_response.text = '{"detail":"Not Found"}'
+    mock_get.return_value = mock_response
+
+    models = utils.fetch_chatgpt_models("test_access_token", "test_account_id")
+    assert models is not None
+    # Should return default models
+    assert "gpt-5.2" in models
+    assert "gpt-4o" in models
 
 
 def test_add_models_to_chatgpt_config(tmp_path):
@@ -262,14 +278,13 @@ def test_add_models_to_chatgpt_config(tmp_path):
         config, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json"
     ):
         models = ["gpt-4o", "gpt-3.5-turbo"]
-        api_key = "sk-test"
 
-        assert utils.add_models_to_extra_config(models, api_key)
+        assert utils.add_models_to_extra_config(models)
 
         loaded = utils.load_chatgpt_models()
         assert "chatgpt-gpt-4o" in loaded
         assert "chatgpt-gpt-3.5-turbo" in loaded
-        assert loaded["chatgpt-gpt-4o"]["type"] == "openai"
+        assert loaded["chatgpt-gpt-4o"]["type"] == "chatgpt_oauth"
         assert loaded["chatgpt-gpt-4o"]["name"] == "gpt-4o"
         assert loaded["chatgpt-gpt-4o"]["oauth_source"] == "chatgpt-oauth-plugin"
 

code_puppy/plugins/chatgpt_oauth/utils.py

@@ -149,6 +149,98 @@ def load_stored_tokens() -> Optional[Dict[str, Any]]:
     return None
 
 
+def get_valid_access_token() -> Optional[str]:
+    """Get a valid access token, refreshing if expired.
+
+    Returns:
+        Valid access token string, or None if not authenticated or refresh failed.
+    """
+    tokens = load_stored_tokens()
+    if not tokens:
+        logger.debug("No stored ChatGPT OAuth tokens found")
+        return None
+
+    access_token = tokens.get("access_token")
+    if not access_token:
+        logger.debug("No access_token in stored tokens")
+        return None
+
+    # Check if token is expired by parsing JWT claims
+    claims = parse_jwt_claims(access_token)
+    if claims:
+        exp = claims.get("exp")
+        if exp and isinstance(exp, (int, float)):
+            # Add 30 second buffer before expiry
+            if time.time() > exp - 30:
+                logger.info("ChatGPT OAuth token expired, attempting refresh")
+                refreshed = refresh_access_token()
+                if refreshed:
+                    return refreshed
+                logger.warning("Token refresh failed")
+                return None
+
+    return access_token
+
+
+def refresh_access_token() -> Optional[str]:
+    """Refresh the access token using the refresh token.
+
+    Returns:
+        New access token if refresh succeeded, None otherwise.
+    """
+    tokens = load_stored_tokens()
+    if not tokens:
+        return None
+
+    refresh_token = tokens.get("refresh_token")
+    if not refresh_token:
+        logger.debug("No refresh_token available")
+        return None
+
+    payload = {
+        "grant_type": "refresh_token",
+        "refresh_token": refresh_token,
+        "client_id": CHATGPT_OAUTH_CONFIG["client_id"],
+    }
+
+    headers = {
+        "Content-Type": "application/x-www-form-urlencoded",
+    }
+
+    try:
+        response = requests.post(
+            CHATGPT_OAUTH_CONFIG["token_url"],
+            data=payload,
+            headers=headers,
+            timeout=30,
+        )
+
+        if response.status_code == 200:
+            new_tokens = response.json()
+            # Merge with existing tokens (preserve account_id, etc.)
+            tokens.update(
+                {
+                    "access_token": new_tokens.get("access_token"),
+                    "refresh_token": new_tokens.get("refresh_token", refresh_token),
+                    "id_token": new_tokens.get("id_token", tokens.get("id_token")),
+                    "last_refresh": datetime.datetime.now(datetime.timezone.utc)
+                    .isoformat()
+                    .replace("+00:00", "Z"),
+                }
+            )
+            if save_tokens(tokens):
+                logger.info("Successfully refreshed ChatGPT OAuth token")
+                return tokens["access_token"]
+        else:
+            logger.error(
+                "Token refresh failed: %s - %s", response.status_code, response.text
+            )
+    except Exception as exc:
+        logger.error("Token refresh error: %s", exc)
+
+    return None
+
+
 def save_tokens(tokens: Dict[str, Any]) -> bool:
     if tokens is None:
         raise TypeError("tokens cannot be None")
@@ -248,103 +340,126 @@ def exchange_code_for_tokens(
     return None
 
 
-def fetch_chatgpt_models(api_key: str) -> Optional[List[str]]:
-    """Fetch available models from OpenAI API.
+# Default models available via ChatGPT Codex API
+# These are the known models that work with ChatGPT OAuth tokens
+# Based on codex-rs CLI and shell-scripts/codex-call.sh
+DEFAULT_CODEX_MODELS = [
+    "gpt-5.2",
+    "gpt-5.2-codex",
+]
+
 
-    Makes a real HTTP GET request to OpenAI's models endpoint and filters
-    the results to include only GPT series models while preserving server order.
+def fetch_chatgpt_models(access_token: str, account_id: str) -> Optional[List[str]]:
+    """Fetch available models from ChatGPT Codex API.
+
+    Attempts to fetch models from the API, but falls back to a default list
+    of known Codex-compatible models if the API is unavailable.
 
     Args:
-        api_key: OpenAI API key for authentication
+        access_token: OAuth access token for authentication
+        account_id: ChatGPT account ID (required for the API)
 
     Returns:
-        List of filtered model IDs preserving server order, or None if request fails
+        List of model IDs, or default list if API fails
     """
-    # Build the models URL, ensuring it ends with /v1/models
-    base_url = CHATGPT_OAUTH_CONFIG["api_base_url"].rstrip("/")
-    models_url = f"{base_url}/v1/models"
+    import platform
 
-    # Blocklist of model IDs to exclude
-    blocklist = {"whisper-1"}
+    # Build the models URL with client version
+    client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.72.0")
+    base_url = CHATGPT_OAUTH_CONFIG["api_base_url"].rstrip("/")
+    models_url = f"{base_url}/models"
+
+    # Build User-Agent to match codex-rs CLI format
+    originator = CHATGPT_OAUTH_CONFIG.get("originator", "codex_cli_rs")
+    os_name = platform.system()
+    if os_name == "Darwin":
+        os_name = "Mac OS"
+    os_version = platform.release()
+    arch = platform.machine()
+    user_agent = (
+        f"{originator}/{client_version} ({os_name} {os_version}; {arch}) "
+        "Terminal_Codex_CLI"
+    )
 
     headers = {
-        "Authorization": f"Bearer {api_key}",
+        "Authorization": f"Bearer {access_token}",
+        "ChatGPT-Account-Id": account_id,
+        "User-Agent": user_agent,
+        "originator": originator,
+        "Accept": "application/json",
     }
 
-    try:
-        response = requests.get(models_url, headers=headers, timeout=30)
-
-        if response.status_code != 200:
-            logger.error(
-                "Failed to fetch models: HTTP %d - %s",
-                response.status_code,
-                response.text,
-            )
-            return None
-
-        # Parse JSON response
-        try:
-            data = response.json()
-            if "data" not in data or not isinstance(data["data"], list):
-                logger.error("Invalid response format: missing 'data' list")
-                return None
-        except (json.JSONDecodeError, ValueError) as exc:
-            logger.error("Failed to parse JSON response: %s", exc)
-            return None
-
-        # Filter models: start with "gpt-" or "o1-" and not in blocklist
-        filtered_models = []
-        seen_models = set()  # For deduplication while preserving order
-
-        for model in data["data"]:
-            # Skip None entries
-            if model is None:
-                continue
+    # Query params
+    params = {"client_version": client_version}
 
-            model_id = model.get("id")
-            if not model_id:
-                continue
-
-            # Skip if already seen (deduplication)
-            if model_id in seen_models:
-                continue
-
-            # Check if model starts with allowed prefixes and not in blocklist
-            if (
-                model_id.startswith("gpt-") or model_id.startswith("o1-")
-            ) and model_id not in blocklist:
-                filtered_models.append(model_id)
-                seen_models.add(model_id)
+    try:
+        response = requests.get(models_url, headers=headers, params=params, timeout=30)
 
-        return filtered_models
+        if response.status_code == 200:
+            # Parse JSON response
+            try:
+                data = response.json()
+                # The response has a "models" key with list of model objects
+                if "models" in data and isinstance(data["models"], list):
+                    models = []
+                    for model in data["models"]:
+                        if model is None:
+                            continue
+                        model_id = (
+                            model.get("slug") or model.get("id") or model.get("name")
+                        )
+                        if model_id:
+                            models.append(model_id)
+                    if models:
+                        return models
+            except (json.JSONDecodeError, ValueError) as exc:
+                logger.warning("Failed to parse models response: %s", exc)
+
+        # API didn't return valid models, use default list
+        logger.info(
+            "Models endpoint returned %d, using default model list",
+            response.status_code,
+        )
 
     except requests.exceptions.Timeout:
-        logger.error("Timeout while fetching models after 30 seconds")
-        return None
+        logger.warning("Timeout fetching models, using default list")
     except requests.exceptions.RequestException as exc:
-        logger.error("Network error while fetching models: %s", exc)
-        return None
+        logger.warning("Network error fetching models: %s, using default list", exc)
     except Exception as exc:
-        logger.error("Unexpected error while fetching models: %s", exc)
-        return None
+        logger.warning("Error fetching models: %s, using default list", exc)
 
+    # Return default models when API fails
+    logger.info("Using default Codex models: %s", DEFAULT_CODEX_MODELS)
+    return DEFAULT_CODEX_MODELS
 
-def add_models_to_extra_config(models: List[str], api_key: str) -> bool:
+
+def add_models_to_extra_config(models: List[str]) -> bool:
     """Add ChatGPT models to chatgpt_models.json configuration."""
     try:
         chatgpt_models = load_chatgpt_models()
         added = 0
         for model_name in models:
             prefixed = f"{CHATGPT_OAUTH_CONFIG['prefix']}{model_name}"
+
+            # Determine supported settings based on model type
+            # All GPT-5.x models support reasoning_effort and verbosity
+            supported_settings = ["reasoning_effort", "verbosity"]
+
+            # Only codex models support xhigh reasoning effort
+            # Regular gpt-5.2 is capped at "high"
+            is_codex = "codex" in model_name.lower()
+
             chatgpt_models[prefixed] = {
-                "type": "openai",
+                "type": "chatgpt_oauth",
                 "name": model_name,
                 "custom_endpoint": {
+                    # Codex API uses chatgpt.com/backend-api/codex, not api.openai.com
                     "url": CHATGPT_OAUTH_CONFIG["api_base_url"],
-                    "api_key": "${" + CHATGPT_OAUTH_CONFIG["api_key_env_var"] + "}",
                 },
                 "context_length": CHATGPT_OAUTH_CONFIG["default_context_length"],
                 "oauth_source": "chatgpt-oauth-plugin",
+                "supported_settings": supported_settings,
+                "supports_xhigh_reasoning": is_codex,
             }
             added += 1
         if save_chatgpt_models(chatgpt_models):

code_puppy/plugins/claude_code_oauth/register_callbacks.py

@@ -180,6 +180,31 @@ def _custom_help() -> List[Tuple[str, str]]:
     ]
 
 
+def _reload_current_agent() -> None:
+    """Reload the current agent so new auth tokens are picked up immediately."""
+    try:
+        from code_puppy.agents import get_current_agent
+
+        current_agent = get_current_agent()
+        if current_agent is None:
+            logger.debug("No current agent to reload")
+            return
+
+        # JSON agents may need to refresh their config before reload
+        if hasattr(current_agent, "refresh_config"):
+            try:
+                current_agent.refresh_config()
+            except Exception:
+                # Non-fatal, continue to reload
+                pass
+
+        current_agent.reload_code_generation_agent()
+        emit_info("Active agent reloaded with new authentication")
+    except Exception as e:
+        emit_warning(f"Authentication succeeded but agent reload failed: {e}")
+        logger.exception("Failed to reload agent after authentication")
+
+
 def _perform_authentication() -> None:
     context = prepare_oauth_context()
     code = _await_callback(context)
@@ -219,6 +244,9 @@ def _perform_authentication() -> None:
             "Claude Code models added to your configuration. Use the `claude-code-` prefix!"
         )
 
+    # Reload the current agent so the new auth token is picked up immediately
+    _reload_current_agent()
+
 
 def _handle_custom_command(command: str, name: str) -> Optional[bool]:
     if not name:

code_puppy/plugins/claude_code_oauth/utils.py

@@ -368,6 +368,7 @@ def add_models_to_extra_config(models: List[str]) -> bool:
                     "temperature",
                     "extended_thinking",
                     "budget_tokens",
+                    "interleaved_thinking",
                 ],
             }
             added += 1

code_puppy/plugins/shell_safety/agent_shell_safety.py

@@ -5,13 +5,12 @@ It's designed to be ultra-lightweight with a concise prompt (<200 tokens) and
 uses structured output for reliable parsing.
 """
 
-import asyncio
 from typing import TYPE_CHECKING, List
 
 from code_puppy.agents.base_agent import BaseAgent
 
 if TYPE_CHECKING:
-    from code_puppy.tools.command_runner import ShellSafetyAssessment
+    pass
 
 
 class ShellSafetyAgent(BaseAgent):
@@ -68,119 +67,3 @@ class ShellSafetyAgent(BaseAgent):
     def get_available_tools(self) -> List[str]:
         """This agent uses no tools - pure reasoning only."""
         return []
-
-    async def assess_command(
-        self, command: str, cwd: str | None = None
-    ) -> "ShellSafetyAssessment":
-        """Assess the safety risk of a shell command.
-
-        Args:
-            command: The shell command to assess
-            cwd: Optional working directory context
-
-        Returns:
-            ShellSafetyAssessment with risk level and reasoning
-
-        Note:
-            On timeout or error, defaults to 'high' risk with error reasoning
-            to fail safe. Optionally uses DBOS for durable execution tracking.
-        """
-        import uuid
-
-        from pydantic_ai import Agent, UsageLimits
-
-        from code_puppy.config import get_use_dbos
-        from code_puppy.model_factory import ModelFactory
-        from code_puppy.tools.command_runner import ShellSafetyAssessment
-
-        try:
-            # Build the assessment prompt
-            prompt = f"Assess this shell command:\n\nCommand: {command}"
-            if cwd:
-                prompt += f"\nWorking directory: {cwd}"
-
-            # Get the current model
-            model_name = self.get_model_name()
-            models_config = ModelFactory.load_config()
-
-            if model_name not in models_config:
-                # Fall back to high risk if model config fails
-                return ShellSafetyAssessment(
-                    risk="high",
-                    reasoning="Model configuration unavailable - failing safe",
-                    is_fallback=True,
-                )
-
-            model = ModelFactory.get_model(model_name, models_config)
-
-            # Handle claude-code models: swap instructions and prepend system prompt
-            from code_puppy.model_utils import prepare_prompt_for_model
-
-            instructions = self.get_system_prompt()
-            prepared = prepare_prompt_for_model(model_name, instructions, prompt)
-            instructions = prepared.instructions
-            prompt = prepared.user_prompt
-
-            from code_puppy.model_factory import make_model_settings
-
-            model_settings = make_model_settings(model_name)
-
-            temp_agent = Agent(
-                model=model,
-                system_prompt=instructions,
-                retries=2,  # Increase from 1 to 2 for better reliability
-                output_type=ShellSafetyAssessment,
-                model_settings=model_settings,
-            )
-
-            # Generate unique agent name and workflow ID for DBOS (if enabled)
-            agent_name = f"shell-safety-{uuid.uuid4().hex[:8]}"
-            workflow_id = f"shell-safety-{uuid.uuid4().hex[:8]}"
-
-            # Wrap with DBOS if enabled (same pattern as agent_tools.py)
-            if get_use_dbos():
-                from pydantic_ai.durable_exec.dbos import DBOSAgent
-
-                dbos_agent = DBOSAgent(temp_agent, name=agent_name)
-                temp_agent = dbos_agent
-
-            # Run the agent as a cancellable task
-            # Import the shared task registry for cancellation support
-            from code_puppy.tools.agent_tools import _active_subagent_tasks
-
-            if get_use_dbos():
-                from dbos import DBOS, SetWorkflowID
-
-                with SetWorkflowID(workflow_id):
-                    task = asyncio.create_task(
-                        temp_agent.run(
-                            prompt,
-                            usage_limits=UsageLimits(request_limit=3),
-                        )
-                    )
-                    _active_subagent_tasks.add(task)
-            else:
-                task = asyncio.create_task(
-                    temp_agent.run(
-                        prompt,
-                        usage_limits=UsageLimits(request_limit=3),
-                    )
-                )
-                _active_subagent_tasks.add(task)
-
-            try:
-                result = await task
-            finally:
-                _active_subagent_tasks.discard(task)
-                if task.cancelled():
-                    if get_use_dbos():
-                        DBOS.cancel_workflow(workflow_id)
-
-            return result.output
-
-        except Exception as e:
-            return ShellSafetyAssessment(
-                risk="high",
-                reasoning=f"Safety assessment failed: {str(e)[:200]} - failing safe",
-                is_fallback=True,
-            )

code_puppy/plugins/shell_safety/register_callbacks.py

@@ -7,12 +7,42 @@ and assesses their safety risk before execution.
 from typing import Any, Dict, Optional
 
 from code_puppy.callbacks import register_callback
-from code_puppy.config import get_safety_permission_level, get_yolo_mode
+from code_puppy.config import (
+    get_global_model_name,
+    get_safety_permission_level,
+    get_yolo_mode,
+)
 from code_puppy.messaging import emit_info
 from code_puppy.plugins.shell_safety.command_cache import (
     cache_assessment,
     get_cached_assessment,
 )
+from code_puppy.tools.command_runner import ShellSafetyAssessment
+
+# OAuth model prefixes - these models have their own safety mechanisms
+OAUTH_MODEL_PREFIXES = (
+    "claude-code-",  # Anthropic OAuth
+    "chatgpt-",  # OpenAI OAuth
+    "gemini-oauth",  # Google OAuth
+)
+
+
+def is_oauth_model(model_name: str | None) -> bool:
+    """Check if the model is an OAuth model that should skip safety checks.
+
+    OAuth models have their own built-in safety mechanisms, so we skip
+    the shell safety callback to avoid redundant checks and potential bugs.
+
+    Args:
+        model_name: The name of the current model
+
+    Returns:
+        True if the model is an OAuth model, False otherwise
+    """
+    if not model_name:
+        return False
+    return model_name.startswith(OAUTH_MODEL_PREFIXES)
+
 
 # Risk level hierarchy for numeric comparison
 # Lower numbers = safer commands, higher numbers = more dangerous
@@ -68,6 +98,11 @@ async def shell_safety_callback(
         None if command is safe to proceed
         Dict with rejection info if command should be blocked
     """
+    # Skip safety checks for OAuth models - they have their own safety mechanisms
+    current_model = get_global_model_name()
+    if is_oauth_model(current_model):
+        return None
+
     # Only check safety in yolo_mode - otherwise user is reviewing manually
     yolo_mode = get_yolo_mode()
     if not yolo_mode:
@@ -108,8 +143,14 @@ async def shell_safety_callback(
         # Create agent and assess command
         agent = ShellSafetyAgent()
 
-        # Run async assessment (we're in an async callback now!)
-        assessment = await agent.assess_command(command, cwd)
+        # Build the assessment prompt with optional cwd context
+        prompt = f"Assess this shell command:\n\nCommand: {command}"
+        if cwd:
+            prompt += f"\nWorking directory: {cwd}"
+
+        # Run async assessment with structured output type
+        result = await agent.run_with_mcp(prompt, output_type=ShellSafetyAssessment)
+        assessment = result.output
 
         # Cache the result for future use, but only if it's not a fallback assessment
         if not getattr(assessment, "is_fallback", False):