mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 18:15:02 +00:00
arc53-machine
805e65eeb4
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.
383 lines
16 KiB
Python
383 lines
16 KiB
Python
"""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
|
|
|
|
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
|
|
from docsgpt.connectors.catalog import ConnectorDefinition
|
|
from docsgpt.connectors.permissions import ACCESS_READ, ACCESS_WRITE, apply_default_permissions
|
|
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
|
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)."""
|
|
|
|
|
|
def _builtin_definition(connection: dict) -> Optional[ConnectorDefinition]:
|
|
"""The connection's built-in connector when its tool is that service's MCP server."""
|
|
definition = catalog.get_definition(catalog.connector_key_for_row(connection))
|
|
return definition if service.builtin_mcp_config(definition) is not None else None
|
|
|
|
|
|
def _write_endpoint_access(action: dict) -> dict:
|
|
"""Stamp an action of a write endpoint: a read only when the server marks it read-only.
|
|
|
|
GitHub marks its read tools, so a name is no evidence there: from the name
|
|
alone ``mark_all_notifications_read`` would pass for a read.
|
|
"""
|
|
annotations = action.get("annotations") if isinstance(action.get("annotations"), dict) else {}
|
|
return {**action, "access": ACCESS_READ if annotations.get("readOnlyHint") is True else ACCESS_WRITE}
|
|
|
|
|
|
def _discover(user_id: str, connection: dict, tool: dict) -> list[dict]:
|
|
from docsgpt.agents.tools.mcp_tool import MCPTool
|
|
|
|
config = {k: v for k, v in (tool.get("config") or {}).items() if k != "encrypted_credentials"}
|
|
if (connection.get("auth_kind") or "") in ("api_key", "none", "oauth"):
|
|
# Pasted keys, or the access token of a built-in OAuth sign-in (GitHub's App).
|
|
config["auth_credentials"] = service.access_credentials(connection)
|
|
elif service.normalize_status(connection) != service.STATUS_CONNECTED:
|
|
raise service.ConnectionUnavailable("Reconnect to continue", connection_id=str(connection["id"]))
|
|
config["connection_id"] = str(connection["id"])
|
|
config["query_mode"] = True
|
|
mcp_tool = MCPTool(config, user_id)
|
|
mcp_tool.discover_tools()
|
|
return mcp_tool.get_actions_metadata()
|
|
|
|
|
|
def _classified(connection: dict, config: dict, actions: list[dict]) -> list[dict]:
|
|
"""``actions`` as read from ``config``'s server, stamped strictly when it is a write endpoint."""
|
|
definition = _builtin_definition(connection)
|
|
if definition is None or not definition.mcp_write_url or config.get("server_url") != definition.mcp_write_url:
|
|
return actions
|
|
return [_write_endpoint_access(action) for action in actions]
|
|
|
|
|
|
def discover_builtin_actions(user_id: str, connection: dict, writes: bool = False) -> list[dict]:
|
|
"""The actions of a built-in connector's MCP server (GitHub's), read with the connection.
|
|
|
|
Args:
|
|
user_id: The connection's owner.
|
|
connection: The connection row.
|
|
writes: Read the endpoint that also offers write actions; the caller
|
|
checked that an admin allows them.
|
|
|
|
Raises:
|
|
ValueError: The connector has no MCP server of its own.
|
|
service.ConnectionUnavailable: The connection needs reconnecting.
|
|
Exception: The server could not be reached or listed.
|
|
"""
|
|
config = service.builtin_mcp_config(_builtin_definition(connection), writes=writes)
|
|
if config is None:
|
|
raise ValueError("This connector has no MCP server")
|
|
return _classified(connection, config, _discover(user_id, connection, {"config": config}))
|
|
|
|
|
|
def _rediscover(user_id: str, connection: dict, tool: dict, config: dict) -> tuple[set, set, dict]:
|
|
"""Re-read one MCP tool's actions from ``config``'s server and store both.
|
|
|
|
Actions that still exist keep the permissions the user chose and the
|
|
fixed values of parameters they still have; new ones get the defaults
|
|
from their annotations (reads always allowed, writes need approval);
|
|
removed ones disappear.
|
|
|
|
Returns:
|
|
``(added, removed, tool)``: action names, and the serialized tool.
|
|
"""
|
|
discovered = _classified(connection, config, _discover(user_id, connection, {**tool, "config": config}))
|
|
fresh = apply_default_permissions("mcp_tool", service._transform_actions(discovered))
|
|
previous = {a.get("name"): a for a in (tool.get("actions") or []) if isinstance(a, dict)}
|
|
merged = []
|
|
for action in fresh:
|
|
old = previous.get(action.get("name"))
|
|
if old is not None:
|
|
action = {
|
|
**carry_pins(old, action),
|
|
"active": old.get("active", True),
|
|
"require_approval": bool(old.get("require_approval")),
|
|
}
|
|
merged.append(action)
|
|
names = {a.get("name") for a in merged}
|
|
fields = {"actions": merged}
|
|
if config != (tool.get("config") or {}):
|
|
fields["config"] = config
|
|
with db_session() as conn:
|
|
repo = UserToolsRepository(conn)
|
|
repo.update(str(tool["id"]), user_id, fields)
|
|
serialized = service.serialize_tool(repo.get_any(str(tool["id"]), user_id))
|
|
return names - set(previous), set(previous) - names, serialized
|
|
|
|
|
|
def refresh_mcp_tools(user_id: str, connection: dict) -> dict:
|
|
"""Re-read the actions of every MCP tool on a connection.
|
|
|
|
A built-in connector's tool is re-read from its own endpoint: the write
|
|
one while it is set up for writes and an admin allows them, the read-only
|
|
one otherwise, which it then keeps.
|
|
|
|
Returns:
|
|
``{"added": [...], "removed": [...], "tools": [...]}``.
|
|
"""
|
|
with db_readonly() as conn:
|
|
tools = [t for t in ConnectorSessionsRepository(conn).list_tools(str(connection["id"]))
|
|
if t.get("name") == "mcp_tool"]
|
|
policies = service.load_policies(conn) if tools else {}
|
|
definition = _builtin_definition(connection)
|
|
added: set[str] = set()
|
|
removed: set[str] = set()
|
|
refreshed = []
|
|
for tool in tools:
|
|
config = dict(tool.get("config") or {})
|
|
if definition is not None:
|
|
config["server_url"] = service.builtin_mcp_url(
|
|
definition, config.get("server_url"), service.writes_allowed(policies, definition.key),
|
|
)
|
|
tool_added, tool_removed, serialized = _rediscover(user_id, connection, tool, config)
|
|
added |= tool_added
|
|
removed |= tool_removed
|
|
refreshed.append(serialized)
|
|
return {"added": sorted(added), "removed": sorted(removed), "tools": refreshed}
|
|
|
|
|
|
def set_builtin_writes(user_id: str, connection: dict, allow: bool) -> dict:
|
|
"""Point a connection's built-in MCP tool (GitHub's) at its write or read-only endpoint.
|
|
|
|
The tool's actions are re-read from the new endpoint, keeping the
|
|
choices made for actions that exist on both.
|
|
|
|
Args:
|
|
user_id: The connection's owner.
|
|
connection: The connection row.
|
|
allow: Let agents make changes (the write endpoint) or only read.
|
|
|
|
Returns:
|
|
``{"added", "removed", "tools", "writes"}``.
|
|
|
|
Raises:
|
|
ValueError: The connector offers no write access.
|
|
NoMcpTool: The connection has no such tool yet.
|
|
service.WritesForbidden: ``allow`` while an admin forbids writes.
|
|
service.ConnectionUnavailable: The connection needs reconnecting.
|
|
"""
|
|
definition = _builtin_definition(connection)
|
|
if definition is None or not definition.mcp_write_url:
|
|
raise ValueError("This connector has no write access to turn on")
|
|
with db_readonly() as conn:
|
|
tools = [t for t in ConnectorSessionsRepository(conn).list_tools(str(connection["id"]))
|
|
if t.get("name") == "mcp_tool"]
|
|
allowed = service.writes_allowed(service.load_policies(conn), definition.key)
|
|
if not tools:
|
|
raise NoMcpTool("Add this connection's tools first")
|
|
if allow and not allowed:
|
|
raise service.WritesForbidden(f"An admin turned off changes through {definition.name}")
|
|
added: set[str] = set()
|
|
removed: set[str] = set()
|
|
refreshed = []
|
|
for tool in tools:
|
|
config = {**(tool.get("config") or {}), **service.builtin_mcp_config(definition, writes=allow)}
|
|
tool_added, tool_removed, serialized = _rediscover(user_id, connection, tool, config)
|
|
added |= tool_added
|
|
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
|