retoor/tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

MCP both directions: client bridge plus --serve-mcp (1.42.0)

Browse Source 
- client: stdio + Streamable HTTP transports, 2025-11-25
  handshake/list/call/ping, vault mcp_servers with sealed headers,
  per-profile enable, description-hash pins, tool cache
- mcp_normalize: pasted JSON/TOML configs (type repair, ${VAR}
  expansion from env then vault secrets, allow/deny lists);
  github URLs resolve via registry/package/README with live
  handshake verification before storing
- bridge: enabled tools enter lazy select_tools as
  mcp_<server>_<tool>, dispatch fallback, approval like non-safe
  shell, drift re-prompt, 64KB spill, per-call audit
- 7 agent tools: mcp_add/list/enable/disable/remove/tools/call
  (native count 42->60 in README, was already stale at 53)
- server: --serve-mcp KEY (+ --mcp-port/--mcp-bind) serves all 60
  native tools plus chat (agent turn, per-session history) and
  completions (stateless, current model) over single-threaded
  HTTP with Bearer auth, sessions, idle expiry; yolo documented
- rotation + plaintext-boot refusal cover sealed headers
- mcp.md contract with drift tests; 47 new tests; live loopback
  verified (client->serve over real HTTP)
This commit is contained in:
retoor
2026-10-18 19:46:55 +02:00
parent 251c3dc2ba
commit 66ed0f9ea6
6 changed files with 2197 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -14,6 +14,7 @@
!cap.md
!skills.md
!settings.md
!mcp.md
# Backup files
*.bak
AGENTS.md
+1
View File
@@ -11,6 +11,7 @@
- The terminal multiplexer (`Mplex*` classes, `tai.py --mplex`, `~/.tai/mplex.toml`) is specified in [mplex.md](mplex.md): read it before touching any `Mplex*` code, and keep its option, command and key tables in step with the code (`MplexSpecTests` fails when they drift).
- Interactive UI flows (`plexer_show`, `PlexerFlow`, `PlexerController.ask`, the `show`/`result`/`cancel` control commands) are specified in [uiflow.md](uiflow.md): read it before touching any flow code, and keep its contract and cap tables in step with the code (`UiflowDocTests` fails when they drift).
- CAP protocol support (`Cap*` classes, `--cap-manifest`, `--cap-speak`, orchestrator drivers) is specified in [cap.md](cap.md): read it before touching any `Cap*` code, and keep its event, error and manifest tables in step with the code (`CapDocTests` fails when they drift).
- MCP support (`Mcp*` classes, `mcp_*` tools, `--serve-mcp`) is specified in [mcp.md](mcp.md): read it before touching any `Mcp*` code, and keep its wire, caps, and tool tables in step with the code (`McpDocTests` fails when they drift).
- Skills (`discover_skills`, `SKILL_BLUEPRINTS`, `BUNDLED_SKILLS`, the `skills` install target) are specified in [skills.md](skills.md): read it before touching skill code or embedded skill text, and keep its bundled-skill and adaptation tables in step with the code (`BundledSkillsTests` fails when they drift).
- Settings (`SETTINGS_CATALOG`, the `settings` table, `get_or_create_settings`/`upsert_settings`/`get_all_settings`, `boot_settings`) are specified in [settings.md](settings.md): read it before touching settings code, and keep its catalog table in step with the code (`SettingsTests` fails when they drift).
README.md
+11 -2
View File
@@ -4,7 +4,7 @@ retoor <retoor@molodetz.nl>
**Zero dependencies. A real agent.**
`tai` is an autonomous AI agent written in Python — the entire product lives in the standalone `tai.py` file (grouped into classes) and imports **only the Python standard library**. No `pip install`, no `requirements.txt`, no build step, no lockfile to rot. It runs as an interactive REPL or a one-shot command, reasons through an OpenAI-compatible backend, acts through 42 tools, keeps per-profile memory in a self-healing SQLite vault, seals that vault at rest, ranks every search with real BM25 (no vector database, no embeddings bill, no GPU), chunks and indexes anything worth remembering for retrieval, and can wall shell/file work off inside a disposable container sandbox.
`tai` is an autonomous AI agent written in Python — the entire product lives in the standalone `tai.py` file (grouped into classes) and imports **only the Python standard library**. No `pip install`, no `requirements.txt`, no build step, no lockfile to rot. It runs as an interactive REPL or a one-shot command, reasons through an OpenAI-compatible backend, acts through 60 tools, keeps per-profile memory in a self-healing SQLite vault, seals that vault at rest, ranks every search with real BM25 (no vector database, no embeddings bill, no GPU), chunks and indexes anything worth remembering for retrieval, and can wall shell/file work off inside a disposable container sandbox.
Nothing here is decorative. Every mechanism below exists because something specific was expensive — context tokens, API calls, attacker surface, disk writes, or your attention — and got engineered out.
@@ -153,7 +153,7 @@ This only ever activates as a genuine last resort: it's wired into `all_candidat
Every tool schema you hand an LLM costs real tokens on every single turn, whether or not it gets used, and a model drowning in fifty tool definitions picks worse than one offered the five that actually matter. tai treats this as a first-class engineering problem, not an afterthought, with three independent layers stacked on top of each other:
**1. Core vs. lazy.** Of the 42 tools, only `CORE_TOOLS` — `shell`, `read_file`, `write_file`, `remember`, `recall`, `load_skill`, `fork`, `poll` — are sent to the model on every turn. The other 34 are *lazy*: their full JSON schema only enters the request payload when they're judged relevant to what's actually being discussed.
**1. Core vs. lazy.** Of the 60 tools, only `CORE_TOOLS` — `shell`, `read_file`, `write_file`, `remember`, `recall`, `load_skill`, `fork`, `poll` — are sent to the model on every turn. The other 51 are *lazy*: their full JSON schema only enters the request payload when they're judged relevant to what's actually being discussed.
**2. Relevance by three independent signals, not a fixed list.** `select_tools()` scans the last 6000 characters of conversation (user, assistant, and tool-result turns — `conversation_text()`) and pulls in a lazy tool's full schema the moment any one of three things matches:
@@ -272,6 +272,14 @@ tai speaks the [CLI Agent Protocol](https://cap-protocol.org/) (draft v1 plus th
**Honest gaps, stated up front.** gRPC is parsed in manifests but not driven (the standard library has no HTTP/2 and no protobuf, so there is no stdlib-only client to write); A2A is driven but not yet served; ACP server argv travels out-of-band in the fleet file because manifests carry none; the reverse-RPC JSON shapes on the speak path are tai-defined, since stream-json specifies no serving direction. TOML manifest and fleet parsing needs Python 3.11+ (`tomllib`); older interpreters get a clean error, not a crash.
## MCP: using servers, being one
tai speaks MCP in both directions, with no SDK: newline-delimited JSON-RPC over a spawned stdio child, or Streamable HTTP via `urllib` with a hand-rolled SSE parse, on the 2025-11-25 baseline. Full contract: [mcp.md](mcp.md); research: `research/mcp-client-support/`.
**As a client.** Paste any server definition into `mcp_add`: a JSON `mcpServers` block (URL form with headers, or command form with argv), a TOML `[mcp_servers.*]` table, or a bare GitHub URL. Configs normalize deterministically (`type` repair, `${VAR}`/`${VAR:-default}` expansion from env then vault secrets at use time, `enabled_tools` allow/deny lists); GitHub URLs resolve through registry, package files, and README run blocks, and every resolution is verified by a live handshake before anything is stored. Connections live in the vault (`mcp_servers`, headers sealed), enablement is per profile and defaults to off, and each server's tools enter the lazy system as `mcp_<server>_<tool>`. Servers are untrusted by design: registration shows the exact spawn command, tool descriptions are hash-pinned with re-approval on drift, calls ask like non-safe shell commands, results over 64KB spill to a record, and every call is audited.
**As a server.** `./tai.py --serve-mcp KEY` exposes all 60 native tools plus `chat` (a full agent turn on a per-session agent, history kept) and `completions` (one stateless completion over a messages array, current model, no tools) over Streamable HTTP on `127.0.0.1:8765` (`--mcp-port`, `--mcp-bind`, `TAI_MCP_PORT`/`TAI_MCP_BIND` to move it). Clients send the key as a Bearer token; miss it and you get 401. The key is the authorization: serve mode runs yolo because no user is present to ask, which is stated loudly at startup. Bridged foreign tools are never re-served.
## Execution environment / sandbox
`/env` shows the execution environment, `/env sandbox` switches shell and file tools into an isolated `tai-box` container (podman or docker, no mounts, no shared filesystem), `/env home` switches back. The box is resource-capped at 2GB RAM, 2 CPUs, 256 processes (`BOX_MEMORY`/`BOX_CPUS`/`BOX_PIDS_LIMIT`), because an uncapped sandbox is just a slower way to fork-bomb the host; pre-existing boxes get memory and CPU applied live (`podman update`), the process cap lands on recreation. The image builds automatically on first use from an embedded Containerfile; every pip requirement the *agent's own features* need (`faster-whisper`, `edge-tts`) lives inside that image, while `tai.py` itself stays dependency-free — the sandbox is where "needs a real package" problems go, never the host process. Sandbox commands need no approval because the container is disposable; this is deliberately a *different axis* from the `--yes`/`--yolo`/`--auto` approval settings, not a replacement for them.
@@ -421,6 +429,7 @@ Unit tests live at the top level (`test_tai.py`, `test_seal.py`); integration te
mplex.md the native multiplexer (`--mplex`): options, commands, keys
uiflow.md interactive UI flows (`plexer_show`): contract and caps
cap.md the CAP protocol binding (`--cap-*`, `--cap`): wire contract
mcp.md MCP client + server (`mcp_*` tools, `--serve-mcp`): wire contract
skills.md skills system and bundled set: provenance, adaptations, install
settings.md vault-backed settings: catalog, precedence, API
research/ archived design research, registry at research/INDEX.md
mcp.md
+202
View File
@@ -0,0 +1,202 @@
# mcp.md
tai speaks MCP in both directions. As a **client** it spawns or dials
third-party MCP servers, bridges their tools into the lazy tool system,
and stores connections in the vault with per-profile enablement. As a
**server** (`tai.py --serve-mcp KEY`) it exposes every native tool plus
`chat` and `completions` over Streamable HTTP. `McpDocTests` fails when
this file drifts from the code. Research: `research/mcp-client-support/`.
## Shared wire subset
Protocol version is `2025-11-25`, sent as `MCP_PROTOCOL_VERSION`.
Client name and server name are both `tai`; server version is `VERSION`.
| Direction | Method | Shape |
|---|---|---|
| C2S | `initialize` | `{protocolVersion, capabilities: {}, clientInfo: {name, version}}`; server answers its version plus `capabilities.tools`, or the client aborts |
| C2S | `notifications/initialized` | no params, no answer |
| both | `ping` | `{}` in, `{}` out |
| C2S | `tools/list` | `{cursor?}` in, `{tools: [...], nextCursor?}` out, followed until no cursor |
| C2S | `tools/call` | `{name, arguments?}` in, `{content: [...], isError?}` out |
| C2S | `notifications/cancelled` | `{requestId, reason?}` on call timeout, then the call fails locally |
JSON-RPC errors (`-32601` unknown, `-32602` bad params, `-32603`
internal) mean protocol failure; a tool that ran and failed arrives
as a normal result with `isError: true`, and tai preserves that split
end to end. Roots, sampling, elicitation, logging display, progress
tokens, and the legacy HTTP+SSE transport are not implemented on
either side; anything asking for them gets `-32601`.
## Caps
| Cap | Value | Meaning |
|---|---|---|
| `MCP_CONNECT_TIMEOUT` | 20s | spawn plus handshake, per server |
| `MCP_CALL_TIMEOUT` | 60s | one `tools/call`, then cancel notification and local failure |
| `MCP_RESULT_MAX` | 64KB | tool result text kept inline; larger spills to a vault record and the model gets the id, never a silent truncation |
| `MCP_HTTP_MAX_BODY` | 10MB | largest request/response body on either HTTP leg |
| `MCP_MAX_SERVERS` | 32 | registered servers; the 33rd registration is refused |
| `MCP_MAX_TOOLS` | 100 | tools bridged per server; extras are listed but not offered |
| `MCP_SESSION_IDLE` | 3600s | idle server-side MCP sessions are dropped |
| `MCP_CHAT_TIMEOUT` | 600s | deadline per served `chat` turn |
## Client
### Transports
stdio: the server command spawns via subprocess with piped stdio,
newline-delimited JSON-RPC, stderr drained to a capped buffer (never
parsed). Shutdown order is close-stdin, wait, terminate, kill. The
exact argv is shown at registration consent and re-shown whenever it
changes. stdio reads need POSIX `select`; the HTTP leg covers the
rest.
Streamable HTTP: one `POST <url>` per request with
`Accept: application/json, text/event-stream`, the stored headers,
and `Mcp-Session-Id` after `initialize`. The reply is a single JSON
message or a `text/event-stream` body that the client parses by hand
(`data:` lines, blank-line frames, `event: message` only). `GET` and
`DELETE` are never sent; sessions end by dropping the id.
### Vault tables
```sql
mcp_servers(name PK, transport, command, args_json, env_json, url,
headers_json, cwd, allow_json, deny_json, source, created, updated)
mcp_server_profiles(server, profile, enabled) PK(server, profile)
mcp_tool_pins(server, tool, hash, updated) PK(server, tool)
mcp_server_tools(server, tool, description, schema_json, updated) PK(server, tool)
```
`headers_json` is sealed at rest whenever the vault seal is enabled.
`source` records what was pasted (`config-json`, `config-toml`,
`github:<owner>/<repo>`, `explicit`). A server with no profile row
defaults to disabled everywhere; enabling is always explicit and per
profile.
### Normalizer (`mcp_normalize`)
Accepted input: a JSON object with a `mcpServers` or `servers`
wrapper (or one bare entry), a TOML document with `[mcp_servers.*]`
tables (3.11+, clean error below), or a bare server name already
stored. Repair rules: an entry with `url` and no `type` gains
`type: http`; command entries keep `command`/`args`/`env`/`cwd`.
`${VAR}` and `${VAR:-default}` expand in `command`, `args`, `env`,
`url`, and `headers` at use time from process env first, then vault
secrets; Cursor `${env:NAME}` rewrites to `${NAME}`; VS Code
`${input:id}` resolves against vault secrets by id and fails the
spawn when missing. Unresolved names warn and keep the literal for
display, but the spawn is refused. `npx` on Windows wraps in
`cmd /c`; `uvx` never does. TOML-only keys map as: `env`/`env_vars`
to env, `url` plus `http_headers`/`env_http_headers` to the HTTP
leg, `enabled` to the registering profile, `enabled_tools`/
`disabled_tools` to a tool allow/deny stored beside the server row.
`timeout`/`startup_timeout` are ignored with a note: tai's connect/call
caps govern.
GitHub URLs (`github.com/<owner>/<repo>` plus anything after) are
not normalizable; they go to the resolver: registry API by repo
path, `server.json`/`mcpName`, `package.json` binaries, README run
block, Glama page, then a live `initialize` + `tools/list` probe.
The resolver is tai itself (the agent loop, usually one
`mcp_add` call after reading): it returns exactly one normalized
entry or a loud failure with the evidence gathered. Every resolved
entry is verified by a live handshake before it is stored.
### Bridge
Enabled servers' tools enter the lazy tool system as
`mcp_<server>_<tool>`, sanitized to `[a-z0-9_]` (anything else
becomes `_`, empties gain a numeric suffix). `select_tools` also
iterates the bridge schemas, so description-word matching works
unchanged; `Tools.dispatch` falls back to the bridge for `mcp_*`
names. Schemas pass the server's `inputSchema` through as
`parameters`, defaulting to `{type: object, properties: {}}`.
Calls ask approval exactly like non-safe shell commands (always in
interactive mode, auto in `--auto`, silent in `--yolo`), and the
consent screen names server, tool, and argument digest.
Trust: tool names, descriptions, and annotations are untrusted text.
At registration tai hashes name plus description plus schema per
tool into `mcp_tool_pins`; added or redefined tools on a later
`tools/list` re-prompt for approval and are withheld until granted,
while pure removals apply silently.
Results over `MCP_RESULT_MAX` spill to a vault record. Every call
lands in the audit trail (`mcp-call`, `server/tool`, argument
digest, outcome, duration, size). `${VAR}` secret values resolve at
spawn/call time and `Store.redact` scrubs known values from results.
### Agent tools
| Tool | Args | Effect |
|---|---|---|
| `mcp_add` | `source`, `name?`, `profile?` | normalize or resolve, consent the spawn, live-verify, store, return the tool list |
| `mcp_list` | - | servers with transport, enabled profiles, tool counts, pin state |
| `mcp_enable` | `server`, `profile?` | enable for a profile (current when omitted) |
| `mcp_disable` | `server`, `profile?` | disable for a profile |
| `mcp_remove` | `server` | delete server, profile rows, and pins; asks first |
| `mcp_tools` | `server?` | list bridged tools with descriptions and pin state |
| `mcp_call` | `server`, `tool`, `args?` | call one tool explicitly (same path as bridged dispatch) |
## Server (`tai.py --serve-mcp KEY`)
One process, one `Config`, one `Store`, single-threaded `HTTPServer`
(no thread pool: sequential requests fit MCP clients and the
no-internal-threads rule). Flags: `--serve-mcp KEY` (required; the
Bearer token clients must send), `--mcp-port` (default from
`TAI_MCP_PORT` else `MCP_DEFAULT_PORT`, 8765), `--mcp-bind`
(default `TAI_MCP_BIND` else `127.0.0.1`).
| Endpoint | Behavior |
|---|---|
| `POST /mcp` | one JSON-RPC message; missing/invalid `Authorization: Bearer KEY` gets 401. Answers `initialize` (new `Mcp-Session-Id`), `notifications/initialized` (none), `ping`, `tools/list` (no pagination, all tools fit), `tools/call`. Always a single JSON reply, never SSE. Unknown methods get `-32601`. |
| `GET /mcp` | 405 (no server-initiated streams in v1) |
| `DELETE /mcp` | drops the `Mcp-Session-Id` session, 200 |
Sessions live in memory: `{agent, profile, created, last_seen}`.
Idle sessions past `MCP_SESSION_IDLE` are dropped on next access.
### Served tools
Every native tool schema is served verbatim (name, description,
`parameters` as `inputSchema`), plus two extras, 62 total:
| Tool | Args | Effect |
|---|---|---|
| `chat` | `message`, `session?`, `profile?`, `max_steps?` | full agent turn on the session's own `Agent` (`run_turn`, capture), history kept in the session; uses the current model with tools |
| `completions` | `messages[]` | one stateless `ChatClient.complete` with the current model, no tools; roles validated (`system`/`user`/`assistant`) |
The key is the authorization: serve mode runs the equivalent of
`--yolo` for tool approvals, documented loudly at startup, because
no user is present to ask. `ask_user` over the server always denies
(no prompter, no tty). Served `chat` agents can themselves use
bridged MCP tools (they are full agents), but bridged foreign tools
are never re-served as MCP tools: transitive serving would launder
untrusted servers through tai's identity.
### Exclusions (both sides)
OAuth browser flows, dynamic client registration, token storage and
refresh; roots/sampling/elicitation serving; legacy HTTP+SSE;
`GET /mcp` streams; file-attachment content beyond text passthrough
(images/audio arrive as their JSON, unchanged); batch JSON-RPC
arrays (one message per POST).
## Code map
`MCP_PROTOCOL_VERSION` pins the wire; `mcp_normalize` plus
`mcp_expand_vars` own pasted-config intake; `McpStdioTransport` and
`McpHttpTransport` own the two legs behind `McpClient`;
`McpBridge` owns enablement, drift pins, approval, caps, spill, and
audit; `McpServeContext` plus `mcp_serve_handler` own the HTTP
server; `mcp_serve_main` owns flags and boot.
## Test hooks
`McpDocTests` asserts this file names every `MCP_*` cap, every
method in the wire table, every vault table, every agent tool, and
both server extras. Transport tests drive a fake MCP server (a
`python3 -c` NDJSON stub for stdio, `http.server` in-test for the
HTTP leg); serve tests drive the real `POST /mcp` with `urllib`
against `ChatClient` fakes, never a live model.
tai.py
+1267 -3
View File
File diff suppressed because it is too large Load Diff
test_tai.py
+715
View File
@@ -19,6 +19,7 @@ import time
import unicodedata
import unittest
import urllib.error
import urllib.request
from datetime import datetime, timedelta, timezone
from unittest import mock
@@ -16030,5 +16031,719 @@ class CapDocTests(unittest.TestCase):
self.assertNotIn("—", self.doc)
class McpBase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.home = os.path.join(self.tmp.name, "home")
os.makedirs(self.home)
self.env = mock.patch.dict(os.environ, {"HOME": self.home, "TAI_HOME": os.path.join(self.home, ".tai")})
self.env.start()
class FakeArgs:
profile = "t"
yes = True
self.config = tai.Config(FakeArgs())
self.seal = tai.Seal(self.config.home, "mcp-test-1")
self.store = tai.Store(self.config, self.seal)
self.app = FakeSecretApp(self.store)
self.app.profile = "t"
self.app.config = self.config
self.tools = tai.Tools(self.app)
self.app.tools = self.tools
def tearDown(self):
self.store.close()
self.env.stop()
self.tmp.cleanup()
os.environ.pop("TAI_HOME", None)
class McpNormalizeTests(unittest.TestCase):
JINA_URL = '{"mcpServers": {"jina-mcp-server": {"url": "https://mcp.jina.ai/v1", "headers": {"Authorization": "Bearer ${JINA_API_KEY}"}}}}'
JINA_CMD = '{"mcpServers": {"jina-mcp-server": {"command": "npx", "args": ["mcp-remote", "https://mcp.jina.ai/v1", "--header", "Authorization: Bearer ${JINA_API_KEY}", "--ignore-tool", "search_images"]}}}'
JINA_TOML = '[mcp_servers.jina-mcp-server]\ncommand = "npx"\nargs = ["-y", "mcp-remote", "https://mcp.jina.ai/v1", "--header", "Authorization: Bearer ${JINA_API_KEY}"]\n'
def test_json_url_form(self):
entry, notes = tai.mcp_normalize(self.JINA_URL)
self.assertEqual(entry["name"], "jina-mcp-server")
self.assertEqual(entry["transport"], "http")
self.assertEqual(entry["url"], "https://mcp.jina.ai/v1")
self.assertEqual(entry["headers"], {"Authorization": "Bearer ${JINA_API_KEY}"})
self.assertIn("entry has url without type, treating as http", notes)
self.assertEqual(entry["source"], "config-json")
def test_json_command_form_keeps_argv(self):
entry, _notes = tai.mcp_normalize(self.JINA_CMD)
self.assertEqual(entry["transport"], "stdio")
self.assertEqual(entry["command"], "npx")
self.assertEqual(entry["args"][:2], ["mcp-remote", "https://mcp.jina.ai/v1"])
self.assertIn("--ignore-tool", entry["args"])
def test_toml_form(self):
entry, _notes = tai.mcp_normalize(self.JINA_TOML)
self.assertEqual(entry["name"], "jina-mcp-server")
self.assertEqual(entry["transport"], "stdio")
self.assertEqual(entry["args"][0], "-y")
self.assertEqual(entry["source"], "config-toml")
def test_bare_and_multi_pick(self):
entry, _notes = tai.mcp_normalize('{"command": "uvx", "args": ["srv"]}', name="solo")
self.assertEqual(entry["name"], "solo")
multi = '{"mcpServers": {"a": {"command": "x"}, "b": {"command": "y"}}}'
with self.assertRaises(tai.McpError):
tai.mcp_normalize(multi)
entry, _notes = tai.mcp_normalize(multi, name="b")
self.assertEqual(entry["name"], "b")
with self.assertRaises(tai.McpError):
tai.mcp_normalize(multi, name="zzz")
def test_rejects_bad_shapes(self):
with self.assertRaises(tai.McpError):
tai.mcp_normalize('{"mcpServers": {"s": {"command": "x", "url": "http://y"}}}')
with self.assertRaises(tai.McpError):
tai.mcp_normalize('{"mcpServers": {"s": {"command": 42}}}')
with self.assertRaises(tai.McpError):
tai.mcp_normalize('{"mcpServers": {"s": {"url": "gopher://x"}}}')
with self.assertRaises(tai.McpError):
tai.mcp_normalize('{"mcpServers": {"s": {"url": " "}}}')
with self.assertRaises(tai.McpError):
tai.mcp_normalize('{"mcpServers": {}}')
with self.assertRaises(tai.McpError):
tai.mcp_normalize('not json or toml at all {{{')
with self.assertRaises(tai.McpError):
tai.mcp_normalize("")
with self.assertRaises(tai.McpError):
tai.mcp_normalize('{"mcpServers": {"bad name!": {"command": "x"}}}')
def test_toml_extras_map(self):
text = ('[mcp_servers.s]\ncommand = "srv"\nenv_vars = {A = "1"}\n'
'enabled = false\nenabled_tools = ["a"]\ndisabled_tools = ["b"]\n'
'timeout = 5\nstartup_timeout = 9\n')
entry, notes = tai.mcp_normalize(text)
self.assertEqual(entry["env"], {"A": "1"})
self.assertIs(entry["enabled"], False)
self.assertEqual(entry["allow"], ["a"])
self.assertEqual(entry["deny"], ["b"])
self.assertTrue(any("timeout fields ignored" in note for note in notes))
def test_env_http_headers_become_placeholders(self):
entry, _notes = tai.mcp_normalize('{"mcpServers": {"s": {"url": "https://x", "env_http_headers": {"Authorization": "TOK"}}}}')
self.assertEqual(entry["headers"], {"Authorization": "${TOK}"})
def test_placeholder_rewrites(self):
text, noted = tai.mcp_rewrite_placeholders("run ${env:FOO} and ${input:bar}")
self.assertEqual(text, "run ${FOO} and ${bar}")
self.assertEqual(len(noted), 2)
def test_expand_vars(self):
lookup = lambda name: {"A": "1"}.get(name)
text, missing = tai.mcp_expand_vars("${A}/${B}/${C:-d}", lookup)
self.assertEqual(text, "1/${B}/d")
self.assertEqual(missing, ["B"])
def test_github_target(self):
self.assertEqual(tai.mcp_github_target("https://github.com/jina-ai/MCP"), ("jina-ai", "MCP"))
self.assertEqual(tai.mcp_github_target("https://github.com/o/r.git"), ("o", "r"))
self.assertIsNone(tai.mcp_github_target("https://example.com/o/r"))
with self.assertRaises(tai.McpError):
tai.mcp_normalize("https://github.com/jina-ai/MCP")
def test_sanitize(self):
self.assertEqual(tai.mcp_sanitize_name("Jina-MCP.Server!"), "jina_mcp_server")
self.assertEqual(tai.mcp_sanitize_name("!!!"), "server")
self.assertEqual(tai.mcp_bridge_name("Srv", "Do Thing"), "mcp_srv_do_thing")
MCP_STUB = r'''
import sys, json
TOOLS = [{"name": "echo", "description": "echo back input", "inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]}},
{"name": "boom", "description": "always fails", "inputSchema": {"type": "object", "properties": {}}}]
for line in sys.stdin:
try:
msg = json.loads(line)
except ValueError:
continue
mid = msg.get("id")
if mid is None:
continue
method = msg.get("method")
if method == "initialize":
result = {"protocolVersion": "2025-11-25", "capabilities": {"tools": {}}, "serverInfo": {"name": "stub", "version": "1"}}
elif method == "tools/list":
result = {"tools": TOOLS}
elif method == "tools/call":
if msg.get("params", {}).get("name") == "boom":
result = {"content": [{"type": "text", "text": "kaput"}], "isError": True}
else:
result = {"content": [{"type": "text", "text": "echo:" + msg.get("params", {}).get("arguments", {}).get("text", "")}]}
elif method == "ping":
result = {}
else:
sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "nope"}}) + "\n")
sys.stdout.flush()
continue
sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": mid, "result": result}) + "\n")
sys.stdout.flush()
'''
class FakeMcpHttpHandler(http.server.BaseHTTPRequestHandler):
seen_session = ""
mode = "json"
def log_message(self, *args):
pass
def do_POST(self):
length = int(self.headers.get("Content-Length") or 0)
message = json.loads(self.rfile.read(length).decode("utf-8") or "{}")
FakeMcpHttpHandler.seen_session = self.headers.get("Mcp-Session-Id") or ""
mid = message.get("id")
method = message.get("method")
if mid is None:
self.send_response(202)
self.send_header("Content-Length", "0")
self.end_headers()
return
if FakeMcpHttpHandler.mode == "error500":
self.send_response(500)
self.send_header("Content-Length", "0")
self.end_headers()
return
if method == "initialize":
result = {"protocolVersion": "2025-11-25", "capabilities": {"tools": {}}, "serverInfo": {"name": "hstub", "version": "1"}}
elif method == "tools/list":
result = {"tools": [{"name": "echo", "description": "echo back input",
"inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}}}]}
elif method == "tools/call":
result = {"content": [{"type": "text", "text": "hecho"}]}
elif method == "ping":
result = {}
else:
payload = {"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "nope"}}
raw = json.dumps(payload).encode("utf-8")
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(raw)))
self.end_headers()
self.wfile.write(raw)
return
payload = {"jsonrpc": "2.0", "id": mid, "result": result}
raw = json.dumps(payload).encode("utf-8")
self.send_response(200)
if FakeMcpHttpHandler.mode == "sse":
body = b"event: message\ndata: " + raw + b"\n\n"
self.send_header("Content-Type", "text/event-stream")
self.send_header("Content-Length", str(len(body)))
self.send_header("Mcp-Session-Id", "sess-1")
self.end_headers()
self.wfile.write(body)
return
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(raw)))
self.send_header("Mcp-Session-Id", "sess-1")
self.end_headers()
self.wfile.write(raw)
class McpTransportTests(McpBase):
def stub_client(self):
transport = tai.McpStdioTransport().spawn([sys.executable, "-c", MCP_STUB], {}, "")
return tai.McpClient(transport)
def test_stdio_full_flow(self):
client = self.stub_client()
try:
result = client.initialize()
self.assertEqual(client.server_version, "2025-11-25")
self.assertEqual(result["serverInfo"]["name"], "stub")
tools = client.list_tools()
self.assertEqual([item["name"] for item in tools], ["echo", "boom"])
out = client.call_tool("echo", {"text": "hi"})
self.assertEqual(out["content"][0]["text"], "echo:hi")
failed = client.call_tool("boom", {})
self.assertTrue(failed["isError"])
client.ping()
with self.assertRaises(tai.McpJsonRpcError) as caught:
client._request("nope", {})
self.assertEqual(caught.exception.code, -32601)
finally:
proc = client.transport.proc
client.close()
self.assertIsNotNone(proc.poll())
def test_stdio_spawn_failure(self):
with self.assertRaises(tai.McpError):
tai.McpStdioTransport().spawn(["/nonexistent-mcp-binary-xyz"], {}, "")
def http_client(self, mode="json", count=2):
FakeMcpHttpHandler.mode = mode
FakeMcpHttpHandler.seen_session = ""
server = http.server.HTTPServer(("127.0.0.1", 0), FakeMcpHttpHandler)
def serve():
for _hit in range(count):
server.handle_request()
thread = threading.Thread(target=serve)
thread.daemon = True
self.addCleanup(server.server_close)
return server, thread
def test_http_json_flow(self):
server, thread = self.http_client()
thread.start()
url = "http://127.0.0.1:%d/" % server.server_port
transport = tai.McpHttpTransport(url)
client = tai.McpClient(transport)
client.initialize()
self.assertEqual(transport.session, "sess-1")
thread.join(timeout=10)
server2, thread2 = self.http_client(count=1)
thread2.start()
transport2 = tai.McpHttpTransport("http://127.0.0.1:%d/" % server2.server_port)
transport2.session = "sess-1"
tai.McpClient(transport2).ping()
thread2.join(timeout=10)
self.assertEqual(FakeMcpHttpHandler.seen_session, "sess-1")
def test_http_sse_flow(self):
server, thread = self.http_client(mode="sse")
thread.start()
client = tai.McpClient(tai.McpHttpTransport("http://127.0.0.1:%d/" % server.server_port))
try:
client.initialize()
self.assertEqual(client.server_version, "2025-11-25")
finally:
client.close()
thread.join(timeout=10)
def test_http_error(self):
server, thread = self.http_client(mode="error500", count=1)
thread.start()
client = tai.McpClient(tai.McpHttpTransport("http://127.0.0.1:%d/" % server.server_port))
try:
with self.assertRaises(tai.McpError):
client.initialize()
finally:
client.close()
thread.join(timeout=10)
def test_sse_parse(self):
payload = tai.McpHttpTransport._parse_sse('event: message\ndata: {"a": 1}\n\n: comment\n\ndata: {"b": 2}\n\n')
self.assertEqual(payload, {"b": 2})
with self.assertRaises(tai.McpError):
tai.McpHttpTransport._parse_sse(": nothing here\n\n")
class McpVaultTests(McpBase):
ENTRY = {"name": "srv", "transport": "http", "command": "", "args": [], "env": {},
"url": "https://x.test/mcp", "headers": {"Authorization": "Bearer s3cret"}, "cwd": "",
"allow": [], "deny": [], "source": "config-json"}
def test_roundtrip_and_sealed_headers(self):
self.store.mcp_save_server(dict(self.ENTRY))
row = self.store.mcp_get_server("srv")
self.assertEqual(row["url"], "https://x.test/mcp")
self.assertEqual(row["headers"], {"Authorization": "Bearer s3cret"})
raw = self.store.db.execute("SELECT headers_json FROM mcp_servers WHERE name = 'srv'").fetchone()[0]
self.assertTrue(raw.startswith(tai.Seal.PREFIX))
self.assertNotIn("s3cret", raw)
self.assertEqual([item["name"] for item in self.store.mcp_list_servers()], ["srv"])
self.store.mcp_remove_server("srv")
self.assertIsNone(self.store.mcp_get_server("srv"))
self.assertEqual(self.store.mcp_list_servers(), [])
def test_enable_defaults_disabled(self):
self.store.mcp_save_server(dict(self.ENTRY))
self.assertFalse(self.store.mcp_server_enabled("srv", "t"))
self.assertEqual(self.store.mcp_enabled_servers("t"), [])
self.store.mcp_set_enabled("srv", "t", True)
self.assertTrue(self.store.mcp_server_enabled("srv", "t"))
self.assertFalse(self.store.mcp_server_enabled("srv", "other"))
self.assertEqual([row["name"] for row in self.store.mcp_enabled_servers("t")], ["srv"])
self.store.mcp_set_enabled("srv", "t", False)
self.assertFalse(self.store.mcp_server_enabled("srv", "t"))
def test_tools_cache_and_pins(self):
self.store.mcp_save_server(dict(self.ENTRY))
tools = [{"name": "b", "description": "bee", "inputSchema": {"type": "object", "properties": {}}},
{"name": "a", "description": "", "inputSchema": {}}]
self.store.mcp_store_tools("srv", tools)
cached = self.store.mcp_cached_tools("srv")
self.assertEqual([item["tool"] for item in cached], ["a", "b"])
pins = self.store.mcp_tool_pins("srv")
self.assertEqual(set(pins), {"a", "b"})
self.assertEqual(pins["b"], tai.mcp_tool_hash("b", "bee", {"type": "object", "properties": {}}))
self.store.mcp_store_tools("srv", tools[:1])
self.assertEqual([item["tool"] for item in self.store.mcp_cached_tools("srv")], ["b"])
def test_bad_name_and_cap(self):
with self.assertRaises(ValueError):
self.store.mcp_save_server(dict(self.ENTRY, name="bad name!"))
for pos in range(tai.MCP_MAX_SERVERS):
self.store.mcp_save_server(dict(self.ENTRY, name="s%d" % pos, headers={}))
with self.assertRaises(ValueError):
self.store.mcp_save_server(dict(self.ENTRY, name="overflow", headers={}))
class FakeMcpClient(object):
def __init__(self, tools, results=None, fail=None):
self.tools = tools
self.results = results or {}
self.fail = fail
self.closed = False
def list_tools(self):
if self.fail == "list":
raise tai.McpError("list down")
return self.tools
def call_tool(self, name, args):
if self.fail == "call":
raise tai.McpError("call down")
return self.results.get(name, {"content": [{"type": "text", "text": "ok:%s" % name}]})
def close(self):
self.closed = True
MCP_FAKE_TOOLS = [{"name": "echo", "description": "echo back input",
"inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]}}]
class McpBridgeTests(McpBase):
def seed(self, tools=None, enabled=True, **over):
entry = dict(McpVaultTests.ENTRY, name="srv")
entry.update(over)
self.store.mcp_save_server(entry)
self.store.mcp_set_enabled("srv", "t", enabled)
self.store.mcp_store_tools("srv", tools if tools is not None else MCP_FAKE_TOOLS)
return entry
def test_schemas_for_enabled_only(self):
self.seed()
schemas = tai.McpBridge(self.app).schemas_for("t")
self.assertEqual(len(schemas), 1)
schema = schemas[0]["function"]
self.assertEqual(schema["name"], "mcp_srv_echo")
self.assertIn("MCP srv/echo", schema["description"])
self.assertEqual(schema["parameters"]["required"], ["text"])
self.assertEqual(tai.McpBridge(self.app).schemas_for("other"), [])
def test_allow_deny_and_cap(self):
many = [{"name": "tool%d" % pos, "description": "d", "inputSchema": {}} for pos in range(tai.MCP_MAX_TOOLS + 5)]
self.seed(tools=many, deny=["tool0"])
schemas = tai.McpBridge(self.app).schemas_for("t")
self.assertEqual(len(schemas), tai.MCP_MAX_TOOLS)
names = {schema["function"]["name"] for schema in schemas}
self.assertNotIn("mcp_srv_tool0", names)
self.seed(tools=many, allow=["tool1"])
schemas = tai.McpBridge(self.app).schemas_for("t")
self.assertEqual([schema["function"]["name"] for schema in schemas], ["mcp_srv_tool1"])
def test_resolve(self):
self.seed()
bridge = tai.McpBridge(self.app)
self.assertEqual(bridge.resolve("mcp_srv_echo"), ("srv", "echo"))
self.assertIsNone(bridge.resolve("mcp_nope_echo"))
self.assertIsNone(bridge.resolve("shell"))
def test_call_happy_path(self):
self.seed()
fake = FakeMcpClient(MCP_FAKE_TOOLS)
with mock.patch.object(tai.McpBridge, "connect", return_value=fake):
out = tai.McpBridge(self.app).call("srv", "echo", {"text": "hi"}, "t")
self.assertEqual(out, "ok:echo")
self.assertTrue(fake.closed)
self.assertTrue(any("mcp call srv/echo" in item for item in self.app.approvals))
found = self.store.db.execute("SELECT action, path FROM audit WHERE action = 'mcp-call'").fetchall()
self.assertEqual(len(found), 1)
self.assertEqual(found[0][1], "srv/echo")
def test_call_guards(self):
self.assertTrue(tai.McpBridge(self.app).call("nope", "x", {}, "t").startswith("error:"))
self.seed(enabled=False)
self.assertIn("disabled", tai.McpBridge(self.app).call("srv", "echo", {}, "t"))
self.seed(deny=["echo"])
self.assertIn("not enabled", tai.McpBridge(self.app).call("srv", "echo", {}, "t"))
self.assertIn("must be an object", tai.McpBridge(self.app).call("srv", "echo", [], "t"))
def test_call_drift_prompts(self):
self.seed()
drifted = MCP_FAKE_TOOLS + [{"name": "newtool", "description": "appeared", "inputSchema": {}}]
fake = FakeMcpClient(drifted)
with mock.patch.object(tai.McpBridge, "connect", return_value=fake):
out = tai.McpBridge(self.app).call("srv", "echo", {}, "t")
self.assertEqual(out, "ok:echo")
self.assertIn("newtool", self.store.mcp_tool_pins("srv"))
self.app.approve = False
drifted2 = drifted + [{"name": "evil", "description": "x", "inputSchema": {}}]
with mock.patch.object(tai.McpBridge, "connect", return_value=FakeMcpClient(drifted2)):
out = tai.McpBridge(self.app).call("srv", "echo", {}, "t")
self.assertIn("not approved", out)
self.assertNotIn("evil", self.store.mcp_tool_pins("srv"))
def test_call_iserror_and_failure(self):
self.seed()
fake = FakeMcpClient(MCP_FAKE_TOOLS, results={"echo": {"content": [{"type": "text", "text": "kaput"}], "isError": True}})
with mock.patch.object(tai.McpBridge, "connect", return_value=fake):
self.assertEqual(tai.McpBridge(self.app).call("srv", "echo", {}, "t"), "error: kaput")
with mock.patch.object(tai.McpBridge, "connect", return_value=FakeMcpClient(MCP_FAKE_TOOLS, fail="call")):
self.assertTrue(tai.McpBridge(self.app).call("srv", "echo", {}, "t").startswith("error:"))
def test_call_spills_large_results(self):
self.seed()
big = "x" * 300
fake = FakeMcpClient(MCP_FAKE_TOOLS, results={"echo": {"content": [{"type": "text", "text": big}]}})
with mock.patch.object(tai, "MCP_RESULT_MAX", 100):
with mock.patch.object(tai.McpBridge, "connect", return_value=fake):
out = tai.McpBridge(self.app).call("srv", "echo", {}, "t")
self.assertIn("spilled to record mem:", out)
self.assertIn(big[:100], out)
self.assertNotIn(big, out)
def test_call_missing_secrets(self):
self.seed(headers={"Authorization": "Bearer ${NOPE_MISSING}"})
out = tai.McpBridge(self.app).call("srv", "echo", {}, "t")
self.assertIn("unresolved secrets", out)
self.assertIn("NOPE_MISSING", out)
class McpToolsTests(McpBase):
def test_add_json_flow(self):
fake = FakeMcpClient(MCP_FAKE_TOOLS)
with mock.patch.object(tai.McpBridge, "connect", return_value=fake):
out = self.tools.dispatch("mcp_add", {"source": McpNormalizeTests.JINA_URL})
self.assertIn("added MCP server 'jina-mcp-server'", out)
self.assertIn("echo back input", out)
self.assertTrue(self.store.mcp_server_enabled("jina-mcp-server", "t"))
self.assertTrue(fake.closed)
def test_add_denied_stores_nothing(self):
self.app.approve = False
out = self.tools.dispatch("mcp_add", {"source": McpNormalizeTests.JINA_URL})
self.assertEqual(out, "denied by user")
self.assertEqual(self.store.mcp_list_servers(), [])
def test_add_verify_failure_stores_nothing(self):
with mock.patch.object(tai.McpBridge, "connect", side_effect=tai.McpError("nope")):
out = self.tools.dispatch("mcp_add", {"source": McpNormalizeTests.JINA_URL})
self.assertIn("verify failed", out)
self.assertEqual(self.store.mcp_list_servers(), [])
def test_add_rejects_garbage(self):
self.assertTrue(self.tools.dispatch("mcp_add", {"source": "hello"}).startswith("error:"))
self.assertTrue(self.tools.dispatch("mcp_add", {}).startswith("error:"))
def test_add_github_guidance(self):
with mock.patch.object(tai.Tools, "mcp_resolve_github", side_effect=tai.McpError("nothing found")):
out = self.tools.dispatch("mcp_add", {"source": "https://github.com/jina-ai/MCP"})
self.assertIn("could not resolve github jina-ai/MCP", out)
self.assertIn("mcpServers", out)
def test_resolve_github_package(self):
payload = json.dumps({"type": "file", "content": base64.b64encode(b'{"name": "@x/y"}').decode("ascii")})
def fake_fetch(args):
if "package.json" in args["url"]:
return payload
return "error: 404"
with mock.patch.object(tai.Tools, "run_web_fetch", side_effect=fake_fetch):
entry, notes = self.tools.mcp_resolve_github("o", "r")
self.assertEqual(entry["command"], "npx")
self.assertEqual(entry["args"], ["-y", "@x/y"])
with mock.patch.object(tai.Tools, "run_web_fetch", return_value="error: 404"):
with self.assertRaises(tai.McpError):
self.tools.mcp_resolve_github("o", "r")
def test_enable_disable_remove_list(self):
self.store.mcp_save_server(dict(McpVaultTests.ENTRY))
self.assertIn("enabled for t", self.tools.dispatch("mcp_enable", {"server": "srv"}))
self.assertIn("srv (http)", self.tools.dispatch("mcp_list", {}))
self.assertIn("disabled for t", self.tools.dispatch("mcp_disable", {"server": "srv"}))
self.assertIn("unknown MCP server", self.tools.dispatch("mcp_enable", {"server": "nope"}))
self.assertIn("removed MCP server", self.tools.dispatch("mcp_remove", {"server": "srv"}))
self.assertEqual(self.tools.dispatch("mcp_list", {}), "no MCP servers registered")
def test_tools_and_call(self):
self.store.mcp_save_server(dict(McpVaultTests.ENTRY))
self.store.mcp_set_enabled("srv", "t", True)
self.store.mcp_store_tools("srv", MCP_FAKE_TOOLS)
out = self.tools.dispatch("mcp_tools", {"server": "srv"})
self.assertIn("srv/echo", out)
fake = FakeMcpClient(MCP_FAKE_TOOLS)
with mock.patch.object(tai.McpBridge, "connect", return_value=fake):
self.assertEqual(self.tools.dispatch("mcp_call", {"server": "srv", "tool": "echo", "args": {}}), "ok:echo")
self.assertEqual(self.tools.dispatch("mcp_srv_echo", {"text": "hi"}), "ok:echo")
self.assertIn("has no tool", self.tools.dispatch("mcp_srv_nope", {}))
self.assertEqual(self.tools.dispatch("mcp_zzz_nope", {}), "error: unknown tool mcp_zzz_nope")
self.assertTrue(self.tools.dispatch("mcp_call", {"server": "srv"}).startswith("error:"))
def test_select_tools_extra(self):
schema = tai.tool_schema("mcp_srv_echo", "MCP srv/echo: echo back input reverberation", {"text": {"type": "string"}}, ["text"])
found = {item["function"]["name"] for item in tai.select_tools("please echo back input", [schema])}
self.assertIn("mcp_srv_echo", found)
found = {item["function"]["name"] for item in tai.select_tools("please echo back input")}
self.assertNotIn("mcp_srv_echo", found)
class McpServeTests(McpBase):
def setUp(self):
super(McpServeTests, self).setUp()
self.ctx = tai.McpServeContext(self.config, self.store, "k")
self.server = http.server.HTTPServer(("127.0.0.1", 0), tai.mcp_serve_handler(self.ctx))
self.thread = threading.Thread(target=self.server.serve_forever)
self.thread.daemon = True
self.thread.start()
self.url = "http://127.0.0.1:%d/mcp" % self.server.server_port
def tearDown(self):
self.server.shutdown()
self.server.server_close()
self.thread.join(timeout=10)
super(McpServeTests, self).tearDown()
def post(self, message, key="k", session="", raw=None):
heads = {"Content-Type": "application/json"}
if key is not None:
heads["Authorization"] = "Bearer " + key
if session:
heads["Mcp-Session-Id"] = session
body = raw if raw is not None else json.dumps(message).encode("utf-8")
request = urllib.request.Request(self.url, data=body, headers=heads, method="POST")
try:
response = urllib.request.urlopen(request, timeout=10)
except urllib.error.HTTPError as exc:
return exc.code, {}
return response.status, json.loads(response.read().decode("utf-8") or "{}")
def handshake(self):
status, reply = self.post({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}})
self.assertEqual(status, 200)
self.assertEqual(reply["result"]["protocolVersion"], tai.MCP_PROTOCOL_VERSION)
self.assertIn("tools", reply["result"]["capabilities"])
sids = list(self.ctx.sessions)
self.assertEqual(len(sids), 1)
return sids[0]
def test_auth_path_and_method_gates(self):
self.assertEqual(self.post({"jsonrpc": "2.0", "id": 1, "method": "ping"}, key="wrong")[0], 401)
self.assertEqual(self.post({"jsonrpc": "2.0", "id": 1, "method": "ping"}, key=None)[0], 401)
request = urllib.request.Request(self.url + "x", data=b"{}",
headers={"Authorization": "Bearer k"}, method="POST")
try:
urllib.request.urlopen(request, timeout=10)
self.fail("expected 404")
except urllib.error.HTTPError as exc:
self.assertEqual(exc.code, 404)
request = urllib.request.Request(self.url, headers={"Authorization": "Bearer k"}, method="GET")
try:
urllib.request.urlopen(request, timeout=10)
self.fail("expected 405")
except urllib.error.HTTPError as exc:
self.assertEqual(exc.code, 405)
self.assertEqual(self.post({}, raw=b"not json")[0], 400)
def test_ping_list_and_unknown(self):
sid = self.handshake()
status, reply = self.post({"jsonrpc": "2.0", "id": 2, "method": "ping", "params": {}}, session=sid)
self.assertEqual(reply["result"], {})
status, reply = self.post({"jsonrpc": "2.0", "id": 3, "method": "tools/list", "params": {}}, session=sid)
names = [item["name"] for item in reply["result"]["tools"]]
self.assertEqual(len(names), len(tai.TOOL_SCHEMAS) + 2)
self.assertIn("chat", names)
self.assertIn("completions", names)
self.assertIn("shell", names)
chat = [item for item in reply["result"]["tools"] if item["name"] == "chat"][0]
self.assertIn("message", chat["inputSchema"]["required"])
status, reply = self.post({"jsonrpc": "2.0", "id": 4, "method": "roots/list", "params": {}}, session=sid)
self.assertEqual(reply["error"]["code"], -32601)
status, reply = self.post({"jsonrpc": "2.0", "id": 5, "method": "tools/call",
"params": {"name": "chat", "arguments": []}}, session=sid)
self.assertEqual(reply["error"]["code"], -32602)
def test_unknown_session_rejected(self):
status, reply = self.post({"jsonrpc": "2.0", "id": 1, "method": "ping", "params": {}}, session="nope")
self.assertIn("error", reply)
def test_call_native_tool(self):
sid = self.handshake()
path = os.path.join(self.tmp.name, "note.txt")
with open(path, "w", encoding="utf-8") as handle:
handle.write("hello")
status, reply = self.post({"jsonrpc": "2.0", "id": 6, "method": "tools/call",
"params": {"name": "read_file", "arguments": {"path": path}}}, session=sid)
self.assertIn("hello", reply["result"]["content"][0]["text"])
self.assertNotIn("isError", reply["result"])
status, reply = self.post({"jsonrpc": "2.0", "id": 7, "method": "tools/call",
"params": {"name": "no_such_tool", "arguments": {}}}, session=sid)
self.assertTrue(reply["result"]["isError"])
def test_chat_and_completions(self):
sid = self.handshake()
with mock.patch.object(tai.Agent, "run_turn", return_value="hi there") as turned:
status, reply = self.post({"jsonrpc": "2.0", "id": 8, "method": "tools/call",
"params": {"name": "chat", "arguments": {"message": "hey"}}}, session=sid)
self.assertEqual(reply["result"]["content"][0]["text"], "hi there")
self.assertTrue(turned.called)
self.assertIsNotNone(self.ctx.sessions[sid].agent)
with mock.patch.object(tai.ChatClient, "complete", return_value="done") as completed:
status, reply = self.post({"jsonrpc": "2.0", "id": 9, "method": "tools/call",
"params": {"name": "completions",
"arguments": {"messages": [{"role": "user", "content": "2+2"}]}}},
session=sid)
self.assertEqual(reply["result"]["content"][0]["text"], "done")
self.assertTrue(completed.called)
for bad in ({}, {"messages": []}, {"messages": [{"role": "bogus", "content": "x"}]},
{"messages": [{"role": "user", "content": 42}]}):
status, reply = self.post({"jsonrpc": "2.0", "id": 10, "method": "tools/call",
"params": {"name": "completions", "arguments": bad}}, session=sid)
self.assertEqual(reply["error"]["code"], -32602)
status, reply = self.post({"jsonrpc": "2.0", "id": 11, "method": "tools/call",
"params": {"name": "chat", "arguments": {"message": " "}}}, session=sid)
self.assertEqual(reply["error"]["code"], -32602)
def test_delete_drops_session(self):
sid = self.handshake()
request = urllib.request.Request(self.url, headers={"Authorization": "Bearer k", "Mcp-Session-Id": sid},
method="DELETE")
response = urllib.request.urlopen(request, timeout=10)
self.assertEqual(response.status, 200)
self.assertNotIn(sid, self.ctx.sessions)
status, reply = self.post({"jsonrpc": "2.0", "id": 2, "method": "ping", "params": {}}, session=sid)
self.assertIn("error", reply)
def test_serve_main_validates_port(self):
args = mock.Mock(mcp_port=99999, mcp_bind=None)
with contextlib.redirect_stderr(io.StringIO()):
self.assertEqual(tai.mcp_serve_main(args, "k"), 2)
class McpDocTests(unittest.TestCase):
def setUp(self):
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "mcp.md"), encoding="utf-8") as handle:
self.doc = handle.read()
def test_doc_names_everything(self):
for const in ("MCP_PROTOCOL_VERSION", "MCP_CONNECT_TIMEOUT", "MCP_CALL_TIMEOUT", "MCP_RESULT_MAX",
"MCP_HTTP_MAX_BODY", "MCP_MAX_SERVERS", "MCP_MAX_TOOLS", "MCP_SESSION_IDLE",
"MCP_CHAT_TIMEOUT", "MCP_DEFAULT_PORT"):
self.assertIn(const, self.doc, const)
self.assertTrue(hasattr(tai, const), const)
for method in ("initialize", "tools/list", "tools/call", "ping", "notifications/cancelled"):
self.assertIn(method, self.doc, method)
for table in ("mcp_servers", "mcp_server_profiles", "mcp_tool_pins", "mcp_server_tools"):
self.assertIn(table, self.doc, table)
for tool in ("mcp_add", "mcp_list", "mcp_enable", "mcp_disable", "mcp_remove", "mcp_tools", "mcp_call"):
self.assertIn(tool, self.doc, tool)
for symbol in ("McpBridge", "McpClient", "McpStdioTransport", "McpHttpTransport", "McpServeContext",
"mcp_normalize", "mcp_serve_main", "chat", "completions", "--serve-mcp", "--mcp-port", "--mcp-bind"):
self.assertIn(symbol, self.doc, symbol)
self.assertNotIn(chr(0x2014), self.doc)
if __name__ == "__main__":
unittest.main()