From e8a909bcd2d0583fe8207042c6f8f56fc8b95ece Mon Sep 17 00:00:00 2001 From: thinkinginmath Date: Mon, 21 Sep 2026 03:32:52 +0000 Subject: [PATCH 1/3] Design for an MCP server, so any agent client can reach a learner's lab An adapter process over the localhost protocol rather than a second runtime: the kernels, prepared labs and runs are process state, so importing them twice would give a tutor its own kernels and hide the cell the learner just ran. Records the parts that are not obvious from the code: local client tokens as a third principal (a process sends no Origin, so today it cannot authenticate at all), why lab runs must not block a tool call to completion, why ANSI stripping is the adapter's job and not a change to what the browser gets, and that a code-scope token is a revocation and audit story rather than a new boundary against someone who can already run Python as the user. Co-Authored-By: Claude Opus 5 (1M context) --- docs/mcp.md | 227 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 docs/mcp.md diff --git a/docs/mcp.md b/docs/mcp.md new file mode 100644 index 0000000..731641f --- /dev/null +++ b/docs/mcp.md @@ -0,0 +1,227 @@ +# MCP Server + +> [!NOTE] +> **Status: proposed.** Nothing in this document is implemented. It records the design and the +> reasons behind it, so the first pull request can be read against something. The endpoints and +> behavior described in [protocol.md](protocol.md) are what exists today. + +An MCP server would let any Model Context Protocol client — Claude Code, Claude Desktop, Cursor, +an agent gateway — run code in a learner's lab, start the lab's approved actions, and read what +came back. For an AI tutor this is the difference between guessing from what a page reports and +reading the learner's actual traceback, in the learner's actual environment. + +MCP is also the reason not to write one integration per client. One adapter reaches all of them. + +## Shape: an adapter, not a second runtime + +The MCP server is a thin process that talks to a running runtime over loopback HTTP. It owns no +kernels, no environments and no runs. + +``` +MCP client ──stdio──> agent-runtime mcp ──HTTP 127.0.0.1:9477──> agent-runtime serve + (adapter, stateless) (kernels, labs, runs) +``` + +This is the load-bearing decision. `kernel_manager` and `lab_manager` are process globals, and +runs live in memory (`labs.py`). A second process that imported them would get its *own* kernels: +the tutor would not see the cell the learner just ran in the browser, two processes would race to +create the same virtual environment and install the same kernel spec, and runs started through one +would be invisible to the other. An adapter has one source of truth by construction. + +When the runtime is not listening, the adapter's tools fail with a message saying to start +`agent-runtime serve`, rather than starting a rival copy. + +## Transport + +Ship **stdio** first. The client configuration is one line, every MCP client supports it, and it +needs no port, no CORS and no DNS-rebinding defense: + +```json +{ "mcpServers": { "agent-runtime": { "command": "agent-runtime", "args": ["mcp"] } } } +``` + +Streamable HTTP at `/mcp` inside the `serve` process is the natural second step. Note before +taking it that the CORS middleware allows any web origin (`server.py`), so that endpoint would be +reachable from any page; the bearer check remains the authorization boundary, but it deserves its +own review. Write the tool bodies against a small `RuntimeClient` interface so swapping transport +does not touch them. + +SDK notes, checked 2026-09-21: the Python package is `mcp` (`pip install "mcp[cli]"`), the server +class is `MCPServer` from `mcp.server`, tools are declared with the `@mcp.tool()` decorator over a +typed function, and the process runs with `mcp.run(transport="stdio")`. On stdio, stdout belongs +to the protocol: log through `logging`, which writes to stderr. + +## Authentication: local client tokens + +A local process is not a browser origin, and today it cannot authenticate at all. `_authenticate` +in `auth.py` reads the `Origin` header; a request without one falls through to the token check, +where `validate_token(token, "")` cannot match, because every token is bound to a paired origin. +The only way in is `--no-pairing`, which is not an answer. + +Add a third principal type with its own store, beside `paired_origins.json`: + +| Principal | Identified by | Granted by | +|---|---|---| +| Loopback page | `Origin` header on a literal loopback host | automatic | +| Paired site | origin-bound bearer token | terminal approval of the origin | +| **Local client** | **bearer token, no `Origin` header** | **`agent-runtime token create`** | + +```python +# auth.py +async def _authenticate(request: Request, authorization: str | None) -> tuple[str, str]: + origin = request.headers.get("origin", "") + if not settings.require_pairing: + return origin, "code" + if origin: + return await _authenticate_browser(origin, authorization) # today's path, unchanged + # No Origin: a local process, not a browser. + client = local_tokens.lookup(_bearer(authorization)) + if client is None: + raise HTTPException(status_code=401, detail="no local client token") + return f"local:{client.name}", client.scope +``` + +Two invariants keep the browser model intact: + +- a paired-origin token never validates as a local client token, and vice versa — separate stores, + looked up on separate branches; +- a local token presented *with* an `Origin` header is rejected. It is not origin-bound, so it must + never be usable from a page. + +Tokens are hashed before storage, like pairing tokens, in `/local_clients.json` +(`0600`). The CLI mirrors `agent-runtime pairing`: + +```bash +agent-runtime token create "Claude Desktop" # prints the token once; defaults to scope actions +agent-runtime token create "tutor" --scope code +agent-runtime token list +agent-runtime token revoke "Claude Desktop" +``` + +Add `GET /runtime/whoami`, returning the caller's principal and scope, so the adapter registers +only the tools its token can use instead of advertising tools that will answer `403`. + +## Tools + +Ten tools, each a verb a model already understands. Deliberately not a mirror of the REST surface: +an agent should not have to orchestrate prepare → run → poll by hand. + +| Tool | Arguments | Scope | Notes | +|---|---|---|---| +| `list_labs` | — | both | environments, live kernels, prepared commits | +| `run_python` | `lab_id`, `code`, `timeout_s` | code | creates the environment and kernel if missing | +| `interrupt_kernel` | `lab_id` | code | | +| `restart_kernel` | `lab_id` | code | | +| `install_packages` | `lab_id`, `packages` | code | returns the tail of pip's output | +| `prepare_lab` | `repo`, `commit` | both | blocks while the terminal asks for approval | +| `run_action` | `lab_id`, `commit`, `action`, `wait_s` | both | returns a run handle and output so far | +| `read_run` | `run_id`, `offset` | both | | +| `stop_run` | `run_id` | both | | +| `get_history` | `lab_id`, `n` | both | see below | + +```python +from mcp.server import MCPServer + +mcp = MCPServer("agent-runtime") + + +@mcp.tool() +async def run_python(lab_id: str, code: str, timeout_s: int = 60) -> list[Content]: + """Run Python in the learner's lab kernel. Variables persist between calls.""" + result = await runtime.post("/cell/run", {"lab_id": lab_id, "code": code}) + return render_outputs(result["outputs"]) +``` + +`run_action` must not block until the action finishes. Actions may run up to 7200 seconds +(`labs.py`), and a tool call held open that long will time out somewhere in the client. It returns +after `wait_s` (20 by default) with the run id, the status and the output so far; the model polls +`read_run` with `next_offset`. That is the same shape the course page already uses. + +## Rendering what comes back + +Three conversions decide whether any of this is usable, and all three belong in the adapter. + +**Strip ANSI.** Lab actions run with `FORCE_COLOR=1` and `CLICOLOR_FORCE=1` on purpose, because the +page that started them renders color (`labs.py`). A model receives escape codes as noise it pays +tokens for. Strip them in the adapter rather than changing what the browser gets. + +**Pass images through.** `display_data` messages are captured with their whole `data` dictionary +(`kernels/ipython.py`), so a matplotlib figure already arrives from `/cell/run` as base64 +`image/png`. Turning those into MCP image content lets a tutor see the learner's plot, which for a +course is a feature rather than a detail. + +**Bound the result.** Cap a tool result at roughly 8–16 KB, keeping the head and tail with a marker +between them. Nothing is lost: the runtime keeps the last 1,000,000 characters of each run and +serves them by offset, so a model that needs more can page through it. + +One rule that is easy to get backwards: **a traceback is a successful tool call.** The learner's +code raising is the signal the tutor exists to read, so it comes back as ordinary content with +`success: false`. MCP tool errors are reserved for the runtime being unreachable, a `403`, or a lab +that does not exist. + +## Session history + +`get_history` has nothing behind it today — the transcript of cells and outputs lives only in the +browser. Add a bounded ring buffer per lab to `KernelManager` (the last ~50 cells: code, outputs, +timestamps) and `GET /labs/{lab_id}/history`. It is the one non-trivial addition on the server +side, and it is what separates a tutor that can run code from a tutor that knows what the learner +has been doing. + +Expose it as a tool *and* as a `runtime://labs/{lab_id}/history` resource. A resource is the +semantically correct shape, but many clients surface resources only when the user attaches one by +hand, while tools are called on the model's own initiative. The tool is the one that will be used. + +## Layout + +``` +src/agent_runtime/mcp/ +├── __init__.py +├── server.py # the MCPServer instance and the tool definitions +├── client.py # RuntimeClient: httpx against 127.0.0.1:9477, token read from disk +└── render.py # outputs -> MCP content: ANSI stripping, truncation, images +src/agent_runtime/auth.py # the local client principal +src/agent_runtime/api/health.py # GET /runtime/whoami +src/agent_runtime/cli.py # the `mcp` and `token` subcommands +``` + +`mcp` goes in an optional dependency group so the base install stays small. The tool functions take +a client, so they unit-test against a fake one; keep a single integration test against a live +`serve`. + +## Security + +This belongs in [security.md](security.md) once it ships, stated plainly. + +A local client token with the `code` scope **is not a new authority boundary**. Anyone who can read +the token file can already run Python as the user. What the token buys is revocation and an audit +trail, not defense against someone who is already on the machine. + +The genuinely new exposure is **prompt injection**. An MCP client is a model reading untrusted +text — course pages, lab output, tracebacks from third-party packages — and then deciding what to +execute. The `actions` scope is a real mitigation here, not a formality: a client holding an +actions-scope token can be talked into running a lab action the user already approved, and nothing +else. So `agent-runtime token create` defaults to `--scope actions`, and `--scope code` is typed +deliberately, against the same blunt warning the pairing prompt prints. + +Route every MCP tool call through `events.py` with the principal attached, so that executions by +`local:claude-desktop` are distinguishable from a page's in the log. + +## Milestones + +1. **Adapter, `run_python`, local client tokens.** End to end, one client, real value. +2. **Lab tools and rendering** — ANSI, truncation, images. +3. **History** — the buffer, the endpoint, the tool and the resource. +4. **Streamable HTTP**, if clients ask for it. + +Add `"mcp"` to the capability list in `/runtime/info` when step 1 lands, and document the client +configuration in this file. + +## Open questions + +- **Elicitation.** MCP can ask the client's user to confirm an action. Approvals live in the + runtime's terminal today, which is the stronger position — the terminal is the user's, the client + may be remote. Worth revisiting only if the terminal proves unreachable in practice. +- **Streaming.** `/cell/run/stream` exists; the MCP equivalent is progress notifications. Polling is + enough for a first version. +- **Whether the adapter may start the runtime.** Convenient, and it puts a process launch behind a + tool call. Leaning no. From 7ab350e1261fe4ef8dc4ab30fab636c798667219 Mon Sep 17 00:00:00 2001 From: thinkinginmath Date: Mon, 21 Sep 2026 03:50:20 +0000 Subject: [PATCH 2/3] An MCP server, so an agent client can run code in a learner's lab Step 1 of docs/mcp.md: `agent-runtime mcp` runs a stdio MCP server exposing one tool, `run_python`. It holds no state - it is an HTTP client of a runtime that is already serving, because kernels and runs are process state and a second copy would give a tutor its own kernels, invisible to the learner's page. A process is not a browser origin, so it could not authenticate at all: no Origin header means no pairing to bind a token to. Local client tokens are that third principal, issued by `agent-runtime token create`, hashed at rest, and answering `local:` from the new `GET /runtime/whoami`. They default to the `actions` scope: an agent client is a model reading untrusted course pages and tracebacks, so `--scope code` is typed deliberately. The store is re-read whenever the file changes, which a live test caught: tokens are issued by the CLI while the server runs, so without it a new token would not work until a restart - and a revoked one would keep working. What the adapter converts, because the runtime serves a browser and a model is not one: ANSI stripped from IPython's colored tracebacks, figures passed through as images, results capped head-and-tail. A traceback is a successful call - it is the thing a tutor exists to read - so only an unreachable runtime, a refusal or an impossible lab id fail the tool. The tool list does not vary with the token's scope, though the design said it would: a client asks for it once, at connect time, when the runtime may not be up. Co-Authored-By: Claude Opus 5 (1M context) --- CLAUDE.md | 11 ++- README.md | 51 ++++++++-- docs/mcp.md | 101 +++++++++++-------- docs/protocol.md | 44 ++++++++- docs/security.md | 32 ++++++ pyproject.toml | 6 ++ src/agent_runtime/auth.py | 82 ++++++++++++---- src/agent_runtime/cli.py | 116 ++++++++++++++++++++++ src/agent_runtime/config.py | 5 + src/agent_runtime/local_tokens.py | 157 ++++++++++++++++++++++++++++++ src/agent_runtime/mcp/__init__.py | 6 ++ src/agent_runtime/mcp/client.py | 119 ++++++++++++++++++++++ src/agent_runtime/mcp/render.py | 135 +++++++++++++++++++++++++ src/agent_runtime/mcp/server.py | 130 +++++++++++++++++++++++++ src/agent_runtime/server.py | 24 ++++- tests/test_auth.py | 97 ++++++++++++++++++ tests/test_local_tokens.py | 155 +++++++++++++++++++++++++++++ tests/test_mcp_render.py | 125 ++++++++++++++++++++++++ tests/test_mcp_server.py | 138 ++++++++++++++++++++++++++ 19 files changed, 1467 insertions(+), 67 deletions(-) create mode 100644 src/agent_runtime/local_tokens.py create mode 100644 src/agent_runtime/mcp/__init__.py create mode 100644 src/agent_runtime/mcp/client.py create mode 100644 src/agent_runtime/mcp/render.py create mode 100644 src/agent_runtime/mcp/server.py create mode 100644 tests/test_local_tokens.py create mode 100644 tests/test_mcp_render.py create mode 100644 tests/test_mcp_server.py diff --git a/CLAUDE.md b/CLAUDE.md index f44f33b..22b6504 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -26,17 +26,23 @@ agent-runtime/ │ ├── serving.py # HTTP + HTTPS listeners with one shutdown │ ├── tls.py # Loopback certificate and macOS keychain trust (Safari) │ ├── labs.py # Named lab actions: lab.toml, pinned checkouts, approvals, runs +│ ├── local_tokens.py # Tokens for local processes, which have no origin to pair │ ├── api/ │ │ ├── health.py # Health check endpoints │ │ ├── kernel.py # Kernel lifecycle endpoints │ │ ├── labs.py # /labs: prepare, runs, stop │ │ └── execute.py # Code execution endpoints +│ ├── mcp/ # MCP server (optional dependency); see docs/mcp.md +│ │ ├── server.py # Tools, and the MCPServer instance +│ │ ├── client.py # HTTP client of a runtime already serving +│ │ └── render.py # Kernel outputs -> MCP content │ └── kernels/ │ ├── manager.py # Kernel lifecycle management │ └── ipython.py # IPython kernel wrapper ├── docs/ │ ├── protocol.md # REST API specification │ ├── security.md # Security model and threat analysis +│ ├── mcp.md # MCP server: design, status and what comes next │ └── integration.md # Browser integration guide ├── tests/ └── pyproject.toml @@ -47,7 +53,9 @@ agent-runtime/ ### Standalone Architecture Agent-runtime uses a public localhost HTTP protocol and has no dependency on a -specific agent SDK. See `docs/integration.md` for browser integration patterns. +specific agent SDK. See `docs/integration.md` for browser integration patterns, +and `docs/mcp.md` for the MCP server, which is an adapter over that same +protocol rather than a second runtime. ### Security Model @@ -55,6 +63,7 @@ specific agent SDK. See `docs/integration.md` for browser integration patterns. - Origin allowlist via pairing flow - Bearer token authentication - Per-lab dependency isolation; code itself is not sandboxed +- Local processes (an MCP client) authenticate with an issued token, not an origin See `docs/security.md` for the full threat model. diff --git a/README.md b/README.md index 004f4bb..65ac48b 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,7 @@ Agent Runtime is the execution layer that powers AI agent code execution. It pro - **Per-lab Python environments** - dependencies are separated between sessions - **Secure localhost protocol** - pairing-based authentication for browser-to-runtime communication - **Streaming execution** - real-time stdout/stderr with interrupt support +- **MCP server** - any agent client can run code in a lab, through `agent-runtime mcp` - **Observability hooks** - structured events for monitoring and debugging ## Trust Model @@ -239,20 +240,57 @@ await executeStreaming('my-session', 'for i in range(5): print(i)', (output) => | `AGENT_RUNTIME_REQUIRE_PAIRING` | `true` | Require origin approval | | `AGENT_RUNTIME_DEBUG` | `false` | Enable debug mode with hot reload | +## Agent Clients (MCP) + +Agent Runtime speaks the Model Context Protocol, so an agent client - Claude Code, Claude +Desktop, Cursor - can run code in a learner's lab and read what came back: the real traceback, +in the real environment. Install the extra and issue that client a token: + +```bash +uv pip install 'agent-runtime[mcp]' +agent-runtime token create "Claude Desktop" --scope code +``` + +The token is shown once. Put it in the client's own configuration: + +```json +{ + "mcpServers": { + "agent-runtime": { + "command": "agent-runtime", + "args": ["mcp"], + "env": { "AGENT_RUNTIME_TOKEN": "" } + } + } +} +``` + +With `agent-runtime serve` running, the client gets a `run_python` tool: a persistent kernel per +lab, figures returned as images, tracebacks returned as text to read rather than as errors. + +`agent-runtime token list` and `agent-runtime token revoke ` manage clients, and take +effect without restarting the runtime. The default scope is `actions`, which can start only the +named actions of labs you have approved; `--scope code` is what runs arbitrary Python, and is +worth the same thought as letting that client type into your terminal. + +See [docs/mcp.md](docs/mcp.md) for the design and what comes next. + ## API Overview ### Runtime Info ``` -GET /runtime/info +GET /runtime/info - version and capabilities +GET /runtime/whoami - who the runtime takes the caller to be, and its scope ``` -Returns: +`/runtime/info` returns: ```json { "runtime_version": "0.1.0", "protocol_version": "2025-01", - "capabilities": ["python", "jupyter", "local_fs"] + "capabilities": ["python", "jupyter", "local_fs", "lab_actions", "pairing_scopes", + "local_clients"] } ``` @@ -309,9 +347,10 @@ Configuration is stored in `~/.agent-runtime/`: ``` ~/.agent-runtime/ -├── config.toml # Runtime configuration -├── paired_origins.json # Approved origins -└── envs/ # Virtual environments +├── config.toml # Runtime configuration +├── paired_origins.json # Approved origins +├── local_clients.json # Tokens for local processes, such as an MCP client +└── envs/ # Virtual environments └── lab-{id}/ └── .venv/ ``` diff --git a/docs/mcp.md b/docs/mcp.md index 731641f..8c78f1d 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -1,9 +1,10 @@ # MCP Server > [!NOTE] -> **Status: proposed.** Nothing in this document is implemented. It records the design and the -> reasons behind it, so the first pull request can be read against something. The endpoints and -> behavior described in [protocol.md](protocol.md) are what exists today. +> **Status: step 1 is implemented**, and this document doubles as its design. Shipped: +> local client tokens, `GET /runtime/whoami`, `agent-runtime token`, `agent-runtime mcp` and +> its one tool, `run_python`. Steps 2 to 4 below are still proposals. Where the implementation +> departed from the design, the section says so. An MCP server would let any Model Context Protocol client — Claude Code, Claude Desktop, Cursor, an agent gateway — run code in a learner's lab, start the lab's approved actions, and read what @@ -37,9 +38,20 @@ Ship **stdio** first. The client configuration is one line, every MCP client sup needs no port, no CORS and no DNS-rebinding defense: ```json -{ "mcpServers": { "agent-runtime": { "command": "agent-runtime", "args": ["mcp"] } } } +{ + "mcpServers": { + "agent-runtime": { + "command": "agent-runtime", + "args": ["mcp"], + "env": { "AGENT_RUNTIME_TOKEN": "" } + } + } +} ``` +`AGENT_RUNTIME_URL` overrides where the adapter looks for the runtime; `--url` and `--token` +do the same on the command line. + Streamable HTTP at `/mcp` inside the `serve` process is the natural second step. Note before taking it that the CORS middleware allows any web origin (`server.py`), so that endpoint would be reachable from any page; the bearer check remains the authorization boundary, but it deserves its @@ -58,7 +70,7 @@ in `auth.py` reads the `Origin` header; a request without one falls through to t where `validate_token(token, "")` cannot match, because every token is bound to a paired origin. The only way in is `--no-pairing`, which is not an answer. -Add a third principal type with its own store, beside `paired_origins.json`: +So there is a third principal type, with its own store beside `paired_origins.json`: | Principal | Identified by | Granted by | |---|---|---| @@ -89,7 +101,12 @@ Two invariants keep the browser model intact: never be usable from a page. Tokens are hashed before storage, like pairing tokens, in `/local_clients.json` -(`0600`). The CLI mirrors `agent-runtime pairing`: +(`0600`). Unlike pairing, which happens through the running server, tokens are issued by a +separate process, so the server re-reads that file whenever it changes: a token issued while +the runtime is serving works immediately, and - the reason the check is on every lookup rather +than only on a miss - a revoked one stops working immediately. + +The CLI mirrors `agent-runtime pairing`: ```bash agent-runtime token create "Claude Desktop" # prints the token once; defaults to scope actions @@ -98,18 +115,26 @@ agent-runtime token list agent-runtime token revoke "Claude Desktop" ``` -Add `GET /runtime/whoami`, returning the caller's principal and scope, so the adapter registers -only the tools its token can use instead of advertising tools that will answer `403`. +`GET /runtime/whoami` returns the caller's principal and scope, for a client to check its token +with. + +The design had the adapter register only the tools its token can use. It does not: a client asks +for the tool list once, at connect time, when the runtime may not be running yet, and a list that +varies with whether the daemon happened to be up is worse than a stable one. A call the scope +forbids fails with the runtime's own message and the command that widens it. ## Tools Ten tools, each a verb a model already understands. Deliberately not a mirror of the REST surface: an agent should not have to orchestrate prepare → run → poll by hand. +`run_python` is implemented; the rest arrive with steps 2 and 3. `list_labs` waits for step 2 +because environments are not exposed over HTTP at all yet. + | Tool | Arguments | Scope | Notes | |---|---|---|---| | `list_labs` | — | both | environments, live kernels, prepared commits | -| `run_python` | `lab_id`, `code`, `timeout_s` | code | creates the environment and kernel if missing | +| `run_python` | `lab_id`, `code` | code | **shipped.** Creates the environment and kernel if missing | | `interrupt_kernel` | `lab_id` | code | | | `restart_kernel` | `lab_id` | code | | | `install_packages` | `lab_id`, `packages` | code | returns the tail of pip's output | @@ -119,18 +144,9 @@ an agent should not have to orchestrate prepare → run → poll by hand. | `stop_run` | `run_id` | both | | | `get_history` | `lab_id`, `n` | both | see below | -```python -from mcp.server import MCPServer - -mcp = MCPServer("agent-runtime") - - -@mcp.tool() -async def run_python(lab_id: str, code: str, timeout_s: int = 60) -> list[Content]: - """Run Python in the learner's lab kernel. Variables persist between calls.""" - result = await runtime.post("/cell/run", {"lab_id": lab_id, "code": code}) - return render_outputs(result["outputs"]) -``` +The SDK is `mcp` (`MCPServer` from `mcp.server`), an optional dependency: install it with +`pip install 'agent-runtime[mcp]'`. Nothing else in the runtime needs it, and a learner who only +opens a course page never runs it. `run_action` must not block until the action finishes. Actions may run up to 7200 seconds (`labs.py`), and a tool call held open that long will time out somewhere in the client. It returns @@ -141,15 +157,20 @@ after `wait_s` (20 by default) with the run id, the status and the output so far Three conversions decide whether any of this is usable, and all three belong in the adapter. -**Strip ANSI.** Lab actions run with `FORCE_COLOR=1` and `CLICOLOR_FORCE=1` on purpose, because the -page that started them renders color (`labs.py`). A model receives escape codes as noise it pays -tokens for. Strip them in the adapter rather than changing what the browser gets. +**Strip ANSI.** IPython colors its tracebacks, and lab actions run with `FORCE_COLOR=1` and +`CLICOLOR_FORCE=1` on purpose, because the page that started them renders color (`labs.py`). A +model receives escape codes as noise it pays tokens for. Strip them in the adapter rather than +changing what the browser gets. **Pass images through.** `display_data` messages are captured with their whole `data` dictionary (`kernels/ipython.py`), so a matplotlib figure already arrives from `/cell/run` as base64 `image/png`. Turning those into MCP image content lets a tutor see the learner's plot, which for a course is a feature rather than a detail. +Note while reading that code that `/cell/run` returns outputs shaped `{"type": ..., "content": +{...}}`, while `docs/protocol.md` and the streaming endpoint document the flatter nbformat-like +`{"output_type": ..., "text": ...}`. The renderer accepts both, but the two should be reconciled. + **Bound the result.** Cap a tool result at roughly 8–16 KB, keeping the head and tail with a marker between them. Nothing is lost: the runtime keeps the last 1,000,000 characters of each run and serves them by offset, so a model that needs more can page through it. @@ -176,21 +197,22 @@ hand, while tools are called on the model's own initiative. The tool is the one ``` src/agent_runtime/mcp/ ├── __init__.py -├── server.py # the MCPServer instance and the tool definitions -├── client.py # RuntimeClient: httpx against 127.0.0.1:9477, token read from disk -└── render.py # outputs -> MCP content: ANSI stripping, truncation, images -src/agent_runtime/auth.py # the local client principal -src/agent_runtime/api/health.py # GET /runtime/whoami -src/agent_runtime/cli.py # the `mcp` and `token` subcommands +├── server.py # the MCPServer instance and the tool definitions +├── client.py # RuntimeClient: httpx against 127.0.0.1:9477 +└── render.py # outputs -> MCP content: ANSI stripping, truncation, images +src/agent_runtime/local_tokens.py # the token store +src/agent_runtime/auth.py # the local client principal +src/agent_runtime/server.py # GET /runtime/whoami +src/agent_runtime/cli.py # the `mcp` and `token` subcommands ``` `mcp` goes in an optional dependency group so the base install stays small. The tool functions take -a client, so they unit-test against a fake one; keep a single integration test against a live -`serve`. +a client, so they are tested against an `httpx.MockTransport` rather than a live server - including +the invariant that the adapter sends no `Origin` header. ## Security -This belongs in [security.md](security.md) once it ships, stated plainly. +Stated plainly in [security.md](security.md). A local client token with the `code` scope **is not a new authority boundary**. Anyone who can read the token file can already run Python as the user. What the token buys is revocation and an audit @@ -208,13 +230,14 @@ Route every MCP tool call through `events.py` with the principal attached, so th ## Milestones -1. **Adapter, `run_python`, local client tokens.** End to end, one client, real value. -2. **Lab tools and rendering** — ANSI, truncation, images. +1. ~~**Adapter, `run_python`, local client tokens.**~~ Done. +2. **Lab tools and rendering** — the lab tools, `list_labs` and the environments it needs. 3. **History** — the buffer, the endpoint, the tool and the resource. 4. **Streamable HTTP**, if clients ask for it. -Add `"mcp"` to the capability list in `/runtime/info` when step 1 lands, and document the client -configuration in this file. +`/runtime/info` advertises `local_clients` rather than the `mcp` the design named: the capability +list tells a *page* what the runtime offers, and the MCP server is a separate process rather than +something a page can reach. ## Open questions @@ -224,4 +247,6 @@ configuration in this file. - **Streaming.** `/cell/run/stream` exists; the MCP equivalent is progress notifications. Polling is enough for a first version. - **Whether the adapter may start the runtime.** Convenient, and it puts a process launch behind a - tool call. Leaning no. + tool call. Still no: an unreachable runtime returns a message naming `agent-runtime serve`. +- **Audit events.** MCP calls reach the runtime as ordinary HTTP requests and are logged as such, + with `local:` as the principal, but they do not yet emit their own events. diff --git a/docs/protocol.md b/docs/protocol.md index fb48667..035a032 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -53,6 +53,25 @@ that resolve to loopback are not included. Non-loopback pairing requires HTTPS; tokens are never issued to a remote plain-HTTP origin. +### Local Clients + +A process on this machine - an MCP server, a script - sends no `Origin` header, so pairing has +nothing to bind a token to and nothing to show the user. Such a client holds a token issued on +the machine instead: + +```bash +agent-runtime token create "Claude Desktop" --scope code +``` + +It sends that token as `Authorization: Bearer ` with no `Origin`, and gets the scope the +token was created with. The two kinds of token are kept apart: a paired site's token is not a +local client token, and a local client token sent *with* an `Origin` is rejected, because it is +bound to no origin. Tokens are stored hashed in `/local_clients.json` (mode `0600`) +and can be listed and revoked with `agent-runtime token list` and `agent-runtime token revoke`, +which take effect without restarting the runtime. + +See [mcp.md](mcp.md) for the MCP server that uses this. + ## Endpoints ### Request Pairing @@ -91,10 +110,33 @@ Returns runtime capabilities and version. { "runtime_version": "0.1.0", "protocol_version": "2025-01", - "capabilities": ["python", "jupyter", "local_fs"] + "capabilities": ["python", "jupyter", "local_fs", "lab_actions", "pairing_scopes", + "local_clients"] } ``` +### Who Am I + +``` +GET /runtime/whoami +Authorization: Bearer +``` + +Who the runtime takes the caller to be, and what that principal may do. A client checks its +token with this instead of discovering the answer when a call fails. + +**Response:** +```json +{ + "principal": "local:Claude Desktop", + "scope": "code" +} +``` + +`principal` is the browser origin, or `local:` for a local client. It is `anonymous` only +when the runtime was started with `--no-pairing`, which identifies nobody. Returns `401` without +usable credentials and `403` for a token that does not match its origin. + ### Health Check ``` diff --git a/docs/security.md b/docs/security.md index 9ad0990..0b0abb3 100644 --- a/docs/security.md +++ b/docs/security.md @@ -177,6 +177,26 @@ run, down to the commit. It does not sandbox them. Approving `python demo.py` tr `demo.py` at that commit, with your permissions, like cloning the repository and running it yourself. A new commit needs a new approval. +### 9. Local Client Tokens + +A process on this machine has no origin to pair, so it authenticates with a token issued by +`agent-runtime token create ` and sent without an `Origin` header. Tokens are hashed +before storage in `/local_clients.json` (mode `0600`), and `agent-runtime token +list` and `token revoke` work against a running runtime, which re-reads the file whenever it +changes. + +**What this is not.** A `code`-scope local token is not a boundary against an attacker who can +already run code as the user: they could start a kernel themselves. What it buys is revocation, +a name per client in the logs, and a scope. + +**What the scope is for.** The new exposure is prompt injection, not token theft. An MCP client +is a model reading untrusted text - course pages, lab output, tracebacks from third-party +packages - and then deciding what to run. A client holding an `actions`-scope token can be +talked into starting a lab action the user already approved, and nothing else. That is why +`agent-runtime token create` defaults to `actions` and `--scope code` has to be typed. + +See [mcp.md](mcp.md). + ## Known Limitations ### 1. No Code Sandboxing @@ -205,6 +225,12 @@ Tokens are stored by the web application. If the web app is compromised, tokens - Tokens are origin-bound (can't be used from other origins) - Users can revoke pairings: `agent-runtime pairing revoke ` +### 4. Agent Clients Read Untrusted Text + +A local client with the `code` scope runs whatever its model decides to run, and that model +reads course pages, lab output and library tracebacks. Treat a `code`-scope token as equivalent +to letting that client type into your terminal, and prefer `actions` where it is enough. + ## Hardening Recommendations ### For Local Development @@ -249,6 +275,10 @@ agent-runtime tls remove # Forget every approved lab version (the next prepare asks again) rm ~/.agent-runtime/labs/approvals.json + +# List and revoke local clients (MCP servers and other local processes) +agent-runtime token list +agent-runtime token revoke "Claude Desktop" ``` ### Cleaning Up Environments @@ -273,6 +303,8 @@ rm -rf ~/.agent-runtime/envs/ - [ ] Input validation on all endpoints - [ ] Environments isolated per lab - [ ] User can revoke pairings +- [ ] User can revoke local client tokens +- [ ] Local client tokens are hashed, 0600, and re-read while serving - [ ] User can delete environments ## Reporting Security Issues diff --git a/pyproject.toml b/pyproject.toml index 94c5b8e..7013f9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,12 +34,18 @@ dependencies = [ ] [project.optional-dependencies] +# The MCP server (agent-runtime mcp); see docs/mcp.md. Optional: the runtime itself does not +# need it, and a learner who only opens a course page never runs it. +mcp = [ + "mcp>=2.2.0", +] dev = [ "pytest>=7.4.0", "pytest-asyncio>=0.23.0", "pytest-cov>=4.1.0", "ruff>=0.1.0", "mypy>=1.8.0", + "mcp>=2.2.0", ] [project.scripts] diff --git a/src/agent_runtime/auth.py b/src/agent_runtime/auth.py index 38d04a2..c5fb745 100644 --- a/src/agent_runtime/auth.py +++ b/src/agent_runtime/auth.py @@ -4,6 +4,7 @@ - Origin allowlist management - Pairing flow for new origins - Token-based authentication +- Tokens for local processes, which have no origin to pair (see local_tokens.py) """ import asyncio @@ -18,6 +19,7 @@ from fastapi import Header, HTTPException, Request from agent_runtime.config import settings +from agent_runtime.local_tokens import local_tokens # What a paired origin may do. "code" runs any Python through kernels (/kernel, /cell) and # includes lab actions; "actions" only runs the named actions of labs the user approved (/labs). @@ -227,33 +229,67 @@ def get_allowed_origins() -> list[str]: return pairing_manager.get_allowed_origins() -async def _authenticate(request: Request, authorization: str | None) -> tuple[str, str]: - """Return the request's origin and the scope it holds, or raise 401/403.""" - origin = request.headers.get("origin", "") - - if not settings.require_pairing: - return origin, "code" +def create_local_token(name: str, scope: str) -> str: + """Issue a token to a local process. Returns the plaintext, which is not stored.""" + if scope not in SCOPES: + raise ValueError(f"scope must be one of {', '.join(SCOPES)}") + return local_tokens.create(name, scope) - if is_loopback_origin(origin): - return origin, "code" - # Require authorization header +def _bearer(authorization: str | None) -> str: + """The token from an Authorization header, or raise 401.""" if not authorization: raise HTTPException( status_code=401, detail="Authorization header required", ) - # Parse Bearer token if not authorization.startswith("Bearer "): raise HTTPException( status_code=401, detail="Invalid authorization format. Use: Bearer ", ) - token = authorization[7:] + return authorization[7:] + + +async def _authenticate(request: Request, authorization: str | None) -> tuple[str, str]: + """Return the caller's principal and the scope it holds, or raise 401/403. + + A principal is a browser origin, or `local:` for a process on this machine. + """ + origin = request.headers.get("origin", "") + + if not settings.require_pairing: + return origin, "code" + + if origin: + return _authenticate_origin(origin, authorization) + + # No Origin header, so not a browser: a process on this machine, holding a token this + # machine issued. The two kinds of token live in separate stores and are looked up on + # separate branches, so a paired site's token is not a local client token, and a local + # client token sent with an Origin takes the branch above and fails there - it is bound + # to no origin. See local_tokens.py. + client = local_tokens.lookup(_bearer(authorization)) + if client is None: + raise HTTPException( + status_code=401, + detail=( + "Unknown local client token. Issue one with " + "'agent-runtime token create '." + ), + ) + + return f"local:{client.name}", client.scope + - if not pairing_manager.validate_token(token, origin): +def _authenticate_origin(origin: str, authorization: str | None) -> tuple[str, str]: + """Authenticate a browser origin: loopback is trusted, anything else needs its token.""" + if is_loopback_origin(origin): + return origin, "code" + + if not pairing_manager.validate_token(_bearer(authorization), origin): raise HTTPException( status_code=403, detail="Invalid or expired token", @@ -268,21 +304,29 @@ async def require_auth( ) -> str: """Dependency for code execution (kernels, cells): requires the "code" scope. - Returns the validated origin. + Returns the validated principal. """ - origin, scope = await _authenticate(request, authorization) + principal, scope = await _authenticate(request, authorization) if scope != "code": raise HTTPException( status_code=403, - detail=f"This origin is paired for '{scope}' only; running code needs 'code'", + detail=f"This client is paired for '{scope}' only; running code needs 'code'", ) - return origin + return principal async def require_lab_auth( request: Request, authorization: str | None = Header(None), ) -> str: - """Dependency for lab actions (/labs): either scope. Returns the validated origin.""" - origin, _scope = await _authenticate(request, authorization) - return origin + """Dependency for lab actions (/labs): either scope. Returns the validated principal.""" + principal, _scope = await _authenticate(request, authorization) + return principal + + +async def require_identity( + request: Request, + authorization: str | None = Header(None), +) -> tuple[str, str]: + """Dependency for /runtime/whoami: the caller's principal and scope, under either scope.""" + return await _authenticate(request, authorization) diff --git a/src/agent_runtime/cli.py b/src/agent_runtime/cli.py index dbf9fb1..8a5b510 100644 --- a/src/agent_runtime/cli.py +++ b/src/agent_runtime/cli.py @@ -4,6 +4,7 @@ import sys from agent_runtime import __version__ +from agent_runtime.auth import SCOPE_DESCRIPTIONS, SCOPES from agent_runtime.config import settings from agent_runtime.observability import setup_logging @@ -93,6 +94,39 @@ def main() -> None: pairing_revoke = pairing_subparsers.add_parser("revoke", help="Revoke a paired origin") pairing_revoke.add_argument("origin", help="Origin to revoke") + # token commands: bearer tokens for local processes, which have no origin to pair + token_parser = subparsers.add_parser( + "token", help="Manage tokens for local clients, such as an MCP server" + ) + token_subparsers = token_parser.add_subparsers(dest="token_command") + + token_create = token_subparsers.add_parser("create", help="Issue a token to a local client") + token_create.add_argument("name", help="What to call this client, e.g. 'Claude Desktop'") + token_create.add_argument( + "--scope", + choices=SCOPES, + default="actions", + help="What the client may do (default: actions, which cannot run arbitrary code)", + ) + token_subparsers.add_parser("list", help="List local clients") + token_revoke = token_subparsers.add_parser("revoke", help="Revoke a local client's token") + token_revoke.add_argument("name", help="The client to revoke") + + # mcp command + mcp_parser = subparsers.add_parser( + "mcp", help="Run the MCP server on stdio, for an agent client to connect to" + ) + mcp_parser.add_argument( + "--url", + default=None, + help="The runtime to talk to (default: $AGENT_RUNTIME_URL, else the local runtime)", + ) + mcp_parser.add_argument( + "--token", + default=None, + help="Local client token (default: $AGENT_RUNTIME_TOKEN)", + ) + args = parser.parse_args() if args.command == "serve": @@ -101,6 +135,10 @@ def main() -> None: run_env(args) elif args.command == "pairing": run_pairing(args) + elif args.command == "token": + run_token(args) + elif args.command == "mcp": + run_mcp(args) elif args.command == "tls": run_tls(args) else: @@ -299,5 +337,83 @@ def run_pairing(args: argparse.Namespace) -> None: sys.exit(1) +def run_token(args: argparse.Namespace) -> None: + """Handle token commands: bearer tokens for local processes.""" + from agent_runtime.auth import create_local_token + from agent_runtime.local_tokens import local_tokens + + if args.token_command == "create": + existed = any(client.name == args.name for client in local_tokens.list_clients()) + try: + token = create_local_token(args.name, args.scope) + except ValueError as error: + print(error) + sys.exit(1) + + if existed: + print(f"Replaced the token for {args.name!r}; the old one no longer works.") + print(f"Token for {args.name!r}, scope {args.scope}:") + print() + print(f" {token}") + print() + print("This is the only time it is shown. It grants the right to") + print(f"{SCOPE_DESCRIPTIONS[args.scope]}.") + print() + print("For an MCP client, put it in that client's own configuration:") + print() + print(' "agent-runtime": {') + print(' "command": "agent-runtime",') + print(' "args": ["mcp"],') + print(f' "env": {{"AGENT_RUNTIME_TOKEN": "{token}"}}') + print(" }") + if args.scope != "code": + print() + print("Note: running Python needs --scope code. This token can only start the") + print("named actions of labs you approve.") + + elif args.token_command == "list": + clients = local_tokens.list_clients() + if not clients: + print("No local clients. Issue a token with 'agent-runtime token create '.") + return + + print("Local clients:") + for client in clients: + created = local_tokens.created_at(client.name) + when = f", created {created[:10]}" if created else "" + print(f" - {client.name} (scope {client.scope}{when})") + + elif args.token_command == "revoke": + if local_tokens.revoke(args.name): + print(f"Revoked the token for {args.name!r}") + else: + print(f"No local client named {args.name!r}") + sys.exit(1) + + else: + print("Usage: agent-runtime token [create|list|revoke]") + sys.exit(1) + + +def run_mcp(args: argparse.Namespace) -> None: + """Run the MCP server on stdio. + + Started by an MCP client, not usually by hand: stdin and stdout carry the protocol. + """ + import os + + try: + from agent_runtime.mcp.server import serve + except ImportError as error: + print( + f"The MCP server needs the 'mcp' package ({error}). Install it with:\n" + " uv pip install 'agent-runtime[mcp]'", + file=sys.stderr, + ) + sys.exit(1) + + serve(base_url=args.url, token=args.token or os.environ.get("AGENT_RUNTIME_TOKEN")) + + if __name__ == "__main__": main() diff --git a/src/agent_runtime/config.py b/src/agent_runtime/config.py index 3134a2e..c372da6 100644 --- a/src/agent_runtime/config.py +++ b/src/agent_runtime/config.py @@ -86,6 +86,11 @@ def paired_origins_file(self) -> Path: """Path to the paired origins file.""" return self.runtime_dir / "paired_origins.json" + @property + def local_clients_file(self) -> Path: + """Tokens issued to local processes that are not browser origins (see local_tokens.py).""" + return self.runtime_dir / "local_clients.json" + def ensure_dirs(self) -> None: """Ensure all required directories exist.""" self.runtime_dir.mkdir(parents=True, exist_ok=True) diff --git a/src/agent_runtime/local_tokens.py b/src/agent_runtime/local_tokens.py new file mode 100644 index 0000000..ef169f9 --- /dev/null +++ b/src/agent_runtime/local_tokens.py @@ -0,0 +1,157 @@ +"""Bearer tokens for local processes, which are not browser origins. + +A browser is identified by the `Origin` header, and a pairing binds a token to that exact +value (see `auth.py`). A process on this machine — an MCP server, a script — sends no origin, +so pairing has nothing to bind to and nothing to show the user in a prompt. Such a client is +identified instead by holding a token this machine issued: + + agent-runtime token create "Claude Desktop" --scope code + +Tokens are stored hashed, like pairing tokens, in `/local_clients.json` (0600). +The plaintext is printed once, at creation, and goes in the client's own configuration. + +This is not a boundary against an attacker who already runs code as the user: they could start +a kernel themselves. What it buys is revocation, an audit trail per client, and a scope, so a +client meant only to run a lab's approved actions cannot be talked into running arbitrary +Python. See `docs/security.md`. +""" + +from __future__ import annotations + +import hashlib +import hmac +import json +import os +import secrets +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any + +from agent_runtime.config import settings + +NAME_MAX_LENGTH = 64 + + +@dataclass(frozen=True) +class LocalClient: + """A local process allowed to call the runtime, and what it may do.""" + + name: str + scope: str + + +def _hash(token: str) -> str: + return hashlib.sha256(token.encode()).hexdigest() + + +class LocalTokenStore: + """The local client tokens on this machine.""" + + def __init__(self) -> None: + self._clients: dict[str, dict[str, Any]] = {} + self._stamp: tuple[int, int, int] | None = None + self._load() + + def _stamp_of_file(self) -> tuple[int, int, int] | None: + try: + info = settings.local_clients_file.stat() + except OSError: + return None + return (info.st_mtime_ns, info.st_size, info.st_ino) + + def _load(self) -> None: + self._stamp = self._stamp_of_file() + path = settings.local_clients_file + if not path.exists(): + self._clients = {} + return + try: + with open(path) as f: + loaded = json.load(f) + except (json.JSONDecodeError, OSError): + self._clients = {} + return + self._clients = loaded if isinstance(loaded, dict) else {} + + def _refresh(self) -> None: + """Pick up what another process wrote. + + The CLI issues and revokes tokens while the runtime is serving, in a different process. + Without this the server would keep whatever it read at startup: a token just created + would not work until a restart, and - the reason this is a stat on every lookup rather + than only on a miss - a token just revoked would keep working. + """ + if self._stamp_of_file() != self._stamp: + self._load() + + def _save(self) -> None: + settings.ensure_dirs() + path = settings.local_clients_file + + # Created unreadable by anyone else, before a token hash is ever written into it. + def private(file: str, flags: int) -> int: + return os.open(file, flags, 0o600) + + with open(path, "w", opener=private) as f: + json.dump(self._clients, f, indent=2) + os.chmod(path, 0o600) # a file from an earlier version may be more permissive + self._stamp = self._stamp_of_file() + + def create(self, name: str, scope: str) -> str: + """Issue a token for `name`, replacing any token it already had. + + Returns the plaintext token; only its hash is kept. + """ + name = name.strip() + if not name or len(name) > NAME_MAX_LENGTH: + raise ValueError(f"a client name must be 1-{NAME_MAX_LENGTH} characters") + + # Whole-file rewrite, so start from what is on disk. Two people issuing tokens at the + # same second is not a workload worth locking for. + self._refresh() + + token = secrets.token_urlsafe(32) + self._clients[name] = { + "created_at": datetime.now(UTC).isoformat(), + "token_hash": _hash(token), + "scope": scope, + } + self._save() + return token + + def lookup(self, token: str | None) -> LocalClient | None: + """The client holding this token, or None.""" + if not token: + return None + + self._refresh() + presented = _hash(token) + for name, record in self._clients.items(): + stored = record.get("token_hash", "") + if stored and hmac.compare_digest(stored, presented): + return LocalClient(name=name, scope=str(record.get("scope", "actions"))) + return None + + def revoke(self, name: str) -> bool: + """Forget a client's token. Returns False if there was no such client.""" + self._refresh() + if name not in self._clients: + return False + del self._clients[name] + self._save() + return True + + def list_clients(self) -> list[LocalClient]: + self._refresh() + return [ + LocalClient(name=name, scope=str(record.get("scope", "actions"))) + for name, record in sorted(self._clients.items()) + ] + + def created_at(self, name: str) -> str: + self._refresh() + return str(self._clients.get(name, {}).get("created_at", "")) + + +# Global local token store +local_tokens = LocalTokenStore() diff --git a/src/agent_runtime/mcp/__init__.py b/src/agent_runtime/mcp/__init__.py new file mode 100644 index 0000000..6455ce9 --- /dev/null +++ b/src/agent_runtime/mcp/__init__.py @@ -0,0 +1,6 @@ +"""An MCP server for the runtime, so agent clients can run code in a learner's lab. + +See `docs/mcp.md` for the design. The adapter holds no state of its own: it is an HTTP client +of a runtime already serving on loopback. Importing `agent_runtime.mcp.server` requires the +optional `mcp` dependency; nothing else in the runtime does. +""" diff --git a/src/agent_runtime/mcp/client.py b/src/agent_runtime/mcp/client.py new file mode 100644 index 0000000..c7e806b --- /dev/null +++ b/src/agent_runtime/mcp/client.py @@ -0,0 +1,119 @@ +"""An HTTP client for a runtime that is already serving on loopback. + +The MCP adapter owns no kernels, environments or runs; it asks the runtime process that does. +That is the point: a second process importing `kernel_manager` would get its own kernels, so a +tutor would not see the cell the learner just ran in the browser (see `docs/mcp.md`). + +This client is not a browser. It sends no `Origin` header, which is exactly how the runtime +tells it apart from a page: it authenticates with a local client token instead of a paired +origin's (see `agent_runtime.local_tokens`). +""" + +from __future__ import annotations + +import os +from typing import Any + +import httpx + +from agent_runtime.config import settings + +DEFAULT_TIMEOUT = 120.0 + + +class RuntimeUnavailableError(Exception): + """No runtime is listening.""" + + +class RuntimeRejectedError(Exception): + """The runtime answered, and said no.""" + + def __init__(self, status: int, detail: str) -> None: + super().__init__(detail) + self.status = status + self.detail = detail + + +def default_base_url() -> str: + """Where the runtime is expected to be listening.""" + from_env = os.environ.get("AGENT_RUNTIME_URL") + if from_env: + return from_env.rstrip("/") + return f"http://{settings.host}:{settings.port}" + + +class RuntimeClient: + """Calls the localhost protocol on behalf of an MCP client.""" + + def __init__( + self, + base_url: str | None = None, + token: str | None = None, + timeout: float = DEFAULT_TIMEOUT, + transport: httpx.AsyncBaseTransport | None = None, + ) -> None: + self.base_url = (base_url or default_base_url()).rstrip("/") + self._token = token + self._timeout = timeout + self._transport = transport + self._client: httpx.AsyncClient | None = None + + def _http(self) -> httpx.AsyncClient: + if self._client is None: + headers = {"Accept": "application/json"} + if self._token: + headers["Authorization"] = f"Bearer {self._token}" + self._client = httpx.AsyncClient( + base_url=self.base_url, + headers=headers, + timeout=self._timeout, + transport=self._transport, + ) + return self._client + + async def request(self, method: str, path: str, json: Any = None) -> Any: + """Call the runtime, or raise RuntimeUnavailableError / RuntimeRejectedError.""" + try: + response = await self._http().request(method, path, json=json) + except httpx.RequestError as error: + raise RuntimeUnavailableError(str(error)) from error + + if response.status_code >= 400: + raise RuntimeRejectedError(response.status_code, _detail(response)) + + if not response.content: + return None + try: + return response.json() + except ValueError: + return response.text + + async def get(self, path: str) -> Any: + return await self.request("GET", path) + + async def post(self, path: str, json: Any = None) -> Any: + return await self.request("POST", path, json=json) + + async def whoami(self) -> dict[str, str] | None: + """The principal and scope this client holds, or None if the runtime cannot say.""" + try: + identity = await self.get("/runtime/whoami") + except (RuntimeUnavailableError, RuntimeRejectedError): + return None + return identity if isinstance(identity, dict) else None + + async def aclose(self) -> None: + if self._client is not None: + await self._client.aclose() + self._client = None + + +def _detail(response: httpx.Response) -> str: + """The runtime's own error message, which is worth more than the status code alone.""" + try: + body = response.json() + except ValueError: + return response.text.strip() or f"HTTP {response.status_code}" + if isinstance(body, dict) and body.get("detail"): + return str(body["detail"]) + return response.text.strip() or f"HTTP {response.status_code}" diff --git a/src/agent_runtime/mcp/render.py b/src/agent_runtime/mcp/render.py new file mode 100644 index 0000000..643319a --- /dev/null +++ b/src/agent_runtime/mcp/render.py @@ -0,0 +1,135 @@ +"""Turn a kernel's outputs into content an MCP client can read. + +Three conversions decide whether any of this is usable, and all three belong here rather than +in the runtime, which serves a browser (see `docs/mcp.md`): + +- **ANSI.** IPython colors its tracebacks, and lab actions run with `FORCE_COLOR=1` because the + page that started them renders color. A model reads escape codes as noise it pays for. +- **Images.** A figure arrives as base64 `image/png` in a `display_data` output. Passed through + as image content, a tutor can see the learner's plot. +- **Size.** A tool result is capped, keeping the head and tail, because an accidental + `print(range(10**6))` should not fill the client's context. + +Kernel outputs come in two shapes: `{"type": ..., "content": {...}}` from `/cell/run`, and the +flatter `{"output_type": ..., "text": ...}` of the streaming endpoint and `docs/protocol.md`. +Both are accepted. +""" + +from __future__ import annotations + +import re +from typing import Any + +from mcp.types import ContentBlock, ImageContent, TextContent + +# CSI sequences, which is what colored terminal output and IPython tracebacks contain. +ANSI_PATTERN = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]") + +MAX_TEXT_CHARS = 12_000 +MAX_IMAGES = 4 +IMAGE_MIME_TYPES = ("image/png", "image/jpeg", "image/gif", "image/webp") + + +def strip_ansi(text: str) -> str: + """Remove terminal color codes.""" + return ANSI_PATTERN.sub("", text) + + +def truncate(text: str, limit: int = MAX_TEXT_CHARS) -> str: + """Keep the head and the tail, which is where the answer usually is.""" + if len(text) <= limit: + return text + head = (limit * 2) // 3 + tail = limit - head + omitted = len(text) - head - tail + return f"{text[:head]}\n\n[... {omitted:,} characters omitted ...]\n\n{text[-tail:]}" + + +def render_execution(result: dict[str, Any]) -> list[ContentBlock]: + """Render a `/cell/run` response. + + A traceback is not a failure of this tool: the learner's code raising is the thing a tutor + exists to read, so it comes back as ordinary content. + """ + texts: list[str] = [] + images: list[ImageContent] = [] + + for output in result.get("outputs") or []: + if not isinstance(output, dict): + continue + kind = str(output.get("type") or output.get("output_type") or "") + content = output.get("content") + body: dict[str, Any] = content if isinstance(content, dict) else output + + if kind == "stream": + texts.append(str(body.get("text", ""))) + elif kind in ("result", "execute_result", "display_data"): + data = body.get("data") + if not isinstance(data, dict): + continue + image = _image(data) + if image is not None: + images.append(image) + else: + texts.append(_plain_text(data)) + elif kind == "error": + texts.append(_error_text(body)) + + error = result.get("error") + if isinstance(error, dict): + texts.append(_error_text(error)) + + blocks: list[ContentBlock] = [] + text = strip_ansi("".join(texts)).strip() + if text: + blocks.append(TextContent(type="text", text=truncate(text))) + + dropped = max(0, len(images) - MAX_IMAGES) + blocks.extend(images[:MAX_IMAGES]) + if dropped: + blocks.append( + TextContent(type="text", text=f"[{dropped} further image(s) not shown]") + ) + + if not blocks: + # Content of some kind, always: an empty result reads like a broken tool. + count = result.get("execution_count") + ran = f" (execution count {count})" if count is not None else "" + blocks.append(TextContent(type="text", text=f"The code ran and produced no output{ran}.")) + + return blocks + + +def _image(data: dict[str, Any]) -> ImageContent | None: + for mime in IMAGE_MIME_TYPES: + payload = data.get(mime) + if isinstance(payload, str) and payload.strip(): + # Jupyter wraps base64 payloads; a client wants them unwrapped. + return ImageContent( + type="image", + data="".join(payload.split()), + mime_type=mime, + ) + return None + + +def _plain_text(data: dict[str, Any]) -> str: + """The text of a result, preferring the representation meant for reading.""" + for mime in ("text/plain", "text/markdown"): + value = data.get(mime) + if isinstance(value, str): + return value if value.endswith("\n") else value + "\n" + for mime, value in data.items(): + if isinstance(value, str) and mime.startswith("text/"): + return value if value.endswith("\n") else value + "\n" + return "" + + +def _error_text(error: dict[str, Any]) -> str: + """A traceback if the kernel sent one, else the exception line.""" + traceback = error.get("traceback") + if isinstance(traceback, list) and traceback: + return "\n".join(str(line) for line in traceback) + "\n" + ename = str(error.get("ename", "Error")) + evalue = str(error.get("evalue", "")) + return f"{ename}: {evalue}\n" if evalue else f"{ename}\n" diff --git a/src/agent_runtime/mcp/server.py b/src/agent_runtime/mcp/server.py new file mode 100644 index 0000000..521a518 --- /dev/null +++ b/src/agent_runtime/mcp/server.py @@ -0,0 +1,130 @@ +"""The MCP server: an adapter between an MCP client and a runtime serving on loopback. + +Step 1 of `docs/mcp.md`: one tool, `run_python`. The lab tools arrive with step 2. + +Two decisions worth knowing when reading this: + +- The tool list does not depend on the token's scope. The design considered registering only + the tools a token can use, but a client asks for the list once, at connect time, when the + runtime may not be running yet. A list that varies with whether the daemon happened to be up + is worse than a stable one, so a call that the scope forbids fails with a message saying so. +- Nothing is called at import. An MCP client starts this process when it starts itself, often + long before the user asks for anything. + +On stdio, stdout carries the protocol: log through `logging`, which writes to stderr. +""" + +from __future__ import annotations + +import logging +import re +from typing import Any + +from mcp.server import MCPServer +from mcp.server.mcpserver.exceptions import ToolError +from mcp.types import ContentBlock + +from agent_runtime import __version__ +from agent_runtime.envs import LAB_ID_PATTERN +from agent_runtime.mcp.client import RuntimeClient, RuntimeRejectedError, RuntimeUnavailableError +from agent_runtime.mcp.render import render_execution + +logger = logging.getLogger(__name__) + +INSTRUCTIONS = """\ +Runs code in a learner's lab through an Agent Runtime on this machine. + +A lab is a Python environment with a persistent kernel: variables, imports and open files +survive between calls, like cells in a notebook. Labs are separate from each other. + +Nothing here is sandboxed. Code runs with the user's own permissions, on their own files, so +treat it as you would treat typing into their terminal.""" + +mcp: MCPServer = MCPServer( + "agent-runtime", + title="Agent Runtime", + version=__version__, + instructions=INSTRUCTIONS, +) + +_runtime: RuntimeClient | None = None + + +def _client() -> RuntimeClient: + if _runtime is None: # pragma: no cover - only reachable if a tool is called before serve() + raise ToolError("This MCP server was not started by 'agent-runtime mcp'.") + return _runtime + + +async def _call(method: str, path: str, json: Any = None) -> Any: + """Call the runtime, turning its refusals into something a model can act on.""" + client = _client() + try: + return await client.request(method, path, json=json) + except RuntimeUnavailableError as error: + logger.info("runtime unavailable at %s: %s", client.base_url, error) + raise ToolError( + f"No Agent Runtime is listening at {client.base_url}. " + "Ask the user to start it with 'agent-runtime serve'." + ) from error + except RuntimeRejectedError as error: + logger.info("runtime rejected %s %s: %s", method, path, error.detail) + raise ToolError(_rejection(error)) from error + + +def _rejection(error: RuntimeRejectedError) -> str: + if error.status == 401: + return ( + f"{error.detail} The user issues one with " + "'agent-runtime token create --scope code' and puts it in this client's " + "configuration as AGENT_RUNTIME_TOKEN." + ) + if error.status == 403: + return ( + f"{error.detail}. The user can widen it with " + "'agent-runtime token create --scope code', which replaces the token." + ) + return error.detail + + +@mcp.tool() +async def run_python(lab_id: str, code: str) -> list[ContentBlock]: + """Run Python in a lab's kernel on the user's machine, and return what it printed. + + The kernel is persistent, so state survives between calls: variables stay bound, imports + stay imported, and a name defined in an earlier call can be used in a later one. Figures + come back as images. + + A traceback is a normal result, not an error of this tool - read it and say what went + wrong. The first call for a lab can take several seconds, because its virtual environment + and kernel are created then. + + Args: + lab_id: which lab (session) to run in. Ask the user if you do not know it. + code: Python source, as you would type into a notebook cell. + """ + if re.fullmatch(LAB_ID_PATTERN, lab_id) is None: + raise ToolError( + "lab_id must be 1-64 letters, digits, dots, underscores or hyphens, " + "starting with a letter or digit." + ) + + result = await _call("POST", "/cell/run", {"lab_id": lab_id, "code": code}) + if not isinstance(result, dict): + raise ToolError("The runtime returned something other than an execution result.") + return render_execution(result) + + +def serve(base_url: str | None = None, token: str | None = None) -> None: + """Run the MCP server on stdio until the client disconnects.""" + global _runtime + + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + _runtime = RuntimeClient(base_url=base_url, token=token) + logger.info("agent-runtime MCP server talking to %s", _runtime.base_url) + if token is None: + logger.info( + "No token set. Loopback pages are trusted without one, but this process is not a " + "page: set AGENT_RUNTIME_TOKEN unless the runtime runs with --no-pairing." + ) + mcp.run(transport="stdio") diff --git a/src/agent_runtime/server.py b/src/agent_runtime/server.py index 83b5258..c6f0942 100644 --- a/src/agent_runtime/server.py +++ b/src/agent_runtime/server.py @@ -3,12 +3,13 @@ from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import asynccontextmanager -from fastapi import FastAPI, Request, Response +from fastapi import Depends, FastAPI, Request, Response from fastapi.middleware.cors import CORSMiddleware from agent_runtime import __protocol_version__, __version__ from agent_runtime import labs as lab_runtime from agent_runtime.api import execute, health, kernel, labs, pairing +from agent_runtime.auth import require_identity from agent_runtime.config import settings from agent_runtime.kernels.manager import kernel_manager @@ -68,10 +69,29 @@ async def runtime_info() -> dict[str, object]: return { "runtime_version": __version__, "protocol_version": __protocol_version__, - "capabilities": ["python", "jupyter", "local_fs", "lab_actions", "pairing_scopes"], + "capabilities": [ + "python", + "jupyter", + "local_fs", + "lab_actions", + "pairing_scopes", + "local_clients", + ], } +@app.get("/runtime/whoami") +async def whoami(identity: tuple[str, str] = Depends(require_identity)) -> dict[str, str]: + """Who the runtime takes the caller to be, and what that principal may do. + + A client uses this to check its token and to report the scope it holds before a call + fails on one. + """ + principal, scope = identity + # An empty principal means nobody was identified, which only happens under --no-pairing. + return {"principal": principal or "anonymous", "scope": scope} + + def main() -> None: """Run the server.""" import uvicorn diff --git a/tests/test_auth.py b/tests/test_auth.py index 31d0b96..a3dbdd7 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -252,3 +252,100 @@ async def test_require_auth_invalid_format(self): assert exc_info.value.status_code == 401 assert "Bearer" in exc_info.value.detail + + +class TestLocalClientAuth: + """A process on this machine sends no Origin; it holds a token this machine issued.""" + + @staticmethod + def _request(origin: str | None = None): + from fastapi import Request + + request = MagicMock(spec=Request) + request.headers = {"origin": origin} if origin else {} + return request + + @staticmethod + def _store(client=None): + store = MagicMock() + store.lookup.return_value = client + return store + + @pytest.mark.asyncio + async def test_a_local_token_names_a_principal_and_its_scope(self): + from agent_runtime.auth import _authenticate + from agent_runtime.local_tokens import LocalClient + + store = self._store(LocalClient(name="Claude Desktop", scope="code")) + with patch("agent_runtime.auth.settings") as mock_settings: + mock_settings.require_pairing = True + with patch("agent_runtime.auth.local_tokens", store): + principal, scope = await _authenticate(self._request(), "Bearer abc") + + assert principal == "local:Claude Desktop" + assert scope == "code" + + @pytest.mark.asyncio + async def test_an_unknown_local_token_is_refused(self): + from fastapi import HTTPException + + from agent_runtime.auth import _authenticate + + with patch("agent_runtime.auth.settings") as mock_settings: + mock_settings.require_pairing = True + with patch("agent_runtime.auth.local_tokens", self._store(None)): + with pytest.raises(HTTPException) as error: + await _authenticate(self._request(), "Bearer abc") + + assert error.value.status_code == 401 + assert "token create" in error.value.detail + + @pytest.mark.asyncio + async def test_a_process_without_a_token_is_refused(self): + from fastapi import HTTPException + + from agent_runtime.auth import _authenticate + + with patch("agent_runtime.auth.settings") as mock_settings: + mock_settings.require_pairing = True + with patch("agent_runtime.auth.local_tokens", self._store(None)): + with pytest.raises(HTTPException) as error: + await _authenticate(self._request(), None) + + assert error.value.status_code == 401 + + @pytest.mark.asyncio + async def test_a_local_token_is_not_usable_from_a_page(self): + """It is bound to no origin, so a request carrying one takes the browser branch.""" + from fastapi import HTTPException + + from agent_runtime.auth import _authenticate + from agent_runtime.local_tokens import LocalClient + + store = self._store(LocalClient(name="tutor", scope="code")) + with patch("agent_runtime.auth.settings") as mock_settings: + mock_settings.require_pairing = True + with patch("agent_runtime.auth.local_tokens", store): + with pytest.raises(HTTPException) as error: + await _authenticate(self._request("https://example.com"), "Bearer abc") + + assert error.value.status_code == 403 + store.lookup.assert_not_called() + + @pytest.mark.asyncio + async def test_an_actions_client_may_not_run_code(self): + from fastapi import HTTPException + + from agent_runtime.auth import require_auth, require_lab_auth + from agent_runtime.local_tokens import LocalClient + + store = self._store(LocalClient(name="tutor", scope="actions")) + with patch("agent_runtime.auth.settings") as mock_settings: + mock_settings.require_pairing = True + with patch("agent_runtime.auth.local_tokens", store): + assert await require_lab_auth(self._request(), "Bearer abc") == "local:tutor" + with pytest.raises(HTTPException) as error: + await require_auth(self._request(), "Bearer abc") + + assert error.value.status_code == 403 + assert "code" in error.value.detail diff --git a/tests/test_local_tokens.py b/tests/test_local_tokens.py new file mode 100644 index 0000000..bdd9dbf --- /dev/null +++ b/tests/test_local_tokens.py @@ -0,0 +1,155 @@ +"""Tests for local client tokens: bearer tokens for processes that have no origin.""" + +import json +import stat +from unittest.mock import MagicMock, patch + +import pytest + +from agent_runtime.local_tokens import LocalTokenStore + + +@pytest.fixture +def store(temp_runtime_dir): + """A store backed by a file in the temporary runtime directory.""" + with patch("agent_runtime.local_tokens.settings") as mock_settings: + mock_settings.local_clients_file = temp_runtime_dir / "local_clients.json" + mock_settings.ensure_dirs = MagicMock() + yield LocalTokenStore() + + +class TestIssuing: + def test_created_token_identifies_its_client(self, store): + token = store.create("Claude Desktop", "code") + + client = store.lookup(token) + assert client is not None + assert client.name == "Claude Desktop" + assert client.scope == "code" + + def test_unknown_and_empty_tokens_identify_nobody(self, store): + store.create("tutor", "actions") + + assert store.lookup("not-a-token") is None + assert store.lookup("") is None + assert store.lookup(None) is None + + def test_creating_again_replaces_the_previous_token(self, store): + first = store.create("tutor", "actions") + second = store.create("tutor", "code") + + assert store.lookup(first) is None + assert store.lookup(second).scope == "code" + assert len(store.list_clients()) == 1 + + def test_names_must_be_present_and_short(self, store): + with pytest.raises(ValueError): + store.create(" ", "code") + with pytest.raises(ValueError): + store.create("x" * 65, "code") + + def test_scope_defaults_to_actions_for_a_record_without_one(self, store, temp_runtime_dir): + path = temp_runtime_dir / "local_clients.json" + path.write_text(json.dumps({"old": {"token_hash": "deadbeef"}})) + + with patch("agent_runtime.local_tokens.settings") as mock_settings: + mock_settings.local_clients_file = path + mock_settings.ensure_dirs = MagicMock() + reloaded = LocalTokenStore() + + assert reloaded.list_clients()[0].scope == "actions" + + +class TestStorage: + def test_the_plaintext_token_is_not_written_to_disk(self, store, temp_runtime_dir): + token = store.create("tutor", "code") + + stored = (temp_runtime_dir / "local_clients.json").read_text() + assert token not in stored + assert json.loads(stored)["tutor"]["token_hash"] + + def test_the_file_is_readable_only_by_its_owner(self, store, temp_runtime_dir): + store.create("tutor", "code") + + mode = (temp_runtime_dir / "local_clients.json").stat().st_mode + assert stat.S_IMODE(mode) == 0o600 + + def test_tokens_survive_a_restart(self, store, temp_runtime_dir): + token = store.create("tutor", "code") + + with patch("agent_runtime.local_tokens.settings") as mock_settings: + mock_settings.local_clients_file = temp_runtime_dir / "local_clients.json" + mock_settings.ensure_dirs = MagicMock() + reloaded = LocalTokenStore() + + assert reloaded.lookup(token).name == "tutor" + + def test_a_damaged_file_does_not_stop_the_runtime(self, temp_runtime_dir): + path = temp_runtime_dir / "local_clients.json" + path.write_text("{ not json") + + with patch("agent_runtime.local_tokens.settings") as mock_settings: + mock_settings.local_clients_file = path + mock_settings.ensure_dirs = MagicMock() + store = LocalTokenStore() + + assert store.list_clients() == [] + + +class TestRevoking: + def test_revoked_tokens_stop_working(self, store): + token = store.create("tutor", "code") + + assert store.revoke("tutor") is True + assert store.lookup(token) is None + + def test_revoking_an_unknown_client_reports_it(self, store): + assert store.revoke("nobody") is False + + def test_clients_are_listed_by_name(self, store): + store.create("b-client", "code") + store.create("a-client", "actions") + + assert [c.name for c in store.list_clients()] == ["a-client", "b-client"] + + +class TestAnotherProcess: + """The CLI issues and revokes tokens while the runtime is serving, in a separate process.""" + + @pytest.fixture + def stores(self, temp_runtime_dir): + """A factory of stores over one file, standing in for separate processes.""" + path = temp_runtime_dir / "local_clients.json" + with patch("agent_runtime.local_tokens.settings") as mock_settings: + mock_settings.local_clients_file = path + mock_settings.ensure_dirs = MagicMock() + yield LocalTokenStore + + def test_a_token_issued_after_startup_works_without_a_restart(self, stores): + serving = stores() # started before the token existed + assert serving.list_clients() == [] + + token = stores().create("tutor", "code") + + assert serving.lookup(token).name == "tutor" + + def test_a_revoked_token_stops_working_without_a_restart(self, stores): + cli = stores() + token = cli.create("tutor", "code") + serving = stores() + assert serving.lookup(token) is not None + + cli.revoke("tutor") + + assert serving.lookup(token) is None + + def test_issuing_a_token_keeps_one_issued_meanwhile(self, stores): + first = stores() + second = stores() + kept = first.create("kept", "actions") + + added = second.create("added", "code") + + assert {c.name for c in stores().list_clients()} == {"kept", "added"} + assert stores().lookup(kept) is not None + assert stores().lookup(added) is not None diff --git a/tests/test_mcp_render.py b/tests/test_mcp_render.py new file mode 100644 index 0000000..c9271d4 --- /dev/null +++ b/tests/test_mcp_render.py @@ -0,0 +1,125 @@ +"""Tests for turning kernel outputs into MCP content.""" + +from agent_runtime.mcp.render import MAX_IMAGES, MAX_TEXT_CHARS, render_execution, strip_ansi + + +def png(data: str = "iVBORw0KGgo=") -> dict: + return {"type": "display_data", "content": {"data": {"image/png": data}}} + + +class TestText: + def test_stream_output_becomes_text(self): + blocks = render_execution( + {"outputs": [{"type": "stream", "content": {"name": "stdout", "text": "hello\n"}}]} + ) + + assert [b.type for b in blocks] == ["text"] + assert blocks[0].text == "hello" + + def test_a_result_uses_its_readable_representation(self): + blocks = render_execution( + {"outputs": [{"type": "result", "content": {"data": {"text/plain": "42"}}}]} + ) + + assert blocks[0].text == "42" + + def test_the_documented_flat_shape_is_accepted_too(self): + blocks = render_execution( + { + "outputs": [ + {"output_type": "stream", "name": "stdout", "text": "one\n"}, + {"output_type": "execute_result", "data": {"text/plain": "two"}}, + ] + } + ) + + assert blocks[0].text == "one\ntwo" + + def test_a_cell_with_no_output_still_says_something(self): + blocks = render_execution({"outputs": [], "error": None, "execution_count": 3}) + + assert len(blocks) == 1 + assert "no output" in blocks[0].text + + def test_long_output_is_truncated_from_the_middle(self): + blocks = render_execution( + {"outputs": [{"type": "stream", "content": {"text": "x" * 40_000}}]} + ) + + text = blocks[0].text + assert len(text) < 40_000 + assert "characters omitted" in text + assert text.startswith("x") and text.endswith("x") + + def test_short_output_is_left_alone(self): + body = "y" * (MAX_TEXT_CHARS - 10) + blocks = render_execution({"outputs": [{"type": "stream", "content": {"text": body}}]}) + + assert blocks[0].text == body + + +class TestErrors: + def test_a_traceback_comes_back_as_content_not_a_failure(self): + blocks = render_execution( + { + "outputs": [], + "success": False, + "error": { + "ename": "ValueError", + "evalue": "bad", + "traceback": ["Traceback (most recent call last)", "ValueError: bad"], + }, + } + ) + + assert "ValueError: bad" in blocks[0].text + + def test_terminal_colors_are_stripped(self): + blocks = render_execution( + { + "outputs": [], + "error": { + "ename": "ValueError", + "evalue": "bad", + "traceback": ["\x1b[0;31mValueError\x1b[0m: bad"], + }, + } + ) + + assert "\x1b" not in blocks[0].text + assert blocks[0].text == "ValueError: bad" + + def test_an_error_without_a_traceback_still_reads(self): + blocks = render_execution({"outputs": [], "error": {"ename": "KeyError", "evalue": "k"}}) + + assert blocks[0].text == "KeyError: k" + + def test_strip_ansi_leaves_ordinary_text(self): + assert strip_ansi("plain text") == "plain text" + + +class TestImages: + def test_a_figure_comes_back_as_an_image(self): + blocks = render_execution({"outputs": [png()]}) + + assert [b.type for b in blocks] == ["image"] + assert blocks[0].mime_type == "image/png" + assert blocks[0].data == "iVBORw0KGgo=" + + def test_wrapped_base64_is_unwrapped(self): + blocks = render_execution({"outputs": [png("iVBO\nRw0K\nGgo=")]}) + + assert blocks[0].data == "iVBORw0KGgo=" + + def test_text_comes_before_images(self): + blocks = render_execution( + {"outputs": [png(), {"type": "stream", "content": {"text": "plotting\n"}}]} + ) + + assert [b.type for b in blocks] == ["text", "image"] + + def test_too_many_images_are_capped_and_counted(self): + blocks = render_execution({"outputs": [png()] * (MAX_IMAGES + 2)}) + + assert sum(b.type == "image" for b in blocks) == MAX_IMAGES + assert "2 further image" in blocks[-1].text diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py new file mode 100644 index 0000000..65c48a0 --- /dev/null +++ b/tests/test_mcp_server.py @@ -0,0 +1,138 @@ +"""Tests for the MCP server: what a client gets back, and what it is told when a call fails.""" + +from collections.abc import Callable +from unittest.mock import patch + +import httpx +import pytest +from mcp.server.mcpserver.exceptions import ToolError + +from agent_runtime.mcp import server +from agent_runtime.mcp.client import RuntimeClient + +CELL_OK = { + "cell_id": "c1", + "success": True, + "outputs": [{"type": "stream", "content": {"name": "stdout", "text": "hello\n"}}], + "error": None, + "execution_count": 1, +} + + +def runtime( + handler: Callable[[httpx.Request], httpx.Response], token: str | None = "t" +) -> RuntimeClient: + return RuntimeClient( + base_url="http://127.0.0.1:9477", + token=token, + transport=httpx.MockTransport(handler), + ) + + +def answering(response: httpx.Response, seen: list[httpx.Request] | None = None): + def handler(request: httpx.Request) -> httpx.Response: + if seen is not None: + seen.append(request) + return response + + return handler + + +class TestRunPython: + async def test_output_comes_back_as_content(self): + ok = answering(httpx.Response(200, json=CELL_OK)) + with patch.object(server, "_runtime", runtime(ok)): + blocks = await server.run_python("demo", "print('hello')") + + assert [b.type for b in blocks] == ["text"] + assert blocks[0].text == "hello" + + async def test_the_code_reaches_the_runtime_without_an_origin(self): + seen: list[httpx.Request] = [] + handler = answering(httpx.Response(200, json=CELL_OK), seen) + + with patch.object(server, "_runtime", runtime(handler)): + await server.run_python("demo", "print('hello')") + + request = seen[0] + assert request.url.path == "/cell/run" + # Not a browser: the runtime tells local clients apart by the absence of this header. + assert "origin" not in {k.lower() for k in request.headers} + assert request.headers["authorization"] == "Bearer t" + + async def test_an_impossible_lab_id_is_refused_before_the_call(self): + seen: list[httpx.Request] = [] + handler = answering(httpx.Response(200, json=CELL_OK), seen) + + with patch.object(server, "_runtime", runtime(handler)): + with pytest.raises(ToolError, match="lab_id"): + await server.run_python("../etc/passwd", "1") + + assert seen == [] + + async def test_a_traceback_is_a_result_not_a_tool_failure(self): + failed = { + "cell_id": "c1", + "success": False, + "outputs": [], + "error": {"ename": "ValueError", "evalue": "bad", "traceback": ["ValueError: bad"]}, + } + + with patch.object(server, "_runtime", runtime(answering(httpx.Response(200, json=failed)))): + blocks = await server.run_python("demo", "raise ValueError(\'bad\')") + + assert "ValueError: bad" in blocks[0].text + + +class TestFailures: + async def test_a_runtime_that_is_not_running_says_how_to_start_it(self): + def refuse(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused", request=request) + + with patch.object(server, "_runtime", runtime(refuse)): + with pytest.raises(ToolError, match="agent-runtime serve"): + await server.run_python("demo", "1 + 1") + + async def test_an_unknown_token_says_how_to_get_one(self): + denied = httpx.Response(401, json={"detail": "Unknown local client token."}) + + with patch.object(server, "_runtime", runtime(answering(denied))): + with pytest.raises(ToolError, match="token create"): + await server.run_python("demo", "1 + 1") + + async def test_a_narrow_scope_explains_the_widening(self): + detail = "This client is paired for 'actions' only; running code needs 'code'" + denied = httpx.Response(403, json={"detail": detail}) + + with patch.object(server, "_runtime", runtime(answering(denied))): + with pytest.raises(ToolError, match="--scope code"): + await server.run_python("demo", "1 + 1") + + async def test_other_refusals_keep_the_runtime_s_own_words(self): + denied = httpx.Response(500, json={"detail": "kernel died"}) + + with patch.object(server, "_runtime", runtime(answering(denied))): + with pytest.raises(ToolError, match="kernel died"): + await server.run_python("demo", "1 + 1") + + +class TestWhoami: + async def test_it_reports_the_principal_and_scope(self): + identity = httpx.Response(200, json={"principal": "local:tutor", "scope": "code"}) + client = runtime(answering(identity)) + + assert await client.whoami() == {"principal": "local:tutor", "scope": "code"} + + async def test_it_stays_quiet_when_the_runtime_cannot_answer(self): + def refuse(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused", request=request) + + assert await runtime(refuse).whoami() is None + + +class TestToolSurface: + async def test_one_tool_is_published_whatever_the_scope(self): + tools = await server.mcp.list_tools() + + assert [tool.name for tool in tools] == ["run_python"] + assert "persistent" in (tools[0].description or "") From 409efc074dcfe9f297c2e329deb8e93ada8b7680 Mon Sep 17 00:00:00 2001 From: thinkinginmath Date: Mon, 21 Sep 2026 08:39:35 +0000 Subject: [PATCH 3/3] Run ruff format, which CI checks and I had not Three files the formatter joins onto one line. No behavior change. Co-Authored-By: Claude Opus 5 (1M context) --- src/agent_runtime/auth.py | 3 +-- src/agent_runtime/mcp/render.py | 4 +--- tests/test_mcp_server.py | 2 +- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/agent_runtime/auth.py b/src/agent_runtime/auth.py index c5fb745..77865d9 100644 --- a/src/agent_runtime/auth.py +++ b/src/agent_runtime/auth.py @@ -276,8 +276,7 @@ async def _authenticate(request: Request, authorization: str | None) -> tuple[st raise HTTPException( status_code=401, detail=( - "Unknown local client token. Issue one with " - "'agent-runtime token create '." + "Unknown local client token. Issue one with 'agent-runtime token create '." ), ) diff --git a/src/agent_runtime/mcp/render.py b/src/agent_runtime/mcp/render.py index 643319a..ae7438c 100644 --- a/src/agent_runtime/mcp/render.py +++ b/src/agent_runtime/mcp/render.py @@ -87,9 +87,7 @@ def render_execution(result: dict[str, Any]) -> list[ContentBlock]: dropped = max(0, len(images) - MAX_IMAGES) blocks.extend(images[:MAX_IMAGES]) if dropped: - blocks.append( - TextContent(type="text", text=f"[{dropped} further image(s) not shown]") - ) + blocks.append(TextContent(type="text", text=f"[{dropped} further image(s) not shown]")) if not blocks: # Content of some kind, always: an empty result reads like a broken tool. diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 65c48a0..0f5f586 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -79,7 +79,7 @@ async def test_a_traceback_is_a_result_not_a_tool_failure(self): } with patch.object(server, "_runtime", runtime(answering(httpx.Response(200, json=failed)))): - blocks = await server.run_python("demo", "raise ValueError(\'bad\')") + blocks = await server.run_python("demo", "raise ValueError('bad')") assert "ValueError: bad" in blocks[0].text