mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-19 22:14:31 +00:00
Connections own their credentials
Migration 0038 moves every stored secret (OAuth tokens, MCP OAuth tokens and client registrations, API keys) into the connection's encrypted envelope, links API-key tools to one connection per distinct credential, allows several accounts per provider, and adds credential_mode to sources and tools. OAuth MCP tools keep resolving each member's own token, as they did before. docsgpt.connectors.service is now the only reader of OAuth tokens: get_valid_token_info refreshes under a row lock and persists rotated refresh tokens, and a revoked grant flags the connection, pauses its sources and notifies the owner. Loaders build from a connection (BaseConnectorLoader.from_connection), so scheduled sync covers Drive, SharePoint and Confluence sources with no browser. S3 and Reddit keys stay on the connection instead of in remote_data. New endpoints: POST /api/connections, /setup, /reconnect, /picker-token, /claim, DELETE /api/connections/<id>, per-action permissions and MCP refresh-tools. Upload, file listing, sync and validate-session take a connection_id; session tokens keep working for this release. The tool executor reads credentials from the resolved connection (owner or member mode) and pauses on a Connect card when a connection needs signing in. docsgpt connectors reencrypt rewrites stored credentials after a key rotation.
This commit is contained in:
1 parent
667d4bbab0
commit
3511004f50
54 files changed
+4007
-754
No files matched your search
@@ -485,6 +485,9 @@ class ToolExecutor:
|
||||
self._tool_to_name: Dict[Tuple[str, str], str] = {}
|
||||
# Filled by the LLMHandler.handle_tool_calls headless loop.
|
||||
self.headless_denials: List[Dict] = []
|
||||
# Per-turn connection resolution for connection-backed tools, keyed
|
||||
# by tool row id, so check_pause and execute share one lookup.
|
||||
self._connections: Dict[str, Any] = {}
|
||||
|
||||
def get_tools(self) -> Dict[str, Dict]:
|
||||
"""Load tool configs from DB based on user context.
|
||||
@@ -785,6 +788,33 @@ class ToolExecutor:
|
||||
return result
|
||||
return resolve_tool_result(result, decision)
|
||||
|
||||
def _resolve_connection(self, tool_data: Dict):
|
||||
"""The connection a connection-backed tool runs with this turn, or None."""
|
||||
if not tool_data.get("connection_id") or tool_data.get("client_side"):
|
||||
return None
|
||||
key = str(tool_data.get("id") or tool_data.get("connection_id"))
|
||||
if key not in self._connections:
|
||||
from docsgpt.connectors.resolve import resolve_connection
|
||||
|
||||
try:
|
||||
self._connections[key] = resolve_connection(tool_data, self.user)
|
||||
except Exception:
|
||||
logger.exception("connection resolution failed for tool %s", key)
|
||||
self._connections[key] = None
|
||||
return self._connections[key]
|
||||
|
||||
@staticmethod
|
||||
def _connection_payload(resolved) -> Dict:
|
||||
"""What the chat's Connect card needs; never an account or a secret."""
|
||||
return {
|
||||
"connector_key": resolved.connector_key,
|
||||
"connector_name": resolved.connector_name,
|
||||
"status": (
|
||||
"missing" if resolved.row is None
|
||||
else (resolved.row.get("status") or "reconnect_needed")
|
||||
),
|
||||
}
|
||||
|
||||
def check_pause(self, tools_dict: Dict, call, llm_class_name: str) -> Optional[Dict]:
|
||||
"""Return a pending-action dict (approval / client / headless_denied) or None.
|
||||
|
||||
@@ -831,6 +861,41 @@ class ToolExecutor:
|
||||
"thought_signature": getattr(call, "thought_signature", None),
|
||||
}
|
||||
|
||||
# A tool whose connection needs signing in pauses on a Connect card
|
||||
# (the approval card's connection variant) instead of failing; the
|
||||
# user connects, then continues, and the pending call resumes.
|
||||
resolved = self._resolve_connection(tool_data)
|
||||
if resolved is not None and not resolved.available:
|
||||
if self.headless:
|
||||
return {
|
||||
"call_id": call_id,
|
||||
"name": llm_name,
|
||||
"tool_name": tool_data.get("name", "unknown"),
|
||||
"tool_id": tool_id,
|
||||
"action_name": action_name,
|
||||
"llm_name": llm_name,
|
||||
"arguments": arguments,
|
||||
"pause_type": "headless_denied",
|
||||
"deny_reason": (
|
||||
f"{resolved.connector_name or 'This service'} needs to be connected "
|
||||
"before this tool can run."
|
||||
),
|
||||
"error_type": "connection_required",
|
||||
"thought_signature": getattr(call, "thought_signature", None),
|
||||
}
|
||||
return {
|
||||
"call_id": call_id,
|
||||
"name": llm_name,
|
||||
"tool_name": tool_data.get("name", "unknown"),
|
||||
"tool_id": tool_id,
|
||||
"action_name": action_name,
|
||||
"llm_name": llm_name,
|
||||
"arguments": arguments,
|
||||
"pause_type": "awaiting_approval",
|
||||
"connection_required": self._connection_payload(resolved),
|
||||
"thought_signature": getattr(call, "thought_signature", None),
|
||||
}
|
||||
|
||||
# Approval required
|
||||
if tool_data["name"] == "api_tool":
|
||||
action_data = tool_data.get("config", {}).get("actions", {}).get(action_name, {})
|
||||
@@ -1285,16 +1350,24 @@ class ToolExecutor:
|
||||
target_dict[param] = value
|
||||
|
||||
# Load tool (with caching)
|
||||
tool = self._get_or_load_tool(
|
||||
tool_data,
|
||||
tool_id,
|
||||
action_name,
|
||||
headers=headers,
|
||||
query_params=query_params,
|
||||
)
|
||||
from docsgpt.connectors.service import ConnectionUnavailable
|
||||
|
||||
connection_error = None
|
||||
try:
|
||||
tool = self._get_or_load_tool(
|
||||
tool_data,
|
||||
tool_id,
|
||||
action_name,
|
||||
headers=headers,
|
||||
query_params=query_params,
|
||||
)
|
||||
except ConnectionUnavailable as exc:
|
||||
tool, connection_error = None, str(exc)
|
||||
|
||||
if tool is None:
|
||||
error_message = (
|
||||
error_message = connection_error and (
|
||||
f"{connection_error}. Ask the user to connect it in Settings > Connectors, then try again."
|
||||
) or (
|
||||
f"Failed to load tool '{tool_data.get('name')}' (tool_id key={tool_id}): missing 'id' on tool row."
|
||||
)
|
||||
logger.error(
|
||||
@@ -1476,7 +1549,10 @@ class ToolExecutor:
|
||||
# silently decrypt-failing. Falls back to self.user for the
|
||||
# agentless path where the tool row carries no user_id.
|
||||
tool_owner = tool_data.get("user_id") or self.user
|
||||
if tool_config.get("encrypted_credentials") and tool_owner:
|
||||
resolved = self._resolve_connection(tool_data)
|
||||
if resolved is not None:
|
||||
self._apply_connection(tool_data, tool_id, tool_config, resolved)
|
||||
elif tool_config.get("encrypted_credentials") and tool_owner:
|
||||
if tool_owner != self.user:
|
||||
# Credential delegation: the invoker is running a shared
|
||||
# tool with the owner's secrets. Audit it (the agent-run
|
||||
@@ -1540,6 +1616,36 @@ class ToolExecutor:
|
||||
|
||||
return tool
|
||||
|
||||
def _apply_connection(self, tool_data: Dict, tool_id: str, tool_config: Dict, resolved) -> None:
|
||||
"""Merge a connection's credentials into ``tool_config``.
|
||||
|
||||
Raises:
|
||||
ConnectionUnavailable: The connection needs signing in again.
|
||||
"""
|
||||
from docsgpt.connectors import service
|
||||
from docsgpt.connectors.resolve import audit_delegation
|
||||
|
||||
if not resolved.available or resolved.row is None:
|
||||
raise service.ConnectionUnavailable(
|
||||
f"{resolved.connector_name or 'This service'} needs to be connected",
|
||||
connection_id=resolved.connection_id,
|
||||
)
|
||||
audit_delegation(
|
||||
resolved,
|
||||
invoker=self.user,
|
||||
resource_type="tool",
|
||||
resource_id=str(tool_data.get("id") or tool_id),
|
||||
agent_id=self.agent_id,
|
||||
)
|
||||
tool_config.pop("encrypted_credentials", None)
|
||||
if (resolved.row.get("auth_kind") or "") in ("api_key", "none"):
|
||||
credentials = service.get_credentials(resolved.row)
|
||||
tool_config.update(credentials)
|
||||
tool_config["auth_credentials"] = credentials
|
||||
if tool_data.get("name") == "mcp_tool":
|
||||
# MCP OAuth tokens are read by connection id inside the tool.
|
||||
tool_config["connection_id"] = resolved.connection_id
|
||||
|
||||
# Keys the client needs that are not part of the fixed shape below. They are
|
||||
# small and optional, and are copied only when present so an ordinary tool
|
||||
# call does not grow null columns in every persisted row.
|
||||
|
||||
@@ -136,6 +136,7 @@ class BraveSearchTool(Tool):
|
||||
return [
|
||||
{
|
||||
"name": "brave_web_search",
|
||||
"access": "read",
|
||||
"description": (
|
||||
"Search the web with Brave Search. Returns result titles, "
|
||||
"URLs, and snippets. Use it for current events or "
|
||||
@@ -163,6 +164,7 @@ class BraveSearchTool(Tool):
|
||||
},
|
||||
{
|
||||
"name": "brave_image_search",
|
||||
"access": "read",
|
||||
"description": (
|
||||
"Search for images with Brave Search. Returns image "
|
||||
"titles, page URLs, and thumbnail URLs."
|
||||
|
||||
@@ -31,6 +31,19 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
_mcp_clients_cache = {}
|
||||
|
||||
_ANNOTATION_HINTS = ("readOnlyHint", "destructiveHint", "idempotentHint", "openWorldHint")
|
||||
|
||||
|
||||
def _annotation_hints(annotations: Any) -> Dict[str, bool]:
|
||||
"""The boolean MCP tool annotation hints, from a model or a dict."""
|
||||
if annotations is None:
|
||||
return {}
|
||||
if hasattr(annotations, "model_dump"):
|
||||
annotations = annotations.model_dump()
|
||||
if not isinstance(annotations, dict):
|
||||
return {}
|
||||
return {k: annotations[k] for k in _ANNOTATION_HINTS if isinstance(annotations.get(k), bool)}
|
||||
|
||||
|
||||
class MCPTool(Tool):
|
||||
"""
|
||||
@@ -38,6 +51,8 @@ class MCPTool(Tool):
|
||||
Connect to remote Model Context Protocol (MCP) servers to access dynamic tools and resources.
|
||||
"""
|
||||
|
||||
connection_id: Optional[str] = None
|
||||
|
||||
def __init__(self, config: Dict[str, Any], user_id: Optional[str] = None):
|
||||
"""
|
||||
Initialize the MCP Tool with configuration.
|
||||
@@ -76,6 +91,10 @@ class MCPTool(Tool):
|
||||
self.oauth_scopes = config.get("oauth_scopes", [])
|
||||
self.oauth_task_id = config.get("oauth_task_id", None)
|
||||
self.oauth_client_name = config.get("oauth_client_name", "DocsGPT-MCP")
|
||||
# The connection whose OAuth tokens this tool uses. Set by the tool
|
||||
# executor for connection-backed tools, so a shared tool in ``owner``
|
||||
# mode signs in with its owner's account rather than the invoker's.
|
||||
self.connection_id = config.get("connection_id")
|
||||
self.redirect_uri = self._resolve_redirect_uri(config.get("redirect_uri"))
|
||||
# Pulled out of ``config`` (rather than left in ``self.config``)
|
||||
# because it is a callable supplied by the OAuth worker — not
|
||||
@@ -125,7 +144,9 @@ class MCPTool(Tool):
|
||||
auth_key = ""
|
||||
if self.auth_type == "oauth":
|
||||
scopes_str = ",".join(self.oauth_scopes) if self.oauth_scopes else "none"
|
||||
oauth_identity = self.user_id or self.oauth_task_id or "anonymous"
|
||||
# A connection-backed tool shares a client only with calls that
|
||||
# use the same connection's tokens.
|
||||
oauth_identity = self.connection_id or self.user_id or self.oauth_task_id or "anonymous"
|
||||
auth_key = (
|
||||
f"oauth:{oauth_identity}:{self.oauth_client_name}:{scopes_str}:{self.redirect_uri}"
|
||||
)
|
||||
@@ -165,6 +186,7 @@ class MCPTool(Tool):
|
||||
redis_client=redis_client,
|
||||
redirect_uri=self.redirect_uri,
|
||||
user_id=self.user_id,
|
||||
connection_id=self.connection_id,
|
||||
)
|
||||
else:
|
||||
auth = DocsGPTOAuth(
|
||||
@@ -175,6 +197,7 @@ class MCPTool(Tool):
|
||||
task_id=self.oauth_task_id,
|
||||
user_id=self.user_id,
|
||||
redirect_publish=self.oauth_redirect_publish,
|
||||
connection_id=self.connection_id,
|
||||
)
|
||||
elif self.auth_type == "bearer":
|
||||
token = self.auth_credentials.get(
|
||||
@@ -245,6 +268,9 @@ class MCPTool(Tool):
|
||||
}
|
||||
if hasattr(tool, "inputSchema"):
|
||||
tool_dict["inputSchema"] = tool.inputSchema
|
||||
annotations = _annotation_hints(getattr(tool, "annotations", None))
|
||||
if annotations:
|
||||
tool_dict["annotations"] = annotations
|
||||
tools_dict.append(tool_dict)
|
||||
elif isinstance(tool, dict):
|
||||
tools_dict.append(tool)
|
||||
@@ -493,7 +519,7 @@ class MCPTool(Tool):
|
||||
|
||||
def _test_oauth_connection(self) -> Dict:
|
||||
storage = DBTokenStorage(
|
||||
server_url=self.server_url, user_id=self.user_id,
|
||||
server_url=self.server_url, user_id=self.user_id, connection_id=self.connection_id,
|
||||
)
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
@@ -580,6 +606,11 @@ class MCPTool(Tool):
|
||||
"description": tool.get("description", ""),
|
||||
"parameters": parameters_schema,
|
||||
}
|
||||
# ``readOnlyHint`` / ``destructiveHint`` decide whether the action
|
||||
# is a read (always allowed) or a write (needs approval).
|
||||
annotations = _annotation_hints(tool.get("annotations"))
|
||||
if annotations:
|
||||
action["annotations"] = annotations
|
||||
actions.append(action)
|
||||
return actions
|
||||
|
||||
@@ -688,6 +719,7 @@ class DocsGPTOAuth(OAuthClientProvider):
|
||||
additional_client_metadata: dict[str, Any] | None = None,
|
||||
skip_redirect_validation: bool = False,
|
||||
redirect_publish=None,
|
||||
connection_id: Optional[str] = None,
|
||||
):
|
||||
self.redirect_uri = redirect_uri
|
||||
self.redis_client = redis_client
|
||||
@@ -717,6 +749,7 @@ class DocsGPTOAuth(OAuthClientProvider):
|
||||
server_url=self.server_base_url,
|
||||
user_id=self.user_id,
|
||||
expected_redirect_uri=None if skip_redirect_validation else redirect_uri,
|
||||
connection_id=connection_id,
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
@@ -836,15 +869,26 @@ class NonInteractiveOAuth(DocsGPTOAuth):
|
||||
|
||||
|
||||
class DBTokenStorage(TokenStorage):
|
||||
"""MCP OAuth tokens and client registration, kept encrypted on the connection.
|
||||
|
||||
Reads and writes go through ``docsgpt.connectors.service``, which stores
|
||||
them in the connection's owner-bound ``encrypted_credentials``. A tool
|
||||
that runs with a specific connection (``owner`` mode on a shared tool)
|
||||
passes its ``connection_id``; otherwise the invoking user's connection
|
||||
for the server's base URL is used.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_url: str,
|
||||
user_id: str,
|
||||
expected_redirect_uri: Optional[str] = None,
|
||||
connection_id: Optional[str] = None,
|
||||
):
|
||||
self.server_url = server_url
|
||||
self.user_id = user_id
|
||||
self.expected_redirect_uri = expected_redirect_uri
|
||||
self.connection_id = connection_id
|
||||
|
||||
@staticmethod
|
||||
def get_base_url(url: str) -> str:
|
||||
@@ -855,26 +899,12 @@ class DBTokenStorage(TokenStorage):
|
||||
return f"mcp:{self.get_base_url(self.server_url)}"
|
||||
|
||||
def _fetch_session_data(self) -> dict:
|
||||
"""Read the JSONB ``session_data`` blob for this MCP server row."""
|
||||
from docsgpt.storage.db.repositories.connector_sessions import (
|
||||
ConnectorSessionsRepository,
|
||||
)
|
||||
from docsgpt.storage.db.session import db_readonly
|
||||
"""The decrypted ``tokens`` / ``client_info`` for this MCP server."""
|
||||
from docsgpt.connectors import service
|
||||
|
||||
base_url = self.get_base_url(self.server_url)
|
||||
with db_readonly() as conn:
|
||||
row = ConnectorSessionsRepository(conn).get_by_user_and_server_url(
|
||||
self.user_id, base_url,
|
||||
)
|
||||
if not row:
|
||||
return {}
|
||||
data = row.get("session_data") or {}
|
||||
if isinstance(data, str):
|
||||
try:
|
||||
data = json.loads(data)
|
||||
except ValueError:
|
||||
return {}
|
||||
return data if isinstance(data, dict) else {}
|
||||
return service.read_mcp_secrets(
|
||||
self.user_id, self.get_base_url(self.server_url), self.connection_id,
|
||||
)
|
||||
|
||||
async def get_tokens(self) -> OAuthToken | None:
|
||||
data = await asyncio.to_thread(self._fetch_session_data)
|
||||
@@ -887,22 +917,19 @@ class DBTokenStorage(TokenStorage):
|
||||
return None
|
||||
|
||||
def _merge(self, patch: dict) -> None:
|
||||
"""Shallow-merge ``patch`` into this row's ``session_data``.
|
||||
"""Merge ``patch`` into the connection's secrets; ``None`` drops a key."""
|
||||
from docsgpt.connectors import service
|
||||
|
||||
Threads ``server_url`` through to the repository so it lands in
|
||||
the scalar column — ``get_by_user_and_server_url`` needs that to
|
||||
resolve the row (``NULL = 'https://...'`` is UNKNOWN in SQL).
|
||||
"""
|
||||
from docsgpt.storage.db.repositories.connector_sessions import (
|
||||
ConnectorSessionsRepository,
|
||||
status = None
|
||||
if patch.get("tokens"):
|
||||
status = service.STATUS_CONNECTED
|
||||
service.update_mcp_secrets(
|
||||
self.user_id,
|
||||
self.get_base_url(self.server_url),
|
||||
patch,
|
||||
connection_id=self.connection_id,
|
||||
status=status,
|
||||
)
|
||||
from docsgpt.storage.db.session import db_session
|
||||
|
||||
base_url = self.get_base_url(self.server_url)
|
||||
with db_session() as conn:
|
||||
ConnectorSessionsRepository(conn).merge_session_data(
|
||||
self.user_id, self._pg_provider(), base_url, patch,
|
||||
)
|
||||
|
||||
def _delete(self) -> None:
|
||||
from docsgpt.storage.db.repositories.connector_sessions import (
|
||||
|
||||
@@ -89,6 +89,7 @@ class NtfyTool(Tool):
|
||||
return [
|
||||
{
|
||||
"name": "ntfy_send_message",
|
||||
"access": "write",
|
||||
"description": (
|
||||
"Send a push notification to an ntfy topic on the "
|
||||
"configured server. Provide the message text; title and "
|
||||
|
||||
@@ -136,6 +136,7 @@ class PostgresTool(Tool):
|
||||
return [
|
||||
{
|
||||
"name": "postgres_execute_sql",
|
||||
"access": "write",
|
||||
"description": "Execute an SQL query against the PostgreSQL database and return the results. Use this tool to interact with the database, e.g., retrieve specific data or perform updates. Only SELECT queries will return data, other queries will return execution status.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
@@ -151,6 +152,7 @@ class PostgresTool(Tool):
|
||||
},
|
||||
{
|
||||
"name": "postgres_get_schema",
|
||||
"access": "read",
|
||||
"description": "Retrieve the schema of the PostgreSQL database, including tables and their columns. Use this to understand the database structure before executing queries. db_name is 'default' if not provided.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
|
||||
@@ -45,6 +45,7 @@ class TelegramTool(Tool):
|
||||
return [
|
||||
{
|
||||
"name": "telegram_send_message",
|
||||
"access": "write",
|
||||
"description": (
|
||||
"Send a text message to the configured Telegram chat via "
|
||||
"the bot. Compose the final message text before sending."
|
||||
@@ -67,6 +68,7 @@ class TelegramTool(Tool):
|
||||
},
|
||||
{
|
||||
"name": "telegram_send_image",
|
||||
"access": "write",
|
||||
"description": (
|
||||
"Send an image to the configured Telegram chat. Requires "
|
||||
"a publicly accessible image URL."
|
||||
|
||||
Reference in new issue
Block a user