diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 37bd72fa3..51dfb527d 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -3511,6 +3511,19 @@ class OpenAIHandlerMixin: _compression_failed = False original_messages = messages # Preserve for 400-retry fallback + # Cross-turn dedup rewrites repeated tool-output spans to bare + # `[↑NL same as msg M]` in-context pointers. Those are recoverable only + # where the model can resolve the reference; on the streaming chat path + # the CCR retrieval tool cannot be injected (this path cannot intercept + # tool calls) and OpenAI-compatible clients never show the model + # numbered messages, so a folded pointer reads as deleted content and + # models retry-loop on the "missing" output. Gate the fold on the same + # recoverability predicate that gates CCR tool injection: the buffered + # (non-streaming) chat path keeps dedup, the streaming path skips it. + _dedup_pointers_recoverable = _should_inject_openai_chat_ccr_tool( + ccr_inject_tool=self.config.ccr_inject_tool, + stream=stream, + ) _decision = CompressionDecision.decide( headers=request.headers, config=self.config, @@ -3565,6 +3578,7 @@ class OpenAIHandlerMixin: ), biases=_hook_biases, compression_policy=compression_policy, + cross_turn_dedup_recoverable=_dedup_pointers_recoverable, # Thread the savings-profile knobs (e.g. # HEADROOM_SAVINGS_PROFILE=agent-90) onto the live # chat-completions path, matching handlers/ @@ -3605,6 +3619,7 @@ class OpenAIHandlerMixin: frozen_message_count=apply_frozen_count, biases=_hook_biases, compression_policy=compression_policy, + cross_turn_dedup_recoverable=_dedup_pointers_recoverable, # Same savings-profile threading as the token-mode # branch above — the non-token chat path must honor # the configured profile too (#1534). diff --git a/headroom/transforms/cold_prefix.py b/headroom/transforms/cold_prefix.py index 3777e1c1d..b58ac4d14 100644 --- a/headroom/transforms/cold_prefix.py +++ b/headroom/transforms/cold_prefix.py @@ -179,7 +179,11 @@ def has_plaintext_reasoning(messages: list[dict[str, Any]]) -> bool: def cold_recompact_messages( - messages: list[dict[str, Any]], *, tokenizer: Any, context: str = "" + messages: list[dict[str, Any]], + *, + tokenizer: Any, + context: str = "", + cross_turn_dedup_recoverable: bool = True, ) -> tuple[list[dict[str, Any]], list[str]]: """Lossless whole-prefix recompaction for a confirmed-cold turn. @@ -191,6 +195,15 @@ def cold_recompact_messages( preserve nothing. Lossless + prefix-monotonic ⇒ deterministic per content ⇒ the recompacted prefix re-caches and stays byte-stable on later warm turns. + ``cross_turn_dedup_recoverable`` is forwarded to the router's dedup gate: + pass False on paths where a bare ``[↑NL same as msg M]`` pointer cannot be + resolved — no CCR retrieval tool can be injected and the client never shows + the model numbered messages (OpenAI chat-completions streaming, e.g. + ``wrap copilot``). The fold is then skipped and the bytes stay verbatim; + the lossless folds still run. The Anthropic cache-mode caller keeps the + default True (the in-context reference resolves there and the retrieval + tool is injectable). + Returns (new_messages, transforms_applied). Fail-open: returns the input unchanged on any error (never breaks the request). """ @@ -200,8 +213,17 @@ def cold_recompact_messages( ContentRouterConfig, ) + # enable_cross_turn_dedup stays on: whether the fold may EMIT pointers + # is decided per-path by the router's recoverability gate below, not + # hardcoded here. router = ContentRouter(ContentRouterConfig(lossless=True, enable_cross_turn_dedup=True)) - res = router.apply(list(messages), tokenizer, frozen_message_count=0, context=context) + res = router.apply( + list(messages), + tokenizer, + frozen_message_count=0, + context=context, + cross_turn_dedup_recoverable=cross_turn_dedup_recoverable, + ) return res.messages, list(res.transforms_applied) except Exception as e: # never break the request log.warning("cold-prefix recompaction failed (%s); leaving prefix unchanged", e) diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index a77446149..6f1fba323 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -1528,6 +1528,11 @@ class ContentRouterConfig: # Runs in both modes: lossless references verbatim/folded content; CCR mode # references the earlier block's kompressed-but-CCR-recoverable form # (deterministic content-hash → stable → still cache-safe, no added loss). + # Per request the fold is skipped when the caller reports the serving path + # cannot resolve the in-context `[↑NL same as msg M]` pointer + # (`apply(cross_turn_dedup_recoverable=False)`, e.g. OpenAI chat-completions + # streaming, where no CCR retrieval tool can be injected and clients never + # show the model numbered messages). enable_cross_turn_dedup: bool = False # Lossless-then-lossy. In lossy mode (not `lossless`), after a byte/data # lossless fold (search/log/text) run the aggressive lossy compressor @@ -4757,6 +4762,20 @@ class ContentRouter(Transform): # pass a policy — ``_record_to_toin`` treats that as "no gate" # to preserve pre-F2.2 behaviour for non-proxy callers. self._runtime_compression_policy = kwargs.get("compression_policy") + # Cross-turn dedup recoverability gate. The fold rewrites a repeated + # span to a bare in-context pointer (``[↑NL same as msg M]``) that names + # Headroom's internal message index. That reference is only resolvable + # where the model can locate the original: on the OpenAI + # chat-completions streaming path (e.g. ``wrap copilot``) no CCR + # retrieval tool can be injected (the path cannot intercept tool calls) + # and the client never shows the model numbered messages, so the + # pointer reads as deleted content and the model retry-loops on the + # "missing" output. Same recoverability posture as the lossy + # ``lossy_unrecoverable_skipped`` guard: when the caller reports the + # path cannot resolve in-context pointers, skip the fold and keep the + # bytes verbatim. Default True: every path that does not opt out keeps + # today's behavior. + dedup_pointers_recoverable = bool(kwargs.get("cross_turn_dedup_recoverable", True)) tokens_before = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) context = kwargs.get("context", "") @@ -5538,7 +5557,7 @@ class ContentRouter(Transform): # later duplicate would carry the same (recoverable) form anyway; dedup # just points to the earlier copy instead of repeating it. Frozen + # cache_control blocks are reference targets only (never rewritten). - if self._cross_turn_dedup_enabled: + if self._cross_turn_dedup_enabled and dedup_pointers_recoverable: transformed_messages = self._cross_turn_dedup_messages( transformed_messages, frozen_message_count, transforms_applied, route_counts ) diff --git a/tests/test_cold_prefix.py b/tests/test_cold_prefix.py new file mode 100644 index 000000000..15d0618e4 --- /dev/null +++ b/tests/test_cold_prefix.py @@ -0,0 +1,69 @@ +"""Cold-prefix recompaction: the cross-turn dedup fold must be path-aware. + +``cold_recompact_messages`` builds its own lossless ContentRouter with +``enable_cross_turn_dedup=True``. The fold rewrites repeated tool-output spans +to bare ``[↑NL same as msg M]`` in-context pointers — unresolvable on paths +where no CCR retrieval tool can be injected and the client never shows the +model numbered messages (OpenAI chat-completions streaming, wrap copilot). +The recompaction therefore takes ``cross_turn_dedup_recoverable`` and forwards +it to the router gate: unrecoverable paths keep the bytes verbatim, the +Anthropic cache-mode caller (default True) keeps folding. +""" + +from headroom.transforms.cold_prefix import cold_recompact_messages + + +def _mk_tok(): + from headroom.providers import OpenAIProvider + from headroom.tokenizer import Tokenizer + + return Tokenizer(OpenAIProvider().get_token_counter("gpt-4o"), "gpt-4o") + + +def _toolmsg(text, tid): + return { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": tid, "content": text}], + } + + +def _conversation(): + span = "\n".join(f" result_{i} = compute_overdraft(business_id={i})" for i in range(12)) + return [ + {"role": "user", "content": "fix the overdraft bug"}, + {"role": "assistant", "content": "cat merge.py"}, + _toolmsg(f"$ cat merge.py\n{span}\n# end", "t1"), + {"role": "assistant", "content": "sed -n range"}, + _toolmsg(f"$ sed -n 1,20p merge.py\n{span}\n# more", "t2"), + ] + + +def test_cold_recompact_folds_by_default(): + # Anthropic cache-mode path (the only caller today): unchanged — the + # repeated span still folds to an in-context pointer. + msgs = _conversation() + out, transforms = cold_recompact_messages(msgs, tokenizer=_mk_tok()) + later = out[-1]["content"][0]["content"] + assert "[↑" in later + assert any("cross_turn_dedup" in t for t in transforms) + + +def test_cold_recompact_unrecoverable_path_keeps_verbatim_bytes(): + # Unresolvable-pointer path (OpenAI chat streaming shape): the fold is + # skipped, the repeated span stays byte-verbatim, and no pointer is + # emitted — while the recompaction itself still runs (message count and + # order unchanged). + msgs = _conversation() + out, transforms = cold_recompact_messages( + msgs, tokenizer=_mk_tok(), cross_turn_dedup_recoverable=False + ) + later = out[-1]["content"][0]["content"] + assert "[↑" not in later + assert ( + later + == "$ sed -n 1,20p merge.py\n" + + "\n".join(f" result_{i} = compute_overdraft(business_id={i})" for i in range(12)) + + "\n# more" + ) + assert not any("cross_turn_dedup" in t for t in transforms) + assert len(out) == len(msgs) diff --git a/tests/test_cross_turn_dedup.py b/tests/test_cross_turn_dedup.py index 53cd9b13a..391f87eb6 100644 --- a/tests/test_cross_turn_dedup.py +++ b/tests/test_cross_turn_dedup.py @@ -399,3 +399,76 @@ def test_dedup_folds_role_function_output(): ] out = _dedup_only(msgs) assert "[↑" in out[2]["content"] + + +# -------------------------------------------------------------------------- +# Recoverability gate (unresolvable-pointer paths). The fold rewrites a +# repeated span to a bare `[↑NL same as msg M]` pointer naming Headroom's +# internal message index. On the OpenAI chat-completions streaming path +# (wrap copilot) no CCR retrieval tool can be injected and the client never +# shows the model numbered messages, so the pointer is unresolvable: the +# model reads it as deleted content and retry-loops. `apply()` therefore +# accepts `cross_turn_dedup_recoverable=False` — the same recoverability +# posture as the lossy `lossy_unrecoverable_skipped` guard — and keeps the +# repeated bytes verbatim. Default True preserves every other path. +# -------------------------------------------------------------------------- +def _apply_with_recoverable(messages, recoverable): + import copy + + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + + r = ContentRouter(ContentRouterConfig(lossless=True, enable_cross_turn_dedup=True)) + return r.apply( + copy.deepcopy(messages), _mk_tok(), cross_turn_dedup_recoverable=recoverable + ).messages + + +def test_apply_unrecoverable_path_keeps_verbatim_bytes(): + # The OpenAI chat streaming shape (role:tool strings): with dedup ENABLED + # but the path flagged unrecoverable, the re-read must NOT fold — the + # request keeps the verbatim bytes, no bare pointer. + span = _readspan() + msgs = [ + {"role": "tool", "tool_call_id": "c1", "content": f"$ cat f.py\n{span}"}, + {"role": "assistant", "content": "again"}, + {"role": "tool", "tool_call_id": "c2", "content": f"$ cat f.py\n{span}"}, + ] + out = _apply_with_recoverable(msgs, recoverable=False) + assert out[2]["content"] == f"$ cat f.py\n{span}" # verbatim, no pointer + assert "[↑" not in out[2]["content"] + + +def test_apply_unrecoverable_gate_also_covers_tool_result_blocks(): + # Anthropic tool_result block shape, same gate: nothing folds when the + # caller reports the pointer is unresolvable on this path. + span = _readspan() + msgs = [_toolmsg(f"a\n{span}", "t1"), _toolmsg(f"b\n{span}", "t2")] + out = _apply_with_recoverable(msgs, recoverable=False) + joined = "".join(b["content"] for m in out for b in m["content"] if isinstance(b, dict)) + assert "[↑" not in joined + assert out[-1]["content"][0]["content"] == f"b\n{span}" # verbatim bytes kept + + +def test_apply_recoverable_default_and_true_still_fold(): + # The recoverable paths (Anthropic, buffered/non-streaming chat — anywhere + # the reference resolves) keep folding: default kwarg-absent behavior is + # unchanged, and an explicit True folds too. + span = _readspan() + msgs = [ + {"role": "tool", "tool_call_id": "c1", "content": f"$ cat f.py\n{span}"}, + {"role": "assistant", "content": "again"}, + {"role": "tool", "tool_call_id": "c2", "content": f"$ cat f.py\n{span}"}, + ] + import copy + + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + + default_out = ( + ContentRouter(ContentRouterConfig(lossless=True, enable_cross_turn_dedup=True)) + .apply(copy.deepcopy(msgs), _mk_tok()) + .messages + ) + assert "[↑" in default_out[2]["content"] # no kwarg -> still folds + + true_out = _apply_with_recoverable(msgs, recoverable=True) + assert "[↑" in true_out[2]["content"] # explicit recoverable -> folds diff --git a/tests/test_openai_chat_dedup_recoverability.py b/tests/test_openai_chat_dedup_recoverability.py new file mode 100644 index 000000000..9b1f61ee0 --- /dev/null +++ b/tests/test_openai_chat_dedup_recoverability.py @@ -0,0 +1,158 @@ +"""OpenAI chat-completions: cross-turn dedup pointers are recoverability-gated. + +The fold rewrites a repeated tool-output span to a bare ``[↑NL same as msg M]`` +pointer naming Headroom's internal message index. On the STREAMING chat path +(``wrap copilot``) the CCR retrieval tool cannot be injected — the path cannot +intercept tool calls — and OpenAI-compatible clients never show the model +numbered messages, so the pointer is unresolvable: models read it as deleted +content and retry-loop. The chat handler therefore threads +``cross_turn_dedup_recoverable=_should_inject_openai_chat_ccr_tool(...)`` into +the router: streaming requests keep the repeated bytes verbatim, while the +buffered (non-streaming) path — where the retrieval tool IS injectable — keeps +folding. + +These tests drive the real ``/v1/chat/completions`` handler through a TestClient +with dedup force-enabled (``HEADROOM_DEDUPE=1``) and capture the exact upstream +request body, the same evidence the proxy logs showed when the bug bit. +""" + +from __future__ import annotations + +import pytest + +fastapi = pytest.importorskip("fastapi") +httpx = pytest.importorskip("httpx") + +from fastapi.responses import StreamingResponse # noqa: E402 +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + +_SPAN = "\n".join(f" result_{i} = compute_overdraft(business_id={i})" for i in range(12)) + + +def _messages() -> list[dict]: + """Two identical multi-line tool outputs — the re-read dedup folds.""" + return [ + {"role": "user", "content": "fix the overdraft bug"}, + {"role": "assistant", "content": "cat merge.py"}, + {"role": "tool", "tool_call_id": "call_1", "content": f"$ cat merge.py\n{_SPAN}\n# end"}, + {"role": "assistant", "content": "sed -n range"}, + {"role": "tool", "tool_call_id": "call_2", "content": f"$ cat merge.py\n{_SPAN}\n# end"}, + ] + + +def _config() -> ProxyConfig: + return ProxyConfig(optimize=True, cache_enabled=False, rate_limit_enabled=False) + + +def _post(client: TestClient, *, stream: bool): + return client.post( + "/v1/chat/completions", + json={"model": "gpt-4o", "messages": _messages(), "stream": stream}, + headers={"Authorization": "******"}, + ) + + +def _sent_text(body: dict) -> str: + """Concatenate the upstream message contents (parsed, so newlines are real).""" + return "\n".join(str(m.get("content", "")) for m in body["messages"]) + + +def test_streaming_chat_keeps_verbatim_bytes_no_dedup_pointer(monkeypatch): + """The bug: a streaming chat request with a repeated span got a bare + ``[↑NL same as msg M]`` pointer the model cannot resolve. Now the upstream + body must carry the repeated bytes verbatim.""" + monkeypatch.setenv("HEADROOM_DEDUPE", "1") # before create_app: router reads env at init + captured: list[dict] = [] + + async def fake_stream(url, headers, body, *args, **kwargs): + captured.append(body) + return StreamingResponse(iter([b"data: {}\n\n"]), media_type="text/event-stream") + + app = create_app(_config()) + with TestClient(app) as client: + client.app.state.proxy._stream_response = fake_stream + resp = _post(client, stream=True) + + assert resp.status_code == 200, resp.text + assert captured, "streaming upstream send was not captured" + sent = _sent_text(captured[0]) + assert "[↑" not in sent # no unresolvable pointer on the streaming path + assert sent.count(_SPAN) == 2 # both copies forwarded byte-verbatim + + +def test_lossless_buffered_chat_also_skips_the_fold(monkeypatch): + """Coupling lock: --lossless forces ccr_inject_tool=False (server.py), so + the recoverability predicate is False for buffered chat too and the fold + is skipped there as well (no retrieval tool exists to redeem anything in + no-CCR mode). Bytes stay verbatim; the conservative direction is intended.""" + monkeypatch.setenv("HEADROOM_DEDUPE", "1") + captured: list[dict] = [] + + async def fake_retry(method, url, headers, body, *args, **kwargs): + captured.append(body) + payload = { + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "done"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 5, "total_tokens": 105}, + } + return httpx.Response(200, json=payload, headers={"content-type": "application/json"}) + + config = ProxyConfig( + optimize=True, lossless=True, cache_enabled=False, rate_limit_enabled=False + ) + app = create_app(config) + with TestClient(app) as client: + client.app.state.proxy._retry_request = fake_retry + resp = _post(client, stream=False) + + assert resp.status_code == 200, resp.text + assert captured, "buffered upstream send was not captured" + sent = _sent_text(captured[0]) + assert "[↑" not in sent # no retrieval tool in lossless mode -> no bare pointer + assert sent.count(_SPAN) == 2 # both copies forwarded byte-verbatim + + +def test_buffered_chat_still_folds_repeated_tool_output(monkeypatch): + """The recoverable counterpart: non-streaming chat can inject the CCR + retrieval tool, so the in-context pointer stays resolvable and the + repeated span still folds (today's behavior, unchanged).""" + monkeypatch.setenv("HEADROOM_DEDUPE", "1") + captured: list[dict] = [] + + async def fake_retry(method, url, headers, body, *args, **kwargs): + captured.append(body) + payload = { + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "done"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 5, "total_tokens": 105}, + } + return httpx.Response(200, json=payload, 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, stream=False) + + assert resp.status_code == 200, resp.text + assert captured, "buffered upstream send was not captured" + sent = _sent_text(captured[0]) + assert "[↑" in sent # fold still fires where the pointer resolves + assert sent.count(_SPAN) == 1 # earliest copy stays as the in-context original