From 427ef2eeabdc07856f421255cd8333aa60d7e558 Mon Sep 17 00:00:00 2001 From: wangzhengzhuo05 <175673456+wangzhengzhuo05@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:21:19 +0800 Subject: [PATCH] fix(agentgateway): bound MCP protocol calls with a wall-clock deadline The httpx.AsyncClient timeout only bounds individual HTTP chunk reads; SSE keep-alives from an unresponsive MCP server reset that timer continuously, so session.initialize() / list_tools() / call_tool() could hang forever instead of failing fast. Add _mcp_call_with_deadline() which wraps a protocol call in asyncio.wait_for, enforcing a deadline that keep-alives cannot extend, and raises AgentGatewaySDKError on timeout. Apply it to initialize, list_tools and call_tool in the LoB flow so tool listing and invocation both surface a clear error when the server stalls. Fixes #313 --- src/sap_cloud_sdk/agentgateway/_lob.py | 65 ++++++++++++++- tests/agentgateway/unit/test_lob.py | 105 +++++++++++++++++++++++++ 2 files changed, 166 insertions(+), 4 deletions(-) diff --git a/src/sap_cloud_sdk/agentgateway/_lob.py b/src/sap_cloud_sdk/agentgateway/_lob.py index 8a495979..b741e967 100644 --- a/src/sap_cloud_sdk/agentgateway/_lob.py +++ b/src/sap_cloud_sdk/agentgateway/_lob.py @@ -9,6 +9,7 @@ import logging import os import uuid +from typing import Any, Awaitable import httpx from mcp import ClientSession @@ -57,6 +58,42 @@ _DESTINATION_INSTANCE = "default" +async def _mcp_call_with_deadline( + operation: str, + awaitable: Awaitable, + *, + timeout: float, + target: str, +) -> Any: + """Run an MCP protocol call under a wall-clock deadline. + + The ``httpx.AsyncClient`` timeout only bounds individual HTTP chunk + reads; an MCP stream's SSE keep-alives reset that timer continuously, so + an unresponsive server can keep ``session.initialize()`` / + ``list_tools()`` / ``call_tool()`` alive forever. Wrapping the protocol + call in ``asyncio.wait_for`` enforces a deadline that keep-alives cannot + extend, letting the caller fail fast instead of hanging. + + Args: + operation: Human-readable protocol operation name (for the error). + awaitable: The protocol coroutine to run under the deadline. + timeout: Wall-clock seconds budget for the call. + target: Human-readable server/tool name (for the error). + + Returns: + The protocol call's result. + + Raises: + AgentGatewaySDKError: If the call does not complete within ``timeout``. + """ + try: + return await asyncio.wait_for(awaitable, timeout=timeout) + except asyncio.TimeoutError as exc: + raise AgentGatewaySDKError( + f"MCP {operation} on '{target}' timed out after {timeout}s" + ) from exc + + def _system_scope_key(tenant_subdomain: str) -> str: """Build the cache scope key for tenant-scoped system auth.""" return f"lob-system::{tenant_subdomain}" @@ -383,9 +420,19 @@ async def list_server_tools( *_, ): async with ClientSession(read, write) as session: - init_result = await session.initialize() + init_result = await _mcp_call_with_deadline( + "initialize", + session.initialize(), + timeout=timeout, + target=fragment_name, + ) server_name = mcp_server_name(init_result) or fragment_name - result = await session.list_tools() + result = await _mcp_call_with_deadline( + "list_tools", + session.list_tools(), + timeout=timeout, + target=fragment_name, + ) tools = result.tools or [] if not tools: logger.info( @@ -514,8 +561,18 @@ async def call_mcp_tool_lob( *_, ): async with ClientSession(read, write) as session: - await session.initialize() - result = await session.call_tool(tool.name, kwargs) + await _mcp_call_with_deadline( + "initialize", + session.initialize(), + timeout=timeout, + target=tool.fragment_name or tool.name, + ) + result = await _mcp_call_with_deadline( + f"call_tool({tool.name})", + session.call_tool(tool.name, kwargs), + timeout=timeout, + target=tool.fragment_name or tool.name, + ) if not result.content: logger.warning( "Tool '%s' on '%s' returned empty content", tool.name, tool.url diff --git a/tests/agentgateway/unit/test_lob.py b/tests/agentgateway/unit/test_lob.py index d570b9cd..7c1729cd 100644 --- a/tests/agentgateway/unit/test_lob.py +++ b/tests/agentgateway/unit/test_lob.py @@ -1,5 +1,6 @@ """Unit tests for LoB agent flow.""" +import asyncio import logging import os from unittest.mock import patch, MagicMock, AsyncMock @@ -904,6 +905,73 @@ async def test_returns_tools_with_server_info_name(self): assert result[0].name == "do-something" assert result[0].server_name == "real-server-name" + @pytest.mark.asyncio + async def test_initialize_timeout_raises_agent_gateway_error(self): + """A stalled MCP server must fail fast instead of hanging forever. + + The httpx client timeout only bounds individual chunk reads; SSE + keep-alives reset it continuously, so the protocol call needs its own + wall-clock deadline (xorbitsai/xagent-issue #313). + """ + with ( + patch("sap_cloud_sdk.agentgateway._lob.httpx.AsyncClient") as mock_http, + patch( + "sap_cloud_sdk.agentgateway._lob.streamable_http_client" + ) as mock_stream, + patch("sap_cloud_sdk.agentgateway._lob.ClientSession") as mock_session, + ): + mock_http.return_value.__aenter__.return_value = AsyncMock() + mock_stream.return_value.__aenter__.return_value = ( + AsyncMock(), + AsyncMock(), + None, + ) + + async def _never_returns(): + await asyncio.sleep(3600) + + mock_session_instance = AsyncMock() + mock_session_instance.initialize = _never_returns + mock_session.return_value.__aenter__.return_value = mock_session_instance + + with pytest.raises(AgentGatewaySDKError, match="timed out"): + await list_server_tools( + "https://example.com/mcp", "token", "my-fragment", 0.05 + ) + + @pytest.mark.asyncio + async def test_list_tools_timeout_raises_agent_gateway_error(self): + """list_tools also gets the protocol-level deadline.""" + with ( + patch("sap_cloud_sdk.agentgateway._lob.httpx.AsyncClient") as mock_http, + patch( + "sap_cloud_sdk.agentgateway._lob.streamable_http_client" + ) as mock_stream, + patch("sap_cloud_sdk.agentgateway._lob.ClientSession") as mock_session, + ): + mock_http.return_value.__aenter__.return_value = AsyncMock() + mock_stream.return_value.__aenter__.return_value = ( + AsyncMock(), + AsyncMock(), + None, + ) + + async def _never_returns(): + await asyncio.sleep(3600) + + mock_init = MagicMock() + mock_init.server_info = None + + mock_session_instance = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=mock_init) + mock_session_instance.list_tools = _never_returns + mock_session.return_value.__aenter__.return_value = mock_session_instance + + with pytest.raises(AgentGatewaySDKError, match="list_tools"): + await list_server_tools( + "https://example.com/mcp", "token", "my-fragment", 0.05 + ) + @pytest.mark.asyncio async def test_falls_back_to_fragment_name_when_server_info_missing(self): """Fall back to fragment_name when server_info or its name is absent.""" @@ -1028,6 +1096,43 @@ async def test_returns_empty_string_when_no_content(self): assert result == "" + @pytest.mark.asyncio + async def test_call_tool_timeout_raises_agent_gateway_error(self): + """Tool invocation also gets the protocol-level deadline.""" + tool = MCPTool( + name="test-tool", + server_name="test-server", + description="Test tool", + input_schema={}, + url="https://example.com/mcp", + fragment_name="mcp-server-a", + ) + + async def _never_returns(*_args, **_kwargs): + await asyncio.sleep(3600) + + with ( + patch("sap_cloud_sdk.agentgateway._lob.httpx.AsyncClient") as mock_http, + patch( + "sap_cloud_sdk.agentgateway._lob.streamable_http_client" + ) as mock_stream, + patch("sap_cloud_sdk.agentgateway._lob.ClientSession") as mock_session, + ): + mock_http.return_value.__aenter__.return_value = AsyncMock() + mock_stream.return_value.__aenter__.return_value = ( + AsyncMock(), + AsyncMock(), + None, + ) + + mock_session_instance = AsyncMock() + mock_session_instance.initialize = AsyncMock() + mock_session_instance.call_tool = _never_returns + mock_session.return_value.__aenter__.return_value = mock_session_instance + + with pytest.raises(AgentGatewaySDKError, match="call_tool\\("): + await call_mcp_tool_lob(tool, "user-auth-token", 0.05) + # ============================================================ # Test: list_a2a_fragments