MCP servers, over streamable HTTP

A server is a row with a URL; its tools are discovered by a button and
cached, then offered beside the built-in ones. Written by hand rather than
taken from the reference SDK, because that SDK's transport does its own
connecting -- and the one thing that must not be bypassed is check_url on
every hop. Owning the transport is the point; the framing beside it is the
small part.

Sessions are per call: initialize, initialized, the call, a best-effort
DELETE. Caching one wants an owner, a TTL, eviction, a lock and a shutdown
hook, and the server may expire it under all of that anyway -- ToolContext
is a session-free snapshot precisely so nothing in a tool holds live state.

A server's names and descriptions reach the model as instructions and are
bounded before they do; what it returns is escaped preformatted text, never
markdown. Tools are namespaced per server, so two servers exposing "search"
do not collide and neither shadows a built-in.

Also: a round's calls now run together under a semaphore, results indexed
so each tool turn stays paired with its call, and generation.status names
what is running -- a remote tool is latency-bound, and a silent pause is
what a hang looks like.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
2026-08-01 16:44:29 +02:00
parent bc84fec21d
commit ecb52e9978
17 changed files with 2270 additions and 28 deletions
+88 -9
View File
@@ -19,7 +19,7 @@ lembas info # paths + counts, useful when confused
lembas secret-key # generate LEMBAS_SECRET_KEY lembas secret-key # generate LEMBAS_SECRET_KEY
lembas create-admin # create or promote an admin lembas create-admin # create or promote an admin
pytest # 590 tests, ~35s pytest # 698 tests, ~40s
# PLAN.md tracks what is and is not built # PLAN.md tracks what is and is not built
ruff check . # lint (line length 100) ruff check . # lint (line length 100)
python scripts/build_artwork.py # regenerate artwork (SVG + PWA icons; python scripts/build_artwork.py # regenerate artwork (SVG + PWA icons;
@@ -82,6 +82,7 @@ src/lembas/
admin_search.py web search provider and credentials admin_search.py web search provider and credentials
admin_prompts.py the prompt fragment editor and its preview admin_prompts.py the prompt fragment editor and its preview
admin_suggestions.py the cards offered on the new-chat screen admin_suggestions.py the cards offered on the new-chat screen
admin_tools.py custom HTTP tools and MCP servers
audio.py transcribe, speak, voice discovery audio.py transcribe, speak, voice discovery
library.py knowledge, notes, skills pages; memory CRUD library.py knowledge, notes, skills pages; memory CRUD
files.py upload, serve, remove attachments files.py upload, serve, remove attachments
@@ -96,6 +97,7 @@ src/lembas/
llm/openai_client.py httpx streaming + model discovery llm/openai_client.py httpx streaming + model discovery
search/ ddgs, SearXNG and Firecrawl behind one shape search/ ddgs, SearXNG and Firecrawl behind one shape
library/ documents, notes, memories, skills, FTS library/ documents, notes, memories, skills, FTS
mcp/ remote MCP servers: framing, transport, rows to tools
audio.py OpenAI-shaped /v1/audio/* client audio.py OpenAI-shaped /v1/audio/* client
fetch.py URL retrieval, HTML to text, the SSRF guard fetch.py URL retrieval, HTML to text, the SSRF guard
sharing.py one visibility rule for every library store sharing.py one visibility rule for every library store
@@ -106,6 +108,8 @@ src/lembas/
suggestions.py new-chat starting points, seeded once suggestions.py new-chat starting points, seeded once
harness.py the operational prompt built from what a model has harness.py the operational prompt built from what a model has
tools.py tool registry, schemas, streamed-call reassembly tools.py tool registry, schemas, streamed-call reassembly
custom_tools.py the admin-defined HTTP tool runner
tool_access.py who may be offered which admin-defined tool
chat.py request building, endpoint resolution, titles chat.py request building, endpoint resolution, titles
markdown.py markdown-it + pygments + nh3 markdown.py markdown-it + pygments + nh3
crypto.py Fernet encrypt/decrypt/mask crypto.py Fernet encrypt/decrypt/mask
@@ -317,6 +321,77 @@ survives. Tools are only offered when search is enabled, the user has
array to an endpoint without support fails the whole request, exactly as images array to an endpoint without support fails the whole request, exactly as images
do without `vision`. do without `vision`.
**The registry is resolved per request, not imported.** `REGISTRY` holds the
built-ins; a custom tool or an MCP tool is a row. `tools.resolve_tools()` returns
a `ToolSet` carrying the schemas *and* the runners, and the runners travel to the
loop on `ToolContext.tools` — because a generation outlives the session that
could look them up. `run_tool` consults that map, so what may be *run* is what
was *offered*. `None` means nobody resolved a set and falls back to the built-ins;
an empty dict is authoritative. Reaching for the global registry instead is how a
model naming a tool its chat was gated out of used to get it run anyway.
**A row-backed tool's family is `gate:slug`.** `custom:weather`, `mcp:github`.
The capability flag and the permission are named after the *gate*
(`tool_custom` / `tools.custom`), so a server advertising forty tools does not
mean forty checkboxes on every model; the full family exists so each row can
carry its own prompt fragment, gated to appear exactly when its tool is offered.
`harness._families` and `admin_prompts` therefore take a `db`. Custom and MCP
deliberately do **not** require `library.use`: an endpoint an administrator wrote
has nothing to do with anyone's own notes.
**An argument may fill a hole; it may never move the target.** A custom tool's
URL is a template. The scheme and host must be literal — checked at save *and*
again at call time, since a row can predate a check — values are escaped for
where they land (`quote(safe="")` in a URL, JSON-escaped in a body, control
characters stripped in a header), and the filled URL's origin is compared with
the template's afterwards. An undeclared `{{name}}` becomes nothing rather than
passing through, which is the opposite of `prompts.substitute` and deliberately
so: a literal `{{x}}` in a URL is not a feature.
**Three places now follow redirects by hand.** `fetch.fetch`,
`custom_tools._send` and `mcp.client.Session._post`, each re-running
`check_url` on every hop. `fetch()` itself is not reusable — GET-only,
bodyless, and it *raises* on any content type that is not HTML or text, which is
every JSON API there is. The duplication is deliberate; bending a page fetcher
into a general HTTP client is not. A secret is dropped when a hop leaves the
origin it was issued for.
**MCP sessions are per call.** Initialize, `notifications/initialized`, the call,
then a best-effort `DELETE`. Caching one would need an owner, a TTL, eviction, a
lock (a round runs its tools concurrently) and a shutdown hook, and the server
can expire it underneath all of that anyway — `ToolContext` is a session-free
snapshot precisely so nothing in a tool holds live state. The cost is one POST in
front of a call that is already a network round trip. `307`/`308` are followed;
`301`/`302`/`303` turn a POST into a GET and are refused rather than guessed at.
**A discovered MCP tool is JSON, not a row.** `McpServer.tools_json` caches
`tools/list`. `Model` is a table because each row carries eight independent admin
decisions; a discovered tool carries one (offered or not, in
`tool_overrides_json`, where absent means on), credentials and guidance are per
server, and the list is replaced wholesale on every refresh — a table would mean
reconciling rows against a cache of somebody else's document.
**An MCP tool has two names.** The server's own, which `tools/call` needs, and
the offered one in the schema — `slug_tool`, lowercased into
`[a-z0-9_-]{1,64}` because endpoints accept less than MCP does. Built-ins claim
their names first and can never be shadowed; a custom tool whose slug collides is
*refused at save*, an MCP tool is renamed silently, since the one that can adapt
should be the one that has to. The rename never leaves `mcp/registry.py`.
**A server's tool metadata is untrusted input that becomes instructions.**
Names, descriptions and schemas from `tools/list` are bounded and sanitised in
`mcp/protocol.clean_tool` before anything reaches a model. What a tool *returns*
is untrusted too, and is rendered as escaped preformatted text — never through
`services/markdown.py`, which is the one path allowed to emit HTML.
**A round's tool calls run together.** `generation._run_calls` gathers them under
a semaphore of four and keeps the results **indexed, not appended as they
finish**: each tool turn must line up with the assistant turn's `tool_calls` or
an endpoint matching on `tool_call_id` pairs the right id with the wrong content.
Safe because `run_tool` never raises and every runner opens its own session.
`generation.status` names what is running, because a remote tool taking seconds
with nothing streaming is exactly what a hang looks like.
**Tool-call arguments arrive in fragments.** `delta.tool_calls` carries an **Tool-call arguments arrive in fragments.** `delta.tool_calls` carries an
`index`, a name that appears once, and an `arguments` string split across `index`, a name that appears once, and an `arguments` string split across
chunks. `tools.ToolCallAccumulator` rejoins them keyed on `index` — not on chunks. `tools.ToolCallAccumulator` rejoins them keyed on `index` — not on
@@ -503,14 +578,18 @@ notes describe the machine.
## Not built yet ## Not built yet
Custom tools and MCP, agentic execution (local subprocess and SSH connection Agentic execution (local subprocess and SSH connection profiles), image
profiles), image generation. Nav entries mark where each one goes. The tool generation. A nav entry marks where each one goes. The tool loop in
loop in `services/generation.py` is what they plug into — a new tool is a `services/generation.py` is what they plug into — a new tool is a `ToolDef`
`ToolDef` in `services/tools.py:REGISTRY` plus a permission and a capability reaching `tools.resolve_tools()` plus a permission and a capability flag, not a
flag, not a new code path. Its guidance is the same shape: a new code path. Its guidance is the same shape: a `prompts.register_source`
`prompts.register_source` yielding one `Fragment` per tool row puts it in the yielding one `Fragment` per row puts it in the harness, on the admin page and in
harness, on the admin page and in the preview without touching the assembler, the preview without touching the assembler, the save handler or a template.
the save handler or a template. Custom HTTP tools and MCP servers are built and are the worked example of both.
**Local MCP is deliberately absent.** Only remote servers over streamable HTTP.
Spawning `npx` 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.
**Unknown is not zero.** `Model.context_length` of 0 means nobody has said how **Unknown is not zero.** `Model.context_length` of 0 means nobody has said how
big the window is, which is different from "small". The context percentage is big the window is, which is different from "small". The context percentage is
+16 -7
View File
@@ -64,6 +64,22 @@ be a different project, not a refactor.
support rejects the whole request rather than ignoring the array support rejects the whole request rather than ignoring the array
- [x] Sources stay in the transcript; results are **not** replayed as context on - [x] Sources stay in the transcript; results are **not** replayed as context on
the next turn, for the same reasons reasoning is not the next turn, for the same reasons reasoning is not
- [x] A round's calls run together, and the reply says which tool is running —
a remote tool taking seconds with nothing streaming looks like a hang
- [x] **Custom HTTP tools** — an administrator describes one call: a JSON Schema,
a URL template, headers, an encrypted secret and how to read the answer.
Arguments may fill a hole but never move the target: the scheme and host
are literal, values are escaped for where they land, and the origin is
pinned afterwards
- [x] **MCP servers** over streamable HTTP — a hand-written client, so that
`check_url` runs on every hop rather than being bypassed by somebody
else's transport. Tools are discovered and cached by a button, namespaced
per server, and a server's own descriptions are bounded before they reach
a model as instructions
- [x] Both gated like the built-ins — a model capability, a permission — and
restrictable to groups, with guidance of their own on `/admin/prompts`
- [x] Local MCP over stdio is deliberately absent: spawning a subprocess is the
agentic-execution feature and wants a confirmation model first
### The library ### The library
- [x] **Knowledge bases** — documents, images and saved web pages, grouped into - [x] **Knowledge bases** — documents, images and saved web pages, grouped into
@@ -155,13 +171,6 @@ be a different project, not a refactor.
In the order they are likely to be worth doing. In the order they are likely to be worth doing.
### Custom tools and MCP servers
An MCP client managing configured servers, their tools surfaced alongside the
built-in ones. The loop they plug into exists now — `services/tools.py` is a
registry of thirteen tools and `services/generation.py` already runs bounded
rounds — so this is a client and an admin screen rather than a change to how
chat works.
### Agentic execution ### Agentic execution
Two modes, as originally specified: Two modes, as originally specified:
- **local** — subprocess on the machine LLeMbas runs on - **local** — subprocess on the machine LLeMbas runs on
+6 -2
View File
@@ -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 - **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 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 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 - **Speech in and out** — dictate a message and have replies read aloud, against
any OpenAI-compatible audio endpoint (whisper.cpp, Speaches, Kokoro…). Each any OpenAI-compatible audio endpoint (whisper.cpp, Speaches, Kokoro…). Each
person picks their own voice person picks their own voice
@@ -79,8 +83,8 @@ runtime. Clone it, `pip install -e .`, run it.
**Planned** **Planned**
Custom tools and MCP servers · agentic execution (local and over SSH) · image Agentic execution (local and over SSH) · image generation · OCR for scanned
generation · OCR for scanned PDFs · semantic search in the library. PDFs · semantic search in the library.
See [PLAN.md](PLAN.md) for what is built, what is not, and why. See [PLAN.md](PLAN.md) for what is built, what is not, and why.
+246
View File
@@ -31,11 +31,15 @@ from lembas.db.models import (
SECRET_PLACEMENTS, SECRET_PLACEMENTS,
CustomTool, CustomTool,
Group, Group,
McpServer,
) )
from lembas.services import custom_tools from lembas.services import custom_tools
from lembas.services import prompts as prompts_service from lembas.services import prompts as prompts_service
from lembas.services import tools as tools_service from lembas.services import tools as tools_service
from lembas.services.crypto import UNCHANGED_SENTINEL, decrypt, keep_or_replace, mask 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 from lembas.web.templating import render
log = logging.getLogger(__name__) 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) log.info("%s updated custom tool %s", user.email, tool.slug)
return _back(f"Saved {tool.name}.") 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}.")
+54 -5
View File
@@ -338,13 +338,19 @@ async def _run(generation: Generation) -> None:
*payload["messages"], *payload["messages"],
tools_service.assistant_turn(calls, "".join(round_text)), tools_service.assistant_turn(calls, "".join(round_text)),
] ]
for call in calls:
outcome = await tools_service.run_tool( generation.status = _tool_status(calls)
tool_context, call["name"], call["arguments"] generation.touch()
) try:
generation.tool_events.append(outcome.event) outcomes = await _run_calls(tool_context, calls)
finally:
generation.status = ""
generation.touch() 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)) messages.append(tools_service.tool_turn(call, outcome.content))
generation.touch()
payload = {**payload, "messages": messages} payload = {**payload, "messages": messages}
@@ -465,6 +471,49 @@ async def _maybe_compact(generation: Generation) -> None:
generation.touch() 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: def _pending_text(db, message: Message) -> str:
"""The user turn this reply is answering, for the size estimate.""" """The user turn this reply is answering, for the size estimate."""
previous = db.scalars( previous = db.scalars(
+26
View File
@@ -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",
]
+273
View File
@@ -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"]
+195
View File
@@ -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",
]
+185
View File
@@ -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"]
+25 -5
View File
@@ -715,12 +715,19 @@ def _family_allowed(
def _row_defs(db: DBSession, user: User | None, *, everything: bool = False) -> list[ToolDef]: def _row_defs(db: DBSession, user: User | None, *, everything: bool = False) -> list[ToolDef]:
"""Tool definitions built from rows, in the order they claim names. """Tool definitions built from rows, in the order they claim names.
Imported here rather than at the top because `custom_tools` needs `ToolDef` Custom tools first, then MCP servers, because a custom tool's name is
from this module. 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 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]: def _book(defs: list[ToolDef]) -> dict[str, ToolDef]:
@@ -944,9 +951,11 @@ def _row_source(db: DBSession):
it -- the override outlives the row. it -- the override outlives the row.
Gated on the tool's own family, so the guidance appears exactly when the 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)): for row in db.scalars(select(CustomTool).order_by(CustomTool.position, CustomTool.slug)):
yield prompts_service.Fragment( yield prompts_service.Fragment(
@@ -959,6 +968,17 @@ def _row_source(db: DBSession):
default=row.guidance or "", 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__ = [ __all__ = [
"FAMILIES", "FAMILIES",
@@ -47,6 +47,10 @@
{{ icon("link", "icon--sm") }} {{ icon("link", "icon--sm") }}
<span class="nav-item__label">Tools</span> <span class="nav-item__label">Tools</span>
</a> </a>
<a class="nav-item {{ 'is-active' if section == 'mcp' }}" href="/admin/mcp">
{{ icon("server", "icon--sm") }}
<span class="nav-item__label">MCP servers</span>
</a>
<a class="nav-item {{ 'is-active' if section == 'prompts' }}" href="/admin/prompts"> <a class="nav-item {{ 'is-active' if section == 'prompts' }}" href="/admin/prompts">
{{ icon("sparkle", "icon--sm") }} {{ icon("sparkle", "icon--sm") }}
<span class="nav-item__label">Prompts</span> <span class="nav-item__label">Prompts</span>
@@ -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.
#}
<div class="model-row {{ 'is-off' if not server.enabled }}" id="mcp-{{ server.id }}">
<div class="model-row__main">
<div class="model-row__title">
<a class="model-row__name" href="/admin/mcp/{{ server.id }}/edit">{{ server.name }}</a>
<span class="badge">{{ tool_count }} tool{{ '' if tool_count == 1 else 's' }}</span>
{% if not server.enabled %}<span class="badge badge--danger">disabled</span>{% endif %}
{% if not server.public %}<span class="badge">restricted</span>{% endif %}
{% if server.allow_private %}<span class="badge">private network</span>{% endif %}
{% if server.protocol_version %}
<span class="badge badge--leaf">MCP {{ server.protocol_version }}</span>
{% endif %}
</div>
<code class="model-row__id">{{ server.slug }} · {{ server.url }}</code>
{% if message %}
<p class="text-xs {{ 'danger' if message_kind == 'error' else 'faint' }}">{{ message }}</p>
{% elif server.last_error %}
<p class="text-xs danger">{{ server.last_error }}</p>
{% endif %}
</div>
<div class="model-row__actions">
<button class="btn btn--sm" type="button"
hx-post="/admin/mcp/{{ server.id }}/test"
hx-target="#mcp-{{ server.id }}" hx-swap="outerHTML">
{{ icon("refresh", "icon--sm") }} Test &amp; refresh
</button>
<a class="btn btn--sm" href="/admin/mcp/{{ server.id }}/edit">Edit</a>
</div>
</div>
+53
View File
@@ -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 %}
<p class="admin-lede">
Remote servers speaking the Model Context Protocol over HTTP. Their tools are
offered beside the built-in ones to models marked <strong>MCP servers</strong>.
The list of tools is discovered and cached — press <strong>Test &amp;
refresh</strong> after adding one, and again whenever the server changes.
</p>
<div class="alert">
{{ icon("shield", "icon--sm") }}
<span>
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.
</span>
</div>
{% if saved %}
<div class="alert alert--success">{{ icon("check", "icon--sm") }} <span>{{ saved }}</span></div>
{% endif %}
<div class="btn-row">
<a class="btn btn--primary" href="/admin/mcp/new">
{{ icon("plus", "icon--sm") }} Add a server
</a>
</div>
{% if not servers %}
<div class="empty" style="padding: var(--sp-10) 0">
{{ icon("server", "empty__mark") }}
<p class="empty__text">
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.
</p>
</div>
{% else %}
<div class="model-rows">
{% for server in servers %}
{% with tool_count = counts[server.id] %}
{% include "admin/_mcp_row.html" %}
{% endwith %}
{% endfor %}
</div>
{% endif %}
{% endblock %}
@@ -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 %}
<nav class="crumbs">
<a class="crumbs__back" href="/admin/mcp">
{{ icon("chevron-right", "icon--sm crumbs__icon") }} All servers
</a>
</nav>
{% if error %}
<div class="alert alert--error">{{ icon("warning", "icon--sm") }} <span>{{ error }}</span></div>
{% endif %}
{% if server.last_error %}
<div class="alert alert--error">
{{ icon("warning", "icon--sm") }}
<span>Last contacted unsuccessfully: {{ server.last_error }}</span>
</div>
{% endif %}
<form method="post" action="{{ '/admin/mcp' if is_new else '/admin/mcp/' ~ server.id }}"
class="form-grid">
<section class="card">
<h2 class="card__title">The server</h2>
<div class="field">
<label class="field__label" for="name">Name</label>
<input class="input" id="name" name="name" value="{{ server.name }}" required
maxlength="120" placeholder="GitHub">
</div>
<div class="field">
<label class="field__label" for="slug">Identifier</label>
<input class="input input--mono" id="slug" name="slug" value="{{ server.slug }}" required
maxlength="24" pattern="[a-z0-9][a-z0-9_\-]*" placeholder="github">
<p class="field__hint">
Prefixed onto every tool name this server offers, so that two servers
both exposing <code>search</code> do not collide.
</p>
</div>
<div class="field">
<label class="field__label" for="url">Endpoint URL</label>
<input class="input input--mono" id="url" name="url" value="{{ server.url }}" required
placeholder="https://mcp.example.com/mcp">
<p class="field__hint">
The streamable-HTTP endpoint itself, the one that accepts a POST. A
server that answers with a redirect to somewhere else will be refused.
</p>
</div>
<div class="field">
<label class="field__label" for="headers">Extra headers</label>
<textarea class="textarea input--mono" id="headers" name="headers" rows="3"
spellcheck="false">{{ headers_text }}</textarea>
<p class="field__hint">One <code>Name: value</code> per line.</p>
</div>
<div class="field">
<label class="field__label" for="timeout">Timeout (seconds)</label>
<input class="input" id="timeout" name="timeout" value="{{ server.timeout }}"
inputmode="numeric">
</div>
<div class="field">
<label class="field__label" for="max_chars">Most characters to keep per call</label>
<input class="input" id="max_chars" name="max_chars" value="{{ server.max_chars }}"
inputmode="numeric">
</div>
</section>
<section class="card">
<h2 class="card__title">Credential</h2>
<div class="field">
<label class="field__label" for="secret_placement">How it is sent</label>
<select class="select" id="secret_placement" name="secret_placement">
{% for value, label in secret_placements %}
<option value="{{ value }}" {{ 'selected' if value == server.secret_placement }}>
{{ label }}
</option>
{% endfor %}
</select>
</div>
<div class="field">
<label class="field__label" for="secret_name">Header or parameter name</label>
<input class="input input--mono" id="secret_name" name="secret_name"
value="{{ server.secret_name }}" maxlength="120">
</div>
<div class="field">
<label class="field__label" for="secret">Secret</label>
<input class="input input--mono" id="secret" name="secret" type="password"
autocomplete="off" placeholder="No secret set"
value="{{ unchanged if server.secret_encrypted else '' }}">
<p class="field__hint">
{% if server.secret_encrypted %}
Currently <code>{{ masked }}</code>. Leave the dots alone to keep it,
or clear the field to remove it.
{% else %}
Encrypted at rest and never shown again.
{% endif %}
</p>
</div>
</section>
{% if tools %}
<section class="card">
<h2 class="card__title">Tools it offers</h2>
<p class="field__hint">
Discovered at the last refresh. Untick one to withhold it — a tool this
server adds later is offered by default.
</p>
<input type="hidden" name="tool_choices" value="1">
<div class="checkbox-row checkbox-row--stacked">
{% for tool in tools %}
<label class="checkbox">
<input type="hidden" name="tool_names" value="{{ tool.name }}">
<input type="checkbox" name="tool_names_on" value="{{ tool.name }}"
{{ 'checked' if tool.on }}>
<span>
<code>{{ tool.offer_name or tool.name }}</code>
{% if tool.offer_name and tool.offer_name != tool.name %}
<span class="faint text-xs">({{ tool.name }} on the server)</span>
{% endif %}
{% if tool.description %}
<br><span class="faint text-xs">{{ tool.description }}</span>
{% endif %}
</span>
</label>
{% endfor %}
</div>
</section>
{% elif not is_new %}
<section class="card">
<h2 class="card__title">Tools it offers</h2>
<p class="field__hint">
Nothing discovered yet. Save, then press <strong>Test &amp; refresh</strong>
on the <a href="/admin/mcp">list</a>.
</p>
</section>
{% endif %}
<section class="card">
<h2 class="card__title">Guidance</h2>
<div class="field">
<textarea class="textarea" name="guidance" rows="4"
placeholder="- Use the GitHub tools for anything about our repositories."
>{{ server.guidance }}</textarea>
<p class="field__hint">
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 %}
<br><strong>Someone has overridden this wording under
<a href="/admin/prompts">Prompts</a></strong> — that is what the model
sees, not this.
{% endif %}
</p>
</div>
</section>
<section class="card">
<h2 class="card__title">Availability</h2>
<div class="field">
<div class="checkbox-row">
<label class="checkbox">
<input type="checkbox" name="enabled" value="true" {{ 'checked' if server.enabled }}>
<span>Enabled — offered in chats</span>
</label>
<label class="checkbox">
<input type="checkbox" name="allow_private" value="true"
{{ 'checked' if server.allow_private }}>
<span>May reach private and loopback addresses</span>
</label>
</div>
<p class="field__hint">
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.
</p>
</div>
<div class="field">
<label class="checkbox">
<input type="checkbox" name="public" value="true" {{ 'checked' if server.public }}>
<span>Available to everyone</span>
</label>
</div>
<div class="field">
<span class="field__label">Groups with access</span>
{% if groups %}
<div class="checkbox-row">
{% for group in groups %}
<label class="checkbox">
<input type="checkbox" name="group_ids" value="{{ group.id }}"
{{ 'checked' if group.id in selected_groups }}>
<span>{{ group.name }}</span>
</label>
{% endfor %}
</div>
<p class="field__hint">Ignored while the server is available to everyone.</p>
{% else %}
<p class="field__hint">
No groups yet — <a href="/admin/groups">create one</a> to restrict access.
</p>
{% endif %}
</div>
<div class="field">
<label class="field__label" for="position">Position</label>
<input class="input" id="position" name="position" value="{{ server.position }}"
inputmode="numeric">
</div>
</section>
<div class="btn-row">
<button class="btn btn--primary" type="submit">
{{ "Add server" if is_new else "Save changes" }}
</button>
<a class="btn btn--ghost" href="/admin/mcp">Back to all servers</a>
{% if not is_new %}
<button class="btn btn--danger" type="submit" formnovalidate
formaction="/admin/mcp/{{ server.id }}/delete"
data-confirm-button="Delete the server “{{ server.name }}”? Chats that used its tools keep their transcripts.">
Delete
</button>
{% endif %}
</div>
</form>
{% endblock %}
+193
View File
@@ -301,3 +301,196 @@ def test_the_secret_survives_a_round_trip_through_the_form(client: TestClient, d
) )
db.refresh(tool) db.refresh(tool)
assert decrypt(tool.secret_encrypted) == "s3cret" 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
+123
View File
@@ -207,3 +207,126 @@ async def _empty_search(_config, _query, *, limit=None):
async def _never_called_title(*_args, **_kwargs): async def _never_called_title(*_args, **_kwargs):
"""Auto-titling makes its own request; these tests are about the tool loop.""" """Auto-titling makes its own request; these tests are about the tool loop."""
return "A title" 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"
+506
View File
@@ -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