Extract tool schema savings policy

headroom/integrations/litellm_callback.py

@@ -104,7 +104,6 @@ async def async_pre_call_hook(
             data, call_type = cache, data
         if data is None:
             return None
-
         if call_type not in ("completion", "acompletion"):
             return data
 

headroom/proxy/server.py

@@ -163,6 +163,7 @@
 from headroom.proxy.savings_tracker import LITELLM_AVAILABLE
 from headroom.proxy.semantic_cache import SemanticCache  # noqa: F401
 from headroom.proxy.ssl_context import build_httpx_verify
+from headroom.proxy.tool_schema_savings_policy import tool_schema_saved_from_tags
 from headroom.proxy.warmup import WarmupRegistry
 from headroom.proxy.ws_session_registry import WebSocketSessionRegistry
 from headroom.subscription.base import get_quota_registry, reset_quota_registry
@@ -2020,22 +2021,7 @@ def _is_known_websocket_callback_failure(context: dict[str, Any]) -> bool:
     )
 
 
-def _tool_schema_saved_from_tags(tags: object) -> int:
-    """Tool-definition tokens Headroom kept out of context for one request by
-    deferring heavy tool schemas: the native tool-search injection
-    (``tool_search_deferred_tokens``) plus any registered turn-hook tools
-    rewrite (``turn_hook_tools_saved_tokens``). Both tags are set only on the
-    path where Headroom performed the deferral, so a client that already had
-    tool search enabled (e.g. Claude Code / Codex) contributes zero here."""
-    if not isinstance(tags, dict):
-        return 0
-    total = 0
-    for key in ("tool_search_deferred_tokens", "turn_hook_tools_saved_tokens"):
-        try:
-            total += int(tags.get(key, 0) or 0)
-        except (TypeError, ValueError):
-            continue
-    return total
+_tool_schema_saved_from_tags = tool_schema_saved_from_tags
 
 
 def create_app(config: ProxyConfig | None = None) -> FastAPI:

headroom/proxy/tool_schema_savings_policy.py

@@ -0,0 +1,26 @@
+"""Tool-schema savings attribution policy for proxy stats."""
+
+from __future__ import annotations
+
+TOOL_SCHEMA_SAVINGS_TAGS: tuple[str, ...] = (
+    "tool_search_deferred_tokens",
+    "turn_hook_tools_saved_tokens",
+)
+
+
+def tool_schema_saved_from_tags(tags: object) -> int:
+    """Return tool-definition tokens Headroom kept out of context for one request.
+
+    The summed tags are set only on paths where Headroom performed the deferral,
+    so clients that already had tool search enabled contribute zero here.
+    """
+    if not isinstance(tags, dict):
+        return 0
+
+    total = 0
+    for key in TOOL_SCHEMA_SAVINGS_TAGS:
+        try:
+            total += int(tags.get(key, 0) or 0)
+        except (TypeError, ValueError):
+            continue
+    return total

tests/test_tool_schema_savings_policy.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from headroom.proxy.tool_schema_savings_policy import (
+    TOOL_SCHEMA_SAVINGS_TAGS,
+    tool_schema_saved_from_tags,
+)
+
+
+def test_tool_schema_saved_from_tags_sums_headroom_deferral_tags() -> None:
+    assert (
+        tool_schema_saved_from_tags(
+            {
+                "tool_search_deferred_tokens": "120",
+                "turn_hook_tools_saved_tokens": 30,
+                "unrelated": 999,
+            }
+        )
+        == 150
+    )
+
+
+def test_tool_schema_saved_from_tags_ignores_invalid_values() -> None:
+    assert (
+        tool_schema_saved_from_tags(
+            {
+                "tool_search_deferred_tokens": "not-an-int",
+                "turn_hook_tools_saved_tokens": None,
+            }
+        )
+        == 0
+    )
+
+
+def test_tool_schema_saved_from_tags_rejects_non_mapping_tags() -> None:
+    assert tool_schema_saved_from_tags(None) == 0
+    assert tool_schema_saved_from_tags([("tool_search_deferred_tokens", 10)]) == 0
+
+
+def test_tool_schema_savings_tags_are_stable() -> None:
+    assert TOOL_SCHEMA_SAVINGS_TAGS == (
+        "tool_search_deferred_tokens",
+        "turn_hook_tools_saved_tokens",
+    )