From c371d5ad602f5ab93645b2db4673ae2c5e9f0575 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Fri, 24 Jul 2026 09:38:52 -0700 Subject: [PATCH] fix(proxy/perf): count turn-hook message folds in token accounting (#2520) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Turn hooks (the `headroom.proxy.turn_hooks` seam used by proxy extensions, e.g. the lossless-guard plugin) fold tool_result / message content in `on_request`, which runs **after** the pipeline has already computed `optimized_tokens`. The saving was recorded to `/stats` via `record_compression`, but was invisible to the `PERF` log line and `headroom perf` (both read the pipeline's `original → optimized` delta). Net effect: a plugin that folded 463 tokens still logged `tok_saved=0`. This makes the per-turn token accounting count the hook's fold too, across all three handler paths. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - **Anthropic Messages handler** (`/v1/messages`): re-count messages right after `run_request_hooks`, regardless of whether the hook replaced the list or mutated it in place. Attribute the fold as a `turn_hook` transform. Only ever lowers `optimized_tokens`. - **OpenAI Chat handler** (`handle_openai_chat`, `/v1/chat/completions`): same re-count. The existing code re-counted hook-modified *tools* but not the *message* fold — this closes that gap and adds the `turn_hook` transform tag. - **OpenAI Responses handler** (`_compress_openai_responses_payload`, `/v1/responses`): the seam previously only wrote hook-modified *tools* back — a folded/replaced `input` list was silently dropped and uncounted. Now snapshot the message-items token count **before** the hook (an in-place fold would corrupt a post-hook baseline), write back a replaced list, and add the fold delta to `tokens_saved` (the same channel the tool-schema savings already ride to `/stats` and `headroom perf`). - Key detail: the identity check `ctx.messages is not ` is insufficient — the lossless-guard plugin mutates messages **in place**, so an identity-gated re-count misses it. The re-count runs unconditionally whenever a hook ran. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ ruff check headroom/proxy/handlers/anthropic.py headroom/proxy/handlers/openai.py \ tests/test_openai_chat_turn_hooks.py tests/test_openai_responses_context_compaction.py All checks passed! $ mypy headroom/proxy/handlers/anthropic.py headroom/proxy/handlers/openai.py Success: no issues found in 2 source files $ pytest tests/test_turn_hooks.py tests/test_openai_chat_turn_hooks.py \ tests/test_openai_responses_context_compaction.py -q tests/test_turn_hooks.py ......... [ 34%] tests/test_openai_chat_turn_hooks.py ..... [ 53%] tests/test_openai_responses_context_compaction.py ............ [100%] 26 passed in 14.05s ``` New regression tests (each fails on the pre-fix code): - `test_in_place_message_fold_is_counted` (chat path) — hook folds message content in place; asserts `turn_hook` in `x-headroom-transforms` and a recorded `tokens_saved > 0`. - `test_responses_turn_hook_message_fold_is_applied_and_counted` (Responses path) — hook folds a `function_call_output` in place; asserts the outbound payload reflects the fold **and** `tokens_saved > 0`. ## Real Behavior Proof - **Environment:** local proxy (`headroom proxy --port 8793 --proxy-extension lossless_guard`), `HEADROOM_LICENSE_DEV=1`, `HEADROOM_PROTECT_TOOL_RESULTS=Bash` (so the fold is purely the plugin's turn hook), model `claude-haiku-4-5`. Request carries a `gh --json` object (folded to TOON) and a `docker pull` log. - **Exact steps:** send the request → read the `PERF` line in `~/.headroom/logs/proxy.log` and `GET /stats`. - **Observed result:** - Before this change: `PERF ... tok_before=607 tok_after=607 tok_saved=0 ... transforms=none` while `/stats` reported `{"lossless_guard": 145}` — i.e. the saving existed but perf showed nothing. - After this change: `PERF ... tok_before=607 tok_after=484 tok_saved=123 ... transforms=turn_hook`, `/stats` still `{"lossless_guard": 145}`. (`123` is the honest whole-request `count_messages` delta; `145` is the per-content-string delta `record_compression` measures — different scopes, both real and positive.) - **Not tested:** the OpenAI Chat and Responses paths were verified by unit test, not a live client run — my live setup routes Claude Code through the Anthropic handler only. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation — N/A (internal accounting; no public API/doc surface) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` ## Additional Notes Behavior is unchanged when no turn hook is registered (`registered_turn_hooks() == []` → the re-count block is skipped), so pure-OSS installs are byte-identical and unaffected. OSS's own pipeline compression was already counted correctly (it runs before the hook); this only surfaces the extension/turn-hook layer. --- headroom/proxy/handlers/anthropic.py | 15 ++++++ headroom/proxy/handlers/openai.py | 46 ++++++++++++++++++- tests/test_openai_chat_turn_hooks.py | 42 +++++++++++++++++ ...est_openai_responses_context_compaction.py | 45 ++++++++++++++++++ 4 files changed, 147 insertions(+), 1 deletion(-) 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)