Skip to content

fix(proxy/memory): harden MemoryToolAdapter._extract_tool_calls against malformed responses#2485

Open
abhay-codes07 wants to merge 1 commit into
headroomlabs-ai:mainfrom
abhay-codes07:fix/memory-tool-adapter-extract-malformed
Open

fix(proxy/memory): harden MemoryToolAdapter._extract_tool_calls against malformed responses#2485
abhay-codes07 wants to merge 1 commit into
headroomlabs-ai:mainfrom
abhay-codes07:fix/memory-tool-adapter-extract-malformed

Conversation

@abhay-codes07

Copy link
Copy Markdown
Contributor

Description

MemoryToolAdapter._extract_tool_calls reads tool calls out of an upstream response for each provider format. Its own _get_tool_name / _get_tool_input helpers were already hardened against malformed upstream data (see test_memory_tool_adapter_null_fields.py: "reachable from the untrusted upstream response"), but _extract_tool_calls itself indexed and chained .get with no guards:

# openai
choices = response.get("choices", [])
if choices:
    message = choices[0].get("message", {})   # choices: [null] -> AttributeError
    ...
# gemini
content = candidates[0].get("content", {})    # candidates: [null] -> AttributeError
parts = content.get("parts", [])              # content: null -> AttributeError
# anthropic
[block for block in content if block.get("type") == "tool_use"]  # null block -> AttributeError

An OpenAI-compatible gateway can send choices: [null] on a content-filtered / usage-only turn (the same shape the sibling MemoryHandler._extract_tool_calls was hardened for — see PR #2483). The truthiness check on choices lets [null] through, and choices[0].get then raises AttributeError. The Gemini branch has the same gap on a null candidate / null content, 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 / parts for type before calling .get, and skip non-dict elements — mirroring MemoryHandler._extract_tool_calls and 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

  • 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/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_calls helpers 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

  • Unit tests pass (pytest)
  • Linting passes (ruff check .)
  • Type checking passes (mypy headroom)
  • New tests added for new functionality
  • Manual testing performed

Test Output

$ python -m pytest tests/test_memory_tool_adapter_null_fields.py -q
7 passed

# with the fix reverted, the three new malformed-shape tests fail with
# AttributeError: 'NoneType' object has no attribute 'get'

$ uvx ruff@0.15.17 check headroom/proxy/memory_tool_adapter.py tests/test_memory_tool_adapter_null_fields.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/memory_tool_adapter.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: called _extract_tool_calls on a bare MemoryToolAdapter (via object.__new__) with choices: [null] / [str] / [{"message": null}] (openai and generic), candidates: [null] / null content / null parts (gemini), and a null Anthropic content block; then reverted memory_tool_adapter.py and re-ran.
  • Observed result: with the fix the malformed shapes return [], a real Gemini functionCall part 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 raise AttributeError: 'NoneType' object has no attribute 'get'. Ran against the actual module via tests/test_memory_tool_adapter_null_fields.py.
  • Not tested: a live OpenAI-compatible gateway emitting choices: [null] routed through the adapter end to end.

Review Readiness

  • I have performed a self-review
  • This PR is ready for human review

Checklist

  • My code follows the project's style guidelines
  • I have performed a self-review of my code
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have updated the CHANGELOG.md if applicable

…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.
Copilot AI review requested due to automatic review settings July 22, 2026 06:47

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@github-actions

Copy link
Copy Markdown
Contributor

PR governance

This PR follows the template and is marked ready for human review.

@github-actions github-actions Bot added the status: ready for review Pull request body is complete and the author marked it ready for human review label Jul 22, 2026
@codecov-commenter

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 95.65217% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
headroom/proxy/memory_tool_adapter.py 95.65% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@JerrettDavis JerrettDavis left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@github-actions github-actions Bot removed the status: ready for review Pull request body is complete and the author marked it ready for human review label Jul 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants