mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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
This commit is contained in:
parent
f0975b8de0
commit
07cf547607
2 changed files with 63 additions and 7 deletions
|
|
@ -56,6 +56,24 @@ class GeminiHandlerMixin:
|
|||
return ANTIGRAVITY_DAILY_API_URL
|
||||
return getattr(self, "CLOUDCODE_API_URL", DEFAULT_CLOUDCODE_API_URL).rstrip("/")
|
||||
|
||||
@staticmethod
|
||||
def _dict_parts(content: Any) -> list[dict]:
|
||||
"""Return the dict entries of a Gemini content's ``parts``.
|
||||
|
||||
``parts`` is request-controlled. ``.get("parts", [])`` only falls back
|
||||
when the key is absent, so a present-but-null ``parts`` returns ``None``
|
||||
(crashing ``for part in parts``), and a list carrying a bare string —
|
||||
which a client that treats ``parts`` as a string array can send —
|
||||
crashes ``part.get(...)`` / ``key in part`` semantics downstream.
|
||||
Coerce to a clean list of dict parts so every caller can iterate safely.
|
||||
"""
|
||||
if not isinstance(content, dict):
|
||||
return []
|
||||
parts = content.get("parts")
|
||||
if not isinstance(parts, list):
|
||||
return []
|
||||
return [part for part in parts if isinstance(part, dict)]
|
||||
|
||||
def _has_non_text_parts(self, content: dict) -> bool:
|
||||
"""Check if a Gemini content entry has non-text parts.
|
||||
|
||||
|
|
@ -73,8 +91,7 @@ class GeminiHandlerMixin:
|
|||
Returns:
|
||||
True if any part contains non-text data.
|
||||
"""
|
||||
parts = content.get("parts", [])
|
||||
for part in parts:
|
||||
for part in self._dict_parts(content):
|
||||
if any(
|
||||
key in part
|
||||
for key in (
|
||||
|
|
@ -110,7 +127,7 @@ class GeminiHandlerMixin:
|
|||
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", []))
|
||||
had_text = any("text" in p for p in self._dict_parts(content))
|
||||
if idx in preserved_indices:
|
||||
result.append(preserved_contents[idx])
|
||||
if had_text:
|
||||
|
|
@ -155,8 +172,8 @@ class GeminiHandlerMixin:
|
|||
|
||||
# Add system instruction as system message
|
||||
if system_instruction:
|
||||
parts = system_instruction.get("parts", [])
|
||||
text_parts = [p.get("text", "") for p in parts if "text" in p]
|
||||
sys_parts = self._dict_parts(system_instruction)
|
||||
text_parts = [p.get("text", "") for p in sys_parts if "text" in p]
|
||||
if text_parts:
|
||||
messages.append({"role": "system", "content": "\n".join(text_parts)})
|
||||
|
||||
|
|
@ -166,12 +183,12 @@ class GeminiHandlerMixin:
|
|||
if self._has_non_text_parts(content):
|
||||
preserved_indices.add(idx)
|
||||
|
||||
role = content.get("role", "user")
|
||||
role = content.get("role", "user") if isinstance(content, dict) else "user"
|
||||
# Map Gemini roles to OpenAI roles
|
||||
if role == "model":
|
||||
role = "assistant"
|
||||
|
||||
parts = content.get("parts", [])
|
||||
parts = self._dict_parts(content)
|
||||
text_parts = [p.get("text", "") for p in parts if "text" in p]
|
||||
|
||||
if text_parts:
|
||||
|
|
|
|||
|
|
@ -124,6 +124,45 @@ class TestFunctionResponseConversion:
|
|||
assert "loop" in text
|
||||
|
||||
|
||||
class TestMalformedPartsToleration:
|
||||
"""A request-controlled `parts` that is null or carries non-dict elements
|
||||
must not crash the compression-path conversion helpers."""
|
||||
|
||||
def test_null_parts_does_not_crash(self, proxy):
|
||||
contents = [
|
||||
{"role": "user", "parts": None},
|
||||
{"role": "user", "parts": [{"text": "real"}]},
|
||||
]
|
||||
messages, preserved = proxy._gemini_contents_to_messages(contents)
|
||||
assert messages == [{"role": "user", "content": "real"}]
|
||||
assert preserved == set()
|
||||
|
||||
def test_string_part_elements_do_not_crash(self, proxy):
|
||||
# A client that treats `parts` as a string array sends bare strings;
|
||||
# they carry no `text` key, so they contribute nothing but must not
|
||||
# crash `.get`.
|
||||
contents = [{"role": "user", "parts": ["bare string", {"text": "kept"}]}]
|
||||
messages, _ = proxy._gemini_contents_to_messages(contents)
|
||||
assert messages == [{"role": "user", "content": "kept"}]
|
||||
|
||||
def test_null_part_element_is_skipped(self, proxy):
|
||||
contents = [{"role": "user", "parts": [None, {"text": "kept"}]}]
|
||||
messages, _ = proxy._gemini_contents_to_messages(contents)
|
||||
assert messages == [{"role": "user", "content": "kept"}]
|
||||
|
||||
def test_has_non_text_parts_tolerates_null_parts(self, proxy):
|
||||
assert proxy._has_non_text_parts({"role": "user", "parts": None}) is False
|
||||
assert proxy._has_non_text_parts({"role": "user", "parts": ["str"]}) is False
|
||||
assert proxy._has_non_text_parts({"parts": [{"inlineData": {"data": "x"}}]}) is True
|
||||
|
||||
def test_non_dict_content_entry_is_tolerated(self, proxy):
|
||||
# A non-dict entry in contents[] is treated as an empty user turn rather
|
||||
# than crashing content.get / the parts iteration.
|
||||
contents = ["not a dict", {"role": "user", "parts": [{"text": "kept"}]}]
|
||||
messages, _ = proxy._gemini_contents_to_messages(contents)
|
||||
assert messages == [{"role": "user", "content": "kept"}]
|
||||
|
||||
|
||||
class TestFunctionResponseWasteParsing:
|
||||
def test_function_response_payload_reaches_waste_signals(self, proxy, tokenizer):
|
||||
contents = [
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue