diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index bc7f809a..2e9faf9f 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -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. diff --git a/docsgpt/agents/tools/brave.py b/docsgpt/agents/tools/brave.py index 532e310f..c58d784a 100644 --- a/docsgpt/agents/tools/brave.py +++ b/docsgpt/agents/tools/brave.py @@ -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." diff --git a/docsgpt/agents/tools/mcp_tool.py b/docsgpt/agents/tools/mcp_tool.py index 061f4c6d..c3cfdc91 100644 --- a/docsgpt/agents/tools/mcp_tool.py +++ b/docsgpt/agents/tools/mcp_tool.py @@ -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 ( diff --git a/docsgpt/agents/tools/ntfy.py b/docsgpt/agents/tools/ntfy.py index 8d52d3f5..d719211f 100644 --- a/docsgpt/agents/tools/ntfy.py +++ b/docsgpt/agents/tools/ntfy.py @@ -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 " diff --git a/docsgpt/agents/tools/postgres.py b/docsgpt/agents/tools/postgres.py index 35aaab81..d465faee 100644 --- a/docsgpt/agents/tools/postgres.py +++ b/docsgpt/agents/tools/postgres.py @@ -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", diff --git a/docsgpt/agents/tools/telegram.py b/docsgpt/agents/tools/telegram.py index cf8ca67a..138a0240 100644 --- a/docsgpt/agents/tools/telegram.py +++ b/docsgpt/agents/tools/telegram.py @@ -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." diff --git a/docsgpt/alembic/versions/0038_connections.py b/docsgpt/alembic/versions/0038_connections.py index 3ce6aa43..42586c75 100644 --- a/docsgpt/alembic/versions/0038_connections.py +++ b/docsgpt/alembic/versions/0038_connections.py @@ -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;") diff --git a/docsgpt/api/connector/connections.py b/docsgpt/api/connector/connections.py index 46cb9edc..32bbd0af 100644 --- a/docsgpt/api/connector/connections.py +++ b/docsgpt/api/connector/connections.py @@ -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//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//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//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//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//tools//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//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) diff --git a/docsgpt/api/connector/routes.py b/docsgpt/api/connector/routes.py index 5240ee58..03822102 100644 --- a/docsgpt/api/connector/routes.py +++ b/docsgpt/api/connector/routes.py @@ -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, diff --git a/docsgpt/api/user/sources/routes.py b/docsgpt/api/user/sources/routes.py index 7ca3b339..8b13abd1 100644 --- a/docsgpt/api/user/sources/routes.py +++ b/docsgpt/api/user/sources/routes.py @@ -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) diff --git a/docsgpt/api/user/sources/upload.py b/docsgpt/api/user/sources/upload.py index f15bee65..52b62fd1 100644 --- a/docsgpt/api/user/sources/upload.py +++ b/docsgpt/api/user/sources/upload.py @@ -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"], diff --git a/docsgpt/api/user/tasks.py b/docsgpt/api/user/tasks.py index b34945fe..94f5fcb6 100644 --- a/docsgpt/api/user/tasks.py +++ b/docsgpt/api/user/tasks.py @@ -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, diff --git a/docsgpt/api/user/tools/mcp.py b/docsgpt/api/user/tools/mcp.py index bcc99084..cd69bf75 100644 --- a/docsgpt/api/user/tools/mcp.py +++ b/docsgpt/api/user/tools/mcp.py @@ -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:"`` - # 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( diff --git a/docsgpt/api/user/tools/routes.py b/docsgpt/api/user/tools/routes.py index c4464b68..db72a847 100644 --- a/docsgpt/api/user/tools/routes.py +++ b/docsgpt/api/user/tools/routes.py @@ -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( diff --git a/docsgpt/cli.py b/docsgpt/cli.py index d837642d..c1264dc7 100644 --- a/docsgpt/cli.py +++ b/docsgpt/cli.py @@ -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="") + 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) diff --git a/docsgpt/connectors/mcp.py b/docsgpt/connectors/mcp.py new file mode 100644 index 00000000..8ae389d7 --- /dev/null +++ b/docsgpt/connectors/mcp.py @@ -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} diff --git a/docsgpt/connectors/resolve.py b/docsgpt/connectors/resolve.py new file mode 100644 index 00000000..8580b842 --- /dev/null +++ b/docsgpt/connectors/resolve.py @@ -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, + }, + ) diff --git a/docsgpt/connectors/service.py b/docsgpt/connectors/service.py index 3bb16f72..52c9f31d 100644 --- a/docsgpt/connectors/service.py +++ b/docsgpt/connectors/service.py @@ -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 diff --git a/docsgpt/llm/handlers/base.py b/docsgpt/llm/handlers/base.py index 0954dd53..a8090056 100644 --- a/docsgpt/llm/handlers/base.py +++ b/docsgpt/llm/handlers/base.py @@ -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) diff --git a/docsgpt/parser/connectors/base.py b/docsgpt/parser/connectors/base.py index 5d2dd706..1b0f3c4d 100644 --- a/docsgpt/parser/connectors/base.py +++ b/docsgpt/parser/connectors/base.py @@ -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]: diff --git a/docsgpt/parser/connectors/confluence/auth.py b/docsgpt/parser/connectors/confluence/auth.py index 431e99c7..6a5cf0fe 100644 --- a/docsgpt/parser/connectors/confluence/auth.py +++ b/docsgpt/parser/connectors/confluence/auth.py @@ -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]: diff --git a/docsgpt/parser/connectors/confluence/loader.py b/docsgpt/parser/connectors/confluence/loader.py index 38c379dc..646913b0 100644 --- a/docsgpt/parser/connectors/confluence/loader.py +++ b/docsgpt/parser/connectors/confluence/loader.py @@ -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") diff --git a/docsgpt/parser/connectors/google_drive/auth.py b/docsgpt/parser/connectors/google_drive/auth.py index ddff3e4e..0a10d203 100644 --- a/docsgpt/parser/connectors/google_drive/auth.py +++ b/docsgpt/parser/connectors/google_drive/auth.py @@ -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. diff --git a/docsgpt/parser/connectors/google_drive/loader.py b/docsgpt/parser/connectors/google_drive/loader.py index 9a8d7a97..f37a6560 100644 --- a/docsgpt/parser/connectors/google_drive/loader.py +++ b/docsgpt/parser/connectors/google_drive/loader.py @@ -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: diff --git a/docsgpt/parser/connectors/share_point/auth.py b/docsgpt/parser/connectors/share_point/auth.py index ec006740..4188f143 100644 --- a/docsgpt/parser/connectors/share_point/auth.py +++ b/docsgpt/parser/connectors/share_point/auth.py @@ -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 diff --git a/docsgpt/parser/connectors/share_point/loader.py b/docsgpt/parser/connectors/share_point/loader.py index 6fa3e2c3..4690b303 100644 --- a/docsgpt/parser/connectors/share_point/loader.py +++ b/docsgpt/parser/connectors/share_point/loader.py @@ -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") diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index 147f4316..4e8d052a 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -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)), ) diff --git a/docsgpt/storage/db/repositories/connector_sessions.py b/docsgpt/storage/db/repositories/connector_sessions.py index a521406c..1ec86927 100644 --- a/docsgpt/storage/db/repositories/connector_sessions.py +++ b/docsgpt/storage/db/repositories/connector_sessions.py @@ -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 = diff --git a/docsgpt/storage/db/repositories/user_tools.py b/docsgpt/storage/db/repositories/user_tools.py index ff235858..0c8e5ca7 100644 --- a/docsgpt/storage/db/repositories/user_tools.py +++ b/docsgpt/storage/db/repositories/user_tools.py @@ -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 diff --git a/docsgpt/worker.py b/docsgpt/worker.py index f3311859..83f7abc0 100755 --- a/docsgpt/worker.py +++ b/docsgpt/worker.py @@ -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( diff --git a/tests/agents/tools/test_mcp_tool.py b/tests/agents/tools/test_mcp_tool.py index d418ef0a..d24da614 100644 --- a/tests/agents/tools/test_mcp_tool.py +++ b/tests/agents/tools/test_mcp_tool.py @@ -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 diff --git a/tests/api/test_connector_oauth_security.py b/tests/api/test_connector_oauth_security.py index f8b2e521..d91641c7 100644 --- a/tests/api/test_connector_oauth_security.py +++ b/tests/api/test_connector_oauth_security.py @@ -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 diff --git a/tests/api/test_connector_routes_happy.py b/tests/api/test_connector_routes_happy.py index 141f87b0..52176f0a 100644 --- a/tests/api/test_connector_routes_happy.py +++ b/tests/api/test_connector_routes_happy.py @@ -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"} diff --git a/tests/api/user/sources/test_upload_idempotency.py b/tests/api/user/sources/test_upload_idempotency.py index f5b069bc..fe251ff4 100644 --- a/tests/api/user/sources/test_upload_idempotency.py +++ b/tests/api/user/sources/test_upload_idempotency.py @@ -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 diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index a3ff2a19..dcf5be29 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -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, diff --git a/tests/connectors/test_connections_api.py b/tests/connectors/test_connections_api.py new file mode 100644 index 00000000..e02144c8 --- /dev/null +++ b/tests/connectors/test_connections_api.py @@ -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 diff --git a/tests/connectors/test_runtime.py b/tests/connectors/test_runtime.py new file mode 100644 index 00000000..b02f2442 --- /dev/null +++ b/tests/connectors/test_runtime.py @@ -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 diff --git a/tests/connectors/test_service_credentials.py b/tests/connectors/test_service_credentials.py new file mode 100644 index 00000000..884cef77 --- /dev/null +++ b/tests/connectors/test_service_credentials.py @@ -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"}}) diff --git a/tests/connectors/test_service_read.py b/tests/connectors/test_service_read.py index 01c8a6d6..2388b6c2 100644 --- a/tests/connectors/test_service_read.py +++ b/tests/connectors/test_service_read.py @@ -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 diff --git a/tests/parser/connectors/confluence/test_auth_session.py b/tests/parser/connectors/confluence/test_auth_session.py index c4d11797..49ef5f43 100644 --- a/tests/parser/connectors/confluence/test_auth_session.py +++ b/tests/parser/connectors/confluence/test_auth_session.py @@ -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" diff --git a/tests/parser/connectors/confluence/test_loader.py b/tests/parser/connectors/confluence/test_loader.py index ba891d66..61cc3433 100644 --- a/tests/parser/connectors/confluence/test_loader.py +++ b/tests/parser/connectors/confluence/test_loader.py @@ -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"}) diff --git a/tests/parser/connectors/test_auth_token_leakage.py b/tests/parser/connectors/test_auth_token_leakage.py index 98d63aa4..d997620a 100644 --- a/tests/parser/connectors/test_auth_token_leakage.py +++ b/tests/parser/connectors/test_auth_token_leakage.py @@ -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 diff --git a/tests/parser/connectors/test_connector_creator.py b/tests/parser/connectors/test_connector_creator.py index 535441db..321aa016 100644 --- a/tests/parser/connectors/test_connector_creator.py +++ b/tests/parser/connectors/test_connector_creator.py @@ -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) diff --git a/tests/parser/connectors/test_google_drive_auth.py b/tests/parser/connectors/test_google_drive_auth.py index 610e6d11..63c50a4a 100644 --- a/tests/parser/connectors/test_google_drive_auth.py +++ b/tests/parser/connectors/test_google_drive_auth.py @@ -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 diff --git a/tests/parser/connectors/test_google_drive_loader.py b/tests/parser/connectors/test_google_drive_loader.py index ed765e60..3a6dcb67 100644 --- a/tests/parser/connectors/test_google_drive_loader.py +++ b/tests/parser/connectors/test_google_drive_loader.py @@ -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 diff --git a/tests/parser/connectors/test_share_point_auth.py b/tests/parser/connectors/test_share_point_auth.py index ecf25a7d..23529ddd 100644 --- a/tests/parser/connectors/test_share_point_auth.py +++ b/tests/parser/connectors/test_share_point_auth.py @@ -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") diff --git a/tests/parser/connectors/test_share_point_loader.py b/tests/parser/connectors/test_share_point_loader.py index b54eb02b..4390bc2d 100644 --- a/tests/parser/connectors/test_share_point_loader.py +++ b/tests/parser/connectors/test_share_point_loader.py @@ -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"): diff --git a/tests/parser/connectors/token_patch.py b/tests/parser/connectors/token_patch.py new file mode 100644 index 00000000..dafcd115 --- /dev/null +++ b/tests/parser/connectors/token_patch.py @@ -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 diff --git a/tests/parser/remote/test_share_point_loader.py b/tests/parser/remote/test_share_point_loader.py index 07944ee6..5ee13392 100644 --- a/tests/parser/remote/test_share_point_loader.py +++ b/tests/parser/remote/test_share_point_loader.py @@ -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): diff --git a/tests/storage/db/test_migration_0038.py b/tests/storage/db/test_migration_0038.py index 5bdee6ff..2740f33a 100644 --- a/tests/storage/db/test_migration_0038.py +++ b/tests/storage/db/test_migration_0038.py @@ -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 diff --git a/tests/test_extra_coverage.py b/tests/test_extra_coverage.py index fcb3333f..4731df1c 100644 --- a/tests/test_extra_coverage.py +++ b/tests/test_extra_coverage.py @@ -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"}, diff --git a/tests/worker/test_ingest_connector.py b/tests/worker/test_ingest_connector.py index 998b6d38..2e566bc6 100644 --- a/tests/worker/test_ingest_connector.py +++ b/tests/worker/test_ingest_connector.py @@ -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") diff --git a/tests/worker/test_remote_worker.py b/tests/worker/test_remote_worker.py index 4f4e0bb6..2eab9b54 100644 --- a/tests/worker/test_remote_worker.py +++ b/tests/worker/test_remote_worker.py @@ -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, diff --git a/tests/worker/test_worker_sse_publishes.py b/tests/worker/test_worker_sse_publishes.py index f50869b5..e8a1ba1e 100644 --- a/tests/worker/test_worker_sse_publishes.py +++ b/tests/worker/test_worker_sse_publishes.py @@ -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")