fix(proxy/perf): count turn-hook message folds in token accounting (#2520)

## 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 <orig>` 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

@@ -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

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
             )

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}

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)