fix: replace bare asserts with ValueError in _normalize_response (tool response parsing)#716
Conversation
_normalize_response() used bare �ssert statements to validate that an MCP tool response contains exactly one text content block before converting it to structuredContent. This has two problems: 1. Python removes �ssert statements when running with -O / PYTHONOPTIMIZE, so the guard disappears entirely in optimized production deployments. 2. When the assert fires, it raises AssertionError with no message, giving the caller no information about which tool, which response, or what the content actually was. MCP tools routinely return multiple content blocks or non-text types (image, blob), both of which are valid per the MCP spec. Replace both asserts with explicit ValueError that include the unexpected content in the message so failures are immediately actionable. Add _manager_test.py with tests for the normal path and both error cases, which also serve as the first test coverage for this module.
wbizmo
left a comment
There was a problem hiding this comment.
Diff matches the description. _manager.py swaps the two bare asserts for ValueErrors that actually tell you what came back, _manager_test.py is new with 5 tests. There's a third file in the raw diff (_transformer.py) but that's not from this PR, main just moved on since this branch was forked, some unrelated attention masking fix that already landed separately.
Ran the tests against the fix, all 5 pass. Then swapped the old assert code back in and ran them again, the two regression tests fail with a plain AssertionError right at the assert lines, then pass again once the fix goes back in. So this is a real regression guard, not tests written to just match whatever the new code does.
Reasoning checks out too. Bare asserts get stripped under python -O / PYTHONOPTIMIZE, and a ValueError that actually shows you the content is a real upgrade from a bare AssertionError.
Approving.
|
Thanks @wbizmo for the thorough review and for verifying the regression tests actually fail on the old code This PR is green across the board (CLA, tests, no conflicts) just needs a look from someone with write access to merge. Tagging @Qwlouse and @osanseviero for visibility since you've both been active on this repo, but happy for any maintainer with bandwidth to take a look Happy to make any changes needed. |
Description
This has two concrete problems:
1. The guards disappear in optimized deployments.
Python silently removes all
assertstatements when running with-OorPYTHONOPTIMIZE=1. This means the validation is completely absent in anyproduction environment using optimized Python, allowing invalid tool
responses to corrupt
structuredContentwith no error.2. The error gives no useful context.
When the assert does fire (in non-optimized mode), the caller sees:
No information about which tool, how many content blocks were returned, or
what the types were. Debugging requires a full stack trace and a debugger.
Root cause
_normalize_responsewas written assuming MCP tools always return a singletext content block. This is a common case but not a contract — the MCP spec
allows tools to return multiple content blocks and non-text types (image,
blob). Both are valid and observable in real MCP server implementations.
Fix
Replace both
assertstatements with explicitValueErrorthat include theactual content in the message:
ValueErroris the correct exception for runtime data validation. UnlikeAssertionError, it is never optimized away, and the message immediatelyidentifies what the tool actually returned.
Testing
Added
gemma/gm/tools/_manager_test.py— the first test file for thismodule — with four tests:
test_normalize_response_passthrough_when_structured_content_presenttest_normalize_response_success_single_text_blocktest_normalize_response_error_single_text_blocktest_normalize_response_raises_for_multiple_content_blocks← catches the bugtest_normalize_response_raises_for_non_text_content_type← catches the bugThe last two tests would have failed against the original code (wrong
exception type) and pass after the fix.