From 3ebef82a52f2b650d1ee1f94dbdeeb31fbd0ae10 Mon Sep 17 00:00:00 2001 From: Kathy Wu Date: Mon, 6 Jul 2026 14:39:40 -0700 Subject: [PATCH] fix: Truncate MCP http debug logs if greater than 1000 chars To prevent OOM issues if HTTP request/response bodies are too long Co-authored-by: Kathy Wu PiperOrigin-RevId: 943502772 --- .../adk/tools/mcp_tool/mcp_session_manager.py | 8 ++++ .../mcp_tool/test_mcp_session_manager.py | 43 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 8407ee7a..3a61929e 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -81,6 +81,8 @@ from .session_context import SessionContext logger = logging.getLogger('google_adk.' + __name__) +_MAX_LOG_BODY_LENGTH = 1000 + def create_mcp_http_client( headers: dict[str, str] | None = None, @@ -271,6 +273,8 @@ class _DebugHttpxClientFactory: request_body = response.request.content.decode( 'utf-8', errors='replace' ) + if len(request_body) > _MAX_LOG_BODY_LENGTH: + request_body = request_body[:_MAX_LOG_BODY_LENGTH] + '... [truncated]' except Exception: # pylint: disable=broad-exception-caught request_body = '' @@ -278,6 +282,10 @@ class _DebugHttpxClientFactory: try: await response.aread() response_body = response.text + if len(response_body) > _MAX_LOG_BODY_LENGTH: + response_body = ( + response_body[:_MAX_LOG_BODY_LENGTH] + '... [truncated]' + ) except Exception as e: # pylint: disable=broad-exception-caught response_body = f'' else: diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index b769eca5..2f6a1130 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -1471,3 +1471,46 @@ class TestDebugHttpxClientFactory: client = debug_factory({"X-Test": "Val"}, None, None) assert client is base_client await base_client.aclose() + + @pytest.mark.asyncio + async def test_response_hook_truncates_large_bodies(self): + """Test that response hook truncates request and response bodies exceeding limit.""" + base_client = httpx.AsyncClient() + base_factory = Mock(return_value=base_client) + debug_factory = _DebugHttpxClientFactory(base_factory) + + # Mock request and response with large content + large_req_body = b"a" * 1500 + large_resp_body = "b" * 1500 + + mock_request = Mock(spec=httpx.Request) + mock_request.method = "POST" + mock_request.content = large_req_body + mock_request.headers = httpx.Headers() + + mock_response = Mock(spec=httpx.Response) + mock_response.url = httpx.URL("https://example.com/large") + mock_response.status_code = 200 + mock_response.request = mock_request + mock_response.headers = httpx.Headers({"content-type": "application/json"}) + mock_response.text = large_resp_body + mock_response.aread = AsyncMock() + + debug_list = [] + token = _http_debug_var.set(debug_list) + try: + await debug_factory._response_hook(mock_response) + finally: + _http_debug_var.reset(token) + + assert len(debug_list) == 1 + record = debug_list[0] + assert len(record["request_body"]) == 1015 # 1000 + len("... [truncated]") + assert record["request_body"].endswith("... [truncated]") + assert record["request_body"].startswith("a" * 1000) + + assert len(record["response_body"]) == 1015 # 1000 + len("... [truncated]") + assert record["response_body"].endswith("... [truncated]") + assert record["response_body"].startswith("b" * 1000) + + await base_client.aclose()