mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-19 18:45:15 +02:00
user_tool_allowlist: extract trust-tool storage into reusable service
This commit is contained in:
parent
31d6b43a42
commit
e99c06c887
2 changed files with 259 additions and 60 deletions
|
|
@ -3071,6 +3071,37 @@ class MCPTrustToolRequest(BaseModel):
|
||||||
tool_name: str
|
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")
|
@router.post("/connectors/mcp/{connector_id}/trust-tool")
|
||||||
async def trust_mcp_tool(
|
async def trust_mcp_tool(
|
||||||
connector_id: int,
|
connector_id: int,
|
||||||
|
|
@ -3080,45 +3111,32 @@ async def trust_mcp_tool(
|
||||||
):
|
):
|
||||||
"""Add a tool to the MCP connector's trusted (always-allow) list.
|
"""Add a tool to the MCP connector's trusted (always-allow) list.
|
||||||
|
|
||||||
Once trusted, the tool executes without HITL approval on subsequent calls.
|
Once trusted, the tool executes without HITL approval on subsequent
|
||||||
Works for both generic MCP_CONNECTOR and OAuth-backed MCP connectors
|
calls. Works for both generic ``MCP_CONNECTOR`` and OAuth-backed MCP
|
||||||
(LINEAR_CONNECTOR, JIRA_CONNECTOR, etc.) by checking for ``server_config``.
|
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:
|
try:
|
||||||
from sqlalchemy import cast
|
search_space_id = await _ensure_mcp_connector_for_user(
|
||||||
from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB
|
session, user_id=user.id, connector_id=connector_id
|
||||||
|
)
|
||||||
result = await session.execute(
|
trusted = await add_user_trust(
|
||||||
select(SearchSourceConnector).filter(
|
session,
|
||||||
SearchSourceConnector.id == connector_id,
|
user_id=user.id,
|
||||||
SearchSourceConnector.user_id == user.id,
|
connector_id=connector_id,
|
||||||
cast(SearchSourceConnector.config, PG_JSONB).has_key("server_config"),
|
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()
|
await session.commit()
|
||||||
|
invalidate_mcp_tools_cache(search_space_id)
|
||||||
from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache
|
|
||||||
|
|
||||||
invalidate_mcp_tools_cache(connector.search_space_id)
|
|
||||||
|
|
||||||
return {"status": "ok", "trusted_tools": trusted}
|
return {"status": "ok", "trusted_tools": trusted}
|
||||||
|
|
||||||
except HTTPException:
|
except HTTPException:
|
||||||
raise
|
raise
|
||||||
|
except LookupError as e:
|
||||||
|
raise HTTPException(status_code=404, detail="MCP connector not found") from e
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to trust MCP tool: {e!s}", exc_info=True)
|
logger.error(f"Failed to trust MCP tool: {e!s}", exc_info=True)
|
||||||
await session.rollback()
|
await session.rollback()
|
||||||
|
|
@ -3137,43 +3155,28 @@ async def untrust_mcp_tool(
|
||||||
"""Remove a tool from the MCP connector's trusted list.
|
"""Remove a tool from the MCP connector's trusted list.
|
||||||
|
|
||||||
The tool will require HITL approval again on subsequent calls.
|
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:
|
try:
|
||||||
from sqlalchemy import cast
|
search_space_id = await _ensure_mcp_connector_for_user(
|
||||||
from sqlalchemy.dialects.postgresql import JSONB as PG_JSONB
|
session, user_id=user.id, connector_id=connector_id
|
||||||
|
)
|
||||||
result = await session.execute(
|
trusted = await remove_user_trust(
|
||||||
select(SearchSourceConnector).filter(
|
session,
|
||||||
SearchSourceConnector.id == connector_id,
|
user_id=user.id,
|
||||||
SearchSourceConnector.user_id == user.id,
|
connector_id=connector_id,
|
||||||
cast(SearchSourceConnector.config, PG_JSONB).has_key("server_config"),
|
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()
|
await session.commit()
|
||||||
|
invalidate_mcp_tools_cache(search_space_id)
|
||||||
from app.agents.new_chat.tools.mcp_tool import invalidate_mcp_tools_cache
|
|
||||||
|
|
||||||
invalidate_mcp_tools_cache(connector.search_space_id)
|
|
||||||
|
|
||||||
return {"status": "ok", "trusted_tools": trusted}
|
return {"status": "ok", "trusted_tools": trusted}
|
||||||
|
|
||||||
except HTTPException:
|
except HTTPException:
|
||||||
raise
|
raise
|
||||||
|
except LookupError as e:
|
||||||
|
raise HTTPException(status_code=404, detail="MCP connector not found") from e
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to untrust MCP tool: {e!s}", exc_info=True)
|
logger.error(f"Failed to untrust MCP tool: {e!s}", exc_info=True)
|
||||||
await session.rollback()
|
await session.rollback()
|
||||||
|
|
|
||||||
196
surfsense_backend/app/services/user_tool_allowlist.py
Normal file
196
surfsense_backend/app/services/user_tool_allowlist.py
Normal file
|
|
@ -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",
|
||||||
|
]
|
||||||
Loading…
Add table
Add a link
Reference in a new issue