Jaroslav Beneš
ecadb66414
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>
507 lines
18 KiB
Python
507 lines
18 KiB
Python
"""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
|