From 07cf5476072a45bac7dd94386de126234a8049e7 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Wed, 22 Jul 2026 18:37:15 +0530 Subject: [PATCH] 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 --- headroom/proxy/handlers/gemini.py | 31 ++++++++++++---- tests/test_gemini_function_response_waste.py | 39 ++++++++++++++++++++ 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index efb9cfe08..dba8bbbb1 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -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: diff --git a/tests/test_gemini_function_response_waste.py b/tests/test_gemini_function_response_waste.py index 638db4f8a..20161e66c 100644 --- a/tests/test_gemini_function_response_waste.py +++ b/tests/test_gemini_function_response_waste.py @@ -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 = [