From 01be269b6864f4918630b4e77f4dbcb962f0e0ed Mon Sep 17 00:00:00 2001 From: Dinesh Date: Fri, 30 Jan 2026 21:00:39 +0530 Subject: [PATCH 1/2] refactor(mcp):used logger to log instead of print and updated unittest case --- .../adk/tools/mcp_tool/mcp_session_manager.py | 7 +++---- .../tools/mcp_tool/test_mcp_session_manager.py | 14 ++++++-------- 2 files changed, 9 insertions(+), 12 deletions(-) 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 2571cb6dbd..cd757e1500 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -434,10 +434,9 @@ async def close(self): await exit_stack.aclose() except Exception as e: # Log the error but don't re-raise to avoid blocking shutdown - print( - 'Warning: Error during MCP session cleanup for' - f' {session_key}: {e}', - file=self._errlog, + logger.warning( + f'Warning: Error during MCP session cleanup for:{session_key}', + exc_info=True, ) finally: del self._sessions[session_key] 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 410e528295..096c5f2df4 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -388,7 +388,8 @@ async def test_close_success(self): assert len(manager._sessions) == 0 @pytest.mark.asyncio - async def test_close_with_errors(self): + @patch("google.adk.tools.mcp_tool.mcp_session_manager.logger") + async def test_close_with_errors(self, mock_logger): """Test cleanup when some sessions fail to close.""" manager = MCPSessionManager(self.mock_stdio_connection_params) @@ -403,9 +404,6 @@ async def test_close_with_errors(self): manager._sessions["session1"] = (session1, exit_stack1) manager._sessions["session2"] = (session2, exit_stack2) - custom_errlog = StringIO() - manager._errlog = custom_errlog - # Should not raise exception await manager.close() @@ -413,10 +411,10 @@ async def test_close_with_errors(self): exit_stack2.aclose.assert_called_once() assert len(manager._sessions) == 0 - # Error should be logged - error_output = custom_errlog.getvalue() - assert "Warning: Error during MCP session cleanup" in error_output - assert "Close error 1" in error_output + # Error should be logged via logger.warning + mock_logger.warning.assert_called() + warning_call = str(mock_logger.warning.call_args) + assert "Warning: Error during MCP session cleanup" in warning_call @pytest.mark.asyncio @patch("google.adk.tools.mcp_tool.mcp_session_manager.stdio_client") From 3a19d076f0ee542960c6a97671958016d8e89c7c Mon Sep 17 00:00:00 2001 From: Dinesh Date: Fri, 30 Jan 2026 21:26:54 +0530 Subject: [PATCH 2/2] refactored with suggested changes --- src/google/adk/tools/mcp_tool/mcp_session_manager.py | 2 +- tests/unittests/tools/mcp_tool/test_mcp_session_manager.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) 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 cd757e1500..a9bfcfbc37 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -435,7 +435,7 @@ async def close(self): except Exception as e: # Log the error but don't re-raise to avoid blocking shutdown logger.warning( - f'Warning: Error during MCP session cleanup for:{session_key}', + f'Error during MCP session cleanup for {session_key}', exc_info=True, ) finally: 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 096c5f2df4..cc2cb487d9 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -412,9 +412,10 @@ async def test_close_with_errors(self, mock_logger): assert len(manager._sessions) == 0 # Error should be logged via logger.warning - mock_logger.warning.assert_called() - warning_call = str(mock_logger.warning.call_args) - assert "Warning: Error during MCP session cleanup" in warning_call + mock_logger.warning.assert_called_once() + args, kwargs = mock_logger.warning.call_args + assert "Error during MCP session cleanup for session1" in args[0] + assert kwargs.get("exc_info") @pytest.mark.asyncio @patch("google.adk.tools.mcp_tool.mcp_session_manager.stdio_client")