From ba7eba875ffabc18b8bd64b587db6e2e05ee9da7 Mon Sep 17 00:00:00 2001 From: forthfate Date: Sat, 26 Sep 2026 18:59:47 +0900 Subject: [PATCH] Retry transient assistant MCP initialization --- backend/app/assistant_graph.py | 67 +++++++++++++++++++++++++-- backend/tests/test_assistant_graph.py | 37 ++++++++++++++- 2 files changed, 98 insertions(+), 6 deletions(-) diff --git a/backend/app/assistant_graph.py b/backend/app/assistant_graph.py index 4959182..15bc023 100644 --- a/backend/app/assistant_graph.py +++ b/backend/app/assistant_graph.py @@ -4,9 +4,12 @@ import asyncio import json -from typing import Any, Callable, TypedDict +import logging +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Literal, TypedDict from coding_agents import execution_environment_context +from httpx import TransportError from langgraph.graph import END, START, StateGraph from mcp import ClientSession from mcp.client.streamable_http import streamablehttp_client @@ -15,19 +18,59 @@ from .assistant_ui import AssistantUiToolExecutor from .providers import AzureOpenAIProvider, BedrockProvider, ModelSettings +logger = logging.getLogger(__name__) + class AssistantState(TypedDict, total=False): prompt: str response: str +@dataclass +class McpConnectionStatus: + """One Assistant turn's bounded MCP connection state.""" + + state: Literal["idle", "connecting", "ready", "retrying", "failed"] = "idle" + attempts: int = 0 + last_error: str = "" + + class LocalMcpTools: """Expose the mounted OpenOrbit MCP server as native model function tools.""" - def __init__(self, url: str = DEFAULT_MCP_URL): + def __init__(self, url: str = DEFAULT_MCP_URL, *, retry_delays: tuple[float, ...] = (0.25, 0.5, 1.0)): self.url = url + self.retry_delays = retry_delays + self.connection = McpConnectionStatus() + self._definitions: list[dict[str, Any]] | None = None - async def _list_tools(self) -> list[dict[str, Any]]: + @staticmethod + def _retryable(error: Exception) -> bool: + """Retry startup and transport failures, never protocol/configuration failures.""" + return isinstance(error, (OSError, TimeoutError, TransportError)) + + async def _retry_read( + self, operation: Callable[[], Awaitable[list[dict[str, Any]]]] + ) -> list[dict[str, Any]]: + attempts = len(self.retry_delays) + 1 + for attempt in range(1, attempts + 1): + self.connection.state = "connecting" if attempt == 1 else "retrying" + self.connection.attempts = attempt + try: + value = await operation() + self.connection.state, self.connection.last_error = "ready", "" + return value + except Exception as error: + self.connection.last_error = str(error) + if not self._retryable(error) or attempt == attempts: + self.connection.state = "failed" + raise + delay = self.retry_delays[attempt - 1] + logger.info("MCP initialization failed; retrying in %.2fs: %s", delay, error) + await asyncio.sleep(delay) + raise RuntimeError("unreachable MCP retry state") + + async def _list_tools_once(self) -> list[dict[str, Any]]: async with streamablehttp_client(self.url) as (read, write, _): async with ClientSession(read, write) as session: await session.initialize() @@ -41,10 +84,18 @@ async def _list_tools(self) -> list[dict[str, Any]]: for tool in result.tools ] + async def _list_tools(self) -> list[dict[str, Any]]: + return await self._retry_read(self._list_tools_once) + async def _call_tool(self, name: str, arguments: dict[str, Any]) -> str: + # Do not retry calls: an action might already have reached the target + # before a connection drops. Retrying it could duplicate a mutation. + self.connection.state = "connecting" + self.connection.attempts += 1 async with streamablehttp_client(self.url) as (read, write, _): async with ClientSession(read, write) as session: await session.initialize() + self.connection.state, self.connection.last_error = "ready", "" result = await session.call_tool(name, arguments) if result.isError: return json.dumps({"error": "MCP tool failed", "content": result.content}, default=str) @@ -55,9 +106,14 @@ def _run(coro: Any) -> Any: return asyncio.run(coro) def definitions(self) -> list[dict[str, Any]]: + if self._definitions is not None: + return self._definitions try: - return self._run(self._list_tools()) - except Exception: + self._definitions = self._run(self._list_tools()) + return self._definitions + except Exception as error: + self.connection.state, self.connection.last_error = "failed", str(error) + logger.info("MCP tools unavailable: %s", error) # The assistant remains usable while the local server is starting. return [] @@ -65,6 +121,7 @@ def execute(self, name: str, arguments: dict[str, Any]) -> str: try: return self._run(self._call_tool(name, arguments)) except Exception as error: + self.connection.state, self.connection.last_error = "failed", str(error) return json.dumps({"error": f"MCP unavailable: {error}"}) diff --git a/backend/tests/test_assistant_graph.py b/backend/tests/test_assistant_graph.py index 4ad16c3..63b0e7d 100644 --- a/backend/tests/test_assistant_graph.py +++ b/backend/tests/test_assistant_graph.py @@ -1,4 +1,4 @@ -from app.assistant_graph import DEFAULT_MCP_URL, OrbitAssistantGraph, build_assistant_prompt +from app.assistant_graph import DEFAULT_MCP_URL, LocalMcpTools, OrbitAssistantGraph, build_assistant_prompt from app.providers import ModelSettings @@ -61,3 +61,38 @@ def test_assistant_prompt_prioritizes_ui_tools_for_browser_ui_requests(): assert "first call ui_get_context" in prompt assert "Do not use file tools to infer browser UI state" in prompt + + +def test_local_mcp_tools_retries_transient_initialization_failures(monkeypatch): + tools = LocalMcpTools(retry_delays=(0, 0)) + attempts = 0 + + async def flaky_list_tools(): + nonlocal attempts + attempts += 1 + if attempts < 3: + raise OSError("connection refused") + return [{"name": "get_status", "description": "status", "parameters": {}}] + + monkeypatch.setattr(tools, "_list_tools_once", flaky_list_tools) + + assert tools.definitions() == [{"name": "get_status", "description": "status", "parameters": {}}] + assert attempts == 3 + assert tools.connection.state == "ready" + assert tools.connection.attempts == 3 + + +def test_local_mcp_tools_does_not_retry_protocol_failures(monkeypatch): + tools = LocalMcpTools(retry_delays=(0, 0)) + attempts = 0 + + async def bad_protocol(): + nonlocal attempts + attempts += 1 + raise ValueError("unsupported transport") + + monkeypatch.setattr(tools, "_list_tools_once", bad_protocol) + + assert tools.definitions() == [] + assert attempts == 1 + assert tools.connection.state == "failed"