Call an MCP server's tools from the server with a connection's sign-in

run_connection_session opens one MCP session signed in with an MCP OAuth
connection's stored tokens and runs a batch of tool calls in it, so a
sync makes hundreds of calls over one sign-in. The tokens only go to the
server the connection signed in to. A sign-in that expired and cannot be
renewed flags the connection for reconnecting, which pauses its sources;
network trouble and rate limits are reported as worth retrying.

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