mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 14:15:52 +00:00
arc53-machine
3cab8af753
Only removing a connection notes it on a kept tool. A config save through update_tool or update_tool_config carries the stored note and ignores the client's copy, so it can neither fake nor clear it; creating a tool, an MCP server save and an agent import start without one. The note is now written even for a connection with no catalog key, and a kept tool its owner gave credentials of its own since runs again.
1705 lines
68 KiB
Python
1705 lines
68 KiB
Python
"""Connection lifecycle: listing, status, and what each connection provides.
|
|
|
|
A connection is one ``connector_sessions`` row: a signed-in OAuth account, an
|
|
MCP server or a set of API credentials. Sources and tools point at it through
|
|
``connection_id``. Everything the API returns about a connection goes through
|
|
:func:`serialize_connection`, which never includes credentials.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sys
|
|
from typing import Any, Iterable, Optional
|
|
|
|
from docsgpt.connectors import catalog
|
|
from docsgpt.connectors.catalog import ConnectorDefinition
|
|
from docsgpt.storage.db.repositories.connector_sessions import (
|
|
ConnectorSessionsRepository,
|
|
owns_connector_session,
|
|
)
|
|
from docsgpt.storage.db.session import db_readonly, db_session
|
|
|
|
STATUS_CONNECTED = "connected"
|
|
STATUS_RECONNECT = "reconnect_needed"
|
|
STATUS_DISCONNECTED = "disconnected"
|
|
STATUS_ERROR = "error"
|
|
STATUS_PENDING = "pending"
|
|
|
|
# Worst first: a card shows the worst status among the user's accounts.
|
|
_STATUS_SEVERITY = {
|
|
STATUS_RECONNECT: 4,
|
|
STATUS_ERROR: 3,
|
|
STATUS_DISCONNECTED: 2,
|
|
STATUS_CONNECTED: 1,
|
|
STATUS_PENDING: 0,
|
|
}
|
|
|
|
|
|
def _json(value: Any) -> Any:
|
|
if isinstance(value, str):
|
|
try:
|
|
return json.loads(value)
|
|
except ValueError:
|
|
return None
|
|
return value
|
|
|
|
|
|
def has_credentials(row: dict) -> bool:
|
|
"""Whether a connection holds usable credentials (without decrypting them)."""
|
|
if row.get("encrypted_credentials"):
|
|
return True
|
|
return _has_plaintext_tokens(row)
|
|
|
|
|
|
def _has_plaintext_tokens(row: dict) -> bool:
|
|
"""Legacy plaintext tokens on a row not yet converted by migration 0040."""
|
|
token_info = _json(row.get("token_info")) or {}
|
|
if isinstance(token_info, dict) and token_info.get("access_token"):
|
|
return True
|
|
session_data = _json(row.get("session_data")) or {}
|
|
tokens = session_data.get("tokens") if isinstance(session_data, dict) else None
|
|
return bool(isinstance(tokens, dict) and tokens.get("access_token"))
|
|
|
|
|
|
def normalize_status(row: dict) -> str:
|
|
"""Map a row's stored status onto the connection status set.
|
|
|
|
Rows written before ``0040_connections`` use ``authorized`` for a
|
|
finished OAuth sign-in, and MCP rows carry no status at all; both are
|
|
read from whether the row holds credentials.
|
|
"""
|
|
raw = (row.get("status") or "").lower()
|
|
if raw in (STATUS_RECONNECT, STATUS_DISCONNECTED, STATUS_ERROR):
|
|
return raw
|
|
if raw in ("authorized", STATUS_CONNECTED, "active"):
|
|
return STATUS_CONNECTED
|
|
if raw == STATUS_PENDING:
|
|
# A legacy row whose plaintext tokens landed after the pending mark.
|
|
return STATUS_CONNECTED if _has_plaintext_tokens(row) else STATUS_PENDING
|
|
return STATUS_CONNECTED if has_credentials(row) else STATUS_PENDING
|
|
|
|
|
|
def worst_status(statuses: Iterable[str]) -> Optional[str]:
|
|
"""The status that needs the most attention, or None for an empty list."""
|
|
ranked = sorted(statuses, key=lambda s: _STATUS_SEVERITY.get(s, 0), reverse=True)
|
|
return ranked[0] if ranked else None
|
|
|
|
|
|
def account_label(row: dict) -> str:
|
|
"""What identifies the account to its owner: an email, a workspace or a host."""
|
|
return (
|
|
row.get("account_label")
|
|
or row.get("user_email")
|
|
or row.get("display_name")
|
|
or catalog.base_url(row.get("server_url")).split("://")[-1]
|
|
or ""
|
|
)
|
|
|
|
|
|
#: Longest name a user can give an account.
|
|
ACCOUNT_NAME_MAX = 80
|
|
|
|
|
|
def account_name(row: dict) -> str:
|
|
"""What the user calls the account, falling back to its label."""
|
|
return (row.get("account_name") or "").strip() or account_label(row)
|
|
|
|
|
|
def rename_connection(conn, row: dict, name: str) -> dict:
|
|
"""Give an account a name (an empty one clears it).
|
|
|
|
The name is only a name: the account is still found by its label when
|
|
its owner signs in again.
|
|
|
|
Args:
|
|
conn: Open connection inside a transaction.
|
|
row: The connection, already authorised for its owner.
|
|
name: The new name; surrounding spaces are dropped.
|
|
|
|
Returns:
|
|
The connection's public shape after the change.
|
|
"""
|
|
repo = ConnectorSessionsRepository(conn)
|
|
repo.update(str(row["id"]), {"account_name": name.strip() or None})
|
|
return serialize_connection(repo.get(str(row["id"])))
|
|
|
|
|
|
def with_account(name: str, account: str) -> str:
|
|
"""``Telegram · Alerts bot``: a tool's name with the account it uses."""
|
|
return f"{name} · {account}" if account and account not in name else name
|
|
|
|
|
|
def account_tool_names(conn, tools: Iterable[dict]) -> dict[str, str]:
|
|
"""Names for tools whose owner has several accounts of the tool's service.
|
|
|
|
With one Telegram bot a tool is just "Telegram"; with two, each tool
|
|
still named after the service becomes "Telegram · <account>", so people
|
|
(and the model) can tell them apart. A name the user chose is kept.
|
|
|
|
Args:
|
|
conn: Open database connection.
|
|
tools: ``user_tools`` rows.
|
|
|
|
Returns:
|
|
Tool id to its name, only for tools whose name changes.
|
|
"""
|
|
repo = ConnectorSessionsRepository(conn)
|
|
connections: dict[str, dict] = {}
|
|
for tool in tools:
|
|
connection_id = str(tool.get("connection_id") or "")
|
|
if connection_id and connection_id not in connections:
|
|
row = repo.get(connection_id)
|
|
if row is not None:
|
|
connections[connection_id] = row
|
|
counts: dict[tuple, int] = {}
|
|
for owner in {row["user_id"] for row in connections.values()}:
|
|
for row in repo.list_for_user(owner):
|
|
if normalize_status(row) != STATUS_PENDING:
|
|
key = (owner, catalog.connector_key_for_row(row))
|
|
counts[key] = counts.get(key, 0) + 1
|
|
names = {}
|
|
for tool in tools:
|
|
row = connections.get(str(tool.get("connection_id") or ""))
|
|
if row is None or counts.get((row["user_id"], catalog.connector_key_for_row(row)), 0) < 2:
|
|
continue
|
|
name = tool.get("custom_name") or tool.get("display_name") or ""
|
|
if name == serialize_connection(row)["name"]:
|
|
names[str(tool["id"])] = with_account(name, account_name(row))
|
|
return names
|
|
|
|
|
|
def _iso(value: Any) -> Optional[str]:
|
|
return value.isoformat() if hasattr(value, "isoformat") else value
|
|
|
|
|
|
def serialize_connection(row: dict, counts: Optional[dict] = None) -> dict:
|
|
"""Public shape of a connection. Never includes tokens or secrets."""
|
|
key = catalog.connector_key_for_row(row)
|
|
definition = catalog.get_definition(key)
|
|
counts = counts or {}
|
|
return {
|
|
"id": str(row["id"]),
|
|
"connector_key": key,
|
|
"name": (definition.name if definition and definition.publisher != "custom" else None)
|
|
or row.get("display_name")
|
|
or (definition.name if definition else key),
|
|
"display_name": row.get("display_name"),
|
|
"icon": definition.icon if definition else "tool_mcp_tool",
|
|
"account_label": account_label(row),
|
|
"account_name": (row.get("account_name") or "").strip() or None,
|
|
"auth_kind": row.get("auth_kind") or (definition.auth_kind if definition else None),
|
|
"status": normalize_status(row),
|
|
"server_url": row.get("server_url"),
|
|
"last_error": row.get("last_error"),
|
|
"created_at": _iso(row.get("created_at")),
|
|
"updated_at": _iso(row.get("updated_at")),
|
|
"last_used_at": _iso(row.get("last_used_at")),
|
|
"source_count": counts.get("sources", 0),
|
|
"tool_count": counts.get("tools", 0),
|
|
"credential_mode": row.get("credential_mode"),
|
|
}
|
|
|
|
|
|
def list_connections(conn, user_id: str) -> list[dict]:
|
|
"""The user's connections, finished sign-ins only, with resource counts."""
|
|
repo = ConnectorSessionsRepository(conn)
|
|
rows = [r for r in repo.list_for_user(user_id) if normalize_status(r) != STATUS_PENDING]
|
|
counts = repo.resource_counts([str(r["id"]) for r in rows])
|
|
return [serialize_connection(r, counts.get(str(r["id"]))) for r in rows]
|
|
|
|
|
|
def serialize_source(row: dict, connection_status: Optional[str] = None) -> dict:
|
|
"""A source linked to a connection, as the connection drawer lists it."""
|
|
metadata = _json(row.get("metadata")) or {}
|
|
sync_state = metadata.get("sync_state") if isinstance(metadata, dict) else None
|
|
if connection_status in (STATUS_RECONNECT, STATUS_DISCONNECTED) and not sync_state:
|
|
sync_state = "paused_reconnect"
|
|
return {
|
|
"id": str(row["id"]),
|
|
"name": row.get("name"),
|
|
"type": row.get("type"),
|
|
"last_sync": _iso(row.get("date")),
|
|
"sync_frequency": row.get("sync_frequency") or "never",
|
|
"sync_state": sync_state or "active",
|
|
}
|
|
|
|
|
|
def serialize_parameters(action: dict, account_parameters: Optional[dict] = None) -> list[dict]:
|
|
"""An action's parameters as the connection drawer shows them.
|
|
|
|
``fixed`` parameters are sent with ``value`` on every call and hidden
|
|
from the model; the others are left to it. A parameter the connection
|
|
itself sets (Telegram's default chat) is fixed with ``set_by: account``.
|
|
"""
|
|
account_parameters = account_parameters or {}
|
|
from docsgpt.agents.tool_pins import is_pinned, iter_parameters
|
|
|
|
schema = action.get("parameters") if isinstance(action.get("parameters"), dict) else {}
|
|
required = set(schema.get("required") or []) if isinstance(schema.get("required"), list) else set()
|
|
parameters = []
|
|
for _section, name, details in iter_parameters(action):
|
|
fixed = is_pinned(details)
|
|
entry = {
|
|
"name": name,
|
|
"description": details.get("description") or "",
|
|
"type": details.get("type") if isinstance(details.get("type"), str) else "string",
|
|
"required": bool(details.get("required")) or name in required,
|
|
"fixed": fixed,
|
|
"value": details.get("value") if fixed else None,
|
|
"set_by": "tool" if fixed else None,
|
|
}
|
|
if name in account_parameters:
|
|
entry.update(fixed=True, value=account_parameters[name], set_by="account")
|
|
parameters.append(entry)
|
|
return parameters
|
|
|
|
|
|
def serialize_tool(
|
|
row: dict, account_parameters: Optional[dict] = None, display_name: Optional[str] = None,
|
|
) -> dict:
|
|
"""A tool linked to a connection, with its actions, permissions and parameters.
|
|
|
|
Args:
|
|
row: The ``user_tools`` row.
|
|
account_parameters: Parameters its connection sets, from
|
|
:func:`connection_parameters`.
|
|
display_name: The name to show instead of the stored one, from
|
|
:func:`account_tool_names`.
|
|
"""
|
|
from docsgpt.connectors.permissions import action_access, action_permission
|
|
|
|
actions = []
|
|
for action in _json(row.get("actions")) or []:
|
|
if not isinstance(action, dict):
|
|
continue
|
|
access = action_access(row.get("name"), action)
|
|
actions.append(
|
|
{
|
|
"name": action.get("name"),
|
|
"description": action.get("description", ""),
|
|
"access": access,
|
|
"permission": action_permission(action),
|
|
"parameters": serialize_parameters(action, account_parameters),
|
|
}
|
|
)
|
|
return {
|
|
"id": str(row["id"]),
|
|
"name": row.get("name"),
|
|
"display_name": display_name or row.get("custom_name") or row.get("display_name") or row.get("name"),
|
|
"status": bool(row.get("status")),
|
|
"credential_mode": row.get("credential_mode") or "owner",
|
|
"actions": actions,
|
|
}
|
|
|
|
|
|
def connection_detail(conn, row: dict) -> dict:
|
|
"""A connection with the sources and tools it feeds."""
|
|
repo = ConnectorSessionsRepository(conn)
|
|
status = normalize_status(row)
|
|
sources = [serialize_source(s, status) for s in repo.list_sources(str(row["id"]))]
|
|
account_parameters = connection_parameters(row) if status == STATUS_CONNECTED else {}
|
|
tool_rows = repo.list_tools(str(row["id"]))
|
|
names = account_tool_names(conn, tool_rows)
|
|
tools = [serialize_tool(t, account_parameters, names.get(str(t["id"]))) for t in tool_rows]
|
|
detail = serialize_connection(row, {"sources": len(sources), "tools": len(tools)})
|
|
detail["sources"] = sources
|
|
detail["tools"] = tools
|
|
detail["writes"] = builtin_writes(catalog.get_definition(detail["connector_key"]), tool_rows)
|
|
return detail
|
|
|
|
|
|
def _card_state(definition: ConnectorDefinition, available: bool, disabled: bool, status: Optional[str]) -> str:
|
|
if disabled:
|
|
return "disabled"
|
|
if status in (STATUS_RECONNECT, STATUS_ERROR):
|
|
return "reconnect"
|
|
if status == STATUS_CONNECTED:
|
|
return "connected"
|
|
if not available:
|
|
return "needs_setup"
|
|
if definition.publisher == "custom":
|
|
return "custom"
|
|
return "available"
|
|
|
|
|
|
def catalog_for_user(conn, user_id: str, *, is_admin: bool, policies: Optional[dict] = None) -> list[dict]:
|
|
"""Catalog entries with availability and the caller's connection summary.
|
|
|
|
Args:
|
|
conn: Open database connection.
|
|
user_id: The caller.
|
|
is_admin: Admins see the names of missing server settings; everyone
|
|
else only learns that setup is needed.
|
|
policies: ``connector_key`` to policy row, when admin policies exist.
|
|
"""
|
|
if policies is None:
|
|
policies = load_policies(conn)
|
|
by_key: dict[str, list[str]] = {}
|
|
for connection in list_connections(conn, user_id):
|
|
by_key.setdefault(connection["connector_key"], []).append(connection["status"])
|
|
|
|
entries = []
|
|
for definition in catalog.all_definitions():
|
|
policy = policies.get(definition.key) or {}
|
|
disabled = not connector_is_enabled(policies, definition.key)
|
|
missing = definition.missing_settings
|
|
available = not missing and not disabled
|
|
statuses = by_key.get(definition.key, [])
|
|
if not available and not statuses:
|
|
# Members only see what they can use, plus what they connected
|
|
# before an admin turned it off (to manage or remove it).
|
|
continue
|
|
status = worst_status(statuses)
|
|
entries.append(
|
|
{
|
|
**definition.to_dict(),
|
|
"available": available,
|
|
"disabled": disabled,
|
|
"needs_setup": bool(missing),
|
|
"missing_settings": missing if is_admin else [],
|
|
"connected_count": sum(1 for s in statuses if s == STATUS_CONNECTED),
|
|
"connection_count": len(statuses),
|
|
"status": status,
|
|
"state": _card_state(definition, available, disabled, status),
|
|
"credential_policy": policy.get("credential_mode") or "choose",
|
|
"writes_allowed": writes_allowed(policies, definition.key),
|
|
}
|
|
)
|
|
return entries
|
|
|
|
|
|
def disconnect(conn, row: dict) -> dict:
|
|
"""Forget a connection's credentials but keep the row and what it feeds.
|
|
|
|
Revokes the grant at the provider where that is supported (Google);
|
|
a failed revocation never blocks forgetting the tokens here. Sources keep
|
|
their indexed content and pause syncing; tools stop working until the
|
|
account is reconnected. An MCP server's client registration is kept so
|
|
reconnecting skips dynamic client registration.
|
|
|
|
Args:
|
|
conn: Open database connection inside a transaction.
|
|
row: The connection row, already authorised for the caller.
|
|
|
|
Returns:
|
|
The connection's public shape after the change.
|
|
"""
|
|
from sqlalchemy import text
|
|
|
|
from docsgpt.security.encryption import CredentialDecryptionError
|
|
|
|
try:
|
|
secrets = read_secrets(row)
|
|
except CredentialDecryptionError:
|
|
secrets = {}
|
|
revoke_at_provider(row, secrets)
|
|
_forget_mcp_clients(row)
|
|
kept = {"client_info": secrets["client_info"]} if secrets.get("client_info") else {}
|
|
write_secrets(
|
|
conn, row, kept, status=STATUS_DISCONNECTED, session_token=None, last_error=None,
|
|
)
|
|
conn.execute(
|
|
text(
|
|
"UPDATE sources SET metadata = metadata || '{\"sync_state\": \"paused_reconnect\"}'::jsonb "
|
|
"WHERE connection_id = CAST(:id AS uuid)"
|
|
),
|
|
{"id": str(row["id"])},
|
|
)
|
|
return serialize_connection(ConnectorSessionsRepository(conn).get(str(row["id"])))
|
|
|
|
|
|
def revoke_at_provider(row: dict, secrets: dict) -> bool:
|
|
"""Best-effort revocation of an OAuth grant. Returns whether it succeeded.
|
|
|
|
Only Google offers a public revocation endpoint for these apps; Microsoft
|
|
and Atlassian tokens are simply forgotten locally.
|
|
"""
|
|
import logging
|
|
|
|
import requests
|
|
|
|
token_info = secrets.get("token_info") or {}
|
|
token = token_info.get("refresh_token") or token_info.get("access_token")
|
|
if row.get("provider") != "google_drive" or not token:
|
|
return False
|
|
try:
|
|
response = requests.post(
|
|
"https://oauth2.googleapis.com/revoke",
|
|
data={"token": token},
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
|
timeout=10,
|
|
)
|
|
return response.status_code == 200
|
|
except requests.RequestException as exc:
|
|
logging.getLogger(__name__).warning("Google token revocation failed: %s", type(exc).__name__)
|
|
return False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Credentials
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# A connection's secrets are one dict, encrypted as a whole into
|
|
# ``encrypted_credentials``:
|
|
#
|
|
# token_info OAuth tokens of an ingest provider (Drive, SharePoint, Confluence)
|
|
# tokens MCP OAuth tokens
|
|
# client_info MCP dynamic client registration (client_id / client_secret)
|
|
# credentials API keys and other pasted secrets, by credential field key
|
|
#
|
|
# This module is the only caller of encrypt_json / decrypt_json for
|
|
# connections, and get_valid_token_info the only reader of OAuth tokens.
|
|
|
|
DECRYPT_ERROR = "Stored credentials could not be decrypted. Reconnect to continue."
|
|
_SECRET_SESSION_KEYS = ("tokens", "client_info")
|
|
|
|
|
|
class ConnectionUnavailable(ValueError):
|
|
"""A connection cannot be used until its owner reconnects it.
|
|
|
|
Attributes:
|
|
connection_id: The connection, when known.
|
|
status: ``reconnect_needed``, ``disconnected`` or ``missing``.
|
|
"""
|
|
|
|
def __init__(self, message: str, *, connection_id: Optional[str] = None, status: str = STATUS_RECONNECT):
|
|
super().__init__(message)
|
|
self.connection_id = connection_id
|
|
self.status = status
|
|
|
|
|
|
class TransientConnectionError(Exception):
|
|
"""The provider failed in a way worth retrying (5xx, rate limit, network)."""
|
|
|
|
|
|
class EncryptionKeyNotConfigured(Exception):
|
|
"""New credentials are refused while a multi-user install uses the public default key."""
|
|
|
|
|
|
def ensure_can_store_credentials() -> None:
|
|
"""Refuse to store new credentials under the public default encryption key.
|
|
|
|
Only multi-user installs (any ``AUTH_TYPE``) are refused; a single-user
|
|
local install keeps working and logs a warning at startup instead.
|
|
|
|
Raises:
|
|
EncryptionKeyNotConfigured: When the key is the default and auth is on.
|
|
"""
|
|
from docsgpt.core.settings import settings
|
|
from docsgpt.security.encryption import is_default_encryption_key
|
|
|
|
if settings.AUTH_TYPE and is_default_encryption_key():
|
|
raise EncryptionKeyNotConfigured("Set ENCRYPTION_SECRET_KEY before connecting services.")
|
|
|
|
|
|
def read_secrets(row: dict) -> dict:
|
|
"""Decrypt a connection's secrets.
|
|
|
|
Rows not yet converted by migration 0040 (or written by an older process
|
|
during a rolling deploy) still carry plaintext ``token_info`` and
|
|
``session_data`` tokens; those are read as they are.
|
|
|
|
Raises:
|
|
CredentialDecryptionError: The envelope cannot be decrypted.
|
|
"""
|
|
from docsgpt.security.encryption import decrypt_json
|
|
|
|
blob = row.get("encrypted_credentials")
|
|
if blob:
|
|
return decrypt_json(blob, row["user_id"])
|
|
secrets: dict = {}
|
|
token_info = _json(row.get("token_info"))
|
|
if isinstance(token_info, dict) and token_info:
|
|
secrets["token_info"] = token_info
|
|
session_data = _json(row.get("session_data")) or {}
|
|
if isinstance(session_data, dict):
|
|
for key in _SECRET_SESSION_KEYS:
|
|
if key in session_data:
|
|
secrets[key] = session_data[key]
|
|
return secrets
|
|
|
|
|
|
def _has_refresh(secrets: dict) -> bool:
|
|
token_info = secrets.get("token_info") or {}
|
|
tokens = secrets.get("tokens") or {}
|
|
return bool(
|
|
(isinstance(token_info, dict) and token_info.get("refresh_token"))
|
|
or (isinstance(tokens, dict) and tokens.get("refresh_token"))
|
|
)
|
|
|
|
|
|
def write_secrets(conn, row: dict, secrets: dict, **fields: Any) -> None:
|
|
"""Replace a connection's secrets and clear every plaintext copy.
|
|
|
|
Args:
|
|
conn: Open connection inside a transaction.
|
|
row: The connection row (only ``id``, ``user_id`` and ``session_data``
|
|
are read).
|
|
secrets: The complete secrets dict; empty clears the credentials.
|
|
**fields: Other columns to update in the same statement.
|
|
"""
|
|
from docsgpt.security.encryption import encrypt_json
|
|
|
|
session_data = _json(row.get("session_data")) or {}
|
|
if isinstance(session_data, dict):
|
|
session_data = {k: v for k, v in session_data.items() if k not in _SECRET_SESSION_KEYS}
|
|
else:
|
|
session_data = {}
|
|
update = {
|
|
"encrypted_credentials": encrypt_json(secrets, row["user_id"]) if secrets else None,
|
|
"has_refresh_token": _has_refresh(secrets),
|
|
"token_info": None,
|
|
"session_data": session_data,
|
|
**fields,
|
|
}
|
|
token_info = secrets.get("token_info")
|
|
if isinstance(token_info, dict) and token_info.get("scopes") and "scopes" not in fields:
|
|
scopes = token_info["scopes"]
|
|
update["scopes"] = scopes.split() if isinstance(scopes, str) else list(scopes)
|
|
ConnectorSessionsRepository(conn).update(str(row["id"]), update)
|
|
|
|
|
|
def credential_hint(credentials: dict) -> str:
|
|
"""``…abcd``: the last four characters of the first long secret, never more."""
|
|
for value in credentials.values():
|
|
if isinstance(value, str) and len(value) >= 8:
|
|
return "…" + value[-4:]
|
|
return "…"
|
|
|
|
|
|
def mark_reconnect_needed(connection_id: str, error: str) -> None:
|
|
"""Flag a connection whose credentials stopped working and tell its owner.
|
|
|
|
Pauses syncing on every source the connection feeds (no retry storm) and
|
|
publishes a ``connection.reconnect_needed`` user event, which the
|
|
frontend shows as a toast with a Reconnect action. Runs in its own
|
|
transaction so it sticks even when the caller's transaction rolls back.
|
|
"""
|
|
from sqlalchemy import text
|
|
|
|
from docsgpt.events.publisher import publish_user_event
|
|
with db_session() as conn:
|
|
repo = ConnectorSessionsRepository(conn)
|
|
row = repo.get(connection_id)
|
|
if row is None:
|
|
return
|
|
already = normalize_status(row) == STATUS_RECONNECT
|
|
counts = repo.resource_counts([connection_id]).get(str(connection_id), {})
|
|
repo.update(connection_id, {"status": STATUS_RECONNECT, "last_error": error[:500]})
|
|
conn.execute(
|
|
text(
|
|
"UPDATE sources SET metadata = metadata || '{\"sync_state\": \"paused_reconnect\"}'::jsonb "
|
|
"WHERE connection_id = CAST(:id AS uuid)"
|
|
),
|
|
{"id": connection_id},
|
|
)
|
|
if not already:
|
|
publish_user_event(
|
|
row["user_id"],
|
|
"connection.reconnect_needed",
|
|
{
|
|
"connection_id": connection_id,
|
|
"connector_key": catalog.connector_key_for_row(row),
|
|
"name": serialize_connection(row)["name"],
|
|
"source_count": counts.get("sources", 0),
|
|
"tool_count": counts.get("tools", 0),
|
|
},
|
|
scope={"kind": "connection", "id": connection_id},
|
|
)
|
|
|
|
|
|
def resume_sources(conn, connection_id: str) -> None:
|
|
"""Lift the reconnect pause from a connection's sources."""
|
|
from sqlalchemy import text
|
|
|
|
conn.execute(
|
|
text(
|
|
"UPDATE sources SET metadata = metadata - 'sync_state' "
|
|
"WHERE connection_id = CAST(:id AS uuid) AND metadata ? 'sync_state'"
|
|
),
|
|
{"id": connection_id},
|
|
)
|
|
|
|
|
|
def load_secrets(row: dict) -> Optional[dict]:
|
|
"""Decrypt a row's secrets; on failure flag it for reconnect and return None."""
|
|
from docsgpt.security.encryption import CredentialDecryptionError
|
|
|
|
try:
|
|
return read_secrets(row)
|
|
except CredentialDecryptionError:
|
|
mark_reconnect_needed(str(row["id"]), DECRYPT_ERROR)
|
|
return None
|
|
|
|
|
|
def _is_auth_failure(exc: BaseException) -> bool:
|
|
"""Whether a refresh failure means the grant is gone (vs. a retryable blip)."""
|
|
import requests
|
|
|
|
response = getattr(exc, "response", None)
|
|
status = getattr(response, "status_code", None)
|
|
if isinstance(exc, requests.exceptions.HTTPError) and status is not None:
|
|
return status in (400, 401, 403)
|
|
if isinstance(exc, (requests.exceptions.ConnectionError, requests.exceptions.Timeout)):
|
|
return False
|
|
message = str(exc).lower()
|
|
if any(word in message for word in ("timed out", "temporarily", "rate limit", "503", "502", "504")):
|
|
return False
|
|
try:
|
|
from google.auth.exceptions import RefreshError, TransportError
|
|
|
|
if isinstance(exc, TransportError):
|
|
return False
|
|
if isinstance(exc, RefreshError):
|
|
return True
|
|
except ImportError: # pragma: no cover - google libs are core deps
|
|
pass
|
|
return True
|
|
|
|
|
|
def connection_id_for_session_token(session_token: Optional[str]) -> str:
|
|
"""The connection behind a legacy browser session token.
|
|
|
|
Raises:
|
|
ConnectionUnavailable: No connection holds that token.
|
|
"""
|
|
from docsgpt.parser.connectors._auth_utils import session_token_fingerprint
|
|
row = None
|
|
if session_token:
|
|
with db_readonly() as conn:
|
|
row = ConnectorSessionsRepository(conn).get_by_session_token(session_token)
|
|
if row is None:
|
|
raise ConnectionUnavailable(
|
|
f"Invalid session token ({session_token_fingerprint(session_token or '')})", status="missing",
|
|
)
|
|
return str(row["id"])
|
|
|
|
|
|
def get_valid_token_info(connection_id: str, *, rejected_access_token: Optional[str] = None) -> dict:
|
|
"""An unexpired OAuth ``token_info`` for an ingest connection.
|
|
|
|
Holds a row lock (``SELECT ... FOR UPDATE``) while it checks expiry and
|
|
refreshes, and writes the rotated refresh token back in the same
|
|
transaction, so two workers refreshing a rotating token (Microsoft,
|
|
Atlassian) never spend the same one twice: the second waits and then
|
|
reads the first one's fresh token.
|
|
|
|
Args:
|
|
connection_id: The connection to use.
|
|
rejected_access_token: An access token the provider just answered
|
|
with 401. The token is refreshed when it is still the stored one;
|
|
when another worker already replaced it, the new one is returned.
|
|
|
|
Returns:
|
|
The decrypted token info, refreshed if it was about to expire.
|
|
|
|
Raises:
|
|
ConnectionUnavailable: The connection is missing, disconnected, cannot
|
|
be decrypted, or its grant was revoked (it is then flagged
|
|
``reconnect_needed``).
|
|
TransientConnectionError: The provider failed in a retryable way.
|
|
"""
|
|
from docsgpt.parser.connectors.connector_creator import ConnectorCreator
|
|
from docsgpt.security.encryption import CredentialDecryptionError
|
|
failure: Optional[str] = None
|
|
transient: Optional[BaseException] = None
|
|
token_info: Optional[dict] = None
|
|
with db_session() as conn:
|
|
repo = ConnectorSessionsRepository(conn)
|
|
row = repo.get_for_update(connection_id)
|
|
if row is None:
|
|
raise ConnectionUnavailable("Connection not found", connection_id=connection_id, status="missing")
|
|
status = normalize_status(row)
|
|
if status == STATUS_DISCONNECTED:
|
|
raise ConnectionUnavailable(
|
|
"Connection is disconnected", connection_id=connection_id, status=STATUS_DISCONNECTED
|
|
)
|
|
try:
|
|
secrets = read_secrets(row)
|
|
except CredentialDecryptionError:
|
|
secrets, failure = {}, DECRYPT_ERROR
|
|
token_info = secrets.get("token_info") if not failure else None
|
|
if not failure and not token_info:
|
|
failure = "No stored sign-in for this connection. Reconnect to continue."
|
|
if not failure:
|
|
auth = ConnectorCreator.create_auth(row["provider"])
|
|
rejected = bool(rejected_access_token) and token_info.get("access_token") == rejected_access_token
|
|
if rejected or auth.is_token_expired(token_info):
|
|
refresh_token = token_info.get("refresh_token")
|
|
if not refresh_token:
|
|
failure = "The sign-in expired and cannot be renewed. Reconnect to continue."
|
|
else:
|
|
try:
|
|
refreshed = auth.refresh_access_token(refresh_token)
|
|
except Exception as exc: # classified below
|
|
if _is_auth_failure(exc):
|
|
failure = f"The provider rejected the stored sign-in: {type(exc).__name__}"
|
|
else:
|
|
transient = exc
|
|
else:
|
|
merged = {**token_info, **{k: v for k, v in refreshed.items() if v is not None}}
|
|
token_info = auth.sanitize_token_info(merged)
|
|
write_secrets(
|
|
conn, row, {**secrets, "token_info": token_info},
|
|
status=STATUS_CONNECTED, last_error=None,
|
|
)
|
|
if not failure and transient is None:
|
|
repo.update(connection_id, {"last_used_at": _now()})
|
|
if failure:
|
|
mark_reconnect_needed(connection_id, failure)
|
|
raise ConnectionUnavailable(failure, connection_id=connection_id)
|
|
if transient is not None:
|
|
raise TransientConnectionError(str(transient)) from transient
|
|
return token_info
|
|
|
|
|
|
def _now():
|
|
import datetime
|
|
|
|
return datetime.datetime.now(datetime.timezone.utc)
|
|
|
|
|
|
def get_credentials(row: dict) -> dict:
|
|
"""The pasted credentials (API keys) of an ``api_key`` connection.
|
|
|
|
Raises:
|
|
ConnectionUnavailable: Disconnected, flagged, or undecryptable.
|
|
"""
|
|
status = normalize_status(row)
|
|
if status in (STATUS_DISCONNECTED, STATUS_RECONNECT):
|
|
raise ConnectionUnavailable(f"Connection is {status}", connection_id=str(row["id"]), status=status)
|
|
secrets = load_secrets(row)
|
|
if secrets is None:
|
|
raise ConnectionUnavailable(DECRYPT_ERROR, connection_id=str(row["id"]))
|
|
return dict(secrets.get("credentials") or {})
|
|
|
|
|
|
def access_credentials(row: dict) -> dict:
|
|
"""What a loader or tool authenticates with through this connection.
|
|
|
|
Pasted credentials for an ``api_key`` connection. An OAuth connection
|
|
whose token a loader or tool sends itself (GitHub's App sign-in) gives
|
|
its current access token as ``access_token``, refreshed first when it
|
|
has expired.
|
|
|
|
Raises:
|
|
ConnectionUnavailable: Disconnected, flagged, undecryptable, or the
|
|
refresh was refused (the connection is then flagged).
|
|
TransientConnectionError: The provider failed in a retryable way.
|
|
"""
|
|
if (row.get("auth_kind") or "") == "oauth":
|
|
status = normalize_status(row)
|
|
if status in (STATUS_DISCONNECTED, STATUS_RECONNECT):
|
|
raise ConnectionUnavailable(f"Connection is {status}", connection_id=str(row["id"]), status=status)
|
|
return {"access_token": get_valid_token_info(str(row["id"])).get("access_token")}
|
|
return get_credentials(row)
|
|
|
|
|
|
def connection_parameters(row: Optional[dict]) -> dict:
|
|
"""Tool parameters a connection sets, from its credential fields.
|
|
|
|
Only fields the catalog marks with ``parameter`` count (Telegram's
|
|
default chat); an empty value sets nothing. A connection that cannot be
|
|
read sets nothing either: the call then fails on its credentials.
|
|
|
|
Args:
|
|
row: The connection row, or None.
|
|
|
|
Returns:
|
|
Parameter name to value.
|
|
"""
|
|
if row is None:
|
|
return {}
|
|
fields = catalog.parameter_fields(catalog.connector_key_for_row(row))
|
|
if not fields:
|
|
return {}
|
|
try:
|
|
credentials = get_credentials(row)
|
|
except ConnectionUnavailable:
|
|
return {}
|
|
return {
|
|
f.parameter: credentials[f.key]
|
|
for f in fields
|
|
if credentials.get(f.key) not in (None, "")
|
|
}
|
|
|
|
|
|
def _api_key_account(
|
|
repo: ConnectorSessionsRepository,
|
|
user_id: str,
|
|
connector_key: str,
|
|
server_url: Optional[str],
|
|
label: str,
|
|
credentials: dict,
|
|
) -> tuple[Optional[dict], str]:
|
|
"""The connection holding exactly ``credentials``, or a free label for a new one.
|
|
|
|
A label (a hint of the key, or what the user typed) is not an identity:
|
|
two different keys can share it. A row under the label is reused only
|
|
when it holds the same credentials; otherwise the label gets a
|
|
``(2)``, ``(3)`` suffix until it names a matching row or no row.
|
|
|
|
Returns:
|
|
``(row, label)``: the row to reuse (None to create one) and its label.
|
|
"""
|
|
from docsgpt.security.encryption import CredentialDecryptionError
|
|
|
|
candidate, suffix = label, 1
|
|
while True:
|
|
existing = repo.find_account(user_id, connector_key, server_url=server_url, account_label=candidate)
|
|
if existing is None:
|
|
return None, candidate
|
|
try:
|
|
if (read_secrets(existing).get("credentials") or {}) == credentials:
|
|
return existing, candidate
|
|
except CredentialDecryptionError:
|
|
pass
|
|
suffix += 1
|
|
candidate = f"{label} ({suffix})"
|
|
|
|
|
|
def create_api_key_connection(
|
|
conn,
|
|
user_id: str,
|
|
definition,
|
|
credentials: dict,
|
|
*,
|
|
label: Optional[str] = None,
|
|
server_url: Optional[str] = None,
|
|
display_name: Optional[str] = None,
|
|
) -> tuple[dict, bool]:
|
|
"""Store a connection for pasted credentials, or reuse the matching one.
|
|
|
|
The same credentials for the same service map to one connection
|
|
("enter secrets once"): adding a second tool or bucket reuses it.
|
|
|
|
Args:
|
|
conn: Open connection inside a transaction.
|
|
user_id: The owner.
|
|
definition: The catalog entry (``api_key`` or ``custom_mcp``).
|
|
credentials: Field key to value, as the user entered them.
|
|
label: What the account is called; defaults to a hint of the key.
|
|
server_url: Base URL, for custom MCP servers.
|
|
display_name: Name shown for custom connectors.
|
|
|
|
Returns:
|
|
``(row, created)``.
|
|
|
|
Raises:
|
|
ConnectorDisabled: An admin turned the connector off. Checked first,
|
|
so callers that fall back to a legacy path on the other errors
|
|
never do so for a disabled connector. A custom MCP server at a
|
|
preset's address is that preset (see
|
|
:func:`catalog.connector_key_for_row`), so its switch applies.
|
|
ValueError: A required credential field is missing.
|
|
EncryptionKeyNotConfigured: See :func:`ensure_can_store_credentials`.
|
|
"""
|
|
preset = catalog.preset_for_url(server_url) if definition.key == "custom_mcp" else None
|
|
ensure_connector_allowed(conn, preset.key if preset else definition.key)
|
|
fields = {f.key: f for f in definition.credential_fields}
|
|
if fields:
|
|
missing = [f.label for f in fields.values() if f.required and not str(credentials.get(f.key) or "").strip()]
|
|
if missing:
|
|
raise ValueError(f"Missing credentials: {', '.join(missing)}")
|
|
credentials = {k: v for k, v in credentials.items() if k in fields and v not in (None, "")}
|
|
else:
|
|
credentials = {k: v for k, v in credentials.items() if v not in (None, "")}
|
|
ensure_can_store_credentials()
|
|
secret_values = {k: v for k, v in credentials.items() if (fields.get(k).secret if fields.get(k) else True)}
|
|
repo = ConnectorSessionsRepository(conn)
|
|
existing, account_label = _api_key_account(
|
|
repo, user_id, definition.key, server_url, label or credential_hint(secret_values or credentials), credentials,
|
|
)
|
|
if existing is not None:
|
|
write_secrets(conn, existing, {"credentials": credentials}, status=STATUS_CONNECTED, last_error=None)
|
|
resume_sources(conn, str(existing["id"]))
|
|
return repo.get(str(existing["id"])), False
|
|
from docsgpt.security.encryption import encrypt_json
|
|
|
|
row = repo.create(
|
|
user_id,
|
|
definition.key,
|
|
connector_key=definition.key,
|
|
auth_kind="api_key",
|
|
display_name=display_name or definition.name,
|
|
account_label=account_label,
|
|
server_url=server_url,
|
|
encrypted_credentials=encrypt_json({"credentials": credentials}, user_id),
|
|
)
|
|
if row is None: # lost a race with an identical insert
|
|
row = repo.find_account(user_id, definition.key, server_url=server_url, account_label=account_label)
|
|
return row, False
|
|
return row, True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OAuth sign-in
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def begin_oauth(conn, user_id: str, provider: str, connection_id: Optional[str] = None) -> dict:
|
|
"""The row an OAuth sign-in writes into when it completes.
|
|
|
|
Reconnecting passes the connection being reconnected. A new sign-in
|
|
gets a pending row with no account label; the callback moves the tokens
|
|
onto the existing row for that account if there is one.
|
|
|
|
Raises:
|
|
ConnectionUnavailable: ``connection_id`` is not the caller's
|
|
connection for ``provider``.
|
|
"""
|
|
from sqlalchemy import text
|
|
|
|
repo = ConnectorSessionsRepository(conn)
|
|
if connection_id:
|
|
row = repo.get_for_user(connection_id, user_id)
|
|
if row is None or row.get("provider") != provider:
|
|
raise ConnectionUnavailable("Connection not found", connection_id=connection_id, status="missing")
|
|
return row
|
|
ensure_connector_allowed(conn, provider)
|
|
definition = catalog.get_definition(provider)
|
|
result = conn.execute(
|
|
text(
|
|
"""
|
|
INSERT INTO connector_sessions (user_id, provider, status, connector_key, auth_kind, display_name)
|
|
VALUES (:user_id, :provider, 'pending', :provider, 'oauth', :display_name)
|
|
ON CONFLICT (user_id, provider, COALESCE(server_url, ''), COALESCE(account_label, ''))
|
|
DO UPDATE SET status = CASE
|
|
WHEN connector_sessions.status IN ('connected', 'authorized') THEN connector_sessions.status
|
|
ELSE 'pending' END
|
|
RETURNING *
|
|
"""
|
|
),
|
|
{"user_id": user_id, "provider": provider, "display_name": definition.name if definition else provider},
|
|
)
|
|
from docsgpt.storage.db.base_repository import row_to_dict
|
|
|
|
return row_to_dict(result.fetchone())
|
|
|
|
|
|
def complete_oauth(conn, state_row: dict, provider: str, token_info: dict, account: str) -> dict:
|
|
"""Store a finished OAuth sign-in and return the connection it belongs to.
|
|
|
|
Signing in to an account that already has a connection updates that
|
|
connection (and drops the pending placeholder), so reconnecting from any
|
|
entry point heals every source and tool of the account. Reconnecting a
|
|
connection but signing in to a different account never rewrites it (its
|
|
sources and tools would silently run as the new account): the new
|
|
account gets its own connection and the original keeps its status.
|
|
|
|
Args:
|
|
conn: Open connection inside a transaction.
|
|
state_row: The row named in the OAuth ``state``.
|
|
provider: ``google_drive``, ``share_point`` or ``confluence``.
|
|
token_info: Sanitised token info from the provider.
|
|
account: The account's email or name, shown as "Connected as".
|
|
|
|
Returns:
|
|
The connection now holding the sign-in.
|
|
"""
|
|
import uuid
|
|
|
|
repo = ConnectorSessionsRepository(conn)
|
|
target = state_row
|
|
definition = catalog.get_definition(provider)
|
|
display_name = definition.name if definition else provider
|
|
existing = repo.find_account(state_row["user_id"], provider, server_url=None, account_label=account)
|
|
if existing is not None and str(existing["id"]) != str(state_row["id"]):
|
|
target = existing
|
|
if not has_credentials(state_row) and normalize_status(state_row) == STATUS_PENDING:
|
|
repo.delete_by_id(str(state_row["id"]))
|
|
elif existing is None and state_row.get("account_label"):
|
|
# The row belongs to another account (a reconnect that signed in as someone else).
|
|
target = repo.create(
|
|
state_row["user_id"], provider, connector_key=provider, auth_kind="oauth",
|
|
display_name=display_name, account_label=account, status=STATUS_PENDING,
|
|
) or repo.find_account(state_row["user_id"], provider, server_url=None, account_label=account)
|
|
write_secrets(
|
|
conn,
|
|
target,
|
|
{"token_info": token_info},
|
|
status=STATUS_CONNECTED,
|
|
# Kept for frontends from before connections; new ones use the id.
|
|
session_token=str(uuid.uuid4()),
|
|
user_email=account,
|
|
account_label=account,
|
|
connector_key=provider,
|
|
auth_kind="oauth",
|
|
display_name=display_name,
|
|
last_error=None,
|
|
)
|
|
resume_sources(conn, str(target["id"]))
|
|
return repo.get(str(target["id"]))
|
|
|
|
|
|
def picker_token(connection_id: str) -> dict:
|
|
"""What a browser-side picker needs: a short-lived access token, never the refresh token."""
|
|
token_info = get_valid_token_info(connection_id)
|
|
return {
|
|
"access_token": token_info.get("access_token"),
|
|
"expiry": token_info.get("expiry"),
|
|
"allows_shared_content": bool(token_info.get("allows_shared_content")),
|
|
}
|
|
|
|
|
|
def claim_session_token(conn, user_id: str, provider: str, session_token: str) -> Optional[dict]:
|
|
"""Link a legacy browser session token to its connection, once.
|
|
|
|
Returns the connection when the token belongs to ``user_id`` and
|
|
``provider``. The token itself keeps working for the session-token
|
|
routes during this release, so an older tab does not break.
|
|
"""
|
|
repo = ConnectorSessionsRepository(conn)
|
|
row = repo.get_by_session_token(session_token)
|
|
if not row or row.get("user_id") != user_id or (row.get("provider") or "").lower() != provider.lower():
|
|
return None
|
|
return row
|
|
|
|
|
|
def resolve_request_connection(user_id: str, provider: Optional[str], data: dict) -> Optional[dict]:
|
|
"""The caller's connection a request names, by ``connection_id`` or legacy ``session_token``.
|
|
|
|
Returns None unless the caller owns the connection and it belongs to
|
|
``provider`` (when given).
|
|
"""
|
|
connection_id = data.get("connection_id")
|
|
session_token = data.get("session_token")
|
|
if not connection_id and not session_token:
|
|
return None
|
|
with db_readonly() as conn:
|
|
repo = ConnectorSessionsRepository(conn)
|
|
if connection_id:
|
|
row = repo.get_for_user(str(connection_id), user_id)
|
|
else:
|
|
row = repo.get_by_session_token(session_token)
|
|
if not owns_connector_session(row, user_id, provider):
|
|
return None
|
|
if row is None:
|
|
return None
|
|
if provider and (row.get("provider") or "").lower() != provider.lower():
|
|
return None
|
|
return row
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MCP OAuth token storage
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def mcp_provider(base_url: str) -> str:
|
|
"""``provider`` value MCP OAuth connections are stored under."""
|
|
return f"mcp:{base_url}"
|
|
|
|
|
|
def _mcp_row(conn, user_id: str, base_url: str, connection_id: Optional[str], *, lock: bool = False):
|
|
from sqlalchemy import text
|
|
|
|
from docsgpt.storage.db.base_repository import row_to_dict
|
|
|
|
repo = ConnectorSessionsRepository(conn)
|
|
if connection_id:
|
|
row = repo.get_for_update(connection_id) if lock else repo.get(connection_id)
|
|
# A connection's tokens only ever go to its own server. Ownership is
|
|
# checked where the id is chosen (the tool executor); clients cannot
|
|
# supply one (see ``_sanitize_mcp_transport``).
|
|
if row is None or row.get("provider") != mcp_provider(base_url):
|
|
return None
|
|
return row
|
|
result = conn.execute(
|
|
text(
|
|
"SELECT * FROM connector_sessions WHERE user_id = :user_id AND provider = :provider "
|
|
"ORDER BY updated_at DESC LIMIT 1" + (" FOR UPDATE" if lock else "")
|
|
),
|
|
{"user_id": user_id, "provider": mcp_provider(base_url)},
|
|
)
|
|
row = result.fetchone()
|
|
return row_to_dict(row) if row is not None else None
|
|
|
|
|
|
def read_mcp_secrets(user_id: str, base_url: str, connection_id: Optional[str] = None) -> dict:
|
|
"""The MCP OAuth ``tokens`` and ``client_info`` of a server connection."""
|
|
with db_readonly() as conn:
|
|
row = _mcp_row(conn, user_id, base_url, connection_id)
|
|
if row is None:
|
|
return {}
|
|
return load_secrets(row) or {}
|
|
|
|
|
|
def update_mcp_secrets(
|
|
user_id: str,
|
|
base_url: str,
|
|
patch: dict,
|
|
*,
|
|
connection_id: Optional[str] = None,
|
|
status: Optional[str] = None,
|
|
) -> dict:
|
|
"""Merge ``patch`` into an MCP connection's secrets (``None`` drops a key).
|
|
|
|
Creates the connection on first use, named after the matching preset or
|
|
the server's host. A write for a named connection never creates one: a
|
|
client or sync still running when its connection was removed would
|
|
otherwise bring it back, holding the tokens it just renewed.
|
|
|
|
Returns:
|
|
The connection row after the update.
|
|
|
|
Raises:
|
|
ConnectionUnavailable: ``connection_id`` names no connection of this
|
|
user for this server (removed, or another server's).
|
|
"""
|
|
from docsgpt.security.encryption import CredentialDecryptionError
|
|
|
|
with db_session() as conn:
|
|
repo = ConnectorSessionsRepository(conn)
|
|
row = _mcp_row(conn, user_id, base_url, connection_id, lock=True)
|
|
if row is None and connection_id:
|
|
raise ConnectionUnavailable("Connection not found", connection_id=connection_id, status="missing")
|
|
if row is None:
|
|
row = repo.merge_session_data(user_id, mcp_provider(base_url), base_url, {})
|
|
try:
|
|
secrets = read_secrets(row)
|
|
except CredentialDecryptionError:
|
|
secrets = {}
|
|
for key, value in patch.items():
|
|
if value is None:
|
|
secrets.pop(key, None)
|
|
else:
|
|
secrets[key] = value
|
|
fields: dict = {}
|
|
if status:
|
|
fields["status"] = status
|
|
if status == STATUS_CONNECTED:
|
|
fields["last_error"] = None
|
|
elif not row.get("status"):
|
|
fields["status"] = STATUS_CONNECTED if secrets.get("tokens") else STATUS_PENDING
|
|
if not row.get("connector_key"):
|
|
preset = catalog.preset_for_url(base_url)
|
|
fields["connector_key"] = preset.key if preset else "custom_mcp"
|
|
fields["auth_kind"] = "mcp_oauth"
|
|
fields["display_name"] = preset.name if preset else base_url.split("://")[-1]
|
|
write_secrets(conn, row, secrets, **fields)
|
|
if fields.get("status") == STATUS_CONNECTED:
|
|
resume_sources(conn, str(row["id"]))
|
|
return repo.get(str(row["id"]))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tools created from a connection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _transform_actions(actions: list) -> list:
|
|
"""``transform_actions`` from the tools API: active, LLM-filled parameters."""
|
|
transformed = []
|
|
for action in actions:
|
|
action = dict(action)
|
|
action["active"] = True
|
|
parameters = action.get("parameters")
|
|
if isinstance(parameters, dict):
|
|
for details in (parameters.get("properties") or {}).values():
|
|
if isinstance(details, dict):
|
|
details["filled_by_llm"] = True
|
|
details["value"] = ""
|
|
transformed.append(action)
|
|
return transformed
|
|
|
|
|
|
def create_tool_for_connection(
|
|
conn,
|
|
user_id: str,
|
|
connection: dict,
|
|
*,
|
|
template: Optional[str] = None,
|
|
display_name: Optional[str] = None,
|
|
config: Optional[dict] = None,
|
|
actions: Optional[list] = None,
|
|
permissions: Optional[dict] = None,
|
|
status: bool = True,
|
|
) -> dict:
|
|
"""Create the tool a connection provides, with read / write defaults.
|
|
|
|
Writes default to "Needs approval", reads to "Always allow"; the
|
|
``permissions`` map (action name to ``always`` / ``ask`` / ``off``)
|
|
overrides them. Secrets are never copied onto the tool: the executor
|
|
reads them from the connection at run time.
|
|
|
|
Args:
|
|
conn: Open connection inside a transaction.
|
|
user_id: The owner of the connection and the new tool.
|
|
connection: The connection row.
|
|
template: ``user_tools`` name; defaults to the connector's first template.
|
|
display_name: Name shown for the tool.
|
|
config: Non-secret tool configuration (an MCP server URL, say).
|
|
actions: Action metadata; defaults to the tool class's own.
|
|
permissions: Per-action permission overrides.
|
|
status: Whether the tool starts enabled.
|
|
|
|
Returns:
|
|
The new ``user_tools`` row.
|
|
"""
|
|
from docsgpt.agents.tools.tool_manager import ToolManager
|
|
from docsgpt.connectors.permissions import apply_default_permissions, apply_permission
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
|
|
key = catalog.connector_key_for_row(connection)
|
|
definition = catalog.get_definition(key)
|
|
template = template or (definition.tool_templates[0] if definition and definition.tool_templates else None)
|
|
if not template:
|
|
raise ValueError(f"Connector {key} provides no tool")
|
|
tool = ToolManager(config={}).tools.get(template)
|
|
if tool is None:
|
|
raise ValueError(f"Unknown tool template: {template}")
|
|
doc = (tool.__doc__ or template).strip().split("\n", 1)
|
|
if actions is None:
|
|
actions = tool.get_actions_metadata()
|
|
actions = apply_default_permissions(template, _transform_actions(actions))
|
|
for index, action in enumerate(actions):
|
|
permission = (permissions or {}).get(action.get("name"))
|
|
if permission:
|
|
actions[index] = apply_permission(action, permission)
|
|
name = display_name or (serialize_connection(connection)["name"] if connection else doc[0].strip())
|
|
return UserToolsRepository(conn).create(
|
|
user_id,
|
|
template,
|
|
config=dict(config or {}),
|
|
custom_name=name,
|
|
display_name=name,
|
|
description=doc[1].strip() if len(doc) > 1 else "",
|
|
config_requirements=tool.get_config_requirements(),
|
|
actions=actions,
|
|
status=status,
|
|
connection_id=str(connection["id"]),
|
|
)
|
|
|
|
|
|
def split_secrets(config: dict, config_requirements: dict) -> tuple[dict, dict]:
|
|
"""``(public, secrets)``: a tool config split by its secret requirements."""
|
|
secret_keys = {k for k, spec in (config_requirements or {}).items() if spec.get("secret")}
|
|
public = {k: v for k, v in config.items() if k not in secret_keys}
|
|
secrets = {k: v for k, v in config.items() if k in secret_keys and v not in (None, "")}
|
|
return public, secrets
|
|
|
|
|
|
def builtin_mcp_config(definition: Optional[ConnectorDefinition], writes: bool = False) -> Optional[dict]:
|
|
"""Tool config of a built-in connector whose tool is its service's MCP server.
|
|
|
|
The token is not part of it: the executor reads it from the connection
|
|
at run time and sends it as a bearer token.
|
|
|
|
Args:
|
|
definition: The connector.
|
|
writes: Use the endpoint that also offers write actions, when the
|
|
connector has one (GitHub's full server).
|
|
|
|
Returns:
|
|
The config, or None when the connector has no such tool.
|
|
"""
|
|
if (
|
|
definition is None
|
|
or definition.publisher != "built_in"
|
|
or not definition.mcp_url
|
|
or "mcp_tool" not in definition.tool_templates
|
|
):
|
|
return None
|
|
url = definition.mcp_write_url if writes and definition.mcp_write_url else definition.mcp_url
|
|
return {"server_url": url, "auth_type": "bearer", "transport_type": "http", "timeout": 30}
|
|
|
|
|
|
def builtin_mcp_url(definition: Optional[ConnectorDefinition], stored_url: Optional[str], allowed: bool) -> str:
|
|
"""The endpoint a built-in connector's MCP tool may call, whatever its config says.
|
|
|
|
Only the connector's own endpoints are ever used: the write endpoint when
|
|
the tool was set up for writes and an admin allows them, the read-only
|
|
one otherwise (another path on the same host included).
|
|
|
|
Args:
|
|
definition: The connector, with ``mcp_url`` set.
|
|
stored_url: The ``server_url`` in the tool's config.
|
|
allowed: Whether an admin allows writes through this connector.
|
|
"""
|
|
if allowed and definition.mcp_write_url and stored_url == definition.mcp_write_url:
|
|
return definition.mcp_write_url
|
|
return definition.mcp_url or ""
|
|
|
|
|
|
def builtin_writes(definition: Optional[ConnectorDefinition], tool_rows: Iterable[dict]) -> Optional[bool]:
|
|
"""Whether a connection's built-in MCP tool is set up for writes.
|
|
|
|
Returns:
|
|
None when the connector offers no writes or the connection has no such
|
|
tool yet; otherwise whether the tool uses the write endpoint.
|
|
"""
|
|
if definition is None or not definition.mcp_write_url or builtin_mcp_config(definition) is None:
|
|
return None
|
|
for tool in tool_rows:
|
|
if tool.get("name") == "mcp_tool":
|
|
return (_json(tool.get("config")) or {}).get("server_url") == definition.mcp_write_url
|
|
return None
|
|
|
|
|
|
def needs_mcp_discovery(conn, connection: dict) -> bool:
|
|
"""Whether creating this connection's tools first needs its MCP server's actions."""
|
|
definition = catalog.get_definition(catalog.connector_key_for_row(connection))
|
|
if builtin_mcp_config(definition) is None:
|
|
return False
|
|
have = {tool.get("name") for tool in ConnectorSessionsRepository(conn).list_tools(str(connection["id"]))}
|
|
return "mcp_tool" not in have
|
|
|
|
|
|
def ensure_connection_tools(
|
|
conn,
|
|
user_id: str,
|
|
connection: dict,
|
|
permissions: Optional[dict] = None,
|
|
mcp_actions: Optional[list] = None,
|
|
mcp_writes: bool = False,
|
|
) -> list[dict]:
|
|
"""Create the connector's tools once; later calls return the existing ones.
|
|
|
|
This is what makes the setup step idempotent for tools: a retried or
|
|
repeated setup never creates a second Telegram tool for the same bot.
|
|
|
|
Args:
|
|
conn: Open connection inside a transaction.
|
|
user_id: The owner.
|
|
connection: The connection row.
|
|
permissions: Per-action permission overrides.
|
|
mcp_actions: The actions of a built-in connector's MCP server
|
|
(GitHub's), discovered by the caller outside this transaction.
|
|
Without them that tool is not created.
|
|
mcp_writes: Point that tool at the endpoint that also offers writes;
|
|
the caller checked that an admin allows them.
|
|
"""
|
|
repo = ConnectorSessionsRepository(conn)
|
|
existing = repo.list_tools(str(connection["id"]))
|
|
key = catalog.connector_key_for_row(connection)
|
|
definition = catalog.get_definition(key)
|
|
if not definition or not definition.tool_templates:
|
|
return existing
|
|
have = {tool.get("name") for tool in existing}
|
|
mcp_config = builtin_mcp_config(definition, writes=mcp_writes)
|
|
created = []
|
|
for template in definition.tool_templates:
|
|
if template in have:
|
|
continue
|
|
if template == "mcp_tool" and mcp_config is not None and mcp_actions is not None:
|
|
created.append(create_tool_for_connection(
|
|
conn, user_id, connection, template=template, config=mcp_config,
|
|
actions=mcp_actions, permissions=permissions,
|
|
))
|
|
continue
|
|
if template in ("mcp_tool", "api_tool"):
|
|
# Other MCP tools are created by the MCP save flow, which has the
|
|
# discovered actions; OpenAPI tools come from an imported spec.
|
|
continue
|
|
created.append(
|
|
create_tool_for_connection(conn, user_id, connection, template=template, permissions=permissions)
|
|
)
|
|
return existing + created
|
|
|
|
|
|
def _forget_mcp_clients(row: dict) -> None:
|
|
"""Drop the MCP clients this process cached for a connection's tokens.
|
|
|
|
A cached client keeps the tokens it signed in with for a few minutes, so
|
|
after the connection is removed or disconnected it would still answer as
|
|
signed in (and save a tool with no connection behind it).
|
|
"""
|
|
if not str(row.get("provider") or "").startswith("mcp:") and row.get("auth_kind") != "mcp_oauth":
|
|
return
|
|
# Not imported here: a process that never loaded the MCP tool has no
|
|
# clients cached, and importing it from this module would be circular.
|
|
mcp_tool = sys.modules.get("docsgpt.agents.tools.mcp_tool")
|
|
if mcp_tool is not None:
|
|
mcp_tool.forget_cached_clients(str(row["id"]), str(row.get("user_id") or ""))
|
|
|
|
|
|
def remove_connection(conn, row: dict, *, sources: str = "keep", tools: str = "delete") -> list[dict]:
|
|
"""Delete a connection, choosing what happens to what it feeds.
|
|
|
|
Args:
|
|
conn: Open connection inside a transaction.
|
|
row: The connection, already authorised for its owner.
|
|
sources: ``keep`` (content stays, syncing stops) or ``delete``.
|
|
tools: ``delete`` (the default: they cannot work without it) or ``keep``.
|
|
|
|
Returns:
|
|
The source rows the caller must delete with their indexes (empty
|
|
when ``sources`` is ``keep``). Index files live outside the database,
|
|
so the caller deletes them after this transaction commits.
|
|
"""
|
|
from sqlalchemy import text
|
|
|
|
from docsgpt.security.encryption import CredentialDecryptionError
|
|
|
|
repo = ConnectorSessionsRepository(conn)
|
|
connection_id = str(row["id"])
|
|
_forget_mcp_clients(row)
|
|
linked_sources = repo.list_sources(connection_id)
|
|
try:
|
|
revoke_at_provider(row, read_secrets(row))
|
|
except CredentialDecryptionError:
|
|
pass
|
|
if tools == "delete":
|
|
# Team grants go with the tools (delete trigger) and chat preferences
|
|
# by FK cascade; the owner's sharing switches have no FK.
|
|
conn.execute(
|
|
text(
|
|
"WITH gone AS (DELETE FROM user_tools WHERE connection_id = CAST(:id AS uuid) "
|
|
"AND user_id = :user_id RETURNING id) "
|
|
"DELETE FROM resource_share_settings WHERE resource_type = 'tool' "
|
|
"AND resource_id IN (SELECT id FROM gone)"
|
|
),
|
|
{"id": connection_id, "user_id": row["user_id"]},
|
|
)
|
|
else:
|
|
# The link itself goes with the row (ON DELETE SET NULL); note which
|
|
# service the tools lost, so an agent can say why they stopped.
|
|
conn.execute(
|
|
text(
|
|
"UPDATE user_tools SET status = false, "
|
|
"config = COALESCE(config, '{}'::jsonb) || jsonb_build_object('removed_connection', "
|
|
"CAST(:marker AS jsonb)) WHERE connection_id = CAST(:id AS uuid)"
|
|
),
|
|
{"id": connection_id, "marker": json.dumps(catalog.connector_key_for_row(row) or True)},
|
|
)
|
|
conn.execute(
|
|
text(
|
|
"UPDATE sources SET sync_frequency = 'never', metadata = metadata - 'sync_state' "
|
|
"WHERE connection_id = CAST(:id AS uuid)"
|
|
),
|
|
{"id": connection_id},
|
|
)
|
|
repo.delete_by_id(connection_id)
|
|
return linked_sources if sources == "delete" else []
|
|
|
|
|
|
def set_tool_permissions(
|
|
conn, user_id: str, connection_id: str, tool_id: str, permissions: dict,
|
|
) -> Optional[dict]:
|
|
"""Apply ``{action: always | ask | off}`` to a tool the user owns.
|
|
|
|
Args:
|
|
conn: Open database connection.
|
|
user_id: The caller, who must own the tool.
|
|
connection_id: The connection the tool must belong to.
|
|
tool_id: The tool to update.
|
|
permissions: Action name to ``always``, ``ask`` or ``off``.
|
|
|
|
Returns:
|
|
The updated tool, or None (nothing written) when the tool is not the
|
|
user's or belongs to another connection.
|
|
"""
|
|
from docsgpt.connectors.permissions import apply_permission
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
|
|
tools = UserToolsRepository(conn)
|
|
tool = tools.get_any(tool_id, user_id)
|
|
if tool is None or tool.get("user_id") != user_id or str(tool.get("connection_id")) != connection_id:
|
|
return None
|
|
actions = [
|
|
apply_permission(action, permissions[action.get("name")]) if action.get("name") in permissions else action
|
|
for action in (_json(tool.get("actions")) or [])
|
|
]
|
|
tools.update(str(tool["id"]), user_id, {"actions": actions})
|
|
return tools.get_any(str(tool["id"]), user_id)
|
|
|
|
|
|
def set_tool_parameters(
|
|
conn, user_id: str, connection_id: str, tool_id: str, action_name: str, pins: dict,
|
|
) -> Optional[dict]:
|
|
"""Fix or release parameters of one action of a tool the user owns.
|
|
|
|
Args:
|
|
conn: Open database connection.
|
|
user_id: The caller, who must own the tool.
|
|
connection_id: The connection the tool must belong to.
|
|
tool_id: The tool to update.
|
|
action_name: The action whose parameters change.
|
|
pins: Parameter name to the value to always use, or None to let the
|
|
model decide.
|
|
|
|
Returns:
|
|
The updated tool, or None (nothing written) when the tool is not the
|
|
user's or belongs to another connection.
|
|
|
|
Raises:
|
|
ValueError: The action or a parameter does not exist, or a value does
|
|
not fit its parameter.
|
|
"""
|
|
from docsgpt.agents.tool_pins import set_pins
|
|
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
|
|
|
tools = UserToolsRepository(conn)
|
|
tool = tools.get_any(tool_id, user_id)
|
|
if tool is None or tool.get("user_id") != user_id or str(tool.get("connection_id")) != connection_id:
|
|
return None
|
|
actions = _json(tool.get("actions")) or []
|
|
if not any(isinstance(a, dict) and a.get("name") == action_name for a in actions):
|
|
raise ValueError(f"Unknown action: {action_name}")
|
|
actions = [
|
|
set_pins(action, pins) if isinstance(action, dict) and action.get("name") == action_name else action
|
|
for action in actions
|
|
]
|
|
tools.update(str(tool["id"]), user_id, {"actions": actions})
|
|
return tools.get_any(str(tool["id"]), user_id)
|
|
|
|
|
|
def reencrypt_all(batch_size: int = 500) -> dict:
|
|
"""Rewrite every connection's credentials with the current key.
|
|
|
|
Run after rotating ENCRYPTION_SECRET_KEY with the old value in
|
|
ENCRYPTION_SECRET_KEY_PREVIOUS; afterwards the previous key can go.
|
|
|
|
Returns:
|
|
``{"rewritten": n, "current": n, "failed": n}``; failed rows (neither
|
|
key opens them) are flagged ``reconnect_needed``.
|
|
"""
|
|
from sqlalchemy import text
|
|
|
|
from docsgpt.security.encryption import CredentialDecryptionError, current_key_id, envelope_key_id
|
|
|
|
key_id = current_key_id()
|
|
counts = {"rewritten": 0, "current": 0, "failed": 0}
|
|
failed_ids: list[str] = []
|
|
last_id = "00000000-0000-0000-0000-000000000000"
|
|
while True:
|
|
with db_session() as conn:
|
|
rows = conn.execute(
|
|
text(
|
|
"SELECT * FROM connector_sessions WHERE encrypted_credentials IS NOT NULL "
|
|
"AND id > CAST(:last AS uuid) ORDER BY id LIMIT :batch FOR UPDATE"
|
|
),
|
|
{"last": last_id, "batch": batch_size},
|
|
).fetchall()
|
|
if not rows:
|
|
break
|
|
from docsgpt.storage.db.base_repository import row_to_dict
|
|
|
|
for raw in rows:
|
|
row = row_to_dict(raw)
|
|
last_id = str(row["id"])
|
|
if envelope_key_id(row["encrypted_credentials"]) == key_id:
|
|
counts["current"] += 1
|
|
continue
|
|
try:
|
|
secrets = read_secrets(row)
|
|
except CredentialDecryptionError:
|
|
counts["failed"] += 1
|
|
failed_ids.append(last_id)
|
|
continue
|
|
write_secrets(conn, row, secrets)
|
|
counts["rewritten"] += 1
|
|
for connection_id in failed_ids:
|
|
mark_reconnect_needed(connection_id, DECRYPT_ERROR)
|
|
return counts
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Admin policy
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class ConnectorDisabled(Exception):
|
|
"""An admin turned this connector (or custom MCP servers) off."""
|
|
|
|
|
|
class WritesForbidden(Exception):
|
|
"""An admin forbade write actions through this connector (GitHub's changes)."""
|
|
|
|
|
|
def custom_mcp_allowed(conn) -> bool:
|
|
"""The instance-wide "Allow custom MCP servers" switch (on unless turned off)."""
|
|
from docsgpt.storage.db.repositories.app_metadata import AppMetadataRepository
|
|
from docsgpt.storage.db.repositories.connector_policies import ALLOW_CUSTOM_MCP_KEY
|
|
|
|
return AppMetadataRepository(conn).get(ALLOW_CUSTOM_MCP_KEY) != "false"
|
|
|
|
|
|
def connector_is_enabled(policies: dict, connector_key: Optional[str]) -> bool:
|
|
"""Whether a connector is switched on.
|
|
|
|
An admin's explicit switch wins. Without one, a connector is on when it
|
|
has the server settings it needs, so one that still needs admin setup
|
|
starts off and members never see it.
|
|
|
|
Args:
|
|
policies: ``connector_key`` to policy row, from :func:`load_policies`.
|
|
connector_key: The catalog key.
|
|
"""
|
|
explicit = (policies.get(connector_key) or {}).get("enabled") if connector_key else None
|
|
if explicit is not None:
|
|
return bool(explicit)
|
|
definition = catalog.get_definition(connector_key) if connector_key else None
|
|
return definition is None or definition.configured
|
|
|
|
|
|
def load_policies(conn) -> dict[str, dict]:
|
|
"""Every connector's policy, with custom MCP and write switches folded in.
|
|
|
|
A connector that offers writes (GitHub) gets ``allow_writes``: False when
|
|
an admin forbade them.
|
|
"""
|
|
from docsgpt.storage.db.repositories.app_metadata import AppMetadataRepository
|
|
from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository, allow_writes_key
|
|
|
|
policies = dict(ConnectorPoliciesRepository(conn).all())
|
|
if not custom_mcp_allowed(conn):
|
|
policies["custom_mcp"] = {**policies.get("custom_mcp", {}), "enabled": False}
|
|
metadata = AppMetadataRepository(conn)
|
|
for definition in catalog.all_definitions():
|
|
if definition.mcp_write_url and metadata.get(allow_writes_key(definition.key)) == "false":
|
|
policies[definition.key] = {**policies.get(definition.key, {}), "allow_writes": False}
|
|
return policies
|
|
|
|
|
|
def writes_allowed(policies: dict, connector_key: Optional[str]) -> bool:
|
|
"""Whether a connection to this connector may be set up for write actions.
|
|
|
|
Args:
|
|
policies: ``connector_key`` to policy row, from :func:`load_policies`.
|
|
connector_key: The catalog key.
|
|
|
|
Returns:
|
|
False for a connector that offers no writes, or whose writes an admin
|
|
forbade.
|
|
"""
|
|
definition = catalog.get_definition(connector_key) if connector_key else None
|
|
if definition is None or not definition.mcp_write_url:
|
|
return False
|
|
return (policies.get(connector_key) or {}).get("allow_writes") is not False
|
|
|
|
|
|
def connector_enabled(conn, row: dict) -> bool:
|
|
"""Whether an admin left on the connector a connection row belongs to.
|
|
|
|
Args:
|
|
conn: Open database connection.
|
|
row: A ``connector_sessions`` row.
|
|
|
|
Returns:
|
|
False when the connector (or custom MCP servers) is turned off.
|
|
"""
|
|
return connector_is_enabled(load_policies(conn), catalog.connector_key_for_row(row))
|
|
|
|
|
|
def ensure_connector_allowed(conn, connector_key: Optional[str]) -> None:
|
|
"""Refuse a new connection to a connector an admin turned off.
|
|
|
|
Raises:
|
|
ConnectorDisabled: The connector, or custom MCP servers, are off.
|
|
"""
|
|
if not connector_key:
|
|
return
|
|
if not connector_is_enabled(load_policies(conn), connector_key):
|
|
raise ConnectorDisabled(f"{connector_key} is turned off by an admin")
|
|
|
|
|
|
def forced_credential_mode(conn, connector_key: Optional[str]) -> Optional[str]:
|
|
"""``owner`` or ``member`` when an admin forces one for this connector."""
|
|
if not connector_key:
|
|
return None
|
|
mode = (load_policies(conn).get(connector_key) or {}).get("credential_mode")
|
|
return mode if mode in ("owner", "member") else None
|