mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 16:15:14 +00:00
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.
This commit is contained in:
arc53-machine committed
1 parent
6eda4a6738
commit
805e65eeb4
2 files changed
+409
-2
No files matched your search
+198
-2
@@ -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
|
||||
@@ -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()) == []
|
||||
Reference in new issue
Block a user