fix(agentgateway): bound MCP protocol calls with a wall-clock deadline

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

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