code-puppy 0.0.361__py3-none-any.whl → 0.0.372__py3-none-any.whl

code_puppy/plugins/antigravity_oauth/transport.py

@@ -23,20 +23,78 @@ from .constants import (
 logger = logging.getLogger(__name__)
 
 
-def _inline_refs(
-    schema: dict, convert_unions: bool = False, simplify_for_claude: bool = False
-) -> dict:
+def _flatten_union_to_object(union_items: list, defs: dict, resolve_fn) -> dict:
+    """Flatten a union of object types into a single object with all properties.
+
+    For discriminated unions like EditFilePayload (ContentPayload | ReplacementsPayload | DeleteSnippetPayload),
+    we merge all object types into one with all properties marked as optional.
+    """
+    merged_properties = {}
+    has_string_type = False
+
+    for item in union_items:
+        if not isinstance(item, dict):
+            continue
+
+        # Resolve $ref first
+        if "$ref" in item:
+            ref_path = item["$ref"]
+            ref_name = None
+            if ref_path.startswith("#/$defs/"):
+                ref_name = ref_path[8:]
+            elif ref_path.startswith("#/definitions/"):
+                ref_name = ref_path[14:]
+            if ref_name and ref_name in defs:
+                item = copy.deepcopy(defs[ref_name])
+            else:
+                continue
+
+        # Check for string type (common fallback)
+        if item.get("type") == "string":
+            has_string_type = True
+            continue
+
+        # Skip null types
+        if item.get("type") == "null":
+            continue
+
+        # Merge properties from object types
+        if item.get("type") == "object" or "properties" in item:
+            props = item.get("properties", {})
+            for prop_name, prop_schema in props.items():
+                if prop_name not in merged_properties:
+                    # Resolve the property schema
+                    merged_properties[prop_name] = resolve_fn(
+                        copy.deepcopy(prop_schema)
+                    )
+
+    if not merged_properties:
+        # No object properties found, return string type as fallback
+        return {"type": "string"} if has_string_type else {"type": "object"}
+
+    # Build merged object - no required fields since any subset is valid
+    result = {
+        "type": "object",
+        "properties": merged_properties,
+    }
+
+    return result
+
+
+def _inline_refs(schema: dict, simplify_unions: bool = False) -> dict:
     """Inline $ref references and transform schema for Antigravity compatibility.
 
     - Inlines $ref references
     - Removes $defs, definitions, $schema, $id
-    - Optionally converts anyOf/oneOf/allOf to any_of/one_of/all_of (only for Gemini)
     - Removes unsupported fields like 'default', 'examples', 'const'
-    - For Claude: simplifies anyOf unions to single types
+    - When simplify_unions=True: flattens anyOf/oneOf unions:
+      - For unions of objects: merges into single object with all properties
+      - For simple unions (string | null): picks first non-null type
+      (required for both Gemini AND Claude - neither supports union types in function schemas!)
 
     Args:
-        convert_unions: If True, convert anyOf->any_of etc. (for Gemini).
-        simplify_for_claude: If True, simplify anyOf to single types.
+        simplify_unions: If True, simplify anyOf/oneOf unions.
+                         Required for Gemini and Claude models.
     """
     if not isinstance(schema, dict):
         return schema
@@ -47,31 +105,73 @@ def _inline_refs(
     # Extract $defs for reference resolution
     defs = schema.pop("$defs", schema.pop("definitions", {}))
 
-    def resolve_refs(
-        obj, convert_unions=convert_unions, simplify_for_claude=simplify_for_claude
-    ):
+    def resolve_refs(obj, simplify_unions=simplify_unions):
         """Recursively resolve $ref references and transform schema."""
         if isinstance(obj, dict):
-            # For Claude: simplify anyOf/oneOf unions to first non-null type
-            if simplify_for_claude:
+            # Handle anyOf/oneOf unions
+            if simplify_unions:
                 for union_key in ["anyOf", "oneOf"]:
                     if union_key in obj:
                         union = obj[union_key]
                         if isinstance(union, list):
-                            # Find first non-null type
+                            # Check if this is a complex union of objects (discriminated union)
+                            # vs a simple nullable type (string | null)
+                            object_count = 0
+                            has_refs = False
+                            for item in union:
+                                if isinstance(item, dict):
+                                    if "$ref" in item:
+                                        has_refs = True
+                                        object_count += 1
+                                    elif (
+                                        item.get("type") == "object"
+                                        or "properties" in item
+                                    ):
+                                        object_count += 1
+
+                            # If multiple objects or has refs, flatten to single object
+                            if object_count > 1 or has_refs:
+                                flattened = _flatten_union_to_object(
+                                    union, defs, resolve_refs
+                                )
+                                # Keep description if present
+                                if "description" in obj:
+                                    flattened["description"] = obj["description"]
+                                return flattened
+
+                            # Simple union - pick first non-null type
                             for item in union:
                                 if (
                                     isinstance(item, dict)
                                     and item.get("type") != "null"
                                 ):
-                                    # Replace the whole object with this type
                                     result = dict(item)
-                                    # Keep description if present
                                     if "description" in obj:
                                         result["description"] = obj["description"]
-                                    return resolve_refs(
-                                        result, convert_unions, simplify_for_claude
-                                    )
+                                    return resolve_refs(result, simplify_unions)
+
+            # Also handle allOf by merging all schemas
+            if simplify_unions and "allOf" in obj:
+                all_of = obj["allOf"]
+                if isinstance(all_of, list):
+                    merged = {}
+                    merged_properties = {}
+                    for item in all_of:
+                        if isinstance(item, dict):
+                            resolved_item = resolve_refs(item, simplify_unions)
+                            # Deep merge properties dicts
+                            if "properties" in resolved_item:
+                                merged_properties.update(
+                                    resolved_item.pop("properties")
+                                )
+                            merged.update(resolved_item)
+                    if merged_properties:
+                        merged["properties"] = merged_properties
+                    # Keep other fields from original object (except allOf)
+                    for k, v in obj.items():
+                        if k != "allOf":
+                            merged[k] = v
+                    return resolve_refs(merged, simplify_unions)
 
             # Check for $ref
             if "$ref" in obj:
@@ -111,34 +211,23 @@ def _inline_refs(
                 ):
                     continue
 
-                # For Claude: skip additionalProperties
-                if simplify_for_claude and key == "additionalProperties":
+                # Skip additionalProperties (not supported by Gemini or Claude)
+                if simplify_unions and key == "additionalProperties":
                     continue
 
-                # Optionally transform union types for Gemini
-                new_key = key
-                if convert_unions:
-                    if key == "anyOf":
-                        new_key = "any_of"
-                    elif key == "oneOf":
-                        new_key = "one_of"
-                    elif key == "allOf":
-                        new_key = "all_of"
-                    elif key == "additionalProperties":
-                        new_key = "additional_properties"
-
-                result[new_key] = resolve_refs(
-                    value, convert_unions, simplify_for_claude
-                )
+                # Skip any remaining union type keys that weren't simplified above
+                # (This shouldn't happen normally, but just in case)
+                if simplify_unions and key in ("anyOf", "oneOf", "allOf"):
+                    continue
+
+                result[key] = resolve_refs(value, simplify_unions)
             return result
         elif isinstance(obj, list):
-            return [
-                resolve_refs(item, convert_unions, simplify_for_claude) for item in obj
-            ]
+            return [resolve_refs(item, simplify_unions) for item in obj]
         else:
             return obj
 
-    return resolve_refs(schema, convert_unions, simplify_for_claude)
+    return resolve_refs(schema, simplify_unions)
 
 
 class UnwrappedResponse(httpx.Response):
@@ -264,17 +353,91 @@ class UnwrappedSSEResponse(httpx.Response):
 
 
 class AntigravityClient(httpx.AsyncClient):
-    """Custom httpx client that handles Antigravity request/response wrapping."""
+    """Custom httpx client that handles Antigravity request/response wrapping.
+
+    Supports proactive token refresh to prevent expiry during long sessions.
+    """
 
     def __init__(
         self,
         project_id: str = "",
         model_name: str = "",
+        refresh_token: str = "",
+        expires_at: Optional[float] = None,
+        on_token_refreshed: Optional[Any] = None,
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
         self.project_id = project_id
         self.model_name = model_name
+        self._refresh_token = refresh_token
+        self._expires_at = expires_at
+        self._on_token_refreshed = on_token_refreshed
+        self._refresh_lock = None  # Lazy init for async lock
+
+    async def _ensure_valid_token(self) -> None:
+        """Proactively refresh the access token if it's expired or about to expire.
+
+        This prevents 401 errors during long-running sessions by checking and
+        refreshing the token BEFORE making requests, not after they fail.
+        """
+        import asyncio
+
+        from .token import is_token_expired, refresh_access_token
+
+        # Skip if no refresh token configured
+        if not self._refresh_token:
+            return
+
+        # Check if token needs refresh (includes 60-second buffer)
+        if not is_token_expired(self._expires_at):
+            return
+
+        # Lazy init the async lock
+        if self._refresh_lock is None:
+            self._refresh_lock = asyncio.Lock()
+
+        async with self._refresh_lock:
+            # Double-check after acquiring lock (another coroutine may have refreshed)
+            if not is_token_expired(self._expires_at):
+                return
+
+            logger.debug("Proactively refreshing Antigravity access token...")
+
+            try:
+                # Run the synchronous refresh in a thread pool to avoid blocking
+                loop = asyncio.get_event_loop()
+                new_tokens = await loop.run_in_executor(
+                    None, refresh_access_token, self._refresh_token
+                )
+
+                if new_tokens:
+                    # Update internal state
+                    new_access_token = new_tokens.access_token
+                    self._expires_at = new_tokens.expires_at
+                    self._refresh_token = new_tokens.refresh_token
+
+                    # Update the Authorization header
+                    self.headers["Authorization"] = f"Bearer {new_access_token}"
+
+                    logger.info(
+                        "Proactively refreshed Antigravity token (expires in %ds)",
+                        int(self._expires_at - __import__("time").time()),
+                    )
+
+                    # Notify callback (e.g., to persist updated tokens)
+                    if self._on_token_refreshed:
+                        try:
+                            self._on_token_refreshed(new_tokens)
+                        except Exception as e:
+                            logger.warning("Token refresh callback failed: %s", e)
+                else:
+                    logger.warning(
+                        "Failed to proactively refresh token - request may fail with 401"
+                    )
+
+            except Exception as e:
+                logger.warning("Proactive token refresh error: %s", e)
 
     def _wrap_request(self, content: bytes, url: str) -> tuple[bytes, str, str, bool]:
         """Wrap request body in Antigravity envelope and transform URL.
@@ -331,15 +494,14 @@ class AntigravityClient(httpx.AsyncClient):
                                         )
 
                                     # Inline $refs and remove $defs from parameters
-                                    # Convert unions (anyOf->any_of) only for Gemini
-                                    # Simplify schemas for Claude (no anyOf, no additionalProperties)
+                                    # Simplify union types (anyOf/oneOf/allOf) for BOTH Gemini and Claude
+                                    # Neither API supports union types in function schemas!
                                     if "parameters" in func_decl:
                                         is_gemini = "gemini" in model.lower()
                                         is_claude = "claude" in model.lower()
                                         func_decl["parameters"] = _inline_refs(
                                             func_decl["parameters"],
-                                            convert_unions=is_gemini,
-                                            simplify_for_claude=is_claude,
+                                            simplify_unions=(is_gemini or is_claude),
                                         )
 
                 # Fix generationConfig for Antigravity compatibility
@@ -425,6 +587,9 @@ class AntigravityClient(httpx.AsyncClient):
         """Override send to intercept at the lowest level with endpoint fallback."""
         import asyncio
 
+        # Proactively refresh token BEFORE making the request
+        await self._ensure_valid_token()
+
         # Transform POST requests to Antigravity format
         if request.method == "POST" and request.content:
             new_content, new_path, new_query, is_claude_thinking = self._wrap_request(
@@ -638,14 +803,36 @@ class AntigravityClient(httpx.AsyncClient):
             return None
 
 
+# Type alias for token refresh callback
+TokenRefreshCallback = Any  # Callable[[OAuthTokens], None]
+
+
 def create_antigravity_client(
     access_token: str,
     project_id: str = "",
     model_name: str = "",
     base_url: str = "https://daily-cloudcode-pa.sandbox.googleapis.com",
     headers: Optional[Dict[str, str]] = None,
+    refresh_token: str = "",
+    expires_at: Optional[float] = None,
+    on_token_refreshed: Optional[TokenRefreshCallback] = None,
 ) -> AntigravityClient:
-    """Create an httpx client configured for Antigravity API."""
+    """Create an httpx client configured for Antigravity API.
+
+    Args:
+        access_token: The OAuth access token for authentication
+        project_id: The GCP project ID
+        model_name: The model name being used
+        base_url: The API base URL
+        headers: Additional headers to include
+        refresh_token: The OAuth refresh token for proactive token refresh
+        expires_at: Unix timestamp when the access token expires
+        on_token_refreshed: Callback called when token is proactively refreshed,
+                           receives OAuthTokens object to persist the new tokens
+
+    Returns:
+        An AntigravityClient configured for API requests with proactive token refresh
+    """
     # Start with Antigravity-specific headers
     default_headers = {
         "Authorization": f"Bearer {access_token}",
@@ -659,6 +846,9 @@ def create_antigravity_client(
     return AntigravityClient(
         project_id=project_id,
         model_name=model_name,
+        refresh_token=refresh_token,
+        expires_at=expires_at,
+        on_token_refreshed=on_token_refreshed,
         base_url=base_url,
         headers=default_headers,
         timeout=httpx.Timeout(180.0, connect=30.0),

code_puppy/plugins/chatgpt_oauth/register_callbacks.py

@@ -6,8 +6,8 @@ import os
 from typing import List, Optional, Tuple
 
 from code_puppy.callbacks import register_callback
-from code_puppy.config import set_model_name
 from code_puppy.messaging import emit_info, emit_success, emit_warning
+from code_puppy.model_switching import set_model_and_reload_agent
 
 from .config import CHATGPT_OAUTH_CONFIG, get_token_storage_path
 from .oauth_flow import run_oauth_flow
@@ -76,7 +76,7 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]:
 
     if name == "chatgpt-auth":
         run_oauth_flow()
-        set_model_name("chatgpt-gpt-5.2-codex")
+        set_model_and_reload_agent("chatgpt-gpt-5.2-codex")
         return True
 
     if name == "chatgpt-status":

code_puppy/plugins/claude_code_oauth/register_callbacks.py

@@ -12,8 +12,8 @@ from typing import Any, Dict, List, Optional, Tuple
 from urllib.parse import parse_qs, urlparse
 
 from code_puppy.callbacks import register_callback
-from code_puppy.config import set_model_name
 from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning
+from code_puppy.model_switching import set_model_and_reload_agent
 
 from ..oauth_puppy_html import oauth_failure_html, oauth_success_html
 from .config import CLAUDE_CODE_OAUTH_CONFIG, get_token_storage_path
@@ -181,31 +181,6 @@ 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)
@@ -245,9 +220,6 @@ 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:
@@ -261,7 +233,7 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]:
                 "Existing Claude Code tokens found. Continuing will overwrite them."
             )
         _perform_authentication()
-        set_model_name("claude-code-claude-opus-4-5-20251101")
+        set_model_and_reload_agent("claude-code-claude-opus-4-5-20251101")
         return True
 
     if name == "claude-code-status":

code_puppy/plugins/claude_code_oauth/utils.py

@@ -21,7 +21,10 @@ from .config import (
     get_token_storage_path,
 )
 
-TOKEN_REFRESH_BUFFER_SECONDS = 60
+# Proactive refresh buffer: refresh tokens 1 hour before expiration
+# This ensures smooth operation by refreshing during model requests,
+# well before the token actually expires
+TOKEN_REFRESH_BUFFER_SECONDS = 3600  # 1 hour
 
 logger = logging.getLogger(__name__)
 

code_puppy/pydantic_patches.py

@@ -121,6 +121,57 @@ def patch_process_message_history() -> None:
         pass
 
 
+def patch_tool_call_json_repair() -> None:
+    """Patch pydantic-ai's _call_tool to auto-repair malformed JSON arguments.
+
+    LLMs sometimes produce slightly broken JSON in tool calls (trailing commas,
+    missing quotes, etc.). This patch intercepts tool calls and runs json_repair
+    on the arguments before validation, preventing unnecessary retries.
+    """
+    try:
+        import json_repair
+        from pydantic_ai._tool_manager import ToolManager
+
+        # Store the original method
+        _original_call_tool = ToolManager._call_tool
+
+        async def _patched_call_tool(
+            self,
+            call,
+            *,
+            allow_partial: bool,
+            wrap_validation_errors: bool,
+            approved: bool,
+        ):
+            """Patched _call_tool that repairs malformed JSON before validation."""
+            # Only attempt repair if args is a string (JSON)
+            if isinstance(call.args, str) and call.args:
+                try:
+                    repaired = json_repair.repair_json(call.args)
+                    if repaired != call.args:
+                        # Update the call args with repaired JSON
+                        call.args = repaired
+                except Exception:
+                    pass  # If repair fails, let original validation handle it
+
+            # Call the original method
+            return await _original_call_tool(
+                self,
+                call,
+                allow_partial=allow_partial,
+                wrap_validation_errors=wrap_validation_errors,
+                approved=approved,
+            )
+
+        # Apply the patch
+        ToolManager._call_tool = _patched_call_tool
+
+    except ImportError:
+        pass  # json_repair or pydantic_ai not available
+    except Exception:
+        pass  # Don't crash on patch failure
+
+
 def apply_all_patches() -> None:
     """Apply all pydantic-ai monkey patches.
 
@@ -129,3 +180,4 @@ def apply_all_patches() -> None:
     patch_user_agent()
     patch_message_history_cleaning()
     patch_process_message_history()
+    patch_tool_call_json_repair()

code_puppy/tools/agent_tools.py

@@ -443,9 +443,9 @@ def register_invoke_agent(agent):
 
         terminal_session_token = set_terminal_session(f"terminal-{session_id}")
 
-        # Set browser session for Camoufox browser tools (qa-kitten, etc.)
+        # Set browser session for browser tools (qa-kitten, etc.)
         # This allows parallel agent invocations to each have their own browser
-        from code_puppy.tools.browser.camoufox_manager import (
+        from code_puppy.tools.browser.browser_manager import (
             set_browser_session,
         )
 
@@ -665,7 +665,7 @@ def register_invoke_agent(agent):
             # Reset terminal session context
             _terminal_session_var.reset(terminal_session_token)
             # Reset browser session context
-            from code_puppy.tools.browser.camoufox_manager import (
+            from code_puppy.tools.browser.browser_manager import (
                 _browser_session_var,
             )
 

code_puppy/tools/browser/__init__.py

@@ -5,7 +5,7 @@ This module provides browser-based terminal automation tools.
 
 from code_puppy.config import get_banner_color
 
-from .camoufox_manager import (
+from .browser_manager import (
     cleanup_all_browsers,
     get_browser_session,
     get_session_browser_manager,

code_puppy/tools/browser/browser_control.py

@@ -7,7 +7,7 @@ from pydantic_ai import RunContext
 from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning
 from code_puppy.tools.common import generate_group_id
 
-from .camoufox_manager import get_session_browser_manager
+from .browser_manager import get_session_browser_manager
 
 
 async def initialize_browser(

code_puppy/tools/browser/browser_interactions.py

@@ -7,7 +7,7 @@ from pydantic_ai import RunContext
 from code_puppy.messaging import emit_error, emit_info, emit_success
 from code_puppy.tools.common import generate_group_id
 
-from .camoufox_manager import get_session_browser_manager
+from .browser_manager import get_session_browser_manager
 
 
 async def click_element(

code_puppy/tools/browser/browser_locators.py

@@ -7,7 +7,7 @@ from pydantic_ai import RunContext
 from code_puppy.messaging import emit_info, emit_success
 from code_puppy.tools.common import generate_group_id
 
-from .camoufox_manager import get_session_browser_manager
+from .browser_manager import get_session_browser_manager
 
 
 async def find_by_role(