headroom/tests/test_memory_tool_adapter_null_fields.py
Abhay Singh ef1e7e403b
fix(proxy/memory): don't crash the memory tool adapter on a null function/arguments (#2270)
## Description

The memory tool adapter crashes when an upstream response carries a tool
call whose `function` or `arguments` field is explicitly `null`.

`_get_tool_name`, `_get_tool_id`, and `_get_tool_input` all read the
nested function like this:

```python
str(tool_call.get("function", {}).get("name", ""))
tool_call.get("function", {}).get("arguments", "{}")
```

Two distinct crashes:

1. **Null `function`** → `AttributeError`. `dict.get("function", {})`
only substitutes `{}` for a *missing* key. A present-but-null `{"id":
"c1", "type": "function", "function": null}` (which upstreams and
gateways emit for partial/streamed tool calls) makes the result `None`,
and `None.get("name")` raises.
2. **Null `arguments`** → `TypeError`. `tool_call.get("function",
{}).get("arguments", "{}")` returns `None` when `arguments` is null, and
`json.loads(None)` raises `TypeError` — which the surrounding `except
json.JSONDecodeError` does **not** catch.

Both parse the untrusted upstream response inside `handle_tool_calls`,
so a single malformed tool call takes down memory tool handling. Notably
`parse_tool_call` in `headroom/ccr/tool_injection.py` already catches
the `json.loads(None)` `TypeError` with an explicit comment, so the
null-arguments hazard is known in the codebase; this path just wasn't
hardened.

## Fix

- Coalesce `function` / `functionCall` with `or {}` so a null value
collapses to `{}`.
- Coalesce the arguments string with `or "{}"` and add `TypeError` to
the `except`, so a null `arguments` yields `{}` instead of crashing.

Real tool calls parse exactly as before.

Closes #

## 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/memory_tool_adapter.py`: coalesce
`function`/`functionCall` (`or {}`) in
`_get_tool_name`/`_get_tool_id`/`_get_tool_input`; coalesce the
arguments string (`or "{}"`) and catch `TypeError`.
- `tests/test_memory_tool_adapter_null_fields.py`: new tests for null
function, null arguments, and that real calls still parse.
- `CHANGELOG.md`: Bug Fixes entry.

## Testing

- [ ] 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/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, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOM-kills this box (ML stack import), so I
reproduced the parse helpers with a dependency-free script and left the
full pytest to CI.
- Exact command / steps: ran `{"function": null}` and `{"function":
{"arguments": null}}` (plus a real `memory_save` call) through the OLD
and NEW `_get_tool_name`/`_get_tool_input` logic.
- Observed result: OLD raises `AttributeError` on the null function and
`TypeError` on the null arguments; NEW returns `""`/`{}` for both and
still parses the real call to `{"content": "hi"}`.
- Not tested: a live upstream emitting a null field; full local `pytest`
deferred to CI (OOM).

## 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
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

The "unit tests pass locally" box is unchecked because the full suite
imports the ML stack, which I can't run here. The parse helpers read
only their `tool_call` argument (no instance state), so the new test
exercises them on a bare instance via `object.__new__` — it runs under
the normal CI pytest job, and the standalone proof above corroborates
it.

---------

Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
2026-07-16 14:37:27 -07:00

35 lines
1.4 KiB
Python

"""A tool call with a null ``function`` / ``arguments`` must not crash the
memory tool adapter's provider-format parsing.
``dict.get("function", {})`` returns ``None`` for a present-but-null key, so the
following ``.get`` raised ``AttributeError``; a null ``arguments`` makes
``json.loads(None)`` raise ``TypeError`` that the bare ``JSONDecodeError`` catch
missed. Both are reachable from the untrusted upstream response. The parse
helpers read only the ``tool_call`` argument, so we exercise them on a bare
instance via ``object.__new__``.
"""
from __future__ import annotations
from headroom.proxy.memory_tool_adapter import MemoryToolAdapter
_adapter = object.__new__(MemoryToolAdapter)
def test_get_tool_name_survives_null_function():
tc = {"id": "c1", "type": "function", "function": None}
assert _adapter._get_tool_name(tc, "openai") == ""
assert _adapter._get_tool_id(tc, "openai") == "c1"
assert _adapter._get_tool_input(tc, "openai") == {}
def test_get_tool_input_survives_null_arguments():
# json.loads(None) raises TypeError, not JSONDecodeError.
tc = {"function": {"name": "memory_save", "arguments": None}}
assert _adapter._get_tool_input(tc, "openai") == {}
def test_get_tool_helpers_still_parse_real_calls():
tc = {"function": {"name": "memory_save", "arguments": '{"content": "hi"}'}}
assert _adapter._get_tool_name(tc, "openai") == "memory_save"
assert _adapter._get_tool_input(tc, "openai") == {"content": "hi"}