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.
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.

1 participant