From 967b67e93fd7f85becd1226bb383ee4649546887 Mon Sep 17 00:00:00 2001 From: DhruvTilva Date: Sun, 5 Jul 2026 11:27:58 +0530 Subject: [PATCH] fix: replace bare asserts with ValueError in _normalize_response _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. --- gemma/gm/tools/_manager.py | 13 ++++- gemma/gm/tools/_manager_test.py | 97 +++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 3 deletions(-) create mode 100644 gemma/gm/tools/_manager_test.py diff --git a/gemma/gm/tools/_manager.py b/gemma/gm/tools/_manager.py index 00012cfa..10028a08 100644 --- a/gemma/gm/tools/_manager.py +++ b/gemma/gm/tools/_manager.py @@ -129,9 +129,16 @@ def _normalize_response( if response.structuredContent is not None: return response - # Assume the response is a single text block. - assert len(response.content) == 1 - assert response.content[0].type == 'text' + 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}' + ) content = response.content[0].text if response.isError: diff --git a/gemma/gm/tools/_manager_test.py b/gemma/gm/tools/_manager_test.py new file mode 100644 index 00000000..b414d669 --- /dev/null +++ b/gemma/gm/tools/_manager_test.py @@ -0,0 +1,97 @@ +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for _manager._normalize_response.""" + +import pytest +from gemma.gm.tools import _manager +from unittest import mock + + +def _make_text_content(text: str): + content = mock.MagicMock() + content.type = 'text' + content.text = text + return content + + +def _make_image_content(): + content = mock.MagicMock() + content.type = 'image' + return content + + +def _make_tool_result( + *, + content, + is_error: bool = False, + structured_content=None, +): + result = mock.MagicMock() + result.content = list(content) + result.isError = is_error + result.structuredContent = structured_content + return result + + +def test_normalize_response_passthrough_when_structured_content_present(): + """If structuredContent is already set, return the response unchanged.""" + result = _make_tool_result( + content=[_make_text_content('hello')], + structured_content={'result': 'already set'}, + ) + out = _manager._normalize_response(result) + assert out is result + + +def test_normalize_response_success_single_text_block(): + """Single text block with no error should map to {'result': }.""" + result = _make_tool_result(content=[_make_text_content('hello world')]) + out = _manager._normalize_response(result) + assert out.structuredContent == {'result': 'hello world'} + + +def test_normalize_response_error_single_text_block(): + """Single text block with isError=True should map to {'error': }.""" + result = _make_tool_result( + content=[_make_text_content('something went wrong')], is_error=True + ) + out = _manager._normalize_response(result) + assert out.structuredContent == {'error': 'something went wrong'} + + +def test_normalize_response_raises_for_multiple_content_blocks(): + """ValueError (not AssertionError) is raised when content has >1 blocks. + + Bare `assert` statements are silently removed under `python -O`, making + this guard disappear in optimized production deployments. The fix + replaces them with explicit ValueError so the check is always active and + the error message identifies what was returned. + """ + result = _make_tool_result( + content=[_make_text_content('block1'), _make_text_content('block2')] + ) + with pytest.raises(ValueError, match='Expected a single content block'): + _manager._normalize_response(result) + + +def test_normalize_response_raises_for_non_text_content_type(): + """ValueError is raised when the single content block is not type 'text'. + + MCP tools may return image or blob content; this path is unsupported and + must raise a clear, actionable error rather than an opaque AssertionError. + """ + result = _make_tool_result(content=[_make_image_content()]) + with pytest.raises(ValueError, match="Expected a text content block"): + _manager._normalize_response(result)