code-puppy 0.0.365__py3-none-any.whl → 0.0.367__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-0.0.365.dist-info → code_puppy-0.0.367.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-puppy
-Version: 0.0.365
+Version: 0.0.367
 Summary: Code generation agent
 Project-URL: repository, https://github.com/mpfaffenberger/code_puppy
 Project-URL: HomePage, https://github.com/mpfaffenberger/code_puppy
@@ -17,7 +17,6 @@ Classifier: Topic :: Software Development :: Code Generators
 Requires-Python: <3.14,>=3.11
 Requires-Dist: dbos>=2.5.0
 Requires-Dist: fastapi>=0.109.0
-Requires-Dist: google-genai>=1.20.0
 Requires-Dist: httpx[http2]>=0.24.1
 Requires-Dist: json-repair>=0.46.2
 Requires-Dist: mcp>=1.9.4
@@ -25,11 +24,12 @@ Requires-Dist: openai>=1.99.1
 Requires-Dist: pillow>=10.0.0
 Requires-Dist: playwright>=1.40.0
 Requires-Dist: prompt-toolkit>=3.0.52
-Requires-Dist: pydantic-ai-slim[anthropic,google,openai]==1.25.0
+Requires-Dist: pydantic-ai-slim[anthropic,openai]==1.26.0
 Requires-Dist: pydantic>=2.4.0
 Requires-Dist: pyfiglet>=0.8.post1
 Requires-Dist: python-dotenv>=1.0.0
 Requires-Dist: rapidfuzz>=3.13.0
+Requires-Dist: requests>=2.28.0
 Requires-Dist: rich>=13.4.2
 Requires-Dist: ripgrep==14.1.0
 Requires-Dist: tenacity>=8.2.0

{code_puppy-0.0.365.dist-info → code_puppy-0.0.367.dist-info}/RECORD

@@ -7,10 +7,11 @@ code_puppy/cli_runner.py,sha256=w5CLKgQYYaT7My3Cga2StXYol-u6DBxNzzUuhhsfhsA,3495
 code_puppy/config.py,sha256=blowBU3bBOdQSuLYKBUrb7f7CxHH_e25a_A4lQGsjgk,53494
 code_puppy/error_logging.py,sha256=a80OILCUtJhexI6a9GM-r5LqIdjvSRzggfgPp2jv1X0,3297
 code_puppy/gemini_code_assist.py,sha256=KGS7sO5OLc83nDF3xxS-QiU6vxW9vcm6hmzilu79Ef8,13867
-code_puppy/http_utils.py,sha256=H3N5Qz2B1CcsGUYOycGWAqoNMr2P1NCVluKX3aRwRqI,10358
+code_puppy/gemini_model.py,sha256=i8XXmx9s1eWEXpJ8U288w0yayTt6Nq8V-hxpUHhti4s,25984
+code_puppy/http_utils.py,sha256=SAH6EOdbR6Cbfmi-4EtHDqRDBUV5bWtGc-5nr44F0Is,10418
 code_puppy/keymap.py,sha256=IvMkTlB_bIqOWpbTpmftkdyjhtD5todXuEIw1zCZ4u0,3584
 code_puppy/main.py,sha256=82r3vZy_XcyEsenLn82BnUusaoyL3Bpm_Th_jKgqecE,273
-code_puppy/model_factory.py,sha256=ZLJFL1LWc-45vR6vb-RarAEjOuC0UVXZ1ButEtsnp08,38636
+code_puppy/model_factory.py,sha256=854Bo8wx59dOirAMhH0YuSXVHs-IxQBT_yrJGyVjKcA,39924
 code_puppy/model_utils.py,sha256=55TKNnGTXQlHJNqejL2PfQqQmChXfzOjJg-hlarfR7w,5551
 code_puppy/models.json,sha256=FMQdE_yvP_8y0xxt3K918UkFL9cZMYAqW1SfXcQkU_k,3105
 code_puppy/models_dev_api.json,sha256=wHjkj-IM_fx1oHki6-GqtOoCrRMR0ScK0f-Iz0UEcy8,548187
@@ -147,7 +148,7 @@ code_puppy/plugins/__init__.py,sha256=gWgrXWoFpl-3Mxz2DAvxKW6SkCWrOnw-hKsY9O7nHc
 code_puppy/plugins/oauth_puppy_html.py,sha256=Wpa-V_NlRiBAvo_OXHuR7wvOH_jSt8L9HSFGiab6xI0,13058
 code_puppy/plugins/antigravity_oauth/__init__.py,sha256=1miHihSqRNXO20Vh_Gn9M3Aa2szh0gtdSCaKKj9nq0Q,362
 code_puppy/plugins/antigravity_oauth/accounts.py,sha256=GQit2-K24bsopmTZyscFUq3M0cAEO5WutHWnipVdgz8,14304
