diff --git a/README.md b/README.md index 7e274d7..dba0286 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,10 @@ runtime. Clone it, `pip install -e .`, run it. - **Web search** — offered to the model as a tool it calls when a question needs it. DuckDuckGo out of the box (no account, no key), or point it at your own SearXNG, or Firecrawl. The sources stay in the transcript +- **Your own tools** — describe an HTTP call in the admin area (a schema, a URL + template, a secret) and a model can make it. Or add an **MCP server** by URL + and its tools appear beside the built-in ones. Both restrictable to groups, + and neither can be pointed at your own network unless you say so - **Speech in and out** — dictate a message and have replies read aloud, against any OpenAI-compatible audio endpoint (whisper.cpp, Speaches, Kokoro…). Each person picks their own voice @@ -79,8 +83,8 @@ runtime. Clone it, `pip install -e .`, run it. **Planned** -Custom tools and MCP servers · agentic execution (local and over SSH) · image -generation · OCR for scanned PDFs · semantic search in the library. +Agentic execution (local and over SSH) · image generation · OCR for scanned +PDFs · semantic search in the library. See [PLAN.md](PLAN.md) for what is built, what is not, and why. diff --git a/src/lembas/api/admin_tools.py b/src/lembas/api/admin_tools.py index 62af912..6a9f55f 100644 --- a/src/lembas/api/admin_tools.py +++ b/src/lembas/api/admin_tools.py @@ -31,11 +31,15 @@ from lembas.db.models import ( SECRET_PLACEMENTS, CustomTool, Group, + McpServer, ) from lembas.services import custom_tools from lembas.services import prompts as prompts_service from lembas.services import tools as tools_service from lembas.services.crypto import UNCHANGED_SENTINEL, decrypt, keep_or_replace, mask +from lembas.services.fetch import FetchError, check_url +from lembas.services.mcp import client as mcp_client +from lembas.services.mcp import registry as mcp_registry from lembas.web.templating import render log = logging.getLogger(__name__) @@ -413,3 +417,245 @@ async def update_tool(request: Request, db: Db, user: AdminUser, tool_id: str) - log.info("%s updated custom tool %s", user.email, tool.slug) return _back(f"Saved {tool.name}.") + + +# --- MCP servers ------------------------------------------------------------- +MCP_SLUG_PATTERN = re.compile(r"^[a-z0-9][a-z0-9_-]{0,23}$") + + +def _server(db: Db, server_id: str) -> McpServer: + server = db.get(McpServer, server_id) + if server is None: + raise HTTPException(status.HTTP_404_NOT_FOUND, "That server no longer exists.") + return server + + +def _mcp_back(message: str = "") -> Response: + target = f"/admin/mcp?saved={message}" if message else "/admin/mcp" + return RedirectResponse(target, status_code=status.HTTP_303_SEE_OTHER) + + +def _populate_server(server: McpServer, form) -> None: + server.name = str(form.get("name") or "").strip()[:120] + server.url = str(form.get("url") or "").strip()[:1000] + server.guidance = str(form.get("guidance") or "").replace("\r\n", "\n").strip() + server.headers_json = _parse_headers(str(form.get("headers") or "")) + + placement = str(form.get("secret_placement") or SECRET_NONE) + server.secret_placement = placement if placement in SECRET_PLACEMENTS else SECRET_NONE + server.secret_name = str(form.get("secret_name") or "Authorization").strip()[:120] + + server.timeout = _number( + form.get("timeout"), default=30, low=mcp_client.MIN_TIMEOUT, high=mcp_client.MAX_TIMEOUT + ) + server.max_chars = _number( + form.get("max_chars"), default=8000, low=mcp_client.MIN_CHARS, high=mcp_client.MAX_CHARS + ) + server.position = _number(form.get("position"), default=server.position or 0, low=0, high=999) + + server.allow_private = "allow_private" in form + server.enabled = "enabled" in form + server.public = "public" in form + + # One checkbox per advertised tool, so an unticked one is absent. The + # stored map holds only the refusals; absent means on. + if "tool_choices" in form: + offered = set(form.getlist("tool_names")) + chosen = set(form.getlist("tool_names_on")) + server.tool_overrides_json = dict.fromkeys(offered - chosen, False) + + +def _server_problem(db: Db, server: McpServer, form, *, existing_id: str = "") -> str: + if not server.name: + return "A server needs a name." + + slug = str(form.get("slug") or "").strip().lower() + if not MCP_SLUG_PATTERN.match(slug): + return ( + "The identifier must be lowercase letters, digits, hyphens or " + "underscores, and at most 24 characters. It prefixes every tool " + "name this server offers." + ) + clash = db.scalar(select(McpServer).where(McpServer.slug == slug)) + if clash is not None and clash.id != existing_id: + return f"There is already a server called “{slug}”." + server.slug = slug + + try: + check_url(server.url, allow_private=True) + except FetchError as exc: + return exc.message + return "" + + +def _server_detail( + request: Request, db: Db, server: McpServer, *, is_new: bool, error: str = "", **extra +): + key = f"tool.mcp_{server.slug}" if server.slug else "" + overrides = server.tool_overrides_json or {} + return render( + request, + "admin/mcp_detail.html", + { + "server": server, + "is_new": is_new, + "error": error, + "groups": list(db.scalars(select(Group).order_by(Group.name))), + "selected_groups": extra.pop( + "selected_groups", {group.id for group in (server.groups if server.id else [])} + ), + "headers_text": extra.pop("headers_text", _headers_text(server.headers_json)), + "tools": [ + {**entry, "on": overrides.get(entry.get("name"), True)} + for entry in (server.tools_json or []) + if isinstance(entry, dict) + ], + "masked": mask(decrypt(server.secret_encrypted)) if server.secret_encrypted else "", + "unchanged": UNCHANGED_SENTINEL, + "secret_placements": SECRET_LABELS, + "prompt_key": key, + "prompt_overridden": key in prompts_service.stored(db), + **extra, + }, + ) + + +@router.get("/admin/mcp") +async def mcp_page(request: Request, db: Db, user: AdminUser, saved: str = ""): + servers = list(db.scalars(select(McpServer).order_by(McpServer.position, McpServer.slug))) + return render( + request, + "admin/mcp.html", + { + "servers": servers, + "counts": {server.id: len(server.tools_json or []) for server in servers}, + "saved": saved, + }, + ) + + +# Registered before /{server_id}, for the reason given above. +@router.get("/admin/mcp/new") +async def new_server_page(request: Request, db: Db, user: AdminUser): + draft = McpServer( + name="", + slug="", + url="https://", + secret_placement=SECRET_NONE, + timeout=30, + max_chars=8000, + enabled=True, + public=True, + position=0, + tools_json=[], + tool_overrides_json={}, + ) + return _server_detail(request, db, draft, is_new=True) + + +@router.post("/admin/mcp") +async def create_server(request: Request, db: Db, user: AdminUser) -> Response: + form = await request.form() + draft = McpServer(headers_json={}, tools_json=[], tool_overrides_json={}) + _populate_server(draft, form) + draft.position = db.scalar(select(func.coalesce(func.max(McpServer.position), -1))) + 1 + + problem = _server_problem(db, draft, form) + if problem: + return _server_detail( + request, + db, + draft, + is_new=True, + error=problem, + headers_text=str(form.get("headers") or ""), + selected_groups=set(form.getlist("group_ids")), + ) + + draft.secret_encrypted = keep_or_replace(str(form.get("secret") or ""), "") + draft.groups = _chosen_groups(db, form, public=draft.public) + db.add(draft) + db.commit() + + # Discovered immediately, the way a new connection's models are: an + # administrator who has just typed a URL wants to know whether it answered. + count, error = await mcp_registry.refresh(db, draft) + log.info("%s added MCP server %s (%d tools)", user.email, draft.slug, count) + if error: + return _mcp_back(f"Added {draft.name}, but it could not be reached: {error}") + return _mcp_back(f"Added {draft.name} — {count} tool(s).") + + +@router.get("/admin/mcp/{server_id}/edit") +async def edit_server_page(request: Request, db: Db, user: AdminUser, server_id: str): + return _server_detail(request, db, _server(db, server_id), is_new=False) + + +@router.post("/admin/mcp/{server_id}/test") +async def test_server(request: Request, db: Db, user: AdminUser, server_id: str): + """Contact the server and cache what it advertises. + + Returns the row fragment, swapped in place, exactly as "Test & refresh" + does for a connection. + """ + server = _server(db, server_id) + count, error = await mcp_registry.refresh(db, server) + message = ( + f"{server.name}: {error}" + if error + else f"{server.name}: found {count} tool{'s' if count != 1 else ''}." + ) + return render( + request, + "admin/_mcp_row.html", + { + "server": server, + "tool_count": len(server.tools_json or []), + "message": message, + "message_kind": "error" if error else "success", + }, + ) + + +@router.post("/admin/mcp/{server_id}/delete") +async def delete_server(db: Db, user: AdminUser, server_id: str) -> Response: + server = _server(db, server_id) + name = server.name + db.delete(server) + db.commit() + log.info("%s deleted MCP server %s", user.email, name) + return _mcp_back(f"Deleted {name}.") + + +@router.post("/admin/mcp/{server_id}") +async def update_server(request: Request, db: Db, user: AdminUser, server_id: str) -> Response: + server = _server(db, server_id) + form = await request.form() + + draft = McpServer(headers_json={}, tools_json=[], position=server.position) + _populate_server(draft, form) + problem = _server_problem(db, draft, form, existing_id=server.id) + if problem: + draft.id = server.id + draft.secret_encrypted = server.secret_encrypted + draft.tools_json = server.tools_json + return _server_detail( + request, + db, + draft, + is_new=False, + error=problem, + headers_text=str(form.get("headers") or ""), + selected_groups=set(form.getlist("group_ids")), + ) + + _populate_server(server, form) + server.slug = draft.slug + server.secret_encrypted = keep_or_replace( + str(form.get("secret") or ""), server.secret_encrypted + ) + server.groups = _chosen_groups(db, form, public=server.public) + db.commit() + + log.info("%s updated MCP server %s", user.email, server.slug) + return _mcp_back(f"Saved {server.name}.") diff --git a/src/lembas/services/generation.py b/src/lembas/services/generation.py index 14bc816..e704738 100644 --- a/src/lembas/services/generation.py +++ b/src/lembas/services/generation.py @@ -338,13 +338,19 @@ async def _run(generation: Generation) -> None: *payload["messages"], tools_service.assistant_turn(calls, "".join(round_text)), ] - for call in calls: - outcome = await tools_service.run_tool( - tool_context, call["name"], call["arguments"] - ) - generation.tool_events.append(outcome.event) + + generation.status = _tool_status(calls) + generation.touch() + try: + outcomes = await _run_calls(tool_context, calls) + finally: + generation.status = "" generation.touch() + + for call, outcome in zip(calls, outcomes, strict=True): + generation.tool_events.append(outcome.event) messages.append(tools_service.tool_turn(call, outcome.content)) + generation.touch() payload = {**payload, "messages": messages} @@ -465,6 +471,49 @@ async def _maybe_compact(generation: Generation) -> None: generation.touch() +# How many of a round's tool calls may be in flight at once. A bound rather +# than none: a model that asks for eight would otherwise open eight sockets and +# eight database sessions at the same moment. +MAX_PARALLEL_TOOLS = 4 + + +def _tool_status(calls: list[dict]) -> str: + """What to show while tools run. + + A remote tool -- an HTTP endpoint, an MCP server -- can take seconds with + nothing streaming, and a silent pause is exactly what a hang looks like. + """ + if len(calls) == 1: + return f"Running {calls[0]['name']}…" + return f"Running {len(calls)} tools…" + + +async def _run_calls(context, calls: list[dict]) -> list: + """Run one round's calls together, results in call order. + + Sequential was right when every tool was a local database read. A remote one + is latency-bound, and three two-second calls in a row are six seconds of a + reply looking hung -- while the model has already been told it may ask for + several at once. + + The results are indexed rather than appended as they finish, because each + tool turn has to line up with the assistant turn's `tool_calls`: an endpoint + matching on `tool_call_id` would otherwise pair the right id with the wrong + content the moment two calls came back out of order. + + Safe to run together because `run_tool` never raises, so no failure cancels + its siblings, and each runner opens its own `session_scope()` against a + database in WAL mode with a busy timeout. + """ + limit = asyncio.Semaphore(MAX_PARALLEL_TOOLS) + + async def one(call: dict): + async with limit: + return await tools_service.run_tool(context, call["name"], call["arguments"]) + + return list(await asyncio.gather(*(one(call) for call in calls))) + + def _pending_text(db, message: Message) -> str: """The user turn this reply is answering, for the size estimate.""" previous = db.scalars( diff --git a/src/lembas/services/mcp/__init__.py b/src/lembas/services/mcp/__init__.py new file mode 100644 index 0000000..903bf3d --- /dev/null +++ b/src/lembas/services/mcp/__init__.py @@ -0,0 +1,26 @@ +"""An MCP client: remote servers over streamable HTTP. + +Three parts. `protocol` is the wire format and nothing else -- it knows no +database and no HTTP. `client` owns the transport, which is where the SSRF guard +lives and the reason none of this comes from the reference SDK. `registry` is +where a row becomes a tool the chat loop can be offered. + +Local stdio servers are deliberately absent. Spawning a subprocess is the +agentic-execution feature, which wants a confirmation model before it does +anything; a URL is a different act with a different blast radius. +""" + +from lembas.services.mcp.client import McpSpec, call_tool, list_tools, spec_from +from lembas.services.mcp.protocol import McpError +from lembas.services.mcp.registry import offer_name, refresh, tool_defs + +__all__ = [ + "McpError", + "McpSpec", + "call_tool", + "list_tools", + "offer_name", + "refresh", + "spec_from", + "tool_defs", +] diff --git a/src/lembas/services/mcp/client.py b/src/lembas/services/mcp/client.py new file mode 100644 index 0000000..8de311d --- /dev/null +++ b/src/lembas/services/mcp/client.py @@ -0,0 +1,273 @@ +"""Talking to a remote MCP server over streamable HTTP. + +One session per call, deliberately. A cached session would need an owner, a +lifetime, eviction, a lock -- a round runs its tools concurrently -- and a +shutdown hook, and the server may expire it underneath all of that anyway. +`ToolContext` is a session-free snapshot precisely so that nothing inside a tool +holds live state. The cost is one extra POST in front of a call that is already +a network round trip inside a reply taking seconds; the upgrade, if it is ever +worth it, is a dict in this module and invisible to everything else. + +Redirects are followed by hand and every hop is re-checked, for the reason +`services/fetch.py` gives: an administrator can point this at any URL, and a +name resolving to 127.0.0.1 walks past any check that only reads the text. +""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass, field +from typing import Any +from urllib.parse import quote, urlparse + +import httpx + +from lembas.db.models import SECRET_BEARER, SECRET_HEADER, SECRET_QUERY, McpServer +from lembas.services import fetch as fetch_service +from lembas.services.crypto import decrypt +from lembas.services.mcp import protocol +from lembas.services.mcp.protocol import McpError + +log = logging.getLogger(__name__) + +# A JSON-RPC message is small; anything this large is a server misbehaving, and +# reading it into memory before parsing is the failure to avoid. +MAX_MESSAGE_BYTES = 4 * 1024 * 1024 + +# tools/list is paged. Both bounds exist because the whole list goes into every +# request as schema. +MAX_PAGES = 10 +MAX_TOOLS = 100 + +MIN_TIMEOUT, MAX_TIMEOUT = 1, 120 +MIN_CHARS, MAX_CHARS = 200, 40_000 + + +@dataclass(frozen=True) +class McpSpec: + """One server, read while the session was open. See `custom_tools.HttpSpec`.""" + + slug: str + name: str + url: str + headers: dict[str, str] = field(default_factory=dict) + secret: str = "" + secret_placement: str = SECRET_BEARER + secret_name: str = "Authorization" + timeout: int = 30 + max_chars: int = 8000 + allow_private: bool = False + + +def spec_from(row: McpServer) -> McpSpec: + return McpSpec( + slug=row.slug, + name=row.name or row.slug, + url=row.url or "", + headers=dict(row.headers_json or {}), + secret=decrypt(row.secret_encrypted), + secret_placement=row.secret_placement, + secret_name=row.secret_name or "Authorization", + timeout=min(max(int(row.timeout or 0), MIN_TIMEOUT), MAX_TIMEOUT), + max_chars=min(max(int(row.max_chars or 0), MIN_CHARS), MAX_CHARS), + allow_private=bool(row.allow_private), + ) + + +def _headers(spec: McpSpec) -> dict[str, str]: + headers = { + "User-Agent": fetch_service.USER_AGENT, + "Content-Type": "application/json", + # Both, because a server may answer either for the same request. + "Accept": "application/json, text/event-stream", + **{str(k): str(v) for k, v in spec.headers.items()}, + } + if spec.secret: + if spec.secret_placement == SECRET_BEARER: + headers[spec.secret_name or "Authorization"] = f"Bearer {spec.secret}" + elif spec.secret_placement == SECRET_HEADER: + headers[spec.secret_name or "Authorization"] = spec.secret + return headers + + +def _endpoint(spec: McpSpec) -> str: + url = fetch_service.check_url(spec.url, allow_private=spec.allow_private) + if spec.secret and spec.secret_placement == SECRET_QUERY: + joiner = "&" if urlparse(url).query else "?" + url = f"{url}{joiner}{quote(spec.secret_name)}={quote(spec.secret, safe='')}" + return url + + +class Session: + """One initialised conversation with a server.""" + + def __init__(self, spec: McpSpec, client: httpx.AsyncClient) -> None: + self.spec = spec + self._client = client + self._url = "" + self._headers = _headers(spec) + self._session_id = "" + self._next_id = 0 + self.protocol_version = "" + self.server_info: dict[str, Any] = {} + + # --- Transport ----------------------------------------------------------- + async def _post(self, message: dict[str, Any]) -> httpx.Response: + current = self._url + headers = dict(self._headers) + if self._session_id: + headers["Mcp-Session-Id"] = self._session_id + if self.protocol_version: + headers["MCP-Protocol-Version"] = self.protocol_version + + origin = (urlparse(current).scheme, urlparse(current).netloc) + for _ in range(fetch_service.MAX_REDIRECTS + 1): + try: + response = await self._client.post(current, json=message, headers=headers) + except httpx.RequestError as exc: + raise McpError(f"Could not reach {self.spec.name}: {exc}") from exc + + if not response.is_redirect: + return response + + location = response.headers.get("location", "") + if not location: + raise McpError(f"{self.spec.name} redirected to nowhere.") + if response.status_code not in (307, 308): + # 301, 302 and 303 turn a POST into a GET, which means nothing + # to a JSON-RPC endpoint. Refused rather than guessed at. + raise McpError( + f"{self.spec.name} answered {response.status_code}, which would " + "turn the request into a GET. Point the URL at the endpoint itself." + ) + + current = fetch_service.check_url( + str(response.url.join(location)), allow_private=self.spec.allow_private + ) + if (urlparse(current).scheme, urlparse(current).netloc) != origin: + headers.pop(self.spec.secret_name or "Authorization", None) + origin = (urlparse(current).scheme, urlparse(current).netloc) + + raise McpError(f"{self.spec.name} redirected too many times.") + + def _read(self, response: httpx.Response, *, request_id: int) -> dict[str, Any]: + if response.status_code >= 400: + raise McpError(f"{self.spec.name} returned HTTP {response.status_code}.") + return protocol.result_of( + response.content[:MAX_MESSAGE_BYTES], + response.headers.get("content-type", ""), + request_id=request_id, + ) + + # --- Lifecycle ----------------------------------------------------------- + async def open(self) -> None: + self._url = _endpoint(self.spec) + self._next_id += 1 + request_id = self._next_id + + response = await self._post( + protocol.request( + "initialize", + { + "protocolVersion": protocol.PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": protocol.CLIENT_INFO, + }, + request_id=request_id, + ) + ) + # Captured before the body is read: a server that issues one expects it + # on everything after this, including the initialized notification. + self._session_id = response.headers.get("mcp-session-id", "") + result = self._read(response, request_id=request_id) + + self.protocol_version = str(result.get("protocolVersion") or protocol.PROTOCOL_VERSION) + info = result.get("serverInfo") + self.server_info = info if isinstance(info, dict) else {} + if self.protocol_version != protocol.PROTOCOL_VERSION: + log.info( + "%s speaks MCP %s, we asked for %s", + self.spec.name, + self.protocol_version, + protocol.PROTOCOL_VERSION, + ) + + await self._post(protocol.notification("notifications/initialized")) + + async def call(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + """One request, re-initialising once if the session has expired.""" + self._next_id += 1 + request_id = self._next_id + response = await self._post(protocol.request(method, params, request_id=request_id)) + + if response.status_code == 404 and self._session_id: + # The server dropped the session. One retry, then it is a failure + # like any other -- a loop here would be a loop against a server + # that has decided to forget us. + log.info("%s expired its session; re-initialising", self.spec.name) + self._session_id = "" + await self.open() + self._next_id += 1 + request_id = self._next_id + response = await self._post(protocol.request(method, params, request_id=request_id)) + + return self._read(response, request_id=request_id) + + async def close(self) -> None: + """Best effort. A server with nothing to clean up answers 405.""" + if not self._session_id: + return + headers = {**self._headers, "Mcp-Session-Id": self._session_id} + if self.protocol_version: + headers["MCP-Protocol-Version"] = self.protocol_version + with suppress(httpx.RequestError, McpError): + await self._client.delete(self._url, headers=headers) + + +@asynccontextmanager +async def session_for(spec: McpSpec) -> AsyncIterator[Session]: + """An initialised session, closed afterwards whatever happened.""" + client = httpx.AsyncClient(timeout=spec.timeout, follow_redirects=False) + session = Session(spec, client) + try: + await session.open() + yield session + finally: + await session.close() + await client.aclose() + + +async def list_tools(spec: McpSpec) -> tuple[list[dict[str, Any]], dict[str, Any], str]: + """Every tool a server advertises, plus what it said about itself. + + Returns (tools, serverInfo, protocolVersion). Bounded at MAX_TOOLS: the list + is sent as schema on every request, so a server offering two hundred is a + server that would fill the window before anything was asked. + """ + tools: list[dict[str, Any]] = [] + async with session_for(spec) as session: + cursor = "" + for _ in range(MAX_PAGES): + result = await session.call("tools/list", {"cursor": cursor} if cursor else {}) + for entry in result.get("tools") or []: + cleaned = protocol.clean_tool(entry) + if cleaned is None: + log.info("%s advertised an unusable tool entry", spec.name) + elif len(tools) < MAX_TOOLS: + tools.append(cleaned) + cursor = str(result.get("nextCursor") or "") + if not cursor or len(tools) >= MAX_TOOLS: + break + return tools, session.server_info, session.protocol_version + + +async def call_tool(spec: McpSpec, name: str, arguments: dict[str, Any]) -> tuple[str, bool]: + """Run one tool. Returns (text, is_error).""" + async with session_for(spec) as session: + result = await session.call("tools/call", {"name": name, "arguments": arguments}) + return protocol.content_to_text(result), bool(result.get("isError")) + + +__all__ = ["McpError", "McpSpec", "Session", "call_tool", "list_tools", "spec_from"] diff --git a/src/lembas/services/mcp/protocol.py b/src/lembas/services/mcp/protocol.py new file mode 100644 index 0000000..ecc4f4c --- /dev/null +++ b/src/lembas/services/mcp/protocol.py @@ -0,0 +1,195 @@ +"""The MCP wire format: JSON-RPC 2.0, and what comes back from a tool call. + +Written out rather than taken from the reference SDK. The client is a few +hundred lines of framing, and the SDK's transport does its own connecting -- +which would mean the one thing that must not be bypassed, `fetch.check_url` on +every hop, being bypassed. Owning the transport is the point; owning the framing +beside it is the small part. + +A response arrives either as one JSON object or as an event stream carrying +several messages, and a server may choose either for the same request. Both are +read here so `client.py` does not have to care which it got. +""" + +from __future__ import annotations + +import json +from contextlib import suppress +from typing import Any + +from lembas import __version__ + +# What we tell a server we speak. A server answering an older version is not +# refused: several in the wild still answer 2024-11-05 and work perfectly. +PROTOCOL_VERSION = "2025-06-18" + +CLIENT_INFO = {"name": "LLeMbas", "version": __version__} + +# A tool's metadata is sent to the model as instructions, so it is bounded here +# rather than trusted. A server advertising a 40 KB description would spend the +# context window before the conversation started. +MAX_DESCRIPTION = 1000 +MAX_SCHEMA_BYTES = 8192 +MAX_NAME = 64 + +# Content that is not text is described rather than forwarded. A tool turn is a +# string, images only reach models marked as having vision, and base64 in a tool +# result is the fastest way to fill a window with nothing. +UNSUPPORTED = "[{kind}: {detail} — not shown to the model]" + + +class McpError(Exception): + """A failed exchange, with a message fit to show an administrator. + + Same contract as `LLMError`, `SearchError` and `FetchError`: the message is + the whole error, and is safe to render. + """ + + def __init__(self, message: str) -> None: + super().__init__(message) + self.message = message + + +def request(method: str, params: dict[str, Any] | None, *, request_id: int) -> dict[str, Any]: + message: dict[str, Any] = {"jsonrpc": "2.0", "id": request_id, "method": method} + if params is not None: + message["params"] = params + return message + + +def notification(method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + message: dict[str, Any] = {"jsonrpc": "2.0", "method": method} + if params is not None: + message["params"] = params + return message + + +def _messages(body: bytes, content_type: str) -> list[dict[str, Any]]: + """Every JSON-RPC message in a response body, whichever framing was used.""" + text = body.decode("utf-8", "replace").strip() + if not text: + return [] + + if "text/event-stream" not in content_type.lower(): + try: + document = json.loads(text) + except json.JSONDecodeError as exc: + raise McpError(f"That server did not answer with JSON: {exc}") from exc + return document if isinstance(document, list) else [document] + + # Frames are `data:` lines gathered until a blank line -- the same framing + # services/sse.py writes. Read here rather than shared, because that module + # is a writer and this is a reader with a size cap. + out: list[dict[str, Any]] = [] + data: list[str] = [] + for line in text.splitlines() + [""]: + if line.startswith("data:"): + data.append(line[5:].lstrip()) + elif not line.strip() and data: + # A frame that is not JSON is a comment or a keep-alive, not a + # message; the stream carries both. + with suppress(json.JSONDecodeError): + out.append(json.loads("\n".join(data))) + data = [] + return out + + +def result_of(body: bytes, content_type: str, *, request_id: int) -> dict[str, Any]: + """The result for one request, out of whatever the server sent back. + + Raises `McpError` on a JSON-RPC error member, because that is a failure the + administrator needs the words of -- "Unknown tool" and "invalid params" are + the two that actually happen. + """ + for message in _messages(body, content_type): + if not isinstance(message, dict) or message.get("id") != request_id: + continue + if "error" in message: + error = message["error"] or {} + code = error.get("code", "") + text = str(error.get("message") or "The server reported an error.") + raise McpError(f"{text}{f' (code {code})' if code != '' else ''}") + result = message.get("result") + return result if isinstance(result, dict) else {} + + raise McpError("That server did not answer the request.") + + +# --- Tool metadata ----------------------------------------------------------- +def clean_tool(entry: Any) -> dict[str, Any] | None: + """One advertised tool, bounded. None if there is nothing usable here.""" + if not isinstance(entry, dict): + return None + name = str(entry.get("name") or "").strip() + if not name or len(name) > MAX_NAME: + return None + + schema = entry.get("inputSchema") + if not isinstance(schema, dict) or schema.get("type") != "object": + schema = {"type": "object", "properties": {}} + elif len(json.dumps(schema)) > MAX_SCHEMA_BYTES: + # Kept callable rather than dropped: a model can still be told the tool + # exists, and an argument it guesses is no worse than not offering it. + schema = {"type": "object", "properties": {}} + + return { + "name": name, + "title": str(entry.get("title") or "")[:MAX_NAME], + "description": str(entry.get("description") or "")[:MAX_DESCRIPTION], + "schema": schema, + } + + +# --- Tool results ------------------------------------------------------------ +def _block_text(block: Any) -> str: + if not isinstance(block, dict): + return "" + kind = str(block.get("type") or "") + + if kind == "text": + return str(block.get("text") or "") + + if kind in ("image", "audio"): + size = len(str(block.get("data") or "")) + detail = f"{block.get('mimeType') or 'unknown type'}, about {size * 3 // 4} bytes" + return UNSUPPORTED.format(kind=kind, detail=detail) + + if kind == "resource": + resource = block.get("resource") + if not isinstance(resource, dict): + return "" + uri = str(resource.get("uri") or "") + if isinstance(resource.get("text"), str): + return f"{uri}\n{resource['text']}" if uri else str(resource["text"]) + detail = f"{uri or 'unnamed'}, {resource.get('mimeType') or 'unknown type'}" + return UNSUPPORTED.format(kind="resource", detail=detail) + + if kind == "resource_link": + return f"{block.get('name') or 'resource'} ({block.get('uri') or ''})".strip() + + # An addition to the protocol degrades to a note rather than to silence: + # a model told nothing came back will say nothing came back. + return UNSUPPORTED.format(kind="content", detail=kind or "no type") + + +def content_to_text(result: dict[str, Any]) -> str: + """A `tools/call` result as the flat text a tool turn carries.""" + blocks = result.get("content") + parts = [text for block in (blocks or []) if (text := _block_text(block).strip())] + + if not parts and isinstance(result.get("structuredContent"), dict | list): + return json.dumps(result["structuredContent"], indent=2, ensure_ascii=False) + + return "\n\n".join(parts) + + +__all__ = [ + "CLIENT_INFO", + "PROTOCOL_VERSION", + "McpError", + "clean_tool", + "content_to_text", + "notification", + "request", + "result_of", +] diff --git a/src/lembas/services/mcp/registry.py b/src/lembas/services/mcp/registry.py new file mode 100644 index 0000000..3c3f66a --- /dev/null +++ b/src/lembas/services/mcp/registry.py @@ -0,0 +1,185 @@ +"""Turning MCP servers into tools a chat can be offered. + +Two names per tool. The server has its own, which is what `tools/call` must be +given; we have an *offered* name, which is what goes in the schema the endpoint +sees. They differ because two servers both exposing `search` would collide, a +server exposing `notes_delete` would shadow a built-in, and endpoints accept a +narrower character set than MCP does. The rename never leaves this module: the +runner closes over the server's own name. + +The advertised list is cached on the row and refreshed by a button, the same +shape as discovering a connection's models. A server is contacted when an +administrator asks, not when a chat starts -- a slow server must not be able to +delay every reply. +""" + +from __future__ import annotations + +import logging +import re +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy.orm import Session as DBSession + +from lembas.db.models import McpServer, User +from lembas.services import tool_access +from lembas.services.fetch import FetchError +from lembas.services.mcp import client +from lembas.services.mcp.protocol import McpError +from lembas.services.tools import FAMILY_MCP, ToolContext, ToolDef, ToolOutcome + +log = logging.getLogger(__name__) + +# What an endpoint will accept as a function name. +FUNCTION_NAME = re.compile(r"^[a-zA-Z0-9_-]{1,64}$") +MAX_NAME = 64 + +# How much of a reply is kept on the message row for the transcript. +MAX_EVENT_CHARS = 2000 +MAX_SUMMARY_CHARS = 200 + + +def offer_name(server_slug: str, tool_name: str, *, taken: set[str]) -> str: + """A name the endpoint will accept, unique across everything offered. + + Truncation can collide where the full names would not, so a numeric suffix + is appended until it does not. Deterministic given a stable iteration order, + which is why rows are walked in (position, slug) order everywhere. + """ + combined = f"{server_slug}_{tool_name}".lower() + cleaned = re.sub(r"_+", "_", re.sub(r"[^a-z0-9_-]", "_", combined)).strip("_") + candidate = (cleaned or "tool")[:MAX_NAME] + + suffix = 2 + while candidate in taken: + tail = f"_{suffix}" + candidate = f"{(cleaned or 'tool')[: MAX_NAME - len(tail)]}{tail}" + suffix += 1 + + taken.add(candidate) + return candidate + + +def _summary(arguments: dict[str, Any]) -> str: + parts = [f"{name}={value!r}" for name, value in arguments.items()] + text = ", ".join(parts) + return text[:MAX_SUMMARY_CHARS] + + +def _runner(spec: client.McpSpec, tool_name: str, offered: str): + async def run(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: + event = { + "name": offered, + "kind": "mcp", + "label": f"{spec.name} · {tool_name}", + "query": _summary(args), + "detail": spec.name, + "results": [], + } + try: + text, failed = await client.call_tool(spec, tool_name, args) + except (McpError, FetchError) as exc: + message = exc.message + log.info("mcp %s/%s failed: %s", spec.slug, tool_name, message) + return ToolOutcome( + f"The {tool_name} tool on {spec.name} failed: {message}", + {**event, "status": "error", "error": message[:200]}, + ) + + text = text.strip()[: spec.max_chars] + if failed: + return ToolOutcome( + f"The tool reported an error:\n{text}" if text else "The tool reported an error.", + {**event, "status": "error", "error": text[:200] or "The tool reported an error."}, + ) + if not text: + return ToolOutcome( + f"{tool_name} returned nothing.", {**event, "status": "ok", "text": ""} + ) + return ToolOutcome(text, {**event, "status": "ok", "text": text[:MAX_EVENT_CHARS]}) + + return run + + +def _offered_tools(server: McpServer) -> list[dict[str, Any]]: + """The cached entries this server is currently willing to offer.""" + overrides = server.tool_overrides_json or {} + # Absent means on, the rule the model capability flags follow: a tool that + # appeared in the last refresh should work rather than silently do nothing. + return [ + entry + for entry in (server.tools_json or []) + if isinstance(entry, dict) and overrides.get(entry.get("name"), True) + ] + + +def tool_defs( + db: DBSession, user: User | None, *, everything: bool = False, taken: set[str] | None = None +) -> list[ToolDef]: + """One `ToolDef` per offerable tool across every visible server.""" + claimed = taken if taken is not None else set() + out: list[ToolDef] = [] + + for server in tool_access.visible_mcp_servers(db, user, everything=everything): + spec = client.spec_from(server) + for entry in _offered_tools(server): + name = str(entry.get("name") or "") + offered = str(entry.get("offer_name") or "") or offer_name( + server.slug, name, taken=claimed + ) + claimed.add(offered) + if not FUNCTION_NAME.match(offered): + continue + out.append( + ToolDef( + name=offered, + family=f"{FAMILY_MCP}:{server.slug}", + description=entry.get("description") or f"{name}, from {server.name}.", + parameters=entry.get("schema") or {"type": "object", "properties": {}}, + run=_runner(spec, name, offered), + ) + ) + + return out + + +async def refresh(db: DBSession, server: McpServer) -> tuple[int, str]: + """Contact a server and cache what it advertises. Returns (count, error). + + Shaped like `api/admin.py:_refresh_models`, including writing `last_error` + and `last_checked_at` on both paths so the row says what happened rather + than only whether it worked. + """ + spec = client.spec_from(server) + try: + advertised, info, version = await client.list_tools(spec) + except (McpError, FetchError) as exc: + server.last_error = exc.message + server.last_checked_at = datetime.now(UTC) + db.commit() + return 0, exc.message + + taken: set[str] = set() + for entry in advertised: + entry["offer_name"] = offer_name(server.slug, entry["name"], taken=taken) + + # Choices about tools that are still advertised survive; ones about tools + # that have gone are dropped rather than left to accumulate. + names = {entry["name"] for entry in advertised} + server.tool_overrides_json = { + name: on for name, on in (server.tool_overrides_json or {}).items() if name in names + } + + server.tools_json = advertised + server.server_info = info + server.protocol_version = version + server.last_error = "" + server.last_checked_at = datetime.now(UTC) + db.commit() + + log.info("mcp %s advertised %d tool(s)", server.slug, len(advertised)) + return len(advertised), "" + + +__all__ = ["offer_name", "refresh", "tool_defs"] diff --git a/src/lembas/services/tools.py b/src/lembas/services/tools.py index 27dbe09..5e01607 100644 --- a/src/lembas/services/tools.py +++ b/src/lembas/services/tools.py @@ -715,12 +715,19 @@ def _family_allowed( def _row_defs(db: DBSession, user: User | None, *, everything: bool = False) -> list[ToolDef]: """Tool definitions built from rows, in the order they claim names. - Imported here rather than at the top because `custom_tools` needs `ToolDef` - from this module. + Custom tools first, then MCP servers, because a custom tool's name is + written by hand and refused if it collides while an MCP tool's is derived + and renamed silently -- the one that can adapt should be the one that has to. + + Imported here rather than at the top because both modules need `ToolDef` + from this one. """ from lembas.services import custom_tools + from lembas.services.mcp import registry as mcp_registry - return custom_tools.tool_defs(db, user, everything=everything) + custom = custom_tools.tool_defs(db, user, everything=everything) + taken = {*REGISTRY, *(tool.name for tool in custom)} + return [*custom, *mcp_registry.tool_defs(db, user, everything=everything, taken=taken)] def _book(defs: list[ToolDef]) -> dict[str, ToolDef]: @@ -944,9 +951,11 @@ def _row_source(db: DBSession): it -- the override outlives the row. Gated on the tool's own family, so the guidance appears exactly when the - tool it describes is offered and never otherwise. + tool it describes is offered and never otherwise. An MCP server gets one + fragment rather than one per advertised tool: forty entries on the prompts + page is a page nobody would read. """ - from lembas.db.models import CustomTool + from lembas.db.models import CustomTool, McpServer for row in db.scalars(select(CustomTool).order_by(CustomTool.position, CustomTool.slug)): yield prompts_service.Fragment( @@ -959,6 +968,17 @@ def _row_source(db: DBSession): default=row.guidance or "", ) + for server in db.scalars(select(McpServer).order_by(McpServer.position, McpServer.slug)): + yield prompts_service.Fragment( + key=f"tool.mcp_{server.slug}", + label=server.name or server.slug, + group=prompts_service.GROUP_TOOLS, + order=600 + server.position, + families=(f"{FAMILY_MCP}:{server.slug}",), + hint=f"Appears when any tool from {server.name or server.slug} is offered.", + default=server.guidance or "", + ) + __all__ = [ "FAMILIES", diff --git a/src/lembas/web/templates/admin/_layout.html b/src/lembas/web/templates/admin/_layout.html index a943f6a..f88de21 100644 --- a/src/lembas/web/templates/admin/_layout.html +++ b/src/lembas/web/templates/admin/_layout.html @@ -47,6 +47,10 @@ {{ icon("link", "icon--sm") }} Tools + + {{ icon("server", "icon--sm") }} + MCP servers + {{ icon("sparkle", "icon--sm") }} Prompts diff --git a/src/lembas/web/templates/admin/_mcp_row.html b/src/lembas/web/templates/admin/_mcp_row.html new file mode 100644 index 0000000..f1ac1f2 --- /dev/null +++ b/src/lembas/web/templates/admin/_mcp_row.html @@ -0,0 +1,36 @@ +{% from "_macros.html" import icon %} +{# + One MCP server in the list. + + Swapped in place by the “Test & refresh” button, so this fragment has to be + able to render on its own as well as inside the list. +#} +
+
+
+ {{ server.name }} + {{ tool_count }} tool{{ '' if tool_count == 1 else 's' }} + {% if not server.enabled %}disabled{% endif %} + {% if not server.public %}restricted{% endif %} + {% if server.allow_private %}private network{% endif %} + {% if server.protocol_version %} + MCP {{ server.protocol_version }} + {% endif %} +
+ {{ server.slug }} · {{ server.url }} + {% if message %} +

