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.

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()) == []