-code_puppy/plugins/antigravity_oauth/antigravity_model.py,sha256=T81dp_j59BAKxcE8lIt3XKJQVB9Wb0IFugcNeR8mdVQ,29195
+code_puppy/plugins/antigravity_oauth/antigravity_model.py,sha256=ZFarvPYgYixQxEmEylIIbsSl7fmL8bYD8RXL_9wzDRA,26281
 code_puppy/plugins/antigravity_oauth/config.py,sha256=BoQgqf5I2XoHWnBBo9vhCIc_XwPj9Mbp0Z95ygWwt78,1362
 code_puppy/plugins/antigravity_oauth/constants.py,sha256=qsrA10JJvzNuY0OobvvwCQcoGpILBninllcUUMKkUrQ,4644
 code_puppy/plugins/antigravity_oauth/oauth.py,sha256=ZHXJtZP63l6brOpX1WdLfuUClIleA79-4y36YUJc6Wo,15137
@@ -155,7 +156,7 @@ code_puppy/plugins/antigravity_oauth/register_callbacks.py,sha256=uKIvfzH-dXj1g_
 code_puppy/plugins/antigravity_oauth/storage.py,sha256=LW1DkY6Z-GRbBDrIitT6glKemZptp3NzldIrLRqTAK0,8971
 code_puppy/plugins/antigravity_oauth/test_plugin.py,sha256=n0kjFG8Vt2n1j0GgTRSdSyhF0t9xxE8Ht60SH5CSwzw,11027
 code_puppy/plugins/antigravity_oauth/token.py,sha256=WbiFCkrZvChpGXvwIYsJMgqU9xdJ81KwR062lFlnL3U,5038
-code_puppy/plugins/antigravity_oauth/transport.py,sha256=GykwfwDmru5bw5WqGyVgxALQ_UyshMSmtDXsUfdgbD4,27729
+code_puppy/plugins/antigravity_oauth/transport.py,sha256=XgeG2Ys_Ui0_9cnrYWN7nENehFv76Rr2_QPAcIbPROI,35723
 code_puppy/plugins/antigravity_oauth/utils.py,sha256=mXHRv0l07r27VjtSsIy9rlpkUheP88RaM4x4M0O1mMY,5401
 code_puppy/plugins/chatgpt_oauth/__init__.py,sha256=Kjc6Hsz1sWvMD2OdAlWZvJRiKJSj4fx22boa-aVFKjA,189
 code_puppy/plugins/chatgpt_oauth/config.py,sha256=H_wAH9Duyn8WH2Kq8oe72uda-_4qu1uXLPun_SDdtsk,2023
@@ -207,10 +208,10 @@ code_puppy/tools/browser/chromium_terminal_manager.py,sha256=w1thQ_ACb6oV45L93TS
 code_puppy/tools/browser/terminal_command_tools.py,sha256=9byOZku-dwvTtCl532xt7Lumed_jTn0sLvUe_X75XCQ,19068
 code_puppy/tools/browser/terminal_screenshot_tools.py,sha256=J_21YO_495NvYgNFu9KQP6VYg2K_f8CtSdZuF94Yhnw,18448
 code_puppy/tools/browser/terminal_tools.py,sha256=F5LjVH3udSCFHmqC3O1UJLoLozZFZsEdX42jOmkqkW0,17853
-code_puppy-0.0.365.data/data/code_puppy/models.json,sha256=FMQdE_yvP_8y0xxt3K918UkFL9cZMYAqW1SfXcQkU_k,3105
-code_puppy-0.0.365.data/data/code_puppy/models_dev_api.json,sha256=wHjkj-IM_fx1oHki6-GqtOoCrRMR0ScK0f-Iz0UEcy8,548187
-code_puppy-0.0.365.dist-info/METADATA,sha256=cMC5fk4Nm2d_ZWj-1oLtK7vVJujRFjMOP-y3nNvRLh8,27611
-code_puppy-0.0.365.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-code_puppy-0.0.365.dist-info/entry_points.txt,sha256=Tp4eQC99WY3HOKd3sdvb22vZODRq0XkZVNpXOag_KdI,91
-code_puppy-0.0.365.dist-info/licenses/LICENSE,sha256=31u8x0SPgdOq3izJX41kgFazWsM43zPEF9eskzqbJMY,1075
-code_puppy-0.0.365.dist-info/RECORD,,
+code_puppy-0.0.367.data/data/code_puppy/models.json,sha256=FMQdE_yvP_8y0xxt3K918UkFL9cZMYAqW1SfXcQkU_k,3105
+code_puppy-0.0.367.data/data/code_puppy/models_dev_api.json,sha256=wHjkj-IM_fx1oHki6-GqtOoCrRMR0ScK0f-Iz0UEcy8,548187
+code_puppy-0.0.367.dist-info/METADATA,sha256=fgyQYowvJcFxYCStXVQwHncMqvLo1JBztvtaAl2U7Yg,27600
+code_puppy-0.0.367.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+code_puppy-0.0.367.dist-info/entry_points.txt,sha256=Tp4eQC99WY3HOKd3sdvb22vZODRq0XkZVNpXOag_KdI,91
+code_puppy-0.0.367.dist-info/licenses/LICENSE,sha256=31u8x0SPgdOq3izJX41kgFazWsM43zPEF9eskzqbJMY,1075
+code_puppy-0.0.367.dist-info/RECORD,,