{{ message }}

+ {% elif server.last_error %} +

{{ server.last_error }}

+ {% endif %} +
+ +
+ + Edit +
+
diff --git a/src/lembas/web/templates/admin/mcp.html b/src/lembas/web/templates/admin/mcp.html new file mode 100644 index 0000000..d23b3a0 --- /dev/null +++ b/src/lembas/web/templates/admin/mcp.html @@ -0,0 +1,53 @@ +{% extends "admin/_layout.html" %} +{% from "_macros.html" import icon %} +{% set section = "mcp" %} + +{% block title %}MCP servers - LLeMbas{% endblock %} +{% block heading %}MCP servers{% endblock %} + +{% block admin_content %} +

+ Remote servers speaking the Model Context Protocol over HTTP. Their tools are + offered beside the built-in ones to models marked MCP servers. + The list of tools is discovered and cached — press Test & + refresh after adding one, and again whenever the server changes. +

+ +
+ {{ icon("shield", "icon--sm") }} + + A server's tool names and descriptions are sent to the model as + instructions, and what it returns is read back as fact. Add servers you + trust, the way you would a dependency. + +
+ +{% if saved %} +
{{ icon("check", "icon--sm") }} {{ saved }}
+{% endif %} + +
+ + {{ icon("plus", "icon--sm") }} Add a server + +
+ +{% if not servers %} +
+ {{ icon("server", "empty__mark") }} +

