mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-06 20:15:17 +02:00
refactor(agents): group MCP tools into shared/tools/mcp/ subpackage
The three MCP siblings (mcp_client/mcp_tool/mcp_tools_cache) served one objective but sat loose at the top of shared/tools. Grouped them into an mcp/ package and dropped the redundant prefix: client.py, tool.py, cache.py. Updated all importers (routes, mcp_tools subagent, e2e fake patch targets, unit test) to the new paths.
This commit is contained in:
parent
8d0090c6a1
commit
c51aca6ccc
9 changed files with 23 additions and 16 deletions
145
surfsense_backend/app/agents/shared/tools/mcp/cache.py
Normal file
145
surfsense_backend/app/agents/shared/tools/mcp/cache.py
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
"""Persist MCP ``list_tools`` results in ``SearchSourceConnector.config.cached_tools``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm.attributes import flag_modified
|
||||
|
||||
from app.db import SearchSourceConnector, async_session_maker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_pending_prefetch_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
|
||||
class CachedMCPToolDef(BaseModel):
|
||||
name: str
|
||||
description: str = ""
|
||||
input_schema: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class CachedMCPTools(BaseModel):
|
||||
discovered_at: datetime
|
||||
server_version: str | None = None
|
||||
server_name: str | None = None
|
||||
transport: str | None = None
|
||||
tools: list[CachedMCPToolDef]
|
||||
|
||||
|
||||
def read_cached_tools(connector: SearchSourceConnector) -> CachedMCPTools | None:
|
||||
"""Return parsed cached tools or ``None`` if missing / corrupt (caller falls back to live discovery)."""
|
||||
cfg = connector.config or {}
|
||||
raw = cfg.get("cached_tools")
|
||||
if not raw or not isinstance(raw, dict):
|
||||
return None
|
||||
|
||||
try:
|
||||
return CachedMCPTools.model_validate(raw)
|
||||
except ValidationError as exc:
|
||||
logger.warning(
|
||||
"MCP connector %d has corrupt cached_tools — falling back to live discovery: %s",
|
||||
connector.id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def write_cached_tools(
|
||||
connector_id: int,
|
||||
tool_definitions: list[dict[str, Any]],
|
||||
*,
|
||||
server_name: str | None = None,
|
||||
server_version: str | None = None,
|
||||
transport: str | None = None,
|
||||
) -> None:
|
||||
"""Best-effort persist; uses its own session so a write failure cannot poison the caller's transaction."""
|
||||
payload = CachedMCPTools(
|
||||
discovered_at=datetime.now(UTC),
|
||||
server_version=server_version,
|
||||
server_name=server_name,
|
||||
transport=transport,
|
||||
tools=[CachedMCPToolDef.model_validate(td) for td in tool_definitions],
|
||||
)
|
||||
|
||||
try:
|
||||
async with async_session_maker() as session:
|
||||
result = await session.execute(
|
||||
select(SearchSourceConnector).filter(
|
||||
SearchSourceConnector.id == connector_id,
|
||||
)
|
||||
)
|
||||
connector = result.scalars().first()
|
||||
if connector is None:
|
||||
return
|
||||
|
||||
cfg = dict(connector.config or {})
|
||||
cfg["cached_tools"] = payload.model_dump(mode="json")
|
||||
connector.config = cfg
|
||||
flag_modified(connector, "config")
|
||||
await session.commit()
|
||||
|
||||
logger.info(
|
||||
"Persisted cached_tools for MCP connector %d (%d tools)",
|
||||
connector_id,
|
||||
len(payload.tools),
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to persist cached_tools for MCP connector %d",
|
||||
connector_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
|
||||
def refresh_mcp_tools_cache_for_connector(
|
||||
connector_id: int,
|
||||
search_space_id: int,
|
||||
) -> None:
|
||||
"""Maintain the MCP tool cache after a single-connector lifecycle event.
|
||||
|
||||
Synchronously evicts the in-process LRU for the connector's search space
|
||||
(LRU keys are per-space, so eviction cannot be scoped finer), then schedules
|
||||
a background live discovery for this connector alone so its persisted
|
||||
``cached_tools`` row is refreshed before the next user query.
|
||||
|
||||
Idempotent. Eviction is best-effort; prefetch is best-effort and only runs
|
||||
when an event loop is available. Neither path raises.
|
||||
"""
|
||||
try:
|
||||
from app.agents.shared.tools.mcp.tool import invalidate_mcp_tools_cache
|
||||
|
||||
invalidate_mcp_tools_cache(search_space_id)
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"MCP in-process cache eviction skipped for space %d",
|
||||
search_space_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return
|
||||
|
||||
task = loop.create_task(_run_connector_prefetch(connector_id))
|
||||
_pending_prefetch_tasks.add(task)
|
||||
task.add_done_callback(_pending_prefetch_tasks.discard)
|
||||
|
||||
|
||||
async def _run_connector_prefetch(connector_id: int) -> None:
|
||||
from app.agents.shared.tools.mcp.tool import discover_single_mcp_connector
|
||||
|
||||
try:
|
||||
await discover_single_mcp_connector(connector_id)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"MCP background prefetch failed for connector_id=%d",
|
||||
connector_id,
|
||||
exc_info=True,
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue