refactor(agent/agent_init): wip checkpoint — phase-helper split of init, param tables, header factory table, compression parse helpers; mixin compaction

2026-09-02 19:07:13 -07:00
parent 113f04616b
commit a5e537cfa9
5 changed files with 1095 additions and 1753 deletions
@@ -11,6 +11,22 @@ from typing import Any, Dict, Optional

 from agent.usage_pricing import normalize_usage

+_SENSITIVE_HOOK_KEYS = {"api_key", "authorization", "proxy_authorization", "cookie", "set_cookie"}
+
+
+def _model_dump(value: Any) -> Any:
+    """``value.model_dump(mode="json")`` with graceful degradation for older pydantic signatures.
+
+    warnings=False: pydantic UserWarnings on generic-union SDK models would leak to the terminal.
+    """
+    try:
+        return value.model_dump(mode="json", warnings=False)
+    except TypeError:
+        try:
+            return value.model_dump(mode="json")
+        except TypeError:
+            return value.model_dump()
+

 class ApiRequestHooksMixin:
    """Hook payload sanitising + ``api_request_error`` dispatch (see module docstring)."""
@@ -44,23 +60,11 @@ class ApiRequestHooksMixin:
        if not isinstance(key, str):
            return False
        lowered = key.lower().replace("-", "_")
-        exact = {
-            "api_key",
-            "authorization",
-            "proxy_authorization",
-            "cookie",
-            "set_cookie",
-        }
-        return lowered in exact or lowered.endswith("_api_key")
+        return lowered in _SENSITIVE_HOOK_KEYS or lowered.endswith("_api_key")

    @classmethod
    def _hook_jsonable(
-        cls,
-        value: Any,
-        *,
-        depth: int = 0,
-        max_depth: int = 8,
-        max_string: int = 8000,
+        cls, value: Any, *, depth: int = 0, max_depth: int = 8, max_string: int = 8000,
        max_sequence: int = 200,
    ) -> Any:
        if depth > max_depth:
@@ -73,6 +77,13 @@ class ApiRequestHooksMixin:
            return value
        if isinstance(value, (bytes, bytearray)):
            return f"<{len(value)} bytes>"
+
+        def recurse(item):
+            return cls._hook_jsonable(
+                item, depth=depth + 1, max_depth=max_depth, max_string=max_string,
+                max_sequence=max_sequence,
+            )
+
        if isinstance(value, dict):
            out: Dict[str, Any] = {}
            for idx, (key, item) in enumerate(value.items()):
@@ -80,111 +91,49 @@ class ApiRequestHooksMixin:
                    out["_truncated_items"] = len(value) - max_sequence
                    break
                str_key = str(key)
-                if cls._is_sensitive_hook_key(str_key):
-                    out[str_key] = "<redacted>"
-                else:
-                    out[str_key] = cls._hook_jsonable(
-                        item,
-                        depth=depth + 1,
-                        max_depth=max_depth,
-                        max_string=max_string,
-                        max_sequence=max_sequence,
-                    )
+                out[str_key] = "<redacted>" if cls._is_sensitive_hook_key(str_key) else recurse(item)
            return out
        if isinstance(value, (list, tuple, set)):
            seq = list(value)
-            out = [
-                cls._hook_jsonable(
-                    item,
-                    depth=depth + 1,
-                    max_depth=max_depth,
-                    max_string=max_string,
-                    max_sequence=max_sequence,
-                )
-                for item in seq[:max_sequence]
-            ]
+            out = [recurse(item) for item in seq[:max_sequence]]
            if len(seq) > max_sequence:
                out.append({"_truncated_items": len(seq) - max_sequence})
            return out
        try:
            if hasattr(value, "model_dump"):
-                try:
-                    # warnings=False: pydantic UserWarnings on generic-union SDK models would leak to the
-                    # terminal.
-                    dumped = value.model_dump(mode="json", warnings=False)
-                except TypeError:
-                    try:
-                        dumped = value.model_dump(mode="json")
-                    except TypeError:
-                        dumped = value.model_dump()
-                return cls._hook_jsonable(
-                    dumped,
-                    depth=depth + 1,
-                    max_depth=max_depth,
-                    max_string=max_string,
-                    max_sequence=max_sequence,
-                )
+                return recurse(_model_dump(value))
        except Exception:
            pass
        try:
            from dataclasses import asdict, is_dataclass
            if is_dataclass(value):
-                return cls._hook_jsonable(
-                    asdict(value),
-                    depth=depth + 1,
-                    max_depth=max_depth,
-                    max_string=max_string,
-                    max_sequence=max_sequence,
-                )
+                return recurse(asdict(value))
        except Exception:
            pass
        if isinstance(value, SimpleNamespace):
-            return cls._hook_jsonable(
-                vars(value),
-                depth=depth + 1,
-                max_depth=max_depth,
-                max_string=max_string,
-                max_sequence=max_sequence,
-            )
+            return recurse(vars(value))
        if hasattr(value, "__dict__"):
            try:
-                public_attrs = {
-                    k: v
-                    for k, v in vars(value).items()
-                    if not str(k).startswith("_")
-                }
-                return cls._hook_jsonable(
-                    public_attrs,
-                    depth=depth + 1,
-                    max_depth=max_depth,
-                    max_string=max_string,
-                    max_sequence=max_sequence,
-                )
+                return recurse({k: v for k, v in vars(value).items() if not str(k).startswith("_")})
            except Exception:
                pass
        return str(value)[:max_string]

    @classmethod
    def _sanitize_hook_payload(cls, value: Any) -> Any:
-        payload = cls._hook_jsonable(value)
+        """JSON-able payload under the size cap: full → reduced caps → truncated preview."""
        limit = cls._hook_payload_max_chars()
-        try:
-            encoded = json.dumps(payload, ensure_ascii=False, default=str)
-        except Exception:
-            return str(payload)[:limit]
-        if len(encoded) <= limit:
-            return payload
-        payload = cls._hook_jsonable(value, max_string=1000, max_sequence=50)
-        try:
-            encoded = json.dumps(payload, ensure_ascii=False, default=str)
-        except Exception:
-            return str(payload)[:limit]
-        if len(encoded) <= limit:
-            return payload
+        encoded = ""
+        for caps in ({}, {"max_string": 1000, "max_sequence": 50}):
+            payload = cls._hook_jsonable(value, **caps)
+            try:
+                encoded = json.dumps(payload, ensure_ascii=False, default=str)
+            except Exception:
+                return str(payload)[:limit]
+            if len(encoded) <= limit:
+                return payload
        return {
-            "_truncated": True,
-            "original_type": type(value).__name__,
-            "preview": encoded[:limit],
+            "_truncated": True, "original_type": type(value).__name__, "preview": encoded[:limit]
        }

    def _api_request_payload_for_hook(self, api_kwargs: Optional[Dict[str, Any]]) -> Dict[str, Any]:
@@ -193,19 +142,10 @@ class ApiRequestHooksMixin:
            for key, value in (api_kwargs or {}).items()
            if key not in {"timeout", "http_client"}
        }
-        return self._sanitize_hook_payload(
-            {
-                "method": "POST",
-                "body": body,
-            }
-        )
+        return self._sanitize_hook_payload({"method": "POST", "body": body})

    def _api_response_payload_for_hook(
-        self,
-        response: Any,
-        assistant_message: Any,
-        *,
-        finish_reason: Optional[str],
+        self, response: Any, assistant_message: Any, *, finish_reason: Optional[str]
    ) -> Dict[str, Any]:
        # Raw provider SDK tool_call objects are handed to the sanitizer on purpose; `_hook_jsonable` must
        # keep normalising them (model_dump / __dict__ / dataclass) or subscribers get str() blobs.
@@ -224,20 +164,10 @@ class ApiRequestHooksMixin:
        )

    def _invoke_api_request_error_hook(
-        self,
-        *,
-        task_id: str,
-        turn_id: str,
-        api_request_id: str,
-        api_call_count: int,
-        api_start_time: float,
-        api_kwargs: Optional[Dict[str, Any]],
-        error_type: str,
-        error_message: str,
-        status_code: Optional[int] = None,
-        retry_count: Optional[int] = None,
-        max_retries: Optional[int] = None,
-        retryable: Optional[bool] = None,
+        self, *, task_id: str, turn_id: str, api_request_id: str, api_call_count: int,
+        api_start_time: float, api_kwargs: Optional[Dict[str, Any]], error_type: str,
+        error_message: str, status_code: Optional[int] = None, retry_count: Optional[int] = None,
+        max_retries: Optional[int] = None, retryable: Optional[bool] = None,
        reason: Optional[str] = None,
    ) -> None:
        # Lazy module import (not from-import) so tests can replace lifecycle dispatch at this call site.
@@ -267,10 +197,7 @@ class ApiRequestHooksMixin:
                max_retries=max_retries,
                retryable=retryable,
                reason=reason,
-                error={
-                    "type": error_type,
-                    "message": error_message,
-                },
+                error={"type": error_type, "message": error_message},
                request=self._api_request_payload_for_hook(api_kwargs),
            )
        except Exception: