diff --git a/docsgpt/connectors/mcp.py b/docsgpt/connectors/mcp.py index 344b3c20..6ca3ddf3 100644 --- a/docsgpt/connectors/mcp.py +++ b/docsgpt/connectors/mcp.py @@ -1,8 +1,18 @@ -"""MCP helpers for connections: re-scanning a server's actions, switching GitHub's writes.""" +"""MCP helpers for connections: re-scanning a server's actions, switching GitHub's +writes, and calling a server's tools. + +Agents call an MCP server's tools through ``MCPTool``. Server-side work that +reads a service through its MCP server (syncing Linear into Knowledge) opens +one session with :func:`run_connection_session` and makes all its calls in +it, signed in with the connection's own OAuth tokens. +""" from __future__ import annotations -from typing import Optional +import asyncio +import concurrent.futures +import json +from typing import Any, Awaitable, Callable, Optional, TypeVar from docsgpt.agents.tool_pins import carry_pins from docsgpt.connectors import catalog, service @@ -12,6 +22,8 @@ from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessions from docsgpt.storage.db.repositories.user_tools import UserToolsRepository from docsgpt.storage.db.session import db_readonly, db_session +T = TypeVar("T") + class NoMcpTool(LookupError): """The connection has no MCP tool yet (it was set up without tools).""" @@ -184,3 +196,187 @@ def set_builtin_writes(user_id: str, connection: dict, allow: bool) -> dict: removed |= tool_removed refreshed.append(serialized) return {"added": sorted(added), "removed": sorted(removed), "tools": refreshed, "writes": allow} + + +# --------------------------------------------------------------------------- +# Calling a server's tools from the server side +# --------------------------------------------------------------------------- + +SIGN_IN_EXPIRED = "The sign-in expired and could not be renewed. Reconnect to continue." +_TRANSIENT_WORDS = ("rate limit", "ratelimit", "too many requests", "temporarily", "timed out", "timeout") + + +class MCPToolError(Exception): + """An MCP tool answered with an error (a missing issue, a refused argument).""" + + +class MCPSession: + """One open MCP session: list the server's tools and call them. + + Made by :func:`run_connection_session`. Every call goes over the same + session, so a sync of hundreds of items signs in once. + """ + + def __init__(self, client: Any): + self._client = client + self._schemas: Optional[dict[str, dict]] = None + + async def input_schema(self, name: str) -> Optional[dict]: + """The input schema of tool ``name`` (``{}`` when it has none), or None when there is no such tool.""" + if self._schemas is None: + tools = await self._client.list_tools() + self._schemas = {tool.name: (getattr(tool, "inputSchema", None) or {}) for tool in tools} + return self._schemas.get(name) + + async def call(self, name: str, arguments: dict) -> Any: + """Call tool ``name`` and return what it answered. + + Returns: + Its structured content when it sends some, else its text parsed + as JSON, else the text itself. + + Raises: + MCPToolError: The tool answered with an error. + service.TransientConnectionError: The error says to try again later. + """ + result = await self._client.call_tool(name, arguments, raise_on_error=False) + if getattr(result, "is_error", False): + message = _result_text(result) or "the tool failed" + if any(word in message.lower() for word in _TRANSIENT_WORDS): + raise service.TransientConnectionError(f"{name}: {message}") + raise MCPToolError(f"{name}: {message}") + return tool_result_data(result) + + +def _result_text(result: Any) -> str: + blocks = getattr(result, "content", None) or [] + return "\n".join(str(block.text) for block in blocks if getattr(block, "text", None)).strip() + + +def tool_result_data(result: Any) -> Any: + """What an MCP tool answered: its structured content, its text parsed as JSON, or the text.""" + structured = getattr(result, "structured_content", None) + if isinstance(structured, dict) and structured: + # A tool without an output schema has other values wrapped as ``{"result": ...}``. + return structured["result"] if set(structured) == {"result"} else structured + text = _result_text(result) + try: + return json.loads(text) + except ValueError: + return text + + +def _client_for(connection: dict, server_url: str, timeout: float) -> Any: + """An MCP client for ``server_url`` that signs in with ``connection``'s stored OAuth tokens. + + The tokens are renewed with their refresh token when they expire; when + that fails the client raises ``MCPReauthorizationRequired`` instead of + starting an interactive sign-in. + """ + from fastmcp import Client + + from docsgpt.agents.tools.mcp_tool import MCPTool, NonInteractiveOAuth + + tool = MCPTool( + {"server_url": server_url, "auth_type": "oauth", "transport_type": "http", "timeout": timeout}, + connection["user_id"], + ) + auth = NonInteractiveOAuth( + mcp_url=tool.server_url, + redirect_uri=tool.redirect_uri, + user_id=connection["user_id"], + connection_id=str(connection["id"]), + ) + return Client(tool._create_transport(), auth=auth, timeout=timeout) + + +def _errors(exc: BaseException): + """``exc`` and every exception it wraps: causes, contexts and group members.""" + seen: set[int] = set() + pending: list[Optional[BaseException]] = [exc] + while pending: + current = pending.pop() + if current is None or id(current) in seen: + continue + seen.add(id(current)) + yield current + pending.extend(getattr(current, "exceptions", None) or ()) + pending.extend((current.__cause__, current.__context__)) + + +def _needs_sign_in(exc: BaseException) -> bool: + from docsgpt.agents.tools.mcp_tool import MCPReauthorizationRequired + + return any( + isinstance(error, MCPReauthorizationRequired) or "oauth session expired" in str(error).lower() + for error in _errors(exc) + ) + + +def _is_transient(exc: BaseException) -> bool: + import httpx + + return any( + isinstance(error, (ConnectionError, TimeoutError, httpx.TransportError)) for error in _errors(exc) + ) + + +def _run_coroutine(coroutine: Awaitable[T]) -> T: + """Run ``coroutine`` to the end, on a thread of its own when this one already runs an event loop.""" + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(coroutine) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(asyncio.run, coroutine).result() + + +def run_connection_session( + connection: dict, + server_url: str, + work: Callable[[MCPSession], Awaitable[T]], + *, + timeout: float = 60.0, +) -> T: + """Open one MCP session signed in with ``connection`` and run ``work`` in it. + + Args: + connection: An MCP OAuth connection (a Linear sign-in). + server_url: The MCP endpoint. It must be the server the connection + signed in to, so its tokens never go anywhere else. + work: Async function given the open :class:`MCPSession`. + timeout: Seconds allowed for each request. + + Returns: + What ``work`` returned. + + Raises: + ValueError: ``server_url`` is not the connection's server. + service.ConnectionUnavailable: The connection needs reconnecting, + also when its sign-in expired and could not be renewed (it is + then flagged, which pauses its sources). + service.TransientConnectionError: The server could not be reached, + or asked to slow down. + MCPToolError: A tool answered with an error. + """ + connection_id = str(connection["id"]) + if catalog.base_url(connection.get("server_url")) != catalog.base_url(server_url): + raise ValueError("A connection's sign-in only goes to its own server") + if service.normalize_status(connection) != service.STATUS_CONNECTED: + raise service.ConnectionUnavailable("Reconnect to continue", connection_id=connection_id) + + async def session_work() -> T: + async with _client_for(connection, server_url, timeout) as client: + return await work(MCPSession(client)) + + try: + return _run_coroutine(session_work()) + except (MCPToolError, service.TransientConnectionError, service.ConnectionUnavailable): + raise + except Exception as exc: + if _needs_sign_in(exc): + service.mark_reconnect_needed(connection_id, SIGN_IN_EXPIRED) + raise service.ConnectionUnavailable(SIGN_IN_EXPIRED, connection_id=connection_id) from exc + if _is_transient(exc): + raise service.TransientConnectionError(f"The MCP server did not answer: {type(exc).__name__}") from exc + raise diff --git a/tests/connectors/test_mcp_session.py b/tests/connectors/test_mcp_session.py new file mode 100644 index 00000000..25387876 --- /dev/null +++ b/tests/connectors/test_mcp_session.py @@ -0,0 +1,211 @@ +"""Calling an MCP server's tools from the server side, signed in with a connection.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +import docsgpt.api.user # noqa: F401 (loads mcp_tool without the circular import) +from docsgpt.agents.tools.mcp_tool import MCPReauthorizationRequired +from docsgpt.connectors import mcp, service + +LINEAR = "https://mcp.linear.app/mcp" + + +def _connection(status="connected", server_url="https://mcp.linear.app") -> dict: + return { + "id": "c1", "user_id": "alice", "provider": "mcp:https://mcp.linear.app", "connector_key": "mcp:linear", + "auth_kind": "mcp_oauth", "status": status, "server_url": server_url, "encrypted_credentials": "v2:x", + } + + +def _text(payload) -> SimpleNamespace: + return SimpleNamespace(type="text", text=payload if isinstance(payload, str) else json.dumps(payload)) + + +def _result(*content, structured=None, error=False) -> SimpleNamespace: + return SimpleNamespace(content=list(content), structured_content=structured, is_error=error) + + +class FakeClient: + """An MCP client that answers from a table of tool results.""" + + def __init__(self, results=None, tools=(), enter_error=None): + self.results = results or {} + self.tools = [SimpleNamespace(name=name, inputSchema=schema) for name, schema in tools] + self.enter_error = enter_error + self.calls: list[tuple[str, dict]] = [] + self.listed = 0 + self.open = False + + async def __aenter__(self): + if self.enter_error: + raise self.enter_error + self.open = True + return self + + async def __aexit__(self, *exc): + self.open = False + + async def list_tools(self): + self.listed += 1 + return self.tools + + async def call_tool(self, name, arguments=None, raise_on_error=True): + assert self.open + self.calls.append((name, dict(arguments or {}))) + result = self.results[name] + if isinstance(result, BaseException): + raise result + return result + + +def _run(client, work, connection=None): + with patch.object(mcp, "_client_for", return_value=client) as make: + value = mcp.run_connection_session(connection or _connection(), LINEAR, work) + return value, make + + +class TestRunConnectionSession: + def test_calls_share_one_session_signed_in_with_the_connection(self): + client = FakeClient({"list_teams": _result(_text({"teams": [{"id": "t1"}]}))}) + + async def work(session): + first = await session.call("list_teams", {"limit": 5}) + second = await session.call("list_teams", {}) + return first, second + + (first, second), make = _run(client, work) + assert first == {"teams": [{"id": "t1"}]} == second + assert client.calls == [("list_teams", {"limit": 5}), ("list_teams", {})] + assert make.call_args.args[0]["id"] == "c1" + assert make.call_args.args[1] == LINEAR + + def test_structured_content_is_preferred_to_text(self): + client = FakeClient({"get_issue": _result(_text("prose"), structured={"id": "ENG-1"})}) + + async def work(session): + return await session.call("get_issue", {"id": "ENG-1"}) + + assert _run(client, work)[0] == {"id": "ENG-1"} + + def test_text_that_is_not_json_comes_back_as_text(self): + client = FakeClient({"search": _result(_text("No results"))}) + + async def work(session): + return await session.call("search", {}) + + assert _run(client, work)[0] == "No results" + + def test_a_tool_error_names_the_tool(self): + client = FakeClient({"get_issue": _result(_text("Entity not found"), error=True)}) + + async def work(session): + return await session.call("get_issue", {"id": "X"}) + + with pytest.raises(mcp.MCPToolError, match="get_issue: Entity not found"): + _run(client, work) + + def test_a_rate_limit_is_worth_retrying(self): + client = FakeClient({"list_issues": _result(_text("Rate limit exceeded, retry later"), error=True)}) + + async def work(session): + return await session.call("list_issues", {}) + + with pytest.raises(service.TransientConnectionError): + _run(client, work) + + def test_input_schemas_are_listed_once(self): + client = FakeClient(tools=[("list_issues", {"properties": {"team": {}}}), ("get_issue", None)]) + + async def work(session): + return ( + await session.input_schema("list_issues"), + await session.input_schema("get_issue"), + await session.input_schema("missing"), + ) + + assert _run(client, work)[0] == ({"properties": {"team": {}}}, {}, None) + assert client.listed == 1 + + @pytest.mark.parametrize("error", [ + MCPReauthorizationRequired("OAuth session expired"), + RuntimeError("Client failed to connect: OAuth session expired — please re-authorize"), + ExceptionGroup("task group", [MCPReauthorizationRequired("OAuth session expired")]), + ]) + def test_a_lost_sign_in_flags_the_connection(self, error): + client = FakeClient(enter_error=error) + + async def work(session): # pragma: no cover - never reached + return None + + with patch.object(service, "mark_reconnect_needed") as flag, pytest.raises(service.ConnectionUnavailable): + _run(client, work) + flag.assert_called_once() + assert flag.call_args.args[0] == "c1" + + def test_a_lost_sign_in_during_a_call_flags_the_connection(self): + client = FakeClient({"list_issues": MCPReauthorizationRequired("OAuth session expired")}) + + async def work(session): + return await session.call("list_issues", {}) + + with patch.object(service, "mark_reconnect_needed") as flag, pytest.raises(service.ConnectionUnavailable): + _run(client, work) + flag.assert_called_once() + + def test_a_connection_that_needs_reconnecting_is_not_used(self): + client = FakeClient() + + async def work(session): # pragma: no cover - never reached + return None + + with pytest.raises(service.ConnectionUnavailable): + _run(client, work, _connection(status="reconnect_needed")) + assert client.calls == [] + + def test_tokens_never_go_to_another_server(self): + async def work(session): # pragma: no cover - never reached + return None + + with pytest.raises(ValueError): + _run(FakeClient(), work, _connection(server_url="https://mcp.notion.com")) + + def test_network_trouble_is_worth_retrying(self): + client = FakeClient({"list_issues": ConnectionError("reset")}) + + async def work(session): + return await session.call("list_issues", {}) + + with patch.object(service, "mark_reconnect_needed") as flag, pytest.raises(service.TransientConnectionError): + _run(client, work) + flag.assert_not_called() + + def test_the_client_signs_in_with_the_connections_stored_tokens(self, monkeypatch): + from fastmcp import Client + + import docsgpt.agents.tools.mcp_tool as mcp_tool + + monkeypatch.setattr(mcp_tool, "validate_url", lambda url, **kw: url) + client = mcp._client_for(_connection(), LINEAR, 30) + assert isinstance(client, Client) + auth = client.transport.auth + assert isinstance(auth, mcp_tool.NonInteractiveOAuth) + assert auth.context.storage.connection_id == "c1" + assert auth.context.storage.user_id == "alice" + + def test_runs_from_inside_a_running_event_loop(self): + import asyncio + + client = FakeClient({"list_teams": _result(_text([]))}) + + async def work(session): + return await session.call("list_teams", {}) + + async def caller(): + return _run(client, work)[0] + + assert asyncio.run(caller()) == []