From 5321a9438865d3bb736b6157724b477ceb1253ad Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Mon, 28 Sep 2026 16:30:48 +0100 Subject: [PATCH] Add the connector catalog and read-only connections API connector_sessions already holds one row per signed-in account or MCP server. Migration 0038 names each row (connector_key, auth_kind, display_name, account_label) and lets sources and tools point at the row they use through connection_id, backfilling the links it can prove. docsgpt/connectors/catalog.py declares every service DocsGPT connects to, with the server settings it needs and what a connection provides. New endpoints: GET /api/connectors/catalog, GET /api/connections, GET /api/connections/ and POST /api/connections//disconnect. /api/available_tools now groups tools into built_in, service and custom, and source listings carry connectionId. --- docsgpt/alembic/versions/0038_connections.py | 136 ++++++ docsgpt/api/connector/connections.py | 106 +++++ docsgpt/api/pat/rules.py | 1 + docsgpt/api/user/sources/routes.py | 8 + docsgpt/api/user/tools/routes.py | 14 + docsgpt/app.py | 4 + docsgpt/connectors/__init__.py | 6 + docsgpt/connectors/catalog.py | 423 ++++++++++++++++++ docsgpt/connectors/permissions.py | 83 ++++ docsgpt/connectors/service.py | 264 +++++++++++ docsgpt/storage/db/models.py | 13 + .../db/repositories/connector_sessions.py | 78 +++- tests/connectors/__init__.py | 0 tests/connectors/test_catalog.py | 116 +++++ tests/connectors/test_permissions.py | 47 ++ tests/connectors/test_service_read.py | 263 +++++++++++ tests/storage/db/test_migration_0038.py | 154 +++++++ 17 files changed, 1714 insertions(+), 2 deletions(-) create mode 100644 docsgpt/alembic/versions/0038_connections.py create mode 100644 docsgpt/api/connector/connections.py create mode 100644 docsgpt/connectors/__init__.py create mode 100644 docsgpt/connectors/catalog.py create mode 100644 docsgpt/connectors/permissions.py create mode 100644 docsgpt/connectors/service.py create mode 100644 tests/connectors/__init__.py create mode 100644 tests/connectors/test_catalog.py create mode 100644 tests/connectors/test_permissions.py create mode 100644 tests/connectors/test_service_read.py create mode 100644 tests/storage/db/test_migration_0038.py diff --git a/docsgpt/alembic/versions/0038_connections.py b/docsgpt/alembic/versions/0038_connections.py new file mode 100644 index 00000000..3ce6aa43 --- /dev/null +++ b/docsgpt/alembic/versions/0038_connections.py @@ -0,0 +1,136 @@ +"""0038 connections — connector_sessions becomes the connections table. + +``connector_sessions`` already holds one row per signed-in account (OAuth +ingest providers) or per MCP server. This migration names what each row is +and lets sources and tools point at the row they use: + +* ``connector_key`` is the catalog entry (``google_drive``, ``custom_mcp``, + ``telegram``), ``auth_kind`` how the row signs in, ``display_name`` and + ``account_label`` what the Connectors page shows. +* ``sources.connection_id`` and ``user_tools.connection_id`` link the + resources a connection feeds. ``ON DELETE SET NULL`` keeps a source's + indexed content when its connection is removed. + +Backfill (idempotent, only fills NULLs): + +1. ``connector_key``, ``auth_kind``, ``display_name`` from ``provider``. +2. ``account_label`` from ``user_email`` for OAuth rows. +3. ``sources.connection_id`` for ``connector:file`` sources, matched to the + 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. + +Revision ID: 0038_connections +Revises: 0037_request_traces +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0038_connections" +down_revision: Union[str, None] = "0037_request_traces" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +_OAUTH_PROVIDERS = { + "google_drive": "Google Drive", + "share_point": "SharePoint", + "confluence": "Confluence", +} + + +def upgrade() -> None: + op.execute( + """ + ALTER TABLE connector_sessions + ADD COLUMN IF NOT EXISTS connector_key TEXT, + ADD COLUMN IF NOT EXISTS display_name TEXT, + ADD COLUMN IF NOT EXISTS account_label TEXT, + ADD COLUMN IF NOT EXISTS auth_kind TEXT, + ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ NOT NULL DEFAULT now(); + """ + ) + op.execute( + "ALTER TABLE sources ADD COLUMN IF NOT EXISTS connection_id UUID " + "REFERENCES connector_sessions(id) ON DELETE SET NULL;" + ) + op.execute( + "ALTER TABLE user_tools ADD COLUMN IF NOT EXISTS connection_id UUID " + "REFERENCES connector_sessions(id) ON DELETE SET NULL;" + ) + op.execute("CREATE INDEX IF NOT EXISTS sources_connection_idx ON sources (connection_id);") + op.execute("CREATE INDEX IF NOT EXISTS user_tools_connection_idx ON user_tools (connection_id);") + + # 1 + 2: name the existing rows. + for provider, name in _OAUTH_PROVIDERS.items(): + op.execute( + f""" + UPDATE connector_sessions SET + connector_key = COALESCE(connector_key, '{provider}'), + auth_kind = COALESCE(auth_kind, 'oauth'), + display_name = COALESCE(display_name, '{name}'), + account_label = COALESCE(account_label, user_email) + WHERE provider = '{provider}'; + """ + ) + op.execute( + """ + UPDATE connector_sessions SET + connector_key = COALESCE(connector_key, 'custom_mcp'), + auth_kind = COALESCE(auth_kind, 'mcp_oauth'), + display_name = COALESCE( + display_name, + regexp_replace(COALESCE(server_url, substr(provider, 5)), '^https?://', '') + ) + WHERE provider LIKE 'mcp:%'; + """ + ) + + # 3: connector sources point at the owner's only session for the provider. + op.execute( + """ + UPDATE sources s SET connection_id = cs.id + FROM connector_sessions cs + WHERE s.connection_id IS NULL + AND s.type = 'connector:file' + AND cs.user_id = s.user_id + AND cs.provider = s.remote_data->>'provider' + AND ( + SELECT count(*) FROM connector_sessions c2 + WHERE c2.user_id = s.user_id AND c2.provider = cs.provider + ) = 1; + """ + ) + + # 4: OAuth MCP tools point at the owner's session for the server's base URL. + op.execute( + """ + UPDATE user_tools t SET connection_id = cs.id + FROM connector_sessions cs + WHERE t.connection_id IS NULL + AND t.name = 'mcp_tool' + AND t.config->>'auth_type' = 'oauth' + AND cs.user_id = t.user_id + AND cs.provider = 'mcp:' || substring(t.config->>'server_url' from '^(https?://[^/]+)'); + """ + ) + + +def downgrade() -> None: + 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;") + op.execute("ALTER TABLE sources DROP COLUMN IF EXISTS connection_id;") + op.execute( + """ + ALTER TABLE connector_sessions + DROP COLUMN IF EXISTS updated_at, + DROP COLUMN IF EXISTS auth_kind, + DROP COLUMN IF EXISTS account_label, + DROP COLUMN IF EXISTS display_name, + DROP COLUMN IF EXISTS connector_key; + """ + ) diff --git a/docsgpt/api/connector/connections.py b/docsgpt/api/connector/connections.py new file mode 100644 index 00000000..46cb9edc --- /dev/null +++ b/docsgpt/api/connector/connections.py @@ -0,0 +1,106 @@ +"""Connectors catalog and connections API. + +``/api/connectors/catalog`` lists every service DocsGPT can connect to, with +whether the server is set up for it and the caller's connection summary. +``/api/connections`` lists and manages the caller's connections. Responses +never include tokens or secrets. +""" + +from __future__ import annotations + +from flask import current_app, jsonify, make_response, request +from flask_restx import Namespace, Resource + +from docsgpt.api import api +from docsgpt.api.user.authz import ROLE_ADMIN, has_role +from docsgpt.connectors import service +from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository +from docsgpt.storage.db.session import db_readonly, db_session + +connections_ns = Namespace("connections", description="Connectors and connections", path="/api") +api.add_namespace(connections_ns) + + +def _user_id() -> str | None: + token = getattr(request, "decoded_token", None) + return token.get("sub") if isinstance(token, dict) else None + + +def _unauthorized(): + return make_response(jsonify({"success": False, "error": "Unauthorized"}), 401) + + +def _not_found(): + return make_response(jsonify({"success": False, "error": "Connection not found"}), 404) + + +@connections_ns.route("/connectors/catalog") +class ConnectorCatalog(Resource): + @api.doc(description="Every connector with its availability and the caller's connection summary") + def get(self): + user_id = _user_id() + if not user_id: + return _unauthorized() + try: + with db_readonly() as conn: + entries = service.catalog_for_user( + conn, user_id, is_admin=has_role(request.decoded_token, ROLE_ADMIN), + ) + except Exception as err: + current_app.logger.error(f"Error building connector catalog: {err}", exc_info=True) + return make_response(jsonify({"success": False, "error": "Failed to load connectors"}), 500) + return make_response(jsonify({"success": True, "connectors": entries}), 200) + + +@connections_ns.route("/connections") +class ConnectionsList(Resource): + @api.doc(description="The caller's connections with status and linked resource counts") + def get(self): + user_id = _user_id() + if not user_id: + return _unauthorized() + try: + with db_readonly() as conn: + connections = service.list_connections(conn, user_id) + except Exception as err: + current_app.logger.error(f"Error listing connections: {err}", exc_info=True) + return make_response(jsonify({"success": False, "error": "Failed to load connections"}), 500) + return make_response(jsonify({"success": True, "connections": connections}), 200) + + +@connections_ns.route("/connections/") +class ConnectionDetail(Resource): + @api.doc(description="One connection with the sources it syncs and the tools it provides") + def get(self, connection_id: str): + user_id = _user_id() + if not user_id: + return _unauthorized() + try: + with db_readonly() as conn: + row = ConnectorSessionsRepository(conn).get_for_user(connection_id, user_id) + if row is None: + return _not_found() + detail = service.connection_detail(conn, row) + except Exception as err: + current_app.logger.error(f"Error loading connection: {err}", exc_info=True) + return make_response(jsonify({"success": False, "error": "Failed to load connection"}), 500) + return make_response(jsonify({"success": True, "connection": detail}), 200) + + +@connections_ns.route("/connections//disconnect") +class ConnectionDisconnect(Resource): + @api.doc(description="Delete a connection's stored credentials; its sources and tools stay") + def post(self, connection_id: str): + user_id = _user_id() + if not user_id: + return _unauthorized() + try: + with db_session() as conn: + row = ConnectorSessionsRepository(conn).get_for_user(connection_id, user_id) + if row is None: + return _not_found() + connection = service.disconnect(conn, row) + except Exception as err: + 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) diff --git a/docsgpt/api/pat/rules.py b/docsgpt/api/pat/rules.py index c1d7683d..b1ebe9e4 100644 --- a/docsgpt/api/pat/rules.py +++ b/docsgpt/api/pat/rules.py @@ -357,6 +357,7 @@ DENIED_PREFIXES = ( "/api/admin/", "/api/auth/oidc/", "/api/connectors/", + "/api/connections", "/api/devices", "/scim/", "/static/", diff --git a/docsgpt/api/user/sources/routes.py b/docsgpt/api/user/sources/routes.py index 6826c9e9..7ca3b339 100644 --- a/docsgpt/api/user/sources/routes.py +++ b/docsgpt/api/user/sources/routes.py @@ -67,6 +67,12 @@ def _get_provider_from_remote_data(remote_data): return None +def _connection_id(row: dict) -> str | None: + """The connection a source syncs from, as a string id, or None.""" + value = row.get("connection_id") + return str(value) if value else None + + @sources_ns.route("/sources") class CombinedJson(Resource): @api.doc(description="Provide JSON file with combined available indexes") @@ -117,6 +123,7 @@ class CombinedJson(Resource): "config": SourceConfig.parse(index.get("config")).model_dump(), "ownership": ownership, "team_access": team_access, + "connectionId": _connection_id(index), } for index in indexes: @@ -207,6 +214,7 @@ class PaginatedSources(Resource): "team_access": ( None if owned else team_shared.get(str(doc["id"])) ), + "connectionId": _connection_id(doc), } ) response = { diff --git a/docsgpt/api/user/tools/routes.py b/docsgpt/api/user/tools/routes.py index ddc56916..c4464b68 100644 --- a/docsgpt/api/user/tools/routes.py +++ b/docsgpt/api/user/tools/routes.py @@ -19,6 +19,7 @@ from docsgpt.api import api from docsgpt.api.pat.rules import filter_listing from docsgpt.api.user.artifacts.authz import Principal, authorize_artifact from docsgpt.api.user.team_sharing import effective_write_owner, visible_with_access +from docsgpt.connectors.catalog import definition_for_tool from docsgpt.core.settings import settings from docsgpt.core.url_validation import SSRFError, validate_url from docsgpt.security.encryption import decrypt_credentials, encrypt_credentials @@ -186,6 +187,10 @@ def transform_actions(actions_metadata): tools_ns = Namespace("tools", description="Tool management operations", path="/api") +# Tools the Connectors page adds through "Add custom connector" rather than +# the Add Tool modal. +_CUSTOM_CONNECTOR_TOOLS = {"mcp_tool": "custom_mcp", "api_tool": "custom_openapi"} + @tools_ns.route("/available_tools") class AvailableTools(Resource): @@ -202,6 +207,13 @@ class AvailableTools(Resource): description = lines[1].strip() if len(lines) > 1 else "" config_req = tool_instance.get_config_requirements() actions = tool_instance.get_actions_metadata() + definition = definition_for_tool(tool_name) + if definition is not None: + group, connector_key = "service", definition.key + elif tool_name in _CUSTOM_CONNECTOR_TOOLS: + group, connector_key = "custom", _CUSTOM_CONNECTOR_TOOLS[tool_name] + else: + group, connector_key = "built_in", None tools_metadata.append( { "name": tool_name, @@ -209,6 +221,8 @@ class AvailableTools(Resource): "description": description, "configRequirements": config_req, "actions": actions, + "group": group, + "connector_key": connector_key, } ) except Exception as err: diff --git a/docsgpt/app.py b/docsgpt/app.py index 88e2d5cd..5a45ab31 100644 --- a/docsgpt/app.py +++ b/docsgpt/app.py @@ -29,6 +29,7 @@ from docsgpt.api.scim import scim_bp # noqa: E402 from docsgpt.api.user.authz import ROLE_USER, resolve_roles # noqa: E402 from docsgpt.api.user.routes import user # noqa: E402 from docsgpt.api.connector.routes import connector # noqa: E402 +from docsgpt.api.connector import connections as _connections # noqa: E402,F401 from docsgpt.api.v1 import v1_bp # noqa: E402 from docsgpt.celery_init import celery # noqa: E402 from docsgpt.core.secret_key import resolve_jwt_secret_key # noqa: E402 @@ -209,6 +210,9 @@ def get_config(): "hybrid_available": settings.VECTOR_STORE == "pgvector", "tts_available": TTSCreator.is_enabled(settings.TTS_PROVIDER), "stt_available": STTCreator.is_enabled(settings.STT_PROVIDER), + # Lets a frontend built before the Connectors page run against this + # backend, and a new frontend hide the page against an older one. + "connectors_enabled": True, } if settings.AUTH_TYPE == "oidc": response["oidc"] = { diff --git a/docsgpt/connectors/__init__.py b/docsgpt/connectors/__init__.py new file mode 100644 index 00000000..fc58c5c0 --- /dev/null +++ b/docsgpt/connectors/__init__.py @@ -0,0 +1,6 @@ +"""Connectors: one place to connect accounts that feed Sources and Tools. + +``catalog`` declares the services, ``service`` manages connections (the rows +of ``connector_sessions``), ``permissions`` classifies tool actions as read +or write. +""" diff --git a/docsgpt/connectors/catalog.py b/docsgpt/connectors/catalog.py new file mode 100644 index 00000000..dcc4e39a --- /dev/null +++ b/docsgpt/connectors/catalog.py @@ -0,0 +1,423 @@ +"""The connector catalog: every service DocsGPT can connect to. + +A connector is something the user connects once (an OAuth sign-in, an API key +or an MCP server). Each connection it produces can feed Sources (content +synced into DocsGPT) and Tools (actions an agent can take). The definitions +here are declarative: they say how a connector signs in, which server +settings it needs, what it can sync and which tools it creates, so the API +and the frontend never hard-code a list of services. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable, Optional +from urllib.parse import urlparse + +import yaml + +from docsgpt.core.settings import settings + +CATEGORIES = ("files", "knowledge", "dev", "messaging", "database", "search", "custom") +AUTH_KINDS = ("oauth", "mcp_oauth", "api_key", "none", "mcp") +PUBLISHERS = ("built_in", "preset", "custom") + + +@dataclass(frozen=True) +class CredentialField: + """One field the user fills in to connect an ``api_key`` connector. + + Attributes: + key: Name the credential is stored and passed to the tool or loader + under, e.g. ``token`` or ``aws_access_key_id``. + label: English label; the frontend shows the translated + ``settings.connectors.fields.`` when it has one. + secret: Masked in the form and never returned by the API. + required: Whether connecting fails without it. + """ + + key: str + label: str + secret: bool = True + required: bool = True + + def to_dict(self) -> dict: + return {"key": self.key, "label": self.label, "secret": self.secret, "required": self.required} + + +@dataclass(frozen=True) +class ConnectorDefinition: + """How one connector signs in and what a connection to it provides. + + Attributes: + key: Catalog key, stored on each connection as ``connector_key``. + name: Display name. + description: One line shown on the catalog card. + icon: Frontend asset key (``assets/connectors/.svg``). + category: One of :data:`CATEGORIES`. + auth_kind: ``oauth`` (server-side OAuth app), ``mcp_oauth`` (MCP + OAuth with dynamic client registration), ``api_key`` (fields the + user pastes), ``none`` or ``mcp`` (a custom MCP server whose auth + the user picks). + capabilities: Any of ``sync``, ``read``, ``write``. + credential_fields: Fields asked for by ``api_key`` connectors. + required_settings: Server settings that must be set for the + connector to be usable, e.g. ``GOOGLE_CLIENT_ID``. + sync_ingestor: Ingest loader a synced source uses (``google_drive``, + ``s3``), or None when the connector cannot sync. + default_sync_frequency: Sync frequency preselected in the wizard. + setup_fields: Per-source fields asked when choosing what to sync + (an S3 bucket, Reddit search queries). + tool_templates: ``user_tools`` names created on connect. + setup: What the wizard does after sign-in: ``tools`` is ``auto`` + (created and enabled), ``ask`` or ``off``; ``sync`` likewise. + mcp_url: MCP endpoint for presets. + publisher: ``built_in``, ``preset`` or ``custom``. + docs_url: Setup guide for admins. + oauth_scopes: Scopes an MCP preset requests. + """ + + key: str + name: str + description: str + icon: str + category: str + auth_kind: str + capabilities: tuple[str, ...] = () + credential_fields: tuple[CredentialField, ...] = () + required_settings: tuple[str, ...] = () + sync_ingestor: Optional[str] = None + default_sync_frequency: str = "weekly" + setup_fields: tuple[CredentialField, ...] = () + tool_templates: tuple[str, ...] = () + setup: dict = field(default_factory=lambda: {"tools": "auto", "sync": "ask"}) + mcp_url: Optional[str] = None + publisher: str = "built_in" + docs_url: Optional[str] = None + oauth_scopes: tuple[str, ...] = () + + @property + def missing_settings(self) -> list[str]: + """Server settings this connector needs that are not set.""" + return [name for name in self.required_settings if not getattr(settings, name, None)] + + @property + def configured(self) -> bool: + """Whether every required server setting is set.""" + return not self.missing_settings + + @property + def mcp_base_url(self) -> Optional[str]: + """``scheme://host`` of the preset's MCP endpoint (the MCP session key).""" + return base_url(self.mcp_url) if self.mcp_url else None + + def to_dict(self) -> dict: + """Serialise the parts the frontend needs (never server secrets).""" + return { + "key": self.key, + "name": self.name, + "description": self.description, + "icon": self.icon, + "category": self.category, + "auth_kind": self.auth_kind, + "capabilities": list(self.capabilities), + "credential_fields": [f.to_dict() for f in self.credential_fields], + "setup_fields": [f.to_dict() for f in self.setup_fields], + "sync_ingestor": self.sync_ingestor, + "default_sync_frequency": self.default_sync_frequency, + "tool_templates": list(self.tool_templates), + "setup": dict(self.setup), + "mcp_url": self.mcp_url, + "publisher": self.publisher, + "docs_url": self.docs_url, + "oauth_scopes": list(self.oauth_scopes), + } + + +def base_url(url: Optional[str]) -> str: + """``scheme://netloc`` of ``url``, the key MCP sessions are stored under.""" + parsed = urlparse(url or "") + if not parsed.scheme or not parsed.netloc: + return "" + return f"{parsed.scheme}://{parsed.netloc}" + + +_DOCS = "https://docs.docsgpt.cloud/Guides/Connectors" + +_BUILT_IN: tuple[ConnectorDefinition, ...] = ( + ConnectorDefinition( + key="google_drive", + name="Google Drive", + description="Sync Docs, Sheets and PDFs into Sources.", + icon="drive", + category="files", + auth_kind="oauth", + capabilities=("sync",), + required_settings=("GOOGLE_CLIENT_ID", "GOOGLE_CLIENT_SECRET"), + sync_ingestor="google_drive", + setup={"tools": "off", "sync": "ask"}, + docs_url=f"{_DOCS}#google-drive", + ), + ConnectorDefinition( + key="share_point", + name="SharePoint", + description="Sync files from SharePoint sites and OneDrive into Sources.", + icon="sharepoint", + category="files", + auth_kind="oauth", + capabilities=("sync",), + required_settings=("MICROSOFT_CLIENT_ID", "MICROSOFT_CLIENT_SECRET"), + sync_ingestor="share_point", + setup={"tools": "off", "sync": "ask"}, + docs_url=f"{_DOCS}#sharepoint-and-onedrive", + ), + ConnectorDefinition( + key="confluence", + name="Confluence", + description="Sync Confluence spaces and pages into Sources.", + icon="confluence", + category="knowledge", + auth_kind="oauth", + capabilities=("sync",), + required_settings=("CONFLUENCE_CLIENT_ID", "CONFLUENCE_CLIENT_SECRET"), + sync_ingestor="confluence", + setup={"tools": "off", "sync": "ask"}, + docs_url=f"{_DOCS}#confluence", + ), + ConnectorDefinition( + key="s3", + name="Amazon S3", + description="Sync documents from an S3 bucket into Sources.", + icon="s3", + category="files", + auth_kind="api_key", + capabilities=("sync",), + credential_fields=( + CredentialField("aws_access_key_id", "Access key ID", secret=False), + CredentialField("aws_secret_access_key", "Secret access key"), + CredentialField("region", "Region", secret=False, required=False), + CredentialField("endpoint_url", "Custom endpoint URL", secret=False, required=False), + ), + sync_ingestor="s3", + setup_fields=( + CredentialField("bucket", "Bucket", secret=False), + CredentialField("prefix", "Path prefix", secret=False, required=False), + ), + setup={"tools": "off", "sync": "ask"}, + docs_url=f"{_DOCS}#amazon-s3", + ), + ConnectorDefinition( + key="reddit", + name="Reddit", + description="Sync Reddit posts that match your searches into Sources.", + icon="reddit", + category="knowledge", + auth_kind="api_key", + capabilities=("sync",), + credential_fields=( + CredentialField("client_id", "Client ID", secret=False), + CredentialField("client_secret", "Client secret"), + CredentialField("user_agent", "User agent", secret=False), + ), + sync_ingestor="reddit", + setup_fields=( + CredentialField("search_queries", "Search queries", secret=False), + CredentialField("number_posts", "Number of posts", secret=False), + ), + setup={"tools": "off", "sync": "ask"}, + ), + ConnectorDefinition( + key="brave", + name="Brave Search", + description="Search the web and images with the Brave Search API.", + icon="tool_brave", + category="search", + auth_kind="api_key", + capabilities=("read",), + credential_fields=(CredentialField("token", "API key"),), + tool_templates=("brave",), + setup={"tools": "auto", "sync": "off"}, + ), + ConnectorDefinition( + key="telegram", + name="Telegram", + description="Send messages and images to a Telegram chat.", + icon="tool_telegram", + category="messaging", + auth_kind="api_key", + capabilities=("write",), + credential_fields=(CredentialField("token", "Bot token"),), + tool_templates=("telegram",), + setup={"tools": "auto", "sync": "off"}, + ), + ConnectorDefinition( + key="ntfy", + name="ntfy", + description="Send push notifications through an ntfy server.", + icon="tool_ntfy", + category="messaging", + auth_kind="api_key", + capabilities=("write",), + credential_fields=(CredentialField("token", "Access token"),), + tool_templates=("ntfy",), + setup={"tools": "auto", "sync": "off"}, + ), + ConnectorDefinition( + key="postgres", + name="PostgreSQL", + description="Read the schema and run SQL against a Postgres database.", + icon="tool_postgres", + category="database", + auth_kind="api_key", + capabilities=("read", "write"), + credential_fields=(CredentialField("token", "Connection string"),), + tool_templates=("postgres",), + setup={"tools": "auto", "sync": "off"}, + ), + ConnectorDefinition( + key="custom_mcp", + name="MCP server", + description="Connect any remote Model Context Protocol server.", + icon="tool_mcp_tool", + category="custom", + auth_kind="mcp", + capabilities=("read", "write"), + tool_templates=("mcp_tool",), + setup={"tools": "auto", "sync": "off"}, + publisher="custom", + ), + ConnectorDefinition( + key="custom_openapi", + name="OpenAPI / REST", + description="Import an OpenAPI spec and call its endpoints as tools.", + icon="tool_api_tool", + category="custom", + auth_kind="none", + capabilities=("read", "write"), + tool_templates=("api_tool",), + setup={"tools": "ask", "sync": "off"}, + publisher="custom", + ), +) + +_PRESETS_FILE = Path(__file__).parent / "presets" / "mcp.yaml" + + +def _load_presets(path: Optional[Path] = None) -> tuple[ConnectorDefinition, ...]: + """Read the curated MCP server presets shipped with the repository.""" + path = path or _PRESETS_FILE + if not path.exists(): + return () + with path.open(encoding="utf-8") as fh: + entries = yaml.safe_load(fh) or [] + presets = [] + for entry in entries: + presets.append( + ConnectorDefinition( + key=entry["key"], + name=entry["name"], + description=entry["description"], + icon=entry.get("icon") or "tool_mcp_tool", + category=entry.get("category", "knowledge"), + auth_kind=entry.get("auth_kind", "mcp_oauth"), + capabilities=tuple(entry.get("capabilities") or ("read", "write")), + tool_templates=("mcp_tool",), + setup={"tools": "auto", "sync": "off"}, + mcp_url=entry["mcp_url"], + publisher="preset", + docs_url=entry.get("docs_url"), + oauth_scopes=tuple(entry.get("oauth_scopes") or ()), + ) + ) + return tuple(presets) + + +_REGISTRY: dict[str, ConnectorDefinition] = {} + + +def _registry() -> dict[str, ConnectorDefinition]: + if not _REGISTRY: + for definition in (*_BUILT_IN, *_load_presets()): + _REGISTRY[definition.key] = definition + return _REGISTRY + + +def all_definitions() -> list[ConnectorDefinition]: + """Every catalog entry, built-ins first, then presets.""" + return list(_registry().values()) + + +def get_definition(key: Optional[str]) -> Optional[ConnectorDefinition]: + """The definition for ``key``, or None when it is not in the catalog.""" + if not key: + return None + return _registry().get(key) + + +def preset_for_url(url: Optional[str]) -> Optional[ConnectorDefinition]: + """The MCP preset whose server shares ``url``'s base URL, if any.""" + target = base_url(url) + if not target: + return None + for definition in _registry().values(): + if definition.publisher == "preset" and definition.mcp_base_url == target: + return definition + return None + + +def definition_for_tool(tool_name: str) -> Optional[ConnectorDefinition]: + """The built-in connector that provides the ``user_tools`` template ``tool_name``.""" + for definition in _BUILT_IN: + if definition.publisher == "built_in" and tool_name in definition.tool_templates: + return definition + return None + + +def connector_key_for_row(row: dict) -> Optional[str]: + """Catalog key for a ``connector_sessions`` row. + + Rows written before ``0038_connections`` carry no key; they are named from + ``provider``. A custom MCP row whose server matches a preset is reported + as that preset. + """ + key = row.get("connector_key") + provider = row.get("provider") or "" + if not key: + if provider.startswith("mcp:"): + key = "custom_mcp" + elif get_definition(provider): + key = provider + if key == "custom_mcp": + preset = preset_for_url(row.get("server_url") or provider[4:]) + if preset: + return preset.key + return key + + +def tool_connector_keys() -> set[str]: + """``user_tools`` names that belong to a built-in service connector.""" + return { + name + for definition in _BUILT_IN + if definition.publisher == "built_in" + for name in definition.tool_templates + } + + +def iter_by_category(definitions: Iterable[ConnectorDefinition]) -> dict[str, list[ConnectorDefinition]]: + """Group definitions by category, keeping :data:`CATEGORIES` order.""" + grouped: dict[str, list[ConnectorDefinition]] = {c: [] for c in CATEGORIES} + for definition in definitions: + grouped.setdefault(definition.category, []).append(definition) + return grouped + + +def reset_registry_for_tests() -> None: + """Drop the cached registry so a test can load a different presets file.""" + _REGISTRY.clear() + + +def to_public(definition: ConnectorDefinition, **extra: Any) -> dict: + """Definition dict plus request-specific fields (availability, counts).""" + return {**definition.to_dict(), **extra} diff --git a/docsgpt/connectors/permissions.py b/docsgpt/connectors/permissions.py new file mode 100644 index 00000000..b5afce78 --- /dev/null +++ b/docsgpt/connectors/permissions.py @@ -0,0 +1,83 @@ +"""Read / write classification and permissions for connection-backed tools. + +Every action a connection's tool offers is either a *read* (it only looks +something up) or a *write* (it changes or sends something). Reads default to +"Always allow"; writes default to "Needs approval". The permission is stored +on the action itself: ``active`` off means "Off", ``require_approval`` means +"Needs approval", neither means "Always allow". +""" + +from __future__ import annotations + +from typing import Optional + +ACCESS_READ = "read" +ACCESS_WRITE = "write" + +PERMISSION_ALWAYS = "always" +PERMISSION_ASK = "ask" +PERMISSION_OFF = "off" +PERMISSIONS = (PERMISSION_ALWAYS, PERMISSION_ASK, PERMISSION_OFF) + +_READ_WORDS = ("search", "query", "find", "list", "get", "read", "fetch", "lookup", "describe", "retrieve") + + +def action_access(tool_name: Optional[str], action: dict) -> str: + """Whether ``action`` reads or writes. + + Order of evidence: an explicit ``access`` on the action metadata, MCP tool + annotations (``readOnlyHint`` / ``destructiveHint``), the HTTP method of an + API tool action, then the action's name. + + Args: + tool_name: The ``user_tools`` name the action belongs to. + action: One entry of the tool's ``actions``. + + Returns: + ``read`` or ``write``. + """ + access = action.get("access") + if access in (ACCESS_READ, ACCESS_WRITE): + return access + annotations = action.get("annotations") or {} + if isinstance(annotations, dict): + if annotations.get("readOnlyHint") is True: + return ACCESS_READ + if annotations.get("destructiveHint") is True or annotations.get("readOnlyHint") is False: + return ACCESS_WRITE + method = (action.get("method") or "").upper() + if tool_name == "api_tool" and method: + return ACCESS_READ if method in ("GET", "HEAD", "OPTIONS") else ACCESS_WRITE + name = (action.get("name") or "").lower() + return ACCESS_READ if any(word in name for word in _READ_WORDS) else ACCESS_WRITE + + +def action_permission(action: dict) -> str: + """The action's permission: ``always``, ``ask`` or ``off``.""" + if action.get("active") is False: + return PERMISSION_OFF + if action.get("require_approval"): + return PERMISSION_ASK + return PERMISSION_ALWAYS + + +def apply_permission(action: dict, permission: str) -> dict: + """Return ``action`` with ``permission`` written onto its flags.""" + if permission not in PERMISSIONS: + raise ValueError(f"Unknown permission: {permission}") + updated = dict(action) + updated["active"] = permission != PERMISSION_OFF + updated["require_approval"] = permission == PERMISSION_ASK + return updated + + +def apply_default_permissions(tool_name: Optional[str], actions: list[dict]) -> list[dict]: + """Stamp ``access`` on each action and default writes to needing approval.""" + stamped = [] + for action in actions: + access = action_access(tool_name, action) + updated = {**action, "access": access} + if access == ACCESS_WRITE: + updated["require_approval"] = True + stamped.append(updated) + return stamped diff --git a/docsgpt/connectors/service.py b/docsgpt/connectors/service.py new file mode 100644 index 00000000..3bb16f72 --- /dev/null +++ b/docsgpt/connectors/service.py @@ -0,0 +1,264 @@ +"""Connection lifecycle: listing, status, and what each connection provides. + +A connection is one ``connector_sessions`` row: a signed-in OAuth account, an +MCP server or a set of API credentials. Sources and tools point at it through +``connection_id``. Everything the API returns about a connection goes through +:func:`serialize_connection`, which never includes credentials. +""" + +from __future__ import annotations + +import json +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 + +STATUS_CONNECTED = "connected" +STATUS_RECONNECT = "reconnect_needed" +STATUS_DISCONNECTED = "disconnected" +STATUS_ERROR = "error" +STATUS_PENDING = "pending" + +# Worst first: a card shows the worst status among the user's accounts. +_STATUS_SEVERITY = { + STATUS_RECONNECT: 4, + STATUS_ERROR: 3, + STATUS_DISCONNECTED: 2, + STATUS_CONNECTED: 1, + STATUS_PENDING: 0, +} + + +def _json(value: Any) -> Any: + if isinstance(value, str): + try: + return json.loads(value) + except ValueError: + return None + return value + + +def has_credentials(row: dict) -> bool: + """Whether a connection holds usable credentials (without decrypting them).""" + if row.get("encrypted_credentials"): + return True + token_info = _json(row.get("token_info")) or {} + if isinstance(token_info, dict) and token_info.get("access_token"): + return True + session_data = _json(row.get("session_data")) or {} + tokens = session_data.get("tokens") if isinstance(session_data, dict) else None + return bool(isinstance(tokens, dict) and tokens.get("access_token")) + + +def normalize_status(row: dict) -> str: + """Map a row's stored status onto the connection status set. + + Rows written before ``0038_connections`` use ``authorized`` for a + finished OAuth sign-in, and MCP rows carry no status at all; both are + read from whether the row holds credentials. + """ + raw = (row.get("status") or "").lower() + if raw in (STATUS_RECONNECT, STATUS_DISCONNECTED, STATUS_ERROR): + return raw + if raw == STATUS_PENDING and not has_credentials(row): + return STATUS_PENDING + if raw in ("authorized", STATUS_CONNECTED, "active") or has_credentials(row): + return STATUS_CONNECTED + return STATUS_PENDING + + +def worst_status(statuses: Iterable[str]) -> Optional[str]: + """The status that needs the most attention, or None for an empty list.""" + ranked = sorted(statuses, key=lambda s: _STATUS_SEVERITY.get(s, 0), reverse=True) + return ranked[0] if ranked else None + + +def account_label(row: dict) -> str: + """What identifies the account to its owner: an email, a workspace or a host.""" + return ( + row.get("account_label") + or row.get("user_email") + or row.get("display_name") + or catalog.base_url(row.get("server_url")).split("://")[-1] + or "" + ) + + +def _iso(value: Any) -> Optional[str]: + return value.isoformat() if hasattr(value, "isoformat") else value + + +def serialize_connection(row: dict, counts: Optional[dict] = None) -> dict: + """Public shape of a connection. Never includes tokens or secrets.""" + key = catalog.connector_key_for_row(row) + definition = catalog.get_definition(key) + counts = counts or {} + return { + "id": str(row["id"]), + "connector_key": key, + "name": (definition.name if definition and definition.publisher != "custom" else None) + or row.get("display_name") + or (definition.name if definition else key), + "display_name": row.get("display_name"), + "icon": definition.icon if definition else "tool_mcp_tool", + "account_label": account_label(row), + "auth_kind": row.get("auth_kind") or (definition.auth_kind if definition else None), + "status": normalize_status(row), + "server_url": row.get("server_url"), + "last_error": row.get("last_error"), + "created_at": _iso(row.get("created_at")), + "updated_at": _iso(row.get("updated_at")), + "last_used_at": _iso(row.get("last_used_at")), + "source_count": counts.get("sources", 0), + "tool_count": counts.get("tools", 0), + "credential_mode": row.get("credential_mode"), + } + + +def list_connections(conn, user_id: str) -> list[dict]: + """The user's connections, finished sign-ins only, with resource counts.""" + repo = ConnectorSessionsRepository(conn) + rows = [r for r in repo.list_for_user(user_id) if normalize_status(r) != STATUS_PENDING] + counts = repo.resource_counts([str(r["id"]) for r in rows]) + return [serialize_connection(r, counts.get(str(r["id"]))) for r in rows] + + +def serialize_source(row: dict, connection_status: Optional[str] = None) -> dict: + """A source linked to a connection, as the connection drawer lists it.""" + metadata = _json(row.get("metadata")) or {} + sync_state = metadata.get("sync_state") if isinstance(metadata, dict) else None + if connection_status in (STATUS_RECONNECT, STATUS_DISCONNECTED) and not sync_state: + sync_state = "paused_reconnect" + return { + "id": str(row["id"]), + "name": row.get("name"), + "type": row.get("type"), + "last_sync": _iso(row.get("date")), + "sync_frequency": row.get("sync_frequency") or "never", + "sync_state": sync_state or "active", + } + + +def serialize_tool(row: dict) -> dict: + """A tool linked to a connection, with its actions and their permissions.""" + from docsgpt.connectors.permissions import action_access, action_permission + + actions = [] + for action in _json(row.get("actions")) or []: + if not isinstance(action, dict): + continue + access = action_access(row.get("name"), action) + actions.append( + { + "name": action.get("name"), + "description": action.get("description", ""), + "access": access, + "permission": action_permission(action), + } + ) + return { + "id": str(row["id"]), + "name": row.get("name"), + "display_name": row.get("custom_name") or row.get("display_name") or row.get("name"), + "status": bool(row.get("status")), + "credential_mode": row.get("credential_mode") or "owner", + "actions": actions, + } + + +def connection_detail(conn, row: dict) -> dict: + """A connection with the sources and tools it feeds.""" + repo = ConnectorSessionsRepository(conn) + status = normalize_status(row) + sources = [serialize_source(s, status) for s in repo.list_sources(str(row["id"]))] + tools = [serialize_tool(t) for t in repo.list_tools(str(row["id"]))] + detail = serialize_connection(row, {"sources": len(sources), "tools": len(tools)}) + detail["sources"] = sources + detail["tools"] = tools + return detail + + +def _card_state(definition: ConnectorDefinition, available: bool, disabled: bool, status: Optional[str]) -> str: + if disabled: + return "disabled" + if status in (STATUS_RECONNECT, STATUS_ERROR): + return "reconnect" + if status == STATUS_CONNECTED: + return "connected" + if not available: + return "needs_setup" + if definition.publisher == "custom": + return "custom" + return "available" + + +def catalog_for_user(conn, user_id: str, *, is_admin: bool, policies: Optional[dict] = None) -> list[dict]: + """Catalog entries with availability and the caller's connection summary. + + Args: + conn: Open database connection. + user_id: The caller. + is_admin: Admins see the names of missing server settings; everyone + else only learns that setup is needed. + policies: ``connector_key`` to policy row, when admin policies exist. + """ + policies = policies or {} + by_key: dict[str, list[str]] = {} + for connection in list_connections(conn, user_id): + by_key.setdefault(connection["connector_key"], []).append(connection["status"]) + + entries = [] + for definition in catalog.all_definitions(): + policy = policies.get(definition.key) or {} + disabled = policy.get("enabled") is False + missing = definition.missing_settings + available = not missing and not disabled + statuses = by_key.get(definition.key, []) + status = worst_status(statuses) + entries.append( + { + **definition.to_dict(), + "available": available, + "disabled": disabled, + "needs_setup": bool(missing), + "missing_settings": missing if is_admin else [], + "connected_count": sum(1 for s in statuses if s == STATUS_CONNECTED), + "connection_count": len(statuses), + "status": status, + "state": _card_state(definition, available, disabled, status), + "credential_policy": policy.get("credential_mode") or "choose", + } + ) + return entries + + +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. + + Args: + conn: Open database connection inside a transaction. + row: The connection row, already authorised for the caller. + + Returns: + The connection's public shape after the change. + """ + repo = ConnectorSessionsRepository(conn) + 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 {}, + }, + ) + return serialize_connection(repo.get(str(row["id"]))) diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index ebd25567..147f4316 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -216,6 +216,9 @@ user_tools_table = Table( Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), Column("legacy_mongo_id", Text), + Column( + "connection_id", UUID(as_uuid=True), ForeignKey("connector_sessions.id", ondelete="SET NULL"), + ), ) token_usage_table = Table( @@ -338,6 +341,9 @@ sources_table = Table( Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), Column("legacy_mongo_id", Text), + Column( + "connection_id", UUID(as_uuid=True), ForeignKey("connector_sessions.id", ondelete="SET NULL"), + ), ) agents_table = Table( @@ -585,6 +591,13 @@ connector_sessions_table = Table( Column("expires_at", DateTime(timezone=True)), Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), Column("legacy_mongo_id", Text), + # Added in ``0038_connections``: each row is a connection (one signed-in + # account, one MCP server or one set of API credentials). + Column("connector_key", Text), + Column("display_name", Text), + Column("account_label", Text), + Column("auth_kind", Text), + Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), ) diff --git a/docsgpt/storage/db/repositories/connector_sessions.py b/docsgpt/storage/db/repositories/connector_sessions.py index 7a95b9d1..a521406c 100644 --- a/docsgpt/storage/db/repositories/connector_sessions.py +++ b/docsgpt/storage/db/repositories/connector_sessions.py @@ -24,12 +24,13 @@ from typing import Any, Optional from sqlalchemy import Connection, text -from docsgpt.storage.db.base_repository import row_to_dict +from docsgpt.storage.db.base_repository import looks_like_uuid, row_to_dict 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", } _UPDATABLE_JSONB = {"session_data", "token_info"} @@ -176,11 +177,83 @@ class ConnectorSessionsRepository: def list_for_user(self, user_id: str) -> list[dict]: result = self._conn.execute( - text("SELECT * FROM connector_sessions WHERE user_id = :user_id"), + text("SELECT * FROM connector_sessions WHERE user_id = :user_id ORDER BY created_at"), {"user_id": user_id}, ) return [row_to_dict(r) for r in result.fetchall()] + 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): + return None + result = self._conn.execute( + text("SELECT * FROM connector_sessions WHERE id = CAST(:id AS uuid)"), + {"id": str(connection_id)}, + ) + row = result.fetchone() + return row_to_dict(row) if row is not None else None + + def get_for_user(self, connection_id: str, user_id: str) -> Optional[dict]: + """Fetch a connection only when ``user_id`` owns it.""" + row = self.get(connection_id) + if row is None or row.get("user_id") != user_id: + return None + return row + + def get_for_update(self, connection_id: str) -> Optional[dict]: + """Fetch and row-lock a connection until the transaction ends. + + Token refresh holds this lock so two workers refreshing a rotating + refresh token (Microsoft, Atlassian) cannot both spend it. + """ + if not looks_like_uuid(connection_id): + return None + result = self._conn.execute( + text("SELECT * FROM connector_sessions WHERE id = CAST(:id AS uuid) FOR UPDATE"), + {"id": str(connection_id)}, + ) + row = result.fetchone() + return row_to_dict(row) if row is not None else None + + def resource_counts(self, connection_ids: list[str]) -> dict[str, dict[str, int]]: + """Number of sources and tools linked to each connection id.""" + ids = [str(i) for i in connection_ids if looks_like_uuid(str(i))] + counts: dict[str, dict[str, int]] = {i: {"sources": 0, "tools": 0} for i in ids} + if not ids: + return counts + for table, key in (("sources", "sources"), ("user_tools", "tools")): + result = self._conn.execute( + text( + f"SELECT connection_id, count(*) FROM {table} " + "WHERE connection_id = ANY(CAST(:ids AS uuid[])) GROUP BY connection_id" + ), + {"ids": ids}, + ) + for connection_id, count in result.fetchall(): + counts[str(connection_id)][key] = int(count) + return counts + + def list_sources(self, connection_id: str) -> list[dict]: + """Sources synced from a connection, newest first.""" + result = self._conn.execute( + text( + "SELECT id, name, type, date, sync_frequency, metadata, remote_data, user_id " + "FROM sources WHERE connection_id = CAST(:id AS uuid) ORDER BY date DESC" + ), + {"id": str(connection_id)}, + ) + return [row_to_dict(r) for r in result.fetchall()] + + def list_tools(self, connection_id: str) -> list[dict]: + """Tools a connection provides, oldest first.""" + result = self._conn.execute( + text( + "SELECT * FROM user_tools WHERE connection_id = CAST(:id AS uuid) ORDER BY created_at" + ), + {"id": str(connection_id)}, + ) + return [row_to_dict(r) for r in result.fetchall()] + def update(self, session_id: str, fields: dict) -> bool: """Partial update by PG UUID.""" filtered = { @@ -198,6 +271,7 @@ class ConnectorSessionsRepository: else: set_clauses.append(f"{col} = :{col}") params[col] = val + set_clauses.append("updated_at = now()") result = self._conn.execute( text( f"UPDATE connector_sessions SET {', '.join(set_clauses)} " diff --git a/tests/connectors/__init__.py b/tests/connectors/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/connectors/test_catalog.py b/tests/connectors/test_catalog.py new file mode 100644 index 00000000..37e3b86e --- /dev/null +++ b/tests/connectors/test_catalog.py @@ -0,0 +1,116 @@ +"""Tests for the connector catalog.""" + +from __future__ import annotations + +import pytest + +from docsgpt.connectors import catalog + + +@pytest.fixture(autouse=True) +def _fresh_registry(): + catalog.reset_registry_for_tests() + yield + catalog.reset_registry_for_tests() + + +class TestDefinitions: + def test_built_ins_are_registered(self): + keys = {d.key for d in catalog.all_definitions()} + assert { + "google_drive", "share_point", "confluence", "s3", "reddit", + "brave", "telegram", "ntfy", "postgres", "custom_mcp", "custom_openapi", + } <= keys + + def test_every_definition_uses_known_values(self): + for definition in catalog.all_definitions(): + assert definition.category in catalog.CATEGORIES, definition.key + assert definition.auth_kind in catalog.AUTH_KINDS, definition.key + assert definition.publisher in catalog.PUBLISHERS, definition.key + assert set(definition.capabilities) <= {"sync", "read", "write"}, definition.key + + def test_tool_connectors_match_tool_config_requirements(self): + """Credential fields must be the keys the tool reads its secrets from.""" + import docsgpt.api.user # noqa: F401 (mcp_tool imports it; load it first) + from docsgpt.agents.tools.tool_manager import ToolManager + + tools = ToolManager(config={}).tools + for definition in catalog.all_definitions(): + if definition.publisher != "built_in": + continue + for tool_name in definition.tool_templates: + requirements = tools[tool_name].get_config_requirements() + secret_keys = {k for k, spec in requirements.items() if spec.get("secret")} + field_keys = {f.key for f in definition.credential_fields} + assert secret_keys <= field_keys, definition.key + + def test_to_dict_has_no_server_secrets(self): + payload = catalog.get_definition("google_drive").to_dict() + assert "required_settings" not in payload + assert payload["capabilities"] == ["sync"] + + +class TestAvailability: + def test_missing_settings_listed(self, monkeypatch): + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", None) + monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", "secret") + definition = catalog.get_definition("google_drive") + assert definition.missing_settings == ["GOOGLE_CLIENT_ID"] + assert not definition.configured + + def test_configured_when_all_set(self, monkeypatch): + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", "id") + monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", "secret") + assert catalog.get_definition("google_drive").configured + + def test_api_key_connectors_need_no_settings(self): + assert catalog.get_definition("telegram").configured + + +class TestRowMapping: + def test_oauth_provider_maps_to_key(self): + assert catalog.connector_key_for_row({"provider": "share_point"}) == "share_point" + + def test_legacy_mcp_row_is_custom(self): + row = {"provider": "mcp:https://mcp.unknown.dev", "server_url": "https://mcp.unknown.dev"} + assert catalog.connector_key_for_row(row) == "custom_mcp" + + def test_stored_key_wins(self): + assert catalog.connector_key_for_row({"provider": "telegram", "connector_key": "telegram"}) == "telegram" + + def test_unknown_provider(self): + assert catalog.connector_key_for_row({"provider": "nope"}) is None + + def test_definition_for_tool(self): + assert catalog.definition_for_tool("telegram").key == "telegram" + assert catalog.definition_for_tool("memory") is None + assert catalog.definition_for_tool("mcp_tool") is None + + +class TestPresets: + def test_presets_load_from_yaml(self, tmp_path, monkeypatch): + presets = tmp_path / "mcp.yaml" + presets.write_text( + "- key: mcp:example\n" + " name: Example\n" + " description: Example records.\n" + " icon: example\n" + " category: knowledge\n" + " mcp_url: https://mcp.example.com/mcp\n" + " auth_kind: mcp_oauth\n" + ) + monkeypatch.setattr(catalog, "_PRESETS_FILE", presets) + definition = catalog.get_definition("mcp:example") + assert definition.publisher == "preset" + assert definition.mcp_base_url == "https://mcp.example.com" + assert catalog.preset_for_url("https://mcp.example.com/other").key == "mcp:example" + row = {"provider": "mcp:https://mcp.example.com", "server_url": "https://mcp.example.com"} + assert catalog.connector_key_for_row(row) == "mcp:example" + + def test_base_url(self): + assert catalog.base_url("https://a.example.com:8443/x/y") == "https://a.example.com:8443" + assert catalog.base_url("not a url") == "" diff --git a/tests/connectors/test_permissions.py b/tests/connectors/test_permissions.py new file mode 100644 index 00000000..bb4318e1 --- /dev/null +++ b/tests/connectors/test_permissions.py @@ -0,0 +1,47 @@ +"""Tests for read / write classification and permissions.""" + +from __future__ import annotations + +import pytest + +from docsgpt.connectors import permissions as p + + +class TestActionAccess: + def test_explicit_access_wins(self): + assert p.action_access("telegram", {"name": "search", "access": "write"}) == "write" + + def test_mcp_annotations(self): + assert p.action_access("mcp_tool", {"name": "create_page", "annotations": {"readOnlyHint": True}}) == "read" + assert p.action_access("mcp_tool", {"name": "search", "annotations": {"destructiveHint": True}}) == "write" + + def test_api_tool_method(self): + assert p.action_access("api_tool", {"name": "x", "method": "get"}) == "read" + assert p.action_access("api_tool", {"name": "x", "method": "POST"}) == "write" + + @pytest.mark.parametrize( + "name,expected", + [("search_pages", "read"), ("list_issues", "read"), ("create_issue", "write"), ("send_message", "write")], + ) + def test_name_heuristic(self, name, expected): + assert p.action_access("mcp_tool", {"name": name}) == expected + + +class TestPermissions: + def test_permission_from_flags(self): + assert p.action_permission({"active": False, "require_approval": True}) == "off" + assert p.action_permission({"active": True, "require_approval": True}) == "ask" + assert p.action_permission({"active": True}) == "always" + + def test_apply_permission(self): + assert p.apply_permission({"name": "a"}, "ask") == {"name": "a", "active": True, "require_approval": True} + assert p.apply_permission({"name": "a"}, "off")["active"] is False + with pytest.raises(ValueError): + p.apply_permission({}, "sometimes") + + def test_defaults_writes_to_approval(self): + stamped = p.apply_default_permissions( + "mcp_tool", [{"name": "search"}, {"name": "delete_page"}] + ) + assert stamped[0] == {"name": "search", "access": "read"} + assert stamped[1]["access"] == "write" and stamped[1]["require_approval"] is True diff --git a/tests/connectors/test_service_read.py b/tests/connectors/test_service_read.py new file mode 100644 index 00000000..01c8a6d6 --- /dev/null +++ b/tests/connectors/test_service_read.py @@ -0,0 +1,263 @@ +"""Tests for the read side of the connection service and its API.""" + +from __future__ import annotations + +import json +from contextlib import contextmanager +from unittest.mock import patch + +import pytest +from flask import Flask +from sqlalchemy import text + +from docsgpt.connectors import service + + +def _session(conn, user="alice", provider="google_drive", **cols) -> str: + cols = {"status": "authorized", "user_email": f"{user}@example.com", **cols} + casts = {"token_info": "jsonb", "session_data": "jsonb"} + names = ", ".join(["user_id", "provider", *cols]) + values = ", ".join( + [":user_id", ":provider", *[f"CAST(:{k} AS {casts[k]})" if k in casts else f":{k}" for k in cols]] + ) + return str( + conn.execute( + text(f"INSERT INTO connector_sessions ({names}) VALUES ({values}) RETURNING id"), + {"user_id": user, "provider": provider, **cols}, + ).scalar() + ) + + +def _source(conn, connection_id, user="alice", name="Handbook") -> str: + return str( + conn.execute( + text( + "INSERT INTO sources (user_id, name, type, sync_frequency, connection_id) " + "VALUES (:u, :n, 'connector:file', 'weekly', CAST(:c AS uuid)) RETURNING id" + ), + {"u": user, "n": name, "c": connection_id}, + ).scalar() + ) + + +def _tool(conn, connection_id, user="alice") -> str: + actions = [ + {"name": "search_pages", "description": "Search", "active": True}, + {"name": "create_page", "description": "Create", "active": True, "require_approval": True}, + ] + return str( + conn.execute( + text( + "INSERT INTO user_tools (user_id, name, display_name, actions, connection_id) " + "VALUES (:u, 'mcp_tool', 'Notion', CAST(:a AS jsonb), CAST(:c AS uuid)) RETURNING id" + ), + {"u": user, "a": json.dumps(actions), "c": connection_id}, + ).scalar() + ) + + +class TestNormalizeStatus: + @pytest.mark.parametrize( + "row,expected", + [ + ({"status": "authorized"}, "connected"), + ({"status": "reconnect_needed"}, "reconnect_needed"), + ({"status": "pending"}, "pending"), + ({"status": "pending", "token_info": {"access_token": "x"}}, "connected"), + ({"status": None, "session_data": {"tokens": {"access_token": "x"}}}, "connected"), + ({"status": None, "session_data": {"client_info": {}}}, "pending"), + ({"status": None, "encrypted_credentials": "v2:..."}, "connected"), + ], + ) + def test_statuses(self, row, expected): + assert service.normalize_status(row) == expected + + def test_worst_status(self): + assert service.worst_status(["connected", "reconnect_needed"]) == "reconnect_needed" + assert service.worst_status([]) is None + + +class TestListing: + def test_lists_finished_connections_with_counts(self, pg_conn): + drive = _session(pg_conn) + _session(pg_conn, provider="confluence", status="pending", user_email=None) + _source(pg_conn, drive) + _source(pg_conn, drive, name="Wiki") + connections = service.list_connections(pg_conn, "alice") + assert [c["connector_key"] for c in connections] == ["google_drive"] + only = connections[0] + assert only["name"] == "Google Drive" + assert only["account_label"] == "alice@example.com" + assert only["status"] == "connected" + assert only["source_count"] == 2 + assert "token_info" not in only and "session_token" not in only + + def test_other_users_rows_hidden(self, pg_conn): + _session(pg_conn, user="bob") + assert service.list_connections(pg_conn, "alice") == [] + + def test_detail_lists_sources_and_tools(self, pg_conn): + from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository + + drive = _session(pg_conn, status="reconnect_needed") + _source(pg_conn, drive) + _tool(pg_conn, drive) + row = ConnectorSessionsRepository(pg_conn).get(drive) + detail = service.connection_detail(pg_conn, row) + assert detail["sources"][0]["sync_state"] == "paused_reconnect" + actions = {a["name"]: a for a in detail["tools"][0]["actions"]} + assert actions["search_pages"] == { + "name": "search_pages", "description": "Search", "access": "read", "permission": "always", + } + assert actions["create_page"]["access"] == "write" + assert actions["create_page"]["permission"] == "ask" + + +class TestCatalogForUser: + def test_states(self, pg_conn, monkeypatch): + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", "id") + monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", "secret") + monkeypatch.setattr(settings, "MICROSOFT_CLIENT_ID", None) + _session(pg_conn) + _session(pg_conn, provider="confluence", status="reconnect_needed") + entries = {e["key"]: e for e in service.catalog_for_user(pg_conn, "alice", is_admin=False)} + assert entries["google_drive"]["state"] == "connected" + assert entries["google_drive"]["connected_count"] == 1 + assert entries["confluence"]["state"] == "reconnect" + assert entries["share_point"]["state"] == "needs_setup" + assert entries["share_point"]["missing_settings"] == [] + assert entries["telegram"]["state"] == "available" + assert entries["custom_mcp"]["state"] == "custom" + + def test_admin_sees_missing_setting_names(self, pg_conn, monkeypatch): + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "MICROSOFT_CLIENT_ID", None) + entries = {e["key"]: e for e in service.catalog_for_user(pg_conn, "alice", is_admin=True)} + assert "MICROSOFT_CLIENT_ID" in entries["share_point"]["missing_settings"] + + def test_policy_disables(self, pg_conn): + entries = { + e["key"]: e + for e in service.catalog_for_user( + pg_conn, "alice", is_admin=False, policies={"telegram": {"enabled": False}} + ) + } + assert entries["telegram"]["state"] == "disabled" + assert entries["telegram"]["available"] is False + + +@contextmanager +def _patched_db(conn): + @contextmanager + def _yield(): + yield conn + + with patch("docsgpt.api.connector.connections.db_readonly", _yield): + yield + + +@pytest.fixture +def app(): + return Flask(__name__) + + +def _call(app, resource, path, token, *args): + with app.test_request_context(path): + from flask import request + + request.decoded_token = token + return resource().get(*args) + + +class TestRoutes: + def test_catalog_requires_auth(self, app): + from docsgpt.api.connector.connections import ConnectorCatalog + + assert _call(app, ConnectorCatalog, "/api/connectors/catalog", None).status_code == 401 + + def test_catalog(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectorCatalog + + with _patched_db(pg_conn): + resp = _call(app, ConnectorCatalog, "/api/connectors/catalog", {"sub": "alice"}) + assert resp.status_code == 200 + keys = {c["key"] for c in resp.get_json()["connectors"]} + assert "telegram" in keys + + def test_list_and_detail(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectionDetail, ConnectionsList + + drive = _session(pg_conn) + with _patched_db(pg_conn): + listed = _call(app, ConnectionsList, "/api/connections", {"sub": "alice"}) + detail = _call(app, ConnectionDetail, f"/api/connections/{drive}", {"sub": "alice"}, drive) + assert listed.get_json()["connections"][0]["id"] == drive + assert detail.get_json()["connection"]["id"] == drive + + def test_detail_of_another_users_connection_is_404(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectionDetail + + drive = _session(pg_conn, user="bob") + with _patched_db(pg_conn): + resp = _call(app, ConnectionDetail, f"/api/connections/{drive}", {"sub": "alice"}, drive) + assert resp.status_code == 404 + + def test_detail_of_bad_id_is_404(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectionDetail + + with _patched_db(pg_conn): + resp = _call(app, ConnectionDetail, "/api/connections/nope", {"sub": "alice"}, "nope") + assert resp.status_code == 404 + + +class TestDisconnect: + def test_clears_credentials_keeps_resources(self, pg_conn): + from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository + + drive = _session( + pg_conn, session_token="tok", token_info=json.dumps({"access_token": "at", "refresh_token": "rt"}), + ) + source = _source(pg_conn, drive) + repo = ConnectorSessionsRepository(pg_conn) + result = service.disconnect(pg_conn, repo.get(drive)) + assert result["status"] == "disconnected" + row = repo.get(drive) + assert row["token_info"] is None and row["session_token"] is None + linked = pg_conn.execute( + text("SELECT connection_id FROM sources WHERE id = CAST(:id AS uuid)"), {"id": source} + ).scalar() + assert str(linked) == drive + + def test_mcp_keeps_client_registration(self, pg_conn): + from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository + + mcp = _session( + pg_conn, + provider="mcp:https://mcp.example.com", + status=None, + session_data=json.dumps({"tokens": {"access_token": "x"}, "client_info": {"client_id": "c"}}), + ) + repo = ConnectorSessionsRepository(pg_conn) + service.disconnect(pg_conn, repo.get(mcp)) + assert repo.get(mcp)["session_data"] == {"client_info": {"client_id": "c"}} + + def test_route_rejects_other_users(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectionDisconnect + + drive = _session(pg_conn, user="bob") + + @contextmanager + def _yield(): + yield pg_conn + + with patch("docsgpt.api.connector.connections.db_session", _yield), app.test_request_context( + f"/api/connections/{drive}/disconnect", method="POST" + ): + from flask import request + + request.decoded_token = {"sub": "alice"} + resp = ConnectionDisconnect().post(drive) + assert resp.status_code == 404 diff --git a/tests/storage/db/test_migration_0038.py b/tests/storage/db/test_migration_0038.py new file mode 100644 index 00000000..5bdee6ff --- /dev/null +++ b/tests/storage/db/test_migration_0038.py @@ -0,0 +1,154 @@ +"""Migration tests for 0038_connections: columns, backfill and round trip.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from sqlalchemy import text + + +pytestmark = pytest.mark.integration + +_0037 = "0037_request_traces" + + +def _alembic_ini() -> Path: + return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini" + + +def _run_alembic(url: str, *args: str) -> None: + subprocess.check_call( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args], + timeout=120, + env={**os.environ, "POSTGRES_URI": url}, + ) + + +def _columns(conn, table: str) -> set[str]: + rows = conn.execute( + text("SELECT column_name FROM information_schema.columns WHERE table_name = :t"), + {"t": table}, + ).fetchall() + return {r[0] for r in rows} + + +def _seed_legacy(conn) -> dict: + """Rows as a pre-0038 install left them.""" + ids = {} + ids["drive"] = conn.execute( + text( + "INSERT INTO connector_sessions (user_id, provider, session_token, user_email, status, token_info) " + "VALUES ('alice', 'google_drive', 'tok-a', 'alice@example.com', 'authorized', " + "CAST(:ti AS jsonb)) RETURNING id" + ), + {"ti": json.dumps({"access_token": "at", "refresh_token": "rt"})}, + ).scalar() + ids["mcp"] = conn.execute( + text( + "INSERT INTO connector_sessions (user_id, provider, server_url, session_data) " + "VALUES ('alice', 'mcp:https://mcp.example.com', 'https://mcp.example.com', CAST(:sd AS jsonb)) " + "RETURNING id" + ), + {"sd": json.dumps({"tokens": {"access_token": "m"}})}, + ).scalar() + ids["source"] = conn.execute( + text( + "INSERT INTO sources (user_id, name, type, remote_data) " + "VALUES ('alice', 'Handbook', 'connector:file', CAST(:rd AS jsonb)) RETURNING id" + ), + {"rd": json.dumps({"provider": "google_drive", "file_ids": ["f1"]})}, + ).scalar() + ids["other_source"] = conn.execute( + text( + "INSERT INTO sources (user_id, name, type, remote_data) " + "VALUES ('bob', 'Bob files', 'connector:file', CAST(:rd AS jsonb)) RETURNING id" + ), + {"rd": json.dumps({"provider": "google_drive"})}, + ).scalar() + ids["mcp_tool"] = conn.execute( + text( + "INSERT INTO user_tools (user_id, name, config) " + "VALUES ('alice', 'mcp_tool', CAST(:c AS jsonb)) RETURNING id" + ), + {"c": json.dumps({"server_url": "https://mcp.example.com/mcp", "auth_type": "oauth"})}, + ).scalar() + ids["bearer_tool"] = conn.execute( + text( + "INSERT INTO user_tools (user_id, name, config) " + "VALUES ('alice', 'mcp_tool', CAST(:c AS jsonb)) RETURNING id" + ), + {"c": json.dumps({"server_url": "https://mcp.example.com/mcp", "auth_type": "bearer"})}, + ).scalar() + return ids + + +class TestMigration0038: + def test_head_has_connection_columns(self, pg_engine): + with pg_engine.connect() as conn: + assert {"connector_key", "display_name", "account_label", "auth_kind", "updated_at"} <= _columns( + conn, "connector_sessions" + ) + assert "connection_id" in _columns(conn, "sources") + assert "connection_id" in _columns(conn, "user_tools") + + def test_backfill_links_legacy_rows(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0037) + with pg_engine.begin() as conn: + ids = _seed_legacy(conn) + _run_alembic(url, "upgrade", "head") + + with pg_engine.connect() as conn: + drive = conn.execute( + text( + "SELECT connector_key, auth_kind, display_name, account_label " + "FROM connector_sessions WHERE id = :id" + ), + {"id": ids["drive"]}, + ).one() + assert tuple(drive) == ("google_drive", "oauth", "Google Drive", "alice@example.com") + + mcp = conn.execute( + text("SELECT connector_key, auth_kind, display_name FROM connector_sessions WHERE id = :id"), + {"id": ids["mcp"]}, + ).one() + assert tuple(mcp) == ("custom_mcp", "mcp_oauth", "mcp.example.com") + + linked = conn.execute( + text("SELECT connection_id FROM sources WHERE id = :id"), {"id": ids["source"]} + ).scalar() + assert linked == ids["drive"] + # Bob has no session for the provider: nothing to link. + assert ( + conn.execute( + text("SELECT connection_id FROM sources WHERE id = :id"), {"id": ids["other_source"]} + ).scalar() + is None + ) + tool = conn.execute( + text("SELECT connection_id FROM user_tools WHERE id = :id"), {"id": ids["mcp_tool"]} + ).scalar() + assert tool == ids["mcp"] + # A bearer MCP tool has no OAuth session to point at. + bearer = conn.execute( + text("SELECT connection_id FROM user_tools WHERE id = :id"), {"id": ids["bearer_tool"]} + ).scalar() + assert bearer is None + + def test_downgrade_then_upgrade_is_idempotent(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0037) + with pg_engine.connect() as conn: + assert "connector_key" not in _columns(conn, "connector_sessions") + assert "connection_id" not in _columns(conn, "sources") + _run_alembic(url, "upgrade", "head") + _run_alembic(url, "downgrade", _0037) + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + assert "connector_key" in _columns(conn, "connector_sessions") + assert "connection_id" in _columns(conn, "user_tools")