mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix: correct preserved-entry index mapping in Gemini content round-trip (#836)
## Summary - `_gemini_contents_to_messages` excludes entries with no text parts (pure `functionCall` / `functionResponse` / image-only) from `messages[]`, but their original `contents[]` indices are stored in `preserved_indices` - After compression, `optimized_contents` has a shorter, different index space — the old restoration loop used raw `orig_idx` to overwrite `optimized_contents[orig_idx]`, silently corrupting text entries at colliding positions and silently dropping preserved entries when `orig_idx >= len(optimized_contents)` - Affects all three Gemini handlers (`generateContent`, `cloudCodeAssist`, `countTokens`) — any agentic session with function calls where compression fires **Concrete failure case:** ``` contents = [user:text, model:functionCall, user:functionResponse, model:text] messages = [user:text, model:text] # only 2 — FC/FR have no text optimized_contents = [user:text, model:text] # positions 0 and 1 old loop: orig_idx=1 → optimized_contents[1] = functionCall ← overwrites model text! orig_idx=2 → 2 < 2 is False → functionResponse silently dropped ``` ## Fix Added `_rebuild_gemini_contents()` helper that walks `original_contents` in order, placing preserved entries at their exact relative positions and consuming optimized text entries sequentially via an iterator. Replaced all three broken loops. ## Test plan - [ ] `TestRebuildGeminiContents::test_text_only_unchanged` — text-only round-trip is identity - [ ] `TestRebuildGeminiContents::test_function_call_sequence_preserved` — functionCall + functionResponse survive at correct positions - [ ] `TestRebuildGeminiContents::test_function_call_at_start` — preserved entry at idx=0 no longer overwrites optimized_contents[0] - [ ] `TestRebuildGeminiContents::test_hybrid_entry_uses_original` — entry with both text and functionCall retains functionCall All 58 tests in `test_google_multimodal.py` pass. Rust CI + mypy + ruff clean. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Ashish
GitHub
parent
693d9d20e2
commit
0ffe2b6ea4
2 changed files with 123 additions and 15 deletions
|
|
@ -73,6 +73,40 @@ class GeminiHandlerMixin:
|
|||
return True
|
||||
return False
|
||||
|
||||
def _rebuild_gemini_contents(
|
||||
self,
|
||||
original_contents: list[dict],
|
||||
preserved_indices: set[int],
|
||||
preserved_contents: dict[int, dict],
|
||||
optimized_contents: list[dict],
|
||||
) -> list[dict]:
|
||||
"""Interleave preserved (non-text) entries back into optimized_contents at their
|
||||
original positions.
|
||||
|
||||
preserved_indices uses original contents[] indices, but optimized_contents uses
|
||||
a different (shorter) index space because entries with no text parts were excluded
|
||||
from the messages[] sent for compression. Using orig_idx directly to overwrite
|
||||
optimized_contents[orig_idx] corrupts or silently drops entries.
|
||||
|
||||
This method walks original_contents in order, placing each position with either
|
||||
the preserved original (for non-text entries) or the next optimized text entry.
|
||||
"""
|
||||
opt_iter = iter(optimized_contents)
|
||||
result: list[dict] = []
|
||||
for idx, content in enumerate(original_contents):
|
||||
had_text = any("text" in p for p in content.get("parts", []))
|
||||
if idx in preserved_indices:
|
||||
result.append(preserved_contents[idx])
|
||||
if had_text:
|
||||
# Entry also produced a message; consume but discard the optimized version
|
||||
next(opt_iter, None)
|
||||
else:
|
||||
opt_entry = next(opt_iter, None)
|
||||
if opt_entry is not None:
|
||||
result.append(opt_entry)
|
||||
# else: dropped by compression — omit
|
||||
return result
|
||||
|
||||
def _gemini_contents_to_messages(
|
||||
self, contents: list[dict], system_instruction: dict | None = None
|
||||
) -> tuple[list[dict], set[int]]:
|
||||
|
|
@ -501,12 +535,9 @@ class GeminiHandlerMixin:
|
|||
optimized_contents, optimized_system = self._messages_to_gemini_contents(
|
||||
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
|
||||
|
||||
optimized_contents = self._rebuild_gemini_contents(
|
||||
contents, preserved_indices, preserved_contents, optimized_contents
|
||||
)
|
||||
body["contents"] = optimized_contents
|
||||
if optimized_system:
|
||||
body["systemInstruction"] = optimized_system
|
||||
|
|
@ -784,9 +815,12 @@ class GeminiHandlerMixin:
|
|||
optimized_contents, optimized_system = self._messages_to_gemini_contents(
|
||||
optimized_messages
|
||||
)
|
||||
for orig_idx, original_content in preserved_contents.items():
|
||||
if orig_idx < len(optimized_contents):
|
||||
optimized_contents[orig_idx] = original_content
|
||||
optimized_contents = self._rebuild_gemini_contents(
|
||||
contents if isinstance(contents, list) else [],
|
||||
preserved_indices,
|
||||
preserved_contents,
|
||||
optimized_contents,
|
||||
)
|
||||
request_payload["contents"] = optimized_contents
|
||||
if not is_antigravity:
|
||||
if optimized_system:
|
||||
|
|
@ -1028,12 +1062,9 @@ class GeminiHandlerMixin:
|
|||
optimized_contents, optimized_system = self._messages_to_gemini_contents(
|
||||
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
|
||||
|
||||
optimized_contents = self._rebuild_gemini_contents(
|
||||
contents, preserved_indices, preserved_contents, optimized_contents
|
||||
)
|
||||
body["contents"] = optimized_contents
|
||||
if optimized_system:
|
||||
body["systemInstruction"] = optimized_system
|
||||
|
|
|
|||
|
|
@ -627,3 +627,80 @@ class TestParametrizedPreservation:
|
|||
"""Parametrized test for preserved indices."""
|
||||
_, preserved_indices = proxy._gemini_contents_to_messages(contents)
|
||||
assert preserved_indices == expected_indices
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Tests for _rebuild_gemini_contents
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestRebuildGeminiContents:
|
||||
"""_rebuild_gemini_contents must re-insert preserved entries at their original positions."""
|
||||
|
||||
def _round_trip(self, proxy, contents):
|
||||
"""Simulate the full compression round-trip for a given contents list.
|
||||
|
||||
Mimics what the handler does: convert → strip system msg → convert back → rebuild.
|
||||
"""
|
||||
messages, preserved_indices = proxy._gemini_contents_to_messages(contents)
|
||||
preserved_contents = {idx: contents[idx] for idx in preserved_indices}
|
||||
optimized_contents, _ = proxy._messages_to_gemini_contents(messages)
|
||||
return proxy._rebuild_gemini_contents(
|
||||
contents, preserved_indices, preserved_contents, optimized_contents
|
||||
)
|
||||
|
||||
def test_text_only_unchanged(self, proxy):
|
||||
"""Text-only round-trip should produce identical contents."""
|
||||
contents = [TEXT_ONLY_CONTENT, MODEL_TEXT_CONTENT]
|
||||
result = self._round_trip(proxy, contents)
|
||||
assert len(result) == 2
|
||||
assert result[0]["parts"][0]["text"] == "Hello, world!"
|
||||
assert result[1]["parts"][0]["text"] == "Hello! How can I help you today?"
|
||||
|
||||
def test_function_call_sequence_preserved(self, proxy):
|
||||
"""functionCall and functionResponse entries must survive and appear at correct positions."""
|
||||
contents = [
|
||||
TEXT_ONLY_CONTENT, # idx 0: text
|
||||
FUNCTION_CALL_CONTENT, # idx 1: functionCall only — no text → preserved
|
||||
FUNCTION_RESPONSE_CONTENT, # idx 2: functionResponse only — no text → preserved
|
||||
MODEL_TEXT_CONTENT, # idx 3: text
|
||||
]
|
||||
result = self._round_trip(proxy, contents)
|
||||
|
||||
assert len(result) == 4, f"Expected 4 entries, got {len(result)}: {result}"
|
||||
# Position 0: original text
|
||||
assert result[0]["parts"][0].get("text") == "Hello, world!"
|
||||
# Position 1: functionCall preserved exactly
|
||||
assert "functionCall" in result[1]["parts"][0], "functionCall missing at position 1"
|
||||
assert result[1]["parts"][0]["functionCall"]["name"] == "get_weather"
|
||||
# Position 2: functionResponse preserved exactly
|
||||
assert "functionResponse" in result[2]["parts"][0], "functionResponse missing at position 2"
|
||||
# Position 3: text preserved
|
||||
assert result[3]["parts"][0].get("text") == "Hello! How can I help you today?"
|
||||
|
||||
def test_function_call_at_start(self, proxy):
|
||||
"""Preserved entry at idx=0 must not overwrite idx=0 of optimized_contents."""
|
||||
contents = [
|
||||
FUNCTION_CALL_CONTENT, # idx 0: no text → preserved
|
||||
TEXT_ONLY_CONTENT, # idx 1: text
|
||||
]
|
||||
result = self._round_trip(proxy, contents)
|
||||
|
||||
assert len(result) == 2
|
||||
assert "functionCall" in result[0]["parts"][0]
|
||||
assert result[1]["parts"][0].get("text") == "Hello, world!"
|
||||
|
||||
def test_hybrid_entry_uses_original(self, proxy):
|
||||
"""Entry with both text and functionCall keeps the original (with functionCall intact)."""
|
||||
contents = [
|
||||
TEXT_ONLY_CONTENT,
|
||||
FUNCTION_CALL_WITH_TEXT_CONTENT, # idx 1: has both text and functionCall → preserved
|
||||
MODEL_TEXT_CONTENT,
|
||||
]
|
||||
result = self._round_trip(proxy, contents)
|
||||
|
||||
assert len(result) == 3
|
||||
# Hybrid entry must come back as the original (functionCall retained)
|
||||
hybrid = result[1]
|
||||
part_keys = {k for p in hybrid["parts"] for k in p}
|
||||
assert "functionCall" in part_keys, "functionCall lost from hybrid entry"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue