diff --git a/surfsense_backend/app/routes/search_source_connectors_routes.py b/surfsense_backend/app/routes/search_source_connectors_routes.py index 9037d275a..e9ffb7050 100644 --- a/surfsense_backend/app/routes/search_source_connectors_routes.py +++ b/surfsense_backend/app/routes/search_source_connectors_routes.py @@ -3071,6 +3071,37 @@ class MCPTrustToolRequest(BaseModel): tool_name: str +async def _ensure_mcp_connector_for_user( + session: AsyncSession, *, user_id, connector_id: int +) -> int: + """Verify ``connector_id`` is an MCP-backed connector owned by ``user_id``. + + The trust-list feature is intentionally MCP-only; native connectors + (Gmail, Calendar, Notion, ...) do not have a "trust this tool" UI. + The JSONB ``has_key("server_config")`` filter is the same MCP marker + used elsewhere in this module. + + Returns the connector's ``search_space_id`` (needed downstream for + MCP tool cache invalidation). Raises ``HTTPException(404)`` when the + connector does not exist, is not owned by the user, or is not + MCP-backed. + """ + from sqlalchemy import cast + from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB + + result = await session.execute( + select(SearchSourceConnector.search_space_id).where( + SearchSourceConnector.id == connector_id, + SearchSourceConnector.user_id == user_id, + cast(SearchSourceConnector.config, PG_JSONB).has_key("server_config"), + ) + ) + search_space_id = result.scalar_one_or_none() + if search_space_id is None: + raise HTTPException(status_code=404, detail="MCP connector not found") + return search_space_id + + @router.post("/connectors/mcp/{connector_id}/trust-tool") async def trust_mcp_tool( connector_id: int, @@ -3080,45 +3111,32 @@ async def trust_mcp_tool( ): """Add a tool to the MCP connector's trusted (always-allow) list. - Once trusted, the tool executes without HITL approval on subsequent calls. - Works for both generic MCP_CONNECTOR and OAuth-backed MCP connectors - (LINEAR_CONNECTOR, JIRA_CONNECTOR, etc.) by checking for ``server_config``. + Once trusted, the tool executes without HITL approval on subsequent + calls. Works for both generic ``MCP_CONNECTOR`` and OAuth-backed MCP + connectors (``LINEAR_CONNECTOR``, ``JIRA_CONNECTOR``, ...) — the + storage primitive is the same JSON list under ``config.trusted_tools``. """ + from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache + from app.services.user_tool_allowlist import add_user_trust + try: - from sqlalchemy import cast - from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB - - result = await session.execute( - select(SearchSourceConnector).filter( - SearchSourceConnector.id == connector_id, - SearchSourceConnector.user_id == user.id, - cast(SearchSourceConnector.config, PG_JSONB).has_key("server_config"), - ) + search_space_id = await _ensure_mcp_connector_for_user( + session, user_id=user.id, connector_id=connector_id + ) + trusted = await add_user_trust( + session, + user_id=user.id, + connector_id=connector_id, + tool_name=body.tool_name, ) - connector = result.scalars().first() - if not connector: - raise HTTPException(status_code=404, detail="MCP connector not found") - - config = dict(connector.config or {}) - trusted: list[str] = list(config.get("trusted_tools", [])) - if body.tool_name not in trusted: - trusted.append(body.tool_name) - config["trusted_tools"] = trusted - connector.config = config - - from sqlalchemy.orm.attributes import flag_modified - - flag_modified(connector, "config") await session.commit() - - from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache - - invalidate_mcp_tools_cache(connector.search_space_id) - + invalidate_mcp_tools_cache(search_space_id) return {"status": "ok", "trusted_tools": trusted} except HTTPException: raise + except LookupError as e: + raise HTTPException(status_code=404, detail="MCP connector not found") from e except Exception as e: logger.error(f"Failed to trust MCP tool: {e!s}", exc_info=True) await session.rollback() @@ -3137,43 +3155,28 @@ async def untrust_mcp_tool( """Remove a tool from the MCP connector's trusted list. The tool will require HITL approval again on subsequent calls. - Works for both generic MCP_CONNECTOR and OAuth-backed MCP connectors. """ + from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache + from app.services.user_tool_allowlist import remove_user_trust + try: - from sqlalchemy import cast - from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB - - result = await session.execute( - select(SearchSourceConnector).filter( - SearchSourceConnector.id == connector_id, - SearchSourceConnector.user_id == user.id, - cast(SearchSourceConnector.config, PG_JSONB).has_key("server_config"), - ) + search_space_id = await _ensure_mcp_connector_for_user( + session, user_id=user.id, connector_id=connector_id + ) + trusted = await remove_user_trust( + session, + user_id=user.id, + connector_id=connector_id, + tool_name=body.tool_name, ) - connector = result.scalars().first() - if not connector: - raise HTTPException(status_code=404, detail="MCP connector not found") - - config = dict(connector.config or {}) - trusted: list[str] = list(config.get("trusted_tools", [])) - if body.tool_name in trusted: - trusted.remove(body.tool_name) - config["trusted_tools"] = trusted - connector.config = config - - from sqlalchemy.orm.attributes import flag_modified - - flag_modified(connector, "config") await session.commit() - - from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache - - invalidate_mcp_tools_cache(connector.search_space_id) - + invalidate_mcp_tools_cache(search_space_id) return {"status": "ok", "trusted_tools": trusted} except HTTPException: raise + except LookupError as e: + raise HTTPException(status_code=404, detail="MCP connector not found") from e except Exception as e: logger.error(f"Failed to untrust MCP tool: {e!s}", exc_info=True) await session.rollback() diff --git a/surfsense_backend/app/services/user_tool_allowlist.py b/surfsense_backend/app/services/user_tool_allowlist.py new file mode 100644 index 000000000..83b075fb7 --- /dev/null +++ b/surfsense_backend/app/services/user_tool_allowlist.py @@ -0,0 +1,196 @@ +"""User-scoped tool allow-list backed by ``SearchSourceConnector.config``. + +Stores the user's "always allow" preferences as a list of tool names under +``connector.config['trusted_tools']``. Storage is per +``(user_id, search_space_id, connector_id)`` — i.e. tied to a specific +connected account inside a specific workspace, exactly what the UI cares +about. + +Callers split into two roles: + +- **Writers** — the ``/connectors/.../trust-tool`` and ``/untrust-tool`` + HTTP routes, and the chat resume handler when it processes a + ``{type: "always"}`` decision. Both call + :func:`add_user_trust` / :func:`remove_user_trust`. The FE button is + the upstream UI trigger but it talks to the routes, never to this + module directly. +- **Reader** — the subagent compile path, which calls + :func:`fetch_user_allowlist_rulesets` and layers the result after the + subagent's coded ruleset. User ``allow`` rules then override coded + ``ask`` via the rule engine's last-match-wins evaluation. + +Coded ``deny`` rules are intentionally **not** overridable by this +allow-list — only ``ask`` can be promoted to ``allow``. The rule engine +enforces this naturally because user rules only ever emit ``allow``. +""" + +from __future__ import annotations + +import uuid +from collections import defaultdict + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm.attributes import flag_modified + +from app.agents.multi_agent_chat.constants import ( + CONNECTOR_TYPE_TO_CONNECTOR_AGENT_MAPS, +) +from app.agents.new_chat.permissions import Rule, Ruleset +from app.db import SearchSourceConnector + +_TRUSTED_TOOLS_KEY = "trusted_tools" + + +async def _load_owned_connector( + session: AsyncSession, + *, + user_id: uuid.UUID, + connector_id: int, +) -> SearchSourceConnector | None: + """Return a connector iff it belongs to ``user_id``, else ``None``. + + Ownership scoping is mandatory: the trust list mutates user-private + data, callers must never write across user boundaries. + """ + result = await session.execute( + select(SearchSourceConnector).where( + SearchSourceConnector.id == connector_id, + SearchSourceConnector.user_id == user_id, + ) + ) + return result.scalars().first() + + +def _read_trusted(connector: SearchSourceConnector) -> list[str]: + config = connector.config or {} + raw = config.get(_TRUSTED_TOOLS_KEY, []) + if not isinstance(raw, list): + return [] + return [str(item) for item in raw if isinstance(item, str)] + + +def _write_trusted(connector: SearchSourceConnector, trusted: list[str]) -> None: + config = dict(connector.config or {}) + config[_TRUSTED_TOOLS_KEY] = trusted + connector.config = config + flag_modified(connector, "config") + + +async def add_user_trust( + session: AsyncSession, + *, + user_id: uuid.UUID, + connector_id: int, + tool_name: str, +) -> list[str]: + """Append ``tool_name`` to the connector's trusted list (idempotent). + + Returns the updated trusted-tools list. Raises ``LookupError`` when + the connector does not exist or is not owned by ``user_id``. + """ + connector = await _load_owned_connector( + session, user_id=user_id, connector_id=connector_id + ) + if connector is None: + raise LookupError( + f"connector {connector_id} not found for user {user_id}" + ) + + trusted = _read_trusted(connector) + if tool_name not in trusted: + trusted.append(tool_name) + _write_trusted(connector, trusted) + await session.flush() + return trusted + + +async def remove_user_trust( + session: AsyncSession, + *, + user_id: uuid.UUID, + connector_id: int, + tool_name: str, +) -> list[str]: + """Remove ``tool_name`` from the connector's trusted list (idempotent). + + Returns the updated trusted-tools list. Raises ``LookupError`` when + the connector does not exist or is not owned by ``user_id``. + """ + connector = await _load_owned_connector( + session, user_id=user_id, connector_id=connector_id + ) + if connector is None: + raise LookupError( + f"connector {connector_id} not found for user {user_id}" + ) + + trusted = _read_trusted(connector) + if tool_name in trusted: + trusted = [t for t in trusted if t != tool_name] + _write_trusted(connector, trusted) + await session.flush() + return trusted + + +async def fetch_user_allowlist_rulesets( + session: AsyncSession, + *, + user_id: uuid.UUID, + search_space_id: int, +) -> dict[str, Ruleset]: + """Project the user's trusted-tool lists into per-subagent rulesets. + + Walks every connector the user owns in this workspace, maps each + ``connector_type`` to its consuming subagent via + :data:`CONNECTOR_TYPE_TO_CONNECTOR_AGENT_MAPS`, and emits one + ``Rule(permission=tool_name, pattern="*", action="allow")`` per + trusted entry. Rules from different connector accounts feeding the + same subagent (e.g. two Linear workspaces) are merged into one + ruleset; duplicates are harmless under last-match-wins. + + Connectors whose type is not mapped (search APIs, Github, etc.) and + connectors with empty trust lists contribute nothing. Subagents + with no trusted tools are absent from the returned dict — callers + should treat ``missing == empty``. + """ + result = await session.execute( + select( + SearchSourceConnector.id, + SearchSourceConnector.connector_type, + SearchSourceConnector.config, + ).where( + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.search_space_id == search_space_id, + ) + ) + + rules_by_subagent: dict[str, list[Rule]] = defaultdict(list) + for _connector_id, connector_type, config in result.all(): + subagent = CONNECTOR_TYPE_TO_CONNECTOR_AGENT_MAPS.get(str(connector_type)) + if subagent is None: + continue + + cfg = config or {} + raw = cfg.get(_TRUSTED_TOOLS_KEY, []) + if not isinstance(raw, list): + continue + + for tool in raw: + if not isinstance(tool, str) or not tool: + continue + rules_by_subagent[subagent].append( + Rule(permission=tool, pattern="*", action="allow") + ) + + return { + subagent: Ruleset(rules=rules, origin=f"user_allowlist:{subagent}") + for subagent, rules in rules_by_subagent.items() + } + + +__all__ = [ + "add_user_trust", + "fetch_user_allowlist_rulesets", + "remove_user_trust", +]