diff --git a/.gitignore b/.gitignore index 0d9ff1c..2cf4aad 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,7 @@ !cap.md !skills.md !settings.md +!mcp.md # Backup files *.bak diff --git a/AGENTS.md b/AGENTS.md index 0caa167..e8deeb6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -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). diff --git a/README.md b/README.md index 67ffec3..e90fe99 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ retoor **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__`. 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 diff --git a/mcp.md b/mcp.md new file mode 100644 index 0000000..7fa3b8c --- /dev/null +++ b/mcp.md @@ -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 ` 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:/`, `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//` 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__`, 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. diff --git a/tai.py b/tai.py index db44a2c..afac34e 100755 --- a/tai.py +++ b/tai.py @@ -14,6 +14,7 @@ import importlib.util import hmac import html import http.client +import http.server import io import itertools import json @@ -102,7 +103,7 @@ def _fromisoformat(text): return parsed.replace(tzinfo=tzinfo) raise ValueError("invalid datetime, use ISO like 2026-10-08T09:00") -VERSION = "1.41.0" +VERSION = "1.42.0" WORKER_STEPS = 12 CREATE_SKILL_STEPS = 40 FORK_MAX_DEPTH = 2 @@ -2358,6 +2359,7 @@ def rotate_seal(config, old_seal, new_passphrase): meta_rows = db.execute("SELECT profile, name, meta FROM secrets WHERE meta IS NOT NULL").fetchall() if has_profile and "meta" in secret_cols else [(None, name, meta) for name, meta in db.execute("SELECT name, meta FROM secrets WHERE meta IS NOT NULL").fetchall()] if "meta" in secret_cols else [] planned_rows = db.execute("SELECT id, prompt FROM schedules").fetchall() if "schedules" in tables else [] setting_rows = db.execute("SELECT key, scope, value FROM settings").fetchall() if "settings" in tables else [] + mcp_rows = db.execute("SELECT name, headers_json FROM mcp_servers").fetchall() if "mcp_servers" in tables else [] db.close() texts = {} for row_id, text in rows: @@ -2377,6 +2379,12 @@ def rotate_seal(config, old_seal, new_passphrase): setting_texts[(key, scope)] = (old_seal.unlock(value), True) else: setting_texts[(key, scope)] = (value, False) + mcp_texts = {} + for name, headers in mcp_rows: + if headers and headers.startswith(Seal.PREFIX): + mcp_texts[name] = (old_seal.unlock(headers), True) + else: + mcp_texts[name] = (headers, False) files = {} for entry in os.listdir(config.profiles_dir): if not entry.endswith((".sys.md", ".session.json", ".bots.json")): @@ -2408,6 +2416,9 @@ def rotate_seal(config, old_seal, new_passphrase): for (key, scope), (plain, was_sealed) in setting_texts.items(): db.execute("UPDATE settings SET value = ? WHERE key = ? AND scope = ?", (fresh.lock(plain) if was_sealed else plain, key, scope)) + for name, (plain, was_sealed) in mcp_texts.items(): + db.execute("UPDATE mcp_servers SET headers_json = ? WHERE name = ?", + (fresh.lock(plain) if was_sealed else plain, name)) db.commit() db.close() for path, plain in files.items(): @@ -5424,6 +5435,10 @@ SCHEMA = ( ("input_history", ("id INTEGER PRIMARY KEY", "profile TEXT", "bot TEXT", "ts TEXT", "text TEXT"), ()), ("shell_sessions", ("sid TEXT PRIMARY KEY", "parent TEXT", "shell TEXT", "pid INTEGER", "started TEXT", "last_seen TEXT", "cwd TEXT", "chain TEXT"), ()), ("settings", ("key TEXT", "scope TEXT", "value TEXT", "created TEXT", "updated TEXT"), ("key", "scope")), + ("mcp_servers", ("name TEXT PRIMARY KEY", "transport TEXT", "command TEXT", "args_json TEXT", "env_json TEXT", "url TEXT", "headers_json TEXT", "cwd TEXT", "allow_json TEXT", "deny_json TEXT", "source TEXT", "created TEXT", "updated TEXT"), ()), + ("mcp_server_profiles", ("server TEXT", "profile TEXT", "enabled INTEGER"), ("server", "profile")), + ("mcp_tool_pins", ("server TEXT", "tool TEXT", "hash TEXT", "updated TEXT"), ("server", "tool")), + ("mcp_server_tools", ("server TEXT", "tool TEXT", "description TEXT", "schema_json TEXT", "updated TEXT"), ("server", "tool")), ) SCHEMA_INDEXES = ( @@ -5512,6 +5527,9 @@ class Store: configured = self.db.execute("SELECT COUNT(*) FROM settings WHERE value LIKE 'tai1$%'").fetchone()[0] if configured: raise SealError("settings are sealed, set TAI_PASSPHRASE") + bridged = self.db.execute("SELECT COUNT(*) FROM mcp_servers WHERE headers_json LIKE 'tai1$%'").fetchone()[0] + if bridged: + raise SealError("mcp headers are sealed, set TAI_PASSPHRASE") self.memdb = None self.mem_events_max = 0 if seal.enabled and self.fts_ok: @@ -6383,6 +6401,100 @@ class Store: raise ValueError("settings upsert failed: %s" % exc) return self.get_or_create_settings(sorted(canonical), scope=scope) + def mcp_parse_json(self, raw, fallback): + try: + data = json.loads(raw or "") + except ValueError: + return fallback + return data if isinstance(data, type(fallback)) else fallback + + def mcp_row(self, found): + name, transport, command, args_json, env_json, url, headers_json, cwd, allow_json, deny_json, source, created, updated = found + if headers_json and headers_json.startswith(Seal.PREFIX): + headers_json = self.seal.safe_unlock(headers_json) + return {"name": name, "transport": transport, "command": command or "", + "args": self.mcp_parse_json(args_json, []), "env": self.mcp_parse_json(env_json, {}), + "url": url or "", "headers": self.mcp_parse_json(headers_json, {}), "cwd": cwd or "", + "allow": self.mcp_parse_json(allow_json, []), "deny": self.mcp_parse_json(deny_json, []), + "source": source or "", "created": created, "updated": updated} + + def mcp_save_server(self, entry): + name = entry.get("name", "") + if not MCP_SERVER_RE.match(name): + raise ValueError("bad MCP server name '%s'" % name) + known = self.db.execute("SELECT created FROM mcp_servers WHERE name = ?", (name,)).fetchone() + if known is None: + count = self.db.execute("SELECT COUNT(*) FROM mcp_servers").fetchone()[0] + if count >= MCP_MAX_SERVERS: + raise ValueError("too many MCP servers (cap %d)" % MCP_MAX_SERVERS) + headers = entry.get("headers") or {} + sealed = self.seal.lock(json.dumps(headers, sort_keys=True)) if headers else "" + now = now_iso() + created = known[0] if known else now + self.db.execute("INSERT OR REPLACE INTO mcp_servers " + "(name, transport, command, args_json, env_json, url, headers_json, cwd, allow_json, deny_json, source, created, updated) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + (name, entry.get("transport"), entry.get("command") or "", json.dumps(entry.get("args") or []), + json.dumps(entry.get("env") or {}), entry.get("url") or "", sealed, entry.get("cwd") or "", + json.dumps(entry.get("allow") or []), json.dumps(entry.get("deny") or []), + entry.get("source") or "", created, now)) + self.db.commit() + + def mcp_get_server(self, name): + found = self.db.execute("SELECT name, transport, command, args_json, env_json, url, headers_json, cwd, allow_json, deny_json, source, created, updated " + "FROM mcp_servers WHERE name = ?", (name,)).fetchone() + return self.mcp_row(found) if found else None + + def mcp_list_servers(self): + found = self.db.execute("SELECT name, transport, command, args_json, env_json, url, headers_json, cwd, allow_json, deny_json, source, created, updated " + "FROM mcp_servers ORDER BY name").fetchall() + return [self.mcp_row(row) for row in found] + + def mcp_remove_server(self, name): + self.db.execute("DELETE FROM mcp_server_tools WHERE server = ?", (name,)) + self.db.execute("DELETE FROM mcp_tool_pins WHERE server = ?", (name,)) + self.db.execute("DELETE FROM mcp_server_profiles WHERE server = ?", (name,)) + self.db.execute("DELETE FROM mcp_servers WHERE name = ?", (name,)) + self.db.commit() + + def mcp_set_enabled(self, server, profile, enabled): + self.db.execute("INSERT OR REPLACE INTO mcp_server_profiles (server, profile, enabled) VALUES (?, ?, ?)", + (server, profile, 1 if enabled else 0)) + self.db.commit() + + def mcp_server_enabled(self, server, profile): + found = self.db.execute("SELECT enabled FROM mcp_server_profiles WHERE server = ? AND profile = ?", + (server, profile)).fetchone() + return bool(found and found[0]) + + def mcp_enabled_servers(self, profile): + return [row for row in self.mcp_list_servers() if self.mcp_server_enabled(row["name"], profile)] + + def mcp_store_tools(self, server, tools): + now = now_iso() + self.db.execute("DELETE FROM mcp_server_tools WHERE server = ?", (server,)) + self.db.execute("DELETE FROM mcp_tool_pins WHERE server = ?", (server,)) + for item in tools: + if not isinstance(item, dict) or not item.get("name"): + continue + tool = item["name"] + desc = item.get("description", "") or "" + schema = item.get("inputSchema") or {} + self.db.execute("INSERT INTO mcp_server_tools (server, tool, description, schema_json, updated) VALUES (?, ?, ?, ?, ?)", + (server, tool, desc, json.dumps(schema), now)) + self.db.execute("INSERT INTO mcp_tool_pins (server, tool, hash, updated) VALUES (?, ?, ?, ?)", + (server, tool, mcp_tool_hash(tool, desc, schema), now)) + self.db.commit() + + def mcp_cached_tools(self, server): + found = self.db.execute("SELECT tool, description, schema_json FROM mcp_server_tools WHERE server = ? ORDER BY tool", + (server,)).fetchall() + return [{"tool": tool, "description": desc or "", "schema": self.mcp_parse_json(schema, {})} for tool, desc, schema in found] + + def mcp_tool_pins(self, server): + found = self.db.execute("SELECT tool, hash FROM mcp_tool_pins WHERE server = ?", (server,)).fetchall() + return dict((tool, digest) for tool, digest in found) + def load_dir_defaults(self, directory): try: row = self.db.execute("SELECT provider, model, bot FROM dir_defaults WHERE dir = ?", (directory,)).fetchone() @@ -8598,6 +8710,13 @@ TOOL_SCHEMAS = [ tool_schema("install", "Manage tai installations: binary, command-not-found hook and plexer auto-wrap hook for bash/zsh/fish, venv, scheduler service, telegram service, container. Status reports what exists. Install adds missing pieces (skipping a shell hook loudly if that shell isn't on the machine), upgrade refreshes in place and restarts services, reinstall rebuilds artifacts, uninstall removes them. Service changes take effect immediately. The vault is never touched by any action. Non-status actions ask the user first.", {"action": {"type": "string"}, "targets": {"type": "array"}}, ["action"]), tool_schema("create_bot", "Create a bot in the current profile: a name plus its own system message and history. Give at least one of description, rules, or behavior; optional nicknames register as @mention aliases (a short form of the name registers automatically). Names are lowercase.", {"name": {"type": "string"}, "description": {"type": "string"}, "rules": {"type": "string"}, "behavior": {"type": "string"}, "nicknames": {"type": "array"}}, ["name"]), tool_schema("search", "Search everything at once with ranked full-text matching (SQLite FTS5 + bm25, no vectors): records, episodic events, the audit trail, and RAG chunks. Chunks are the retrieval-grade granular hits (e.g. one slice of an exported terminal log) and come back with their actual text, not just a snippet; use the chunk's parent id with record_read to see more surrounding context. One call replaces paging through each store separately. expand adds one graph hop per record hit.", {"query": {"type": "string"}, "kinds": {"type": "array"}, "limit": {"type": "integer"}, "expand": {"type": "boolean"}}, ["query"]), + tool_schema("mcp_add", "Register an MCP server from pasted JSON/TOML config or a github URL. Normalizes the config, shows the exact spawn for approval, verifies with a live handshake, stores the connection, and enables it for the profile. Github URLs resolve via registry, package files, and README run blocks, verified live before storing.", {"source": {"type": "string"}, "name": {"type": "string"}, "profile": {"type": "string"}}, ["source"]), + tool_schema("mcp_list", "List registered MCP servers with transport, enabled profiles, tool counts, and origin.", {}, []), + tool_schema("mcp_enable", "Enable an MCP server for a profile (current when omitted).", {"server": {"type": "string"}, "profile": {"type": "string"}}, ["server"]), + tool_schema("mcp_disable", "Disable an MCP server for a profile (current when omitted).", {"server": {"type": "string"}, "profile": {"type": "string"}}, ["server"]), + tool_schema("mcp_remove", "Delete an MCP server with its profile rows, pins, and cached tools. Asks first.", {"server": {"type": "string"}}, ["server"]), + tool_schema("mcp_tools", "List cached MCP tools with descriptions and pin state, for one server or all.", {"server": {"type": "string"}}, []), + tool_schema("mcp_call", "Call one MCP tool explicitly: full path with drift check, approval, caps, and audit. Bridged mcp__ names call the same path.", {"server": {"type": "string"}, "tool": {"type": "string"}, "args": {"type": "object"}}, ["server", "tool"]), ] @@ -8683,6 +8802,13 @@ TOOL_TAGS = { "install": ("install", "uninstall", "reinstall", "upgrade", "setup", "service", "systemd", "venv", "hook", "deploy", "zsh", "fish"), "create_bot": ("bot", "bots", "persona", "character", "mention"), "search": ("search", "find", "lookup", "everything"), + "mcp_add": ("mcp", "server", "register", "connect", "integrate"), + "mcp_list": ("mcp", "servers", "connections"), + "mcp_enable": ("mcp", "enable", "activate"), + "mcp_disable": ("mcp", "disable", "deactivate"), + "mcp_remove": ("mcp", "remove", "delete", "unregister"), + "mcp_tools": ("mcp", "tools"), + "mcp_call": ("mcp", "call", "invoke"), } @@ -8708,7 +8834,7 @@ def conversation_text(messages): return "\n".join(parts)[-6000:] -def select_tools(text): +def select_tools(text, extra=None): words = set(re.findall(r"[a-z0-9]+", (text or "").lower())) picked = [] for schema in TOOL_SCHEMAS: @@ -8728,6 +8854,15 @@ def select_tools(text): desc = set(word for word in re.findall(r"[a-z0-9]+", schema["function"]["description"].lower()) if len(word) >= 6) if desc & words: picked.append(schema) + for schema in extra or []: + name = schema["function"]["name"] + tokens = [item for item in name.split("_") if len(item) >= 4 and item not in LAZY_NAME_STOP] + if tokens and set(tokens) & words: + picked.append(schema) + continue + desc = set(word for word in re.findall(r"[a-z0-9]+", schema["function"]["description"].lower()) if len(word) >= 6) + if desc & words: + picked.append(schema) return picked @@ -16628,6 +16763,13 @@ class Tools: "install": self.run_install, "create_bot": self.run_create_bot, "search": self.run_search, + "mcp_add": self.run_mcp_add, + "mcp_list": self.run_mcp_list, + "mcp_enable": self.run_mcp_enable, + "mcp_disable": self.run_mcp_disable, + "mcp_remove": self.run_mcp_remove, + "mcp_tools": self.run_mcp_tools, + "mcp_call": self.run_mcp_call, } self.secret_grants = set() self.read_files = set() @@ -16644,6 +16786,193 @@ class Tools: except BackendError as exc: return "", "error: " + short_error(exc) + def mcp_schemas(self): + if self.app.store is None: + return [] + try: + return McpBridge(self.app).schemas_for(self.active_profile) + except (sqlite3.Error, OSError, ValueError): + return [] + + def mcp_resolve_github(self, owner, repo): + """Deterministic half of GitHub resolution. Returns (entry-ish, notes) or raises McpError.""" + notes = [] + base = "https://api.github.com/repos/%s/%s/contents/" % (owner, repo) + + def fetch(path): + out = self.run_web_fetch({"url": base + path, "raw": True}) + if not isinstance(out, str) or out.startswith("error:") or len(out) > 200000: + return None + try: + data = json.loads(out) + except ValueError: + return None + if not isinstance(data, dict) or data.get("type") != "file" or not data.get("content"): + return None + try: + return base64.b64decode(data["content"]).decode("utf-8", "replace") + except (ValueError, TypeError): + return None + + for candidate in ("mcp.json", ".mcp.json", "server.json"): + text = fetch(candidate) + if text: + notes.append("found %s, normalizing it" % candidate) + entry, sub = mcp_normalize(text) + return entry, notes + sub + package = fetch("package.json") + if package: + try: + meta = json.loads(package) + except ValueError: + meta = {} + if isinstance(meta, dict): + name = meta.get("name", "") + if name: + notes.append("package %s has no run contract, trying npx" % name) + return {"name": mcp_sanitize_name(repo), "transport": "stdio", "command": "npx", + "args": ["-y", name], "env": {}, "url": "", "headers": {}, "cwd": "", + "allow": [], "deny": [], "enabled": None, "source": "github:%s/%s" % (owner, repo)}, notes + for candidate in ("README.md", "readme.md"): + text = fetch(candidate) + if text: + found = re.search(r"(npx|uvx)\s+([^\s\"'`]+)", text) + if found: + notes.append("run block found in %s" % candidate) + return {"name": mcp_sanitize_name(repo), "transport": "stdio", "command": found.group(1), + "args": ["-y", found.group(2)] if found.group(1) == "npx" else [found.group(2)], + "env": {}, "url": "", "headers": {}, "cwd": "", "allow": [], "deny": [], + "enabled": None, "source": "github:%s/%s" % (owner, repo)}, notes + raise McpError("no run contract found (checked mcp.json, server.json, package.json, README)") + + def run_mcp_add(self, args): + source = str(args.get("source") or "").strip() + profile = str(args.get("profile") or self.active_profile) + if not source: + return "error: source must be pasted JSON/TOML config or a github URL" + target = mcp_github_target(source) if ("\n" not in source and "{" not in source and "[" not in source) else None + notes = [] + try: + if target: + entry, notes = self.mcp_resolve_github(target[0], target[1]) + entry["source"] = "github:%s/%s" % (target[0], target[1]) + else: + text, rewritten = mcp_rewrite_placeholders(source) + notes += rewritten + entry, sub = mcp_normalize(text, str(args.get("name") or "") or None) + notes += sub + except McpError as exc: + if target: + return ("error: could not resolve github %s/%s: %s\n" + "finish it by hand: read the repo README for the run command, then call mcp_add " + "with explicit JSON like {\"mcpServers\": {\"name\": {\"command\": \"npx\", \"args\": [...]}}}" % (target[0], target[1], exc)) + return "error: %s" % exc + name = str(args.get("name") or entry["name"] or "") + if not MCP_SERVER_RE.match(name): + return "error: bad server name '%s'" % name + entry["name"] = name + if entry["transport"] == "stdio": + shown = " ".join([entry["command"]] + entry["args"]) + flags = [word for word in ("sudo", "rm -rf", "curl", "wget", "ssh", "nc ") if word in shown] + prompt = "mcp add stdio '%s'%s (runs with your privileges)" % (shown[:300], " FLAGS: %s" % ",".join(flags) if flags else "") + else: + prompt = "mcp add http %s (headers sealed, secrets resolve at call time)" % entry["url"][:300] + if notes: + prompt += " [%s]" % "; ".join(notes)[:300] + if not self.app.ask_approval(prompt): + return "denied by user" + bridge = McpBridge(self.app) + try: + client = bridge.connect(entry, profile) + except McpError as exc: + return "error: verify failed, nothing stored: %s" % exc + try: + try: + tools = client.list_tools() + except McpError as exc: + return "error: verify failed, nothing stored: %s" % exc + finally: + client.close() + try: + self.app.store.mcp_save_server(entry) + except ValueError as exc: + return "error: %s" % exc + self.app.store.mcp_set_enabled(name, profile, entry["enabled"] if entry["enabled"] is not None else True) + self.app.store.mcp_store_tools(name, tools) + lines = ["added MCP server '%s' (%s, %d tools, enabled for %s)" % (name, entry["transport"], len(tools), profile)] + lines += ["- %s: %s" % (item.get("name", ""), (item.get("description", "") or "")[:120]) for item in tools[:MCP_MAX_TOOLS]] + if len(tools) > MCP_MAX_TOOLS: + lines.append("- ...and %d more (cap %d offered)" % (len(tools) - MCP_MAX_TOOLS, MCP_MAX_TOOLS)) + return "\n".join(lines) + + def run_mcp_list(self, args): + rows = self.app.store.mcp_list_servers() + if not rows: + return "no MCP servers registered" + lines = [] + for row in rows: + if row["transport"] == "stdio": + detail = " ".join([row["command"]] + row["args"])[:120] + else: + detail = row["url"][:120] + enabled = [found[0] for found in + self.app.store.db.execute("SELECT profile FROM mcp_server_profiles WHERE server = ? AND enabled = 1", (row["name"],)).fetchall()] + cached = self.app.store.mcp_cached_tools(row["name"]) + lines.append("%s (%s) %s | profiles: %s | tools: %d | from %s" + % (row["name"], row["transport"], detail, ", ".join(sorted(enabled)) or "none", + len(cached), row["source"] or "?")) + return "\n".join(lines) + + def run_mcp_enable(self, args): + server = str(args.get("server") or "") + profile = str(args.get("profile") or self.active_profile) + if self.app.store.mcp_get_server(server) is None: + return "error: unknown MCP server '%s'" % server + self.app.store.mcp_set_enabled(server, profile, True) + return "MCP server '%s' enabled for %s" % (server, profile) + + def run_mcp_disable(self, args): + server = str(args.get("server") or "") + profile = str(args.get("profile") or self.active_profile) + if self.app.store.mcp_get_server(server) is None: + return "error: unknown MCP server '%s'" % server + self.app.store.mcp_set_enabled(server, profile, False) + return "MCP server '%s' disabled for %s" % (server, profile) + + def run_mcp_remove(self, args): + server = str(args.get("server") or "") + if self.app.store.mcp_get_server(server) is None: + return "error: unknown MCP server '%s'" % server + if not self.app.ask_approval("mcp remove '%s' (server, profile rows, pins, cached tools)" % server): + return "denied by user" + self.app.store.mcp_remove_server(server) + return "removed MCP server '%s'" % server + + def run_mcp_tools(self, args): + server = str(args.get("server") or "") + rows = [self.app.store.mcp_get_server(server)] if server else self.app.store.mcp_list_servers() + rows = [row for row in rows if row is not None] + if server and not rows: + return "error: unknown MCP server '%s'" % server + if not rows: + return "no MCP servers registered" + lines = [] + for row in rows: + pins = self.app.store.mcp_tool_pins(row["name"]) + for item in self.app.store.mcp_cached_tools(row["name"]): + lines.append("%s/%s%s: %s" % (row["name"], item["tool"], + "" if item["tool"] in pins else " (unpinned)", + (item["description"] or "")[:150])) + return "\n".join(lines) or "no cached tools (call mcp_add to verify)" + + def run_mcp_call(self, args): + server = str(args.get("server") or "") + tool = str(args.get("tool") or "") + call_args = args.get("args") or {} + if not server or not tool: + return "error: mcp_call needs server and tool" + return McpBridge(self.app).call(server, tool, call_args, self.active_profile) + def dispatch(self, name, raw_arguments): self.calls += 1 try: @@ -16651,6 +16980,11 @@ class Tools: except (TypeError, ValueError): return "error: arguments are not valid JSON" handler = self.handlers.get(name) + if handler is None and name.startswith(MCP_TOOL_PREFIX): + resolved = McpBridge(self.app).resolve(name) + if resolved is not None: + server, tool = resolved + handler = lambda call_args: McpBridge(self.app).call(server, tool, call_args, self.active_profile) if handler is None: return "error: unknown tool " + name try: @@ -19080,7 +19414,7 @@ class Agent: self.timed_out = True last_text = (last_text + "\n[time limit reached]").strip() break - payload = select_tools(conversation_text(self.messages)) + payload = select_tools(conversation_text(self.messages), self.tools.mcp_schemas()) tokens = conversation_tokens(self.messages, payload) if tokens > CONTEXT_CAP * COMPACT_RATIO: if loud: @@ -34472,6 +34806,928 @@ def grr_tui(): # --------------------------------------------------------------------------- +# MCP (Model Context Protocol), client side and server side. See mcp.md. +MCP_PROTOCOL_VERSION = "2025-11-25" +MCP_CLIENT_NAME = "tai" +MCP_TOOL_PREFIX = "mcp_" +MCP_CONNECT_TIMEOUT = 20 +MCP_CALL_TIMEOUT = 60 +MCP_RESULT_MAX = 65536 +MCP_HTTP_MAX_BODY = 10485760 +MCP_MAX_SERVERS = 32 +MCP_MAX_TOOLS = 100 +MCP_SESSION_IDLE = 3600 +MCP_CHAT_TIMEOUT = 600 +MCP_DEFAULT_PORT = 8765 +MCP_SERVER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$") +MCP_VAR_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}") +MCP_GITHUB_RE = re.compile(r"github\.com/([^/\s?#]+)/([^/\s?#]+?)(?:\.git)?(?:[/?#]|$)", re.IGNORECASE) + + +class McpError(Exception): + pass + + +class McpJsonRpcError(McpError): + def __init__(self, code, message): + super(McpJsonRpcError, self).__init__("json-rpc %s: %s" % (code, message)) + self.code = code + self.message = message + + +def mcp_sanitize_name(text): + cleaned = re.sub(r"[^a-z0-9_]", "_", (text or "").lower()).strip("_") + return cleaned or "server" + + +def mcp_github_target(source): + match = MCP_GITHUB_RE.search(source or "") + if not match: + return None + return match.group(1), match.group(2) + + +def mcp_expand_vars(text, lookup): + """Expand ${VAR} and ${VAR:-default} via lookup(name) -> value or None. + + Returns (expanded, missing): names with neither a value nor a default + keep their literal and land in missing.""" + missing = [] + + def fill(match): + name, default = match.group(1), match.group(2) + value = lookup(name) + if value is None: + if default is not None: + return default + missing.append(name) + return match.group(0) + return value + + return MCP_VAR_RE.sub(fill, text or ""), missing + + +def mcp_rewrite_placeholders(text): + """Rewrite foreign placeholder spellings to ${NAME}. Returns (text, noted).""" + noted = [] + if "${env:" in (text or ""): + text = re.sub(r"\$\{env:([A-Za-z_][A-Za-z0-9_]*)\}", r"${\1}", text) + noted.append("rewrote ${env:NAME} to ${NAME}") + if "${input:" in (text or ""): + text = re.sub(r"\$\{input:([^}]+)\}", r"${\1}", text) + noted.append("rewrote ${input:id} to ${id} (resolves against vault secrets)") + return text, noted + + +def mcp_entry_from_mapping(name, raw): + """Normalize one raw server mapping to a stored entry. Raises McpError.""" + if not isinstance(raw, dict): + raise McpError("server '%s' is not an object" % name) + entry = {"name": name, "transport": None, "command": "", "args": [], "env": {}, + "url": "", "headers": {}, "cwd": "", "allow": [], "deny": [], + "enabled": None, "source": "config"} + notes = [] + command = raw.get("command") or "" + url = raw.get("url") or "" + if command and url: + raise McpError("server '%s' sets both command and url" % name) + if command: + if not isinstance(command, str): + raise McpError("server '%s' command must be a string" % name) + entry["transport"] = "stdio" + entry["command"] = command + args = raw.get("args", []) + if not isinstance(args, list) or any(not isinstance(item, str) for item in args): + raise McpError("server '%s' args must be a string list" % name) + entry["args"] = list(args) + env = raw.get("env") or raw.get("env_vars") or {} + if not isinstance(env, dict) or any(not isinstance(k, str) or not isinstance(v, str) for k, v in env.items()): + raise McpError("server '%s' env must be a string map" % name) + entry["env"] = dict(env) + cwd = raw.get("cwd") or raw.get("working_directory") or "" + if cwd and not isinstance(cwd, str): + raise McpError("server '%s' cwd must be a string" % name) + entry["cwd"] = cwd + elif url: + if not isinstance(url, str): + raise McpError("server '%s' url must be a string" % name) + probe = (url.replace("${", "x").split() or [""])[0] + if not re.match(r"https?://[^/\s]+", probe): + raise McpError("server '%s' url must be http(s)" % name) + if not raw.get("type"): + notes.append("entry has url without type, treating as http") + entry["transport"] = "http" + entry["url"] = url + headers = raw.get("headers") or raw.get("http_headers") or {} + if not isinstance(headers, dict): + raise McpError("server '%s' headers must be a string map" % name) + entry["headers"] = dict((str(k), str(v)) for k, v in headers.items()) + indirect = raw.get("env_http_headers") or {} + if not isinstance(indirect, dict): + raise McpError("server '%s' env_http_headers must be a string map" % name) + for key, var in indirect.items(): + if not isinstance(var, str): + raise McpError("server '%s' env_http_headers values must be var names" % name) + entry["headers"][str(key)] = "${%s}" % var + else: + raise McpError("server '%s' needs command or url" % name) + for key in ("enabled_tools", "disabled_tools"): + items = raw.get(key) or [] + if not isinstance(items, list) or any(not isinstance(item, str) for item in items): + raise McpError("server '%s' %s must be a string list" % (name, key)) + entry["allow"] = list(raw.get("enabled_tools") or []) + entry["deny"] = list(raw.get("disabled_tools") or []) + if raw.get("enabled") is not None: + entry["enabled"] = bool(raw.get("enabled")) + if raw.get("timeout") is not None or raw.get("startup_timeout") is not None: + notes.append("timeout fields ignored, tai caps apply (connect %ds, call %ds)" + % (MCP_CONNECT_TIMEOUT, MCP_CALL_TIMEOUT)) + return entry, notes + + +def mcp_normalize(source, name=None): + """Normalize pasted JSON/TOML config text to (entry, notes). Raises McpError.""" + text = (source or "").strip() + if not text: + raise McpError("nothing to normalize") + if mcp_github_target(text) and "\n" not in text.strip() and len(text.strip()) < 200 and "{" not in text and "[" not in text: + raise McpError("github URL needs the resolver, not the normalizer") + entries = None + if text.startswith("{"): + try: + data = json.loads(text) + except ValueError as exc: + raise McpError("invalid JSON: %s" % short_error(exc)) + if not isinstance(data, dict): + raise McpError("config JSON must be an object") + wrapped = data.get("mcpServers", data.get("servers")) + if isinstance(wrapped, dict): + entries = wrapped + elif "command" in data or "url" in data: + entries = {name or "server": data} + else: + raise McpError("JSON needs mcpServers/servers or a bare command/url entry") + else: + try: + import tomllib + except ImportError: + raise McpError("TOML needs Python 3.11+") + try: + data = tomllib.loads(text) + except (ValueError, UnicodeDecodeError) as exc: + raise McpError("config is neither JSON nor TOML: %s" % short_error(exc)) + wrapped = data.get("mcp_servers", data.get("servers")) + if isinstance(wrapped, dict): + entries = wrapped + else: + raise McpError("TOML needs [mcp_servers.NAME] tables") + if name is not None: + if name not in entries: + raise McpError("no server '%s' in config (%s)" % (name, ", ".join(sorted(entries)) or "empty")) + entries = {name: entries[name]} + if len(entries) != 1: + raise McpError("config names %d servers (%s), pass name to pick one" + % (len(entries), ", ".join(sorted(entries)))) + single, raw = next(iter(entries.items())) + if not MCP_SERVER_RE.match(single): + raise McpError("server name '%s' must match [A-Za-z0-9][A-Za-z0-9_.-]{0,63}" % single) + entry, notes = mcp_entry_from_mapping(single, raw) + entry["source"] = "config-toml" if not text.startswith("{") else "config-json" + return entry, notes + + +class McpStdioTransport(object): + """One MCP server child over newline-delimited JSON-RPC on pipes.""" + + def __init__(self): + self.proc = None + self.buf = b"" + + def spawn(self, argv, env, cwd, timeout=MCP_CONNECT_TIMEOUT): + full_env = dict(os.environ) + full_env.update(env or {}) + try: + self.proc = subprocess.Popen(argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, env=full_env, cwd=cwd or None) + except OSError as exc: + raise McpError("spawn failed: %s" % short_error(exc)) + try: + os.set_blocking(self.proc.stdout.fileno(), False) + except (OSError, ValueError): + pass + return self + + def _read_line(self, deadline): + while b"\n" not in self.buf: + if self.proc.poll() is not None: + raise McpError("server exited during read (code %s)" % self.proc.returncode) + remaining = deadline - time.time() + if remaining <= 0: + raise McpError("read timed out") + try: + ready, _, _ = select.select([self.proc.stdout.fileno()], [], [], min(remaining, 1.0)) + except (OSError, ValueError): + raise McpError("stdio servers need POSIX select") + if not ready: + continue + try: + chunk = os.read(self.proc.stdout.fileno(), 65536) + except BlockingIOError: + continue + except OSError as exc: + raise McpError("read failed: %s" % short_error(exc)) + if not chunk: + raise McpError("server closed stdout") + self.buf += chunk + if len(self.buf) > MCP_HTTP_MAX_BODY: + raise McpError("server line exceeds %d bytes" % MCP_HTTP_MAX_BODY) + line, self.buf = self.buf.split(b"\n", 1) + return line.decode("utf-8", "replace") + + def send(self, message, timeout=MCP_CALL_TIMEOUT): + raw = (json.dumps(message) + "\n").encode("utf-8") + try: + self.proc.stdin.write(raw) + self.proc.stdin.flush() + except (OSError, ValueError) as exc: + raise McpError("write failed: %s" % short_error(exc)) + if "id" not in message: + return None + deadline = time.time() + timeout + while True: + try: + data = json.loads(self._read_line(deadline)) + except ValueError: + continue + if not isinstance(data, dict) or data.get("id") != message["id"]: + continue + return data + + def close(self): + proc, self.proc = self.proc, None + self.buf = b"" + if proc is None: + return + try: + if proc.stdin is not None: + proc.stdin.close() + except OSError: + pass + try: + proc.wait(timeout=2) + except subprocess.TimeoutExpired: + pass + if proc.poll() is None: + cap_proc_kill(proc) + cap_proc_reap(proc) + + +class McpHttpTransport(object): + """Streamable HTTP leg: POST per request, single JSON or SSE reply.""" + + def __init__(self, url, headers=None): + self.url = url + self.headers = dict(headers or {}) + self.session = "" + + def _post(self, message, timeout): + body = json.dumps(message).encode("utf-8") + heads = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"} + heads.update(self.headers) + if self.session: + heads["Mcp-Session-Id"] = self.session + request = urllib.request.Request(self.url, data=body, headers=heads, method="POST") + try: + response = urllib.request.urlopen(request, timeout=timeout) + except urllib.error.HTTPError as exc: + raise McpError("http %s" % exc.code) + except (urllib.error.URLError, OSError) as exc: + raise McpError("http failed: %s" % short_error(exc)) + raw = response.read(MCP_HTTP_MAX_BODY + 1) + if len(raw) > MCP_HTTP_MAX_BODY: + raise McpError("http body exceeds %d bytes" % MCP_HTTP_MAX_BODY) + session = response.headers.get("Mcp-Session-Id") or "" + if session: + self.session = session + if not raw: + return None + kind = (response.headers.get("Content-Type") or "").split(";")[0].strip() + if kind == "text/event-stream": + return self._parse_sse(raw.decode("utf-8", "replace")) + try: + return json.loads(raw.decode("utf-8", "replace")) + except ValueError: + raise McpError("http reply is not JSON") + + @staticmethod + def _parse_sse(text): + """Parse SSE frames, return the last data: JSON payload.""" + payload = None + for frame in text.split("\n\n"): + data = [line[5:].strip() for line in frame.split("\n") if line.startswith("data:")] + if not data: + continue + try: + payload = json.loads("\n".join(data)) + except ValueError: + continue + if payload is None: + raise McpError("SSE reply held no JSON") + return payload + + def send(self, message, timeout=MCP_CALL_TIMEOUT): + data = self._post(message, timeout) + if "id" not in message: + return None + if data is None: + raise McpError("empty http reply") + if not isinstance(data, dict) or data.get("id") != message["id"]: + raise McpError("http reply id mismatch") + return data + + def close(self): + self.session = "" + + +class McpClient(object): + """Minimal MCP client over one transport: handshake, tools, call, ping.""" + + def __init__(self, transport): + self.transport = transport + self.next_id = 0 + self.server_version = "" + self.server_info = {} + + def _request(self, method, params=None, timeout=MCP_CALL_TIMEOUT): + self.next_id += 1 + message = {"jsonrpc": "2.0", "id": self.next_id, "method": method} + if params is not None: + message["params"] = params + try: + reply = self.transport.send(message, timeout) + except McpError: + raise + if not isinstance(reply, dict): + raise McpError("reply is not an object") + if "error" in reply and isinstance(reply["error"], dict): + raise McpJsonRpcError(reply["error"].get("code", -32603), reply["error"].get("message", "error")) + if "result" not in reply: + raise McpError("reply has neither result nor error") + return reply["result"] + + def _notify(self, method, params=None): + message = {"jsonrpc": "2.0", "method": method} + if params is not None: + message["params"] = params + self.transport.send(message) + + def initialize(self): + try: + result = self._request("initialize", {"protocolVersion": MCP_PROTOCOL_VERSION, "capabilities": {}, + "clientInfo": {"name": MCP_CLIENT_NAME, "version": VERSION}}, + timeout=MCP_CONNECT_TIMEOUT) + except McpError as exc: + raise McpError("initialize failed: %s" % exc) + if not isinstance(result, dict) or not isinstance(result.get("protocolVersion"), str): + raise McpError("initialize: bad handshake reply") + if not isinstance(result.get("capabilities"), dict) or "tools" not in result["capabilities"]: + raise McpError("initialize: server offers no tools capability") + self.server_version = result["protocolVersion"] + self.server_info = result.get("serverInfo") or {} + self._notify("notifications/initialized") + return result + + def list_tools(self): + tools = [] + cursor = None + for _page in range(20): + params = {} if cursor is None else {"cursor": cursor} + result = self._request("tools/list", params) + if not isinstance(result, dict) or not isinstance(result.get("tools"), list): + raise McpError("tools/list: bad reply") + tools.extend(item for item in result["tools"] if isinstance(item, dict) and item.get("name")) + cursor = result.get("nextCursor") + if not cursor: + break + return tools + + def call_tool(self, name, args=None): + try: + return self._request("tools/call", {"name": name, "arguments": args or {}}) + except McpError as exc: + try: + self._notify("notifications/cancelled", {"requestId": self.next_id, "reason": "timeout"}) + except McpError: + pass + raise McpError("call %s failed: %s" % (name, exc)) + + def ping(self): + self._request("ping", {}) + + def close(self): + self.transport.close() + + +def mcp_tool_hash(name, description, schema): + blob = "%s\n%s\n%s" % (name, description or "", json.dumps(schema or {}, sort_keys=True)) + return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16] + + +def mcp_bridge_name(server, tool): + return MCP_TOOL_PREFIX + mcp_sanitize_name(server) + "_" + mcp_sanitize_name(tool) + + +class McpBridge(object): + """Client-side bridge: vault servers in, lazy tool schemas and calls out.""" + + def __init__(self, app): + self.app = app + + @property + def store(self): + return self.app.store + + def lookup_var(self, name, profile): + if name in os.environ and os.environ[name]: + return os.environ[name] + try: + return self.store.load_secret(name, profile) + except (sqlite3.Error, OSError): + return None + + def expand_entry(self, row, profile): + """Expand placeholders in a stored server row. Returns (ready, missing).""" + missing = [] + ready = dict(row) + lookup = lambda name: self.lookup_var(name, profile) + if row["transport"] == "stdio": + argv = [row["command"]] + row["args"] + done = [] + for part in argv: + text, miss = mcp_expand_vars(part, lookup) + missing += miss + done.append(text) + ready["argv"] = done + env = {} + for key, value in row["env"].items(): + text, miss = mcp_expand_vars(value, lookup) + missing += miss + env[key] = text + ready["env"] = env + else: + url, miss = mcp_expand_vars(row["url"], lookup) + missing += miss + ready["url"] = url + heads = {} + for key, value in row["headers"].items(): + text, miss = mcp_expand_vars(value, lookup) + missing += miss + heads[key] = text + ready["headers"] = heads + return ready, sorted(set(missing)) + + def connect(self, row, profile): + ready, missing = self.expand_entry(row, profile) + if missing: + raise McpError("unresolved secrets: %s (set env or store_secret)" % ", ".join(missing)) + if row["transport"] == "stdio": + transport = McpStdioTransport().spawn(ready["argv"], ready["env"], row["cwd"]) + else: + transport = McpHttpTransport(ready["url"], ready["headers"]) + client = McpClient(transport) + try: + client.initialize() + except McpError: + client.close() + raise + return client + + def cached_tools(self, server): + return self.store.mcp_cached_tools(server) + + def schemas_for(self, profile): + """Lazy schemas for enabled servers from the vault cache (no spawn).""" + found = [] + for row in self.store.mcp_enabled_servers(profile): + tools = [item for item in self.cached_tools(row["name"]) if self.tool_allowed(row, item["tool"])] + for item in tools[:MCP_MAX_TOOLS]: + params = item["schema"] or {"type": "object", "properties": {}} + if not isinstance(params, dict) or params.get("type") != "object": + params = {"type": "object", "properties": {}} + desc = item["description"] or "MCP tool %s" % item["tool"] + found.append(tool_schema(mcp_bridge_name(row["name"], item["tool"]), + "MCP %s/%s: %s" % (row["name"], item["tool"], desc), + params.get("properties") or {}, params.get("required") or [])) + return found + + @staticmethod + def tool_allowed(row, tool): + if row["deny"] and tool in row["deny"]: + return False + if row["allow"] and tool not in row["allow"]: + return False + return True + + def resolve(self, name): + """Split a bridged tool name back to (server, tool).""" + if not name.startswith(MCP_TOOL_PREFIX): + return None + rest = name[len(MCP_TOOL_PREFIX):] + for row in self.store.mcp_list_servers(): + prefix = mcp_sanitize_name(row["name"]) + "_" + if rest.startswith(prefix): + return row["name"], rest[len(prefix):] + return None + + def check_drift(self, server, fresh, profile): + """Compare fresh tools/list against pins. Returns (ok, note).""" + pins = self.store.mcp_tool_pins(server) + cached = dict((item["tool"], item) for item in self.cached_tools(server)) + changed, added = [], [] + for item in fresh: + tool = item.get("name", "") + digest = mcp_tool_hash(tool, item.get("description", ""), item.get("inputSchema")) + if tool not in pins: + added.append(tool) + elif pins[tool] != digest: + changed.append(tool) + if not added and not changed: + return True, "" + detail = "MCP server '%s' changed tools" % server + if added: + detail += "; new: %s" % ", ".join(sorted(added)[:10]) + if changed: + detail += "; redefined: %s" % ", ".join(sorted(changed)[:10]) + if not self.app.ask_approval(detail + " (approve to accept and re-pin)"): + return False, "tool list changed and was not approved" + self.store.mcp_store_tools(server, fresh) + return True, "re-pinned %d tools" % len(fresh) + + def result_text(self, result): + parts = [] + for item in (result or {}).get("content", []) or []: + if not isinstance(item, dict): + continue + kind = item.get("type", "") + if kind == "text": + parts.append(item.get("text", "")) + elif kind in ("image", "audio"): + parts.append("[%s %d bytes b64]" % (kind, len(item.get("data", "")))) + elif kind == "resource_link": + parts.append("[resource %s]" % item.get("uri", "")) + else: + parts.append("[%s]" % (kind or "content")) + structured = (result or {}).get("structuredContent") + if structured is not None: + parts.append(json.dumps(structured)[:4000]) + return "\n".join(parts) + + def call(self, server, tool, args, profile): + """Full call path: enable check, connect, drift check, approval, call, cap, audit.""" + if not isinstance(args, dict): + return "error: MCP tool args must be an object" + row = self.store.mcp_get_server(server) + if row is None: + return "error: unknown MCP server '%s'" % server + if not self.store.mcp_server_enabled(server, profile): + return "error: MCP server '%s' is disabled for profile '%s'" % (server, profile) + if not self.tool_allowed(row, tool): + return "error: MCP tool '%s' is not enabled on '%s'" % (tool, server) + started = time.time() + try: + client = self.connect(row, profile) + except McpError as exc: + return "error: %s" % exc + try: + try: + fresh = client.list_tools() + except McpError as exc: + return "error: tools/list failed: %s" % exc + names = [item.get("name", "") for item in fresh] + if tool not in names: + return "error: MCP server '%s' has no tool '%s'" % (server, tool) + ok, note = self.check_drift(server, fresh, profile) + if not ok: + return "error: %s" % note + digest = hashlib.sha256(json.dumps(args, sort_keys=True).encode("utf-8")).hexdigest()[:16] + if not self.app.ask_approval("mcp call %s/%s (args %s)" % (server, tool, digest)): + return "denied by user" + try: + result = client.call_tool(tool, args) + except McpError as exc: + self.audit(server, tool, digest, "error", started, profile) + return "error: %s" % exc + text = self.result_text(result) + if (result or {}).get("isError"): + self.audit(server, tool, digest, "isError", started, profile) + return "error: %s" % (text or "tool failed") + if len(text) > MCP_RESULT_MAX: + mem_id = self.spill(server, tool, text, profile) + self.audit(server, tool, digest, "spilled", started, profile) + return "output spilled to record %s (%d bytes); first %d bytes follow:\n%s" % ( + mem_id, len(text), MCP_RESULT_MAX, text[:MCP_RESULT_MAX]) + self.audit(server, tool, digest, "ok", started, profile) + return text + finally: + client.close() + + def spill(self, server, tool, text, profile): + try: + record_id = self.store.add_record("output", "mcp %s/%s output" % (server, tool), text, ["mcp"], profile) + self.store.add_chunks("record", record_id, chunk_lines(text), profile) + return record_id + except (sqlite3.Error, OSError, ValueError): + return "mem:?" + + def audit(self, server, tool, digest, outcome, started, profile): + try: + actor = self.app.tools.audit_actor() if getattr(self.app, "tools", None) else "agent" + self.store.audit_event(actor, "mcp-call", "%s/%s" % (server, tool), + "%s args=%s %.1fs" % (outcome, digest, time.time() - started), + tags=["mcp"], profile=profile) + except (sqlite3.Error, OSError): + pass + + +class McpServeSession(object): + def __init__(self, agent, profile): + self.agent = agent + self.profile = profile + now = time.time() + self.created = now + self.last_seen = now + + +def mcp_serve_tools(): + """Native schemas plus chat/completions as MCP tool dicts.""" + tools = [{"name": schema["function"]["name"], "description": schema["function"]["description"], + "inputSchema": schema["function"]["parameters"]} for schema in TOOL_SCHEMAS] + tools.append({"name": "chat", "description": "Full tai agent turn with tools, on the session agent. History is kept per MCP session.", + "inputSchema": {"type": "object", + "properties": {"message": {"type": "string"}, "profile": {"type": "string"}, + "max_steps": {"type": "integer"}}, + "required": ["message"]}}) + tools.append({"name": "completions", "description": "One stateless model completion with the current model, no tools.", + "inputSchema": {"type": "object", + "properties": {"messages": {"type": "array", + "items": {"type": "object"}}}, + "required": ["messages"]}}) + return tools + + +class McpServeContext(object): + def __init__(self, config, store, key): + self.config = config + self.store = store + self.key = key + self.sessions = {} + + def authed(self, headers): + presented = (headers.get("Authorization") or headers.get("authorization") or "") + if presented.lower().startswith("bearer "): + presented = presented[7:] + return bool(self.key) and hmac.compare_digest(presented.strip(), self.key) + + def prune(self): + now = time.time() + for sid in [key for key, session in self.sessions.items() if now - session.last_seen > MCP_SESSION_IDLE]: + del self.sessions[sid] + + def agent_for(self, sid, profile): + session = self.sessions.get(sid) + if session is None or session.agent is None: + agent = Agent(self.config, self.store, yolo=True, auto=True) + agent.quiet = True + if session is None: + session = self.sessions[sid] = McpServeSession(agent, profile or "default") + else: + session.agent = agent + session.profile = profile or session.profile or "default" + session.last_seen = time.time() + return session.agent + + +def mcp_serve_handler(ctx): + tools = mcp_serve_tools() + + class Handler(http.server.BaseHTTPRequestHandler): + server_version = "tai-mcp/" + VERSION + + def log_message(self, *args): + sys.stderr.write("mcp %s %s\n" % (self.command, self.path.split("?")[0])) + + def _body(self): + try: + length = int(self.headers.get("Content-Length") or 0) + except ValueError: + length = 0 + if length > MCP_HTTP_MAX_BODY: + return None + return self.rfile.read(length) if length else b"" + + def _send_json(self, payload, session=""): + 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))) + if session: + self.send_header("Mcp-Session-Id", session) + self.end_headers() + self.wfile.write(raw) + + def _reply(self, answer, session=""): + if answer is None: + self.send_response(202) + self.send_header("Content-Length", "0") + self.end_headers() + return + self._send_json(answer, session) + + def do_GET(self): + self.send_response(405) + self.send_header("Content-Length", "0") + self.end_headers() + + def do_DELETE(self): + if self.path.split("?")[0] != "/mcp" or not ctx.authed(self.headers): + self.send_response(401 if not ctx.authed(self.headers) else 404) + self.send_header("Content-Length", "0") + self.end_headers() + return + sid = self.headers.get("Mcp-Session-Id") or "" + ctx.sessions.pop(sid, None) + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + def do_POST(self): + ctx.prune() + if self.path.split("?")[0] != "/mcp": + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + return + if not ctx.authed(self.headers): + self.send_response(401) + self.send_header("Content-Length", "0") + self.end_headers() + return + raw = self._body() + if raw is None: + self.send_response(413) + self.send_header("Content-Length", "0") + self.end_headers() + return + try: + message = json.loads(raw.decode("utf-8") or "{}") + except ValueError: + self.send_response(400) + self.send_header("Content-Length", "0") + self.end_headers() + return + if not isinstance(message, dict) or message.get("jsonrpc") != "2.0": + self.send_response(400) + self.send_header("Content-Length", "0") + self.end_headers() + return + sid = self.headers.get("Mcp-Session-Id") or "" + method = message.get("method", "") + if method != "initialize" and sid not in ctx.sessions: + self._reply({"jsonrpc": "2.0", "id": message.get("id"), + "error": {"code": -32603, "message": "unknown session, initialize first"}}) + return + answer, session = self.dispatch_message(message, sid) + self._reply(answer, session) + + def dispatch_message(self, message, sid): + method = message.get("method", "") + msg_id = message.get("id") + params = message.get("params") or {} + + def fail(code, text): + if msg_id is None: + return None, "" + return {"jsonrpc": "2.0", "id": msg_id, "error": {"code": code, "message": text}}, sid + + def ok(result): + if msg_id is None: + return None, "" + return {"jsonrpc": "2.0", "id": msg_id, "result": result}, sid + + if method == "initialize": + fresh = uuid.uuid4().hex + ctx.sessions[fresh] = McpServeSession(None, "") + return ({"jsonrpc": "2.0", "id": msg_id, + "result": {"protocolVersion": MCP_PROTOCOL_VERSION, + "capabilities": {"tools": {}}, + "serverInfo": {"name": MCP_CLIENT_NAME, "version": VERSION}}}, fresh) + if method == "notifications/initialized": + return None, "" + if method == "ping": + return ok({}) + if method == "tools/list": + return ok({"tools": tools}) + if method != "tools/call": + return fail(-32601, "unknown method %s" % method) + name = params.get("name", "") + args = params.get("arguments") or {} + if not isinstance(args, dict): + return fail(-32602, "arguments must be an object") + if name == "chat": + return self.run_chat(msg_id, sid, args) + if name == "completions": + return self.run_completions(msg_id, sid, args) + agent = ctx.agent_for(sid, "") + try: + text = agent.tools.dispatch(name, args) + except Denied as exc: + return ok({"content": [{"type": "text", "text": "denied: %s" % exc}], "isError": True}) + except Exception as exc: + return ok({"content": [{"type": "text", "text": "error: %s" % short_error(exc)}], "isError": True}) + text = text if isinstance(text, str) else json.dumps(text) + failed = text.startswith("error:") or text.startswith("denied") + result = {"content": [{"type": "text", "text": text}]} + if failed: + result["isError"] = True + return ok(result) + + def run_chat(self, msg_id, sid, args): + message = args.get("message", "") + if not isinstance(message, str) or not message.strip(): + return {"jsonrpc": "2.0", "id": msg_id, + "error": {"code": -32602, "message": "chat needs a message string"}}, sid + try: + steps = int(args.get("max_steps", 0) or 0) + except (TypeError, ValueError): + steps = 0 + agent = ctx.agent_for(sid, args.get("profile") or "") + agent.deadline = time.time() + MCP_CHAT_TIMEOUT + try: + text = agent.run_turn(message, capture=True, max_steps=steps or MAX_STEPS) + except BackendError as exc: + return {"jsonrpc": "2.0", "id": msg_id, + "result": {"content": [{"type": "text", "text": "error: %s" % short_error(exc)}], + "isError": True}}, sid + except Exception as exc: + return {"jsonrpc": "2.0", "id": msg_id, + "result": {"content": [{"type": "text", "text": "error: %s" % short_error(exc)}], + "isError": True}}, sid + finally: + agent.deadline = None + return {"jsonrpc": "2.0", "id": msg_id, + "result": {"content": [{"type": "text", "text": text}]}}, sid + + def run_completions(self, msg_id, sid, args): + messages = args.get("messages") + if not isinstance(messages, list) or not messages: + return {"jsonrpc": "2.0", "id": msg_id, + "error": {"code": -32602, "message": "completions needs a messages array"}}, sid + cleaned = [] + for item in messages: + if not isinstance(item, dict) or item.get("role") not in ("system", "user", "assistant"): + return {"jsonrpc": "2.0", "id": msg_id, + "error": {"code": -32602, "message": "bad message entry (role/content required)"}}, sid + if not isinstance(item.get("content"), str): + return {"jsonrpc": "2.0", "id": msg_id, + "error": {"code": -32602, "message": "bad message entry (role/content required)"}}, sid + cleaned.append({"role": item["role"], "content": item["content"]}) + try: + reply = ChatClient(ctx.config).complete(cleaned) + except BackendError as exc: + return {"jsonrpc": "2.0", "id": msg_id, + "result": {"content": [{"type": "text", "text": "error: %s" % short_error(exc)}], + "isError": True}}, sid + text = reply if isinstance(reply, str) else json.dumps(reply) + return {"jsonrpc": "2.0", "id": msg_id, + "result": {"content": [{"type": "text", "text": text}]}}, sid + + return Handler + + +def mcp_serve_main(args, key): + try: + port = int(getattr(args, "mcp_port", None) or os.environ.get("TAI_MCP_PORT") or MCP_DEFAULT_PORT) + except ValueError: + print("error: bad port, use 1-65535", file=sys.stderr) + return 2 + if not 1 <= port <= 65535: + print("error: bad port, use 1-65535", file=sys.stderr) + return 2 + host = getattr(args, "mcp_bind", None) or os.environ.get("TAI_MCP_BIND") or "127.0.0.1" + agent = boot(args) + ctx = McpServeContext(agent.config, agent.store, key) + server = http.server.HTTPServer((host, port), mcp_serve_handler(ctx)) + server.timeout = 30 + print(paint("serving MCP on http://%s:%d/mcp (%d tools + chat + completions)" % (host, port, len(TOOL_SCHEMAS)), Ansi.GREEN)) + print(paint("key is the authorization: served tools run yolo, no approvals asked", Ansi.YELLOW)) + try: + server.serve_forever() + except KeyboardInterrupt: + pass + server.server_close() + return 0 + + # CAP (CLI Agent Protocol), agent side and orchestrator side. See cap.md. # Phase 1: events, errors, manifest registry. Later phases add --cap-speak # plus the PTY/stream-json/ACP/A2A drivers and orchestration on top. @@ -37733,6 +38989,9 @@ def build_parser(): parser.add_argument("--cap-permission-mode", default=None, choices=("none", "confirm", "interactive"), help="speak session permission mode (overrides CAP_SESSION_CONFIG permission_mode)") parser.add_argument("--cap-max-turns", default=None, type=int, metavar="N", help="speak session turn cap (overrides CAP_SESSION_CONFIG max_turns)") parser.add_argument("--cap", nargs=argparse.REMAINDER, default=None, metavar="COMMAND", help="orchestrator: resolve, chat or run fleet commands (see cap.md)") + parser.add_argument("--serve-mcp", default=None, metavar="KEY", help="serve MCP over HTTP (Streamable HTTP): all native tools plus chat/completions, Bearer KEY required (see mcp.md)") + parser.add_argument("--mcp-port", default=None, type=int, metavar="PORT", help="serve-mcp port (default TAI_MCP_PORT else 8765)") + parser.add_argument("--mcp-bind", default=None, metavar="HOST", help="serve-mcp bind address (default TAI_MCP_BIND else 127.0.0.1)") parser.add_argument("--telegram", action="store_true") parser.add_argument("--scheduler", action="store_true") parser.add_argument("--install", action="store_true") @@ -37867,6 +39126,11 @@ def main(argv=None): return cap_speak_main(args) if args.cap is not None: return cap_cli_main(args.cap) + if args.serve_mcp is not None: + if not args.serve_mcp: + print("error: --serve-mcp needs a KEY clients must send as Bearer token", file=sys.stderr) + return 2 + return mcp_serve_main(args, args.serve_mcp) if args.update_models: return update_models_command(args) if args.demo: diff --git a/test_tai.py b/test_tai.py index aecc1b9..bf2abc8 100644 --- a/test_tai.py +++ b/test_tai.py @@ -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()