+ No servers yet. You will need the URL of an MCP endpoint that speaks + streamable HTTP — local ones launched as a subprocess are not supported. +

+
+{% else %} + +
+ {% for server in servers %} + {% with tool_count = counts[server.id] %} + {% include "admin/_mcp_row.html" %} + {% endwith %} + {% endfor %} +
+{% endif %} +{% endblock %} diff --git a/src/lembas/web/templates/admin/mcp_detail.html b/src/lembas/web/templates/admin/mcp_detail.html new file mode 100644 index 0000000..578d5d1 --- /dev/null +++ b/src/lembas/web/templates/admin/mcp_detail.html @@ -0,0 +1,241 @@ +{% extends "admin/_layout.html" %} +{% from "_macros.html" import icon %} +{% set section = "mcp" %} + +{% block title %}{{ "New server" if is_new else server.name }} - LLeMbas{% endblock %} +{% block heading %}{{ "New MCP server" if is_new else server.name }}{% endblock %} + +{% block admin_content %} + + +{% if error %} +
{{ icon("warning", "icon--sm") }} {{ error }}
+{% endif %} + +{% if server.last_error %} +
+ {{ icon("warning", "icon--sm") }} + Last contacted unsuccessfully: {{ server.last_error }} +
+{% endif %} + +
+ +
+

