From 89319fbcaddb4be2ea11e87858ed3bd0fcf9dca5 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Sun, 19 Jul 2026 05:17:59 +0530 Subject: [PATCH] fix(ccr): guard empty/malformed OpenAI choices in _extract_assistant_message (#2389) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `CCRResponseHandler._extract_assistant_message` extracts the assistant message from an upstream response while building the CCR retrieval-continuation history. The OpenAI branch is not defensive about an empty or malformed `choices` array: ```python elif provider == "openai": message = response.get("choices", [{}])[0].get("message", {}) ``` `response.get("choices", [{}])` only falls back to `[{}]` when the key is **absent**. When `choices` is present but empty (`[]`) or carries a null first element (`[null]`), this raises on the success path: - `choices: []` → `[][0]` → `IndexError` - `choices: [null]` → `None.get(...)` → `AttributeError` OpenAI-compatible gateways can return those shapes on content-filtered or usage-only responses. The sibling **Google** branch a few lines below already guards this (`candidates = response.get("candidates", []); if candidates: ... else: parts = []`), and so does `ccr/tool_calls.py` (it checks `isinstance(choices, list)`, non-empty, and `isinstance(first_choice, dict)`). Only this OpenAI branch was missed. ## Fix Guard the list and the first element the same way the siblings do: ```python elif provider == "openai": choices = response.get("choices") first = choices[0] if isinstance(choices, list) and choices else {} message = first.get("message", {}) if isinstance(first, dict) else {} return { "role": "assistant", "content": message.get("content"), "tool_calls": message.get("tool_calls"), } ``` A well-formed response is unaffected; an empty/null/absent `choices` now yields `{"role": "assistant", "content": None, "tool_calls": None}` instead of raising. ## 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/ccr/response_handler.py`: guard empty/non-list `choices` and a non-dict first element in the OpenAI branch of `_extract_assistant_message`. - `tests/test_ccr_response_handler.py`: add `TestExtractAssistantMessageEdgeCases` (empty `choices`, `[null]`, absent, and the normal case). - `CHANGELOG.md`: Bug Fixes 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 $ uvx ruff@0.15.17 check headroom/ccr/response_handler.py tests/test_ccr_response_handler.py All checks passed! $ uvx ruff@0.15.17 format --check headroom/ccr/response_handler.py tests/test_ccr_response_handler.py 2 files already formatted # Verified against the REAL imported module (headroom.ccr.response_handler is # light — no ML imports), so this ran locally in the project venv: $ python -c "from headroom.ccr.response_handler import CCRResponseHandler as H; h=H(); \ assert h._extract_assistant_message({'choices': []}, 'openai') == {'role':'assistant','content':None,'tool_calls':None}" # (no IndexError; normal case still extracts content/tool_calls) ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, project venv (`uv sync --extra proxy`), `uvx ruff@0.15.17`. - Exact command / steps: imported the real `CCRResponseHandler` and called `_extract_assistant_message` with `{"choices": []}`, `{"choices": [null]}`, `{}` (absent), and a normal `{"choices": [{"message": {...}}]}`. - Observed result: the OLD code raised `IndexError` on `[]` and `AttributeError` on `[null]`; the NEW code returns `{"role": "assistant", "content": None, "tool_calls": None}` for all three malformed shapes and still extracts `content`/`tool_calls` from a well-formed response. Because `response_handler` has no ML imports, this ran against the actual module, not a replica. - Not tested: a live CCR retrieval round trip through a gateway that emits empty choices; the added unit tests drive `_extract_assistant_message` directly. ## 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 - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes `headroom/ccr/response_handler.py` is a light module (no ML imports), so unlike most of my recent PRs I verified the fix by importing the real class in the project venv (output above), in addition to the added unit tests. This aligns the OpenAI branch with the already-defensive Google branch and `ccr/tool_calls.py`. --- headroom/ccr/response_handler.py | 10 +++++++++- tests/test_ccr_response_handler.py | 27 +++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/headroom/ccr/response_handler.py b/headroom/ccr/response_handler.py index f81fe2236..af9241e82 100644 --- a/headroom/ccr/response_handler.py +++ b/headroom/ccr/response_handler.py @@ -379,7 +379,15 @@ class CCRResponseHandler: "content": response.get("content", []), } elif provider == "openai": - message = response.get("choices", [{}])[0].get("message", {}) + # Guard an empty/malformed ``choices`` the same way the Google branch + # below (and ccr/tool_calls.py) already do: ``response.get("choices", + # [{}])`` only falls back when the key is absent, so a present-but- + # empty ``choices: []`` (or ``[null]``) — which OpenAI-compatible + # gateways can send on a content-filtered/usage-only response — made + # ``[0]`` raise IndexError (or ``.get`` raise on a non-dict). + choices = response.get("choices") + first = choices[0] if isinstance(choices, list) and choices else {} + message = first.get("message", {}) if isinstance(first, dict) else {} return { "role": "assistant", "content": message.get("content"), diff --git a/tests/test_ccr_response_handler.py b/tests/test_ccr_response_handler.py index 2e850ac1b..d7278b9ab 100644 --- a/tests/test_ccr_response_handler.py +++ b/tests/test_ccr_response_handler.py @@ -719,3 +719,30 @@ class TestExtractAssistantMessage: assert message["role"] == "assistant" assert message["content"] == "Hello" assert message["tool_calls"] == [{"id": "123"}] + + +class TestExtractAssistantMessageEdgeCases: + """Regression: `_extract_assistant_message` must not crash on an empty or + malformed OpenAI `choices` array (OpenAI-compatible gateways can send + `choices: []` or `[null]` on content-filtered / usage-only responses).""" + + def test_openai_empty_choices_does_not_crash(self): + handler = CCRResponseHandler() + msg = handler._extract_assistant_message({"choices": []}, "openai") + assert msg == {"role": "assistant", "content": None, "tool_calls": None} + + def test_openai_null_first_choice_does_not_crash(self): + handler = CCRResponseHandler() + msg = handler._extract_assistant_message({"choices": [None]}, "openai") + assert msg == {"role": "assistant", "content": None, "tool_calls": None} + + def test_openai_absent_choices_does_not_crash(self): + handler = CCRResponseHandler() + msg = handler._extract_assistant_message({}, "openai") + assert msg == {"role": "assistant", "content": None, "tool_calls": None} + + def test_openai_normal_choice_still_extracts(self): + handler = CCRResponseHandler() + resp = {"choices": [{"message": {"content": "hi", "tool_calls": [{"id": "1"}]}}]} + msg = handler._extract_assistant_message(resp, "openai") + assert msg == {"role": "assistant", "content": "hi", "tool_calls": [{"id": "1"}]}