headroom

akadmin/headroom

Fork 0

mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00

Commit graph

Author	SHA1	Message	Date
Abhay Singh	4056117d90	fix(proxy/gemini): preserve non-text content across the compression round-trip (#2079 ) ## 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 <chopratejas@gmail.com>	2026-07-13 09:46:35 -04:00
Ashish	0ffe2b6ea4	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>	2026-06-10 21:10:56 -05:00
chopratejas	dd4c65e565	Fix Google Gemini conversion to preserve non-text content (Gap #4 ) Previously, _gemini_contents_to_messages() only extracted text parts, silently dropping non-text content like images, file references, and function calls. This caused data loss when processing multimodal content. Changes: - Add _has_non_text_parts() helper to detect inlineData, fileData, functionCall, and functionResponse parts - Modify _gemini_contents_to_messages() to return tuple of (messages, preserved_indices) tracking which content entries have non-text parts - Update all call sites to preserve non-text content entries: - handle_google_batch_create - handle_gemini_generate_content - handle_gemini_count_tokens - _store_google_batch_context - Skip compression entirely when all content has non-text parts - Restore preserved entries after compression/optimization Add comprehensive test coverage (54 tests) verifying: - Detection of all non-text part types - Correct index tracking in preserved_indices - Role mapping and message conversion - Realistic conversation flows with images, function calls, PDFs	2026-01-24 12:20:13 -08:00

Author

SHA1

Message

Date

Abhay Singh

4056117d90

fix(proxy/gemini): preserve non-text content across the compression round-trip (#2079 )

## 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 <chopratejas@gmail.com>

2026-07-13 09:46:35 -04:00

Ashish

0ffe2b6ea4

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>

2026-06-10 21:10:56 -05:00

chopratejas

dd4c65e565

Fix Google Gemini conversion to preserve non-text content (Gap #4 )

Previously, _gemini_contents_to_messages() only extracted text parts,
silently dropping non-text content like images, file references, and
function calls. This caused data loss when processing multimodal content.

Changes:
- Add _has_non_text_parts() helper to detect inlineData, fileData,
  functionCall, and functionResponse parts
- Modify _gemini_contents_to_messages() to return tuple of (messages,
  preserved_indices) tracking which content entries have non-text parts
- Update all call sites to preserve non-text content entries:
  - handle_google_batch_create
  - handle_gemini_generate_content
  - handle_gemini_count_tokens
  - _store_google_batch_context
- Skip compression entirely when all content has non-text parts
- Restore preserved entries after compression/optimization

Add comprehensive test coverage (54 tests) verifying:
- Detection of all non-text part types
- Correct index tracking in preserved_indices
- Role mapping and message conversion
- Realistic conversation flows with images, function calls, PDFs

2026-01-24 12:20:13 -08:00

3 commits