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:
arc53-machine committed 2026-09-28 17:03:09 +01:00
1 parent 667d4bbab0
commit 3511004f50
54 files changed
+4007 -754

No files matched your search

+115 -9
View File
@@ -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.
+2
View File
@@ -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."
+62 -35
View File
@@ -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 (
+1
View File
@@ -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 "
+2
View File
@@ -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",
+2
View File
@@ -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."
+296 -1
View File
@@ -11,7 +11,16 @@ and lets sources and tools point at the row they use:
resources a connection feeds. ``ON DELETE SET NULL`` keeps a source's
indexed content when its connection is removed.
Backfill (idempotent, only fills NULLs):
Connections own their credentials. Every secret of a connection (OAuth
tokens, MCP OAuth tokens and the dynamic client registration, API keys)
moves into ``encrypted_credentials``, a v2 envelope from
``docsgpt.security.encryption`` bound to the owner. Plain columns
(``status``, ``has_refresh_token``, ``scopes``, ``expires_at``) keep status
checks from ever decrypting. The unique index gains ``account_label`` so one
user can connect several accounts of the same service, and
``credential_mode`` says whose account a shared tool or source uses.
Backfill (idempotent, only fills NULLs or unconverted rows):
1. ``connector_key``, ``auth_kind``, ``display_name`` from ``provider``.
2. ``account_label`` from ``user_email`` for OAuth rows.
@@ -19,6 +28,17 @@ Backfill (idempotent, only fills NULLs):
owner's only row for ``remote_data->>'provider'``.
4. ``user_tools.connection_id`` for OAuth MCP tools, matched to the owner's
row for the tool's server base URL.
5. API-key tools (Brave, Telegram, ntfy, PostgreSQL, custom MCP with a key,
bearer token or basic auth) get one connection per distinct credential,
re-encrypted into v2. The tool keeps its v1 copy for one release so a
rollback still works; the executor prefers the connection.
6. ``token_info`` and the secret parts of ``session_data`` (``tokens``,
``client_info``) are encrypted into ``encrypted_credentials`` and removed
from the plaintext columns. This needs ``ENCRYPTION_SECRET_KEY`` set to
the value the app will run with. Downgrade decrypts them back.
7. ``credential_mode`` is ``owner`` everywhere except OAuth MCP tools, which
resolved each invoking member's own token before this migration and keep
doing so (``member``); owners can switch them in the share dialog.
Revision ID: 0038_connections
Revises: 0037_request_traces
@@ -42,7 +62,22 @@ _OAUTH_PROVIDERS = {
}
_TOOL_CONNECTORS = {
"brave": "Brave Search",
"telegram": "Telegram",
"ntfy": "ntfy",
"postgres": "PostgreSQL",
}
_MCP_SECRET_AUTH = ("api_key", "bearer", "basic")
_BATCH = 500
def upgrade() -> None:
_upgrade_links()
_upgrade_credentials()
def _upgrade_links() -> None:
op.execute(
"""
ALTER TABLE connector_sessions
@@ -119,7 +154,267 @@ def upgrade() -> None:
)
def _credential_hint(credentials: dict) -> str:
"""``…abcd``: the last four characters of the first secret, never more."""
for value in credentials.values():
if isinstance(value, str) and len(value) >= 8:
return "\u2026" + value[-4:]
return "\u2026"
def _upgrade_credentials() -> None:
op.execute(
"""
ALTER TABLE connector_sessions
ADD COLUMN IF NOT EXISTS encrypted_credentials TEXT,
ADD COLUMN IF NOT EXISTS has_refresh_token BOOLEAN NOT NULL DEFAULT false,
ADD COLUMN IF NOT EXISTS scopes JSONB NOT NULL DEFAULT '[]'::jsonb,
ADD COLUMN IF NOT EXISTS last_error TEXT,
ADD COLUMN IF NOT EXISTS last_used_at TIMESTAMPTZ;
"""
)
for table in ("sources", "user_tools"):
op.execute(
f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS credential_mode TEXT NOT NULL DEFAULT 'owner' "
f"CONSTRAINT {table}_credential_mode_chk CHECK (credential_mode IN ('owner', 'member'));"
)
op.execute("DROP INDEX IF EXISTS connector_sessions_user_endpoint_uidx;")
op.execute(
"CREATE UNIQUE INDEX IF NOT EXISTS connector_sessions_account_uidx ON connector_sessions "
"(user_id, provider, COALESCE(server_url, ''), COALESCE(account_label, ''));"
)
# Statuses become connected / pending / reconnect_needed / disconnected /
# error. MCP rows never had one; name them before their tokens are sealed.
op.execute("UPDATE connector_sessions SET status = 'connected' WHERE status = 'authorized';")
op.execute(
"""
UPDATE connector_sessions SET status =
CASE WHEN session_data ? 'tokens' THEN 'connected' ELSE 'pending' END
WHERE provider LIKE 'mcp:%' AND status IS NULL;
"""
)
bind = op.get_bind()
_encrypt_session_secrets(bind)
_link_api_key_tools(bind)
# 7: keep OAuth MCP tools on each member's own token, as before.
op.execute(
"""
UPDATE user_tools SET credential_mode = 'member'
WHERE name = 'mcp_tool' AND config->>'auth_type' = 'oauth';
"""
)
def _encrypt_session_secrets(bind) -> None:
"""6: move plaintext tokens into the owner-bound v2 envelope."""
import json
from sqlalchemy import text
from docsgpt.security.encryption import encrypt_json
while True:
rows = bind.execute(
text(
"""
SELECT id, user_id, token_info, session_data FROM connector_sessions
WHERE encrypted_credentials IS NULL
AND (token_info IS NOT NULL OR session_data ? 'tokens' OR session_data ? 'client_info')
LIMIT :batch
"""
),
{"batch": _BATCH},
).fetchall()
if not rows:
return
for row in rows:
token_info = row.token_info if isinstance(row.token_info, dict) else None
session_data = dict(row.session_data or {})
secrets = {}
if token_info:
secrets["token_info"] = token_info
for key in ("tokens", "client_info"):
if key in session_data:
secrets[key] = session_data.pop(key)
tokens = secrets.get("tokens") if isinstance(secrets.get("tokens"), dict) else {}
has_refresh = bool((token_info or {}).get("refresh_token") or tokens.get("refresh_token"))
scopes = (token_info or {}).get("scopes") or []
if isinstance(scopes, str):
scopes = scopes.split()
bind.execute(
text(
"""
UPDATE connector_sessions SET
encrypted_credentials = :blob,
has_refresh_token = :has_refresh,
scopes = CAST(:scopes AS jsonb),
token_info = NULL,
session_data = CAST(:session_data AS jsonb)
WHERE id = :id
"""
),
{
"blob": encrypt_json(secrets, row.user_id),
"has_refresh": has_refresh,
"scopes": json.dumps(list(scopes)),
"session_data": json.dumps(session_data),
"id": row.id,
},
)
def _link_api_key_tools(bind) -> None:
"""5: one connection per distinct API credential, linked from its tools."""
from urllib.parse import urlparse
from sqlalchemy import text
from docsgpt.security.encryption import decrypt_credentials, encrypt_json
rows = bind.execute(
text(
"""
SELECT id, user_id, name, custom_name, display_name, config FROM user_tools
WHERE connection_id IS NULL
AND config ? 'encrypted_credentials'
AND (
name = ANY(:names)
OR (name = 'mcp_tool' AND config->>'auth_type' = ANY(:mcp_auth))
)
"""
),
{"names": list(_TOOL_CONNECTORS), "mcp_auth": list(_MCP_SECRET_AUTH)},
).fetchall()
for row in rows:
config = row.config or {}
credentials = decrypt_credentials(config.get("encrypted_credentials") or "", row.user_id)
if not credentials:
# Written with a different key; the tool keeps its v1 copy.
continue
if row.name == "mcp_tool":
parsed = urlparse(config.get("server_url") or "")
server_url = f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else None
connector_key = "custom_mcp"
display_name = row.custom_name or row.display_name or parsed.netloc or "MCP server"
else:
server_url = None
connector_key = row.name
display_name = _TOOL_CONNECTORS[row.name]
label = _credential_hint(credentials)
existing = bind.execute(
text(
"""
SELECT id, encrypted_credentials FROM connector_sessions
WHERE user_id = :user_id AND provider = :provider
AND COALESCE(server_url, '') = COALESCE(:server_url, '')
AND COALESCE(account_label, '') = :label
"""
),
{"user_id": row.user_id, "provider": connector_key, "server_url": server_url, "label": label},
).fetchone()
if existing is None:
connection_id = bind.execute(
text(
"""
INSERT INTO connector_sessions (
user_id, provider, server_url, connector_key, display_name, account_label,
auth_kind, status, encrypted_credentials, session_data
) VALUES (
:user_id, :provider, :server_url, :provider, :display_name, :label,
'api_key', 'connected', :blob, '{}'::jsonb
) RETURNING id
"""
),
{
"user_id": row.user_id,
"provider": connector_key,
"server_url": server_url,
"display_name": display_name,
"label": label,
"blob": encrypt_json({"credentials": credentials}, row.user_id),
},
).scalar()
else:
connection_id = existing.id
bind.execute(
text("UPDATE user_tools SET connection_id = :cid WHERE id = :id"),
{"cid": connection_id, "id": row.id},
)
def _downgrade_credentials() -> None:
"""Decrypt the envelopes back into the pre-0038 plaintext columns."""
import json
from sqlalchemy import text
from docsgpt.security.encryption import CredentialDecryptionError, decrypt_json
bind = op.get_bind()
# Tools keep their v1 secrets; connections made only for them go away.
bind.execute(text("UPDATE user_tools SET connection_id = NULL WHERE connection_id IN "
"(SELECT id FROM connector_sessions WHERE auth_kind = 'api_key')"))
bind.execute(text("DELETE FROM connector_sessions WHERE auth_kind = 'api_key'"))
rows = bind.execute(
text(
"SELECT id, user_id, session_data, encrypted_credentials FROM connector_sessions "
"WHERE encrypted_credentials IS NOT NULL"
)
).fetchall()
for row in rows:
try:
secrets = decrypt_json(row.encrypted_credentials, row.user_id)
except CredentialDecryptionError:
continue
session_data = dict(row.session_data or {})
for key in ("tokens", "client_info"):
if key in secrets:
session_data[key] = secrets[key]
bind.execute(
text(
"UPDATE connector_sessions SET token_info = CAST(:token_info AS jsonb), "
"session_data = CAST(:session_data AS jsonb) WHERE id = :id"
),
{
"token_info": json.dumps(secrets["token_info"]) if "token_info" in secrets else None,
"session_data": json.dumps(session_data),
"id": row.id,
},
)
# Several accounts per provider cannot survive the old unique index: keep
# the most recently updated one.
op.execute(
"""
DELETE FROM connector_sessions c USING connector_sessions newer
WHERE c.user_id = newer.user_id AND c.provider = newer.provider
AND COALESCE(c.server_url, '') = COALESCE(newer.server_url, '')
AND (c.updated_at, c.id) < (newer.updated_at, newer.id);
"""
)
op.execute("DROP INDEX IF EXISTS connector_sessions_account_uidx;")
op.execute(
"CREATE UNIQUE INDEX IF NOT EXISTS connector_sessions_user_endpoint_uidx "
"ON connector_sessions (user_id, COALESCE(server_url, ''), provider);"
)
for table in ("sources", "user_tools"):
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS credential_mode;")
op.execute(
"""
ALTER TABLE connector_sessions
DROP COLUMN IF EXISTS last_used_at,
DROP COLUMN IF EXISTS last_error,
DROP COLUMN IF EXISTS scopes,
DROP COLUMN IF EXISTS has_refresh_token,
DROP COLUMN IF EXISTS encrypted_credentials;
"""
)
def downgrade() -> None:
_downgrade_credentials()
op.execute("DROP INDEX IF EXISTS user_tools_connection_idx;")
op.execute("DROP INDEX IF EXISTS sources_connection_idx;")
op.execute("ALTER TABLE user_tools DROP COLUMN IF EXISTS connection_id;")
+345 -1
View File
@@ -11,12 +11,16 @@ from __future__ import annotations
from flask import current_app, jsonify, make_response, request
from flask_restx import Namespace, Resource
import uuid
from docsgpt.api import api
from docsgpt.api.user.authz import ROLE_ADMIN, has_role
from docsgpt.connectors import service
from docsgpt.connectors import catalog, service
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
from docsgpt.storage.db.session import db_readonly, db_session
_FREQUENCIES = ("never", "daily", "weekly", "monthly")
connections_ns = Namespace("connections", description="Connectors and connections", path="/api")
api.add_namespace(connections_ns)
@@ -52,8 +56,64 @@ class ConnectorCatalog(Resource):
return make_response(jsonify({"success": True, "connectors": entries}), 200)
def _json_body() -> dict:
body = request.get_json(silent=True)
return body if isinstance(body, dict) else {}
def _owned(conn, connection_id: str, user_id: str):
return ConnectorSessionsRepository(conn).get_for_user(connection_id, user_id)
def _error(message: str, status: int, **extra):
return make_response(jsonify({"success": False, "error": message, **extra}), status)
@connections_ns.route("/connections")
class ConnectionsList(Resource):
@api.doc(
description=(
"Create a connection from pasted credentials: "
"{connector_key, credentials, label?}. Same credentials reuse the same connection."
)
)
def post(self):
user_id = _user_id()
if not user_id:
return _unauthorized()
body = _json_body()
definition = catalog.get_definition(body.get("connector_key"))
if definition is None or definition.auth_kind != "api_key":
return _error("This connector does not take pasted credentials", 400)
if definition.missing_settings:
return _error("This connector needs admin setup", 400, code="needs_setup")
credentials = body.get("credentials")
if not isinstance(credentials, dict):
return _error("credentials must be an object", 400)
try:
with db_session() as conn:
row, created = service.create_api_key_connection(
conn, user_id, definition, credentials, label=(body.get("label") or None),
)
except service.EncryptionKeyNotConfigured as err:
return _error(str(err), 400, code="encryption_key_default")
except ValueError as err:
return _error(str(err), 400)
except Exception as err:
current_app.logger.error(f"Error creating connection: {err}", exc_info=True)
return _error("Failed to create connection", 500)
return make_response(
jsonify(
{
"success": True,
"created": created,
"connection": service.serialize_connection(row),
"setup": dict(definition.setup),
}
),
201 if created else 200,
)
@api.doc(description="The caller's connections with status and linked resource counts")
def get(self):
user_id = _user_id()
@@ -86,6 +146,35 @@ class ConnectionDetail(Resource):
return make_response(jsonify({"success": False, "error": "Failed to load connection"}), 500)
return make_response(jsonify({"success": True, "connection": detail}), 200)
@api.doc(
description=(
"Remove a connection: {sources: keep | delete, tools: delete | keep}. "
"Kept sources keep their content and stop syncing."
)
)
def delete(self, connection_id: str):
from docsgpt.api.user.sources.routes import delete_source
user_id = _user_id()
if not user_id:
return _unauthorized()
body = _json_body()
sources_mode = body.get("sources", "keep")
tools_mode = body.get("tools", "delete")
if sources_mode not in ("keep", "delete") or tools_mode not in ("keep", "delete"):
return _error("sources and tools must be keep or delete", 400)
try:
with db_session() as conn:
row = _owned(conn, connection_id, user_id)
if row is None:
return _not_found()
to_delete = service.remove_connection(conn, row, sources=sources_mode, tools=tools_mode)
failed = [str(doc["id"]) for doc in to_delete if not delete_source(user_id, doc)]
except Exception as err:
current_app.logger.error(f"Error removing connection: {err}", exc_info=True)
return _error("Failed to remove connection", 500)
return make_response(jsonify({"success": True, "failed_sources": failed}), 200)
@connections_ns.route("/connections/<string:connection_id>/disconnect")
class ConnectionDisconnect(Resource):
@@ -104,3 +193,258 @@ class ConnectionDisconnect(Resource):
current_app.logger.error(f"Error disconnecting connection: {err}", exc_info=True)
return make_response(jsonify({"success": False, "error": "Failed to disconnect"}), 500)
return make_response(jsonify({"success": True, "connection": connection}), 200)
@connections_ns.route("/connections/<string:connection_id>/setup")
class ConnectionSetup(Resource):
@api.doc(
description=(
"Apply the connect wizard's choices: {create_tools, tool_permissions?, "
"sync?: {items, frequency, name?}}. Honours an Idempotency-Key header for the sync."
)
)
def post(self, connection_id: str):
user_id = _user_id()
if not user_id:
return _unauthorized()
body = _json_body()
try:
with db_session() as conn:
row = _owned(conn, connection_id, user_id)
if row is None:
return _not_found()
if service.normalize_status(row) != service.STATUS_CONNECTED:
return _error("Reconnect before setting up", 409, code="reconnect")
tools = []
if body.get("create_tools", True):
tools = service.ensure_connection_tools(
conn, user_id, row, permissions=body.get("tool_permissions") or None,
)
tool_payload = [service.serialize_tool(tool) for tool in tools]
sources = []
if body.get("sync"):
started = _start_sync(user_id, row, body["sync"])
if isinstance(started, tuple):
return _error(*started)
sources.append(started)
except Exception as err:
current_app.logger.error(f"Error setting up connection: {err}", exc_info=True)
return _error("Failed to set up connection", 500)
return make_response(jsonify({"success": True, "tools": tool_payload, "sources": sources}), 200)
def _start_sync(user_id: str, row: dict, sync: dict):
"""Queue the first ingest of a source synced from ``row``.
Returns the source summary, or ``(message, status)`` on a bad request.
"""
from docsgpt.api.user.sources.upload import (
_claim_task_or_get_cached,
_derive_source_id,
_read_idempotency_key,
_scoped_idempotency_key,
)
from docsgpt.api.user.tasks import ingest_connector_task, ingest_remote
definition = catalog.get_definition(catalog.connector_key_for_row(row))
if definition is None or not definition.sync_ingestor:
return ("This connector does not sync content", 400)
items = sync.get("items") or {}
if not isinstance(items, dict):
return ("items must be an object", 400)
frequency = sync.get("frequency") or definition.default_sync_frequency
if frequency not in _FREQUENCIES:
return ("Unknown sync frequency", 400)
name = (sync.get("name") or "").strip() or definition.name
idempotency_key, _ = _read_idempotency_key()
scoped_key = _scoped_idempotency_key(idempotency_key, user_id)
task_id = None
if scoped_key:
task_id, cached = _claim_task_or_get_cached(scoped_key, "connection_setup_sync")
if cached is not None:
return {"id": cached.get("source_id"), "task_id": cached.get("task_id"), "name": name}
source_id = str(_derive_source_id(scoped_key)) if scoped_key else str(uuid.uuid4())
options = {"task_id": task_id} if task_id else {}
if definition.auth_kind == "oauth":
file_ids = [str(i) for i in items.get("file_ids") or [] if i]
folder_ids = [str(i) for i in items.get("folder_ids") or [] if i]
if not file_ids and not folder_ids:
return ("Pick at least one file or folder", 400)
task = ingest_connector_task.apply_async(
kwargs={
"job_name": name,
"user": user_id,
"source_type": definition.sync_ingestor,
"connection_id": str(row["id"]),
"file_ids": file_ids,
"folder_ids": folder_ids,
"recursive": bool(items.get("recursive", True)),
"sync_frequency": frequency,
"idempotency_key": scoped_key,
"source_id": source_id,
},
**options,
)
else:
fields = {f.key for f in definition.setup_fields}
source_data = {k: v for k, v in items.items() if k in fields and v not in (None, "")}
missing = [f.label for f in definition.setup_fields if f.required and f.key not in source_data]
if missing:
return (f"Missing: {', '.join(missing)}", 400)
task = ingest_remote.apply_async(
kwargs={
"source_data": source_data,
"job_name": name,
"user": user_id,
"loader": definition.sync_ingestor,
"connection_id": str(row["id"]),
"sync_frequency": frequency,
"idempotency_key": scoped_key,
"source_id": source_id,
},
**options,
)
return {"id": source_id, "task_id": task_id or task.id, "name": name, "sync_frequency": frequency}
@connections_ns.route("/connections/<string:connection_id>/reconnect")
class ConnectionReconnect(Resource):
@api.doc(
description=(
"OAuth: returns an authorization URL for the same account. "
"API key: accepts {credentials} and replaces the stored ones."
)
)
def post(self, connection_id: str):
from docsgpt.api.connector.routes import build_authorization
user_id = _user_id()
if not user_id:
return _unauthorized()
body = _json_body()
try:
with db_readonly() as conn:
row = _owned(conn, connection_id, user_id)
if row is None:
return _not_found()
key = catalog.connector_key_for_row(row)
definition = catalog.get_definition(key)
auth_kind = row.get("auth_kind") or (definition.auth_kind if definition else None)
if auth_kind == "oauth":
started = build_authorization(row["provider"], user_id, connection_id)
return make_response(jsonify({"success": True, "kind": "oauth", **started}), 200)
if auth_kind == "mcp_oauth":
# The MCP client runs the OAuth dance (dynamic registration,
# PKCE); the frontend starts it through /api/mcp_server/test.
return make_response(
jsonify({"success": True, "kind": "mcp_oauth", "server_url": row.get("server_url")}), 200,
)
credentials = body.get("credentials")
if not isinstance(credentials, dict) or not credentials:
return _error("credentials are required", 400)
service.ensure_can_store_credentials()
with db_session() as conn:
locked = ConnectorSessionsRepository(conn).get_for_update(connection_id)
stored = service.load_secrets(locked) or {}
merged = {**(stored.get("credentials") or {}), **{k: v for k, v in credentials.items() if v}}
service.write_secrets(
conn, locked, {**stored, "credentials": merged},
status=service.STATUS_CONNECTED, last_error=None,
)
service.resume_sources(conn, connection_id)
connection = service.serialize_connection(ConnectorSessionsRepository(conn).get(connection_id))
except service.EncryptionKeyNotConfigured as err:
return _error(str(err), 400, code="encryption_key_default")
except Exception as err:
current_app.logger.error(f"Error reconnecting: {err}", exc_info=True)
return _error("Failed to reconnect", 500)
return make_response(jsonify({"success": True, "kind": "api_key", "connection": connection}), 200)
@connections_ns.route("/connections/<string:connection_id>/picker-token")
class ConnectionPickerToken(Resource):
@api.doc(description="A short-lived access token for a browser-side file picker. Owner only; never a refresh token.")
def post(self, connection_id: str):
user_id = _user_id()
if not user_id:
return _unauthorized()
with db_readonly() as conn:
row = _owned(conn, connection_id, user_id)
if row is None or (row.get("auth_kind") or "oauth") != "oauth":
return _not_found()
try:
token = service.picker_token(connection_id)
except service.ConnectionUnavailable:
return _error("Reconnect to continue", 409, code="reconnect")
except service.TransientConnectionError:
return _error("The provider is not responding. Try again.", 503)
return make_response(jsonify({"success": True, **token}), 200)
@connections_ns.route("/connections/claim")
class ConnectionClaim(Resource):
@api.doc(
description=(
"One-time link of a legacy browser session token ({provider, session_token}) "
"to the caller's connection. Removed next release."
)
)
def post(self):
user_id = _user_id()
if not user_id:
return _unauthorized()
body = _json_body()
provider, token = body.get("provider"), body.get("session_token")
if not provider or not token:
return _error("provider and session_token are required", 400)
with db_readonly() as conn:
row = service.claim_session_token(conn, user_id, str(provider), str(token))
if row is None:
return _not_found()
return make_response(jsonify({"success": True, "connection_id": str(row["id"])}), 200)
@connections_ns.route("/connections/<string:connection_id>/tools/<string:tool_id>/permissions")
class ConnectionToolPermissions(Resource):
@api.doc(description="Set per-action permissions: {permissions: {action: always | ask | off}}")
def put(self, connection_id: str, tool_id: str):
from docsgpt.connectors.permissions import PERMISSIONS
user_id = _user_id()
if not user_id:
return _unauthorized()
permissions = _json_body().get("permissions")
if not isinstance(permissions, dict) or any(p not in PERMISSIONS for p in permissions.values()):
return _error("permissions must map action names to always, ask or off", 400)
with db_session() as conn:
row = _owned(conn, connection_id, user_id)
if row is None:
return _not_found()
tool = service.set_tool_permissions(conn, user_id, tool_id, permissions)
if tool is None or str(tool.get("connection_id")) != connection_id:
return _not_found()
payload = service.serialize_tool(tool)
return make_response(jsonify({"success": True, "tool": payload}), 200)
@connections_ns.route("/connections/<string:connection_id>/refresh-tools")
class ConnectionRefreshTools(Resource):
@api.doc(description="MCP: re-scan the server's actions and return what was added and removed")
def post(self, connection_id: str):
from docsgpt.connectors.mcp import refresh_mcp_tools
user_id = _user_id()
if not user_id:
return _unauthorized()
with db_readonly() as conn:
row = _owned(conn, connection_id, user_id)
if row is None:
return _not_found()
try:
diff = refresh_mcp_tools(user_id, row)
except service.ConnectionUnavailable:
return _error("Reconnect to continue", 409, code="reconnect")
except Exception as err:
current_app.logger.error(f"Error refreshing MCP tools: {err}", exc_info=True)
return _error("Failed to refresh tools", 502)
return make_response(jsonify({"success": True, **diff}), 200)
+108 -109
View File
@@ -1,7 +1,6 @@
import base64
import html
import json
import uuid
from typing import Optional
from urllib.parse import urlencode, urlsplit
@@ -20,11 +19,11 @@ from docsgpt.api import api
from docsgpt.api.user.tasks import (
ingest_connector_task,
)
from docsgpt.connectors import service
from docsgpt.core.settings import settings
from docsgpt.parser.connectors.connector_creator import ConnectorCreator
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
owns_connector_session,
)
from docsgpt.storage.db.repositories.sources import SourcesRepository
from docsgpt.storage.db.session import db_readonly, db_session
@@ -100,6 +99,7 @@ def _js_literal(value) -> str:
def _render_callback_page(
status: str, message: str, provider_raw: str, session_token: str = "", user_email: str = "",
connection_id: str = "",
):
"""Popup page that reports an OAuth result to the opener on allowed origins only."""
status = status if status in ("success", "error", "cancelled") else "error"
@@ -110,7 +110,14 @@ def _render_callback_page(
)
payload = None
if provider_key and status == "success" and session_token:
payload = {"type": f"{provider_key}_auth_success", "session_token": session_token, "user_email": user_email}
payload = {
"type": f"{provider_key}_auth_success",
# The connection id is what current frontends use; the session
# token is kept for one release for frontends from before it.
"connection_id": connection_id,
"session_token": session_token,
"user_email": user_email,
}
elif provider_key and status == "error":
# The frontend shows its own localized failure message; cancellations are
# reported when the popup closes.
@@ -170,6 +177,27 @@ def _render_callback_page(
def build_authorization(provider: str, user_id: str, connection_id: Optional[str] = None) -> dict:
"""Start an OAuth sign-in for ``provider`` and return its authorization URL.
Raises:
service.EncryptionKeyNotConfigured: See ``ensure_can_store_credentials``.
service.ConnectionUnavailable: ``connection_id`` is not the caller's.
"""
service.ensure_can_store_credentials()
with db_session() as conn:
session_row = service.begin_oauth(conn, user_id, provider, connection_id)
state = base64.urlsafe_b64encode(
json.dumps({"provider": provider, "object_id": str(session_row["id"])}).encode()
).decode()
auth = ConnectorCreator.create_auth(provider)
return {
"authorization_url": auth.get_authorization_url(state=state),
"state": state,
"callback_origin": _origin_of(settings.CONNECTOR_REDIRECT_BASE_URI),
}
@connectors_ns.route("/api/connectors/auth")
class ConnectorAuth(Resource):
@api.doc(description="Get connector OAuth authorization URL", params={"provider": "Connector provider (e.g., google_drive)"})
@@ -186,20 +214,14 @@ class ConnectorAuth(Resource):
if not decoded_token:
return make_response(jsonify({"success": False, "error": "Unauthorized"}), 401)
user_id = decoded_token.get('sub')
with db_session() as conn:
session_row = ConnectorSessionsRepository(conn).upsert(
user_id, provider, status="pending",
try:
started = build_authorization(provider, user_id, request.args.get("connection_id") or None)
except service.EncryptionKeyNotConfigured as err:
return make_response(
jsonify({"success": False, "error": str(err), "code": "encryption_key_default"}), 400,
)
session_pg_id = str(session_row["id"])
state_dict = {
"provider": provider,
"object_id": session_pg_id,
}
state = base64.urlsafe_b64encode(json.dumps(state_dict).encode()).decode()
auth = ConnectorCreator.create_auth(provider)
authorization_url = auth.get_authorization_url(state=state)
except service.ConnectionUnavailable:
return make_response(jsonify({"success": False, "error": "Connection not found"}), 404)
# The popup drops results for origins outside the allowlist, which the
# user only sees as a cancelled sign-in; name the missing origin here.
request_origin = _origin_of(request.headers.get("Origin"))
@@ -208,12 +230,7 @@ class ConnectorAuth(Resource):
f"Connector sign-in requested from {request_origin}, which cannot receive the result; "
"add it to CONNECTOR_ALLOWED_ORIGINS"
)
return make_response(jsonify({
"success": True,
"authorization_url": authorization_url,
"state": state,
"callback_origin": _origin_of(settings.CONNECTOR_REDIRECT_BASE_URI),
}), 200)
return make_response(jsonify({"success": True, **started}), 200)
except Exception as e:
current_app.logger.error(f"Error generating connector auth URL: {e}", exc_info=True)
return make_response(jsonify({"success": False, "error": "Failed to generate authorization URL"}), 500)
@@ -269,13 +286,11 @@ class ConnectorsCallback(Resource):
auth = ConnectorCreator.create_auth(provider)
token_info = auth.exchange_code_for_tokens(authorization_code)
session_token = str(uuid.uuid4())
try:
if provider == "google_drive":
credentials = auth.create_credentials_from_token_info(token_info)
service = auth.build_drive_service(credentials)
user_info = service.about().get(fields="user").execute()
drive_service = auth.build_drive_service(credentials)
user_info = drive_service.about().get(fields="user").execute()
user_email = user_info.get('user', {}).get('emailAddress', 'Connected User')
else:
user_email = token_info.get('user_info', {}).get('email', 'Connected User')
@@ -288,29 +303,23 @@ class ConnectorsCallback(Resource):
# ``object_id`` in the OAuth state is the PG session row
# UUID (new flow) or a legacy Mongo ObjectId (pre-cutover
# issued state). Try UUID update first; fall back to
# legacy id path.
patch = {
"session_token": session_token,
"token_info": sanitized_token_info,
"user_email": user_email,
"status": "authorized",
}
# issued state).
with db_session() as conn:
repo = ConnectorSessionsRepository(conn)
if state_object_id:
value = str(state_object_id)
updated = False
if len(value) == 36 and "-" in value:
updated = repo.update(value, patch)
if not updated:
repo.update_by_legacy_id(value, patch)
value = str(state_object_id or "")
state_row = repo.get(value) or (repo.get_by_legacy_id(value) if value else None)
if state_row is None or state_row.get("provider") != provider:
raise ValueError("OAuth state names no pending connection")
connection = service.complete_oauth(
conn, state_row, provider, sanitized_token_info, user_email,
)
# Render instead of redirecting so the session token never
# lands in a URL (browser history, access logs, Referer).
return _render_callback_page(
"success", "Authentication successful", provider,
session_token=session_token, user_email=user_email,
session_token=connection.get("session_token") or "", user_email=user_email,
connection_id=str(connection["id"]),
)
except Exception as e:
@@ -333,7 +342,8 @@ class ConnectorsCallback(Resource):
class ConnectorFiles(Resource):
@api.expect(api.model("ConnectorFilesModel", {
"provider": fields.String(required=True),
"session_token": fields.String(required=True),
"connection_id": fields.String(required=False),
"session_token": fields.String(required=False, description="Legacy; use connection_id"),
"folder_id": fields.String(required=False),
"limit": fields.Integer(required=False),
"page_token": fields.String(required=False),
@@ -344,26 +354,29 @@ class ConnectorFiles(Resource):
try:
data = request.get_json()
provider = data.get('provider')
session_token = data.get('session_token')
limit = data.get('limit', 10)
if not provider or not session_token:
return make_response(jsonify({"success": False, "error": "provider and session_token are required"}), 400)
if not provider or not (data.get('connection_id') or data.get('session_token')):
return make_response(
jsonify({"success": False, "error": "provider and connection_id are required"}), 400,
)
decoded_token = request.decoded_token
if not decoded_token:
return make_response(jsonify({"success": False, "error": "Unauthorized"}), 401)
user = decoded_token.get('sub')
with db_readonly() as conn:
session = ConnectorSessionsRepository(conn).get_by_session_token(
session_token,
)
if not owns_connector_session(session, user, provider):
session = service.resolve_request_connection(user, provider, data)
if session is None:
return make_response(jsonify({"success": False, "error": "Invalid or unauthorized session"}), 401)
loader = ConnectorCreator.create_connector(provider, session_token)
try:
loader = ConnectorCreator.create_connector(provider, connection_id=str(session["id"]))
except service.ConnectionUnavailable:
return make_response(
jsonify({"success": False, "error": "Reconnect to continue", "reconnect": True}), 401,
)
generic_keys = {'provider', 'session_token'}
generic_keys = {'provider', 'session_token', 'connection_id'}
input_config = {
k: v for k, v in data.items() if k not in generic_keys
}
@@ -408,65 +421,49 @@ class ConnectorFiles(Resource):
@connectors_ns.route("/api/connectors/validate-session")
class ConnectorValidateSession(Resource):
@api.expect(api.model("ConnectorValidateSessionModel", {"provider": fields.String(required=True), "session_token": fields.String(required=True)}))
@api.doc(description="Validate connector session token and return user info and access token")
@api.expect(api.model("ConnectorValidateSessionModel", {
"provider": fields.String(required=True),
"connection_id": fields.String(required=False),
"session_token": fields.String(required=False, description="Legacy; use connection_id"),
}))
@api.doc(description="Validate a connection and return the account and a short-lived access token")
def post(self):
try:
data = request.get_json()
data = request.get_json() or {}
provider = data.get('provider')
session_token = data.get('session_token')
if not provider or not session_token:
return make_response(jsonify({"success": False, "error": "provider and session_token are required"}), 400)
if not provider or not (data.get('connection_id') or data.get('session_token')):
return make_response(
jsonify({"success": False, "error": "provider and connection_id are required"}), 400,
)
decoded_token = request.decoded_token
if not decoded_token:
return make_response(jsonify({"success": False, "error": "Unauthorized"}), 401)
user = decoded_token.get('sub')
with db_readonly() as conn:
session = ConnectorSessionsRepository(conn).get_by_session_token(
session_token,
)
if not owns_connector_session(session, user, provider) or not session.get("token_info"):
session = service.resolve_request_connection(user, provider, data)
if session is None:
return make_response(jsonify({"success": False, "error": "Invalid or expired session"}), 401)
token_info = session["token_info"]
auth = ConnectorCreator.create_auth(provider)
is_expired = auth.is_token_expired(token_info)
if is_expired and token_info.get('refresh_token'):
try:
refreshed_token_info = auth.refresh_access_token(token_info.get('refresh_token'))
sanitized_token_info = auth.sanitize_token_info(refreshed_token_info)
with db_session() as conn:
repo = ConnectorSessionsRepository(conn)
row = repo.get_by_session_token(session_token)
if row:
repo.update(str(row["id"]), {"token_info": sanitized_token_info})
token_info = sanitized_token_info
is_expired = False
except Exception as refresh_error:
current_app.logger.error(f"Failed to refresh token: {refresh_error}")
if is_expired:
try:
token = service.picker_token(str(session["id"]))
except service.ConnectionUnavailable:
return make_response(jsonify({
"success": False,
"expired": True,
"error": "Session token has expired. Please reconnect."
}), 401)
except service.TransientConnectionError:
return make_response(
jsonify({"success": False, "error": "The provider is not responding. Try again."}), 503,
)
_base_fields = {"access_token", "refresh_token", "token_uri", "expiry"}
provider_extras = {k: v for k, v in token_info.items() if k not in _base_fields}
response_data = {
return make_response(jsonify({
"success": True,
"expired": False,
"user_email": session.get('user_email', 'Connected User'),
"access_token": token_info.get('access_token'),
**provider_extras,
}
return make_response(jsonify(response_data), 200)
"connection_id": str(session["id"]),
"user_email": session.get('account_label') or session.get('user_email') or 'Connected User',
**token,
}), 200)
except Exception as e:
current_app.logger.error(f"Error validating connector session: {e}", exc_info=True)
return make_response(jsonify({"success": False, "error": "Failed to validate session"}), 500)
@@ -483,15 +480,13 @@ class ConnectorDisconnect(Resource):
try:
data = request.get_json()
provider = data.get('provider')
session_token = data.get('session_token')
if not provider:
return make_response(jsonify({"success": False, "error": "provider is required"}), 400)
if session_token:
session = service.resolve_request_connection(decoded_token.get('sub'), provider, data)
if session is not None:
with db_session() as conn:
ConnectorSessionsRepository(conn).delete_by_session_token(
session_token, decoded_token.get('sub'),
)
service.disconnect(conn, session)
return make_response(jsonify({"success": True}), 200)
except Exception as e:
@@ -506,7 +501,10 @@ class ConnectorSync(Resource):
"ConnectorSyncModel",
{
"source_id": fields.String(required=True, description="Source ID to sync"),
"session_token": fields.String(required=True, description="Authentication token")
"connection_id": fields.String(
required=False, description="Connection to sync with; defaults to the source's own",
),
"session_token": fields.String(required=False, description="Legacy; use connection_id")
},
)
)
@@ -517,16 +515,15 @@ class ConnectorSync(Resource):
return make_response(jsonify({"success": False}), 401)
try:
data = request.get_json()
data = request.get_json() or {}
source_id = data.get('source_id')
session_token = data.get('session_token')
if not all([source_id, session_token]):
if not source_id:
return make_response(
jsonify({
"success": False,
"error": "source_id and session_token are required"
}),
"error": "source_id is required"
}),
400
)
user_id = decoded_token.get('sub')
@@ -562,9 +559,11 @@ class ConnectorSync(Resource):
400
)
with db_readonly() as conn:
session = ConnectorSessionsRepository(conn).get_by_session_token(session_token)
if not owns_connector_session(session, user_id, source_type):
lookup = dict(data)
if not (lookup.get('connection_id') or lookup.get('session_token')) and source.get('connection_id'):
lookup['connection_id'] = str(source['connection_id'])
session = service.resolve_request_connection(user_id, source_type, lookup)
if session is None:
return make_response(
jsonify({"success": False, "error": "Invalid or unauthorized session"}),
401,
@@ -580,7 +579,7 @@ class ConnectorSync(Resource):
job_name=source.get('name'),
user=decoded_token.get('sub'),
source_type=source_type,
session_token=session_token,
connection_id=str(session["id"]),
file_ids=file_ids,
folder_ids=folder_ids,
recursive=recursive,
+58 -45
View File
@@ -231,6 +231,63 @@ class PaginatedSources(Resource):
return make_response(jsonify({"success": False}), 400)
def delete_source(user: str, doc: dict) -> bool:
"""Delete a source's index, stored files and row. Returns whether it worked.
Args:
user: The owner deleting it.
doc: The source row, already authorised for ``user``.
"""
storage = StorageCreator.get_storage()
resolved_id = str(doc["id"])
source_id = resolved_id
try:
if settings.VECTOR_STORE == "faiss":
index_path = f"indexes/{resolved_id}"
# index.pkl is the legacy sidecar; index.json the current one.
# Older sources have only the former, so clear whichever exist.
for index_file in ("index.faiss", "index.json", "index.pkl"):
if storage.file_exists(f"{index_path}/{index_file}"):
storage.delete_file(f"{index_path}/{index_file}")
else:
vectorstore = VectorCreator.create_vectorstore(
settings.VECTOR_STORE, source_id=source_id
)
vectorstore.delete_index()
if "file_path" in doc and doc["file_path"]:
file_path = doc["file_path"]
if storage.is_directory(file_path):
files = storage.list_files(file_path)
for f in files:
storage.delete_file(f)
else:
storage.delete_file(file_path)
except FileNotFoundError:
pass
except Exception as err:
current_app.logger.error(
f"Error deleting files and indexes: {err}", exc_info=True
)
return False
try:
with db_session() as conn:
SourcesRepository(conn).delete(resolved_id, user)
record_event(
conn,
"source.deleted",
actor=user,
source_id=resolved_id,
name=doc.get("name"),
)
except Exception as err:
current_app.logger.error(
f"Error deleting source row: {err}", exc_info=True
)
return False
return True
@sources_ns.route("/delete_old")
class DeleteOldIndexes(Resource):
@api.doc(
@@ -255,51 +312,7 @@ class DeleteOldIndexes(Resource):
return make_response(jsonify({"success": False}), 400)
if not doc:
return make_response(jsonify({"status": "not found"}), 404)
storage = StorageCreator.get_storage()
resolved_id = str(doc["id"])
try:
if settings.VECTOR_STORE == "faiss":
index_path = f"indexes/{resolved_id}"
# index.pkl is the legacy sidecar; index.json the current one.
# Older sources have only the former, so clear whichever exist.
for index_file in ("index.faiss", "index.json", "index.pkl"):
if storage.file_exists(f"{index_path}/{index_file}"):
storage.delete_file(f"{index_path}/{index_file}")
else:
vectorstore = VectorCreator.create_vectorstore(
settings.VECTOR_STORE, source_id=resolved_id
)
vectorstore.delete_index()
if "file_path" in doc and doc["file_path"]:
file_path = doc["file_path"]
if storage.is_directory(file_path):
files = storage.list_files(file_path)
for f in files:
storage.delete_file(f)
else:
storage.delete_file(file_path)
except FileNotFoundError:
pass
except Exception as err:
current_app.logger.error(
f"Error deleting files and indexes: {err}", exc_info=True
)
return make_response(jsonify({"success": False}), 400)
try:
with db_session() as conn:
SourcesRepository(conn).delete(resolved_id, user)
record_event(
conn,
"source.deleted",
actor=user,
source_id=resolved_id,
name=doc.get("name"),
)
except Exception as err:
current_app.logger.error(
f"Error deleting source row: {err}", exc_info=True
)
if not delete_source(user, doc):
return make_response(jsonify({"success": False}), 400)
return make_response(jsonify({"success": True}), 200)
+52 -12
View File
@@ -28,7 +28,6 @@ from docsgpt.security.zip_archive import (
)
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
owns_connector_session,
)
from docsgpt.storage.db.repositories.idempotency import IdempotencyRepository
from docsgpt.storage.db.repositories.sources import SourcesRepository
@@ -451,6 +450,43 @@ class UploadFile(Resource):
return make_response(jsonify(response_payload), 200)
def _remote_credentials(user, source, config):
"""Split an S3 / Reddit request into loader config and the connection holding its keys.
A request naming a ``connection_id`` uses that connection's stored keys.
A request carrying keys (the form before connections) stores them on a
connection, so they are entered once and never land in
``sources.remote_data``. When a multi-user install still runs on the
public default encryption key, the keys stay with the source as before.
Returns:
``(source_data, connection_id, error_response)``.
"""
from docsgpt.connectors import catalog, service
definition = catalog.get_definition(source)
credential_keys = {f.key for f in definition.credential_fields}
public = {k: v for k, v in config.items() if k not in credential_keys and k != "connection_id"}
connection_id = config.get("connection_id")
if connection_id:
with db_readonly() as conn:
row = ConnectorSessionsRepository(conn).get_for_user(str(connection_id), user)
if row is None or catalog.connector_key_for_row(row) != source:
return None, None, make_response(
jsonify({"success": False, "error": "Invalid or unauthorized connection"}), 401,
)
return public, str(row["id"]), None
provided = {k: config[k] for k in credential_keys if config.get(k) not in (None, "")}
if not provided:
return config, None, None
try:
with db_session() as conn:
row, _ = service.create_api_key_connection(conn, user, definition, provided)
except (service.EncryptionKeyNotConfigured, ValueError):
return config, None, None
return public, str(row["id"]), None
@sources_upload_ns.route("/remote")
class UploadRemote(Resource):
@api.expect(
@@ -520,32 +556,35 @@ class UploadRemote(Resource):
try:
config = json.loads(data["data"])
source_data = None
connection_id = None
if data["source"] == "github":
source_data = config.get("repo_url")
elif data["source"] in ["crawler", "url", "sitemap"]:
source_data = config.get("url")
elif data["source"] == "reddit":
source_data = config
elif data["source"] == "s3":
source_data = config
elif data["source"] in ("reddit", "s3"):
source_data, connection_id, error = _remote_credentials(user, data["source"], config)
if error is not None:
if scoped_key:
_release_claim(scoped_key)
return error
elif data["source"] in ConnectorCreator.get_supported_connectors():
session_token = config.get("session_token")
if not session_token:
if not (config.get("connection_id") or config.get("session_token")):
if scoped_key:
_release_claim(scoped_key)
return make_response(
jsonify(
{
"success": False,
"error": f"Missing session_token in {data['source']} configuration",
"error": f"Missing connection_id in {data['source']} configuration",
}
),
400,
)
with db_readonly() as conn:
connector_session = ConnectorSessionsRepository(conn).get_by_session_token(session_token)
if not owns_connector_session(connector_session, user, data["source"]):
from docsgpt.connectors import service as connection_service
connector_session = connection_service.resolve_request_connection(user, data["source"], config)
if connector_session is None:
if scoped_key:
_release_claim(scoped_key)
return make_response(
@@ -576,7 +615,7 @@ class UploadRemote(Resource):
"job_name": data["name"],
"user": user,
"source_type": data["source"],
"session_token": session_token,
"connection_id": str(connector_session["id"]),
"file_ids": file_ids,
"folder_ids": folder_ids,
"recursive": config.get("recursive", False),
@@ -611,6 +650,7 @@ class UploadRemote(Resource):
remote_kwargs = {
"kwargs": {
"source_data": source_data,
"connection_id": connection_id,
"job_name": data["name"],
"user": user,
"loader": data["source"],
+24 -2
View File
@@ -79,6 +79,12 @@ DURABLE_TASK = dict(
)
def _connection_unavailable():
from docsgpt.connectors.service import ConnectionUnavailable
return ConnectionUnavailable
def durable_task(**overrides) -> Dict:
"""Return ``DURABLE_TASK`` with per-task overrides applied.
@@ -167,13 +173,16 @@ def ingest(
@with_idempotency(task_name="ingest_remote", on_poison=_emit_ingest_poison_event)
def ingest_remote(
self, source_data, job_name, user, loader,
config=None, idempotency_key=None, source_id=None,
config=None, idempotency_key=None, source_id=None, connection_id=None,
sync_frequency="never",
):
resp = remote_worker(
self, source_data, job_name, user, loader,
sync_frequency=sync_frequency,
config=config,
idempotency_key=idempotency_key,
source_id=source_id,
connection_id=connection_id,
)
return resp
@@ -249,6 +258,15 @@ def schedule_syncs(self, frequency):
return resp
@celery.task(bind=True, acks_late=True, autoretry_for=(Exception,), max_retries=3, retry_backoff=60,
dont_autoretry_for=(_connection_unavailable(),))
def sync_connector_source(self, source_id):
"""Re-sync one connector source from its connection, with no browser involved."""
from docsgpt.worker import sync_connector_source as run
return run(self, source_id)
@celery.task(bind=True)
def sync_source(
self,
@@ -408,7 +426,9 @@ except Exception:
pass
@celery.task(**DURABLE_TASK)
# A revoked or disconnected connection will not heal by retrying; the
# service has already paused the source and told its owner to reconnect.
@celery.task(**durable_task(dont_autoretry_for=(DocumentParseError, _connection_unavailable())))
@with_idempotency(
task_name="ingest_connector_task", on_poison=_emit_ingest_poison_event,
)
@@ -428,6 +448,7 @@ def ingest_connector_task(
config=None,
idempotency_key=None,
source_id=None,
connection_id=None,
):
from docsgpt.worker import ingest_connector
@@ -437,6 +458,7 @@ def ingest_connector_task(
user,
source_type,
session_token=session_token,
connection_id=connection_id,
file_ids=file_ids,
folder_ids=folder_ids,
recursive=recursive,
+70 -36
View File
@@ -75,6 +75,43 @@ def _validate_mcp_server_url(config: dict) -> None:
raise ValueError(f"Invalid server URL: {exc}") from exc
def _mcp_connection(user, config, auth_type, auth_credentials, display_name):
"""The connection an MCP tool runs with; created on first save.
OAuth servers already have one (the sign-in stored its tokens there).
Key, bearer and basic auth store their secret on a connection; servers
with no auth get a credential-less connection so they still appear on
the Connectors page. Returns None when a multi-user install runs on the
default encryption key, which keeps the legacy per-tool secret.
"""
from docsgpt.connectors import catalog, service
base = catalog.base_url(config.get("server_url"))
if not base:
return None
if auth_type == "oauth":
with db_readonly() as conn:
row = service._mcp_row(conn, user, base, None)
return str(row["id"]) if row else None
definition = catalog.get_definition("custom_mcp")
host = base.split("://")[-1]
try:
with db_session() as conn:
if auth_credentials:
row, _ = service.create_api_key_connection(
conn, user, definition, auth_credentials, server_url=base, display_name=display_name,
)
else:
repo = ConnectorSessionsRepository(conn)
row = repo.find_account(user, "custom_mcp", server_url=base, account_label=host) or repo.create(
user, "custom_mcp", connector_key="custom_mcp", auth_kind="none",
display_name=display_name, account_label=host, server_url=base,
)
except service.EncryptionKeyNotConfigured:
return None
return str(row["id"]) if row else None
@tools_mcp_ns.route("/mcp_server/test")
class TestMCPServerConfig(Resource):
@api.expect(
@@ -283,9 +320,19 @@ class MCPServerSave(Resource):
"redirect_uri",
]:
storage_config.pop(field, None)
transformed_actions = transform_actions(actions_metadata)
from docsgpt.connectors.permissions import apply_default_permissions
transformed_actions = apply_default_permissions(
"mcp_tool", transform_actions(actions_metadata),
)
display_name = data["displayName"]
connection_id = _mcp_connection(
user, storage_config, auth_type, auth_credentials, display_name,
) or (str(existing_doc["connection_id"]) if existing_doc and existing_doc.get("connection_id") else None)
if connection_id and auth_type != "oauth":
# The secret lives on the connection only.
storage_config.pop("encrypted_credentials", None)
description = f"MCP Server: {storage_config.get('server_url', 'Unknown')}"
status_bool = bool(data.get("status", True))
@@ -301,6 +348,7 @@ class MCPServerSave(Resource):
"config": storage_config,
"actions": transformed_actions,
"status": status_bool,
"connection_id": connection_id,
},
)
saved_id = str(existing_doc["id"])
@@ -328,6 +376,7 @@ class MCPServerSave(Resource):
"config": storage_config,
"actions": transformed_actions,
"status": status_bool,
"connection_id": connection_id,
},
)
saved_id = str(existing_by_name["id"])
@@ -347,6 +396,7 @@ class MCPServerSave(Resource):
config_requirements={},
actions=transformed_actions,
status=status_bool,
connection_id=connection_id,
)
saved_id = str(created["id"])
response_data = {
@@ -461,50 +511,34 @@ class MCPAuthStatus(Resource):
jsonify({"success": True, "statuses": {}}), 200
)
oauth_server_urls: dict = {}
from docsgpt.connectors import service
# Read from connection status alone: status checks never
# decrypt credentials.
statuses: dict = {}
for tool in mcp_tools:
tool_id = str(tool["id"])
config = tool.get("config") or {}
auth_type = config.get("auth_type", "none")
if auth_type == "oauth":
server_url = config.get("server_url", "")
if server_url:
parsed = urlparse(server_url)
base_url = f"{parsed.scheme}://{parsed.netloc}"
oauth_server_urls[tool_id] = base_url
row = None
if tool.get("connection_id"):
row = sessions_repo.get(str(tool["connection_id"]))
elif auth_type == "oauth" and config.get("server_url"):
parsed = urlparse(config["server_url"])
row = sessions_repo.get_by_user_provider(
user, service.mcp_provider(f"{parsed.scheme}://{parsed.netloc}"),
)
if row is not None:
connected = service.normalize_status(row) == service.STATUS_CONNECTED
if auth_type == "oauth" or not connected:
statuses[tool_id] = "connected" if connected else "needs_auth"
else:
statuses[tool_id] = "needs_auth"
statuses[tool_id] = "configured"
elif auth_type == "oauth":
statuses[tool_id] = "needs_auth"
else:
statuses[tool_id] = "configured"
if oauth_server_urls:
# Look up a session per distinct base URL. MCP sessions
# are stored with ``provider = "mcp:<server_url>"``
# and the URL in ``server_url``; reuse the repo's
# per-URL accessor rather than an ad-hoc $in query.
url_has_tokens: dict = {}
for base_url in set(oauth_server_urls.values()):
session = sessions_repo.get_by_user_and_server_url(
user, base_url,
)
tokens = (
(session or {}).get("session_data", {}) or {}
).get("tokens", {}) or {}
# MCP code also stashes tokens into token_info on
# the row; consider either present as "connected".
token_info = (session or {}).get("token_info") or {}
url_has_tokens[base_url] = bool(
tokens.get("access_token")
or token_info.get("access_token")
)
for tool_id, base_url in oauth_server_urls.items():
if url_has_tokens.get(base_url):
statuses[tool_id] = "connected"
else:
statuses[tool_id] = "needs_auth"
return make_response(jsonify({"success": True, "statuses": statuses}), 200)
except Exception as e:
current_app.logger.error(
+88
View File
@@ -271,6 +271,10 @@ class GetTools(Resource):
):
tool_copy["config"]["has_encrypted_credentials"] = True
tool_copy["config"].pop("encrypted_credentials", None)
if tool_copy.get("connection_id"):
# The secret lives on the connection; the form must not
# ask for it again.
tool_copy.setdefault("config", {})["has_encrypted_credentials"] = True
tool_copy["ownership"] = ownership
return tool_copy
@@ -383,6 +387,11 @@ class CreateTool(Resource):
f"Error getting tool actions: {err}", exc_info=True
)
return make_response(jsonify({"success": False}), 400)
definition = definition_for_tool(data["name"])
if definition is not None:
connected = _create_connected_tool(user, data, definition, tool_instance)
if connected is not None:
return connected
try:
config_requirements = tool_instance.get_config_requirements()
if config_requirements:
@@ -422,6 +431,74 @@ class CreateTool(Resource):
return make_response(jsonify({"id": new_id}), 200)
def _create_connected_tool(user, data, definition, tool_instance):
"""Create a service tool whose secret lives on a connection, not the tool.
Uses ``connection_id`` when given, otherwise stores the pasted secret on
a connection (reusing an identical one). Returns None to fall back to the
legacy path when a multi-user install still runs on the default key.
"""
from docsgpt.connectors import catalog as connector_catalog
from docsgpt.connectors import service as connection_service
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
config_requirements = tool_instance.get_config_requirements()
public, secrets = connection_service.split_secrets(data.get("config") or {}, config_requirements)
connection_id = data.get("connection_id")
if not connection_id:
validation_errors = _validate_config(data.get("config") or {}, config_requirements)
if validation_errors:
return make_response(
jsonify({"success": False, "message": "Validation failed", "errors": validation_errors}), 400,
)
try:
with db_session() as conn:
if connection_id:
connection = ConnectorSessionsRepository(conn).get_for_user(str(connection_id), user)
if connection is None or connector_catalog.connector_key_for_row(connection) != definition.key:
return make_response(jsonify({"success": False, "message": "Connection not found"}), 404)
else:
connection, _ = connection_service.create_api_key_connection(conn, user, definition, secrets)
created = connection_service.create_tool_for_connection(
conn,
user,
connection,
template=data["name"],
display_name=data.get("customName") or data.get("displayName") or definition.name,
config=public,
status=bool(data.get("status", True)),
)
except connection_service.EncryptionKeyNotConfigured:
return None
except ValueError as err:
return make_response(jsonify({"success": False, "message": str(err)}), 400)
return make_response(jsonify({"id": str(created["id"]), "connection_id": str(connection["id"])}), 200)
def _update_connection_secrets(conn, user, tool_doc, config, config_requirements):
"""Write changed secrets onto the tool's connection; return the tool's public config.
Returns None when the caller does not own the connection (an editor on a
team share may change actions, never the owner's credentials).
"""
from docsgpt.connectors import service as connection_service
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
public, secrets = connection_service.split_secrets(config or {}, config_requirements)
if secrets:
connection = ConnectorSessionsRepository(conn).get_for_user(str(tool_doc["connection_id"]), user)
if connection is None:
return None
stored = connection_service.read_secrets(connection)
credentials = {**(stored.get("credentials") or {}), **secrets}
connection_service.write_secrets(
conn, connection, {**stored, "credentials": credentials},
status=connection_service.STATUS_CONNECTED, last_error=None,
)
connection_service.resume_sources(conn, str(connection["id"]))
return public
@tools_ns.route("/update_tool")
class UpdateTool(Resource):
@api.expect(
@@ -528,6 +605,17 @@ class UpdateTool(Resource):
)
existing_config = tool_doc.get("config", {}) or {}
has_existing_secrets = "encrypted_credentials" in existing_config
if tool_doc.get("connection_id"):
has_existing_secrets = True
new_config = _update_connection_secrets(
conn, user, tool_doc, data["config"], config_requirements,
)
if new_config is None:
return make_response(
jsonify({"success": False, "message": "Only the owner can change the credentials"}),
403,
)
data = {**data, "config": new_config}
if config_requirements:
validation_errors = _validate_config(
+30
View File
@@ -285,6 +285,28 @@ def _add_deploy_commands(commands) -> None:
dev.set_defaults(func=_deploy("dev"), deploy=True)
def _connectors(args: argparse.Namespace) -> int:
"""``docsgpt connectors reencrypt``: move every credential onto the current key."""
if getattr(args, "connectors_action", None) != "reencrypt":
print("usage: docsgpt connectors reencrypt", file=sys.stderr)
return 2
from docsgpt.connectors.service import reencrypt_all
counts = reencrypt_all()
print(
f"docsgpt: re-encrypted {counts['rewritten']} connection(s), "
f"{counts['current']} already current, {counts['failed']} unreadable",
file=sys.stderr,
)
if counts["failed"]:
print(
"docsgpt: unreadable connections were marked 'Reconnect needed'; "
"their owners must reconnect them.",
file=sys.stderr,
)
return 1 if counts["failed"] else 0
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(prog="docsgpt", description="DocsGPT: private AI for agents, assistants and search.")
parser.add_argument("--version", action="version", version=f"docsgpt {__version__}")
@@ -315,6 +337,14 @@ def build_parser() -> argparse.ArgumentParser:
migrate.add_argument("--no-create", dest="create_db", action="store_false", help="fail instead of creating a missing database")
migrate.set_defaults(func=_migrate)
connectors = commands.add_parser("connectors", help="manage stored connector credentials")
connector_actions = connectors.add_subparsers(dest="connectors_action", metavar="<action>")
connector_actions.add_parser(
"reencrypt",
help="rewrite every stored credential with ENCRYPTION_SECRET_KEY (after a key rotation)",
)
connectors.set_defaults(func=_connectors)
for name, (module, help_text) in SCRIPTS.items():
commands.add_parser(name, help=f"{help_text} (docsgpt.scripts.{module})", add_help=False)
+63
View File
@@ -0,0 +1,63 @@
"""MCP helpers for connections: re-scanning a server's actions."""
from __future__ import annotations
from docsgpt.connectors import service
from docsgpt.connectors.permissions import 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
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"):
config["auth_credentials"] = service.get_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 refresh_mcp_tools(user_id: str, connection: dict) -> dict:
"""Re-read the actions of every MCP tool on a connection.
Actions that still exist keep the permissions the user chose; new ones
get the defaults from their annotations (reads always allowed, writes
need approval); removed ones disappear.
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"]
added: set[str] = set()
removed: set[str] = set()
refreshed = []
for tool in tools:
fresh = apply_default_permissions("mcp_tool", service._transform_actions(_discover(user_id, connection, tool)))
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 = {
**action,
"active": old.get("active", True),
"require_approval": bool(old.get("require_approval")),
}
merged.append(action)
names = {a.get("name") for a in merged}
added |= names - set(previous)
removed |= set(previous) - names
with db_session() as conn:
repo = UserToolsRepository(conn)
repo.update(str(tool["id"]), user_id, {"actions": merged})
refreshed.append(service.serialize_tool(repo.get_any(str(tool["id"]), user_id)))
return {"added": sorted(added), "removed": sorted(removed), "tools": refreshed}
+128
View File
@@ -0,0 +1,128 @@
"""Which connection a shared tool or source uses at runtime.
A resource that points at a connection runs either with its owner's account
(``owner`` mode, the default for new shares) or with the invoking member's
own account for the same service (``member`` mode). Resolution never returns
credentials; callers read them from the resolved row through
``docsgpt.connectors.service``.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Optional
from docsgpt.connectors import catalog, service
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
from docsgpt.storage.db.session import db_readonly
logger = logging.getLogger(__name__)
MODE_OWNER = "owner"
MODE_MEMBER = "member"
@dataclass(frozen=True)
class ResolvedConnection:
"""The connection a call runs with, or why there is none.
Attributes:
row: The ``connector_sessions`` row, None when missing.
available: Whether it can be used right now.
connector_key: Catalog key of the service.
connector_name: Name shown to the user ("Notion", or a custom label).
delegated: The row belongs to someone other than the invoker.
"""
row: Optional[dict]
available: bool
connector_key: Optional[str]
connector_name: Optional[str]
delegated: bool = False
@property
def connection_id(self) -> Optional[str]:
return str(self.row["id"]) if self.row else None
def _name_for(row: Optional[dict], fallback_key: Optional[str]) -> Optional[str]:
if row is not None:
return service.serialize_connection(row)["name"]
definition = catalog.get_definition(fallback_key)
return definition.name if definition else None
def resolve_connection(resource: dict, invoker_user_id: Optional[str]) -> Optional[ResolvedConnection]:
"""Pick the connection a tool or source uses for ``invoker_user_id``.
``owner`` mode uses ``resource.connection_id``. ``member`` mode uses the
invoker's own connection to the same service (and, for MCP, the same
server), falling back to the owner's when the invoker is the owner.
Args:
resource: A ``user_tools`` or ``sources`` row.
invoker_user_id: Who is running it.
Returns:
None when the resource has no connection at all; otherwise the
resolution, possibly with ``available=False``.
"""
connection_id = resource.get("connection_id")
if not connection_id:
return None
mode = resource.get("credential_mode") or MODE_OWNER
owner = resource.get("user_id")
with db_readonly() as conn:
repo = ConnectorSessionsRepository(conn)
owned = repo.get(str(connection_id))
if owned is not None and owner and owned.get("user_id") != owner:
# A resource may only point at its own owner's connection.
logger.warning(
"resource %s points at a connection it does not own", resource.get("id"),
)
owned = None
row = owned
if mode == MODE_MEMBER and invoker_user_id and invoker_user_id != owner:
row = _member_connection(repo, owned, invoker_user_id)
key = catalog.connector_key_for_row(row or owned or {})
available = row is not None and service.normalize_status(row) == service.STATUS_CONNECTED
return ResolvedConnection(
row=row,
available=available,
connector_key=key,
connector_name=_name_for(row or owned, key),
delegated=bool(row and invoker_user_id and row.get("user_id") != invoker_user_id),
)
def _member_connection(repo: ConnectorSessionsRepository, owned: Optional[dict], invoker: str) -> Optional[dict]:
"""The invoker's own connection to the service the owner's connection is for."""
if owned is None:
return None
for row in repo.list_for_user(invoker):
if row.get("provider") != owned.get("provider"):
continue
if (row.get("server_url") or "") != (owned.get("server_url") or ""):
continue
if service.normalize_status(row) == service.STATUS_CONNECTED:
return row
return None
def audit_delegation(resolved: ResolvedConnection, *, invoker: Optional[str], resource_type: str,
resource_id: Optional[str], agent_id: Optional[str] = None) -> None:
"""Log a call that runs with someone else's account (``owner`` mode)."""
if not resolved.delegated or resolved.row is None:
return
logger.info(
"tool_credential_delegation",
extra={
"invoker": invoker,
"tool_owner": resolved.row.get("user_id"),
"connection_id": resolved.connection_id,
"resource_type": resource_type,
"resource_id": resource_id,
"agent_id": agent_id,
},
)
+925 -19
View File
@@ -13,7 +13,11 @@ 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
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"
@@ -44,6 +48,11 @@ 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 0038."""
token_info = _json(row.get("token_info")) or {}
if isinstance(token_info, dict) and token_info.get("access_token"):
return True
@@ -62,11 +71,12 @@ def normalize_status(row: dict) -> str:
raw = (row.get("status") or "").lower()
if raw in (STATUS_RECONNECT, STATUS_DISCONNECTED, STATUS_ERROR):
return raw
if raw == STATUS_PENDING and not has_credentials(row):
return STATUS_PENDING
if raw in ("authorized", STATUS_CONNECTED, "active") or has_credentials(row):
if raw in ("authorized", STATUS_CONNECTED, "active"):
return STATUS_CONNECTED
return STATUS_PENDING
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]:
@@ -237,9 +247,11 @@ def catalog_for_user(conn, user_id: str, *, is_admin: bool, policies: Optional[d
def disconnect(conn, row: dict) -> dict:
"""Forget a connection's credentials but keep the row and what it feeds.
Sources keep their indexed content and stop syncing; tools stop working
until the account is reconnected. An MCP server's client registration is
kept so reconnecting skips dynamic client registration.
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.
@@ -248,17 +260,911 @@ def disconnect(conn, row: dict) -> dict:
Returns:
The connection's public shape after the change.
"""
repo = ConnectorSessionsRepository(conn)
from sqlalchemy import text
from docsgpt.security.encryption import CredentialDecryptionError
try:
secrets = read_secrets(row)
except CredentialDecryptionError:
secrets = {}
revoke_at_provider(row, secrets)
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 0038 (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):
session_data = {k: v for k, v in session_data.items() if k != "tokens"}
repo.update(
str(row["id"]),
{
"status": STATUS_DISCONNECTED,
"token_info": None,
"session_token": None,
"session_data": session_data or {},
},
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"))
)
return serialize_connection(repo.get(str(row["id"])))
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
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"],
},
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 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:
ValueError: A required credential field is missing.
EncryptionKeyNotConfigured: See :func:`ensure_can_store_credentials`.
"""
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)}
account_label = label or credential_hint(secret_values or credentials)
repo = ConnectorSessionsRepository(conn)
existing = repo.find_account(user_id, definition.key, server_url=server_url, account_label=account_label)
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
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.
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".
"""
import uuid
repo = ConnectorSessionsRepository(conn)
target = state_row
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"]))
definition = catalog.get_definition(provider)
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=definition.name if definition else provider,
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")
with db_readonly() as conn:
repo = ConnectorSessionsRepository(conn)
if connection_id:
row = repo.get_for_user(str(connection_id), user_id)
elif session_token:
row = repo.get_by_session_token(session_token)
if not owns_connector_session(row, user_id, provider):
return None
else:
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:
return repo.get_for_update(connection_id) if lock else repo.get(connection_id)
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.
Returns:
The connection row after the update.
"""
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:
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 ensure_connection_tools(conn, user_id: str, connection: dict, permissions: Optional[dict] = None) -> 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.
"""
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}
created = []
for template in definition.tool_templates:
if template in have or template in ("mcp_tool", "api_tool"):
# 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 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"])
linked_sources = repo.list_sources(connection_id)
try:
revoke_at_provider(row, read_secrets(row))
except CredentialDecryptionError:
pass
if tools == "delete":
conn.execute(
text("DELETE FROM user_tools WHERE connection_id = CAST(:id AS uuid) AND user_id = :user_id"),
{"id": connection_id, "user_id": row["user_id"]},
)
else:
conn.execute(
text("UPDATE user_tools SET status = false WHERE connection_id = CAST(:id AS uuid)"),
{"id": connection_id},
)
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, tool_id: str, permissions: dict) -> Optional[dict]:
"""Apply ``{action: always | ask | off}`` to a tool the user owns."""
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:
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 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
+4
View File
@@ -1240,6 +1240,10 @@ class LLMHandler(ABC):
# can wire the sticky "don't ask again" button.
if pause_info.get("device_id"):
pause_data["device_id"] = pause_info["device_id"]
# A connection-backed tool whose account needs signing in: the
# approval card becomes a Connect card.
if pause_info.get("connection_required"):
pause_data["connection_required"] = pause_info["connection_required"]
trace_unexecuted_tool_call(call, pause_data)
yield {"type": "tool_call", "data": pause_data}
pending_actions.append(pause_info)
+42 -5
View File
@@ -6,7 +6,7 @@ interface for external knowledge base connectors.
"""
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Optional, Tuple
from docsgpt.parser.schema.base import Document
@@ -88,18 +88,55 @@ class BaseConnectorLoader(ABC):
Abstract base class for connector loaders.
Defines the minimal interface that all connector loader
implementations must follow.
implementations must follow. A loader reads its OAuth tokens through
``docsgpt.connectors.service`` from the connection it was built for,
either directly (``connection_id``, what background sync uses) or through
a legacy browser ``session_token`` that names the connection.
"""
connection_id: Optional[str] = None
@abstractmethod
def __init__(self, session_token: str):
def __init__(self, session_token: Optional[str] = None, *, connection_id: Optional[str] = None):
"""
Initialize the connector loader.
Args:
session_token: Authentication session token
session_token: Legacy browser session token naming the connection.
connection_id: The connection to read tokens from.
"""
pass
@classmethod
def from_connection(cls, connection_id: str) -> "BaseConnectorLoader":
"""Build a loader that reads its tokens from ``connection_id``."""
return cls(connection_id=connection_id)
def _load_token_info(
self, session_token: Optional[str], connection_id: Optional[str],
) -> Tuple[str, Dict[str, Any]]:
"""Resolve the connection and return ``(connection_id, token_info)``.
Raises:
ValueError: The connection is missing or needs reconnecting.
"""
from docsgpt.connectors import service
resolved = connection_id or service.connection_id_for_session_token(session_token)
self.connection_id = resolved
return resolved, service.get_valid_token_info(resolved)
def _refresh_rejected_token(self, access_token: Optional[str]) -> Dict[str, Any]:
"""Token info after the provider answered 401 to ``access_token``.
Refreshes under the connection's row lock and persists the rotated
refresh token, or returns the token another worker already renewed.
"""
from docsgpt.connectors import service
if not self.connection_id:
raise ValueError("Loader has no connection to refresh")
return service.get_valid_token_info(self.connection_id, rejected_access_token=access_token)
@abstractmethod
def load_data(self, inputs: Dict[str, Any]) -> List[Document]:
@@ -6,7 +6,6 @@ from urllib.parse import urlencode
import requests
from docsgpt.core.settings import settings
from docsgpt.parser.connectors._auth_utils import session_token_fingerprint
from docsgpt.parser.connectors.base import BaseConnectorAuth
logger = logging.getLogger(__name__)
@@ -152,32 +151,6 @@ class ConfluenceAuth(BaseConnectorAuth):
except Exception:
return True
def get_token_info_from_session(self, session_token: str) -> Dict[str, Any]:
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
)
from docsgpt.storage.db.session import db_readonly
with db_readonly() as conn:
session = ConnectorSessionsRepository(conn).get_by_session_token(
session_token
)
if not session:
raise ValueError(
f"Invalid session token ({session_token_fingerprint(session_token)})"
)
token_info = session.get("token_info")
if not token_info:
raise ValueError("Session missing token information")
required = ["access_token", "refresh_token", "cloud_id"]
missing = [f for f in required if not token_info.get(f)]
if missing:
raise ValueError(f"Missing required token fields: {missing}")
return token_info
def sanitize_token_info(
self, token_info: Dict[str, Any], **extra_fields
) -> Dict[str, Any]:
+6 -24
View File
@@ -44,12 +44,8 @@ def _retry_on_auth_failure(func):
"Auth failure in %s, refreshing token and retrying", func.__name__
)
try:
new_token_info = self.auth.refresh_access_token(self.refresh_token)
new_token_info = self._refresh_rejected_token(self.access_token)
self.access_token = new_token_info["access_token"]
self.refresh_token = new_token_info.get(
"refresh_token", self.refresh_token
)
self._persist_refreshed_tokens(new_token_info)
except Exception as refresh_err:
raise ValueError(
f"Authentication failed and could not be refreshed: {refresh_err}"
@@ -62,13 +58,15 @@ def _retry_on_auth_failure(func):
class ConfluenceLoader(BaseConnectorLoader):
def __init__(self, session_token: str):
def __init__(self, session_token: Optional[str] = None, *, connection_id: Optional[str] = None):
self.auth = ConfluenceAuth()
self.session_token = session_token
token_info = self.auth.get_token_info_from_session(session_token)
_, token_info = self._load_token_info(session_token, connection_id)
missing = [f for f in ("access_token", "cloud_id") if not token_info.get(f)]
if missing:
raise ValueError(f"Missing required token fields: {missing}")
self.access_token = token_info["access_token"]
self.refresh_token = token_info["refresh_token"]
self.cloud_id = token_info["cloud_id"]
self.base_url = API_V2.format(cloud_id=self.cloud_id)
@@ -81,22 +79,6 @@ class ConfluenceLoader(BaseConnectorLoader):
"Accept": "application/json",
}
def _persist_refreshed_tokens(self, token_info: Dict[str, Any]) -> None:
try:
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
)
from docsgpt.storage.db.session import db_session
sanitized = self.auth.sanitize_token_info(token_info)
with db_session() as conn:
repo = ConnectorSessionsRepository(conn)
session = repo.get_by_session_token(self.session_token)
if session:
repo.update(str(session["id"]), {"token_info": sanitized})
except Exception as e:
logger.warning("Failed to persist refreshed tokens: %s", e)
@_retry_on_auth_failure
def load_data(self, inputs: Dict[str, Any]) -> List[Document]:
folder_id = inputs.get("folder_id")
@@ -8,7 +8,6 @@ from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from docsgpt.core.settings import settings
from docsgpt.parser.connectors._auth_utils import session_token_fingerprint
from docsgpt.parser.connectors.base import BaseConnectorAuth
@@ -213,39 +212,6 @@ class GoogleDriveAuth(BaseConnectorAuth):
return True
def get_token_info_from_session(self, session_token: str) -> Dict[str, Any]:
try:
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
)
from docsgpt.storage.db.session import db_readonly
with db_readonly() as conn:
session = ConnectorSessionsRepository(conn).get_by_session_token(
session_token
)
if not session:
raise ValueError(
f"Invalid session token ({session_token_fingerprint(session_token)})"
)
token_info = session.get("token_info")
if not token_info:
raise ValueError("Session missing token information")
required_fields = ["access_token", "refresh_token"]
missing_fields = [field for field in required_fields if field not in token_info or not token_info.get(field)]
if missing_fields:
raise ValueError(f"Missing required token fields: {missing_fields}")
if 'token_uri' not in token_info:
token_info['token_uri'] = 'https://oauth2.googleapis.com/token'
return token_info
except Exception as e:
raise ValueError(f"Failed to retrieve Google Drive token information: {str(e)}")
def validate_credentials(self, credentials: Credentials) -> bool:
"""
Validate Google Drive credentials by making a test API call.
@@ -48,11 +48,13 @@ class GoogleDriveLoader(BaseConnectorLoader):
'application/vnd.google-apps.spreadsheet': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
}
def __init__(self, session_token: str):
def __init__(self, session_token: Optional[str] = None, *, connection_id: Optional[str] = None):
self.auth = GoogleDriveAuth()
self.session_token = session_token
token_info = self.auth.get_token_info_from_session(session_token)
_, token_info = self._load_token_info(session_token, connection_id)
# Google refresh tokens do not rotate, so the credentials object may
# renew the access token in memory for the rest of this run.
self.credentials = self.auth.create_credentials_from_token_info(token_info)
try:
@@ -5,7 +5,6 @@ from typing import Optional, Dict, Any
from msal import ConfidentialClientApplication
from docsgpt.core.settings import settings
from docsgpt.parser.connectors._auth_utils import session_token_fingerprint
from docsgpt.parser.connectors.base import BaseConnectorAuth
logger = logging.getLogger(__name__)
@@ -76,41 +75,6 @@ class SharePointAuth(BaseConnectorAuth):
return self.map_token_response(result)
def get_token_info_from_session(self, session_token: str) -> Dict[str, Any]:
try:
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
)
from docsgpt.storage.db.session import db_readonly
with db_readonly() as conn:
session = ConnectorSessionsRepository(conn).get_by_session_token(
session_token
)
if not session:
raise ValueError(
f"Invalid session token ({session_token_fingerprint(session_token)})"
)
token_info = session.get("token_info")
if not token_info:
raise ValueError("Session missing token information")
required_fields = ["access_token", "refresh_token"]
missing_fields = [field for field in required_fields if field not in token_info or not token_info.get(field)]
if missing_fields:
raise ValueError(f"Missing required token fields: {missing_fields}")
if 'token_uri' not in token_info:
token_info['token_uri'] = f"https://login.microsoftonline.com/{settings.MICROSOFT_TENANT_ID}/oauth2/v2.0/token"
return token_info
except Exception as e:
logger.error("Failed to retrieve token from session: %s", e)
raise ValueError(f"Failed to retrieve SharePoint token information: {str(e)}")
def is_token_expired(self, token_info: Dict[str, Any]) -> bool:
if not token_info:
return True
+14 -10
View File
@@ -26,8 +26,7 @@ def _retry_on_auth_failure(func):
if e.response is not None and e.response.status_code in (401, 403):
logging.info(f"Auth failure in {func.__name__}, refreshing token and retrying")
try:
new_token_info = self.auth.refresh_access_token(self.refresh_token)
self.access_token = new_token_info.get('access_token')
self._apply_token_info(self._refresh_rejected_token(self.access_token))
except Exception as refresh_error:
raise ValueError(
f"Authentication failed and could not be refreshed: {refresh_error}"
@@ -63,13 +62,12 @@ class SharePointLoader(BaseConnectorLoader):
GRAPH_API_BASE = "https://graph.microsoft.com/v1.0"
def __init__(self, session_token: str):
def __init__(self, session_token: Optional[str] = None, *, connection_id: Optional[str] = None):
self.auth = SharePointAuth()
self.session_token = session_token
token_info = self.auth.get_token_info_from_session(session_token)
self.access_token = token_info.get('access_token')
self.refresh_token = token_info.get('refresh_token')
_, token_info = self._load_token_info(session_token, connection_id)
self._apply_token_info(token_info)
self.allows_shared_content = token_info.get('allows_shared_content', False)
if not self.access_token:
@@ -77,6 +75,10 @@ class SharePointLoader(BaseConnectorLoader):
self.next_page_token = None
def _apply_token_info(self, token_info: Dict[str, Any]) -> None:
self.access_token = token_info.get('access_token')
self.expiry = token_info.get('expiry')
def _get_headers(self) -> Dict[str, str]:
return {
'Authorization': f'Bearer {self.access_token}',
@@ -87,12 +89,14 @@ class SharePointLoader(BaseConnectorLoader):
if not self.access_token:
raise ValueError("No access token available")
token_info = {'access_token': self.access_token, 'expiry': None}
if self.auth.is_token_expired(token_info):
# The connection service refreshes under a row lock and stores the
# rotated refresh token; refreshing here would spend it and lose it.
if self.auth.is_token_expired({'access_token': self.access_token, 'expiry': self.expiry}):
logging.info("Token expired, attempting refresh")
try:
new_token_info = self.auth.refresh_access_token(self.refresh_token)
self.access_token = new_token_info.get('access_token')
from docsgpt.connectors import service
self._apply_token_info(service.get_valid_token_info(self.connection_id))
except Exception:
raise ValueError("Failed to refresh access token")
+10
View File
@@ -219,6 +219,8 @@ user_tools_table = Table(
Column(
"connection_id", UUID(as_uuid=True), ForeignKey("connector_sessions.id", ondelete="SET NULL"),
),
# Whose account a shared resource runs with: the owner's, or each member's.
Column("credential_mode", Text, nullable=False, server_default="owner"),
)
token_usage_table = Table(
@@ -344,6 +346,8 @@ sources_table = Table(
Column(
"connection_id", UUID(as_uuid=True), ForeignKey("connector_sessions.id", ondelete="SET NULL"),
),
# Whose account a shared resource runs with: the owner's, or each member's.
Column("credential_mode", Text, nullable=False, server_default="owner"),
)
agents_table = Table(
@@ -598,6 +602,12 @@ connector_sessions_table = Table(
Column("account_label", Text),
Column("auth_kind", Text),
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
# Every secret of the connection, as one owner-bound v2 envelope.
Column("encrypted_credentials", Text),
Column("has_refresh_token", Boolean, nullable=False, server_default="false"),
Column("scopes", JSONB, nullable=False, server_default="[]"),
Column("last_error", Text),
Column("last_used_at", DateTime(timezone=True)),
)
@@ -9,8 +9,13 @@ Shape notes:
unique constraint on ``session_token``.
* MCP sessions key off ``server_url`` instead — a single user may have
multiple MCP servers, one row each. The composite unique index
``(user_id, COALESCE(server_url, ''), provider)`` makes both patterns
coexist without collision.
``(user_id, provider, COALESCE(server_url, ''), COALESCE(account_label, ''))``
makes both patterns coexist, and lets one user connect several accounts
of the same provider (each row carries its own ``account_label``).
* Every secret lives in ``encrypted_credentials``, written and read only by
``docsgpt.connectors.service``; ``token_info`` and the ``tokens`` /
``client_info`` keys of ``session_data`` are legacy plaintext that
migration 0038 moved into it.
* ``session_data`` remains a catch-all JSONB for driver-specific state
(tokens that don't fit anywhere else, per-provider scratch data).
Promoted columns (``session_token``, ``user_email``, ``status``,
@@ -31,8 +36,9 @@ from docsgpt.storage.db.serialization import PGNativeJSONEncoder
_UPDATABLE_SCALARS = {
"server_url", "session_token", "user_email", "status", "expires_at",
"connector_key", "display_name", "account_label", "auth_kind",
"encrypted_credentials", "has_refresh_token", "last_error", "last_used_at",
}
_UPDATABLE_JSONB = {"session_data", "token_info"}
_UPDATABLE_JSONB = {"session_data", "token_info", "scopes"}
def _jsonb(value: Any) -> Any:
@@ -71,7 +77,8 @@ class ConnectorSessionsRepository:
) -> dict:
"""Insert or update a connector session row.
Conflict key is ``(user_id, COALESCE(server_url, ''), provider)``
Conflict key is the account index
``(user_id, provider, COALESCE(server_url, ''), COALESCE(account_label, ''))``
so MCP rows (per-server) and OAuth rows (per-provider) both get
idempotent upsert semantics.
"""
@@ -87,7 +94,7 @@ class ConnectorSessionsRepository:
:status, CAST(:token_info AS jsonb),
CAST(:session_data AS jsonb), :expires_at, :legacy_mongo_id
)
ON CONFLICT (user_id, COALESCE(server_url, ''), provider)
ON CONFLICT (user_id, provider, COALESCE(server_url, ''), COALESCE(account_label, ''))
DO UPDATE SET
session_token = COALESCE(EXCLUDED.session_token, connector_sessions.session_token),
user_email = COALESCE(EXCLUDED.user_email, connector_sessions.user_email),
@@ -182,6 +189,77 @@ class ConnectorSessionsRepository:
)
return [row_to_dict(r) for r in result.fetchall()]
def create(
self,
user_id: str,
provider: str,
*,
connector_key: str,
auth_kind: str,
display_name: Optional[str] = None,
account_label: Optional[str] = None,
server_url: Optional[str] = None,
status: str = "connected",
encrypted_credentials: Optional[str] = None,
has_refresh_token: bool = False,
) -> Optional[dict]:
"""Insert a connection; return None when that account already exists."""
result = self._conn.execute(
text(
"""
INSERT INTO connector_sessions (
user_id, provider, server_url, connector_key, auth_kind, display_name,
account_label, status, encrypted_credentials, has_refresh_token, session_data
)
VALUES (
:user_id, :provider, :server_url, :connector_key, :auth_kind, :display_name,
:account_label, :status, :encrypted_credentials, :has_refresh_token, '{}'::jsonb
)
ON CONFLICT (user_id, provider, COALESCE(server_url, ''), COALESCE(account_label, ''))
DO NOTHING
RETURNING *
"""
),
{
"user_id": user_id,
"provider": provider,
"server_url": server_url,
"connector_key": connector_key,
"auth_kind": auth_kind,
"display_name": display_name,
"account_label": account_label,
"status": status,
"encrypted_credentials": encrypted_credentials,
"has_refresh_token": has_refresh_token,
},
)
row = result.fetchone()
return row_to_dict(row) if row is not None else None
def find_account(
self, user_id: str, provider: str, *, server_url: Optional[str], account_label: Optional[str],
) -> Optional[dict]:
"""The connection for one account, matching the account unique index."""
result = self._conn.execute(
text(
"SELECT * FROM connector_sessions WHERE user_id = :user_id AND provider = :provider "
"AND COALESCE(server_url, '') = COALESCE(:server_url, '') "
"AND COALESCE(account_label, '') = COALESCE(:account_label, '')"
),
{"user_id": user_id, "provider": provider, "server_url": server_url, "account_label": account_label},
)
row = result.fetchone()
return row_to_dict(row) if row is not None else None
def delete_by_id(self, connection_id: str) -> bool:
"""Delete a connection row. Linked sources and tools keep existing (SET NULL)."""
if not looks_like_uuid(connection_id):
return False
result = self._conn.execute(
text("DELETE FROM connector_sessions WHERE id = CAST(:id AS uuid)"), {"id": str(connection_id)},
)
return result.rowcount > 0
def get(self, connection_id: str) -> Optional[dict]:
"""Fetch a connection by id, whoever owns it. Callers authorise."""
if not looks_like_uuid(connection_id):
@@ -237,7 +315,7 @@ class ConnectorSessionsRepository:
"""Sources synced from a connection, newest first."""
result = self._conn.execute(
text(
"SELECT id, name, type, date, sync_frequency, metadata, remote_data, user_id "
"SELECT id, name, type, date, sync_frequency, metadata, remote_data, user_id, file_path "
"FROM sources WHERE connection_id = CAST(:id AS uuid) ORDER BY date DESC"
),
{"id": str(connection_id)},
@@ -337,7 +415,7 @@ class ConnectorSessionsRepository:
Notes:
The conflict target matches the table's composite unique
constraint ``(user_id, COALESCE(server_url, ''), provider)``
index ``(user_id, provider, COALESCE(server_url, ''), COALESCE(account_label, ''))``
so MCP's per-URL rows and OAuth's single-row-per-user rows
both upsert idempotently.
"""
@@ -358,7 +436,7 @@ class ConnectorSessionsRepository:
:user_id, :provider, :server_url,
CAST(:patch AS jsonb)
)
ON CONFLICT (user_id, COALESCE(server_url, ''), provider)
ON CONFLICT (user_id, provider, COALESCE(server_url, ''), COALESCE(account_label, ''))
DO UPDATE SET
server_url = COALESCE(EXCLUDED.server_url, connector_sessions.server_url),
session_data =
+16 -4
View File
@@ -23,8 +23,11 @@ from docsgpt.storage.db.base_repository import looks_like_uuid, row_to_dict
_JSONB_COLUMNS = {"config", "config_requirements", "actions"}
_SCALAR_COLUMNS = {"name", "custom_name", "display_name", "description", "status"}
_ALLOWED_COLUMNS = _SCALAR_COLUMNS | _JSONB_COLUMNS
_SCALAR_COLUMNS = {"name", "custom_name", "display_name", "description", "status", "credential_mode"}
# Set by server code only (tool creation, connection setup); route handlers
# must not pass client input here.
_UUID_COLUMNS = {"connection_id"}
_ALLOWED_COLUMNS = _SCALAR_COLUMNS | _JSONB_COLUMNS | _UUID_COLUMNS
def _encode_jsonb(value: Any) -> Any:
@@ -60,6 +63,8 @@ class UserToolsRepository:
status: bool = True,
extra: Optional[dict] = None,
legacy_mongo_id: Optional[str] = None,
connection_id: Optional[str] = None,
credential_mode: str = "owner",
) -> dict:
"""Insert a new tool row. ``extra`` is merged into the config JSONB."""
cfg = config or {}
@@ -70,14 +75,16 @@ class UserToolsRepository:
"""
INSERT INTO user_tools (
user_id, name, custom_name, display_name, description,
config, config_requirements, actions, status, legacy_mongo_id
config, config_requirements, actions, status, legacy_mongo_id,
connection_id, credential_mode
)
VALUES (
:user_id, :name, :custom_name, :display_name, :description,
CAST(:config AS jsonb),
CAST(:config_requirements AS jsonb),
CAST(:actions AS jsonb),
:status, :legacy_mongo_id
:status, :legacy_mongo_id,
CAST(:connection_id AS uuid), :credential_mode
)
RETURNING *
"""
@@ -93,6 +100,8 @@ class UserToolsRepository:
"actions": _encode_jsonb(actions or []),
"status": status,
"legacy_mongo_id": legacy_mongo_id,
"connection_id": str(connection_id) if connection_id else None,
"credential_mode": credential_mode,
},
)
return row_to_dict(result.fetchone())
@@ -213,6 +222,9 @@ class UserToolsRepository:
if col in _JSONB_COLUMNS:
set_clauses.append(f"{col} = CAST(:{col} AS jsonb)")
params[col] = _encode_jsonb(val)
elif col in _UUID_COLUMNS:
set_clauses.append(f"{col} = CAST(:{col} AS uuid)")
params[col] = str(val) if val else None
else:
set_clauses.append(f"{col} = :{col}")
params[col] = val
+143 -11
View File
@@ -1237,6 +1237,7 @@ def remote_worker(
config=None,
idempotency_key=None,
source_id=None,
connection_id=None,
):
safe_user = safe_filename(user)
full_path = os.path.join(directory, safe_user, uuid.uuid4().hex)
@@ -1283,7 +1284,14 @@ def remote_worker(
self.update_state(state="PROGRESS", meta={"current": 1})
logging.info("Initializing remote loader with type: %s", loader)
remote_loader = RemoteCreator.create_loader(loader)
raw_docs = remote_loader.load_data(source_data)
loader_input = source_data
if connection_id:
loader_input = _with_connection_credentials(source_data, connection_id)
if loader_input is None:
from docsgpt.connectors.service import ConnectionUnavailable
raise ConnectionUnavailable("Reconnect to continue", connection_id=str(connection_id))
raw_docs = remote_loader.load_data(loader_input)
cfg = SourceConfig.parse(config)
chunker = ChunkerCreator.create_chunker(
@@ -1414,6 +1422,8 @@ def remote_worker(
f"Failed to update last_sync for source {source_id_for_events}: {upd_err}"
)
upload_index(full_path, file_data)
if connection_id:
_link_source_to_connection(source_id_for_events, str(connection_id))
publish_user_event(
user,
"source.ingest.completed",
@@ -1467,6 +1477,7 @@ def sync(
retriever,
doc_id=None,
directory="temp",
connection_id=None,
):
try:
remote_worker(
@@ -1480,6 +1491,7 @@ def sync(
sync_frequency,
"sync",
doc_id,
connection_id=connection_id,
)
except Exception as e:
logging.error(f"Error during sync: {e}", exc_info=True)
@@ -1494,7 +1506,7 @@ def sync_worker(self, frequency):
with db_readonly() as conn:
result = conn.execute(
sql_text(
"SELECT id, name, user_id, type, remote_data, retriever "
"SELECT id, name, user_id, type, remote_data, retriever, connection_id "
"FROM sources WHERE sync_frequency = :freq"
),
{"freq": frequency},
@@ -1511,10 +1523,17 @@ def sync_worker(self, frequency):
sync_counts["total_sync_count"] += 1
# Connector sources have no RemoteCreator loader and need an OAuth
# token to sync, which a scheduled task lacks — skip them.
# Connector sources sync from their connection, whose token the
# worker can refresh. Legacy ones with no connection still need the
# browser, so they are skipped as before.
if source_type and source_type.startswith("connector"):
sync_counts["sync_skipped"] += 1
if doc.get("connection_id"):
from docsgpt.api.user.tasks import sync_connector_source as sync_task
sync_task.delay(doc_id)
sync_counts["sync_dispatched"] += 1
else:
sync_counts["sync_skipped"] += 1
continue
source_data = normalize_remote_data(source_type, doc.get("remote_data"))
@@ -1525,7 +1544,8 @@ def sync_worker(self, frequency):
continue
resp = sync(
self, source_data, name, user, source_type, frequency, retriever, doc_id
self, source_data, name, user, source_type, frequency, retriever, doc_id,
connection_id=str(doc["connection_id"]) if doc.get("connection_id") else None,
)
sync_counts[
"sync_success" if resp["status"] == "success" else "sync_failure"
@@ -1533,7 +1553,7 @@ def sync_worker(self, frequency):
return {
key: sync_counts[key]
for key in [
"total_sync_count", "sync_success", "sync_failure", "sync_skipped",
"total_sync_count", "sync_success", "sync_failure", "sync_skipped", "sync_dispatched",
]
}
@@ -2161,6 +2181,113 @@ def _webhook_tool_allowlist(agent_config):
return []
def _with_connection_credentials(source_data, connection_id: str):
"""Loader input with the connection's stored keys merged in, or None.
S3 and Reddit sources made from a connection keep their keys on the
connection only, never in ``sources.remote_data``. Returns None when the
connection is gone or needs reconnecting.
"""
from docsgpt.connectors import service
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
with db_readonly() as conn:
row = ConnectorSessionsRepository(conn).get(str(connection_id))
if row is None:
return None
try:
credentials = service.get_credentials(row)
except service.ConnectionUnavailable:
return None
as_text = isinstance(source_data, str)
data = json.loads(source_data) if as_text else dict(source_data or {})
data.update(credentials)
return json.dumps(data) if as_text else data
def _link_source_to_connection(source_id: str, connection_id: str) -> None:
"""Point a source at the connection it syncs from, and lift any reconnect pause."""
from sqlalchemy import text as sql_text
try:
with db_session() as conn:
conn.execute(
sql_text(
"UPDATE sources SET connection_id = CAST(:cid AS uuid), "
"metadata = metadata - 'sync_state' WHERE id = CAST(:sid AS uuid)"
),
{"cid": str(connection_id), "sid": str(source_id)},
)
except Exception:
logging.warning("Could not link source %s to connection %s", source_id, connection_id, exc_info=True)
def sync_connector_source(self, source_id: str) -> Dict[str, Any]:
"""Re-download and re-index a connector source from its connection.
Runs as the connection owner with no browser: the connection service
refreshes the token under a row lock. When the grant was revoked the
service flags the connection and pauses its sources, and this returns
``paused`` rather than failing again on every schedule.
Args:
self: The bound Celery task.
source_id: The source to sync.
Returns:
``{"status": "success" | "paused" | "skipped"}`` plus the ingest result.
"""
from docsgpt.connectors.service import ConnectionUnavailable, normalize_status
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
with db_readonly() as conn:
from sqlalchemy import text as sql_text
row = conn.execute(
sql_text(
"SELECT id, name, user_id, remote_data, retriever, sync_frequency, config, connection_id "
"FROM sources WHERE id = CAST(:id AS uuid)"
),
{"id": str(source_id)},
).fetchone()
source = dict(row._mapping) if row else None
connection = (
ConnectorSessionsRepository(conn).get(str(source["connection_id"]))
if source and source.get("connection_id")
else None
)
if not source or not connection:
return {"status": "skipped"}
if normalize_status(connection) != "connected":
return {"status": "paused"}
remote_data = source.get("remote_data") or {}
if isinstance(remote_data, str):
try:
remote_data = json.loads(remote_data)
except json.JSONDecodeError:
remote_data = {}
provider = remote_data.get("provider") or connection.get("provider")
try:
result = ingest_connector(
self,
source.get("name"),
source.get("user_id"),
provider,
connection_id=str(connection["id"]),
file_ids=remote_data.get("file_ids") or [],
folder_ids=remote_data.get("folder_ids") or [],
recursive=remote_data.get("recursive", True),
retriever=source.get("retriever") or "classic",
operation_mode="sync",
doc_id=str(source["id"]),
sync_frequency=source.get("sync_frequency") or "never",
config=source.get("config") or None,
)
except ConnectionUnavailable:
return {"status": "paused"}
return {"status": "success", "result": result}
def ingest_connector(
self,
job_name: str,
@@ -2177,6 +2304,7 @@ def ingest_connector(
config=None,
idempotency_key=None,
source_id=None,
connection_id=None,
) -> Dict[str, Any]:
"""
Ingestion for internal knowledge bases (GoogleDrive, etc.).
@@ -2185,7 +2313,8 @@ def ingest_connector(
job_name: Name of the ingestion job
user: User identifier
source_type: Type of remote source ("google_drive", "dropbox", etc.)
session_token: Authentication token for the service
session_token: Legacy browser session token naming the connection
connection_id: The connection whose account the files are read with
file_ids: List of file IDs to download
folder_ids: List of folder IDs to download
recursive: Whether to recursively download folders
@@ -2247,8 +2376,8 @@ def ingest_connector(
meta={"current": 10, "status": "Initializing connector"},
)
if not session_token:
raise ValueError(f"{source_type} connector requires session_token")
if not session_token and not connection_id:
raise ValueError(f"{source_type} connector requires a connection")
if not ConnectorCreator.is_supported(source_type):
raise ValueError(
@@ -2256,8 +2385,9 @@ def ingest_connector(
)
remote_loader = ConnectorCreator.create_connector(
source_type, session_token
source_type, session_token, connection_id=connection_id
)
connection_id = remote_loader.connection_id
# Create a clean config for storage
api_source_config = {
@@ -2411,6 +2541,8 @@ def ingest_connector(
)
upload_index(vector_store_path, file_data)
if connection_id:
_link_source_to_connection(source_id_for_events, connection_id)
# Ensure we mark the task as complete
self.update_state(
+11 -1
View File
@@ -1010,6 +1010,10 @@ class TestDBTokenStorage:
monkeypatch.setattr(session_mod, "db_session", _yield)
monkeypatch.setattr(session_mod, "db_readonly", _yield)
import docsgpt.connectors.service as service_mod
monkeypatch.setattr(service_mod, "db_session", _yield)
monkeypatch.setattr(service_mod, "db_readonly", _yield)
def test_get_base_url(self):
from docsgpt.agents.tools.mcp_tool import DBTokenStorage
@@ -1132,7 +1136,13 @@ class TestDBTokenStorage:
# ``server_url`` must NOT be duplicated inside the JSONB blob.
session_data = row["session_data"] or {}
assert "server_url" not in session_data
assert session_data.get("tokens", {}).get("access_token") == "at"
# Tokens live only in the encrypted envelope, never in plaintext.
assert "tokens" not in session_data
assert row["encrypted_credentials"].startswith("v2:")
from docsgpt.connectors.service import read_secrets
assert read_secrets(row)["tokens"]["access_token"] == "at"
assert row["status"] == "connected"
def test_clear_removes_row(self, monkeypatch, pg_conn):
from mcp.shared.auth import OAuthToken
+5 -1
View File
@@ -21,7 +21,11 @@ def _patch_db(conn, module="docsgpt.api.connector.routes"):
def _yield():
yield conn
with patch(f"{module}.db_session", _yield), patch(f"{module}.db_readonly", _yield):
with patch(f"{module}.db_session", _yield), patch(f"{module}.db_readonly", _yield), patch(
"docsgpt.api.connector.routes.db_readonly", _yield
), patch("docsgpt.connectors.service.db_session", _yield), patch(
"docsgpt.connectors.service.db_readonly", _yield
):
yield
+7 -1
View File
@@ -30,6 +30,10 @@ def _patch_db(conn):
"docsgpt.api.connector.routes.db_session", _yield
), patch(
"docsgpt.api.connector.routes.db_readonly", _yield
), patch(
"docsgpt.connectors.service.db_session", _yield
), patch(
"docsgpt.connectors.service.db_readonly", _yield
):
yield
@@ -487,8 +491,10 @@ class TestConnectorSync:
def test_returns_400_missing_fields(self, app):
from docsgpt.api.connector.routes import ConnectorSync
# The source's own connection is used when none is named, so only
# the source id is required.
with app.test_request_context(
"/api/connectors/sync", method="POST", json={"source_id": "x"}
"/api/connectors/sync", method="POST", json={}
):
from flask import request
request.decoded_token = {"sub": "u"}
@@ -24,6 +24,10 @@ def _patch_db(conn):
"docsgpt.api.user.sources.upload.db_session", _yield
), patch(
"docsgpt.api.user.sources.upload.db_readonly", _yield
), patch(
"docsgpt.connectors.service.db_session", _yield
), patch(
"docsgpt.connectors.service.db_readonly", _yield
):
yield
+5 -1
View File
@@ -67,7 +67,8 @@ class TestIngestRemoteTask:
mock_worker.assert_called_once_with(
ANY, {"url": "http://x"}, "job1", "user1", "web",
config=None, idempotency_key=None, source_id=None,
sync_frequency="never",
config=None, idempotency_key=None, source_id=None, connection_id=None,
)
assert result == {"status": "ok"}
@@ -223,6 +224,7 @@ class TestIngestConnectorTask:
"user1",
"gdrive",
session_token=None,
connection_id=None,
file_ids=None,
folder_ids=None,
recursive=True,
@@ -248,6 +250,7 @@ class TestIngestConnectorTask:
"user1",
"sharepoint",
session_token="tok",
connection_id=None,
file_ids=["f1"],
folder_ids=["d1"],
recursive=False,
@@ -263,6 +266,7 @@ class TestIngestConnectorTask:
"user1",
"sharepoint",
session_token="tok",
connection_id=None,
file_ids=["f1"],
folder_ids=["d1"],
recursive=False,
+323
View File
@@ -0,0 +1,323 @@
"""Tests for the connection write endpoints."""
from __future__ import annotations
import json
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
from sqlalchemy import text
from docsgpt.connectors import service
from docsgpt.security.encryption import encrypt_json
@pytest.fixture
def app():
return Flask(__name__)
@contextmanager
def _db(conn):
@contextmanager
def _yield():
yield conn
with patch.multiple("docsgpt.api.connector.connections", db_session=_yield, db_readonly=_yield), \
patch.multiple("docsgpt.connectors.service", db_session=_yield, db_readonly=_yield), \
patch.multiple("docsgpt.api.connector.routes", db_session=_yield, db_readonly=_yield):
yield
def _call(app, resource, method, path, user="alice", body=None, headers=None, args=()):
with app.test_request_context(path, method=method.upper(), json=body, headers=headers or {}):
from flask import request
request.decoded_token = {"sub": user} if user else None
return getattr(resource(), method)(*args)
def _connection(conn, user="alice", provider="telegram", auth_kind="api_key", secrets=None, status="connected",
**cols) -> str:
values = {
"user_id": user, "provider": provider, "connector_key": provider, "auth_kind": auth_kind,
"status": status, "account_label": cols.pop("account_label", "…abcd"),
"encrypted_credentials": encrypt_json(secrets, user) if secrets is not None else None, **cols,
}
names = ", ".join(values)
params = ", ".join(f":{k}" for k in values)
return str(conn.execute(
text(f"INSERT INTO connector_sessions ({names}) VALUES ({params}) RETURNING id"), values,
).scalar())
class TestCreate:
def test_telegram_zero_question_setup(self, app, pg_conn):
"""Appendix C of the connectors spec: create, then set up, tools appear."""
from docsgpt.api.connector.connections import ConnectionSetup, ConnectionsList
with _db(pg_conn):
created = _call(app, ConnectionsList, "post", "/api/connections",
body={"connector_key": "telegram", "credentials": {"token": "123:abcdefgh"}})
assert created.status_code == 201
payload = created.get_json()
assert payload["setup"] == {"tools": "auto", "sync": "off"}
assert "token" not in json.dumps(payload)
cid = payload["connection"]["id"]
setup = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup",
body={"create_tools": True}, args=[cid])
assert setup.status_code == 200
tools = setup.get_json()["tools"]
assert [t["name"] for t in tools] == ["telegram"]
actions = {a["name"]: a for a in tools[0]["actions"]}
assert actions["telegram_send_message"]["access"] == "write"
assert actions["telegram_send_message"]["permission"] == "ask"
def test_setup_is_idempotent_for_tools(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionSetup
cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}})
with _db(pg_conn):
for _ in range(2):
_call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup",
body={"create_tools": True}, args=[cid])
count = pg_conn.execute(
text("SELECT count(*) FROM user_tools WHERE connection_id = CAST(:c AS uuid)"), {"c": cid},
).scalar()
assert count == 1
def test_rejects_oauth_connector(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionsList
with _db(pg_conn):
resp = _call(app, ConnectionsList, "post", "/api/connections",
body={"connector_key": "google_drive", "credentials": {}})
assert resp.status_code == 400
def test_default_key_refused(self, app, pg_conn, monkeypatch):
from docsgpt.api.connector.connections import ConnectionsList
from docsgpt.core.settings import settings
from docsgpt.security.encryption import DEFAULT_ENCRYPTION_KEY
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY", DEFAULT_ENCRYPTION_KEY)
monkeypatch.setattr(settings, "AUTH_TYPE", "oidc")
with _db(pg_conn):
resp = _call(app, ConnectionsList, "post", "/api/connections",
body={"connector_key": "brave", "credentials": {"token": "some-long-key"}})
assert resp.status_code == 400
assert resp.get_json()["code"] == "encryption_key_default"
class TestSetupSync:
def test_oauth_source_is_queued_with_connection(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionSetup
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={"token_info": {}})
task = MagicMock(id="task-1")
with _db(pg_conn), patch("docsgpt.api.user.tasks.ingest_connector_task.apply_async",
return_value=task) as apply:
resp = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", body={
"create_tools": True,
"sync": {"items": {"folder_ids": ["f1"]}, "frequency": "weekly", "name": "Handbook"},
}, args=[cid])
assert resp.status_code == 200
source = resp.get_json()["sources"][0]
assert source["name"] == "Handbook" and source["sync_frequency"] == "weekly"
kwargs = apply.call_args.kwargs["kwargs"]
assert kwargs["connection_id"] == cid
assert kwargs["folder_ids"] == ["f1"]
assert "session_token" not in kwargs
# Drive has no tools; nothing was created.
assert resp.get_json()["tools"] == []
def test_nothing_picked_is_a_bad_request(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionSetup
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={"token_info": {}})
with _db(pg_conn):
resp = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup",
body={"sync": {"items": {}}}, args=[cid])
assert resp.status_code == 400
def test_s3_keys_stay_on_the_connection(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionSetup
cid = _connection(pg_conn, provider="s3", secrets={"credentials": {
"aws_access_key_id": "AKIA", "aws_secret_access_key": "shh-secret"}})
with _db(pg_conn), patch("docsgpt.api.user.tasks.ingest_remote.apply_async",
return_value=MagicMock(id="t")) as apply:
resp = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", body={
"sync": {"items": {"bucket": "docs", "aws_secret_access_key": "smuggled"}},
}, args=[cid])
assert resp.status_code == 200
kwargs = apply.call_args.kwargs["kwargs"]
assert kwargs["source_data"] == {"bucket": "docs"}
assert kwargs["connection_id"] == cid
def test_flagged_connection_must_reconnect_first(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionSetup
cid = _connection(pg_conn, status="reconnect_needed", secrets={})
with _db(pg_conn):
resp = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup",
body={}, args=[cid])
assert resp.status_code == 409
class TestReconnect:
def test_api_key_reconnect_replaces_and_resumes(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionReconnect
cid = _connection(pg_conn, status="reconnect_needed", secrets={"credentials": {"token": "old"}})
pg_conn.execute(text(
"INSERT INTO sources (user_id, name, connection_id, metadata) "
"VALUES ('alice', 's', CAST(:c AS uuid), '{\"sync_state\": \"paused_reconnect\"}')"
), {"c": cid})
with _db(pg_conn):
resp = _call(app, ConnectionReconnect, "post", f"/api/connections/{cid}/reconnect",
body={"credentials": {"token": "new"}}, args=[cid])
assert resp.status_code == 200
assert resp.get_json()["connection"]["status"] == "connected"
row = pg_conn.execute(text("SELECT * FROM connector_sessions WHERE id = CAST(:c AS uuid)"),
{"c": cid}).one()._mapping
assert service.read_secrets(dict(row))["credentials"]["token"] == "new"
meta = pg_conn.execute(text("SELECT metadata FROM sources WHERE user_id = 'alice'")).scalar()
assert "sync_state" not in meta
def test_oauth_reconnect_returns_authorization_url(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionReconnect
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={})
fake_auth = MagicMock()
fake_auth.get_authorization_url.return_value = "https://accounts.example/auth"
with _db(pg_conn), patch("docsgpt.api.connector.routes.ConnectorCreator.create_auth",
return_value=fake_auth):
resp = _call(app, ConnectionReconnect, "post", f"/api/connections/{cid}/reconnect",
body={}, args=[cid])
assert resp.get_json()["authorization_url"] == "https://accounts.example/auth"
state = fake_auth.get_authorization_url.call_args.kwargs["state"]
import base64
assert json.loads(base64.urlsafe_b64decode(state))["object_id"] == cid
def test_cannot_reconnect_someone_elses(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionReconnect
cid = _connection(pg_conn, user="bob", secrets={})
with _db(pg_conn):
resp = _call(app, ConnectionReconnect, "post", f"/api/connections/{cid}/reconnect",
body={"credentials": {"token": "x"}}, args=[cid])
assert resp.status_code == 404
class TestPickerToken:
def test_returns_access_token_only(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionPickerToken
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth",
secrets={"token_info": {"access_token": "at", "refresh_token": "rt"}})
fake_auth = MagicMock()
fake_auth.is_token_expired.return_value = False
with _db(pg_conn), patch("docsgpt.parser.connectors.connector_creator.ConnectorCreator.create_auth",
return_value=fake_auth):
resp = _call(app, ConnectionPickerToken, "post", f"/api/connections/{cid}/picker-token",
args=[cid])
payload = resp.get_json()
assert payload["access_token"] == "at"
assert "rt" not in json.dumps(payload)
def test_owner_only(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionPickerToken
cid = _connection(pg_conn, user="bob", provider="google_drive", auth_kind="oauth", secrets={})
with _db(pg_conn):
resp = _call(app, ConnectionPickerToken, "post", f"/api/connections/{cid}/picker-token",
args=[cid])
assert resp.status_code == 404
class TestClaim:
def test_claims_own_token(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionClaim
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={}, session_token="legacy")
with _db(pg_conn):
resp = _call(app, ConnectionClaim, "post", "/api/connections/claim",
body={"provider": "google_drive", "session_token": "legacy"})
assert resp.get_json()["connection_id"] == cid
def test_other_users_token_is_not_found(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionClaim
_connection(pg_conn, user="bob", provider="google_drive", auth_kind="oauth", secrets={}, session_token="b")
with _db(pg_conn):
resp = _call(app, ConnectionClaim, "post", "/api/connections/claim",
body={"provider": "google_drive", "session_token": "b"})
assert resp.status_code == 404
class TestPermissions:
def test_set_permissions(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionToolPermissions
cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}})
with _db(pg_conn):
tool = service.ensure_connection_tools(pg_conn, "alice", service.ConnectorSessionsRepository(
pg_conn).get(cid))[0]
resp = _call(app, ConnectionToolPermissions, "put",
f"/api/connections/{cid}/tools/{tool['id']}/permissions",
body={"permissions": {"telegram_send_image": "off", "telegram_send_message": "always"}},
args=[cid, str(tool["id"])])
actions = {a["name"]: a["permission"] for a in resp.get_json()["tool"]["actions"]}
assert actions == {"telegram_send_message": "always", "telegram_send_image": "off"}
def test_rejects_unknown_permission(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionToolPermissions
cid = _connection(pg_conn, secrets={})
with _db(pg_conn):
resp = _call(app, ConnectionToolPermissions, "put", "/x", body={"permissions": {"a": "maybe"}},
args=[cid, "t"])
assert resp.status_code == 400
class TestDelete:
def test_delete_with_source_removal(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionDetail
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={})
pg_conn.execute(text("INSERT INTO sources (user_id, name, connection_id) VALUES ('alice', 's', CAST(:c AS uuid))"),
{"c": cid})
with _db(pg_conn), patch("docsgpt.connectors.service.revoke_at_provider"), patch(
"docsgpt.api.user.sources.routes.delete_source", return_value=True
) as delete_source:
resp = _call(app, ConnectionDetail, "delete", f"/api/connections/{cid}",
body={"sources": "delete"}, args=[cid])
assert resp.status_code == 200
assert delete_source.call_count == 1
def test_bad_mode(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionDetail
cid = _connection(pg_conn, secrets={})
with _db(pg_conn):
resp = _call(app, ConnectionDetail, "delete", "/x", body={"sources": "burn"}, args=[cid])
assert resp.status_code == 400
class TestNoSecretsInResponses:
def test_list_and_detail(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionDetail, ConnectionsList
cid = _connection(pg_conn, secrets={"credentials": {"token": "super-secret-value"}},
session_token="sess-tok")
with _db(pg_conn):
listed = _call(app, ConnectionsList, "get", "/api/connections")
detail = _call(app, ConnectionDetail, "get", f"/api/connections/{cid}", args=[cid])
for resp in (listed, detail):
body = json.dumps(resp.get_json())
assert "super-secret-value" not in body
assert "sess-tok" not in body
assert "v2:" not in body
+254
View File
@@ -0,0 +1,254 @@
"""Runtime use of connections: tool execution, sharing modes and scheduled sync."""
from __future__ import annotations
import logging
from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from sqlalchemy import text
from docsgpt.security.encryption import encrypt_json
@contextmanager
def _service_db(conn):
@contextmanager
def _yield():
yield conn
with patch.multiple("docsgpt.connectors.service", db_session=_yield, db_readonly=_yield), \
patch.multiple("docsgpt.connectors.resolve", db_readonly=_yield):
yield
def _connection(conn, user="alice", provider="telegram", status="connected", secrets=None,
auth_kind="api_key", server_url=None) -> str:
return str(conn.execute(
text(
"INSERT INTO connector_sessions (user_id, provider, connector_key, auth_kind, status, "
"account_label, server_url, encrypted_credentials) VALUES (:u, :p, :p, :a, :s, :l, :url, :e) RETURNING id"
),
{"u": user, "p": provider, "a": auth_kind, "s": status, "l": f"{user}-label", "url": server_url,
"e": encrypt_json(secrets or {"credentials": {"token": f"{user}-token"}}, user)},
).scalar())
def _tool(connection_id, *, user="alice", name="telegram", mode="owner", tool_id="tool-1"):
return {
"id": tool_id,
"user_id": user,
"name": name,
"config": {},
"actions": [{"name": "telegram_send_message", "active": True, "require_approval": False}],
"connection_id": connection_id,
"credential_mode": mode,
}
def _call(action="telegram_send_message"):
return SimpleNamespace(id="call-1", name=action, arguments="{}", thought_signature=None)
def _executor(user="alice", headless=False):
from docsgpt.agents.tool_executor import ToolExecutor
executor = ToolExecutor(user=user, headless=headless)
return executor
def _pause(executor, tool):
with patch("docsgpt.agents.tool_executor.ToolActionParser") as parser:
parser.return_value.parse_args.return_value = ("t1", "telegram_send_message", {})
return executor.check_pause({"t1": tool}, _call(), "OpenAILLM")
class TestResolution:
def test_owner_mode_uses_owner_connection(self, pg_conn):
from docsgpt.connectors.resolve import resolve_connection
cid = _connection(pg_conn)
with _service_db(pg_conn):
resolved = resolve_connection(_tool(cid), "bob")
assert resolved.available and resolved.connection_id == cid
assert resolved.delegated is True
assert resolved.connector_name == "Telegram"
def test_member_mode_uses_invokers_connection(self, pg_conn):
from docsgpt.connectors.resolve import resolve_connection
owner = _connection(pg_conn)
bobs = _connection(pg_conn, user="bob")
with _service_db(pg_conn):
resolved = resolve_connection(_tool(owner, mode="member"), "bob")
assert resolved.connection_id == bobs
assert resolved.delegated is False
def test_member_mode_without_own_connection_is_unavailable(self, pg_conn):
from docsgpt.connectors.resolve import resolve_connection
owner = _connection(pg_conn)
with _service_db(pg_conn):
resolved = resolve_connection(_tool(owner, mode="member"), "bob")
assert resolved.available is False
assert resolved.connector_name == "Telegram"
def test_member_mode_matches_mcp_server(self, pg_conn):
from docsgpt.connectors.resolve import resolve_connection
owner = _connection(pg_conn, provider="mcp:https://a.example.com", auth_kind="mcp_oauth",
server_url="https://a.example.com", secrets={"tokens": {"access_token": "x"}})
_connection(pg_conn, user="bob", provider="mcp:https://b.example.com", auth_kind="mcp_oauth",
server_url="https://b.example.com", secrets={"tokens": {"access_token": "y"}})
with _service_db(pg_conn):
resolved = resolve_connection(_tool(owner, name="mcp_tool", mode="member"), "bob")
assert resolved.available is False
def test_resource_cannot_borrow_another_users_connection(self, pg_conn):
from docsgpt.connectors.resolve import resolve_connection
mallorys_target = _connection(pg_conn, user="victim")
with _service_db(pg_conn):
resolved = resolve_connection(_tool(mallorys_target, user="mallory"), "mallory")
assert resolved.available is False and resolved.row is None
class TestExecutor:
def test_credentials_come_from_the_connection(self, pg_conn):
cid = _connection(pg_conn)
executor = _executor()
tool = _tool(cid)
with _service_db(pg_conn), patch("docsgpt.agents.tool_executor.ToolManager") as manager:
executor._get_or_load_tool(tool, "t1", "telegram_send_message")
config = manager.return_value.load_tool.call_args.kwargs["tool_config"]
assert config["token"] == "alice-token"
assert "encrypted_credentials" not in config
def test_owner_mode_delegation_is_audited(self, pg_conn, caplog):
cid = _connection(pg_conn)
executor = _executor(user="bob")
with _service_db(pg_conn), patch("docsgpt.agents.tool_executor.ToolManager"), \
caplog.at_level(logging.INFO, logger="docsgpt.connectors.resolve"):
executor._get_or_load_tool(_tool(cid), "t1", "telegram_send_message")
record = next(r for r in caplog.records if r.message == "tool_credential_delegation")
assert record.connection_id == cid and record.invoker == "bob"
def test_needs_reconnect_pauses_on_connect_card(self, pg_conn):
cid = _connection(pg_conn, status="reconnect_needed")
with _service_db(pg_conn):
pause = _pause(_executor(), _tool(cid))
assert pause["pause_type"] == "awaiting_approval"
assert pause["connection_required"] == {
"connector_key": "telegram", "connector_name": "Telegram", "status": "reconnect_needed",
}
def test_member_without_connection_pauses(self, pg_conn):
cid = _connection(pg_conn)
with _service_db(pg_conn):
pause = _pause(_executor(user="bob"), _tool(cid, mode="member"))
assert pause["connection_required"]["status"] == "missing"
def test_headless_run_is_denied_not_paused(self, pg_conn):
cid = _connection(pg_conn, status="disconnected")
with _service_db(pg_conn):
pause = _pause(_executor(headless=True), _tool(cid))
assert pause["pause_type"] == "headless_denied"
assert pause["error_type"] == "connection_required"
def test_connected_tool_does_not_pause(self, pg_conn):
cid = _connection(pg_conn)
with _service_db(pg_conn):
assert _pause(_executor(), _tool(cid)) is None
def test_mcp_tool_gets_connection_id_not_tokens(self, pg_conn):
cid = _connection(pg_conn, provider="mcp:https://m.example.com", auth_kind="mcp_oauth",
server_url="https://m.example.com", secrets={"tokens": {"access_token": "secret"}})
tool = {**_tool(cid, name="mcp_tool"), "config": {"server_url": "https://m.example.com/mcp",
"auth_type": "oauth"}}
with _service_db(pg_conn), patch("docsgpt.agents.tool_executor.ToolManager") as manager:
_executor()._get_or_load_tool(tool, "t1", "search")
config = manager.return_value.load_tool.call_args.kwargs["tool_config"]
assert config["connection_id"] == cid
assert "secret" not in str(config)
class TestScheduledSync:
def test_connector_sources_with_a_connection_are_dispatched(self, pg_conn):
from docsgpt import worker
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth",
secrets={"token_info": {"access_token": "a"}})
pg_conn.execute(text(
"INSERT INTO sources (user_id, name, type, sync_frequency, connection_id, remote_data) "
"VALUES ('alice', 'Drive', 'connector:file', 'weekly', CAST(:c AS uuid), '{\"provider\": \"google_drive\"}')"
), {"c": cid})
pg_conn.execute(text(
"INSERT INTO sources (user_id, name, type, sync_frequency) VALUES ('alice', 'Old', 'connector:file', 'weekly')"
))
@contextmanager
def _yield():
yield pg_conn
with patch.object(worker, "db_readonly", _yield), patch(
"docsgpt.api.user.tasks.sync_connector_source.delay"
) as delay:
counts = worker.sync_worker(MagicMock(), "weekly")
assert delay.call_count == 1
assert counts["sync_dispatched"] == 1
assert counts["sync_skipped"] == 1
def test_paused_connection_is_not_synced(self, pg_conn):
from docsgpt import worker
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", status="reconnect_needed",
secrets={"token_info": {}})
source = pg_conn.execute(text(
"INSERT INTO sources (user_id, name, type, sync_frequency, connection_id) "
"VALUES ('alice', 'Drive', 'connector:file', 'weekly', CAST(:c AS uuid)) RETURNING id"
), {"c": cid}).scalar()
@contextmanager
def _yield():
yield pg_conn
with patch.object(worker, "db_readonly", _yield), patch.object(worker, "ingest_connector") as ingest:
result = worker.sync_connector_source(MagicMock(), str(source))
assert result == {"status": "paused"}
ingest.assert_not_called()
def test_sync_runs_as_the_connection_without_a_browser(self, pg_conn):
from docsgpt import worker
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth",
secrets={"token_info": {"access_token": "a"}})
source = pg_conn.execute(text(
"INSERT INTO sources (user_id, name, type, sync_frequency, connection_id, remote_data) "
"VALUES ('alice', 'Drive', 'connector:file', 'daily', CAST(:c AS uuid), "
"'{\"provider\": \"google_drive\", \"folder_ids\": [\"f\"], \"recursive\": false}') RETURNING id"
), {"c": cid}).scalar()
@contextmanager
def _yield():
yield pg_conn
with patch.object(worker, "db_readonly", _yield), patch.object(
worker, "ingest_connector", return_value={"tokens": 1}
) as ingest:
result = worker.sync_connector_source(MagicMock(), str(source))
assert result["status"] == "success"
kwargs = ingest.call_args.kwargs
assert kwargs["connection_id"] == cid
assert kwargs["operation_mode"] == "sync"
assert kwargs["folder_ids"] == ["f"] and kwargs["recursive"] is False
assert "session_token" not in kwargs
@pytest.mark.parametrize("secret_key", ["token_info", "tokens", "client_info", "encrypted_credentials",
"client_secret", "refresh_token", "access_token"])
def test_redaction_covers_connection_secrets(secret_key):
from docsgpt.storage.db.redaction import REDACTED, redact_secrets
assert redact_secrets({secret_key: {"x": "y"}})[secret_key] == REDACTED
@@ -0,0 +1,420 @@
"""Tests for connection credentials: storage, refresh, failure handling."""
from __future__ import annotations
import json
import threading
import time
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
import pytest
from sqlalchemy import text
from docsgpt.connectors import service
from docsgpt.security.encryption import encrypt_json
def _patch_service_db(conn):
@contextmanager
def _yield():
yield conn
return patch.multiple("docsgpt.connectors.service", db_session=_yield, db_readonly=_yield)
def _connection(conn, *, user="alice", provider="google_drive", secrets=None, status="connected", **extra) -> str:
cols = {
"user_id": user,
"provider": provider,
"connector_key": provider,
"auth_kind": extra.pop("auth_kind", "oauth"),
"status": status,
"account_label": extra.pop("account_label", f"{user}@example.com"),
"encrypted_credentials": encrypt_json(secrets, user) if secrets is not None else None,
**extra,
}
names = ", ".join(cols)
values = ", ".join(f":{k}" for k in cols)
return str(
conn.execute(
text(f"INSERT INTO connector_sessions ({names}) VALUES ({values}) RETURNING id"), cols,
).scalar()
)
def _source(conn, connection_id, user="alice") -> str:
return str(
conn.execute(
text(
"INSERT INTO sources (user_id, name, type, connection_id) "
"VALUES (:u, 'Docs', 'connector:file', CAST(:c AS uuid)) RETURNING id"
),
{"u": user, "c": connection_id},
).scalar()
)
def _row(conn, connection_id):
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
return ConnectorSessionsRepository(conn).get(connection_id)
class _FakeAuth:
"""A provider whose refresh rotates the refresh token on every call."""
def __init__(self, *, expired=True, fail=None, delay=0.0):
self.expired = expired
self.fail = fail
self.delay = delay
self.calls = []
self._lock = threading.Lock()
def is_token_expired(self, token_info):
return self.expired and token_info.get("access_token") == "old-at"
def refresh_access_token(self, refresh_token):
with self._lock:
self.calls.append(refresh_token)
n = len(self.calls)
if self.delay:
time.sleep(self.delay)
if self.fail:
raise self.fail
if refresh_token != "rt-0":
# A reused rotating refresh token is rejected, like Microsoft's.
raise ValueError("invalid_grant: refresh token already used")
return {"access_token": f"new-at-{n}", "refresh_token": f"rt-{n}", "expiry": None}
def sanitize_token_info(self, token_info, **extra):
return {k: token_info.get(k) for k in ("access_token", "refresh_token", "expiry", "cloud_id")}
def _auth(fake):
return patch("docsgpt.parser.connectors.connector_creator.ConnectorCreator.create_auth", return_value=fake)
class TestSecrets:
def test_write_and_read(self, pg_conn):
cid = _connection(pg_conn, secrets=None, status="pending")
service.write_secrets(pg_conn, _row(pg_conn, cid), {"token_info": {"access_token": "a", "refresh_token": "r"}})
row = _row(pg_conn, cid)
assert row["token_info"] is None
assert row["has_refresh_token"] is True
assert "access_token" not in json.dumps(row["session_data"])
assert service.read_secrets(row) == {"token_info": {"access_token": "a", "refresh_token": "r"}}
def test_legacy_plaintext_is_read(self, pg_conn):
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
row = ConnectorSessionsRepository(pg_conn).upsert(
"alice", "share_point", token_info={"access_token": "p"}, status="authorized",
)
assert service.read_secrets(row)["token_info"]["access_token"] == "p"
def test_decrypt_failure_flags_reconnect(self, pg_conn):
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "a"}})
pg_conn.execute(
text("UPDATE connector_sessions SET encrypted_credentials = 'v2:ffffffff:AAAA' WHERE id = CAST(:i AS uuid)"),
{"i": cid},
)
with _patch_service_db(pg_conn), patch("docsgpt.events.publisher.publish_user_event"):
assert service.load_secrets(_row(pg_conn, cid)) is None
row = _row(pg_conn, cid)
assert row["status"] == "reconnect_needed"
assert row["last_error"] == service.DECRYPT_ERROR
class TestGetValidTokenInfo:
def test_returns_unexpired_token(self, pg_conn):
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "fresh", "refresh_token": "rt-0"}})
fake = _FakeAuth(expired=True)
with _patch_service_db(pg_conn), _auth(fake):
info = service.get_valid_token_info(cid)
assert info["access_token"] == "fresh"
assert fake.calls == []
def test_refreshes_and_persists_rotated_token(self, pg_conn):
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "old-at", "refresh_token": "rt-0"}})
fake = _FakeAuth()
with _patch_service_db(pg_conn), _auth(fake):
info = service.get_valid_token_info(cid)
assert info["access_token"] == "new-at-1"
stored = service.read_secrets(_row(pg_conn, cid))["token_info"]
assert stored["refresh_token"] == "rt-1"
def test_rejected_access_token_forces_refresh(self, pg_conn):
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "fresh", "refresh_token": "rt-0"}})
fake = _FakeAuth(expired=False)
with _patch_service_db(pg_conn), _auth(fake):
info = service.get_valid_token_info(cid, rejected_access_token="fresh")
assert info["access_token"] == "new-at-1"
def test_rejected_token_already_replaced_is_not_refreshed_again(self, pg_conn):
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "newer", "refresh_token": "rt-0"}})
fake = _FakeAuth(expired=False)
with _patch_service_db(pg_conn), _auth(fake):
info = service.get_valid_token_info(cid, rejected_access_token="stale")
assert info["access_token"] == "newer"
assert fake.calls == []
def test_revoked_grant_pauses_sources_and_notifies(self, pg_conn):
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "old-at", "refresh_token": "rt-0"}})
source = _source(pg_conn, cid)
fake = _FakeAuth(fail=ValueError("Error refreshing token: invalid_grant"))
with _patch_service_db(pg_conn), _auth(fake), patch(
"docsgpt.events.publisher.publish_user_event"
) as publish:
with pytest.raises(service.ConnectionUnavailable):
service.get_valid_token_info(cid)
assert _row(pg_conn, cid)["status"] == "reconnect_needed"
meta = pg_conn.execute(
text("SELECT metadata FROM sources WHERE id = CAST(:i AS uuid)"), {"i": source}
).scalar()
assert meta["sync_state"] == "paused_reconnect"
publish.assert_called_once()
assert publish.call_args.args[1] == "connection.reconnect_needed"
# No account data or secrets in the event.
assert "alice@example.com" not in json.dumps(publish.call_args.args[2])
def test_second_failure_does_not_notify_again(self, pg_conn):
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "old-at", "refresh_token": "rt-0"}})
fake = _FakeAuth(fail=ValueError("invalid_grant"))
with _patch_service_db(pg_conn), _auth(fake), patch(
"docsgpt.events.publisher.publish_user_event"
) as publish:
for _ in range(2):
with pytest.raises(service.ConnectionUnavailable):
service.get_valid_token_info(cid)
assert publish.call_count == 1
def test_transient_failure_keeps_status(self, pg_conn):
import requests
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "old-at", "refresh_token": "rt-0"}})
fake = _FakeAuth(fail=requests.exceptions.ConnectionError("network down"))
with _patch_service_db(pg_conn), _auth(fake):
with pytest.raises(service.TransientConnectionError):
service.get_valid_token_info(cid)
assert _row(pg_conn, cid)["status"] == "connected"
def test_disconnected_connection_is_unavailable(self, pg_conn):
cid = _connection(pg_conn, secrets={}, status="disconnected")
with _patch_service_db(pg_conn):
with pytest.raises(service.ConnectionUnavailable) as exc:
service.get_valid_token_info(cid)
assert exc.value.status == "disconnected"
def test_concurrent_refresh_spends_the_refresh_token_once(self, pg_engine):
"""Two workers with an expiring rotating token: one refresh, both succeed."""
with pg_engine.begin() as conn:
cid = _connection(conn, secrets={"token_info": {"access_token": "old-at", "refresh_token": "rt-0"}})
@contextmanager
def _session():
with pg_engine.begin() as conn:
yield conn
fake = _FakeAuth(delay=0.3)
results, errors = [], []
def worker():
try:
results.append(service.get_valid_token_info(cid)["access_token"])
except Exception as exc: # pragma: no cover - reported below
errors.append(exc)
with patch.multiple("docsgpt.connectors.service", db_session=_session, db_readonly=_session), _auth(fake):
threads = [threading.Thread(target=worker) for _ in range(2)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
assert errors == []
assert fake.calls == ["rt-0"]
assert results == ["new-at-1", "new-at-1"]
class TestApiKeyConnections:
def test_create_and_reuse(self, pg_conn):
from docsgpt.connectors import catalog
telegram = catalog.get_definition("telegram")
row, created = service.create_api_key_connection(pg_conn, "alice", telegram, {"token": "123456:ABCDEFG"})
assert created is True
assert row["account_label"] == "…DEFG"
assert row["status"] == "connected"
again, created = service.create_api_key_connection(pg_conn, "alice", telegram, {"token": "123456:ABCDEFG"})
assert created is False and again["id"] == row["id"]
assert service.get_credentials(again) == {"token": "123456:ABCDEFG"}
def test_missing_field(self, pg_conn):
from docsgpt.connectors import catalog
with pytest.raises(ValueError):
service.create_api_key_connection(pg_conn, "alice", catalog.get_definition("telegram"), {})
def test_default_key_refused_with_auth(self, pg_conn, monkeypatch):
from docsgpt.connectors import catalog
from docsgpt.core.settings import settings
from docsgpt.security.encryption import DEFAULT_ENCRYPTION_KEY
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY", DEFAULT_ENCRYPTION_KEY)
monkeypatch.setattr(settings, "AUTH_TYPE", "session_jwt")
with pytest.raises(service.EncryptionKeyNotConfigured):
service.create_api_key_connection(
pg_conn, "alice", catalog.get_definition("brave"), {"token": "secret-key-1234"},
)
def test_default_key_allowed_single_user(self, pg_conn, monkeypatch):
from docsgpt.connectors import catalog
from docsgpt.core.settings import settings
from docsgpt.security.encryption import DEFAULT_ENCRYPTION_KEY
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY", DEFAULT_ENCRYPTION_KEY)
monkeypatch.setattr(settings, "AUTH_TYPE", None)
row, _ = service.create_api_key_connection(
pg_conn, "local", catalog.get_definition("brave"), {"token": "secret-key-1234"},
)
assert row["encrypted_credentials"].startswith("v2:")
def test_flagged_connection_credentials_unavailable(self, pg_conn):
cid = _connection(
pg_conn, provider="telegram", auth_kind="api_key", status="reconnect_needed",
secrets={"credentials": {"token": "t"}},
)
with pytest.raises(service.ConnectionUnavailable):
service.get_credentials(_row(pg_conn, cid))
class TestOAuthCompletion:
def test_new_account(self, pg_conn):
state_row = service.begin_oauth(pg_conn, "alice", "google_drive")
assert state_row["status"] == "pending"
row = service.complete_oauth(pg_conn, state_row, "google_drive", {"access_token": "a"}, "a@example.com")
assert row["id"] == state_row["id"]
assert row["account_label"] == "a@example.com"
assert row["status"] == "connected"
assert row["session_token"]
def test_same_account_updates_existing_connection(self, pg_conn):
first = service.complete_oauth(
pg_conn, service.begin_oauth(pg_conn, "alice", "google_drive"), "google_drive",
{"access_token": "a"}, "a@example.com",
)
source = _source(pg_conn, str(first["id"]))
pg_conn.execute(
text("UPDATE connector_sessions SET status = 'reconnect_needed' WHERE id = CAST(:i AS uuid)"),
{"i": str(first["id"])},
)
pg_conn.execute(
text("UPDATE sources SET metadata = '{\"sync_state\": \"paused_reconnect\"}' WHERE id = CAST(:i AS uuid)"),
{"i": source},
)
pending = service.begin_oauth(pg_conn, "alice", "google_drive")
again = service.complete_oauth(pg_conn, pending, "google_drive", {"access_token": "b"}, "a@example.com")
assert again["id"] == first["id"]
assert again["status"] == "connected"
assert _row(pg_conn, str(pending["id"])) is None
meta = pg_conn.execute(
text("SELECT metadata FROM sources WHERE id = CAST(:i AS uuid)"), {"i": source}
).scalar()
assert "sync_state" not in meta
def test_second_account_is_a_second_connection(self, pg_conn):
first = service.complete_oauth(
pg_conn, service.begin_oauth(pg_conn, "alice", "google_drive"), "google_drive",
{"access_token": "a"}, "a@example.com",
)
second = service.complete_oauth(
pg_conn, service.begin_oauth(pg_conn, "alice", "google_drive"), "google_drive",
{"access_token": "b"}, "b@example.com",
)
assert first["id"] != second["id"]
assert len(service.list_connections(pg_conn, "alice")) == 2
def test_reconnect_must_name_own_connection(self, pg_conn):
cid = _connection(pg_conn, user="bob", secrets={})
with pytest.raises(service.ConnectionUnavailable):
service.begin_oauth(pg_conn, "alice", "google_drive", cid)
class TestReencrypt:
def test_rewrites_rows_on_previous_key(self, pg_conn, monkeypatch):
from docsgpt.core.settings import settings
from docsgpt.security.encryption import current_key_id, envelope_key_id
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY", "old-key")
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "a"}})
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY", "new-key")
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY_PREVIOUS", "old-key")
with _patch_service_db(pg_conn):
counts = service.reencrypt_all()
assert counts == {"rewritten": 1, "current": 0, "failed": 0}
row = _row(pg_conn, cid)
assert envelope_key_id(row["encrypted_credentials"]) == current_key_id()
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY_PREVIOUS", None)
assert service.read_secrets(row)["token_info"]["access_token"] == "a"
def test_unreadable_rows_flagged(self, pg_conn, monkeypatch):
from docsgpt.core.settings import settings
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY", "lost-key")
cid = _connection(pg_conn, secrets={"token_info": {"access_token": "a"}})
monkeypatch.setattr(settings, "ENCRYPTION_SECRET_KEY", "new-key")
with _patch_service_db(pg_conn), patch("docsgpt.events.publisher.publish_user_event"):
counts = service.reencrypt_all()
assert counts["failed"] == 1
assert _row(pg_conn, cid)["status"] == "reconnect_needed"
class TestRemove:
def test_keep_sources_delete_tools(self, pg_conn):
cid = _connection(pg_conn, provider="telegram", auth_kind="api_key", secrets={"credentials": {"token": "t"}})
source = _source(pg_conn, cid)
pg_conn.execute(
text(
"INSERT INTO user_tools (user_id, name, connection_id) VALUES ('alice', 'telegram', CAST(:c AS uuid))"
),
{"c": cid},
)
with patch("docsgpt.connectors.service.revoke_at_provider") as revoke:
to_delete = service.remove_connection(pg_conn, _row(pg_conn, cid))
assert to_delete == []
revoke.assert_called_once()
assert _row(pg_conn, cid) is None
assert pg_conn.execute(text("SELECT count(*) FROM user_tools WHERE user_id = 'alice'")).scalar() == 0
kept = pg_conn.execute(
text("SELECT connection_id, sync_frequency FROM sources WHERE id = CAST(:i AS uuid)"), {"i": source}
).one()
assert kept.connection_id is None and kept.sync_frequency == "never"
def test_delete_sources_returns_them(self, pg_conn):
cid = _connection(pg_conn, secrets={})
source = _source(pg_conn, cid)
with patch("docsgpt.connectors.service.revoke_at_provider"):
to_delete = service.remove_connection(pg_conn, _row(pg_conn, cid), sources="delete")
assert [str(s["id"]) for s in to_delete] == [source]
class TestRevoke:
def test_google_revocation_posts_token(self):
response = MagicMock(status_code=200)
with patch("requests.post", return_value=response) as post:
ok = service.revoke_at_provider(
{"provider": "google_drive"}, {"token_info": {"refresh_token": "rt", "access_token": "at"}},
)
assert ok
assert post.call_args.kwargs["data"] == {"token": "rt"}
def test_other_providers_are_local_only(self):
with patch("requests.post") as post:
assert not service.revoke_at_provider({"provider": "share_point"}, {"token_info": {"access_token": "a"}})
post.assert_not_called()
def test_failure_never_raises(self):
import requests
with patch("requests.post", side_effect=requests.exceptions.Timeout()):
assert not service.revoke_at_provider({"provider": "google_drive"}, {"token_info": {"access_token": "a"}})
+3 -1
View File
@@ -242,7 +242,9 @@ class TestDisconnect:
)
repo = ConnectorSessionsRepository(pg_conn)
service.disconnect(pg_conn, repo.get(mcp))
assert repo.get(mcp)["session_data"] == {"client_info": {"client_id": "c"}}
row = repo.get(mcp)
assert row["session_data"] == {}
assert service.read_secrets(row) == {"client_info": {"client_id": "c"}}
def test_route_rejects_other_users(self, app, pg_conn):
from docsgpt.api.connector.connections import ConnectionDisconnect
@@ -3,7 +3,6 @@
from contextlib import contextmanager
from unittest.mock import patch
import pytest
@contextmanager
@@ -18,82 +17,3 @@ def _patch_db(conn):
yield
class TestGetTokenInfoFromSession:
def test_invalid_session_token_raises(self, pg_conn):
from docsgpt.parser.connectors.confluence.auth import (
ConfluenceAuth,
)
auth = ConfluenceAuth.__new__(ConfluenceAuth)
with _patch_db(pg_conn), pytest.raises(ValueError):
auth.get_token_info_from_session("no-such-token")
def test_missing_token_info_raises(self, pg_conn):
from docsgpt.parser.connectors.confluence.auth import (
ConfluenceAuth,
)
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
)
repo = ConnectorSessionsRepository(pg_conn)
repo.upsert("u", "confluence", status="authorized")
# Set session_token but no token_info
session = repo.get_by_user_provider("u", "confluence")
repo.update(str(session["id"]), {"session_token": "tok-no-info"})
auth = ConfluenceAuth.__new__(ConfluenceAuth)
with _patch_db(pg_conn), pytest.raises(ValueError):
auth.get_token_info_from_session("tok-no-info")
def test_missing_required_fields_raises(self, pg_conn):
from docsgpt.parser.connectors.confluence.auth import (
ConfluenceAuth,
)
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
)
repo = ConnectorSessionsRepository(pg_conn)
repo.upsert("u", "confluence", status="authorized")
session = repo.get_by_user_provider("u", "confluence")
repo.update(
str(session["id"]),
{
"session_token": "tok-partial",
"token_info": {"access_token": "at"}, # missing refresh + cloud_id
},
)
auth = ConfluenceAuth.__new__(ConfluenceAuth)
with _patch_db(pg_conn), pytest.raises(ValueError):
auth.get_token_info_from_session("tok-partial")
def test_complete_token_info_returned(self, pg_conn):
from docsgpt.parser.connectors.confluence.auth import (
ConfluenceAuth,
)
from docsgpt.storage.db.repositories.connector_sessions import (
ConnectorSessionsRepository,
)
repo = ConnectorSessionsRepository(pg_conn)
repo.upsert("u", "confluence", status="authorized")
session = repo.get_by_user_provider("u", "confluence")
repo.update(
str(session["id"]),
{
"session_token": "tok-good",
"token_info": {
"access_token": "at",
"refresh_token": "rt",
"cloud_id": "cid-1",
},
},
)
auth = ConfluenceAuth.__new__(ConfluenceAuth)
with _patch_db(pg_conn):
got = auth.get_token_info_from_session("tok-good")
assert got["access_token"] == "at"
assert got["cloud_id"] == "cid-1"
@@ -2,6 +2,7 @@
import os
import tempfile
from tests.parser.connectors.token_patch import patch_tokens
from unittest.mock import MagicMock, patch
import pytest
@@ -25,11 +26,12 @@ def _make_loader(token_info=None):
with patch("docsgpt.parser.connectors.confluence.loader.ConfluenceAuth") as MockAuth:
mock_auth = MagicMock()
mock_auth.get_token_info_from_session.return_value = token_info
token_info = token_info
MockAuth.return_value = mock_auth
from docsgpt.parser.connectors.confluence.loader import ConfluenceLoader
loader = ConfluenceLoader("session_tok")
with patch_tokens(token_info):
loader = ConfluenceLoader("session_tok")
loader.auth = mock_auth
return loader
@@ -69,7 +71,7 @@ class TestConfluenceLoaderInit:
def test_init_sets_attributes(self, loader):
assert loader.session_token == "session_tok"
assert loader.access_token == "test_at"
assert loader.refresh_token == "test_rt"
assert loader.connection_id == "conn-1"
assert loader.cloud_id == "test_cloud"
assert loader.next_page_token is None
assert "test_cloud" in loader.base_url
@@ -834,18 +836,16 @@ class TestRetryOnAuthFailure:
# Second call succeeds
return _mock_response({"results": [], "_links": {}})
loader.auth.refresh_access_token = MagicMock(return_value={
"access_token": "new_at",
"refresh_token": "new_rt",
})
loader._persist_refreshed_tokens = MagicMock()
with patch("requests.get", side_effect=flaky_request):
with patch(
"docsgpt.connectors.service.get_valid_token_info",
return_value={"access_token": "new_at", "cloud_id": "test_cloud"},
) as tokens, patch("requests.get", side_effect=flaky_request):
loader.load_data({})
assert loader.auth.refresh_access_token.called
# The rejected token goes to the service, which refreshes under the
# connection's row lock and stores the rotated refresh token.
tokens.assert_called_once_with("conn-1", rejected_access_token="test_at")
assert loader.access_token == "new_at"
assert loader._persist_refreshed_tokens.called
@pytest.mark.unit
def test_raises_non_auth_http_error(self, loader):
@@ -867,11 +867,9 @@ class TestRetryOnAuthFailure:
err = requests.exceptions.HTTPError(response=resp)
err.response = resp
loader.auth.refresh_access_token = MagicMock(
side_effect=Exception("refresh failed")
)
with patch("requests.get", side_effect=err):
with patch(
"docsgpt.connectors.service.get_valid_token_info", side_effect=Exception("refresh failed"),
), patch("requests.get", side_effect=err):
with pytest.raises(ValueError, match="Authentication failed"):
loader.load_data({})
@@ -890,29 +888,12 @@ class TestRetryOnAuthFailure:
raise err
return _mock_response({"results": [], "_links": {}})
loader.auth.refresh_access_token = MagicMock(return_value={
"access_token": "new_at",
"refresh_token": "new_rt",
})
loader._persist_refreshed_tokens = MagicMock()
with patch("requests.get", side_effect=flaky):
with patch(
"docsgpt.connectors.service.get_valid_token_info",
return_value={"access_token": "new_at", "cloud_id": "test_cloud"},
) as tokens, patch("requests.get", side_effect=flaky):
loader.load_data({})
assert loader.auth.refresh_access_token.called
assert tokens.called
# ---------------------------------------------------------------------------
# _persist_refreshed_tokens
# ---------------------------------------------------------------------------
class TestPersistRefreshedTokens:
pass
@pytest.mark.unit
def test_logs_warning_on_failure(self, loader):
loader.auth.sanitize_token_info = MagicMock(side_effect=Exception("db error"))
# Should not raise, just log a warning
loader._persist_refreshed_tokens({"access_token": "at"})
@@ -51,8 +51,7 @@ def _patches(session_return: Optional[Dict[str, Any]]):
)
return (
patch(
"docsgpt.storage.db.repositories.connector_sessions."
"ConnectorSessionsRepository",
"docsgpt.connectors.service.ConnectorSessionsRepository",
fake_repo_cls,
),
patch(
@@ -88,57 +87,18 @@ class TestSessionTokenFingerprint:
assert session_token_fingerprint("a") != session_token_fingerprint("b")
class TestConfluenceAuthDoesNotLeakToken:
class TestConnectionLookupDoesNotLeakToken:
"""A session token that names no connection is reported by fingerprint only."""
@pytest.mark.unit
def test_invalid_session_does_not_interpolate_token(self) -> None:
from docsgpt.parser.connectors.confluence.auth import ConfluenceAuth
from docsgpt.connectors import service
auth = ConfluenceAuth.__new__(ConfluenceAuth)
repo_patch, ctx_patch = _patches(None)
with repo_patch, ctx_patch:
with pytest.raises(ValueError) as excinfo:
auth.get_token_info_from_session(SECRET_TOKEN)
service.connection_id_for_session_token(SECRET_TOKEN)
message = str(excinfo.value)
assert SECRET_TOKEN not in message
assert session_token_fingerprint(SECRET_TOKEN) in message
class TestGoogleDriveAuthDoesNotLeakToken:
@pytest.mark.unit
def test_invalid_session_does_not_interpolate_token(self) -> None:
from docsgpt.parser.connectors.google_drive.auth import GoogleDriveAuth
auth = GoogleDriveAuth.__new__(GoogleDriveAuth)
repo_patch, ctx_patch = _patches(None)
with repo_patch, ctx_patch:
with pytest.raises(ValueError) as excinfo:
auth.get_token_info_from_session(SECRET_TOKEN)
# The Google Drive module wraps the inner ValueError in a broad
# ``except Exception as e: raise ValueError(... {str(e)})`` block,
# so the outer message still carries the fingerprint from the
# inner raise but must never carry the raw token.
message = str(excinfo.value)
assert SECRET_TOKEN not in message
assert session_token_fingerprint(SECRET_TOKEN) in message
class TestSharePointAuthDoesNotLeakToken:
@pytest.mark.unit
def test_invalid_session_does_not_interpolate_token(self) -> None:
from docsgpt.parser.connectors.share_point.auth import SharePointAuth
auth = SharePointAuth.__new__(SharePointAuth)
repo_patch, ctx_patch = _patches(None)
with repo_patch, ctx_patch:
with pytest.raises(ValueError) as excinfo:
auth.get_token_info_from_session(SECRET_TOKEN)
# SharePoint also wraps the inner ValueError. Same invariants.
message = str(excinfo.value)
assert SECRET_TOKEN not in message
assert session_token_fingerprint(SECRET_TOKEN) in message
@@ -1,5 +1,6 @@
"""Tests for ConnectorCreator factory class."""
from tests.parser.connectors.token_patch import patch_tokens
from unittest.mock import patch, MagicMock
import pytest
@@ -74,7 +75,7 @@ class TestConnectorCreator:
def test_create_connector_google_drive(self):
with patch("docsgpt.parser.connectors.google_drive.loader.GoogleDriveAuth") as MockAuth:
mock_auth_instance = MagicMock()
mock_auth_instance.get_token_info_from_session.return_value = {
token_info = {
"access_token": "at", "refresh_token": "rt"
}
mock_creds = MagicMock()
@@ -84,7 +85,8 @@ class TestConnectorCreator:
mock_auth_instance.build_drive_service.return_value = MagicMock()
MockAuth.return_value = mock_auth_instance
loader = self.ConnectorCreator.create_connector("google_drive", "session_tok")
with patch_tokens(token_info):
loader = self.ConnectorCreator.create_connector("google_drive", "session_tok")
from docsgpt.parser.connectors.google_drive.loader import GoogleDriveLoader
assert isinstance(loader, GoogleDriveLoader)
@@ -92,11 +94,12 @@ class TestConnectorCreator:
def test_create_connector_share_point(self):
with patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth") as MockAuth:
mock_auth_instance = MagicMock()
mock_auth_instance.get_token_info_from_session.return_value = {
token_info = {
"access_token": "at", "refresh_token": "rt"
}
MockAuth.return_value = mock_auth_instance
loader = self.ConnectorCreator.create_connector("share_point", "session_tok")
with patch_tokens(token_info):
loader = self.ConnectorCreator.create_connector("share_point", "session_tok")
from docsgpt.parser.connectors.share_point.loader import SharePointLoader
assert isinstance(loader, SharePointLoader)
@@ -364,71 +364,6 @@ class _FakeReadonlyCtx:
return False
class TestGetTokenInfoFromSession:
def _patches(self, session_return):
fake_repo_cls = type(
"FakeRepo",
(_FakeRepo,),
{"_session": session_return},
)
return (
patch(
"docsgpt.storage.db.repositories.connector_sessions.ConnectorSessionsRepository",
fake_repo_cls,
),
patch(
"docsgpt.storage.db.session.db_readonly",
lambda: _FakeReadonlyCtx(),
),
)
@pytest.mark.unit
def test_valid_session(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches({
"session_token": "st",
"token_info": {"access_token": "at", "refresh_token": "rt"},
})
with repo_patch, ctx_patch:
result = auth.get_token_info_from_session("st")
assert result["access_token"] == "at"
assert result["token_uri"] == "https://oauth2.googleapis.com/token"
@pytest.mark.unit
def test_session_not_found_raises(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches(None)
with repo_patch, ctx_patch:
with pytest.raises(ValueError, match="Failed to retrieve Google Drive token"):
auth.get_token_info_from_session("bad_token")
@pytest.mark.unit
def test_session_missing_token_info_raises(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches({"session_token": "st"})
with repo_patch, ctx_patch:
with pytest.raises(ValueError, match="Failed to retrieve Google Drive token"):
auth.get_token_info_from_session("st")
@pytest.mark.unit
def test_missing_required_fields_raises(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches({
"session_token": "st",
"token_info": {"access_token": "at"},
})
with repo_patch, ctx_patch:
with pytest.raises(ValueError, match="Failed to retrieve Google Drive token"):
auth.get_token_info_from_session("st")
@pytest.mark.unit
def test_empty_token_info_raises(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches({
"session_token": "st",
"token_info": None,
})
with repo_patch, ctx_patch:
with pytest.raises(ValueError, match="Failed to retrieve Google Drive token"):
auth.get_token_info_from_session("st")
class TestValidateCredentials:
@pytest.mark.unit
@@ -1,5 +1,6 @@
"""Tests for GoogleDriveLoader."""
from tests.parser.connectors.token_patch import patch_tokens
from unittest.mock import MagicMock, patch
import pytest
@@ -11,7 +12,7 @@ def _make_loader(service=None):
"""Create a GoogleDriveLoader with mocked dependencies."""
with patch("docsgpt.parser.connectors.google_drive.loader.GoogleDriveAuth") as MockAuth:
mock_auth = MagicMock()
mock_auth.get_token_info_from_session.return_value = {
token_info = {
"access_token": "at",
"refresh_token": "rt",
}
@@ -24,7 +25,8 @@ def _make_loader(service=None):
MockAuth.return_value = mock_auth
from docsgpt.parser.connectors.google_drive.loader import GoogleDriveLoader
loader = GoogleDriveLoader("session_tok")
with patch_tokens(token_info):
loader = GoogleDriveLoader("session_tok")
return loader
@@ -51,7 +53,7 @@ class TestGoogleDriveLoaderInit:
def test_init_service_failure_sets_none(self):
with patch("docsgpt.parser.connectors.google_drive.loader.GoogleDriveAuth") as MockAuth:
mock_auth = MagicMock()
mock_auth.get_token_info_from_session.return_value = {
token_info = {
"access_token": "at", "refresh_token": "rt"
}
mock_creds = MagicMock()
@@ -61,7 +63,8 @@ class TestGoogleDriveLoaderInit:
MockAuth.return_value = mock_auth
from docsgpt.parser.connectors.google_drive.loader import GoogleDriveLoader
loader = GoogleDriveLoader("st")
with patch_tokens(token_info):
loader = GoogleDriveLoader("st")
assert loader.service is None
@@ -285,63 +285,3 @@ class _FakeReadonlyCtx:
return False
class TestGetTokenInfoFromSession:
def _patches(self, session_return):
fake_repo_cls = type(
"FakeRepo",
(_FakeRepo,),
{"_session": session_return},
)
return (
patch(
"docsgpt.storage.db.repositories.connector_sessions.ConnectorSessionsRepository",
fake_repo_cls,
),
patch(
"docsgpt.storage.db.session.db_readonly",
lambda: _FakeReadonlyCtx(),
),
)
@pytest.mark.unit
def test_valid_session(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches({
"session_token": "st",
"token_info": {"access_token": "at", "refresh_token": "rt"},
})
with repo_patch, ctx_patch:
result = auth.get_token_info_from_session("st")
assert result["access_token"] == "at"
assert "token_uri" in result
@pytest.mark.unit
def test_session_not_found_raises(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches(None)
with repo_patch, ctx_patch:
with pytest.raises(ValueError, match="Failed to retrieve SharePoint token"):
auth.get_token_info_from_session("bad")
@pytest.mark.unit
def test_missing_token_info_raises(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches({"session_token": "st"})
with repo_patch, ctx_patch:
with pytest.raises(ValueError, match="Failed to retrieve SharePoint token"):
auth.get_token_info_from_session("st")
@pytest.mark.unit
def test_empty_token_info_raises(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches({"session_token": "st", "token_info": None})
with repo_patch, ctx_patch:
with pytest.raises(ValueError, match="Failed to retrieve SharePoint token"):
auth.get_token_info_from_session("st")
@pytest.mark.unit
def test_missing_required_fields_raises(self, auth, mock_settings):
repo_patch, ctx_patch = self._patches({
"session_token": "st",
"token_info": {"access_token": "at"},
})
with repo_patch, ctx_patch:
with pytest.raises(ValueError, match="Failed to retrieve SharePoint token"):
auth.get_token_info_from_session("st")
@@ -1,6 +1,7 @@
"""Tests for SharePointLoader."""
import os
from tests.parser.connectors.token_patch import patch_tokens
from unittest.mock import MagicMock, patch
import pytest
@@ -13,7 +14,7 @@ def _make_loader(access_token="at", refresh_token="rt", allows_shared=False):
"""Create a SharePointLoader with mocked dependencies."""
with patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth") as MockAuth:
mock_auth = MagicMock()
mock_auth.get_token_info_from_session.return_value = {
token_info = {
"access_token": access_token,
"refresh_token": refresh_token,
"allows_shared_content": allows_shared,
@@ -22,7 +23,8 @@ def _make_loader(access_token="at", refresh_token="rt", allows_shared=False):
MockAuth.return_value = mock_auth
from docsgpt.parser.connectors.share_point.loader import SharePointLoader
loader = SharePointLoader("session_tok")
with patch_tokens(token_info):
loader = SharePointLoader("session_tok")
return loader
@@ -42,7 +44,7 @@ class TestSharePointLoaderInit:
def test_init_sets_attributes(self, loader):
assert loader.session_token == "session_tok"
assert loader.access_token == "at"
assert loader.refresh_token == "rt"
assert loader.connection_id == "conn-1"
assert loader.allows_shared_content is False
assert loader.next_page_token is None
@@ -50,7 +52,7 @@ class TestSharePointLoaderInit:
def test_no_access_token_raises(self):
with patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth") as MockAuth:
mock_auth = MagicMock()
mock_auth.get_token_info_from_session.return_value = {
token_info = {
"access_token": None,
"refresh_token": "rt",
}
@@ -58,7 +60,8 @@ class TestSharePointLoaderInit:
from docsgpt.parser.connectors.share_point.loader import SharePointLoader
with pytest.raises(ValueError, match="No access token"):
SharePointLoader("st")
with patch_tokens(token_info):
SharePointLoader("st")
class TestGetHeaders:
@@ -86,16 +89,19 @@ class TestEnsureValidToken:
@pytest.mark.unit
def test_expired_token_refreshes(self, loader):
loader.auth.is_token_expired.return_value = True
loader.auth.refresh_access_token.return_value = {"access_token": "new_at"}
loader._ensure_valid_token()
with patch(
"docsgpt.connectors.service.get_valid_token_info", return_value={"access_token": "new_at"},
) as tokens:
loader._ensure_valid_token()
assert loader.access_token == "new_at"
tokens.assert_called_once_with("conn-1")
@pytest.mark.unit
def test_refresh_failure_raises(self, loader):
loader.auth.is_token_expired.return_value = True
loader.auth.refresh_access_token.side_effect = Exception("fail")
with pytest.raises(ValueError, match="Failed to refresh"):
loader._ensure_valid_token()
with patch("docsgpt.connectors.service.get_valid_token_info", side_effect=Exception("fail")):
with pytest.raises(ValueError, match="Failed to refresh"):
loader._ensure_valid_token()
class TestGetItemUrl:
@@ -305,7 +311,7 @@ class TestLoadFileById:
return mock_resp_ok
loader._process_file = MagicMock(return_value=Document(text="", doc_id="f1", extra_info={}))
loader.auth.refresh_access_token.return_value = {"access_token": "new_at"}
loader._refresh_rejected_token = MagicMock(return_value={"access_token": "new_at"})
with patch("docsgpt.parser.connectors.share_point.loader.requests.get", side_effect=get_side_effect):
doc = loader._load_file_by_id("f1")
@@ -481,7 +487,7 @@ class TestDownloadFileContent:
raise http_error
return mock_resp_ok
loader.auth.refresh_access_token.return_value = {"access_token": "new_at"}
loader._refresh_rejected_token = MagicMock(return_value={"access_token": "new_at"})
with patch("docsgpt.parser.connectors.share_point.loader.requests.get", side_effect=get_side_effect):
content = loader._download_file_content("f1")
@@ -1004,7 +1010,7 @@ class TestRetryOnAuthFailureDecorator:
raise http_error
return mock_resp_ok
loader.auth.refresh_access_token.return_value = {"access_token": "new_at"}
loader._refresh_rejected_token = MagicMock(return_value={"access_token": "new_at"})
with patch("docsgpt.parser.connectors.share_point.loader.requests.get", side_effect=get_side_effect):
content = loader._download_file_content("f1")
@@ -1018,7 +1024,7 @@ class TestRetryOnAuthFailureDecorator:
mock_resp_403.status_code = 403
http_error = real_requests.exceptions.HTTPError(response=mock_resp_403)
loader.auth.refresh_access_token.side_effect = Exception("refresh fail")
loader._refresh_rejected_token = MagicMock(side_effect=Exception("refresh fail"))
with patch("docsgpt.parser.connectors.share_point.loader.requests.get", side_effect=http_error):
with pytest.raises(ValueError, match="could not be refreshed"):
+23
View File
@@ -0,0 +1,23 @@
"""Stand-in for the connection service in connector loader tests."""
from __future__ import annotations
from contextlib import ExitStack, contextmanager
from unittest.mock import patch
@contextmanager
def patch_tokens(token_info: dict, connection_id: str = "conn-1"):
"""Make every loader resolve ``connection_id`` and read ``token_info``.
Yields the ``get_valid_token_info`` mock so a test can assert on or
change what a refresh returns.
"""
with ExitStack() as stack:
stack.enter_context(
patch("docsgpt.connectors.service.connection_id_for_session_token", return_value=connection_id)
)
tokens = stack.enter_context(
patch("docsgpt.connectors.service.get_valid_token_info", return_value=token_info)
)
yield tokens
@@ -66,7 +66,8 @@ class TestSharePointLoaderLoadFileById:
"""Test _load_file_by_id method."""
@patch("docsgpt.parser.connectors.share_point.loader.requests.get")
@patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth.get_token_info_from_session")
@patch("docsgpt.connectors.service.connection_id_for_session_token", new=lambda token: "conn-1")
@patch("docsgpt.connectors.service.get_valid_token_info")
@patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth.__init__", return_value=None)
@patch("docsgpt.parser.connectors.share_point.loader.SharePointLoader._ensure_valid_token")
def test_load_file_by_id_includes_size_in_select(self, mock_ensure_token, mock_auth_init, mock_get_token, mock_get):
@@ -97,7 +98,8 @@ class TestSharePointLoaderLoadFileById:
assert "size" in params["$select"]
@patch("docsgpt.parser.connectors.share_point.loader.requests.get")
@patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth.get_token_info_from_session")
@patch("docsgpt.connectors.service.connection_id_for_session_token", new=lambda token: "conn-1")
@patch("docsgpt.connectors.service.get_valid_token_info")
@patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth.__init__", return_value=None)
@patch("docsgpt.parser.connectors.share_point.loader.SharePointLoader._ensure_valid_token")
def test_load_file_by_id_returns_document_with_size(self, mock_ensure_token, mock_auth_init, mock_get_token, mock_get):
@@ -134,7 +136,8 @@ class TestSharePointLoaderListItems:
"""Test _list_items_in_parent method."""
@patch("docsgpt.parser.connectors.share_point.loader.requests.get")
@patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth.get_token_info_from_session")
@patch("docsgpt.connectors.service.connection_id_for_session_token", new=lambda token: "conn-1")
@patch("docsgpt.connectors.service.get_valid_token_info")
@patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth.__init__", return_value=None)
@patch("docsgpt.parser.connectors.share_point.loader.SharePointLoader._ensure_valid_token")
def test_list_items_includes_size_in_select(self, mock_ensure_token, mock_auth_init, mock_get_token, mock_get):
@@ -169,7 +172,8 @@ class TestSharePointLoaderListItems:
assert "size" in params["$select"]
@patch("docsgpt.parser.connectors.share_point.loader.requests.get")
@patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth.get_token_info_from_session")
@patch("docsgpt.connectors.service.connection_id_for_session_token", new=lambda token: "conn-1")
@patch("docsgpt.connectors.service.get_valid_token_info")
@patch("docsgpt.parser.connectors.share_point.loader.SharePointAuth.__init__", return_value=None)
@patch("docsgpt.parser.connectors.share_point.loader.SharePointLoader._ensure_valid_token")
def test_list_items_folders_include_size(self, mock_ensure_token, mock_auth_init, mock_get_token, mock_get):
+174
View File
@@ -152,3 +152,177 @@ class TestMigration0038:
with pg_engine.connect() as conn:
assert "connector_key" in _columns(conn, "connector_sessions")
assert "connection_id" in _columns(conn, "user_tools")
def _seed_secrets(conn) -> dict:
"""Plaintext tokens and v1 tool secrets, as a pre-0038 install stored them."""
from docsgpt.security.encryption import encrypt_credentials
ids = {}
ids["drive"] = conn.execute(
text(
"INSERT INTO connector_sessions (user_id, provider, session_token, user_email, status, token_info) "
"VALUES ('carol', 'google_drive', 'tok-c', 'carol@example.com', 'authorized', CAST(:ti AS jsonb)) "
"RETURNING id"
),
{"ti": json.dumps({"access_token": "plain-at", "refresh_token": "plain-rt", "scopes": ["drive"]})},
).scalar()
ids["mcp"] = conn.execute(
text(
"INSERT INTO connector_sessions (user_id, provider, server_url, session_data) "
"VALUES ('carol', 'mcp:https://mcp.notion.com', 'https://mcp.notion.com', CAST(:sd AS jsonb)) "
"RETURNING id"
),
{"sd": json.dumps({
"tokens": {"access_token": "mcp-at", "refresh_token": "mcp-rt"},
"client_info": {"client_id": "cid", "client_secret": "dcr-secret"},
"other": 1,
})},
).scalar()
ids["mcp_pending"] = conn.execute(
text(
"INSERT INTO connector_sessions (user_id, provider, server_url, session_data) "
"VALUES ('carol', 'mcp:https://half.example.com', 'https://half.example.com', CAST(:sd AS jsonb)) "
"RETURNING id"
),
{"sd": json.dumps({"client_info": {"client_id": "x"}})},
).scalar()
for name, token in (("telegram_a", "111111:AAAAAAAA"), ("telegram_b", "111111:AAAAAAAA")):
ids[name] = conn.execute(
text("INSERT INTO user_tools (user_id, name, config) VALUES ('carol', 'telegram', CAST(:c AS jsonb)) "
"RETURNING id"),
{"c": json.dumps({"encrypted_credentials": encrypt_credentials({"token": token}, "carol")})},
).scalar()
ids["mcp_key_tool"] = conn.execute(
text("INSERT INTO user_tools (user_id, name, config) VALUES ('carol', 'mcp_tool', CAST(:c AS jsonb)) "
"RETURNING id"),
{"c": json.dumps({
"server_url": "https://tools.example.com/mcp", "auth_type": "bearer",
"encrypted_credentials": encrypt_credentials({"bearer_token": "bearer-secret-9999"}, "carol"),
})},
).scalar()
ids["mcp_oauth_tool"] = conn.execute(
text("INSERT INTO user_tools (user_id, name, config) VALUES ('carol', 'mcp_tool', CAST(:c AS jsonb)) "
"RETURNING id"),
{"c": json.dumps({"server_url": "https://mcp.notion.com/mcp", "auth_type": "oauth"})},
).scalar()
return ids
class TestMigration0038Credentials:
def _upgrade_with(self, pg_engine, seed):
url = pg_engine.url.render_as_string(hide_password=False)
_run_alembic(url, "downgrade", _0037)
with pg_engine.begin() as conn:
ids = seed(conn)
_run_alembic(url, "upgrade", "head")
return url, ids
def test_tokens_are_encrypted_and_plaintext_removed(self, pg_engine):
from docsgpt.connectors.service import read_secrets
_, ids = self._upgrade_with(pg_engine, _seed_secrets)
with pg_engine.connect() as conn:
rows = {
r.id: dict(r._mapping)
for r in conn.execute(text("SELECT * FROM connector_sessions WHERE user_id = 'carol'"))
}
dump = json.dumps([
{k: v for k, v in r.items() if k not in ("encrypted_credentials",)} for r in rows.values()
], default=str)
# Nothing readable is left outside the envelope.
for secret in ("plain-at", "plain-rt", "mcp-at", "mcp-rt", "dcr-secret", "AAAAAAAA", "bearer-secret"):
assert secret not in dump
drive = rows[ids["drive"]]
assert drive["token_info"] is None
assert drive["has_refresh_token"] is True
assert drive["scopes"] == ["drive"]
assert drive["status"] == "connected"
assert read_secrets(drive)["token_info"]["refresh_token"] == "plain-rt"
mcp = rows[ids["mcp"]]
assert mcp["session_data"] == {"other": 1}
assert mcp["status"] == "connected"
assert read_secrets(mcp)["client_info"]["client_secret"] == "dcr-secret"
assert rows[ids["mcp_pending"]]["status"] == "pending"
def test_api_key_tools_share_one_connection(self, pg_engine):
from docsgpt.connectors.service import read_secrets
_, ids = self._upgrade_with(pg_engine, _seed_secrets)
with pg_engine.connect() as conn:
links = dict(conn.execute(
text("SELECT id, connection_id FROM user_tools WHERE user_id = 'carol'")
).fetchall())
assert links[ids["telegram_a"]] is not None
assert links[ids["telegram_a"]] == links[ids["telegram_b"]]
telegram = dict(conn.execute(
text("SELECT * FROM connector_sessions WHERE id = :i"), {"i": links[ids["telegram_a"]]}
).one()._mapping)
assert telegram["auth_kind"] == "api_key"
assert telegram["account_label"] == "…AAAA"
assert read_secrets(telegram) == {"credentials": {"token": "111111:AAAAAAAA"}}
bearer = dict(conn.execute(
text("SELECT * FROM connector_sessions WHERE id = :i"), {"i": links[ids["mcp_key_tool"]]}
).one()._mapping)
assert (bearer["connector_key"], bearer["server_url"]) == ("custom_mcp", "https://tools.example.com")
# Rollback safety: the tool keeps its v1 copy for a release.
config = conn.execute(
text("SELECT config FROM user_tools WHERE id = :i"), {"i": ids["telegram_a"]}
).scalar()
assert "encrypted_credentials" in config
def test_oauth_mcp_tools_keep_member_credentials(self, pg_engine):
_, ids = self._upgrade_with(pg_engine, _seed_secrets)
with pg_engine.connect() as conn:
modes = dict(conn.execute(
text("SELECT id, credential_mode FROM user_tools WHERE user_id = 'carol'")
).fetchall())
assert modes[ids["mcp_oauth_tool"]] == "member"
assert modes[ids["telegram_a"]] == "owner"
def test_multiple_accounts_per_provider_allowed(self, pg_engine):
with pg_engine.begin() as conn:
for label in ("a@example.com", "b@example.com"):
conn.execute(
text("INSERT INTO connector_sessions (user_id, provider, account_label, status) "
"VALUES ('dan', 'google_drive', :l, 'connected')"),
{"l": label},
)
with pg_engine.connect() as conn:
assert conn.execute(
text("SELECT count(*) FROM connector_sessions WHERE user_id = 'dan'")
).scalar() == 2
def test_downgrade_restores_plaintext(self, pg_engine):
url, ids = self._upgrade_with(pg_engine, _seed_secrets)
_run_alembic(url, "downgrade", _0037)
with pg_engine.connect() as conn:
token_info = conn.execute(
text("SELECT token_info FROM connector_sessions WHERE id = :i"), {"i": ids["drive"]}
).scalar()
session_data = conn.execute(
text("SELECT session_data FROM connector_sessions WHERE id = :i"), {"i": ids["mcp"]}
).scalar()
api_rows = conn.execute(
text("SELECT count(*) FROM connector_sessions WHERE provider = 'telegram'")
).scalar()
assert token_info["refresh_token"] == "plain-rt"
assert session_data["tokens"]["access_token"] == "mcp-at"
assert api_rows == 0
_run_alembic(url, "upgrade", "head")
def test_second_upgrade_is_a_no_op(self, pg_engine):
from docsgpt.connectors.service import read_secrets
url, ids = self._upgrade_with(pg_engine, _seed_secrets)
_run_alembic(url, "downgrade", _0037)
_run_alembic(url, "upgrade", "head")
with pg_engine.connect() as conn:
drive = dict(conn.execute(
text("SELECT * FROM connector_sessions WHERE id = :i"), {"i": ids["drive"]}
).one()._mapping)
count = conn.execute(
text("SELECT count(*) FROM connector_sessions WHERE provider = 'telegram'")
).scalar()
assert read_secrets(drive)["token_info"]["access_token"] == "plain-at"
assert count == 1
+9 -1
View File
@@ -44,6 +44,10 @@ def _patch_upload_db(conn):
"docsgpt.api.user.sources.upload.db_session", _yield
), patch(
"docsgpt.api.user.sources.upload.db_readonly", _yield
), patch(
"docsgpt.connectors.service.db_session", _yield
), patch(
"docsgpt.connectors.service.db_readonly", _yield
):
yield
@@ -58,6 +62,10 @@ def _patch_conn_db(conn):
"docsgpt.api.connector.routes.db_session", _yield
), patch(
"docsgpt.api.connector.routes.db_readonly", _yield
), patch(
"docsgpt.connectors.service.db_session", _yield
), patch(
"docsgpt.connectors.service.db_readonly", _yield
):
yield
@@ -210,7 +218,7 @@ class TestConnectorExtra:
yield
with patch(
"docsgpt.api.connector.routes.db_session", _broken
"docsgpt.connectors.service.db_readonly", _broken
), app.test_request_context(
"/api/connectors/disconnect", method="POST",
json={"provider": "x", "session_token": "y"},
+2 -1
View File
@@ -24,6 +24,7 @@ def _mock_connector_pipeline(monkeypatch):
from docsgpt import worker
fake_connector = MagicMock(name="connector")
fake_connector.connection_id = None
fake_connector.download_to_directory.return_value = {
"files_downloaded": 1,
"empty_result": False,
@@ -34,7 +35,7 @@ def _mock_connector_pipeline(monkeypatch):
monkeypatch.setattr(
worker.ConnectorCreator,
"create_connector",
staticmethod(lambda source_type, session_token: fake_connector),
staticmethod(lambda source_type, session_token=None, connection_id=None: fake_connector),
)
fake_reader = MagicMock(name="reader")
+2 -1
View File
@@ -125,7 +125,8 @@ class TestSyncWorker:
captured: list[dict] = []
def _fake_sync(self, source_data, name_job, user, loader,
sync_frequency, retriever, doc_id=None, directory="temp"):
sync_frequency, retriever, doc_id=None, directory="temp",
connection_id=None):
captured.append({
"name_job": name_job,
"user": user,
+2 -1
View File
@@ -470,6 +470,7 @@ def _stub_connector_pipeline(
from docsgpt import worker
fake_connector = MagicMock(name="connector")
fake_connector.connection_id = None
fake_connector.download_to_directory.return_value = {
"files_downloaded": files_downloaded,
"empty_result": empty_result,
@@ -482,7 +483,7 @@ def _stub_connector_pipeline(
monkeypatch.setattr(
worker.ConnectorCreator,
"create_connector",
staticmethod(lambda source_type, session_token: fake_connector),
staticmethod(lambda source_type, session_token=None, connection_id=None: fake_connector),
)
fake_reader = MagicMock(name="reader")