fix(proxy/memory): harden MemoryToolAdapter._extract_tool_calls against malformed responses#2485
Conversation
…st malformed responses
_extract_tool_calls reads from the untrusted upstream response (its own
_get_tool_* helpers are already hardened for that reason), but its four
branches indexed and chained .get without guards: choices[0].get on a
choices: [null] turn, candidates[0].get("content", {}).get("parts") on a
null candidate/content, and block.get on a non-dict content block all raised
AttributeError. An OpenAI-compatible gateway sends choices: [null] on a
content-filtered / usage-only turn. Guard the first choice/candidate and the
nested content/message/parts for type, and skip non-dict elements, mirroring
the sibling MemoryHandler._extract_tool_calls. Factor the OpenAI and
Anthropic extraction into shared helpers so the generic fallback stays in
lockstep.
PR governanceThis PR follows the template and is marked ready for human review. |
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
JerrettDavis
left a comment
There was a problem hiding this comment.
Thanks for tightening the outer response shapes here. I found one remaining malformed-response path in the OpenAI branch that still escapes the new guard:
response = {"choices": [{"message": {"tool_calls": [None]}}]}
adapter._extract_tool_calls(response, "openai") # returns [None]
adapter.has_memory_tool_calls(response, "openai") # AttributeError: 'NoneType' object has no attribute 'get'The new _extract_openai_tool_calls guards choices[0] and message, but it still returns list(message.get("tool_calls", []) or []) without filtering the tool-call elements themselves. Since the downstream paths (has_memory_tool_calls, process_tool_responses) immediately pass each entry into _get_tool_name / _get_tool_input, a null or string element in tool_calls still crashes the same untrusted upstream-response path this PR is hardening.
Please filter OpenAI tool_calls to dict entries, similar to the Anthropic/Gemini branches, and add a regression for choices[0].message.tool_calls: [None, {valid call}] proving the malformed entry is skipped and the valid one still parses.
Other checked state: PR is mergeable and visible CI is green. I did not run the full test file after finding this blocker; I verified the failing shape locally with the snippet above.
Description
MemoryToolAdapter._extract_tool_callsreads tool calls out of an upstream response for each provider format. Its own_get_tool_name/_get_tool_inputhelpers were already hardened against malformed upstream data (seetest_memory_tool_adapter_null_fields.py: "reachable from the untrusted upstream response"), but_extract_tool_callsitself indexed and chained.getwith no guards:An OpenAI-compatible gateway can send
choices: [null]on a content-filtered / usage-only turn (the same shape the siblingMemoryHandler._extract_tool_callswas hardened for — see PR #2483). The truthiness check onchoiceslets[null]through, andchoices[0].getthen raisesAttributeError. The Gemini branch has the same gap on a null candidate / nullcontent, and the Anthropic branch (and the generic fallback) crash on a non-dict content block.Fix
Guard the first choice/candidate and the nested
message/content/partsfor type before calling.get, and skip non-dict elements — mirroringMemoryHandler._extract_tool_callsand this adapter's already-hardened_get_tool_*helpers. The OpenAI and Anthropic extraction is factored into two small helpers so the explicit branches and the generic fallback stay in lockstep. Well-formed responses are unchanged.Type of Change
Changes Made
headroom/proxy/memory_tool_adapter.py: type-guard the OpenAI, Gemini, and Anthropic branches (and the generic fallback) of_extract_tool_calls; add_extract_anthropic_tool_uses/_extract_openai_tool_callshelpers shared by the explicit branches and the fallback.tests/test_memory_tool_adapter_null_fields.py: regressions for malformed OpenAI choices, malformed Gemini candidates/content/parts (plus a null part skipped), a non-dict Anthropic block, and a valid OpenAI call still parsed.Testing
pytest)ruff check .)mypy headroom)Test Output
Real Behavior Proof
uv sync --extra proxy),uvx ruff@0.15.17/uvx mypy@1.20.2, pytest in the venv._extract_tool_callson a bareMemoryToolAdapter(viaobject.__new__) withchoices: [null]/[str]/[{"message": null}](openai and generic),candidates: [null]/ null content / null parts (gemini), and a null Anthropic content block; then revertedmemory_tool_adapter.pyand re-ran.[], a real GeminifunctionCallpart is still returned (null part skipped), and a valid OpenAI tool call still parses; with the fix reverted the openai[null], gemini null-content, and non-dict-block cases raiseAttributeError: 'NoneType' object has no attribute 'get'. Ran against the actual module viatests/test_memory_tool_adapter_null_fields.py.choices: [null]routed through the adapter end to end.Review Readiness
Checklist