Skip to content

fix: replace bare asserts with ValueError in _normalize_response (tool response parsing)#716

Open
DhruvTilva wants to merge 1 commit into
google-deepmind:mainfrom
DhruvTilva:fix/tool-response-assert-to-valueerror
Open

fix: replace bare asserts with ValueError in _normalize_response (tool response parsing)#716
DhruvTilva wants to merge 1 commit into
google-deepmind:mainfrom
DhruvTilva:fix/tool-response-assert-to-valueerror

Conversation

@DhruvTilva

Copy link
Copy Markdown

Description

## What fixes

`_normalize_response()` in `gemma/gm/tools/_manager.py` validates that an
MCP tool response contains exactly one text content block before converting
it to `structuredContent`. This validation was implemented with bare
`assert` statements:

```python
# Before
assert len(response.content) == 1
assert response.content[0].type == 'text'

This has two concrete problems:

1. The guards disappear in optimized deployments.
Python silently removes all assert statements when running with -O or
PYTHONOPTIMIZE=1. This means the validation is completely absent in any
production environment using optimized Python, allowing invalid tool
responses to corrupt structuredContent with no error.

2. The error gives no useful context.
When the assert does fire (in non-optimized mode), the caller sees:

AssertionError

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_response was written assuming MCP tools always return a single
text 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 assert statements with explicit ValueError that include the
actual content in the message:

# After
if len(response.content) != 1:
    raise ValueError(
        f'Expected a single content block in the tool response, got'
        f' {len(response.content)}: {response.content!r}'
    )
if response.content[0].type != 'text':
    raise ValueError(
        f'Expected a text content block in the tool response, got'
        f' type {response.content[0].type!r}: {response.content[0]!r}'
    )

ValueError is the correct exception for runtime data validation. Unlike
AssertionError, it is never optimized away, and the message immediately
identifies what the tool actually returned.

Testing

Added gemma/gm/tools/_manager_test.py — the first test file for this
module — with four tests:

  • test_normalize_response_passthrough_when_structured_content_present
  • test_normalize_response_success_single_text_block
  • test_normalize_response_error_single_text_block
  • test_normalize_response_raises_for_multiple_content_blocks ← catches the bug
  • test_normalize_response_raises_for_non_text_content_type ← catches the bug

The last two tests would have failed against the original code (wrong
exception type) and pass after the fix.

_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 wbizmo left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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.

@DhruvTilva

Copy link
Copy Markdown
Author

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.

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.

2 participants