The server

+ +
+ + +
+ +
+ + +

+ Prefixed onto every tool name this server offers, so that two servers + both exposing search do not collide. +

+
+ +
+ + +

+ The streamable-HTTP endpoint itself, the one that accepts a POST. A + server that answers with a redirect to somewhere else will be refused. +

+
+ +
+ + +

One Name: value per line.

+
+ +
+ + +
+ +
+ + +
+
+ +
+

Credential

+ +
+ + +
+ +
+ + +
+ +
+ + +

+ {% if server.secret_encrypted %} + Currently {{ masked }}. Leave the dots alone to keep it, + or clear the field to remove it. + {% else %} + Encrypted at rest and never shown again. + {% endif %} +

+
+
+ + {% if tools %} +
+

Tools it offers

+

+ Discovered at the last refresh. Untick one to withhold it — a tool this + server adds later is offered by default. +

+ + +
+ {% for tool in tools %} + + {% endfor %} +
+
+ {% elif not is_new %} +
+

Tools it offers

+

+ Nothing discovered yet. Save, then press Test & refresh + on the list. +

+
+ {% endif %} + +
+

Guidance

+ +
+ +

+ Added to the system message whenever any tool from this server is + offered. One piece of guidance for the server, not one per tool — + the tools carry their own descriptions. + {% if prompt_overridden %} +
Someone has overridden this wording under + Prompts — that is what the model + sees, not this. + {% endif %} +

