mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
2 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Abhay Singh
|
07cf547607
|
fix(proxy/gemini): tolerate malformed parts on the compression path (#2486)
## Description
Three helpers on the Gemini compression path read a content entry's
`parts` and iterate it without type guards:
```python
# _has_non_text_parts
parts = content.get("parts", [])
for part in parts: ...
# _rebuild_gemini_contents
had_text = any("text" in p for p in content.get("parts", []))
# _gemini_contents_to_messages
parts = content.get("parts", [])
text_parts = [p.get("text", "") for p in parts if "text" in p]
```
`parts` is request-controlled and `.get("parts", [])` only falls back
when the key is absent, so:
- a present-but-null `parts` returns `None`, and `for part in None` /
`any(... for p in None)` raises `TypeError`;
- a list carrying a bare string (a client that treats `parts` as a
string array) makes `p.get("text", "")` raise `AttributeError`, while
`"text" in p` silently does substring matching first;
- a null element in the list crashes the same way.
Any of these 500s the request on the compression path, on data that
parsed as valid JSON.
## Fix
Route all three helpers through a shared `_dict_parts(content)` that
returns the dict entries of `parts`, coercing a non-dict content or a
non-list `parts` to an empty list and dropping non-dict elements.
`_gemini_contents_to_messages` also reads `role` defensively for a
non-dict content entry. Conversion now degrades gracefully (the
malformed part contributes nothing) instead of raising. Well-formed
requests are unchanged.
## 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
- `headroom/proxy/handlers/gemini.py`: add `_dict_parts`; use it in
`_has_non_text_parts`, `_rebuild_gemini_contents`, and
`_gemini_contents_to_messages`; read `role` defensively for a non-dict
content entry.
- `tests/test_gemini_function_response_waste.py`: regressions for null
`parts`, bare-string part elements, a null part element,
`_has_non_text_parts` on malformed parts, and a non-dict content entry.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ python -m pytest tests/test_gemini_function_response_waste.py -q
16 passed
# with the fix reverted, the new malformed-parts tests fail with
# TypeError: 'NoneType' object is not iterable (and AttributeError on string parts)
$ uvx ruff@0.15.17 check headroom/proxy/handlers/gemini.py tests/test_gemini_function_response_waste.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/handlers/gemini.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, project venv (`uv sync --extra
proxy`), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`, pytest in the venv.
- Exact command / steps: built a real `HeadroomProxy` and called
`_gemini_contents_to_messages` / `_has_non_text_parts` with contents
carrying `parts: null`, `parts: ["bare string", {text}]`, `parts: [null,
{text}]`, and a non-dict content entry; then reverted `gemini.py` and
re-ran.
- Observed result: with the fix each malformed shape converts without
raising and the valid text part is still emitted (`[{"role": "user",
"content": "kept"}]`); with the fix reverted the null-`parts` and
null-element cases raise `TypeError: 'NoneType' object is not iterable`
and the string-element case raises `AttributeError: 'str' object has no
attribute 'get'`. Ran against the actual module via
`tests/test_gemini_function_response_waste.py`.
- Not tested: a live Gemini request with malformed `parts` routed
through the full proxy compression pipeline end to end.
## 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
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable
|
||
|
Focused Instability
|
9b0c840dd7
|
fix(gemini): surface functionResponse payloads to waste-signal detection (#897)
## Problem Fixes #819. Gemini `functionResponse` parts are preserved verbatim on the wire (by design — they are never compressed), but their payloads never reached `parse_messages`: `_gemini_contents_to_messages` only extracts `text` parts. Tool output — where most waste lives — contributed nothing to waste detection on either Gemini path, so `json_bloat`, `repetition`, and the new `reread` signal (#853/#854) were all blind to it. ## Fix (telemetry-only) 1. **`_gemini_contents_to_messages(..., include_function_responses=True)`** — new keyword-only flag. When set, each `functionResponse` payload is additionally emitted as a `role="tool"` message (dict payloads JSON-serialized, strings passed through, missing/`None` responses skipped). `preserved_indices` semantics are unchanged: the entries are still restored verbatim on the wire. 2. **`TransformPipeline.apply(..., waste_messages=...)`** — new optional kwarg (popped before transforms, like `record_metrics`). When provided, the waste-signal parse runs over this richer list instead of the transform input. Transforms, token accounting, and savings deltas are untouched — this is why the richer list is not simply fed to the pipeline: compressed copies of preserved entries are discarded on rebuild, which would corrupt savings reporting. 3. Both Gemini `generateContent` paths (native + Cloud Code Assist) build the enriched list and pass it through. The existing `role="tool"` parsing from #815 handles the rest: tool_result blocks, waste flags, and reread grouping all apply. ## Tests `tests/test_gemini_function_response_waste.py` — 11 new tests: - conversion: default unchanged (regression), dict/string payloads, missing response skipped, text-before-tool ordering, preserved_indices unchanged, circular-reference fallback - parsing: functionResponse payload produces tool_result blocks + `json_bloat`; identical payloads far apart count as `reread` - pipeline: `waste_messages` overrides the waste source, does not affect transform output/token counts, falls back to transform input when absent Full local sweep of touched suites: gemini multimodal, parser, safety rails, canonical pipeline — green. The 13 failures in `test_proxy_gemini_*_integration.py` are credential-dependent and identical on clean `main`. ## Live proof Mock Gemini upstream on a real port, proxy with `optimize=True`; conversation with a large functionResponse payload served twice (5 messages apart) plus compressible model text: ``` waste_signals: { "json_bloat": 35003, "reread": 11673, ... } PROOF OK: waste visible, wire verbatim ``` Upstream received both `functionResponse` entries byte-identical to the client request. ## Known limitations / follow-ups - The Cloud Code Assist path passes `waste_messages` but does not yet consume `result.waste_signals` into a recorded outcome (pre-existing gap; the native path records it). - Requests where **all** content entries are preserved (pure functionResponse/media conversations) early-exit before the pipeline and still produce no waste signals. - Codex/Responses-API counterpart is #820 (separate PR). Co-authored-by: integration-check <integration@local> |