From 4056117d90468619dbc2684448ebd76c48c41921 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Mon, 13 Jul 2026 19:16:35 +0530 Subject: [PATCH] fix(proxy/gemini): preserve non-text content across the compression round-trip (#2079) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Two related content-loss bugs in the Gemini `contents[]` <-> `messages[]` compression round-trip. Both drop or misplace real user content that entries with **non-text** parts should carry through untouched. They share the same theme (non-text preservation), so they're bundled here as two commits. ### 1. Google batch handler restores preserved entries by the wrong index (`handlers/batch.py`) The `batchGenerateContent` handler restored preserved (non-text) entries with the raw-index loop that commit #836 (`_rebuild_gemini_contents`) replaced in the three non-batch Gemini handlers: ```python for orig_idx, original_content in preserved_contents.items(): if orig_idx < len(optimized_contents): optimized_contents[orig_idx] = original_content ``` `preserved_indices` are indices into the **original** `contents[]`, but `optimized_contents` is a **shorter** list (text-less entries produce no message). Indexing `optimized_contents` by `orig_idx` overwrites the wrong entry and drops any preserved entry whose original index is past the optimized length. For: ```python [user text, model functionCall, user functionResponse, model text] ``` the batch was forwarded to Google as **two** entries: the model's answer overwritten by the functionCall, and the functionResponse dropped. Unlike `gemini.py` there is no `if optimized_messages != messages` gate, so it runs on every mixed batch item. **Fix:** use the shared `_rebuild_gemini_contents` interleaving helper. ### 2. Code-execution parts not detected as non-text (`handlers/gemini.py`) `_has_non_text_parts` only recognized `inlineData`/`fileData`/`functionCall`/`functionResponse`. Gemini's code-execution feature emits `executableCode` and `codeExecutionResult` parts, echoed back in `contents[]` on later turns. Because they weren't detected: - a mixed `text`+`executableCode` entry lost its code payload (only the text survived the round-trip); - a text-less `executableCode`+`codeExecutionResult` entry was treated as a phantom in `_rebuild_gemini_contents` — it consumed the next optimized message, dropping the whole code turn and shifting a following user turn into the model's role slot (corrupting role alternation). **Fix:** add both keys to the non-text detection so those entries are preserved verbatim. Closes: no issue filed — both found while auditing the Gemini contents<->messages round-trip. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/proxy/handlers/batch.py`: use `_rebuild_gemini_contents` instead of the raw-index restore loop. - `headroom/proxy/handlers/gemini.py`: recognize `executableCode` / `codeExecutionResult` in `_has_non_text_parts`. - `tests/test_proxy_handlers_batch.py`: add `test_handle_google_batch_create_preserves_functioncall_response_order`, driving the handler with the **real** Gemini converters (the existing batch tests stub them, which hid the bug); mix `GeminiHandlerMixin` into the shared `DummyBatchHandler` so `_rebuild_gemini_contents` is available. - `tests/test_google_multimodal.py`: extend the parametrized `test_each_non_text_key_detected` to the two new keys, and add `test_code_execution_entry_survives`. ## Testing - [x] New regression tests added (`tests/test_proxy_handlers_batch.py`, `tests/test_google_multimodal.py`) - [x] Linting/formatting clean — run with the CI-pinned `ruff==0.15.17` - [ ] Full `pytest` deferred to CI (local-OOM reason below). ```text $ uvx ruff@0.15.17 check headroom/proxy/handlers/batch.py headroom/proxy/handlers/gemini.py \ tests/test_proxy_handlers_batch.py tests/test_google_multimodal.py All checks passed! ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.10, headroom from this branch. Importing `headroom` pulls in the torch/transformers stack and a full `pytest` gets OOM-killed on this box, so I verified the interleaving/detection with dependency-free scripts (replicating the Gemini converters, the old loop, and `_rebuild_gemini_contents`) and left the full pytest to CI. - Exact command / steps: ran two standalone scripts. Script 1 rebuilds a Gemini batch request with `preserved_indices` holding a `functionCall`/`functionResponse` pair and compares the old raw-index loop against `_rebuild_gemini_contents`. Script 2 feeds a `codeExecutionResult` entry through `_has_non_text_parts` and the preserve path with and without the two new allowlist keys. Also ran `uvx ruff@0.15.17 check` on the changed files and tests. - Observed result: the old batch loop drops the `functionResponse` and overwrites the answer (4 parts collapse to 2); `_rebuild_gemini_contents` keeps all 4. Without the new keys the code-execution entry is dropped/shifted (2 parts, code absent); with them it survives intact (3 parts, code present). Lint clean. See the two blocks below. Batch fix (bug #1): ```text preserved_indices: [1, 2] OLD result parts: ['text', 'functionCall'] len 2 NEW result parts: ['text', 'functionCall', 'functionResponse', 'text'] len 4 GEMINI BATCH REBUILD FIX VERIFIED (old drops response + overwrites answer; new keeps all 4) ``` Code-execution fix (bug #2): ```text (b) OLD len=2 NEW len=3 (a) OLD has code=False NEW has code=True GEMINI CODE-EXECUTION PRESERVE FIX VERIFIED (old drops/shifts; new keeps intact) ``` - Not tested: a live Google/Gemini round-trip (handlers stubbed, as the existing tests do). Full local `pytest` deferred to CI (OOM, per above). ## 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 - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes — ran lint + standalone logic checks; full pytest deferred to CI (local OOM, disclosed above) - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes - Two small behavioral changes (one loop -> shared helper, two keys added to an allowlist) plus regression tests; no new dependencies. Both complete/extend the non-text preservation the non-batch handlers already do (the #836 line). - @JerrettDavis tagging you since you reviewed the recent Gemini fixes. Both of these drop content (functionResponse/images on batch; code-execution on the normal round-trip), so they seemed worth surfacing together. Thanks. --------- Co-authored-by: Tejas Chopra --- CHANGELOG.md | 2 + headroom/proxy/handlers/batch.py | 15 +++-- headroom/proxy/handlers/gemini.py | 11 +++- tests/test_google_multimodal.py | 35 ++++++++++- tests/test_proxy_handlers_batch.py | 97 +++++++++++++++++++++++++++++- 5 files changed, 153 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de9c10ce2..b418d508e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -102,6 +102,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Bug Fixes +* **proxy/batch:** stop corrupting Google `batchGenerateContent` requests whose contents interleave text turns with text-less entries (functionCall/functionResponse/images). The batch handler restored preserved (non-text) entries with the raw-index loop that #836 replaced everywhere else — indexing the shorter `optimized_contents` (text-less entries produce no message) by the original `contents[]` index, which overwrites the wrong entry and drops any preserved entry whose original index is past the optimized length. A request like `[user text, model functionCall, user functionResponse, model text]` was forwarded to Google as two entries: the model's answer overwritten by the functionCall and the functionResponse dropped. The batch handler now uses the shared `_rebuild_gemini_contents` interleaving helper, so all entries survive in order. +* **proxy/gemini:** preserve Gemini code-execution parts (`executableCode` / `codeExecutionResult`) across the compression round-trip. `_has_non_text_parts` only recognized `inlineData`/`fileData`/`functionCall`/`functionResponse`, so a content entry carrying code-execution parts was not marked as preserved. A mixed `text`+`executableCode` entry lost its code payload (only the text survived), and a text-less code-execution entry was treated as a phantom that dropped the entire turn and shifted a neighboring message into the wrong role slot. Both keys are now recognized so those entries are preserved verbatim. * **cache/ccr:** don't evict a live entry when a duplicate hash is re-stored at capacity. `CompressionStore.store` ran `_evict_if_needed()` before checking whether the key already existed, so re-storing an already-present hash while the store was full evicted the oldest *distinct* entry to "make room" and then merely overwrote the existing key in place — no room was ever needed. The store dropped below `max_entries` and a live, never-retrieved entry was destroyed, so its `<>` marker (still in the conversation) resolved to a 404. The CCR mirror bridge re-stores the same `explicit_hash` every turn a marker is re-encountered, so this fired routinely. Eviction now runs only for a genuinely new key. * **tokenizers:** recurse into a native `tool_result` whose content is a list of blocks instead of JSON-serializing it. `_count_content_parts` counted a `tool_result` with list content via `_count_serialized` (json.dumps + sample), so a base64 image nested in a tool result (computer-use / MCP screenshot tools) was priced as text — a ~50-200x overcount (a ~200KB screenshot read as ~70K tokens instead of ~1600). It now recurses into the nested blocks, matching the sibling Strands `toolResult` branch, so the image is priced structurally. The overcount made a single screenshot appear to blow past the model's context window and triggered unnecessary/over-aggressive compression. * **tokenizers:** price dense scripts (CJK/Kana/Hangul) in the fixed-ratio estimator path. `EstimatingTokenCounter.count_text` applied the CJK correction only on the auto-detect path; its fixed-ratio early return divided by the Latin ratio with no adjustment. The registry builds every provider-calibrated counter with a fixed ratio (Anthropic 3.5, Google 4.0, Cohere 4.0, Moonshot 3.1), and the Anthropic/Gemini proxy handlers count via `get_tokenizer(model).count_messages`, so a CJK-heavy context read as ~40-55% of its true token size — it could fall under the size/backpressure gates and skip compression, and every `x-headroom-tokens-before` metric for CJK traffic was materially wrong. The fixed-ratio path now applies the same dense-script split. diff --git a/headroom/proxy/handlers/batch.py b/headroom/proxy/handlers/batch.py index 07f74defa..928467be2 100644 --- a/headroom/proxy/handlers/batch.py +++ b/headroom/proxy/handlers/batch.py @@ -225,10 +225,17 @@ class BatchHandlerMixin: optimized_messages ) - # Restore preserved content entries that had non-text parts - for orig_idx, original_content in preserved_contents.items(): - if orig_idx < len(optimized_contents): - optimized_contents[orig_idx] = original_content + # Restore preserved (non-text) entries at their ORIGINAL positions. + # preserved_indices are indices into the original contents[], but + # optimized_contents lives in a shorter index space (text-less + # entries produced no message), so indexing it by orig_idx + # overwrites the wrong entry and drops any preserved entry whose + # original index is >= len(optimized_contents). Use the shared + # interleaving helper the non-batch Gemini handlers already use + # (#836). + optimized_contents = self._rebuild_gemini_contents( + contents, preserved_indices, preserved_contents, optimized_contents + ) # Create compressed batch request compressed_req_content = {**req_content, "contents": optimized_contents} diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index 378c269a6..f9583b9a2 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -58,6 +58,8 @@ class GeminiHandlerMixin: - fileData: File references (URI + MIME type) - functionCall: Function calls from model - functionResponse: Responses to function calls + - executableCode / codeExecutionResult: Gemini code-execution parts, + echoed back in contents[] on later turns Args: content: A single Gemini content entry with 'parts' list. @@ -69,7 +71,14 @@ class GeminiHandlerMixin: for part in parts: if any( key in part - for key in ("inlineData", "fileData", "functionCall", "functionResponse") + for key in ( + "inlineData", + "fileData", + "functionCall", + "functionResponse", + "executableCode", + "codeExecutionResult", + ) ): return True return False diff --git a/tests/test_google_multimodal.py b/tests/test_google_multimodal.py index 03e583a8e..21627d1a1 100644 --- a/tests/test_google_multimodal.py +++ b/tests/test_google_multimodal.py @@ -173,7 +173,16 @@ class TestHasNonTextParts: @pytest.mark.parametrize( "non_text_key", - ["inlineData", "fileData", "functionCall", "functionResponse"], + [ + "inlineData", + "fileData", + "functionCall", + "functionResponse", + # Gemini code-execution parts, echoed back in contents[] on later + # turns; previously not detected, so they were dropped on round-trip. + "executableCode", + "codeExecutionResult", + ], ) def test_each_non_text_key_detected(self, proxy, non_text_key): """Each non-text part type is correctly detected.""" @@ -657,6 +666,30 @@ class TestRebuildGeminiContents: assert result[0]["parts"][0]["text"] == "Hello, world!" assert result[1]["parts"][0]["text"] == "Hello! How can I help you today?" + def test_code_execution_entry_survives(self, proxy): + """A text-less code-execution entry (executableCode + codeExecutionResult) + between two text turns must survive the round-trip at its position, and + not shift a neighboring turn. Before the fix it was not detected as + non-text, so it was dropped and the following user turn was misplaced.""" + code_entry = { + "role": "model", + "parts": [ + {"executableCode": {"language": "PYTHON", "code": "x = 1"}}, + {"codeExecutionResult": {"outcome": "OUTCOME_OK", "output": "1"}}, + ], + } + contents = [ + {"role": "user", "parts": [{"text": "Question 1"}]}, + code_entry, + {"role": "user", "parts": [{"text": "Question 2"}]}, + ] + + result = self._round_trip(proxy, contents) + + assert len(result) == 3 + assert result[1] == code_entry # preserved verbatim, in place + assert result[2]["parts"][0]["text"] == "Question 2" + def test_function_call_sequence_preserved(self, proxy): """functionCall and functionResponse entries must survive and appear at correct positions.""" contents = [ diff --git a/tests/test_proxy_handlers_batch.py b/tests/test_proxy_handlers_batch.py index 43812848d..30011017e 100644 --- a/tests/test_proxy_handlers_batch.py +++ b/tests/test_proxy_handlers_batch.py @@ -7,6 +7,7 @@ from types import SimpleNamespace import pytest from headroom.proxy.handlers import batch as batch_module +from headroom.proxy.handlers.gemini import GeminiHandlerMixin class FakeResponse: @@ -72,7 +73,10 @@ class FakeMetrics: self.failed_calls.append(kwargs) -class DummyBatchHandler(batch_module.BatchHandlerMixin): +class DummyBatchHandler(batch_module.BatchHandlerMixin, GeminiHandlerMixin): + # GeminiHandlerMixin supplies the real _rebuild_gemini_contents (and the + # other content helpers); the two converter methods below intentionally + # override the mixin's for the stub-based tests. OPENAI_API_URL = "https://openai.example" GEMINI_API_URL = "https://gemini.example" @@ -939,6 +943,97 @@ async def test_handle_google_batch_create_covers_passthrough_revert_and_store_fa assert optimized["systemInstruction"] == {"parts": [{"text": "sys"}]} +@pytest.mark.asyncio +async def test_handle_google_batch_create_preserves_functioncall_response_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A batch request that interleaves text turns with text-less + functionCall/functionResponse entries must reach Google with all entries + intact and in order. The old raw-index restore loop overwrote the model's + answer with the functionCall and dropped the functionResponse.""" + + class RealConvHandler(batch_module.BatchHandlerMixin, GeminiHandlerMixin): + # Real Gemini converters + _rebuild_gemini_contents (no stubs), so the + # actual index interleaving runs. + GEMINI_API_URL = "https://gemini.example" + + def __init__(self) -> None: + self.http_client = FakeHttpClient() + self.metrics = FakeMetrics() + self.config = SimpleNamespace( + optimize=True, ccr_inject_tool=False, ccr_inject_system_instructions=False + ) + self.openai_provider = SimpleNamespace(get_context_limit=lambda m: 8192) + # No-op pipeline: return the messages unchanged, no token inflation. + self.openai_pipeline = SimpleNamespace( + apply=lambda **kw: SimpleNamespace( + messages=kw["messages"], timing={}, tokens_before=100, tokens_after=100 + ) + ) + self.captured_body: dict | None = None + + async def _next_request_id(self) -> str: + return "req-1" + + async def _record_request_outcome(self, outcome) -> None: # noqa: ANN001 + pass + + def _extract_tags(self, headers: dict) -> dict[str, str]: + return {} + + async def _run_compression_in_executor(self, fn, *, timeout): # noqa: ANN001, ANN201 + return fn() + + async def _store_google_batch_context(self, *a, **k) -> None: # noqa: ANN002, ANN003 + pass + + async def _retry_request(self, method, url, headers, body, **kwargs): # noqa: ANN001, ANN201 + # Capture the (in-place mutated) forwarded batch body for assertions. + self.captured_body = body + return FakeResponse(status_code=200, content=b"{}", json_data={"name": "batches/1"}) + + handler = RealConvHandler() + + contents = [ + {"role": "user", "parts": [{"text": "What's the weather in Paris?"}]}, + { + "role": "model", + "parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}}], + }, + { + "role": "user", + "parts": [{"functionResponse": {"name": "get_weather", "response": {"temp_c": 18}}}], + }, + {"role": "model", "parts": [{"text": "It's 18C and cloudy in Paris."}]}, + ] + batch_body = { + "batch": { + "input_config": { + "requests": {"requests": [{"request": {"contents": contents}, "metadata": {}}]} + } + } + } + + async def payload(request): # noqa: ANN001, ANN201 + return batch_body + + monkeypatch.setattr("headroom.proxy.helpers._read_request_json", payload) + + resp = await handler.handle_google_batch_create(FakeRequest("{}"), "gemini-pro") + assert resp.status_code == 200 + + out = handler.captured_body["batch"]["input_config"]["requests"]["requests"][0]["request"][ + "contents" + ] + # All four entries survive in order. The old loop produced only two, dropping + # the functionResponse and overwriting the model answer with the functionCall. + assert len(out) == 4 + assert "text" in out[0]["parts"][0] + assert out[1]["parts"][0].get("functionCall", {}).get("name") == "get_weather" + assert out[2]["parts"][0].get("functionResponse", {}).get("name") == "get_weather" + assert "Paris" in out[3]["parts"][0]["text"] + + @pytest.mark.asyncio async def test_google_batch_passthrough_without_body_and_query_variants() -> None: handler = DummyBatchHandler()