code-puppy 0.0.348__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
@@ -393,6 +555,7 @@ class AntigravityClient(httpx.AsyncClient):
                 "request": original_body,
                 "userAgent": "antigravity",
                 "requestId": request_id,
+                "requestType": "agent",
             }
 
             # Transform URL to Antigravity format
@@ -424,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(
@@ -637,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}",
@@ -658,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/plugins/frontend_emitter/__init__.py

@@ -0,0 +1,25 @@
+"""Frontend emitter plugin for Code Puppy.
+
+This plugin provides event emission capabilities for frontend integration,
+allowing WebSocket handlers to subscribe to real-time events from the
+agent system including tool calls, streaming events, and agent invocations.
+
+Usage:
+    from code_puppy.plugins.frontend_emitter.emitter import (
+        emit_event,
+        subscribe,
+        unsubscribe,
+        get_recent_events,
+    )
+
+    # Subscribe to events
+    queue = subscribe()
+
+    # Process events in your WebSocket handler
+    while True:
+        event = await queue.get()
+        await websocket.send_json(event)
+
+    # Clean up
+    unsubscribe(queue)
+"""

code_puppy/plugins/frontend_emitter/emitter.py

@@ -0,0 +1,121 @@
+"""Event emitter for frontend integration.
+
+Provides a global event queue that WebSocket handlers can subscribe to.
+Events are JSON-serializable dicts with type, timestamp, and data.
+"""
+
+import asyncio
+import logging
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Set
+from uuid import uuid4
+
+from code_puppy.config import (
+    get_frontend_emitter_enabled,
+    get_frontend_emitter_max_recent_events,
+    get_frontend_emitter_queue_size,
+)
+
+logger = logging.getLogger(__name__)
+
+# Global state for event distribution
+_subscribers: Set[asyncio.Queue[Dict[str, Any]]] = set()
+_recent_events: List[Dict[str, Any]] = []  # Keep last N events for new subscribers
+
+
+def emit_event(event_type: str, data: Any = None) -> None:
+    """Emit an event to all subscribers.
+
+    Creates a structured event dict with unique ID, type, timestamp, and data,
+    then broadcasts it to all active subscriber queues.
+
+    Args:
+        event_type: Type of event (e.g., "tool_call_start", "stream_token")
+        data: Event data payload - should be JSON-serializable
+    """
+    # Early return if emitter is disabled
+    if not get_frontend_emitter_enabled():
+        return
+
+    event: Dict[str, Any] = {
+        "id": str(uuid4()),
+        "type": event_type,
+        "timestamp": datetime.now(timezone.utc).isoformat(),
+        "data": data or {},
+    }
+
+    # Store in recent events for replay to new subscribers
+    max_recent = get_frontend_emitter_max_recent_events()
+    _recent_events.append(event)
+    if len(_recent_events) > max_recent:
+        _recent_events.pop(0)
+
+    # Broadcast to all active subscribers
+    for subscriber_queue in _subscribers.copy():
+        try:
+            subscriber_queue.put_nowait(event)
+        except asyncio.QueueFull:
+            logger.warning(f"Subscriber queue full, dropping event: {event_type}")
+        except Exception as e:
+            logger.error(f"Failed to emit event to subscriber: {e}")
+
+
+def subscribe() -> asyncio.Queue[Dict[str, Any]]:
+    """Subscribe to events.
+
+    Creates and returns a new async queue that will receive all future events.
+    The queue has a configurable max size (via frontend_emitter_queue_size)
+    to prevent unbounded memory growth if the subscriber is slow to process events.
+
+    Returns:
+        An asyncio.Queue that will receive event dictionaries.
+    """
+    queue_size = get_frontend_emitter_queue_size()
+    queue: asyncio.Queue[Dict[str, Any]] = asyncio.Queue(maxsize=queue_size)
+    _subscribers.add(queue)
+    logger.debug(f"New subscriber added, total subscribers: {len(_subscribers)}")
+    return queue
+
+
+def unsubscribe(queue: asyncio.Queue[Dict[str, Any]]) -> None:
+    """Unsubscribe from events.
+
+    Removes the queue from the subscriber set. Safe to call even if the queue
+    was never subscribed or already unsubscribed.
+
+    Args:
+        queue: The queue returned from subscribe()
+    """
+    _subscribers.discard(queue)
+    logger.debug(f"Subscriber removed, remaining subscribers: {len(_subscribers)}")
+
+
+def get_recent_events() -> List[Dict[str, Any]]:
+    """Get recent events for new subscribers.
+
+    Returns a copy of the most recent events (up to frontend_emitter_max_recent_events).
+    Useful for allowing new WebSocket connections to "catch up" on
+    recent activity.
+
+    Returns:
+        A list of recent event dictionaries.
+    """
+    return _recent_events.copy()
+
+
+def get_subscriber_count() -> int:
+    """Get the current number of active subscribers.
+
+    Returns:
+        Number of active subscriber queues.
+    """
+    return len(_subscribers)
+
+
+def clear_recent_events() -> None:
+    """Clear the recent events buffer.
+
+    Useful for testing or resetting state.
+    """
+    _recent_events.clear()
+    logger.debug("Recent events cleared")