From b692eaf7adedb095440a9323bdcc4966415f3654 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Mon, 28 Sep 2026 17:30:53 +0100 Subject: [PATCH] Add MCP presets, admin connector policies and connector attribution Six vendor-run MCP servers (Notion, Linear, Atlassian, Sentry, Asana, Stripe) join the catalog as presets from docsgpt/connectors/presets/ mcp.yaml; existing connections to those servers show under them. connector_policies lets an admin turn a connector off or force whose account every share uses, and app_metadata's connectors.allow_custom_mcp turns custom MCP servers off. Both are enforced on the server: new connections, OAuth sign-ins, MCP test and save, and the tools of a disabled connector stop resolving. GET/PUT /api/admin/connectors reads and changes them. A shared tool's owner picks owner or member credentials with PUT /api/connections/tools//credential-mode, and a member running the owner's account always confirms write actions. Tool-call events and retrieved chunks name the connector they came from (key and display name, never the account), and tool calls keep those fields when the conversation is reloaded. --- docsgpt/agents/tool_executor.py | 19 ++ docsgpt/alembic/versions/0038_connections.py | 19 ++ docsgpt/api/admin/__init__.py | 1 + docsgpt/api/admin/connectors.py | 130 ++++++++ docsgpt/api/connector/connections.py | 35 ++ docsgpt/api/connector/routes.py | 2 + docsgpt/api/user/sources/upload.py | 2 + docsgpt/api/user/tools/mcp.py | 31 ++ docsgpt/api/user/tools/routes.py | 2 + docsgpt/connectors/attribution.py | 68 ++++ docsgpt/connectors/presets/mcp.yaml | 71 ++++ docsgpt/connectors/resolve.py | 11 +- docsgpt/connectors/service.py | 53 ++- docsgpt/retriever/base.py | 11 + docsgpt/retriever/classic_rag.py | 6 +- docsgpt/retriever/graph_rag.py | 2 +- docsgpt/storage/db/models.py | 11 + .../db/repositories/connector_policies.py | 58 ++++ .../test_policies_and_attribution.py | 315 ++++++++++++++++++ tests/storage/db/test_migration_0038.py | 1 + 20 files changed, 844 insertions(+), 4 deletions(-) create mode 100644 docsgpt/api/admin/connectors.py create mode 100644 docsgpt/connectors/attribution.py create mode 100644 docsgpt/connectors/presets/mcp.yaml create mode 100644 docsgpt/storage/db/repositories/connector_policies.py create mode 100644 tests/connectors/test_policies_and_attribution.py diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 181612ae..778f4a1a 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -932,6 +932,14 @@ class ToolExecutor: or require_approval ) + # A member running someone else's account (a shared tool in owner + # mode) always confirms write actions, whatever the owner chose for + # themselves. + if not require_approval and resolved is not None and resolved.delegated: + from docsgpt.connectors.permissions import ACCESS_WRITE, action_access + + require_approval = action_access(tool_data.get("name"), action_data) == ACCESS_WRITE + if require_approval: if self.headless: tool_row_id = str(tool_data.get("id") or tool_id) @@ -1286,6 +1294,12 @@ class ToolExecutor: "arguments": call_args, } tool_data = tools_dict[tool_id] + # Name the service a connection-backed tool used, so the chip can show + # its logo ("Searched Notion"). Never the account behind it. + resolved = self._resolve_connection(tool_data) + if resolved is not None and resolved.connector_key: + tool_call_data["connector_key"] = resolved.connector_key + tool_call_data["connector_name"] = resolved.connector_name # Surface the device id on remote_device tool-call events so the # approval UI can wire up the sticky "don't ask again" button. if tool_data.get("name") == "remote_device": @@ -1331,6 +1345,11 @@ class ToolExecutor: else next(action for action in tool_data["actions"] if action["name"] == action_name) ) + if "connector_key" in tool_call_data: + from docsgpt.connectors.permissions import action_access + + tool_call_data["access"] = action_access(tool_data.get("name"), action_data) + query_params, headers, body, parameters = {}, {}, {}, {} param_types = { "query_params": query_params, diff --git a/docsgpt/alembic/versions/0038_connections.py b/docsgpt/alembic/versions/0038_connections.py index b9da954d..962e5c4a 100644 --- a/docsgpt/alembic/versions/0038_connections.py +++ b/docsgpt/alembic/versions/0038_connections.py @@ -40,6 +40,11 @@ Backfill (idempotent, only fills NULLs or unconverted rows): resolved each invoking member's own token before this migration and keep doing so (``member``); owners can switch them in the share dialog. +``connector_policies`` holds the admin's per-connector switches (enabled, +forced credential mode). The instance-wide "Allow custom MCP servers" switch +is the ``connectors.allow_custom_mcp`` key in ``app_metadata`` (absent means +allowed). + Revision ID: 0038_connections Revises: 0037_request_traces """ @@ -75,6 +80,19 @@ _BATCH = 500 def upgrade() -> None: _upgrade_links() _upgrade_credentials() + op.execute( + """ + CREATE TABLE IF NOT EXISTS connector_policies ( + connector_key TEXT PRIMARY KEY, + enabled BOOLEAN NOT NULL DEFAULT true, + credential_mode TEXT NOT NULL DEFAULT 'choose' + CONSTRAINT connector_policies_credential_mode_chk + CHECK (credential_mode IN ('choose', 'owner', 'member')), + updated_by TEXT, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + """ + ) def _upgrade_links() -> None: @@ -433,6 +451,7 @@ def _restore_account_index() -> None: def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS connector_policies;") _downgrade_credentials() op.execute("DROP INDEX IF EXISTS user_tools_connection_idx;") op.execute("DROP INDEX IF EXISTS sources_connection_idx;") diff --git a/docsgpt/api/admin/__init__.py b/docsgpt/api/admin/__init__.py index 1dcdd686..bf47f250 100644 --- a/docsgpt/api/admin/__init__.py +++ b/docsgpt/api/admin/__init__.py @@ -1,5 +1,6 @@ from .routes import admin_ns from . import quotas # noqa: F401 (registers the quota resources on admin_ns) from . import activity # noqa: F401 (registers the activity resources on admin_ns) +from . import connectors # noqa: F401 (registers the connector policy resources on admin_ns) __all__ = ["admin_ns"] diff --git a/docsgpt/api/admin/connectors.py b/docsgpt/api/admin/connectors.py new file mode 100644 index 00000000..5b2689f0 --- /dev/null +++ b/docsgpt/api/admin/connectors.py @@ -0,0 +1,130 @@ +"""Admin > Connectors: turn connectors on or off and set sharing policy. + +Every resource here is behind ``@admin_required`` (the frontend guard is +cosmetic). The page also shows what each OAuth connector still needs from +the server (settings and the redirect URI to register), never their values. +""" + +from __future__ import annotations + +import logging + +from flask import jsonify, make_response, request +from flask_restx import Resource +from sqlalchemy import text + +from docsgpt.api.admin.routes import _actor, admin_ns +from docsgpt.api.user.authz import admin_required +from docsgpt.connectors import catalog, service +from docsgpt.core.settings import settings +from docsgpt.security.encryption import is_default_encryption_key +from docsgpt.storage.db.repositories.app_metadata import AppMetadataRepository +from docsgpt.storage.db.repositories.connector_policies import ( + ALLOW_CUSTOM_MCP_KEY, + CREDENTIAL_POLICIES, + ConnectorPoliciesRepository, +) +from docsgpt.storage.db.session import db_readonly, db_session + +logger = logging.getLogger(__name__) + + +def _connection_counts(conn) -> dict[str, int]: + rows = conn.execute( + text( + "SELECT connector_key, provider, server_url, count(*) AS n FROM connector_sessions " + "WHERE COALESCE(status, '') <> 'pending' GROUP BY connector_key, provider, server_url" + ) + ).fetchall() + counts: dict[str, int] = {} + for row in rows: + key = catalog.connector_key_for_row(dict(row._mapping)) + if key: + counts[key] = counts.get(key, 0) + int(row.n) + return counts + + +def _mcp_redirect_uri() -> str: + from docsgpt.agents.tools.mcp_tool import MCPTool + + return MCPTool._resolve_redirect_uri(MCPTool.__new__(MCPTool), None) + + +@admin_ns.route("/admin/connectors") +class AdminConnectorsResource(Resource): + @admin_required + def get(self): + """Every connector with its policy, setup state and connection count.""" + with db_readonly() as conn: + policies = ConnectorPoliciesRepository(conn).all() + allow_custom = service.custom_mcp_allowed(conn) + counts = _connection_counts(conn) + connectors = [] + for definition in catalog.all_definitions(): + policy = policies.get(definition.key) or {} + connectors.append( + { + "key": definition.key, + "name": definition.name, + "icon": definition.icon, + "category": definition.category, + "publisher": definition.publisher, + "auth_kind": definition.auth_kind, + "enabled": policy.get("enabled", True), + "credential_mode": policy.get("credential_mode", "choose"), + "configured": definition.configured, + "required_settings": [ + {"name": name, "set": bool(getattr(settings, name, None))} + for name in definition.required_settings + ], + "connection_count": counts.get(definition.key, 0), + "docs_url": definition.docs_url, + "mcp_url": definition.mcp_url, + } + ) + return make_response( + jsonify( + { + "success": True, + "connectors": connectors, + "allow_custom_mcp": allow_custom, + "default_encryption_key": is_default_encryption_key(), + "oauth_redirect_uri": settings.CONNECTOR_REDIRECT_BASE_URI, + "mcp_redirect_uri": _mcp_redirect_uri(), + } + ), + 200, + ) + + @admin_required + def put(self): + """Update policies: {policies: {key: {enabled?, credential_mode?}}, allow_custom_mcp?}.""" + body = request.get_json(silent=True) or {} + updates = body.get("policies") or {} + if not isinstance(updates, dict): + return make_response(jsonify({"success": False, "message": "policies must be an object"}), 400) + for key, change in updates.items(): + if catalog.get_definition(key) is None or not isinstance(change, dict): + return make_response(jsonify({"success": False, "message": f"Unknown connector: {key}"}), 400) + mode = change.get("credential_mode") + if mode is not None and mode not in CREDENTIAL_POLICIES: + return make_response(jsonify({"success": False, "message": "Unknown credential mode"}), 400) + if "enabled" in change and not isinstance(change["enabled"], bool): + return make_response(jsonify({"success": False, "message": "enabled must be true or false"}), 400) + allow_custom = body.get("allow_custom_mcp") + if allow_custom is not None and not isinstance(allow_custom, bool): + return make_response(jsonify({"success": False, "message": "allow_custom_mcp must be a boolean"}), 400) + actor = _actor() + with db_session() as conn: + repo = ConnectorPoliciesRepository(conn) + for key, change in updates.items(): + repo.upsert( + key, + enabled=change.get("enabled"), + credential_mode=change.get("credential_mode"), + updated_by=actor, + ) + if allow_custom is not None: + AppMetadataRepository(conn).set(ALLOW_CUSTOM_MCP_KEY, "true" if allow_custom else "false") + logger.info("connector_policies_updated", extra={"admin": actor, "connectors": sorted(updates)}) + return self.get() diff --git a/docsgpt/api/connector/connections.py b/docsgpt/api/connector/connections.py index 32bbd0af..daf1e06e 100644 --- a/docsgpt/api/connector/connections.py +++ b/docsgpt/api/connector/connections.py @@ -97,6 +97,8 @@ class ConnectionsList(Resource): ) except service.EncryptionKeyNotConfigured as err: return _error(str(err), 400, code="encryption_key_default") + except service.ConnectorDisabled as err: + return _error(str(err), 403, code="disabled") except ValueError as err: return _error(str(err), 400) except Exception as err: @@ -448,3 +450,36 @@ class ConnectionRefreshTools(Resource): 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) + + +@connections_ns.route("/connections/tools//credential-mode") +class ToolCredentialMode(Resource): + @api.doc( + description=( + "Whose account a shared connection-backed tool uses: {mode: owner | member}. " + "Owner only; refused when an admin forces a mode for the connector." + ) + ) + def put(self, tool_id: str): + from docsgpt.connectors.resolve import MODE_MEMBER, MODE_OWNER + from docsgpt.storage.db.repositories.user_tools import UserToolsRepository + + user_id = _user_id() + if not user_id: + return _unauthorized() + mode = _json_body().get("mode") + if mode not in (MODE_OWNER, MODE_MEMBER): + return _error("mode must be owner or member", 400) + with db_session() as conn: + tools = UserToolsRepository(conn) + tool = tools.get_any(tool_id, user_id) + if tool is None or tool.get("user_id") != user_id or not tool.get("connection_id"): + return _error("Tool not found", 404) + connection = ConnectorSessionsRepository(conn).get(str(tool["connection_id"])) + forced = service.forced_credential_mode( + conn, catalog.connector_key_for_row(connection) if connection else None, + ) + if forced and forced != mode: + return _error("An admin sets this for every share", 409, code="forced", mode=forced) + tools.update(str(tool["id"]), user_id, {"credential_mode": mode}) + return make_response(jsonify({"success": True, "mode": mode}), 200) diff --git a/docsgpt/api/connector/routes.py b/docsgpt/api/connector/routes.py index 03822102..2d124df8 100644 --- a/docsgpt/api/connector/routes.py +++ b/docsgpt/api/connector/routes.py @@ -222,6 +222,8 @@ class ConnectorAuth(Resource): ) except service.ConnectionUnavailable: return make_response(jsonify({"success": False, "error": "Connection not found"}), 404) + except service.ConnectorDisabled as err: + return make_response(jsonify({"success": False, "error": str(err), "code": "disabled"}), 403) # 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")) diff --git a/docsgpt/api/user/sources/upload.py b/docsgpt/api/user/sources/upload.py index 52b62fd1..9465f4de 100644 --- a/docsgpt/api/user/sources/upload.py +++ b/docsgpt/api/user/sources/upload.py @@ -482,6 +482,8 @@ def _remote_credentials(user, source, config): try: with db_session() as conn: row, _ = service.create_api_key_connection(conn, user, definition, provided) + except service.ConnectorDisabled as err: + return None, None, make_response(jsonify({"success": False, "error": str(err)}), 403) except (service.EncryptionKeyNotConfigured, ValueError): return config, None, None return public, str(row["id"]), None diff --git a/docsgpt/api/user/tools/mcp.py b/docsgpt/api/user/tools/mcp.py index cd69bf75..6b2567eb 100644 --- a/docsgpt/api/user/tools/mcp.py +++ b/docsgpt/api/user/tools/mcp.py @@ -112,6 +112,31 @@ def _mcp_connection(user, config, auth_type, auth_credentials, display_name): return str(row["id"]) if row else None +def _mcp_policy_error(config: dict): + """A 403 when an admin turned this MCP server's connector off, else None. + + A preset's own switch applies to its server; any other server is a + custom connector and needs "Allow custom MCP servers". + """ + from docsgpt.connectors import catalog, service + + preset = catalog.preset_for_url(config.get("server_url")) + key = preset.key if preset else "custom_mcp" + try: + with db_readonly() as conn: + service.ensure_connector_allowed(conn, key) + except service.ConnectorDisabled: + return make_response( + jsonify({"success": False, "error": "This MCP server is turned off by an admin", "code": "disabled"}), + 403, + ) + except Exception: + # The switch is an admin preference; when the policy table cannot be + # read, saving the tool (which needs the database) fails on its own. + current_app.logger.warning("Could not read connector policies", exc_info=True) + return None + + @tools_mcp_ns.route("/mcp_server/test") class TestMCPServerConfig(Resource): @api.expect( @@ -147,6 +172,9 @@ class TestMCPServerConfig(Resource): ) _validate_mcp_server_url(config) + policy_error = _mcp_policy_error(config) + if policy_error is not None: + return policy_error auth_credentials = _extract_auth_credentials(config) test_config = config.copy() @@ -243,6 +271,9 @@ class MCPServerSave(Resource): ) _validate_mcp_server_url(config) + policy_error = _mcp_policy_error(config) + if policy_error is not None: + return policy_error auth_credentials = _extract_auth_credentials(config) auth_type = config.get("auth_type", "none") diff --git a/docsgpt/api/user/tools/routes.py b/docsgpt/api/user/tools/routes.py index db72a847..93c26e5b 100644 --- a/docsgpt/api/user/tools/routes.py +++ b/docsgpt/api/user/tools/routes.py @@ -470,6 +470,8 @@ def _create_connected_tool(user, data, definition, tool_instance): ) except connection_service.EncryptionKeyNotConfigured: return None + except connection_service.ConnectorDisabled as err: + return make_response(jsonify({"success": False, "message": str(err)}), 403) 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) diff --git a/docsgpt/connectors/attribution.py b/docsgpt/connectors/attribution.py new file mode 100644 index 00000000..9719a8a6 --- /dev/null +++ b/docsgpt/connectors/attribution.py @@ -0,0 +1,68 @@ +"""Which connector a retrieved chunk came from, for "From Google Drive" citations. + +Only the connector's key and display name are attached, never the account +behind the connection, so shared conversations show the same attribution +without revealing whose Drive it was. +""" + +from __future__ import annotations + +import logging +from typing import Optional + +from sqlalchemy import text + +from docsgpt.storage.db.base_repository import looks_like_uuid +from docsgpt.storage.db.session import db_readonly + +logger = logging.getLogger(__name__) + + +def connector_labels(source_ids: list[str]) -> dict[str, dict]: + """``{source_id: {"connector_key", "connector_name"}}`` for synced sources. + + Sources that do not come from a connection are left out. Never raises: + attribution is decoration, and a lookup failure must not break an answer. + """ + from docsgpt.connectors import service + + ids = [str(i) for i in source_ids if i and looks_like_uuid(str(i))] + if not ids: + return {} + try: + with db_readonly() as conn: + rows = conn.execute( + text( + "SELECT s.id AS source_id, cs.* FROM sources s " + "JOIN connector_sessions cs ON cs.id = s.connection_id " + "WHERE s.id = ANY(CAST(:ids AS uuid[]))" + ), + {"ids": ids}, + ).fetchall() + except Exception: + logger.warning("connector attribution lookup failed", exc_info=True) + return {} + labels = {} + for row in rows: + data = dict(row._mapping) + public = service.serialize_connection({**data, "id": data["id"]}) + labels[str(data["source_id"])] = { + "connector_key": public["connector_key"], + "connector_name": public["name"], + } + return labels + + +class ConnectorLabelCache: + """Per-search memo so each source is looked up once.""" + + def __init__(self) -> None: + self._labels: dict[str, dict] = {} + + def for_source(self, source_id: Optional[str]) -> dict: + if not source_id: + return {} + key = str(source_id) + if key not in self._labels: + self._labels[key] = connector_labels([key]).get(key, {}) + return self._labels[key] diff --git a/docsgpt/connectors/presets/mcp.yaml b/docsgpt/connectors/presets/mcp.yaml new file mode 100644 index 00000000..c9642c80 --- /dev/null +++ b/docsgpt/connectors/presets/mcp.yaml @@ -0,0 +1,71 @@ +# Curated remote MCP servers shown as Connectors catalog cards. +# +# Inclusion rules: the vendor operates the server, it has a stable public URL, +# and it signs in with OAuth plus dynamic client registration (or needs no +# auth), so connecting takes no admin setup. Each entry was checked for OAuth +# authorization-server metadata with a registration endpoint and for a 401 +# challenge on the MCP endpoint. Keep the list short; anything else is a +# custom connector. +# +# Keys look like ``mcp:``. The frontend reads ``settings.connectors +# .descriptions.mcp_`` for a translated description. + +- key: mcp:notion + name: Notion + description: Search, read and update Notion pages and databases. + icon: notion + category: knowledge + mcp_url: https://mcp.notion.com/mcp + auth_kind: mcp_oauth + capabilities: [read, write] + docs_url: https://developers.notion.com/docs/mcp + +- key: mcp:linear + name: Linear + description: Find, create and update Linear issues and projects. + icon: linear + category: dev + mcp_url: https://mcp.linear.app/mcp + auth_kind: mcp_oauth + capabilities: [read, write] + docs_url: https://linear.app/docs/mcp + +- key: mcp:atlassian + name: Atlassian (Jira and Confluence) + description: Search and update Jira issues and Confluence pages. + icon: atlassian + category: dev + mcp_url: https://mcp.atlassian.com/v1/mcp + auth_kind: mcp_oauth + capabilities: [read, write] + docs_url: https://support.atlassian.com/rovo/docs/getting-started-with-the-atlassian-remote-mcp-server/ + +- key: mcp:sentry + name: Sentry + description: Look up Sentry issues, events and releases. + icon: sentry + category: dev + mcp_url: https://mcp.sentry.dev/mcp + auth_kind: mcp_oauth + capabilities: [read, write] + docs_url: https://docs.sentry.io/product/sentry-mcp/ + +- key: mcp:asana + name: Asana + description: Find and update Asana tasks and projects. + icon: asana + category: knowledge + mcp_url: https://mcp.asana.com/v2/mcp + auth_kind: mcp_oauth + capabilities: [read, write] + docs_url: https://developers.asana.com/docs/using-asanas-mcp-server + +- key: mcp:stripe + name: Stripe + description: Look up customers, payments and subscriptions in Stripe. + icon: stripe + category: dev + mcp_url: https://mcp.stripe.com + auth_kind: mcp_oauth + capabilities: [read, write] + docs_url: https://docs.stripe.com/mcp diff --git a/docsgpt/connectors/resolve.py b/docsgpt/connectors/resolve.py index 8580b842..490d7c7f 100644 --- a/docsgpt/connectors/resolve.py +++ b/docsgpt/connectors/resolve.py @@ -76,6 +76,11 @@ def resolve_connection(resource: dict, invoker_user_id: Optional[str]) -> Option with db_readonly() as conn: repo = ConnectorSessionsRepository(conn) owned = repo.get(str(connection_id)) + owned_key = catalog.connector_key_for_row(owned) if owned else None + policy = service.load_policies(conn).get(owned_key) or {} if owned_key else {} + if policy.get("credential_mode") in (MODE_OWNER, MODE_MEMBER): + # An admin forces whose account every share of this connector uses. + mode = policy["credential_mode"] 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( @@ -86,7 +91,11 @@ def resolve_connection(resource: dict, invoker_user_id: Optional[str]) -> Option 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 + available = ( + row is not None + and service.normalize_status(row) == service.STATUS_CONNECTED + and policy.get("enabled") is not False + ) return ResolvedConnection( row=row, available=available, diff --git a/docsgpt/connectors/service.py b/docsgpt/connectors/service.py index 52c9f31d..5e00d795 100644 --- a/docsgpt/connectors/service.py +++ b/docsgpt/connectors/service.py @@ -214,7 +214,8 @@ def catalog_for_user(conn, user_id: str, *, is_admin: bool, policies: Optional[d else only learns that setup is needed. policies: ``connector_key`` to policy row, when admin policies exist. """ - policies = policies or {} + if policies is None: + policies = load_policies(conn) by_key: dict[str, list[str]] = {} for connection in list_connections(conn, user_id): by_key.setdefault(connection["connector_key"], []).append(connection["status"]) @@ -686,6 +687,7 @@ def create_api_key_connection( else: credentials = {k: v for k, v in credentials.items() if v not in (None, "")} ensure_can_store_credentials() + ensure_connector_allowed(conn, definition.key) 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) @@ -736,6 +738,7 @@ def begin_oauth(conn, user_id: str, provider: str, connection_id: Optional[str] if row is None or row.get("provider") != provider: raise ConnectionUnavailable("Connection not found", connection_id=connection_id, status="missing") return row + ensure_connector_allowed(conn, provider) definition = catalog.get_definition(provider) result = conn.execute( text( @@ -1168,3 +1171,51 @@ def reencrypt_all(batch_size: int = 500) -> dict: for connection_id in failed_ids: mark_reconnect_needed(connection_id, DECRYPT_ERROR) return counts + + +# --------------------------------------------------------------------------- +# Admin policy +# --------------------------------------------------------------------------- + + +class ConnectorDisabled(Exception): + """An admin turned this connector (or custom MCP servers) off.""" + + +def custom_mcp_allowed(conn) -> bool: + """The instance-wide "Allow custom MCP servers" switch (on unless turned off).""" + from docsgpt.storage.db.repositories.app_metadata import AppMetadataRepository + from docsgpt.storage.db.repositories.connector_policies import ALLOW_CUSTOM_MCP_KEY + + return AppMetadataRepository(conn).get(ALLOW_CUSTOM_MCP_KEY) != "false" + + +def load_policies(conn) -> dict[str, dict]: + """Every connector's policy, with custom MCP folded in from its switch.""" + from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository + + policies = dict(ConnectorPoliciesRepository(conn).all()) + if not custom_mcp_allowed(conn): + policies["custom_mcp"] = {**policies.get("custom_mcp", {}), "enabled": False} + return policies + + +def ensure_connector_allowed(conn, connector_key: Optional[str]) -> None: + """Refuse a new connection to a connector an admin turned off. + + Raises: + ConnectorDisabled: The connector, or custom MCP servers, are off. + """ + if not connector_key: + return + policy = load_policies(conn).get(connector_key) or {} + if policy.get("enabled") is False: + raise ConnectorDisabled(f"{connector_key} is turned off by an admin") + + +def forced_credential_mode(conn, connector_key: Optional[str]) -> Optional[str]: + """``owner`` or ``member`` when an admin forces one for this connector.""" + if not connector_key: + return None + mode = (load_policies(conn).get(connector_key) or {}).get("credential_mode") + return mode if mode in ("owner", "member") else None diff --git a/docsgpt/retriever/base.py b/docsgpt/retriever/base.py index a1b49d9c..897ab622 100644 --- a/docsgpt/retriever/base.py +++ b/docsgpt/retriever/base.py @@ -2,6 +2,17 @@ from abc import ABC, abstractmethod class BaseRetriever(ABC): + @property + def _connector_labels(self): + """Per-retriever memo of each source's connector ("From Google Drive").""" + cache = self.__dict__.get("_connector_label_cache") + if cache is None: + from docsgpt.connectors.attribution import ConnectorLabelCache + + cache = ConnectorLabelCache() + self.__dict__["_connector_label_cache"] = cache + return cache + def __init__(self): pass diff --git a/docsgpt/retriever/classic_rag.py b/docsgpt/retriever/classic_rag.py index 76f343c3..3e8779be 100644 --- a/docsgpt/retriever/classic_rag.py +++ b/docsgpt/retriever/classic_rag.py @@ -385,7 +385,11 @@ class ClassicRAG(BaseRetriever): doc_tokens = num_tokens_from_string(doc_text_with_header) if cumulative_tokens + doc_tokens < token_budget: - entry = {"text": page_content, **labels} + entry = { + "text": page_content, + **labels, + **self._connector_labels.for_source(vectorstore_id), + } if self.include_scores: entry["score"] = score entry["score_kind"] = score_kind diff --git a/docsgpt/retriever/graph_rag.py b/docsgpt/retriever/graph_rag.py index 4b12e1b3..48fa0fd5 100644 --- a/docsgpt/retriever/graph_rag.py +++ b/docsgpt/retriever/graph_rag.py @@ -584,7 +584,7 @@ class GraphRAGRetriever(BaseRetriever): doc_tokens = num_tokens_from_string(f"{labels['filename']}\n{text}") if cumulative_tokens + doc_tokens >= token_budget: break - docs.append({"text": text, **labels}) + docs.append({"text": text, **labels, **self._connector_labels.for_source(source_id)}) cumulative_tokens += doc_tokens return docs diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index 4e8d052a..59568d22 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -611,6 +611,17 @@ connector_sessions_table = Table( ) +connector_policies_table = Table( + "connector_policies", + metadata, + Column("connector_key", Text, primary_key=True), + Column("enabled", Boolean, nullable=False, server_default="true"), + Column("credential_mode", Text, nullable=False, server_default="choose"), + Column("updated_by", Text), + Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), +) + + # --- Conversations, messages, workflows ------------------------------------- conversations_table = Table( diff --git a/docsgpt/storage/db/repositories/connector_policies.py b/docsgpt/storage/db/repositories/connector_policies.py new file mode 100644 index 00000000..2373b8cd --- /dev/null +++ b/docsgpt/storage/db/repositories/connector_policies.py @@ -0,0 +1,58 @@ +"""Repository for ``connector_policies``: the admin's per-connector switches.""" + +from __future__ import annotations + +from typing import Optional + +from sqlalchemy import Connection, text + +from docsgpt.storage.db.base_repository import row_to_dict + +CREDENTIAL_POLICIES = ("choose", "owner", "member") +ALLOW_CUSTOM_MCP_KEY = "connectors.allow_custom_mcp" + + +class ConnectorPoliciesRepository: + """Whether a connector is enabled and which credential mode it forces.""" + + def __init__(self, conn: Connection) -> None: + self._conn = conn + + def all(self) -> dict[str, dict]: + """Every stored policy, by connector key. Missing keys use the defaults.""" + result = self._conn.execute(text("SELECT * FROM connector_policies")) + return {row["connector_key"]: row for row in (row_to_dict(r) for r in result.fetchall())} + + def get(self, connector_key: str) -> Optional[dict]: + row = self._conn.execute( + text("SELECT * FROM connector_policies WHERE connector_key = :key"), {"key": connector_key}, + ).fetchone() + return row_to_dict(row) if row is not None else None + + def upsert( + self, + connector_key: str, + *, + enabled: Optional[bool] = None, + credential_mode: Optional[str] = None, + updated_by: Optional[str] = None, + ) -> dict: + """Set one connector's policy; fields left None keep their value.""" + if credential_mode is not None and credential_mode not in CREDENTIAL_POLICIES: + raise ValueError(f"unknown credential mode: {credential_mode!r}") + row = self._conn.execute( + text( + """ + INSERT INTO connector_policies (connector_key, enabled, credential_mode, updated_by) + VALUES (:key, COALESCE(:enabled, true), COALESCE(:mode, 'choose'), :by) + ON CONFLICT (connector_key) DO UPDATE SET + enabled = COALESCE(:enabled, connector_policies.enabled), + credential_mode = COALESCE(:mode, connector_policies.credential_mode), + updated_by = :by, + updated_at = now() + RETURNING * + """ + ), + {"key": connector_key, "enabled": enabled, "mode": credential_mode, "by": updated_by}, + ).fetchone() + return row_to_dict(row) diff --git a/tests/connectors/test_policies_and_attribution.py b/tests/connectors/test_policies_and_attribution.py new file mode 100644 index 00000000..a3a5f1a7 --- /dev/null +++ b/tests/connectors/test_policies_and_attribution.py @@ -0,0 +1,315 @@ +"""Admin connector policies, sharing modes and connector attribution.""" + +from __future__ import annotations + +import json +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from flask import Flask +from sqlalchemy import text + +from docsgpt.connectors import catalog, service +from docsgpt.security.encryption import encrypt_json + + +@pytest.fixture +def app(): + return Flask(__name__) + + +@contextmanager +def _db(conn): + @contextmanager + def _yield(): + yield conn + + modules = ( + "docsgpt.connectors.service", + "docsgpt.connectors.resolve", + "docsgpt.connectors.attribution", + "docsgpt.api.connector.connections", + "docsgpt.api.admin.connectors", + "docsgpt.api.user.tools.mcp", + ) + patches = [] + for module in modules: + for name in ("db_session", "db_readonly"): + target = f"{module}.{name}" + try: + patches.append(patch(target, _yield)) + patches[-1].start() + except AttributeError: + patches.pop() + try: + yield + finally: + for p in patches: + p.stop() + + +def _connection(conn, user="alice", provider="telegram", status="connected", auth_kind="api_key", + server_url=None, secrets=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}@example.com", "url": server_url, + "e": encrypt_json(secrets or {"credentials": {"token": "tok"}}, user)}, + ).scalar()) + + +def _call(app, resource, method, path, *, user="alice", body=None, roles=None, args=()): + with app.test_request_context(path, method=method.upper(), json=body): + from flask import request + + request.decoded_token = {"sub": user, "roles": roles or ["user"]} if user else None + return getattr(resource(), method)(*args) + + +class TestPresets: + def test_presets_are_in_the_catalog(self): + presets = [d for d in catalog.all_definitions() if d.publisher == "preset"] + assert {p.key for p in presets} >= {"mcp:notion", "mcp:linear"} + for preset in presets: + assert preset.mcp_url.startswith("https://") + assert preset.auth_kind in ("mcp_oauth", "none") + assert preset.tool_templates == ("mcp_tool",) + + def test_existing_mcp_connection_maps_to_its_preset(self): + row = {"provider": "mcp:https://mcp.linear.app", "server_url": "https://mcp.linear.app", + "connector_key": "custom_mcp"} + assert catalog.connector_key_for_row(row) == "mcp:linear" + + +class TestAdminPolicies: + def test_requires_admin(self, app, pg_conn): + from docsgpt.api.admin.connectors import AdminConnectorsResource + + with _db(pg_conn): + resp = _call(app, AdminConnectorsResource, "get", "/api/admin/connectors") + assert resp.status_code == 403 + + def test_lists_setup_state_without_values(self, app, pg_conn, monkeypatch): + from docsgpt.api.admin.connectors import AdminConnectorsResource + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", "secret-client-id") + monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", None) + _connection(pg_conn, provider="google_drive", auth_kind="oauth") + with _db(pg_conn): + resp = _call(app, AdminConnectorsResource, "get", "/api/admin/connectors", roles=["admin"]) + payload = resp.get_json() + drive = next(c for c in payload["connectors"] if c["key"] == "google_drive") + assert drive["required_settings"] == [ + {"name": "GOOGLE_CLIENT_ID", "set": True}, {"name": "GOOGLE_CLIENT_SECRET", "set": False}, + ] + assert drive["connection_count"] == 1 + assert "secret-client-id" not in json.dumps(payload) + assert payload["allow_custom_mcp"] is True + + def test_disable_connector_and_custom_mcp(self, app, pg_conn): + from docsgpt.api.admin.connectors import AdminConnectorsResource + + with _db(pg_conn): + resp = _call(app, AdminConnectorsResource, "put", "/api/admin/connectors", roles=["admin"], body={ + "policies": {"telegram": {"enabled": False, "credential_mode": "member"}}, + "allow_custom_mcp": False, + }) + assert resp.status_code == 200 + entries = {e["key"]: e for e in service.catalog_for_user(pg_conn, "bob", is_admin=False)} + assert entries["telegram"]["state"] == "disabled" + assert entries["custom_mcp"]["state"] == "disabled" + with pytest.raises(service.ConnectorDisabled): + service.create_api_key_connection( + pg_conn, "bob", catalog.get_definition("telegram"), {"token": "long-enough-token"}, + ) + + def test_rejects_unknown_connector(self, app, pg_conn): + from docsgpt.api.admin.connectors import AdminConnectorsResource + + with _db(pg_conn): + resp = _call(app, AdminConnectorsResource, "put", "/api/admin/connectors", roles=["admin"], + body={"policies": {"nope": {"enabled": False}}}) + assert resp.status_code == 400 + + def test_custom_mcp_switch_is_enforced_server_side(self, app, pg_conn): + from docsgpt.api.user.tools.mcp import TestMCPServerConfig + from docsgpt.storage.db.repositories.app_metadata import AppMetadataRepository + + AppMetadataRepository(pg_conn).set("connectors.allow_custom_mcp", "false") + with _db(pg_conn), patch("docsgpt.api.user.tools.mcp._validate_mcp_server_url"): + custom = _call(app, TestMCPServerConfig, "post", "/api/mcp_server/test", + body={"config": {"server_url": "https://tools.example.com/mcp"}}) + assert custom.status_code == 403 + + def test_disabled_connector_stops_existing_tools(self, pg_conn): + from docsgpt.connectors.resolve import resolve_connection + from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository + + cid = _connection(pg_conn) + ConnectorPoliciesRepository(pg_conn).upsert("telegram", enabled=False) + with _db(pg_conn): + resolved = resolve_connection({"connection_id": cid, "user_id": "alice"}, "alice") + assert resolved.available is False + + def test_forced_member_mode_overrides_the_share(self, pg_conn): + from docsgpt.connectors.resolve import resolve_connection + from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository + + owner = _connection(pg_conn) + bobs = _connection(pg_conn, user="bob") + ConnectorPoliciesRepository(pg_conn).upsert("telegram", credential_mode="member") + with _db(pg_conn): + resolved = resolve_connection( + {"connection_id": owner, "user_id": "alice", "credential_mode": "owner"}, "bob", + ) + assert resolved.connection_id == bobs + + +class TestCredentialMode: + def _tool(self, conn, cid, user="alice"): + return str(conn.execute( + text("INSERT INTO user_tools (user_id, name, connection_id) VALUES (:u, 'telegram', CAST(:c AS uuid)) " + "RETURNING id"), + {"u": user, "c": cid}, + ).scalar()) + + def test_owner_sets_member_mode(self, app, pg_conn): + from docsgpt.api.connector.connections import ToolCredentialMode + + tool = self._tool(pg_conn, _connection(pg_conn)) + with _db(pg_conn): + resp = _call(app, ToolCredentialMode, "put", "/x", body={"mode": "member"}, args=[tool]) + assert resp.status_code == 200 + mode = pg_conn.execute(text("SELECT credential_mode FROM user_tools WHERE id = CAST(:i AS uuid)"), + {"i": tool}).scalar() + assert mode == "member" + + def test_other_users_cannot(self, app, pg_conn): + from docsgpt.api.connector.connections import ToolCredentialMode + + tool = self._tool(pg_conn, _connection(pg_conn)) + with _db(pg_conn): + resp = _call(app, ToolCredentialMode, "put", "/x", user="bob", body={"mode": "member"}, args=[tool]) + assert resp.status_code == 404 + + def test_forced_policy_wins(self, app, pg_conn): + from docsgpt.api.connector.connections import ToolCredentialMode + from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository + + tool = self._tool(pg_conn, _connection(pg_conn)) + ConnectorPoliciesRepository(pg_conn).upsert("telegram", credential_mode="owner") + with _db(pg_conn): + resp = _call(app, ToolCredentialMode, "put", "/x", body={"mode": "member"}, args=[tool]) + assert resp.status_code == 409 + assert resp.get_json()["mode"] == "owner" + + +def _executor(user): + from docsgpt.agents.tool_executor import ToolExecutor + + return ToolExecutor(user=user) + + +def _pause(executor, tool, action_name): + with patch("docsgpt.agents.tool_executor.ToolActionParser") as parser: + parser.return_value.parse_args.return_value = ("t1", action_name, {}) + return executor.check_pause( + {"t1": tool}, SimpleNamespace(id="c1", name=action_name, thought_signature=None), "OpenAILLM", + ) + + +class TestSharedWrites: + def _tool(self, cid): + return { + "id": "tool-1", + "user_id": "alice", + "name": "postgres", + "config": {}, + "connection_id": cid, + "credential_mode": "owner", + "actions": [ + {"name": "postgres_execute_sql", "access": "write", "require_approval": False, "active": True}, + {"name": "postgres_get_schema", "access": "read", "require_approval": False, "active": True}, + ], + } + + def test_member_on_owners_account_must_approve_writes(self, pg_conn): + cid = _connection(pg_conn, provider="postgres") + with _db(pg_conn): + pause = _pause(_executor("bob"), self._tool(cid), "postgres_execute_sql") + assert pause["pause_type"] == "awaiting_approval" + + def test_member_reads_are_not_gated(self, pg_conn): + cid = _connection(pg_conn, provider="postgres") + with _db(pg_conn): + assert _pause(_executor("bob"), self._tool(cid), "postgres_get_schema") is None + + def test_owner_keeps_their_own_choice(self, pg_conn): + cid = _connection(pg_conn, provider="postgres") + with _db(pg_conn): + assert _pause(_executor("alice"), self._tool(cid), "postgres_execute_sql") is None + + +class TestAttribution: + def test_tool_call_events_name_the_connector_not_the_account(self, pg_conn): + cid = _connection(pg_conn, provider="mcp:https://mcp.notion.com", auth_kind="mcp_oauth", + server_url="https://mcp.notion.com", secrets={"tokens": {"access_token": "t"}}) + # MCP rows are stored as custom_mcp and named after their preset. + pg_conn.execute(text("UPDATE connector_sessions SET connector_key = 'custom_mcp' WHERE id = CAST(:i AS uuid)"), + {"i": cid}) + executor = _executor("alice") + tool = { + "id": "tool-1", "user_id": "alice", "name": "mcp_tool", "connection_id": cid, + "config": {"server_url": "https://mcp.notion.com/mcp", "auth_type": "oauth"}, + "actions": [{"name": "search", "active": True, "parameters": {"properties": {}}, + "annotations": {"readOnlyHint": True}}], + } + fake_tool = SimpleNamespace(execute_action=lambda *a, **k: "ok", config={}) + call = SimpleNamespace(id="c1", name="search", arguments="{}", thought_signature=None) + with _db(pg_conn), patch("docsgpt.agents.tool_executor.ToolActionParser") as parser, patch.object( + type(executor), "_get_or_load_tool", return_value=fake_tool + ), patch("docsgpt.agents.tool_executor._record_proposed", return_value=True), patch( + "docsgpt.agents.tool_executor._mark_executed" + ): + parser.return_value.parse_args.return_value = ("t1", "search", {}) + events = list(executor._execute({"t1": tool}, call, "OpenAILLM")) + final = events[-1]["data"] + assert final["connector_key"] == "mcp:notion" + assert final["connector_name"] == "Notion" + assert final["access"] == "read" + assert "alice@example.com" not in json.dumps(events) + + def test_built_in_tools_carry_no_connector(self): + executor = _executor("alice") + assert executor._resolve_connection({"id": "x", "name": "memory"}) is None + + def test_citations_name_the_connector(self, pg_conn): + from docsgpt.connectors.attribution import connector_labels + + cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth") + source = str(pg_conn.execute(text( + "INSERT INTO sources (user_id, name, connection_id) VALUES ('alice', 'Handbook', CAST(:c AS uuid)) " + "RETURNING id" + ), {"c": cid}).scalar()) + upload = str(pg_conn.execute(text( + "INSERT INTO sources (user_id, name) VALUES ('alice', 'Upload') RETURNING id" + )).scalar()) + with _db(pg_conn): + labels = connector_labels([source, upload, "not-a-uuid"]) + assert labels == {source: {"connector_key": "google_drive", "connector_name": "Google Drive"}} + + def test_classic_retriever_stamps_connector(self, pg_conn): + from docsgpt.retriever.classic_rag import ClassicRAG + + retriever = ClassicRAG.__new__(ClassicRAG) + with patch("docsgpt.connectors.attribution.connector_labels", + return_value={"s1": {"connector_key": "google_drive", "connector_name": "Google Drive"}}) as lookup: + first = retriever._connector_labels.for_source("s1") + again = retriever._connector_labels.for_source("s1") + assert first == again == {"connector_key": "google_drive", "connector_name": "Google Drive"} + assert lookup.call_count == 1 diff --git a/tests/storage/db/test_migration_0038.py b/tests/storage/db/test_migration_0038.py index 2740f33a..5b84179b 100644 --- a/tests/storage/db/test_migration_0038.py +++ b/tests/storage/db/test_migration_0038.py @@ -95,6 +95,7 @@ class TestMigration0038: ) assert "connection_id" in _columns(conn, "sources") assert "connection_id" in _columns(conn, "user_tools") + assert {"connector_key", "enabled", "credential_mode"} <= _columns(conn, "connector_policies") def test_backfill_links_legacy_rows(self, pg_engine): url = pg_engine.url.render_as_string(hide_password=False)