+
+
+ +
+

Availability

+ +
+
+ + +
+

+ Tick the second only for a server on your own network. It is what stops + this being aimed at LLeMbas itself, a router, or a metadata endpoint. +

+
+ +
+ +
+ +
+ Groups with access + {% if groups %} +
+ {% for group in groups %} + + {% endfor %} +
+

Ignored while the server is available to everyone.

+ {% else %} +

+ No groups yet — create one to restrict access. +

+ {% endif %} +
+ +
+ + +
+
+ +
+ + Back to all servers + {% if not is_new %} + + {% endif %} +
+
+{% endblock %} diff --git a/tests/test_admin_tools.py b/tests/test_admin_tools.py index 24a0233..e9726dc 100644 --- a/tests/test_admin_tools.py +++ b/tests/test_admin_tools.py @@ -301,3 +301,196 @@ def test_the_secret_survives_a_round_trip_through_the_form(client: TestClient, d ) db.refresh(tool) assert decrypt(tool.secret_encrypted) == "s3cret" + + +# --- MCP servers ------------------------------------------------------------- +def _mcp_form(**overrides) -> dict: + base = { + "name": "GitHub", + "slug": "github", + "url": "https://mcp.test/rpc", + "headers": "", + "secret_placement": "none", + "secret_name": "Authorization", + "timeout": "30", + "max_chars": "8000", + "position": "0", + "enabled": "true", + "public": "true", + } + base.update(overrides) + return {key: value for key, value in base.items() if value is not None} + + +@pytest.fixture +def fake_mcp(mock_http, monkeypatch): + """A server that answers initialize and tools/list. + + Invented hostnames resolve to a public address; a literal IP is handed back + as itself, so a test about a private address is still testing one. + """ + import ipaddress + + def resolve(host, *_args, **_kwargs): + try: + ipaddress.ip_address(host) + except ValueError: + return [(2, 1, 6, "", ("93.184.216.34", 80))] + return [(2, 1, 6, "", (host, 80))] + + monkeypatch.setattr("socket.getaddrinfo", resolve) + + def handler(request: httpx.Request) -> httpx.Response: + if request.method != "POST": + return httpx.Response(405) + message = json.loads(request.content) + if message.get("method") == "notifications/initialized": + return httpx.Response(202) + if message.get("method") == "initialize": + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": message["id"], + "result": {"protocolVersion": "2025-06-18", "serverInfo": {"name": "fake"}}, + }, + headers={"mcp-session-id": "s-1"}, + ) + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": message["id"], + "result": {"tools": [{"name": "search", "description": "Search."}]}, + }, + ) + + mock_http(handler) + return handler + + +def test_the_mcp_pages_are_refused_to_a_plain_user(client: TestClient, plain_user): + assert client.get("/admin/mcp").status_code == 403 + assert client.get("/admin/mcp/new").status_code == 403 + assert client.post("/admin/mcp", data=_mcp_form()).status_code == 403 + + +def test_mcp_new_is_not_parsed_as_a_server_id(client: TestClient, registered): + response = client.get("/admin/mcp/new") + assert response.status_code == 200 + assert "New MCP server" in response.text + + +def test_adding_a_server_discovers_its_tools(client: TestClient, db, registered, fake_mcp): + from lembas.db.models import McpServer + + client.post("/admin/mcp", data=_mcp_form(), follow_redirects=False) + + server = db.scalar(select(McpServer)) + assert server.slug == "github" + assert [entry["name"] for entry in server.tools_json] == ["search"] + assert server.tools_json[0]["offer_name"] == "github_search" + assert server.protocol_version == "2025-06-18" + + +def test_refreshing_a_server_swaps_its_row(client: TestClient, db, registered, fake_mcp): + from lembas.db.models import McpServer + + client.post("/admin/mcp", data=_mcp_form(), follow_redirects=False) + server = db.scalar(select(McpServer)) + + response = client.post(f"/admin/mcp/{server.id}/test") + assert response.status_code == 200 + assert "found 1 tool." in response.text + assert f'id="mcp-{server.id}"' in response.text, "the fragment must render standalone" + + +def test_a_server_that_cannot_be_reached_says_so( + client: TestClient, db, registered, mock_http, monkeypatch +): + from lembas.db.models import McpServer + + monkeypatch.setattr( + "socket.getaddrinfo", lambda *a, **k: [(2, 1, 6, "", ("93.184.216.34", 80))] + ) + mock_http(lambda _r: httpx.Response(500, text="down")) + + response = client.post("/admin/mcp", data=_mcp_form(), follow_redirects=False) + assert response.status_code == 303 + server = db.scalar(select(McpServer)) + assert server is not None, "the row is still saved so the URL can be corrected" + assert server.last_error + + +def test_a_private_url_is_not_contacted_unless_the_box_is_ticked( + client: TestClient, db, registered, fake_mcp +): + """The row still saves, so the URL can be corrected -- but the discovery + that runs straight after it is refused, and the row says why.""" + from lembas.db.models import McpServer + + client.post( + "/admin/mcp", data=_mcp_form(url="http://127.0.0.1:9000/rpc"), follow_redirects=False + ) + server = db.scalar(select(McpServer)) + assert server.allow_private is False + assert "private or local" in server.last_error + assert server.tools_json == [] + + client.post( + f"/admin/mcp/{server.id}", + data=_mcp_form(url="http://127.0.0.1:9000/rpc", allow_private="true"), + follow_redirects=False, + ) + client.post(f"/admin/mcp/{server.id}/test") + db.refresh(server) + assert server.last_error == "" + assert [entry["name"] for entry in server.tools_json] == ["search"] + + +def test_a_duplicate_mcp_slug_is_refused(client: TestClient, db, registered, fake_mcp): + from lembas.db.models import McpServer + + client.post("/admin/mcp", data=_mcp_form(), follow_redirects=False) + response = client.post("/admin/mcp", data=_mcp_form(name="Other"), follow_redirects=False) + assert "already a server" in response.text + assert len(list(db.scalars(select(McpServer)))) == 1 + + +def test_a_url_that_is_not_http_is_refused(client: TestClient, db, registered): + from lembas.db.models import McpServer + + response = client.post( + "/admin/mcp", data=_mcp_form(url="ftp://mcp.test/rpc"), follow_redirects=False + ) + assert "http" in response.text + assert db.scalar(select(McpServer)) is None + + +def test_unticking_a_tool_withholds_it(client: TestClient, db, registered, fake_mcp): + from lembas.db.models import McpServer, User + from lembas.services.mcp import registry as mcp_registry + + client.post("/admin/mcp", data=_mcp_form(), follow_redirects=False) + server = db.scalar(select(McpServer)) + + client.post( + f"/admin/mcp/{server.id}", + data={**_mcp_form(), "tool_choices": "1", "tool_names": "search"}, + follow_redirects=False, + ) + db.refresh(server) + assert server.tool_overrides_json == {"search": False} + + owner = db.scalar(select(User)) + assert mcp_registry.tool_defs(db, owner) == [] + + +def test_deleting_a_server_removes_it(client: TestClient, db, registered, fake_mcp): + from lembas.db.models import McpServer + + client.post("/admin/mcp", data=_mcp_form(), follow_redirects=False) + server = db.scalar(select(McpServer)) + + client.post(f"/admin/mcp/{server.id}/delete", follow_redirects=False) + assert db.scalar(select(McpServer)) is None diff --git a/tests/test_generation_tools.py b/tests/test_generation_tools.py index 025a1bf..bc8638d 100644 --- a/tests/test_generation_tools.py +++ b/tests/test_generation_tools.py @@ -207,3 +207,126 @@ async def _empty_search(_config, _query, *, limit=None): async def _never_called_title(*_args, **_kwargs): """Auto-titling makes its own request; these tests are about the tool loop.""" return "A title" + + +# --- Progress and concurrency ------------------------------------------------ +async def test_the_status_names_the_running_tool_and_is_cleared(db, user_id, monkeypatch): + """A remote tool can take seconds with nothing streaming, and a silent + pause is exactly what a hang looks like.""" + settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH) + chat_id, message_id = _chat_with_tools(db, user_id) + + seen: list[str] = [] + + async def fake_search(_config, _query, *, limit=None): + seen.append(generation.status) + return [] + + monkeypatch.setattr("lembas.services.search.run", fake_search) + monkeypatch.setattr( + generation_service, + "stream_chat", + _stub_stream( + [[_tool_call_chunk("web_search", '{"query": "mallorn"}')], [_text_chunk("Done.")]], + [], + ), + ) + + generation = generation_service.Generation(chat_id=chat_id, message_id=message_id) + await generation_service._run(generation) + + assert seen == ["Running web_search…"] + assert generation.status == "", "and it is cleared once they are done" + + +async def test_results_stay_paired_with_their_calls_when_run_together(db, user_id, monkeypatch): + """Indexed rather than appended as they finish: an endpoint matching on + tool_call_id would otherwise pair the right id with the wrong content.""" + import asyncio + + settings_store.update(db, {"enabled": True}, key=settings_store.SEARCH) + chat_id, message_id = _chat_with_tools(db, user_id) + + async def slow_first(_config, query, *, limit=None): + # The first call finishes last, which is the whole point of the test. + await asyncio.sleep(0.02 if query == "first" else 0) + return [SearchResult(f"result for {query}", f"https://t.test/{query}", "")] + + monkeypatch.setattr("lembas.services.search.run", slow_first) + + two_calls = { + "choices": [ + {"delta": {"tool_calls": [ + {"index": 0, "id": "a", "function": { + "name": "web_search", "arguments": '{"query": "first"}'}}, + {"index": 1, "id": "b", "function": { + "name": "web_search", "arguments": '{"query": "second"}'}}, + ]}} + ] + } + payloads: list[dict] = [] + monkeypatch.setattr( + generation_service, + "stream_chat", + _stub_stream([[two_calls], [_text_chunk("Done.")]], payloads), + ) + + generation = generation_service.Generation(chat_id=chat_id, message_id=message_id) + await generation_service._run(generation) + + turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"] + assert [turn["tool_call_id"] for turn in turns] == ["a", "b"] + assert "first" in turns[0]["content"] and "second" in turns[1]["content"] + # And the transcript keeps the same order. + assert [event["query"] for event in generation.tool_events] == ["first", "second"] + + +async def test_a_custom_tool_runs_inside_the_loop(db, user_id, monkeypatch, mock_http): + """End to end: a row becomes an offered schema, the model calls it, and the + result comes back in the next request's messages.""" + import httpx + + from lembas.db.models import CustomTool + + monkeypatch.setattr( + "socket.getaddrinfo", lambda *a, **k: [(2, 1, 6, "", ("93.184.216.34", 80))] + ) + mock_http(lambda _r: httpx.Response(200, json={"summary": "Sunny in Minas Tirith."})) + + db.add( + CustomTool( + slug="weather", + name="Weather", + description="Look up the weather.", + url_template="https://api.test/{{city}}", + parameters_json={"type": "object", "properties": {"city": {"type": "string"}}}, + response_mode="json", + response_path="summary", + ) + ) + db.commit() + + chat_id, message_id = _chat_with_tools(db, user_id) + payloads: list[dict] = [] + monkeypatch.setattr( + generation_service, + "stream_chat", + _stub_stream( + [ + [_tool_call_chunk("weather", '{"city": "Minas Tirith"}')], + [_text_chunk("It is sunny.")], + ], + payloads, + ), + ) + + generation = generation_service.Generation(chat_id=chat_id, message_id=message_id) + await generation_service._run(generation) + + offered = {tool["function"]["name"] for tool in payloads[0]["tools"]} + assert "weather" in offered + + tool_turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"] + assert tool_turns[0]["content"] == "Sunny in Minas Tirith." + assert generation.tool_events[0]["kind"] == "custom" + assert generation.tool_events[0]["label"] == "Weather" diff --git a/tests/test_mcp.py b/tests/test_mcp.py new file mode 100644 index 0000000..c3aa271 --- /dev/null +++ b/tests/test_mcp.py @@ -0,0 +1,506 @@ +"""The MCP client: framing, the session lifecycle, and what comes back. + +The fake server below answers the way a real one does -- JSON-RPC over POST, +either as one JSON object or as an event stream, with a session id it expects +echoed. Everything is driven through `mock_http`, so no test touches a network. +""" + +from __future__ import annotations + +import json + +import httpx +import pytest +from sqlalchemy import select + +from lembas.db.models import McpServer +from lembas.services.mcp import client, protocol, registry +from lembas.services.mcp.protocol import McpError + + +@pytest.fixture(autouse=True) +def dns(monkeypatch): + """`mcp.test` does not exist and check_url resolves for real.""" + import ipaddress + import socket + + def resolve(host, *_args, **_kwargs): + try: + ipaddress.ip_address(host) + except ValueError: + return [(2, 1, 6, "", ("93.184.216.34", 80))] + return [(2, 1, 6, "", (host, 80))] + + monkeypatch.setattr(socket, "getaddrinfo", resolve) + + +def _spec(**overrides) -> client.McpSpec: + return client.McpSpec( + **{"slug": "github", "name": "GitHub", "url": "https://mcp.test/rpc", **overrides} + ) + + +class FakeServer: + """One server's worth of behaviour, as an httpx handler.""" + + def __init__(self, *, tools=None, stream=False, session_id="s-1", pages=None): + self.tools = tools if tools is not None else [ + {"name": "search", "description": "Search things.", "inputSchema": {"type": "object"}} + ] + self.stream = stream + self.session_id = session_id + self.pages = pages + self.requests: list[dict] = [] + self.headers: list[httpx.Headers] = [] + self.result_for = {} + + def _frame(self, message: dict) -> httpx.Response: + if self.stream: + body = "".join(f"data: {line}\n" for line in json.dumps(message).splitlines()) + "\n" + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + return httpx.Response(200, json=message) + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.headers.append(request.headers) + if request.method == "DELETE": + return httpx.Response(405) + + message = json.loads(request.content) + self.requests.append(message) + method, request_id = message.get("method"), message.get("id") + + if method == "notifications/initialized": + return httpx.Response(202) + + if method == "initialize": + headers = {"mcp-session-id": self.session_id} if self.session_id else {} + response = self._frame( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "protocolVersion": "2025-06-18", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "fake", "version": "1"}, + }, + } + ) + return httpx.Response( + 200, content=response.content, headers={**dict(response.headers), **headers} + ) + + if method == "tools/list": + if self.pages is not None: + cursor = (message.get("params") or {}).get("cursor", "") + index = int(cursor or 0) + page = self.pages[index] + result = {"tools": page} + if index + 1 < len(self.pages): + result["nextCursor"] = str(index + 1) + return self._frame({"jsonrpc": "2.0", "id": request_id, "result": result}) + return self._frame( + {"jsonrpc": "2.0", "id": request_id, "result": {"tools": self.tools}} + ) + + if method == "tools/call": + name = (message.get("params") or {}).get("name") + result = self.result_for.get( + name, {"content": [{"type": "text", "text": "It worked."}]} + ) + return self._frame({"jsonrpc": "2.0", "id": request_id, "result": result}) + + return self._frame( + {"jsonrpc": "2.0", "id": request_id, "error": {"code": -32601, "message": "Unknown"}} + ) + + +# --- The lifecycle ----------------------------------------------------------- +async def test_initialize_is_followed_by_the_initialized_notification(mock_http): + server = FakeServer() + mock_http(server) + + tools, info, version = await client.list_tools(_spec()) + methods = [request.get("method") for request in server.requests] + assert methods[:3] == ["initialize", "notifications/initialized", "tools/list"] + assert info == {"name": "fake", "version": "1"} + assert version == "2025-06-18" + assert [tool["name"] for tool in tools] == ["search"] + + +async def test_the_session_id_is_echoed_on_every_later_request(mock_http): + server = FakeServer(session_id="abc-123") + mock_http(server) + + await client.list_tools(_spec()) + assert "mcp-session-id" not in server.headers[0] + assert all(headers["mcp-session-id"] == "abc-123" for headers in server.headers[1:]) + assert all("mcp-protocol-version" in headers for headers in server.headers[1:]) + + +async def test_a_server_that_issues_no_session_is_fine(mock_http): + mock_http(FakeServer(session_id="")) + tools, _info, _version = await client.list_tools(_spec()) + assert tools + + +async def test_an_expired_session_is_reinitialised_once(mock_http): + """A 404 on a request carrying a session id means the server forgot us.""" + server = FakeServer() + calls = {"list": 0} + inner = server.__call__ + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + message = json.loads(request.content) + if message.get("method") == "tools/list": + calls["list"] += 1 + if calls["list"] == 1: + return httpx.Response(404, text="no such session") + return inner(request) + + mock_http(handler) + tools, _info, _version = await client.list_tools(_spec()) + assert tools, "the retry after re-initialising must succeed" + assert [r.get("method") for r in server.requests].count("initialize") == 2 + + +async def test_a_json_response_and_an_event_stream_parse_the_same(mock_http): + plain = FakeServer(stream=False) + mock_http(plain) + from_json, _i, _v = await client.list_tools(_spec()) + + streamed = FakeServer(stream=True) + mock_http(streamed) + from_stream, _i, _v = await client.list_tools(_spec()) + + assert from_json == from_stream + + +async def test_tools_list_follows_the_cursor(mock_http): + server = FakeServer( + pages=[ + [{"name": "one", "inputSchema": {"type": "object"}}], + [{"name": "two", "inputSchema": {"type": "object"}}], + ] + ) + mock_http(server) + tools, _i, _v = await client.list_tools(_spec()) + assert [tool["name"] for tool in tools] == ["one", "two"] + + +async def test_a_jsonrpc_error_becomes_a_readable_message(mock_http): + def handler(request: httpx.Request) -> httpx.Response: + message = json.loads(request.content) + if message.get("method") == "notifications/initialized": + return httpx.Response(202) + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": message.get("id"), + "error": {"code": -32602, "message": "invalid params"}, + }, + ) + + mock_http(handler) + with pytest.raises(McpError, match="invalid params"): + await client.list_tools(_spec()) + + +async def test_an_http_error_is_reported_with_its_status(mock_http): + mock_http(lambda _r: httpx.Response(502, text="bad gateway")) + with pytest.raises(McpError, match="502"): + await client.list_tools(_spec()) + + +# --- Reaching the network ---------------------------------------------------- +async def test_the_url_is_checked_before_the_first_request(mock_http): + server = FakeServer() + mock_http(server) + from lembas.services.fetch import FetchError + + with pytest.raises(FetchError, match="private or local"): + await client.list_tools(_spec(url="http://127.0.0.1:9000/rpc")) + assert server.requests == [], "nothing may be sent before the address is checked" + + +async def test_a_302_is_refused_and_a_307_is_followed_and_checked(mock_http): + """301, 302 and 303 turn a POST into a GET, which means nothing to a + JSON-RPC endpoint.""" + inner = FakeServer() + + def redirect_302(request: httpx.Request) -> httpx.Response: + return httpx.Response(302, headers={"location": "https://elsewhere.test/rpc"}) + + mock_http(redirect_302) + with pytest.raises(McpError, match="turn the request into a GET"): + await client.list_tools(_spec()) + + def redirect_307(request: httpx.Request) -> httpx.Response: + if request.url.host == "mcp.test": + return httpx.Response(307, headers={"location": "https://elsewhere.test/rpc"}) + return inner(request) + + mock_http(redirect_307) + tools, _i, _v = await client.list_tools(_spec()) + assert tools + + +async def test_a_redirect_to_a_private_address_is_refused(mock_http, monkeypatch): + def resolve(host, *_args, **_kwargs): + if host == "inside.test": + return [(2, 1, 6, "", ("10.0.0.5", 80))] + return [(2, 1, 6, "", ("93.184.216.34", 80))] + + monkeypatch.setattr("socket.getaddrinfo", resolve) + mock_http(lambda _r: httpx.Response(307, headers={"location": "https://inside.test/rpc"})) + + from lembas.services.fetch import FetchError + + with pytest.raises(FetchError, match="private or local"): + await client.list_tools(_spec()) + + +async def test_the_secret_reaches_the_server(mock_http): + server = FakeServer() + mock_http(server) + await client.list_tools(_spec(secret="tok", secret_placement="bearer")) + assert server.headers[0]["authorization"] == "Bearer tok" + + +# --- Reading a tool result --------------------------------------------------- +async def test_text_blocks_are_joined(mock_http): + server = FakeServer() + server.result_for["search"] = { + "content": [{"type": "text", "text": "one"}, {"type": "text", "text": "two"}] + } + mock_http(server) + text, failed = await client.call_tool(_spec(), "search", {}) + assert text == "one\n\ntwo" + assert failed is False + + +async def test_an_image_block_is_described_rather_than_forwarded(mock_http): + """A tool turn is a string, images only reach models marked as having + vision, and base64 in a tool result fills a window with nothing.""" + server = FakeServer() + server.result_for["search"] = { + "content": [{"type": "image", "mimeType": "image/png", "data": "A" * 400}] + } + mock_http(server) + text, _failed = await client.call_tool(_spec(), "search", {}) + assert "A" * 40 not in text + assert "image/png" in text + assert "not shown to the model" in text + + +async def test_a_resource_block_with_text_keeps_it(mock_http): + server = FakeServer() + server.result_for["search"] = { + "content": [ + {"type": "resource", "resource": {"uri": "file:///a.txt", "text": "contents"}} + ] + } + mock_http(server) + text, _failed = await client.call_tool(_spec(), "search", {}) + assert "file:///a.txt" in text and "contents" in text + + +async def test_a_resource_block_with_no_text_names_its_type(mock_http): + server = FakeServer() + server.result_for["search"] = { + "content": [ + {"type": "resource", "resource": {"uri": "file:///a.bin", "mimeType": "application/x"}} + ] + } + mock_http(server) + text, _failed = await client.call_tool(_spec(), "search", {}) + assert "application/x" in text + + +async def test_an_unknown_block_type_degrades_to_a_note(mock_http): + server = FakeServer() + server.result_for["search"] = {"content": [{"type": "hologram"}]} + mock_http(server) + text, _failed = await client.call_tool(_spec(), "search", {}) + assert "hologram" in text + + +async def test_structured_content_is_used_when_there_is_no_text(mock_http): + server = FakeServer() + server.result_for["search"] = {"content": [], "structuredContent": {"count": 3}} + mock_http(server) + text, _failed = await client.call_tool(_spec(), "search", {}) + assert json.loads(text) == {"count": 3} + + +async def test_is_error_is_reported_as_a_failure(mock_http): + server = FakeServer() + server.result_for["search"] = { + "content": [{"type": "text", "text": "no such repo"}], + "isError": True, + } + mock_http(server) + _text, failed = await client.call_tool(_spec(), "search", {}) + assert failed is True + + +# --- Bounding what a server can say ------------------------------------------ +def test_a_giant_description_is_truncated(): + cleaned = protocol.clean_tool({"name": "x", "description": "d" * 9000}) + assert len(cleaned["description"]) == protocol.MAX_DESCRIPTION + + +def test_a_schema_that_is_not_an_object_is_replaced(): + cleaned = protocol.clean_tool({"name": "x", "inputSchema": "nope"}) + assert cleaned["schema"] == {"type": "object", "properties": {}} + + +def test_a_giant_schema_is_replaced(): + huge = {"type": "object", "properties": {f"p{i}": {"type": "string"} for i in range(1000)}} + cleaned = protocol.clean_tool({"name": "x", "inputSchema": huge}) + assert cleaned["schema"] == {"type": "object", "properties": {}} + + +def test_an_entry_with_no_name_is_dropped(): + assert protocol.clean_tool({"description": "orphan"}) is None + assert protocol.clean_tool("not a dict") is None + + +# --- Namespacing ------------------------------------------------------------- +def test_tool_names_are_namespaced_and_within_the_charset(): + taken: set[str] = set() + assert registry.offer_name("github", "search", taken=taken) == "github_search" + assert registry.offer_name("github", "Create Issue!", taken=taken) == "github_create_issue" + assert all(registry.FUNCTION_NAME.match(name) for name in taken) + + +def test_a_name_colliding_with_a_builtin_is_renamed(): + from lembas.services import tools as tools_service + + taken = set(tools_service.REGISTRY) + name = registry.offer_name("", "web_search", taken=taken) + assert name != "web_search" + assert name not in tools_service.REGISTRY + + +def test_a_truncation_collision_is_disambiguated(): + """Two long names can collide once cut to 64 characters where the full ones + would not have.""" + taken: set[str] = set() + first = registry.offer_name("srv", "a" * 200, taken=taken) + second = registry.offer_name("srv", "a" * 201, taken=taken) + assert first != second + assert len(first) <= 64 and len(second) <= 64 + + +# --- Rows become tools ------------------------------------------------------- +def _server(db, **overrides) -> McpServer: + row = McpServer( + **{ + "slug": "github", + "name": "GitHub", + "url": "https://mcp.test/rpc", + "tools_json": [ + { + "name": "search", + "offer_name": "github_search", + "description": "Search things.", + "schema": {"type": "object", "properties": {}}, + } + ], + **overrides, + } + ) + db.add(row) + db.commit() + return row + + +def test_a_cached_tool_becomes_a_definition(db, user_id): + from lembas.db.models import User + + _server(db) + defs = registry.tool_defs(db, db.get(User, user_id)) + assert [tool.name for tool in defs] == ["github_search"] + assert defs[0].family == "mcp:github" + + +def test_a_tool_switched_off_is_not_offered(db, user_id): + from lembas.db.models import User + + _server(db, tool_overrides_json={"search": False}) + assert registry.tool_defs(db, db.get(User, user_id)) == [] + + +def test_a_tool_absent_from_the_overrides_is_on(db, user_id): + """Absent means on, the rule the model capability flags follow.""" + from lembas.db.models import User + + _server(db, tool_overrides_json={"something_else": False}) + assert len(registry.tool_defs(db, db.get(User, user_id))) == 1 + + +async def test_refresh_caches_the_tools_and_clears_the_last_error(db, mock_http): + mock_http(FakeServer()) + server = _server(db, tools_json=[], last_error="it was broken") + + count, error = await registry.refresh(db, server) + assert (count, error) == (1, "") + db.refresh(server) + assert server.tools_json[0]["offer_name"] == "github_search" + assert server.last_error == "" + assert server.protocol_version == "2025-06-18" + assert server.server_info["name"] == "fake" + + +async def test_a_failed_refresh_records_it_and_keeps_the_cached_tools(db, mock_http): + mock_http(lambda _r: httpx.Response(500, text="down")) + server = _server(db) + + count, error = await registry.refresh(db, server) + assert count == 0 and "500" in error + db.refresh(server) + assert server.last_error + assert server.last_checked_at is not None + assert len(server.tools_json) == 1, "a bad refresh must not empty the list" + + +async def test_a_refresh_keeps_a_choice_about_a_tool_that_survives(db, mock_http): + mock_http(FakeServer(tools=[{"name": "search"}, {"name": "gone"}])) + server = _server(db, tool_overrides_json={"search": False, "vanished": False}) + await registry.refresh(db, server) + + db.refresh(server) + assert server.tool_overrides_json == {"search": False} + + +async def test_a_failing_tool_call_reaches_the_model_as_words(db, user_id, mock_http): + from lembas.db.models import User + + mock_http(lambda _r: httpx.Response(500, text="down")) + _server(db) + tool = registry.tool_defs(db, db.get(User, user_id))[0] + + outcome = await tool.run(None, {}) + assert outcome.event["status"] == "error" + assert "GitHub" in outcome.content + + +def test_two_servers_exposing_the_same_tool_do_not_collide(db, user_id): + from lembas.db.models import User + + _server(db) + _server(db, slug="gitlab", name="GitLab", tools_json=[{"name": "search", "schema": {}}]) + names = [tool.name for tool in registry.tool_defs(db, db.get(User, user_id))] + assert len(names) == len(set(names)) == 2 + + +def test_a_disabled_server_offers_nothing(db, user_id): + from lembas.db.models import User + + _server(db, enabled=False) + assert registry.tool_defs(db, db.get(User, user_id)) == [] + assert db.scalar(select(McpServer)) is not None