diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 314262d42..1162fd2c2 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2400,6 +2400,21 @@ class AnthropicHandlerMixin: if _req_ctx.tools is not body.get("tools"): tools = _req_ctx.tools body["tools"] = tools + # Turn hooks (e.g. lossless-guard) fold messages AFTER the pipeline's + # token accounting, and may mutate them IN PLACE (identity unchanged), + # so their savings were invisible to the PERF line / `headroom perf` + # (record_compression /stats already counts them). Re-count regardless + # of replace-vs-in-place so original->optimized reflects the fold too. + # tokenizer is initialized → count_messages is a pure CPU call here. + # Only ever lowers optimized_tokens. + try: + _hooked_tokens = tokenizer.count_messages(optimized_messages) + if _hooked_tokens < optimized_tokens: + optimized_tokens = _hooked_tokens + tokens_saved = max(0, original_tokens - optimized_tokens) + transforms_applied.append("turn_hook") + except Exception: + logger.debug("turn-hook token re-count skipped", exc_info=True) # Output shaping (opt-in via HEADROOM_OUTPUT_SHAPER): verbosity # steering appended to the system-prompt tail + effort routing on diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 3830d5a35..93c5147e5 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -2359,16 +2359,48 @@ class OpenAIHandlerMixin: if registered_turn_hooks(): if working is payload: working = copy.deepcopy(payload) + # Match the ctx's `input or messages or []` truthy fallback so we write + # back / re-count the SAME list the hook was handed. + _msg_key = ( + "input" + if working.get("input") + else ("messages" if working.get("messages") else None) + ) + _msgs_before = (working.get(_msg_key) if _msg_key else None) or [] + # Snapshot the pre-hook count BEFORE run_request_hooks — a hook may fold + # in place, which would corrupt a post-hook "before" measurement. + _mt_before = 0 + if _msg_key: + try: + _mt_before = self.openai_provider.get_token_counter(model).count_text( + _json_debug_dumps(_msgs_before) + ) + except Exception: + _mt_before = 0 _req_ctx = TurnContext( provider="openai", model=str(model), - messages=working.get("input") or working.get("messages") or [], + messages=_msgs_before, tools=working.get("tools"), config=getattr(self, "config", None), ) run_request_hooks(_req_ctx) if _req_ctx.tools is not working.get("tools"): working["tools"] = _req_ctx.tools + # A hook may also fold the messages (replace or in-place). Write back a + # replaced list — previously dropped on this path — then re-count so the + # message-fold saving is both applied AND recorded in tokens_saved. + if _msg_key and _req_ctx.messages is not _msgs_before: + working[_msg_key] = _req_ctx.messages + if _msg_key and _mt_before: + try: + _mt_after = self.openai_provider.get_token_counter(model).count_text( + _json_debug_dumps(working.get(_msg_key) or []) + ) + if _mt_after < _mt_before: + tokens_saved += _mt_before - _mt_after + except Exception: + pass modified = True transforms.append("openai:responses:turn_hook") @@ -3468,6 +3500,18 @@ class OpenAIHandlerMixin: if _th_ctx.tools is not _th_tools_before: tools = _th_ctx.tools body["tools"] = tools + # Message folds land AFTER the accounting above, and a hook may mutate + # messages IN PLACE (identity unchanged), so re-count regardless or + # `headroom perf` sees 0 for them (record_compression /stats already + # does). tokenizer is initialized → pure CPU. Only lowers the count. + try: + _th_msg_after = tokenizer.count_messages(body["messages"]) + if _th_msg_after < optimized_tokens: + optimized_tokens = _th_msg_after + tokens_saved = max(0, original_tokens - optimized_tokens) + transforms_applied.append("turn_hook") + except Exception: + logger.debug("turn-hook token re-count skipped", exc_info=True) _th_tok_after = ( tokenizer.count_text(json.dumps(_th_ctx.tools, default=str)) if _th_ctx.tools else 0 ) diff --git a/tests/test_openai_chat_turn_hooks.py b/tests/test_openai_chat_turn_hooks.py index 15f6b27d0..0d93a0873 100644 --- a/tests/test_openai_chat_turn_hooks.py +++ b/tests/test_openai_chat_turn_hooks.py @@ -266,6 +266,48 @@ def test_in_place_shrink_hook_is_counted(): assert ts["tokens"] > 0 and ts["requests"] >= 1, ts +def test_in_place_message_fold_is_counted(): + """A hook may fold MESSAGE content in place (e.g. lossless-guard collapsing a + tool_result), which lands after the pipeline's token accounting. The saving + must be re-counted regardless of object identity, else `headroom perf` shows + 0 for it — regression for identity-gated message-token accounting.""" + + class MessageFold: + name = "msgfold" + + def on_request(self, ctx): + # Fold a big message's content IN PLACE (mutate the dict, no reassign + # of ctx.messages), so the list object identity is unchanged. + for m in ctx.messages: + if isinstance(m.get("content"), str) and len(m["content"]) > 200: + m["content"] = "FOLDED" + + register_turn_hook(MessageFold()) + + async def fake_retry(method, url, headers, body, *args, **kwargs): + return httpx.Response( + 200, json=_final_response(), headers={"content-type": "application/json"} + ) + + app = create_app(_config()) + with TestClient(app) as client: + client.app.state.proxy._retry_request = fake_retry + resp = _post( + client, + { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "pad " * 500}], # big, foldable + "stream": False, + }, + ) + assert resp.status_code == 200, resp.text + # the message fold is attributed even though ctx.messages identity is unchanged + assert "turn_hook" in resp.headers.get("x-headroom-transforms", "") + # ...and the request's recorded token saving reflects it (was 0 pre-fix) + logs = client.app.state.proxy.logger.get_recent(5) + assert any(int(lg.get("tokens_saved", 0) or 0) > 0 for lg in logs), logs + + def test_direct_path_noop_when_no_hook_registered(): # No hook registered -> byte-identical passthrough, single upstream call. calls = {"n": 0} diff --git a/tests/test_openai_responses_context_compaction.py b/tests/test_openai_responses_context_compaction.py index 7802caef2..e29b81dac 100644 --- a/tests/test_openai_responses_context_compaction.py +++ b/tests/test_openai_responses_context_compaction.py @@ -460,3 +460,48 @@ def test_responses_memory_tools_allow_default_and_stored_requests() -> None: assert _responses_request_allows_memory_tool_continuation(default_store_payload) is True assert "store" not in default_store_payload + + +def test_responses_turn_hook_message_fold_is_applied_and_counted() -> None: + """On the Responses path a turn hook may fold the `input` items (in place), + not just tools. The fold must be written back to the outbound payload AND its + token saving added to tokens_saved — before, this path only wrote tools back, + so a message fold was silently dropped and uncounted.""" + from headroom.proxy.turn_hooks import clear_turn_hooks, register_turn_hook + + class FoldInput: + name = "fold_input" + + def on_request(self, ctx: Any) -> None: + # Fold a big function_call_output IN PLACE (mutate the dict; identity + # of ctx.messages is unchanged) — the case an identity gate would miss. + for item in ctx.messages: + if isinstance(item, dict) and isinstance(item.get("output"), str): + item["output"] = "folded" + + router = ContentRouter(ContentRouterConfig()) + handler = _HandlerHarness(router) + payload: dict[str, Any] = { + "type": "response.create", + "model": "gpt-5.5", + "input": [ + { + "type": "function_call_output", + "call_id": "c1", + "output": " ".join(["compressible"] * 300), + } + ], + } + + clear_turn_hooks() + register_turn_hook(FoldInput()) + try: + working, _modified, tokens_saved, *_ = handler._compress_openai_responses_payload( + payload, model="gpt-5.5", request_id="hr_test" + ) + finally: + clear_turn_hooks() + + assert working["input"][0]["output"] == "folded" # fold applied to the outbound payload + assert tokens_saved > 0 # ...and the message-fold saving is counted + assert payload["input"][0]["output"] != "folded" # original untouched (deep-copied)