diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..66a0aab --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,15 @@ +# Agent instructions + +## Agent skills + +### Issue tracker + +Issues and build specifications are tracked in GitHub Issues for `celados/mcpx`. See `docs/agents/issue-tracker.md`. + +### Triage labels + +Use the canonical triage labels defined in `docs/agents/triage-labels.md`. + +### Domain docs + +This is a single-context repository. Read `CONTEXT.md` and relevant ADRs under `docs/adr/`. See `docs/agents/domain.md`. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 120000 index 0000000..47dc3e3 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +AGENTS.md \ No newline at end of file diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 0000000..740c194 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,55 @@ +--- +type: Context +title: MCPX Domain Language +description: Canonical vocabulary for the user-local MCP runtime and its command interface. +--- + +# MCPX + +MCPX gives local commands one consistent authority for discovering, authenticating, and invoking registered MCP servers. + +## Language + +**MCP Runtime**: +The single user-local authority for credential lifecycle, active MCP sessions, and tool-call coordination. +_Avoid_: Daemon, session manager, connection pool when referring to the whole authority + +**CLI Adapter**: +A short-lived command interface that submits intent to the MCP Runtime and renders its response. +_Avoid_: Client control plane, auth owner + +**Declared Registry**: +The persistent set of MCP server declarations chosen by the user, excluding credentials and observed runtime state. +_Avoid_: Runtime registry, session config, schema cache + +**Credential Store**: +The durable, sensitive authorization material that lets MCPX act for the user without embedding secrets in server declarations. +_Avoid_: Auth config, token registry + +**Credential Identity**: +The stable authorization identity shared by every server declaration and call that must use the same grant. +_Avoid_: Server key, access token, session ID + +**Authentication Flow**: +One explicitly requested, shared attempt to make a Credential Identity usable; ordinary Calls report that authentication is required instead of starting the flow. +_Avoid_: Per-command login, token retry + +**Caller Input**: +A request/response on the active Runtime connection for secret or interactive data. The CLI Adapter renders the prompt, while the MCP Runtime retains Authentication Flow ownership. +_Avoid_: Daemon stdin, CLI-owned OAuth flow + +**Call**: +One caller-owned request that moves from accepted through queued or active work to exactly one terminal outcome. +_Avoid_: Detached job, background task + +**Schema Cache**: +Rebuildable knowledge MCPX has observed about server tools and their schemas. +_Avoid_: Declared tools, registry schema + +**Active State**: +Ephemeral facts about work currently coordinated by the MCP Runtime. +_Avoid_: Persistent config, runtime registry + +**Durable Operational State**: +Rebuildable or continuity-preserving coordination data that survives MCP Runtime restarts but is neither user intent nor authorization material. +_Avoid_: Declared Registry, Credential Store diff --git a/docs/adr/0001-mcpxd-is-the-user-local-mcp-runtime.md b/docs/adr/0001-mcpxd-is-the-user-local-mcp-runtime.md new file mode 100644 index 0000000..a75efd1 --- /dev/null +++ b/docs/adr/0001-mcpxd-is-the-user-local-mcp-runtime.md @@ -0,0 +1,9 @@ +--- +type: Decision +title: mcpxd Is the User-Local MCP Runtime +status: accepted +--- + +# mcpxd Is the User-Local MCP Runtime + +MCPX treats `mcpxd` as the single user-local authority for credential lifecycle, active MCP sessions, and tool-call coordination; the CLI is a stateless adapter. This supersedes the V1/V2 split in which the CLI owned authentication and registry-derived runtime state while the daemon owned only pooled connections, because that split permits competing control decisions across concurrent CLI processes. diff --git a/docs/adr/0002-runtime-is-the-sole-declared-registry-writer.md b/docs/adr/0002-runtime-is-the-sole-declared-registry-writer.md new file mode 100644 index 0000000..316d5a4 --- /dev/null +++ b/docs/adr/0002-runtime-is-the-sole-declared-registry-writer.md @@ -0,0 +1,9 @@ +--- +type: Decision +title: The MCP Runtime Is the Sole Declared Registry Writer +status: accepted +--- + +# The MCP Runtime Is the Sole Declared Registry Writer + +All command-driven changes to the Declared Registry go through the MCP Runtime, which is its sole writer. The persisted registry remains inspectable and backup-friendly but is not a supported direct-edit interface; configuration-as-code should use an explicit apply or import operation so external intent cannot race runtime writes. diff --git a/docs/adr/0003-separate-declared-and-observed-state.md b/docs/adr/0003-separate-declared-and-observed-state.md new file mode 100644 index 0000000..90523d6 --- /dev/null +++ b/docs/adr/0003-separate-declared-and-observed-state.md @@ -0,0 +1,9 @@ +--- +type: Decision +title: Separate Declared and Observed State +status: accepted +--- + +# Separate Declared and Observed State + +MCPX stores user declarations, authorization material, rebuildable schema knowledge, durable operational state, and active in-memory coordination as distinct classes of state. Derived schema or runtime status must not be written into the Declared Registry, because incidental observations must never contend with or overwrite user intent. diff --git a/docs/adr/0004-single-flight-authentication-per-credential-identity.md b/docs/adr/0004-single-flight-authentication-per-credential-identity.md new file mode 100644 index 0000000..1568bd0 --- /dev/null +++ b/docs/adr/0004-single-flight-authentication-per-credential-identity.md @@ -0,0 +1,9 @@ +--- +type: Decision +title: Single-Flight Authentication per Credential Identity +status: accepted +--- + +# Single-Flight Authentication per Credential Identity + +The MCP Runtime permits at most one refresh or interactive Authentication Flow for a Credential Identity. Concurrent callers share its result; disconnecting a caller removes only that waiter, while the final waiter leaving cancels the flow, so concurrent commands cannot rotate the same grant independently or open duplicate authorization pages. diff --git a/docs/adr/0005-cli-disconnect-cancels-its-call.md b/docs/adr/0005-cli-disconnect-cancels-its-call.md new file mode 100644 index 0000000..e542291 --- /dev/null +++ b/docs/adr/0005-cli-disconnect-cancels-its-call.md @@ -0,0 +1,9 @@ +--- +type: Decision +title: CLI Disconnect Cancels Its Call +status: accepted +--- + +# CLI Disconnect Cancels Its Call + +A normal Call is owned by its originating CLI connection: disconnect removes a queued Call or requests cancellation of an active Call, and the MCP Runtime never writes its result to a dead connection. Cancellation is best-effort and does not promise to undo external side effects; detached Calls require a future explicit interface rather than arising accidentally from disconnects. diff --git a/docs/adr/0006-each-call-owns-its-lifecycle.md b/docs/adr/0006-each-call-owns-its-lifecycle.md new file mode 100644 index 0000000..10d1415 --- /dev/null +++ b/docs/adr/0006-each-call-owns-its-lifecycle.md @@ -0,0 +1,9 @@ +--- +type: Decision +title: Each Call Owns Its Lifecycle +status: accepted +--- + +# Each Call Owns Its Lifecycle + +Each Call is the sole authority for its queued, active, and terminal transitions, caller-disconnect cancellation, and lifecycle cleanup. A shared session decides when queued work may proceed but cannot transition or cancel the Call itself, because splitting lifecycle ownership between the CLI connection and session queue permits activation after disconnect and cancellation after completion. diff --git a/docs/adr/0007-authentication-is-explicit.md b/docs/adr/0007-authentication-is-explicit.md new file mode 100644 index 0000000..1d76fe4 --- /dev/null +++ b/docs/adr/0007-authentication-is-explicit.md @@ -0,0 +1,13 @@ +--- +type: Decision +title: Authentication Is Explicit +status: accepted +--- + +# Authentication Is Explicit + +An ordinary Call never starts an Authentication Flow. When its Credential Identity is unusable, the MCP Runtime returns `reauth-required`; only an explicit `mcpx @refresh` operation may start or join the single shared flow, so agent and script invocations cannot unexpectedly open a browser or wait for interactive input. + +When a provider requires a manually registered OAuth client, the Runtime sends a typed Caller Input request on the active `@refresh` connection. The CLI Adapter owns terminal prompting only; the Runtime owns metadata discovery, single-flight coordination, callback lifetime, token exchange, and persistence. The daemon never reads its ignored stdin, and the CLI never becomes an authentication state machine. + +Caller Input is independently cancellable. If its caller disconnects while another waiter remains, the Authentication Flow offers the input request to a surviving waiter; Runtime shutdown aborts the flow and the CLI prompt signal before returning its terminal outcome. diff --git a/docs/agents/domain.md b/docs/agents/domain.md new file mode 100644 index 0000000..76c1df8 --- /dev/null +++ b/docs/agents/domain.md @@ -0,0 +1,28 @@ +--- +type: Reference +title: Domain Documentation +description: Rules for consuming mcpx domain language and architectural decisions. +--- + +# Domain Documentation + +This repository uses a single domain context. + +## Before Exploring + +Read: + +- `/CONTEXT.md`, when present. +- Relevant ADRs under `/docs/adr/`. + +Missing files are not errors. `/domain-modeling` creates them lazily when terminology or durable decisions are resolved. + +## Vocabulary + +Use terms exactly as defined in `CONTEXT.md`. Do not introduce synonyms for established concepts. + +If a required concept is absent, either reconsider the new term or record the gap through `/domain-modeling`. + +## Architectural Decisions + +Surface conflicts with existing ADRs explicitly. Do not silently override an accepted decision. diff --git a/docs/agents/issue-tracker.md b/docs/agents/issue-tracker.md new file mode 100644 index 0000000..e6cf9d0 --- /dev/null +++ b/docs/agents/issue-tracker.md @@ -0,0 +1,26 @@ +--- +type: Reference +title: Issue Tracker +description: GitHub issue workflow used by engineering skills in celados/mcpx. +--- + +# Issue Tracker + +Issues and build specifications live in `celados/mcpx` GitHub Issues. Use `gh` from this repository so it infers the remote. + +## Operations + +- Create, read, comment, label, and close issues with `gh issue`. +- When a skill says “publish to the issue tracker,” create a GitHub issue. +- When a skill says “fetch the relevant ticket,” read the issue body, comments, and labels. +- Pull requests are not a triage request surface. +- Resolve a bare issue or PR number before operating on it. + +## Dependencies + +Use GitHub native sub-issues and issue dependencies. + +- A dependency points from the blocked issue to the blocker. +- Use the blocker’s numeric database ID, not its issue number or node ID. +- Fall back to a `Blocked by: #...` line only when native dependencies are unavailable. +- A ticket is ready only when every blocker is closed. diff --git a/docs/agents/triage-labels.md b/docs/agents/triage-labels.md new file mode 100644 index 0000000..9ad6d77 --- /dev/null +++ b/docs/agents/triage-labels.md @@ -0,0 +1,15 @@ +--- +type: Reference +title: Triage Labels +description: Canonical triage roles and their GitHub label mappings. +--- + +# Triage Labels + +| Role | GitHub label | Meaning | +| ----------------- | ----------------- | -------------------------------- | +| `needs-triage` | `needs-triage` | Maintainer evaluation required | +| `needs-info` | `needs-info` | Waiting for reporter information | +| `ready-for-agent` | `ready-for-agent` | Fully specified and agent-ready | +| `ready-for-human` | `ready-for-human` | Human implementation required | +| `wontfix` | `wontfix` | Will not be actioned | diff --git a/docs/mcpxd-bdd.md b/docs/mcpxd-bdd.md index 0a05d82..5f63087 100644 --- a/docs/mcpxd-bdd.md +++ b/docs/mcpxd-bdd.md @@ -1,5 +1,14 @@ +--- +type: Specification +title: mcpxd BDD Spec +status: superseded +superseded_by: docs/specs/mcp-runtime-upgrade.md +--- + # mcpxd BDD Spec +> Superseded for architecture and ownership. Retained only as historical regression evidence. + ## Purpose `mcpxd` is a user-local daemon for reusing stdio MCP server sessions across diff --git a/docs/mcpxd-v2-bdd.md b/docs/mcpxd-v2-bdd.md index 371cf9d..f188d2c 100644 --- a/docs/mcpxd-v2-bdd.md +++ b/docs/mcpxd-v2-bdd.md @@ -1,5 +1,14 @@ +--- +type: Specification +title: mcpxd V2 BDD Spec - HTTP routing and first-class notifications +status: superseded +superseded_by: docs/specs/mcp-runtime-upgrade.md +--- + # mcpxd V2 BDD Spec - HTTP routing and first-class notifications +> Superseded for architecture and ownership. Retained only as historical regression evidence. + V2 extends V1 in three orthogonal directions: 1. Route HTTP MCP servers through the daemon (session-id preservation + connection reuse). diff --git a/docs/specs/mcp-runtime-upgrade.md b/docs/specs/mcp-runtime-upgrade.md new file mode 100644 index 0000000..5e7d37c --- /dev/null +++ b/docs/specs/mcp-runtime-upgrade.md @@ -0,0 +1,343 @@ +--- +type: Specification +title: MCP Runtime Structural Upgrade +description: > + Proposed executable contract for moving registry, credentials, sessions, + authentication, and caller-owned Calls behind one user-local MCP Runtime. +status: accepted +version: 1 +issues: [16] +--- + +# MCP Runtime Structural Upgrade + +## Status and Sources + +This specification is the authority for production implementation. It +supersedes the architectural assumptions in `docs/mcpxd-bdd.md` and +`docs/mcpxd-v2-bdd.md`; their verified transport and rendering scenarios remain +regression evidence. + +Normative decisions: + +- `docs/adr/0001-mcpxd-is-the-user-local-mcp-runtime.md` +- `docs/adr/0002-runtime-is-the-sole-declared-registry-writer.md` +- `docs/adr/0003-separate-declared-and-observed-state.md` +- `docs/adr/0004-single-flight-authentication-per-credential-identity.md` +- `docs/adr/0005-cli-disconnect-cancels-its-call.md` +- `docs/adr/0006-each-call-owns-its-lifecycle.md` +- `docs/adr/0007-authentication-is-explicit.md` +- `prototypes/mcp-cancellation/findings.md` +- GitHub Issue #16 + +## Required Outcomes + +1. The CLI Adapter is short-lived. It discovers command shape from the MCP + Runtime, submits one operation, renders frames, and exits after its terminal + frame without retaining MCP transports, OAuth servers, timers, workers, or + signal handlers. +2. The MCP Runtime is the only process that coordinates Calls, MCP sessions, + Credential Identities, Authentication Flows, and command-driven registry + writes. +3. A Call is caller-owned from acceptance through exactly one terminal outcome. + Disconnect cancels only that Call and never evicts a healthy shared MCP + session. +4. Ordinary Calls never start authentication. An unusable Credential Identity + produces `reauth-required`; explicit `mcpx @refresh` is the only operation + that may start or join an Authentication Flow. +5. Declared Registry, Credential Store, Schema Cache, Active State, and any + Durable Operational State remain distinct. + +## Product Contract + +### Ordinary Call + +- A Call with usable credentials may be queued and executed. +- A Call with missing, expired, rejected, or otherwise unusable credentials + terminates with error code `reauth-required`. +- The Call does not open a browser, prompt for OAuth data, refresh a token, or + wait for a future Authentication Flow. +- After the user runs `mcpx @refresh`, they explicitly retry the original Call. + +### Explicit Refresh + +- `mcpx @refresh` may refresh schemas and make Credential Identities usable. +- Concurrent refresh operations for the same Credential Identity join one + Authentication Flow and observe the same terminal result. +- Exactly one browser authorization flow may exist per Credential Identity. +- Disconnect removes that refresh caller as a waiter. The final waiter leaving + cancels the Authentication Flow and closes its local callback resources. + +### Disconnect + +- Disconnect before activation removes the queued Call without contacting the + MCP server. +- Disconnect during execution requests MCP cancellation through the active + Call's `AbortSignal`. +- Cancellation is best-effort and does not promise rollback of tool side + effects. +- No result or error is written to a dead CLI connection. +- Detached Calls are unsupported. + +## Runtime Protocol + +The upgrade is a clean protocol break. Increment `DAEMON_PROTOCOL_VERSION` and +use the existing incompatible-daemon stop/start path; do not add V2 shims or +feature negotiation. + +One CLI connection performs a handshake followed by one operation. Every +operation has a caller-generated `requestId`. Runtime frames repeat that ID: + +```ts +type RuntimeFrame = + | { requestId: string; kind: 'event'; event: RuntimeEvent } + | { requestId: string; kind: 'result'; result: unknown } + | { + requestId: string + kind: 'error' + error: { code: RuntimeErrorCode; message: string } + } +``` + +The terminal frame is exactly one `result` or `error`. Events are optional and +precede the terminal frame. JSON Lines remains the framing format. + +Supported operation intents: + +- `registrySnapshot`: return declarations plus cached command schemas required + to build the CLI router, without credential material or runtime session data. +- `call`: identify the registered server and tool by name and provide tool + input. The CLI does not send server configuration or authorization headers. +- `addServer`: discover and persist one declaration through the Runtime. +- `removeServers`: remove declarations and unreferenced credential material + through the Runtime. +- `refreshServers`: explicitly refresh schemas and, where required, start or + join Authentication Flows. +- `status`: return redacted operational status. +- `stop`: stop the Runtime after coordinated cleanup. + +## Module Design + +### External Seam: MCP Runtime + +The Runtime is one deep module. Its external interface accepts a validated +operation intent and a caller adapter. Callers do not coordinate credentials, +queues, sessions, persistence, or cancellation. + +```ts +type RuntimeCaller = { + id: string + onDisconnect: (listener: () => void) => () => void + requestInput: ( + request: RuntimeInputRequest, + signal?: AbortSignal, + ) => Promise + send: (frame: RuntimeFrame) => Promise +} + +type McpRuntime = { + handle: (intent: RuntimeIntent, caller: RuntimeCaller) => Promise +} +``` + +The Unix socket is the production caller adapter. Tests use an in-memory caller +adapter. This is a real seam because both adapters exercise the same interface. +During explicit authentication, an `input-required` event may request manual +OAuth client data. The CLI responds on the same connection with a correlated +input frame; this does not count as a second operation. No other intent is +accepted after the first operation. +Terminal or disconnected CLI connections abort any active input provider, and +an Authentication Flow may transfer input ownership to another surviving +waiter for the same Credential Identity. + +### Call Lifecycle + +Each accepted tool invocation creates one internal Call object with: + +- state `accepted | queued | active | terminal`; +- one caller-disconnect subscription; +- one `AbortController` created before queueing; +- one terminal result chosen by an idempotent transition. + +Only the Call object may transition its state or remove its disconnect +subscription. The session queue grants permission to activate by asking the +Call to transition; it cannot mutate Call state directly. + +Terminal transition order is normative: + +1. atomically mark the Call terminal; +2. remove the caller-disconnect subscription; +3. release queue/session accounting; +4. send a terminal frame only if the caller is still connected. + +This order prevents SDK 1.29.0 from emitting stale cancellation after normal +completion. + +### Session Pool and Queue + +The session pool owns MCP connections keyed by the existing stable server key. +Each managed session contains an explicit FIFO of Call entries and one drain +loop. Do not retain the current promise-chain queue, because a queued entry +cannot be removed from a chained promise without leaving hidden work behind. + +The drain loop: + +1. skips Calls already terminal because their caller disconnected; +2. asks the next Call to transition from queued to active; +3. invokes the MCP tool with the Call signal and existing timeout/progress + options; +4. completes session accounting in `finally`; +5. continues with the next entry regardless of the prior outcome. + +One MCP request remains active per server key. Cancellation, timeout, tool +failure, and late response do not evict the connection. Authentication failure, +protocol failure, explicit eviction, or connection closure may rebuild it under +their own policies. + +Runtime shutdown closes admission before awaiting any store read or connection, +cancels active and queued Calls plus Authentication Flows, closes connections +that resolve late, and only then completes the stop operation. + +### Credential Coordination + +Credential coordination is internal to the Runtime: + +- ordinary Calls perform a read-only usability check; +- unusable credentials return `reauth-required` without mutation; +- explicit refresh operations join a single in-memory promise per Credential + Identity; +- each refresh caller is a waiter with independent disconnect cleanup; +- the final waiter leaving aborts the shared flow; +- successful rotation is persisted once before all waiters are released. + +The existing cross-process token-cache lock is migration scaffolding, not the +final coordination model. Remove it only after every credential mutation has +moved behind the Runtime and regression tests prove one refresh across +concurrent CLI processes. + +### State Stores + +The Runtime owns three persistence modules with narrow internal interfaces: + +- Registry Store: declarations selected by the user; +- Credential Store: OAuth tokens, client secrets, bearer material, raw HTTP + headers, and stdio environment values; +- Schema Cache: rebuildable tool schemas, discovery timestamps, refresh status, + and dirty markers. + +Active Calls, waiters, queues, sessions, and Authentication Flows stay in +memory. No durable Call or detached job store is introduced. + +## Persistence Migration + +The current combined registry is migrated once by the Runtime: + +1. acquire Runtime startup ownership; +2. read the current registry and credential files; +3. split declarations from cached schemas and refresh observations; +4. write each new store atomically; +5. retain one recoverable backup of the pre-migration registry; +6. start serving operations only after every write succeeds. + +After migration, the Runtime uses only the new store contracts. The CLI has no +legacy read/write path, and the Runtime does not dual-write old and new shapes. + +## SDK Cancellation Adapter + +For an active Call, pass its signal in the third `Client.callTool` argument +alongside timeout and progress options. + +MCPX-owned state determines whether the terminal cause is caller disconnect or +timeout. Do not infer it from SDK error code `-32001`, because SDK 1.29.0 uses +that code for both. A late-response `client.onerror` for a cancelled request is +diagnostic only and does not mark the session unhealthy. + +Streamable HTTP cancellation sends an MCP cancellation request on the existing +session; it does not abort the transport-wide fetch controller. Do not close or +rebuild the HTTP transport solely because one Call was cancelled. + +## Integration of Existing Dirty Work + +- Preserve the OAuth callback server `unref()` mitigation until Runtime-owned + Authentication Flow tests prove callback cleanup on success, error, timeout, + and final-waiter disconnect. +- Preserve the concurrent refresh regression fixture and its rotating-token + behavior. +- Replace the cross-process token-cache lock only when Runtime sole-writer and + single-flight tests cover the same race. Do not layer a second coordination + mechanism over it indefinitely. +- Keep all OAuth and HTTP acceptance fixtures local. Never require a real + Cloudflare, PostHog, or other remote MCP endpoint. + +## Acceptance Criteria + +### Issue #16 Process Lifecycle + +- Launch at least five concurrent CLI processes against a deterministic local + HTTP MCP fixture; every process preserves its output and exit code and exits + within a bounded post-result deadline. +- Repeat for tool failure, caller cancellation, and broken output pipe. +- After every CLI terminal frame, no CLI-owned timer, listener, server, worker, + pipe, or MCP transport keeps that CLI process alive. +- The shared Runtime may remain alive while idle but must not busy-spin. + +### Caller-Owned Cancellation + +- A queued Call whose socket closes never reaches the fixture. +- An active stdio Call whose socket closes emits `notifications/cancelled`, + settles promptly, and leaves the same child usable. +- An active Streamable HTTP Call whose socket closes emits cancellation on the + same session and leaves that session usable. +- Respect, ignore, late-response, response-race, and complete-before-disconnect + outcomes preserve subsequent request correlation. +- Normal completion removes the disconnect listener before a later socket close + can emit stale cancellation. + +### Explicit Authentication and Single-Flight + +- An ordinary Call with unusable OAuth credentials returns `reauth-required` + without hitting an authorization or token endpoint and without opening a + browser. +- Five concurrent explicit refresh CLI processes for one Credential Identity + produce one token refresh or one interactive flow, persist one result, and all + exit. +- Disconnecting one refresh waiter does not cancel the flow while another + waiter remains. +- Disconnecting the final waiter closes callback resources and cancels the + flow. +- A 401 invalidates usability, applies the existing session-eviction policy, + and returns `reauth-required`; it does not start authentication. + +### State Ownership + +- CLI code has no command-driven write path to Registry Store, Credential + Store, or Schema Cache. +- Registry snapshots never contain tokens, client secrets, raw headers, session + IDs, stdio environment values, queues, or health state. +- Migration preserves declarations and cached command availability while + splitting their storage. +- Interrupted migration leaves the pre-migration data recoverable and never + exposes a partially migrated runtime. + +### Regression Gates + +For every implementation ticket: + +1. write a failing test through the owning module interface; +2. run the focused test; +3. run the relevant integration suite; +4. run the full test suite and typecheck; +5. run formatter checks and `git diff --check`. + +The final upgrade additionally runs all local process-level, stdio, Streamable +HTTP, notification, OAuth, migration, and protocol-mismatch cases. + +## Non-Goals + +- concurrent MCP requests on one session; +- detached or durable Calls; +- background authentication initiated by an ordinary Call; +- real OAuth or remote MCP acceptance dependencies; +- preserving daemon protocol V2 compatibility; +- durable notification subscriptions; +- persisting Active State. diff --git a/package.json b/package.json index e83a016..1742fe3 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "mcpx", - "version": "0.9.15", + "version": "0.10.0", "license": "MIT", "bin": { "mcpx": "./src/main.ts" @@ -10,6 +10,8 @@ "build": "./scripts/build.sh", "check": "bun run typecheck && bun test && bunx oxfmt --check package.json tsconfig.json README.md .gitignore .oxfmtrc.json src tests skills docs", "format": "bunx oxfmt package.json tsconfig.json README.md .gitignore .oxfmtrc.json src tests skills docs", + "prototype:cancellation": "bun prototypes/mcp-cancellation/run.ts", + "prototype:cancellation:evidence": "bun prototypes/mcp-cancellation/run.ts --evidence", "typecheck": "tsgo --project tsconfig.json --noEmit", "test": "bun test" }, diff --git a/prototypes/mcp-cancellation/README.md b/prototypes/mcp-cancellation/README.md new file mode 100644 index 0000000..f07568f --- /dev/null +++ b/prototypes/mcp-cancellation/README.md @@ -0,0 +1,35 @@ +--- +type: Prototype +title: MCP Call Cancellation +description: Throwaway runner for pinned SDK cancellation behavior over local stdio and Streamable HTTP fixtures. +--- + +# MCP Call Cancellation + +This throwaway prototype asks whether one caller-owned `AbortSignal` can cancel +`Client.callTool` promptly without sacrificing a reusable stdio or Streamable +HTTP connection. It drives four server outcomes: respect cancellation, ignore +and reply late, reply as cancellation races, and complete before the caller +aborts. + +No fixture performs OAuth or connects beyond loopback and a spawned local stdio +process. + +## Run + +Interactive state view: + +```sh +bun run prototype:cancellation +``` + +Deterministic evidence capture: + +```sh +bun run prototype:cancellation:evidence +``` + +The evidence command writes `evidence.json` beside this file. + +See [`findings.md`](findings.md) for the observed verdict and the proposed MCPX +adapter boundary. diff --git a/prototypes/mcp-cancellation/evidence.json b/prototypes/mcp-cancellation/evidence.json new file mode 100644 index 0000000..2b23f39 --- /dev/null +++ b/prototypes/mcp-cancellation/evidence.json @@ -0,0 +1,482 @@ +{ + "question": "Can one caller-owned AbortSignal promptly cancel Client.callTool while preserving stdio and Streamable HTTP connection reuse?", + "runtime": { + "bun": "1.3.14", + "sdk": "1.29.0", + "sdkLock": "@modelcontextprotocol/sdk@1.29.0" + }, + "abortReason": "originating CLI socket closed", + "transports": [ + { + "transport": "stdio", + "scenarios": [ + { + "scenario": "acknowledge", + "settlement": { + "status": "rejected", + "elapsedMs": 16, + "error": { + "name": "McpError", + "code": -32001, + "message": "MCP error -32001: Error: originating CLI socket closed" + } + }, + "reuseBeforeLate": "echo-ok", + "reuseAfterLate": "echo-ok", + "fixtureEvents": [ + { + "event": "message", + "method": "initialize", + "id": 0 + }, + { + "event": "message", + "method": "notifications/initialized" + }, + { + "event": "message", + "method": "tools/call", + "id": 1 + }, + { + "event": "call", + "id": 1, + "scenario": "acknowledge" + }, + { + "event": "message", + "method": "notifications/cancelled" + }, + { + "event": "cancel", + "requestId": 1, + "reason": "Error: originating CLI socket closed", + "scenario": "acknowledge", + "pending": true + }, + { + "event": "work-stopped", + "requestId": 1, + "scenario": "acknowledge" + }, + { + "event": "message", + "method": "tools/call", + "id": 2 + }, + { + "event": "message", + "method": "tools/call", + "id": 3 + } + ], + "clientErrors": [] + }, + { + "scenario": "ignore", + "settlement": { + "status": "rejected", + "elapsedMs": 15, + "error": { + "name": "McpError", + "code": -32001, + "message": "MCP error -32001: Error: originating CLI socket closed" + } + }, + "reuseBeforeLate": "echo-ok", + "reuseAfterLate": "echo-ok", + "fixtureEvents": [ + { + "event": "message", + "method": "tools/call", + "id": 4 + }, + { + "event": "call", + "id": 4, + "scenario": "ignore" + }, + { + "event": "message", + "method": "notifications/cancelled" + }, + { + "event": "cancel", + "requestId": 4, + "reason": "Error: originating CLI socket closed", + "scenario": "ignore", + "pending": true + }, + { + "event": "message", + "method": "tools/call", + "id": 5 + }, + { + "event": "response", + "id": 4, + "scenario": "ignore" + }, + { + "event": "message", + "method": "tools/call", + "id": 6 + } + ], + "clientErrors": [ + "Received a response for an unknown message ID: {\"jsonrpc\":\"2.0\",\"id\":4,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"ignore\"}]}}" + ] + }, + { + "scenario": "race", + "settlement": { + "status": "rejected", + "elapsedMs": 15, + "error": { + "name": "McpError", + "code": -32001, + "message": "MCP error -32001: Error: originating CLI socket closed" + } + }, + "reuseBeforeLate": "echo-ok", + "reuseAfterLate": "echo-ok", + "fixtureEvents": [ + { + "event": "message", + "method": "tools/call", + "id": 7 + }, + { + "event": "call", + "id": 7, + "scenario": "race" + }, + { + "event": "message", + "method": "notifications/cancelled" + }, + { + "event": "cancel", + "requestId": 7, + "reason": "Error: originating CLI socket closed", + "scenario": "race", + "pending": true + }, + { + "event": "message", + "method": "tools/call", + "id": 8 + }, + { + "event": "response", + "id": 7, + "scenario": "race" + }, + { + "event": "message", + "method": "tools/call", + "id": 9 + } + ], + "clientErrors": [ + "Received a response for an unknown message ID: {\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"race\"}]}}" + ] + }, + { + "scenario": "complete-before-cancel", + "settlement": { + "status": "fulfilled", + "elapsedMs": 6, + "text": "complete-before-cancel" + }, + "reuseBeforeLate": "echo-ok", + "reuseAfterLate": "echo-ok", + "fixtureEvents": [ + { + "event": "message", + "method": "tools/call", + "id": 10 + }, + { + "event": "call", + "id": 10, + "scenario": "complete-before-cancel" + }, + { + "event": "response", + "id": 10, + "scenario": "complete-before-cancel" + }, + { + "event": "message", + "method": "notifications/cancelled" + }, + { + "event": "cancel", + "requestId": 10, + "reason": "Error: originating CLI socket closed", + "scenario": "complete-before-cancel", + "pending": false + }, + { + "event": "message", + "method": "tools/call", + "id": 11 + }, + { + "event": "message", + "method": "tools/call", + "id": 12 + } + ], + "clientErrors": [] + } + ] + }, + { + "transport": "streamable-http", + "sessionIds": [ + null, + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session", + "local-session" + ], + "scenarios": [ + { + "scenario": "acknowledge", + "settlement": { + "status": "rejected", + "elapsedMs": 16, + "error": { + "name": "McpError", + "code": -32001, + "message": "MCP error -32001: Error: originating CLI socket closed" + } + }, + "reuseBeforeLate": "echo-ok", + "reuseAfterLate": "echo-ok", + "fixtureEvents": [ + { + "event": "message", + "method": "tools/call", + "id": 1 + }, + { + "event": "call", + "id": 1, + "scenario": "acknowledge", + "sessionId": "local-session" + }, + { + "event": "message", + "method": "notifications/cancelled" + }, + { + "event": "cancel", + "requestId": 1, + "reason": "Error: originating CLI socket closed", + "scenario": "acknowledge", + "pending": true, + "originalRequestAborted": false + }, + { + "event": "work-stopped", + "requestId": 1, + "scenario": "acknowledge" + }, + { + "event": "message", + "method": "tools/call", + "id": 2 + }, + { + "event": "message", + "method": "tools/call", + "id": 3 + } + ], + "clientErrors": [] + }, + { + "scenario": "ignore", + "settlement": { + "status": "rejected", + "elapsedMs": 16, + "error": { + "name": "McpError", + "code": -32001, + "message": "MCP error -32001: Error: originating CLI socket closed" + } + }, + "reuseBeforeLate": "echo-ok", + "reuseAfterLate": "echo-ok", + "fixtureEvents": [ + { + "event": "message", + "method": "tools/call", + "id": 4 + }, + { + "event": "call", + "id": 4, + "scenario": "ignore", + "sessionId": "local-session" + }, + { + "event": "message", + "method": "notifications/cancelled" + }, + { + "event": "cancel", + "requestId": 4, + "reason": "Error: originating CLI socket closed", + "scenario": "ignore", + "pending": true, + "originalRequestAborted": false + }, + { + "event": "message", + "method": "tools/call", + "id": 5 + }, + { + "event": "response", + "id": 4, + "scenario": "ignore" + }, + { + "event": "message", + "method": "tools/call", + "id": 6 + } + ], + "clientErrors": [ + "Received a response for an unknown message ID: {\"jsonrpc\":\"2.0\",\"id\":4,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"ignore\"}]}}" + ] + }, + { + "scenario": "race", + "settlement": { + "status": "rejected", + "elapsedMs": 16, + "error": { + "name": "McpError", + "code": -32001, + "message": "MCP error -32001: Error: originating CLI socket closed" + } + }, + "reuseBeforeLate": "echo-ok", + "reuseAfterLate": "echo-ok", + "fixtureEvents": [ + { + "event": "message", + "method": "tools/call", + "id": 7 + }, + { + "event": "call", + "id": 7, + "scenario": "race", + "sessionId": "local-session" + }, + { + "event": "message", + "method": "notifications/cancelled" + }, + { + "event": "cancel", + "requestId": 7, + "reason": "Error: originating CLI socket closed", + "scenario": "race", + "pending": true, + "originalRequestAborted": false + }, + { + "event": "response", + "id": 7, + "scenario": "race" + }, + { + "event": "message", + "method": "tools/call", + "id": 8 + }, + { + "event": "message", + "method": "tools/call", + "id": 9 + } + ], + "clientErrors": [ + "Received a response for an unknown message ID: {\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"race\"}]}}" + ] + }, + { + "scenario": "complete-before-cancel", + "settlement": { + "status": "fulfilled", + "elapsedMs": 6, + "text": "complete-before-cancel" + }, + "reuseBeforeLate": "echo-ok", + "reuseAfterLate": "echo-ok", + "fixtureEvents": [ + { + "event": "message", + "method": "tools/call", + "id": 10 + }, + { + "event": "call", + "id": 10, + "scenario": "complete-before-cancel", + "sessionId": "local-session" + }, + { + "event": "response", + "id": 10, + "scenario": "complete-before-cancel" + }, + { + "event": "message", + "method": "notifications/cancelled" + }, + { + "event": "cancel", + "requestId": 10, + "reason": "Error: originating CLI socket closed", + "scenario": "complete-before-cancel", + "pending": false, + "originalRequestAborted": false + }, + { + "event": "message", + "method": "tools/call", + "id": 11 + }, + { + "event": "message", + "method": "tools/call", + "id": 12 + } + ], + "clientErrors": [] + } + ] + } + ] +} diff --git a/prototypes/mcp-cancellation/findings.md b/prototypes/mcp-cancellation/findings.md new file mode 100644 index 0000000..55770c6 --- /dev/null +++ b/prototypes/mcp-cancellation/findings.md @@ -0,0 +1,112 @@ +--- +type: Research +title: MCP Call Cancellation Findings +description: Observed SDK 1.29.0 cancellation behavior and a proposed MCPX adapter boundary. +--- + +# MCP Call Cancellation Findings + +## Evidence Boundary + +- Runtime: Bun `1.3.14` and `@modelcontextprotocol/sdk@1.29.0`. +- The SDK version is exact in `bun.lock`. The repository does not currently pin + a Bun version in `package.json`, `.tool-versions`, or a mise config, so Bun + `1.3.14` is the observed executor version rather than a repository-declared + pin. +- Fixtures are a spawned local stdio process and a loopback-only Bun HTTP + server. No OAuth or remote MCP endpoint is involved. +- Raw observations: `evidence.json`. +- Reproduction: `bun run prototype:cancellation:evidence`. + +The fixture deliberately controls both compliant and adversarial server +outcomes. SDK implementation checks use the installed files at: + +- `node_modules/@modelcontextprotocol/sdk/dist/esm/shared/protocol.d.ts:61` +- `node_modules/@modelcontextprotocol/sdk/dist/esm/shared/protocol.js:636` +- `node_modules/@modelcontextprotocol/sdk/dist/esm/client/streamableHttp.js:288` + +## Confirmed SDK Behavior + +1. `RequestOptions.signal` is the caller cancellation option for + `Client.callTool`. Aborting it removes the response/progress handlers, clears + the timeout, sends `notifications/cancelled` with the original JSON-RPC + `requestId` and stringified abort reason, and rejects the caller promptly. +2. SDK 1.29.0 rejects an aborted call with `McpError` code `-32001` + (`RequestTimeout`), not an `AbortError`, unless the abort reason is already an + `McpError`. This conflicts with the `RequestOptions.signal` declaration text + that says an `AbortError` is raised. MCPX must not classify cancellation from + the thrown error code alone. +3. `timeout` reaches the same internal cancellation path, including emitting + `notifications/cancelled`, but it represents deadline expiry rather than CLI + ownership loss. It is not a substitute for the caller signal. +4. Stdio cancellation does not close the child transport. A server can respect + cancellation, ignore it, or race a response; a second call remains correctly + correlated in every observed case. +5. An ignored or racing late response is discarded because its response handler + was removed. SDK 1.29.0 reports it through `client.onerror` as an unknown + message ID, but subsequent calls remain usable. +6. Streamable HTTP request sends use the transport's single lifecycle + `AbortController`, not `RequestOptions.signal`. Caller cancellation therefore + does not abort the original fetch. It sends a separate + `notifications/cancelled` POST on the same `mcp-session-id`; the session and + transport remain reusable. Only transport `close()` aborts the transport-wide + controller. +7. In every HTTP scenario, calls before and after a late response reused + `local-session`. The fixture observed `originalRequestAborted: false` when it + received cancellation. +8. The SDK does not remove the request's abort listener after a normal response. + Aborting the controller after fulfillment still emits a stale + `notifications/cancelled` for the completed request. The server safely saw it + as non-pending, but MCPX should prevent it. +9. A cancellation notification has no protocol acknowledgement. In this report, + “acknowledge” means the server observes the notification and stops work. The + SDK server implementation maps an incoming cancellation to the matching + request handler's `AbortSignal`. + +## Outcome Matrix + +| Server outcome | Caller result | Late response | Same connection reusable | +| -------------------------------- | ------------------------------------------------ | ------------------------------------- | ----------------------------------- | +| Respects cancellation | Rejected promptly with `McpError(-32001)` | None | Yes | +| Ignores cancellation | Rejected promptly with `McpError(-32001)` | Reported as unknown ID | Yes, before and after late response | +| Responds as cancellation arrives | Cancellation wins in this deterministic ordering | Reported as unknown ID | Yes | +| Completes before cancellation | Fulfilled normally | A stale cancellation is still emitted | Yes | + +Observed cancellation settlement was 15–20ms in four consecutive local runs. +That timing demonstrates prompt settlement under this fixture; it is not an API +latency guarantee. + +## Recommended MCPX Policy + +Use one caller-owned `AbortController` per Call, independent of the shared MCP +connection: + +```ts +type CallCancellation = { + signal: AbortSignal + wasCallerDisconnected: () => boolean +} +``` + +- Create the controller when the Runtime accepts the Call, so the same signal + can remove a queued Call or cancel an active SDK request. +- While the Call is non-terminal, listen to the originating CLI socket's + `close` event and abort with a stable internal reason such as + `originating CLI socket closed`. +- Pass `signal` alongside the existing timeout in the third `callTool` argument. +- Mark the Call terminal and remove the socket listener before returning or + writing a result. This prevents SDK 1.29.0 from emitting stale cancellation + after normal completion. +- Determine the Call's terminal state from MCPX-owned state + (`wasCallerDisconnected()` or an explicit cancellation cause), not from SDK + error code `-32001`, because timeout and caller abort share that code. +- Never write a result to a closed socket. Treat cancellation as best-effort and + do not claim rollback of tool side effects. +- Do not evict a stdio process, HTTP session, or transport solely because one + Call was cancelled or produced a late-response `client.onerror`. Eviction + remains a connection-health decision. + +The current Runtime has no socket lifecycle hook in `handleConnection`, and +`callToolOnConnectedSession` only passes timeout/progress options. Production +implementation should add the cancellation boundary there without moving +caller ownership into the shared `ManagedSession`. diff --git a/prototypes/mcp-cancellation/run.ts b/prototypes/mcp-cancellation/run.ts new file mode 100644 index 0000000..c9ac8c9 --- /dev/null +++ b/prototypes/mcp-cancellation/run.ts @@ -0,0 +1,460 @@ +import { Client } from '@modelcontextprotocol/sdk/client/index.js' +import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js' +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js' +import fs from 'node:fs/promises' +import path from 'node:path' +import { createInterface } from 'node:readline/promises' + +type TransportName = 'stdio' | 'streamable-http' +type ScenarioName = 'acknowledge' | 'ignore' | 'race' | 'complete-before-cancel' + +type FixtureEvent = Record & { event: string } + +type Settlement = { + status: 'fulfilled' | 'rejected' + elapsedMs: number + text?: string + error?: { name: string; code?: number; message: string } +} + +type ScenarioEvidence = { + scenario: ScenarioName + settlement: Settlement + reuseBeforeLate: string + reuseAfterLate: string + fixtureEvents: FixtureEvent[] + clientErrors: string[] +} + +type TransportEvidence = { + transport: TransportName + sessionIds?: Array + scenarios: ScenarioEvidence[] +} + +type Evidence = { + question: string + runtime: { bun: string; sdk: string; sdkLock: string } + abortReason: string + transports: TransportEvidence[] +} + +type Fixture = { + client: Client + events: FixtureEvent[] + errors: string[] + sessionIds?: Array + close: () => Promise +} + +type PendingHttpCall = { + id: number | string + scenario: ScenarioName + controller: ReadableStreamDefaultController + timer?: Timer + requestAborted: boolean +} + +const prototypeDir = import.meta.dir +const scenarios: ScenarioName[] = [ + 'acknowledge', + 'ignore', + 'race', + 'complete-before-cancel', +] +const abortReason = 'originating CLI socket closed' +const encoder = new TextEncoder() + +async function main(): Promise { + if (process.argv.includes('--evidence')) { + const evidence = await runMatrix() + const outputPath = path.join(prototypeDir, 'evidence.json') + await fs.writeFile(outputPath, `${JSON.stringify(evidence, null, '\t')}\n`) + console.log(JSON.stringify(summary(evidence), null, 2)) + console.log(`Evidence written to ${outputPath}`) + return + } + + await runInteractive() +} + +async function runInteractive(): Promise { + const input = createInterface({ + input: process.stdin, + output: process.stdout, + }) + let state: Evidence | undefined + try { + while (true) { + console.clear() + console.log('\x1b[1mMCP cancellation prototype\x1b[0m') + console.log( + state + ? JSON.stringify(summary(state), null, 2) + : '\x1b[2mNo run yet. All state is in memory.\x1b[0m', + ) + console.log('\n\x1b[1m[a]\x1b[0m run all \x1b[1m[q]\x1b[0m quit') + const action = (await input.question('> ')).trim().toLowerCase() + if (action === 'q') return + if (action === 'a') state = await runMatrix() + } + } finally { + input.close() + } +} + +async function runMatrix(): Promise { + const sdkPackage = await Bun.file( + path.join( + prototypeDir, + '../../node_modules/@modelcontextprotocol/sdk/package.json', + ), + ).json() + const transports = await Promise.all([ + runTransport('stdio'), + runTransport('streamable-http'), + ]) + return { + question: + 'Can one caller-owned AbortSignal promptly cancel Client.callTool while preserving stdio and Streamable HTTP connection reuse?', + runtime: { + bun: Bun.version, + sdk: sdkPackage.version, + sdkLock: '@modelcontextprotocol/sdk@1.29.0', + }, + abortReason, + transports, + } +} + +async function runTransport( + transport: TransportName, +): Promise { + const fixture = + transport === 'stdio' ? await startStdioFixture() : await startHttpFixture() + try { + const evidence: ScenarioEvidence[] = [] + for (const scenario of scenarios) { + evidence.push(await runScenario(fixture, scenario)) + } + return { + transport, + sessionIds: fixture.sessionIds, + scenarios: evidence, + } + } finally { + await fixture.close() + } +} + +async function runScenario( + fixture: Fixture, + scenario: ScenarioName, +): Promise { + const eventStart = fixture.events.length + const errorStart = fixture.errors.length + const controller = new AbortController() + const startedAt = performance.now() + const promise = fixture.client.callTool( + { name: 'controlled', arguments: { scenario } }, + undefined, + { signal: controller.signal, timeout: 1_000 }, + ) + + let settlement: Settlement + if (scenario === 'complete-before-cancel') { + settlement = await settle(promise, startedAt) + controller.abort(new Error(abortReason)) + } else { + await Bun.sleep(15) + controller.abort(new Error(abortReason)) + settlement = await settle(promise, startedAt) + } + + const reuseBeforeLate = await echo(fixture.client) + await Bun.sleep(scenario === 'ignore' ? 110 : 35) + const reuseAfterLate = await echo(fixture.client) + + return { + scenario, + settlement, + reuseBeforeLate, + reuseAfterLate, + fixtureEvents: fixture.events.slice(eventStart), + clientErrors: fixture.errors.slice(errorStart), + } +} + +async function settle( + promise: Promise, + startedAt: number, +): Promise { + try { + const result = await promise + return { + status: 'fulfilled', + elapsedMs: Math.round(performance.now() - startedAt), + text: resultText(result), + } + } catch (error) { + const value = error as Error & { code?: number } + return { + status: 'rejected', + elapsedMs: Math.round(performance.now() - startedAt), + error: { + name: value.name, + code: value.code, + message: value.message, + }, + } + } +} + +async function echo(client: Client): Promise { + const result = await client.callTool( + { name: 'echo', arguments: {} }, + undefined, + { timeout: 1_000 }, + ) + return resultText(result) +} + +function resultText(result: unknown): string { + if (!result || typeof result !== 'object' || !('content' in result)) + return JSON.stringify(result) + const content = result.content + if (!Array.isArray(content)) return JSON.stringify(result) + const first = content[0] + return first && typeof first === 'object' && 'text' in first + ? String(first.text) + : JSON.stringify(result) +} + +async function startStdioFixture(): Promise { + const events: FixtureEvent[] = [] + const errors: string[] = [] + const transport = new StdioClientTransport({ + command: process.execPath, + args: [path.join(prototypeDir, 'stdio-fixture.mjs')], + stderr: 'pipe', + }) + const client = new Client({ + name: 'cancellation-prototype', + version: '1.0.0', + }) + client.onerror = (error) => errors.push(error.message) + await client.connect(transport) + transport.stderr?.on('data', (chunk: Buffer) => { + for (const line of chunk.toString('utf8').split('\n')) { + if (!line) continue + events.push(JSON.parse(line)) + } + }) + return { + client, + events, + errors, + close: async () => { + await client.close() + }, + } +} + +async function startHttpFixture(): Promise { + const events: FixtureEvent[] = [] + const errors: string[] = [] + const sessionIds: Array = [] + const pending = new Map() + const completed = new Map() + const server = Bun.serve({ + hostname: '127.0.0.1', + port: 0, + async fetch(request) { + if (request.method === 'GET') return new Response(null, { status: 405 }) + if (request.method !== 'POST') return new Response(null, { status: 405 }) + + const message = (await request.json()) as { + id?: number | string + method: string + params?: Record + } + const sessionId = request.headers.get('mcp-session-id') + sessionIds.push(sessionId) + events.push({ event: 'message', method: message.method, id: message.id }) + + if (message.method === 'initialize') { + return jsonResponse( + message.id, + { + protocolVersion: message.params?.protocolVersion, + capabilities: { tools: {} }, + serverInfo: { + name: 'cancellation-http-fixture', + version: '1.0.0', + }, + }, + { 'mcp-session-id': 'local-session' }, + ) + } + if (message.method === 'notifications/initialized') { + return new Response(null, { + status: 202, + headers: { 'mcp-session-id': 'local-session' }, + }) + } + if (message.method === 'notifications/cancelled') { + const requestId = message.params?.requestId as number | string + const entry = pending.get(requestId) + events.push({ + event: 'cancel', + requestId, + reason: message.params?.reason, + scenario: entry?.scenario ?? completed.get(requestId) ?? 'unknown', + pending: Boolean(entry), + originalRequestAborted: entry?.requestAborted ?? false, + }) + if (entry?.scenario === 'acknowledge') { + clearTimeout(entry.timer) + pending.delete(requestId) + entry.controller.close() + events.push({ + event: 'work-stopped', + requestId, + scenario: entry.scenario, + }) + } + if (entry?.scenario === 'race') { + clearTimeout(entry.timer) + queueMicrotask(() => + completeHttpCall(entry, pending, completed, events), + ) + } + return new Response(null, { + status: 202, + headers: { 'mcp-session-id': 'local-session' }, + }) + } + if (message.method !== 'tools/call') + return new Response(null, { status: 202 }) + if (message.params?.name === 'echo') { + return jsonResponse(message.id, toolResult('echo-ok'), { + 'mcp-session-id': 'local-session', + }) + } + + const scenario = message.params?.arguments?.scenario as ScenarioName + events.push({ event: 'call', id: message.id, scenario, sessionId }) + if (scenario === 'complete-before-cancel') { + await Bun.sleep(5) + completed.set(message.id!, scenario) + events.push({ event: 'response', id: message.id, scenario }) + return jsonResponse(message.id, toolResult(scenario), { + 'mcp-session-id': 'local-session', + }) + } + + let controller!: ReadableStreamDefaultController + const stream = new ReadableStream({ + start(value) { + controller = value + }, + }) + const entry: PendingHttpCall = { + id: message.id!, + scenario, + controller, + requestAborted: request.signal.aborted, + } + request.signal.addEventListener('abort', () => { + entry.requestAborted = true + events.push({ event: 'request-aborted', id: entry.id, scenario }) + }) + pending.set(entry.id, entry) + entry.timer = setTimeout( + () => completeHttpCall(entry, pending, completed, events), + scenario === 'ignore' ? 80 : 2_000, + ) + return new Response(stream, { + status: 200, + headers: { + 'content-type': 'text/event-stream', + 'mcp-session-id': 'local-session', + }, + }) + }, + }) + + const transport = new StreamableHTTPClientTransport( + new URL(`http://127.0.0.1:${server.port}/mcp`), + ) + const client = new Client({ + name: 'cancellation-prototype', + version: '1.0.0', + }) + client.onerror = (error) => errors.push(error.message) + await client.connect(transport) + return { + client, + events, + errors, + sessionIds, + close: async () => { + await client.close() + server.stop(true) + }, + } +} + +function completeHttpCall( + entry: PendingHttpCall, + pending: Map, + completed: Map, + events: FixtureEvent[], +): void { + pending.delete(entry.id) + completed.set(entry.id, entry.scenario) + events.push({ event: 'response', id: entry.id, scenario: entry.scenario }) + entry.controller.enqueue( + encoder.encode( + `event: message\ndata: ${JSON.stringify({ + jsonrpc: '2.0', + id: entry.id, + result: toolResult(entry.scenario), + })}\n\n`, + ), + ) + entry.controller.close() +} + +function toolResult(text: string): Record { + return { content: [{ type: 'text', text }] } +} + +function jsonResponse( + id: number | string | undefined, + result: unknown, + extraHeaders: Record = {}, +): Response { + return Response.json( + { jsonrpc: '2.0', id, result }, + { headers: extraHeaders }, + ) +} + +function summary(evidence: Evidence): Record { + return { + runtime: evidence.runtime, + transports: evidence.transports.map((transport) => ({ + transport: transport.transport, + scenarios: transport.scenarios.map((scenario) => ({ + scenario: scenario.scenario, + settlement: scenario.settlement, + reuseBeforeLate: scenario.reuseBeforeLate, + reuseAfterLate: scenario.reuseAfterLate, + clientErrors: scenario.clientErrors, + })), + })), + } +} + +await main() diff --git a/prototypes/mcp-cancellation/stdio-fixture.mjs b/prototypes/mcp-cancellation/stdio-fixture.mjs new file mode 100644 index 0000000..2bdd987 --- /dev/null +++ b/prototypes/mcp-cancellation/stdio-fixture.mjs @@ -0,0 +1,128 @@ +import readline from 'node:readline' + +const pending = new Map() +const completed = new Map() + +function event(value) { + process.stderr.write(`${JSON.stringify(value)}\n`) +} + +function send(value) { + process.stdout.write(`${JSON.stringify(value)}\n`) +} + +function result(id, text) { + send({ + jsonrpc: '2.0', + id, + result: { content: [{ type: 'text', text }] }, + }) +} + +function complete(entry, text) { + pending.delete(entry.id) + completed.set(entry.id, entry.scenario) + event({ event: 'response', id: entry.id, scenario: entry.scenario }) + result(entry.id, text) +} + +function startCall(message) { + const scenario = message.params?.arguments?.scenario + const entry = { id: message.id, scenario, timer: undefined } + pending.set(message.id, entry) + event({ event: 'call', id: message.id, scenario }) + + if (scenario === 'complete-before-cancel') { + entry.timer = setTimeout(() => complete(entry, scenario), 5) + return + } + if (scenario === 'ignore') { + entry.timer = setTimeout(() => complete(entry, scenario), 80) + return + } + + // A long fallback makes a missing cancellation observable without hanging forever. + entry.timer = setTimeout(() => complete(entry, `${scenario}-fallback`), 2_000) +} + +function cancel(message) { + const requestId = message.params?.requestId + const entry = pending.get(requestId) + event({ + event: 'cancel', + requestId, + reason: message.params?.reason, + scenario: entry?.scenario ?? completed.get(requestId) ?? 'unknown', + pending: Boolean(entry), + }) + if (!entry) return + if (entry.scenario === 'acknowledge') { + clearTimeout(entry.timer) + pending.delete(requestId) + event({ event: 'work-stopped', requestId, scenario: entry.scenario }) + return + } + if (entry.scenario === 'race') { + clearTimeout(entry.timer) + queueMicrotask(() => complete(entry, entry.scenario)) + } +} + +function handle(message) { + event({ event: 'message', method: message.method, id: message.id }) + if (message.method === 'initialize') { + send({ + jsonrpc: '2.0', + id: message.id, + result: { + protocolVersion: message.params.protocolVersion, + capabilities: { tools: {} }, + serverInfo: { name: 'cancellation-stdio-fixture', version: '1.0.0' }, + }, + }) + return + } + if (message.method === 'notifications/initialized') return + if (message.method === 'tools/list') { + send({ + jsonrpc: '2.0', + id: message.id, + result: { + tools: [ + { name: 'echo', inputSchema: { type: 'object', properties: {} } }, + { name: 'fail', inputSchema: { type: 'object', properties: {} } }, + { + name: 'controlled', + inputSchema: { + type: 'object', + properties: { scenario: { type: 'string' } }, + required: ['scenario'], + }, + }, + ], + }, + }) + return + } + if (message.method === 'notifications/cancelled') { + cancel(message) + return + } + if (message.method !== 'tools/call') return + if (message.params?.name === 'echo') { + result(message.id, 'echo-ok') + return + } + if (message.params?.name === 'fail') { + send({ + jsonrpc: '2.0', + id: message.id, + error: { code: -32000, message: 'fixture failure' }, + }) + return + } + startCall(message) +} + +const lines = readline.createInterface({ input: process.stdin }) +lines.on('line', (line) => handle(JSON.parse(line))) diff --git a/src/authentication-coordinator.ts b/src/authentication-coordinator.ts new file mode 100644 index 0000000..813c197 --- /dev/null +++ b/src/authentication-coordinator.ts @@ -0,0 +1,209 @@ +import type { RuntimeCaller } from './runtime-caller' + +import { RuntimeOperationError } from './runtime-call' + +export type AuthenticationFlow = { + start: ( + signal: AbortSignal, + requestInput: RuntimeCaller['requestInput'], + ) => Promise + persist: (value: T) => Promise +} + +export type AuthenticationWaiterOutcome = + | { status: 'completed' } + | { status: 'disconnected' } + +type Waiter = { + caller: RuntimeCaller + unsubscribe: () => void + resolve: (outcome: AuthenticationWaiterOutcome) => void + reject: (error: Error) => void +} + +type ActiveFlow = { + controller: AbortController + waiters: Map + timer: Timer + finished: boolean + flow: AuthenticationFlow + run?: Promise +} + +const DEFAULT_AUTH_TIMEOUT_MS = 5 * 60 * 1000 + +export class AuthenticationCoordinator { + readonly #timeoutMs: number + readonly #flows = new Map>() + readonly #running = new Set>() + #accepting = true + + constructor(options: { timeoutMs?: number } = {}) { + this.#timeoutMs = options.timeoutMs ?? DEFAULT_AUTH_TIMEOUT_MS + } + + join( + identity: string, + caller: RuntimeCaller, + flow: AuthenticationFlow, + ): Promise { + if (!this.#accepting) { + return Promise.reject( + new RuntimeOperationError('cancelled', 'MCP Runtime is stopping.'), + ) + } + let entry = this.#flows.get(identity) as ActiveFlow | undefined + let created = false + if (!entry) { + entry = this.#createFlow(identity, flow) + this.#flows.set(identity, entry as ActiveFlow) + this.#running.add(entry as ActiveFlow) + created = true + } + + const waiterPromise = new Promise( + (resolve, reject) => { + const waiter: Waiter = { + caller, + resolve, + reject, + unsubscribe: () => {}, + } + entry.waiters.set(caller.id, waiter) + const unsubscribe = caller.onDisconnect(() => { + this.#removeWaiter(identity, entry, waiter) + }) + waiter.unsubscribe = unsubscribe + if (!entry.waiters.has(caller.id)) unsubscribe() + }, + ) + // Register the first waiter before starting code that may throw synchronously. + if (created) entry.run = this.#runFlow(identity, entry) + return waiterPromise + } + + activeFlows(): number { + return this.#running.size + } + + async close(): Promise { + this.#accepting = false + const entries = [...this.#running] + for (const entry of entries) { + entry.controller.abort( + new AuthenticationCancelled('MCP Runtime is stopping.'), + ) + } + await Promise.all(entries.map((entry) => entry.run).filter(Boolean)) + } + + #createFlow(identity: string, flow: AuthenticationFlow): ActiveFlow { + const controller = new AbortController() + const timer = setTimeout(() => { + controller.abort( + new AuthenticationTimeout(`Authentication Flow ${identity} timed out.`), + ) + }, this.#timeoutMs) + timer.unref() + return { + controller, + waiters: new Map(), + timer, + finished: false, + flow, + } + } + + async #runFlow(identity: string, entry: ActiveFlow): Promise { + let failure: Error | undefined + try { + const value = await entry.flow.start(entry.controller.signal, (request) => + this.#requestInput(entry, request), + ) + await entry.flow.persist(value) + } catch (error) { + failure = authenticationError(error) + } finally { + entry.finished = true + this.#running.delete(entry as ActiveFlow) + clearTimeout(entry.timer) + if (this.#flows.get(identity) === entry) this.#flows.delete(identity) + } + + const waiters = [...entry.waiters.values()] + entry.waiters.clear() + for (const waiter of waiters) { + waiter.unsubscribe() + if (failure) waiter.reject(failure) + else waiter.resolve({ status: 'completed' }) + } + } + + async #requestInput( + entry: ActiveFlow, + request: Parameters[0], + ): Promise { + while (!entry.controller.signal.aborted) { + const waiter = entry.waiters.values().next().value as Waiter | undefined + if (!waiter) + throw new AuthenticationCancelled( + 'Authentication Flow has no caller for interactive input.', + ) + try { + return await waiter.caller.requestInput( + request, + entry.controller.signal, + ) + } catch (error) { + // A surviving waiter can take over interaction from a disconnected caller. + if (entry.waiters.has(waiter.caller.id)) throw error + } + } + throw new AuthenticationCancelled('Authentication Flow was cancelled.') + } + + #removeWaiter( + identity: string, + entry: ActiveFlow, + waiter: Waiter, + ): void { + if (!entry.waiters.delete(waiter.caller.id)) return + waiter.unsubscribe() + waiter.resolve({ status: 'disconnected' }) + if (entry.waiters.size === 0 && !entry.finished) { + if (this.#flows.get(identity) === entry) this.#flows.delete(identity) + entry.controller.abort( + new AuthenticationCancelled( + 'Authentication Flow lost its final caller.', + ), + ) + } + } +} + +class AuthenticationTimeout extends Error { + constructor(message: string) { + super(message) + this.name = 'AuthenticationTimeout' + } +} + +class AuthenticationCancelled extends Error { + constructor(message: string) { + super(message) + this.name = 'AuthenticationCancelled' + } +} + +function authenticationError(error: unknown): Error { + if (error instanceof AuthenticationTimeout) { + return new RuntimeOperationError('timeout', error.message) + } + if (error instanceof AuthenticationCancelled) { + return new RuntimeOperationError('cancelled', error.message) + } + return new RuntimeOperationError( + 'operation-failed', + error instanceof Error ? error.message : String(error), + ) +} diff --git a/src/bearer.ts b/src/bearer.ts index 0df2843..a4d49e8 100644 --- a/src/bearer.ts +++ b/src/bearer.ts @@ -1,22 +1,7 @@ import { createHash } from 'node:crypto' -import fs from 'node:fs/promises' -import path from 'node:path' import type { AuthDiscovery, BearerCredential } from './types' -import { daemonDir } from './daemon-paths' - -type BearerCursorState = { - version: 1 - cursors: Record -} - -const CURSOR_FILE = 'bearer-cursors.json' -const LOCK_DIR = 'bearer-cursors.lock' -const LOCK_RETRY_MS = 25 -const LOCK_TIMEOUT_MS = 1_000 -const LOCK_STALE_MS = 30_000 - export function authFromBearerValues( values: string | string[] | undefined, ): AuthDiscovery | undefined { @@ -57,18 +42,11 @@ export function resolveBearerHeaderForProbe( } export async function resolveBearerHeader( - serverUrl: string, + _serverUrl: string, auth: Extract, ): Promise { - if (auth.credentials.length === 0) - throw new Error('Bearer auth requires at least one credential.') - const credential = - auth.credentials.length === 1 - ? auth.credentials[0] - : await selectRoundRobinCredential(serverUrl, auth) - if (!credential) - throw new Error('Bearer auth requires at least one credential.') - return normalizeBearerToken(resolveBearerCredential(credential)) + // Runtime call activation owns round-robin; discovery only probes one identity. + return resolveBearerHeaderForProbe(auth) } function parseBearerCredential(value: string): BearerCredential { @@ -94,9 +72,8 @@ function parseEnvReference(value: string): string | undefined { function resolveBearerCredential(credential: BearerCredential): string { if (credential.kind === 'literal') return credential.value const value = process.env[credential.name] - if (!value) { + if (!value) throw new Error(`Bearer env reference "${credential.name}" is not set.`) - } return value } @@ -104,105 +81,6 @@ function normalizeBearerToken(value: string): string { return value.startsWith('Bearer ') ? value : `Bearer ${value}` } -async function selectRoundRobinCredential( - serverUrl: string, - auth: Extract, -): Promise { - return withCursorLock(async () => { - const state = await readCursorState() - const key = cursorKey(serverUrl, auth) - const current = state.cursors[key] ?? 0 - const index = current % auth.credentials.length - state.cursors[key] = (index + 1) % auth.credentials.length - await writeCursorState(state) - const credential = auth.credentials[index] - if (!credential) - throw new Error('Bearer auth requires at least one credential.') - return credential - }) -} - -async function withCursorLock(run: () => Promise): Promise { - const lockPath = path.join(daemonDir(), LOCK_DIR) - const deadline = Date.now() + LOCK_TIMEOUT_MS - - while (true) { - try { - await fs.mkdir(daemonDir(), { recursive: true, mode: 0o700 }) - await fs.mkdir(lockPath, { recursive: false }) - break - } catch (error) { - if ((error as NodeJS.ErrnoException).code !== 'EEXIST') { - throw error - } - if (await isStaleLock(lockPath)) { - // A crashed CLI can leave the mkdir lock behind; reclaim it after the normal call window. - await fs.rm(lockPath, { recursive: true, force: true }) - continue - } - if (Date.now() >= deadline) { - throw new Error('Timed out waiting for bearer round-robin state lock.') - } - await sleep(LOCK_RETRY_MS) - } - } - - try { - return await run() - } finally { - await fs.rm(lockPath, { recursive: true, force: true }) - } -} - -async function isStaleLock(lockPath: string): Promise { - try { - const stat = await fs.stat(lockPath) - return Date.now() - stat.mtimeMs > LOCK_STALE_MS - } catch { - return false - } -} - -async function readCursorState(): Promise { - try { - const raw = await fs.readFile(cursorPath(), 'utf8') - const parsed = JSON.parse(raw) as BearerCursorState - if ( - parsed.version !== 1 || - !parsed.cursors || - typeof parsed.cursors !== 'object' - ) { - throw new Error(`Invalid bearer cursor state at ${cursorPath()}.`) - } - return { version: 1, cursors: parsed.cursors } - } catch (error) { - if ((error as NodeJS.ErrnoException).code === 'ENOENT') - return { version: 1, cursors: {} } - throw error - } -} - -async function writeCursorState(state: BearerCursorState): Promise { - await fs.mkdir(daemonDir(), { recursive: true, mode: 0o700 }) - await fs.writeFile(cursorPath(), `${JSON.stringify(state, null, 2)}\n`, { - encoding: 'utf8', - mode: 0o600, - }) -} - -function cursorPath(): string { - return path.join(daemonDir(), CURSOR_FILE) -} - -function cursorKey( - serverUrl: string, - auth: Extract, -): string { - return createHash('sha256') - .update(JSON.stringify({ serverUrl, credentials: bearerAuthRef(auth) })) - .digest('hex') -} - function hashSecret(value: string): string { return createHash('sha256').update(value).digest('hex').slice(0, 16) } @@ -210,7 +88,3 @@ function hashSecret(value: string): string { function isEnvName(value: string): boolean { return /^[A-Za-z_][A-Za-z0-9_]*$/.test(value) } - -async function sleep(ms: number): Promise { - await new Promise((resolve) => setTimeout(resolve, ms)) -} diff --git a/src/config.ts b/src/config.ts deleted file mode 100644 index d6e965d..0000000 --- a/src/config.ts +++ /dev/null @@ -1,106 +0,0 @@ -import fs from 'node:fs/promises' -import { homedir } from 'node:os' -import path from 'node:path' - -import type { RegistryConfig, ServerConfig } from './types' - -import { assignCommandNames } from './names' - -const CONFIG_PATH = path.join('.agents', 'mcpx', 'servers.json') - -export function getRegistryConfigPath(): string { - return path.join(homedir(), CONFIG_PATH) -} - -export async function readRegistryConfig(): Promise { - const filePath = getRegistryConfigPath() - try { - const raw = await fs.readFile(filePath, 'utf8') - const parsed = JSON.parse(raw) as RegistryConfig - if ( - parsed.version !== 1 || - !parsed.servers || - typeof parsed.servers !== 'object' - ) { - throw new Error(`Invalid mcpx registry config at ${filePath}.`) - } - return normalizeRegistryConfig(parsed) - } catch (error) { - if ((error as NodeJS.ErrnoException).code === 'ENOENT') { - return { version: 1, servers: {} } - } - throw error - } -} - -export async function writeRegistryConfig( - config: RegistryConfig, -): Promise { - const filePath = getRegistryConfigPath() - await fs.mkdir(path.dirname(filePath), { recursive: true }) - await fs.writeFile(filePath, `${JSON.stringify(config, null, 2)}\n`, 'utf8') -} - -export async function upsertServerConfig( - name: string, - server: ServerConfig, -): Promise { - const config = await readRegistryConfig() - config.servers[name] = server - await writeRegistryConfig(config) - return config -} - -export async function removeServerConfig( - name: string, -): Promise { - const config = await readRegistryConfig() - const removed = removeServerFromConfig(config, name) - if (removed) { - await writeRegistryConfig(config) - } - return removed -} - -export function removeServerFromConfig( - config: RegistryConfig, - name: string, -): ServerConfig | undefined { - const removed = config.servers[name] - if (!removed) return undefined - delete config.servers[name] - return removed -} - -export function normalizeRegistryConfig( - config: RegistryConfig, -): RegistryConfig { - const servers: Record = {} - - for (const [name, server] of Object.entries(config.servers)) { - servers[name] = normalizeServerConfig(server) - } - - return { - version: config.version, - servers, - } -} - -function normalizeServerConfig(server: ServerConfig): ServerConfig { - if (!server.tools || server.tools.length === 0) return server - - const commandNames = assignCommandNames(server.tools.map((tool) => tool.name)) - return { - ...server, - tools: server.tools.map((tool) => { - const { outputSchema: _outputSchema, ...rest } = tool as typeof tool & { - outputSchema?: unknown - } - return { - ...rest, - commandName: commandNames.get(tool.name) ?? tool.name, - } - }), - } -} diff --git a/src/daemon-client.ts b/src/daemon-client.ts index a122f97..955eaae 100644 --- a/src/daemon-client.ts +++ b/src/daemon-client.ts @@ -1,151 +1,18 @@ import net from 'node:net' -import type { McpTool, ServerConfig } from './types' - import { requestJsonLine } from './daemon-io' import { ensureDaemonDir, daemonSocketPath } from './daemon-paths' import { DAEMON_PROTOCOL_VERSION, - buildServerKey, helloMessage, - notificationModeFromEnv, - type ClientMessage, type DaemonMessage, - type DaemonStatus, } from './daemon-protocol' -import { daemonOutputEnvelope } from './daemon-result' -import { resolveHeadersWithState } from './headers' import { MCPX_VERSION } from './version' const START_TIMEOUT_MS = 3_000 const CONNECT_RETRY_MS = 50 -export async function listToolsViaDaemon( - server: ServerConfig, - serverName: string, -): Promise { - const context = await daemonRequestContext(server) - if (context.authRefreshed) { - await evictDaemonSession( - context.serverKey, - 'auth-refreshed', - process.argv[1] ?? import.meta.path, - ) - } - const message: ClientMessage = { - op: 'listTools', - callId: crypto.randomUUID(), - serverName, - serverKey: context.serverKey, - server, - } - if (context.headers) message.headers = context.headers - const result = await requestDaemon( - message, - process.argv[1] ?? import.meta.path, - ) - return result as McpTool[] -} - -export async function callToolViaDaemon( - server: ServerConfig, - serverName: string, - toolName: string, - input: Record, -): Promise { - const context = await daemonRequestContext(server) - if (context.authRefreshed) { - await evictDaemonSession( - context.serverKey, - 'auth-refreshed', - process.argv[1] ?? import.meta.path, - ) - } - const message: ClientMessage = { - op: 'call', - callId: crypto.randomUUID(), - serverName, - serverKey: context.serverKey, - server, - toolName, - input, - notificationMode: notificationModeFromEnv(), - } - if (context.headers) message.headers = context.headers - const response = await requestDaemonMessage( - message, - process.argv[1] ?? import.meta.path, - ) - if ( - response.ok && - (response.notifications?.length || response.toolsChanged) - ) { - return daemonOutputEnvelope({ - result: response.result, - notifications: response.notifications ?? [], - toolsChanged: response.toolsChanged === true, - }) - } - return response.ok ? response.result : undefined -} - -export async function daemonStatus(mainPath: string): Promise { - return requestDaemon({ op: 'status' }, mainPath, { - start: false, - }) as Promise -} - -export async function stopDaemon(mainPath: string): Promise { - return requestDaemon({ op: 'stop' }, mainPath, { start: false }) -} - -async function evictDaemonSession( - serverKey: string, - reason: 'auth-refreshed' | 'unauthorized' | 'manual', - mainPath: string, -): Promise { - await requestDaemon({ op: 'evictSession', serverKey, reason }, mainPath) -} - -async function daemonRequestContext(server: ServerConfig): Promise<{ - serverKey: string - headers?: Record - authRefreshed: boolean -}> { - const serverKey = buildServerKey(server) - if (server.transport === 'stdio') return { serverKey, authRefreshed: false } - const resolved = await resolveHeadersWithState(server) - return { - serverKey, - headers: resolved.headers, - authRefreshed: resolved.authRefreshed, - } -} - -async function requestDaemon( - message: ClientMessage, - mainPath: string, - options: { start?: boolean } = {}, -): Promise { - const response = await requestDaemonMessage(message, mainPath, options) - if (response.ok) return response.result ?? response - throw new Error(response.error.message) -} - -async function requestDaemonMessage( - message: ClientMessage, - mainPath: string, - options: { start?: boolean } = {}, -): Promise { - const start = options.start ?? true - if (start) await ensureDaemon(mainPath) - return withDaemonConnection(async (socket) => { - await sendAndExpectOk(socket, helloMessage()) - return sendAndExpectDaemonMessage(socket, message) - }) -} - -async function ensureDaemon(mainPath: string): Promise { +export async function ensureDaemon(mainPath: string): Promise { const state = await probeDaemon() if (state === 'compatible') return if (state === 'incompatible') { @@ -159,10 +26,7 @@ async function ensureDaemon(mainPath: string): Promise { await ensureDaemonDir() Bun.spawn([process.execPath, mainPath, '@daemon', 'server'], { - env: { - ...process.env, - MCPX_DAEMON_SERVER: '1', - }, + env: { ...process.env, MCPX_DAEMON_SERVER: '1' }, stdin: 'ignore', stdout: 'ignore', stderr: 'ignore', @@ -170,111 +34,72 @@ async function ensureDaemon(mainPath: string): Promise { const deadline = Date.now() + START_TIMEOUT_MS while (Date.now() < deadline) { - if (await canHandshake()) return - await sleep(CONNECT_RETRY_MS) + if ((await probeDaemon()) === 'compatible') return + await Bun.sleep(CONNECT_RETRY_MS) } throw new Error('mcpxd did not start before the startup timeout.') } -async function canHandshake(): Promise { - return (await probeDaemon()) === 'compatible' +export async function connectDaemonSocket(): Promise { + return new Promise((resolve, reject) => { + const socket = net.createConnection(daemonSocketPath()) + socket.once('connect', () => resolve(socket)) + socket.once('error', reject) + }) } async function probeDaemon(): Promise< 'compatible' | 'incompatible' | 'missing' > { + let socket: net.Socket | undefined try { - const result = await withDaemonConnection((socket) => - sendAndExpectOk(socket, helloMessage()), - ) - const protocolVersion = - typeof result === 'object' && - result !== null && - 'protocolVersion' in result - ? result.protocolVersion - : undefined - const daemonVersion = - typeof result === 'object' && result !== null && 'version' in result - ? result.version - : undefined - return protocolVersion === DAEMON_PROTOCOL_VERSION && - daemonVersion === MCPX_VERSION + socket = await connectDaemonSocket() + const response = await requestJsonLine(socket, helloMessage()) + if (!isDaemonMessage(response) || !response.ok) return 'incompatible' + const result = response.result as + | { protocolVersion?: unknown; version?: unknown } + | undefined + return response.protocolVersion === DAEMON_PROTOCOL_VERSION && + result?.protocolVersion === DAEMON_PROTOCOL_VERSION && + result.version === MCPX_VERSION ? 'compatible' : 'incompatible' } catch { return 'missing' + } finally { + socket?.destroy() } } async function stopIncompatibleDaemon(): Promise { - await withDaemonConnection(async (socket) => { - await sendAndExpectOk(socket, helloMessage(), { - allowProtocolMismatch: true, - }) - await sendAndExpectOk(socket, { op: 'stop' }) - }) - const deadline = Date.now() + START_TIMEOUT_MS - while (Date.now() < deadline) { - if ((await probeDaemon()) === 'missing') return - await sleep(CONNECT_RETRY_MS) - } - throw new Error('Incompatible mcpxd did not stop before the timeout.') -} - -async function connectSocket(): Promise { - return new Promise((resolve, reject) => { - const socket = net.createConnection(daemonSocketPath()) - socket.once('connect', () => resolve(socket)) - socket.once('error', reject) - }) -} - -async function sendAndExpectOk( - socket: net.Socket, - message: ClientMessage, - options: { allowProtocolMismatch?: boolean } = {}, -): Promise { - const response = await requestJsonLine(socket, message) - if (isDaemonMessage(response) && response.ok) - return response.result ?? response - if (isDaemonMessage(response) && !response.ok) { + const socket = await connectDaemonSocket() + try { + const hello = await requestJsonLine(socket, helloMessage()) if ( - options.allowProtocolMismatch && - response.error.code === 'protocol-mismatch' + isDaemonMessage(hello) && + hello.ok && + hello.protocolVersion === DAEMON_PROTOCOL_VERSION ) { - return response + await requestJsonLine(socket, { + requestId: crypto.randomUUID(), + op: 'stop', + }) + } else { + // V2 accepts its stop command after returning protocol-mismatch. + await requestJsonLine(socket, { op: 'stop' }) } - throw new Error(response.error.message) - } - throw new Error('Invalid mcpxd response.') -} - -async function sendAndExpectDaemonMessage( - socket: net.Socket, - message: ClientMessage, -): Promise { - const response = await requestJsonLine(socket, message) - if (isDaemonMessage(response) && response.ok) return response - if (isDaemonMessage(response) && !response.ok) - throw new Error(response.error.message) - throw new Error('Invalid mcpxd response.') -} - -async function withDaemonConnection( - run: (socket: net.Socket) => Promise, -): Promise { - const socket = await connectSocket() - try { - return await run(socket) } finally { socket.end() } + + const deadline = Date.now() + START_TIMEOUT_MS + while (Date.now() < deadline) { + if ((await probeDaemon()) === 'missing') return + await Bun.sleep(CONNECT_RETRY_MS) + } + throw new Error('Incompatible mcpxd did not stop before the timeout.') } function isDaemonMessage(value: unknown): value is DaemonMessage { return !!value && typeof value === 'object' && 'ok' in value } - -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)) -} diff --git a/src/daemon-protocol.ts b/src/daemon-protocol.ts index b842783..51479d0 100644 --- a/src/daemon-protocol.ts +++ b/src/daemon-protocol.ts @@ -1,16 +1,13 @@ -import { createHash } from 'node:crypto' - -import type { ServerConfig } from './types' - -import { bearerAuthRef } from './bearer' +import { + DAEMON_PROTOCOL_VERSION, + type NotificationMode, +} from './runtime-protocol' import { MCPX_VERSION } from './version' -export const DAEMON_PROTOCOL_VERSION = 2 -export const DAEMON_ENV = 'MCPX_DAEMON_SERVER' -export const DISABLE_DAEMON_ENV = 'MCPX_DISABLE_DAEMON' -export const NOTIFICATION_MODE_ENV = 'MCPX_NOTIFICATION_MODE' +export { DAEMON_PROTOCOL_VERSION } +export type { NotificationMode } -export type NotificationMode = 'buffer' | 'discard' +export const NOTIFICATION_MODE_ENV = 'MCPX_NOTIFICATION_MODE' export type McpNotification = | { @@ -27,71 +24,16 @@ export type McpNotification = | { method: '$oversize'; params: { savedTo: string } } | { method: string; params?: unknown } -export type DaemonStatus = { - pid: number +export type ClientMessage = { + op: 'hello' protocolVersion: number - version: string - activeServers: number - servers: { - serverKey: string - transport: 'stdio' | 'http' - labels: string[] - pid?: number | null - url?: string - activeCalls: number - queuedCalls: number - idleMs: number - evictCount: number - hasRetainedSessionId: boolean - sessionIdHash?: string - }[] + clientVersion: string } -export type ClientMessage = - | { op: 'hello'; protocolVersion: number; clientVersion: string } - | { - op: 'listTools' - callId: string - serverName: string - serverKey: string - server: ServerConfig - headers?: Record - } - | { - op: 'call' - callId: string - serverName: string - serverKey: string - server: ServerConfig - headers?: Record - toolName: string - input: Record - notificationMode?: NotificationMode - } - | { op: 'status' } - | { op: 'stop' } - | { - op: 'evictSession' - serverKey: string - reason?: 'auth-refreshed' | 'unauthorized' | 'manual' - } - export type DaemonMessage = - | { - ok: true - protocolVersion?: number - result?: unknown - notifications?: McpNotification[] - toolsChanged?: boolean - } + | { ok: true; protocolVersion?: number; result?: unknown } | { ok: false; error: { code: string; message: string } } -export function shouldUseDaemon(): boolean { - return ( - process.env[DAEMON_ENV] !== '1' && process.env[DISABLE_DAEMON_ENV] !== '1' - ) -} - export function notificationModeFromEnv(): NotificationMode { const raw = process.env[NOTIFICATION_MODE_ENV] if (!raw || raw === 'buffer') return 'buffer' @@ -101,29 +43,6 @@ export function notificationModeFromEnv(): NotificationMode { ) } -export function buildServerKey(server: ServerConfig): string { - const payload = - server.transport === 'stdio' - ? stableJson({ - command: server.command, - args: server.args ?? [], - env: server.env ?? {}, - cwd: server.cwd ?? null, - }) - : stableJson({ - type: 'http', - url: server.url, - authKind: server.auth?.kind ?? null, - authRef: - server.auth?.kind === 'bearer' - ? bearerAuthRef(server.auth) - : server.auth?.kind === 'oauth-token' - ? server.auth.tokenKey - : null, - }) - return createHash('sha256').update(payload).digest('hex').slice(0, 16) -} - export function helloMessage(): ClientMessage { return { op: 'hello', @@ -131,19 +50,3 @@ export function helloMessage(): ClientMessage { clientVersion: MCPX_VERSION, } } - -function stableJson(value: unknown): string { - return JSON.stringify(sortValue(value)) -} - -function sortValue(value: unknown): unknown { - if (Array.isArray(value)) return value.map(sortValue) - if (!value || typeof value !== 'object') return value - - const record = value as Record - const sorted: Record = {} - for (const key of Object.keys(record).sort()) { - sorted[key] = sortValue(record[key]) - } - return sorted -} diff --git a/src/daemon-server.ts b/src/daemon-server.ts index 43ce564..bc3d845 100644 --- a/src/daemon-server.ts +++ b/src/daemon-server.ts @@ -1,71 +1,28 @@ -import { createHash } from 'node:crypto' import fs from 'node:fs/promises' import net, { type Socket } from 'node:net' -import { tmpdir } from 'node:os' -import path from 'node:path' -import type { McpTool, ServerConfig } from './types' +import type { SocketRuntimeCaller } from './runtime-caller' import { readJsonLines, requestJsonLine, writeJsonLine } from './daemon-io' -import { - ensureDaemonDir, - daemonLogPath, - daemonSocketPath, - serverLogPath, -} from './daemon-paths' +import { daemonDir, ensureDaemonDir, daemonSocketPath } from './daemon-paths' import { DAEMON_PROTOCOL_VERSION, type ClientMessage, type DaemonMessage, - type DaemonStatus, - type McpNotification, } from './daemon-protocol' -import { - connectMcpClient, - listAllMcpTools, - toolCallRequestOptions, - type McpConnection, -} from './mcp-client' -import { - createNotificationBuffer, - type NotificationBuffer, -} from './notifications' +import { McpRuntime } from './runtime' +import { createSocketRuntimeCaller } from './runtime-caller' +import { isRuntimeInputFrame, parseRuntimeIntent } from './runtime-protocol' +import { openRuntimeStores } from './runtime-stores' import { MCPX_VERSION } from './version' -const CHILD_IDLE_TTL_MS = 15 * 60 * 1000 -const DAEMON_IDLE_TTL_MS = 30 * 60 * 1000 -const CLEANUP_INTERVAL_MS = 30 * 1000 -const LOG_MAX_BYTES = 10 * 1024 * 1024 -const EVICT_DEADLINE_MS = 5_000 - -type ConnectedSession = McpConnection - -type ManagedSession = { - serverKey: string - labels: Set - server: ServerConfig - headers?: Record | undefined - connection?: ConnectedSession - connecting?: Promise - queue: Promise - activeCalls: number - queuedCalls: number - lastUsedAt: number - evictCount: number - lastSessionId?: string | undefined - currentBuffer?: NotificationBuffer | undefined - pendingToolsChanged: boolean -} - -type DaemonCallResult = { - result: unknown - notifications: ReturnType - toolsChanged: boolean -} - -const sessions = new Map() +let runtimePromise: Promise | undefined let stopping = false -let lastDaemonActivity = Date.now() +let lastActivityAt = Date.now() + +const SESSION_IDLE_MS = 15 * 60 * 1000 +const DAEMON_IDLE_MS = 30 * 60 * 1000 +const CLEANUP_INTERVAL_MS = 30 * 1000 export async function runDaemonServer(): Promise { await ensureDaemonDir() @@ -73,7 +30,11 @@ export async function runDaemonServer(): Promise { if (await isLiveSocket(socketPath)) return await fs.rm(socketPath, { force: true }).catch(() => {}) - const server = net.createServer(handleConnection) + const server = net.createServer((socket) => handleConnection(server, socket)) + const cleanupTimer = setInterval(() => { + void cleanupIdleRuntime(server) + }, CLEANUP_INTERVAL_MS) + cleanupTimer.unref() try { await listen(server, socketPath) } catch (error) { @@ -85,454 +46,147 @@ export async function runDaemonServer(): Promise { } throw error } - await logDaemon(`mcpxd started pid=${process.pid}`) - - const cleanupTimer = setInterval(() => { - void cleanupIdleSessions(server) - }, CLEANUP_INTERVAL_MS) - await new Promise((resolve) => { - server.on('close', resolve) - }) + await new Promise((resolve) => server.on('close', resolve)) clearInterval(cleanupTimer) } -function handleConnection(socket: Socket): void { +function handleConnection(server: net.Server, socket: Socket): void { + let phase: 'awaiting-hello' | 'ready' | 'complete' = 'awaiting-hello' + let pending = Promise.resolve() + let activeCaller: SocketRuntimeCaller | undefined readJsonLines( socket, (message) => { - void handleMessage(socket, message) - }, - (error) => { - writeJsonLine(socket, errorResponse('invalid-json', error.message)) - }, - ) -} - -async function handleMessage(socket: Socket, message: unknown): Promise { - lastDaemonActivity = Date.now() - if (!isClientMessage(message)) { - writeJsonLine( - socket, - errorResponse('invalid-message', 'Invalid mcpxd message.'), - ) - return - } - - try { - if ( - message.op === 'hello' && - message.protocolVersion !== DAEMON_PROTOCOL_VERSION - ) { - writeJsonLine( - socket, - errorResponse( - 'protocol-mismatch', - `Unsupported mcpxd protocol ${message.protocolVersion}; expected ${DAEMON_PROTOCOL_VERSION}.`, - ), - ) - return - } - if (stopping && message.op !== 'hello' && message.op !== 'stop') { - writeJsonLine( - socket, - errorResponse('daemon-stopping', 'mcpxd is stopping.'), - ) - return - } - - switch (message.op) { - case 'hello': - writeJsonLine(socket, { - ok: true, - protocolVersion: DAEMON_PROTOCOL_VERSION, - result: { - protocolVersion: DAEMON_PROTOCOL_VERSION, - version: MCPX_VERSION, - }, - } satisfies DaemonMessage) - return - case 'listTools': - writeJsonLine(socket, okResponse(await listTools(message))) - return - case 'call': - writeJsonLine(socket, callResponse(await callTool(message))) - return - case 'status': - writeJsonLine(socket, okResponse(status())) + lastActivityAt = Date.now() + if (activeCaller && isRuntimeInputFrame(message)) { + activeCaller.receiveInput(message) return - case 'evictSession': - writeJsonLine(socket, okResponse(await evictSession(message))) - return - case 'stop': - writeJsonLine(socket, okResponse({ stopping: true })) - await stopDaemon() - return - } - } catch (error) { - const messageText = error instanceof Error ? error.message : String(error) - writeJsonLine(socket, errorResponse('operation-failed', messageText)) - } -} - -async function listTools( - message: Extract, -): Promise { - return enqueue( - message.serverKey, - message.serverName, - message.server, - message.headers, - async (session) => listAllMcpTools((await ensureConnected(session)).client), - ) -} - -async function callTool( - message: Extract, -): Promise { - return enqueue( - message.serverKey, - message.serverName, - message.server, - message.headers, - async (session) => { - const buffer = createNotificationBuffer() - const bufferNotifications = message.notificationMode !== 'discard' - session.currentBuffer = bufferNotifications ? buffer : undefined - try { - const result = await callToolWithRetainedSessionFallback( - session, - message, - bufferNotifications ? buffer : undefined, - ) - const notifications = await flushNotifications(buffer) - const toolsChanged = - buffer.toolsChanged() || session.pendingToolsChanged - session.pendingToolsChanged = false - return { result, notifications, toolsChanged } - } catch (error) { - if (session.server.transport === 'http' && isUnauthorizedError(error)) { - session.lastSessionId = - session.connection?.sessionId() ?? session.lastSessionId - await closeSession(session) - session.evictCount += 1 - } - throw error - } finally { - delete session.currentBuffer } - }, - ) -} - -async function callToolWithRetainedSessionFallback( - session: ManagedSession, - message: Extract, - buffer: NotificationBuffer | undefined, -): Promise { - try { - return await callToolOnConnectedSession(session, message, buffer) - } catch (error) { - if ( - session.server.transport === 'http' && - session.lastSessionId && - isRetainedSessionRejected(error) - ) { - await closeSession(session) - delete session.lastSessionId - return callToolOnConnectedSession(session, message, buffer) - } - throw error - } -} + pending = pending.then(async () => { + if (phase === 'complete') { + writeJsonLine( + socket, + errorResponse( + 'connection-complete', + 'A Runtime connection accepts exactly one operation.', + ), + ) + return + } + if (phase === 'awaiting-hello') { + if (!isHelloMessage(message)) { + writeJsonLine( + socket, + errorResponse( + 'handshake-required', + 'A Runtime connection must begin with a handshake.', + ), + ) + return + } + if (message.protocolVersion !== DAEMON_PROTOCOL_VERSION) { + phase = 'complete' + writeJsonLine( + socket, + errorResponse( + 'protocol-mismatch', + `Unsupported mcpxd protocol ${message.protocolVersion}; expected ${DAEMON_PROTOCOL_VERSION}.`, + ), + ) + return + } + phase = 'ready' + writeJsonLine(socket, { + ok: true, + protocolVersion: DAEMON_PROTOCOL_VERSION, + result: { + protocolVersion: DAEMON_PROTOCOL_VERSION, + version: MCPX_VERSION, + }, + } satisfies DaemonMessage) + return + } -async function callToolOnConnectedSession( - session: ManagedSession, - message: Extract, - buffer: NotificationBuffer | undefined, -): Promise { - const connection = await ensureConnected(session) - const options = toolCallRequestOptions() - if (buffer) { - options.onprogress = (progress) => { - buffer.add({ - method: 'notifications/progress', - params: { progressToken: message.callId, ...progress }, + phase = 'complete' + const intent = parseRuntimeIntent(message) + if ('code' in intent) { + writeJsonLine(socket, errorResponse(intent.code, intent.message)) + return + } + if (stopping && intent.op !== 'stop') { + writeJsonLine( + socket, + errorResponse('operation-failed', 'MCP Runtime is stopping.'), + ) + return + } + if (intent.op === 'stop') stopping = true + + const caller = createSocketRuntimeCaller(intent.requestId, socket) + activeCaller = caller + try { + const runtime = await getRuntime() + await runtime.handle(intent, caller) + if (intent.op === 'stop') await stopRuntime(server) + } catch (error) { + await caller + .send({ + requestId: intent.requestId, + kind: 'error', + error: { + code: 'operation-failed', + message: error instanceof Error ? error.message : String(error), + }, + }) + .catch(() => {}) + } }) - } - } - return connection.client.callTool( - { - name: message.toolName, - arguments: message.input, }, - undefined, - options, + (error) => { + writeJsonLine(socket, errorResponse('invalid-json', error.message)) + }, ) } -async function enqueue( - serverKey: string, - serverName: string, - serverConfig: ServerConfig, - headers: Record | undefined, - run: (session: ManagedSession) => Promise, -): Promise { - const session = getSession(serverKey, serverName, serverConfig, headers) - if (headers) { - session.headers = headers - session.connection?.updateHeaders(headers) - } - session.queuedCalls += 1 - - const task = session.queue.then(async () => { - session.queuedCalls -= 1 - session.activeCalls += 1 - try { - return await run(session) - } finally { - session.activeCalls -= 1 - session.lastUsedAt = Date.now() - lastDaemonActivity = Date.now() - } - }) - - session.queue = task.catch(() => {}) - return task -} - -function getSession( - serverKey: string, - serverName: string, - serverConfig: ServerConfig, - headers: Record | undefined, -): ManagedSession { - const existing = sessions.get(serverKey) - if (existing) { - existing.labels.add(serverName) - existing.server = serverConfig - if (headers) existing.headers = headers - return existing - } - - const session: ManagedSession = { - serverKey, - labels: new Set([serverName]), - server: serverConfig, - headers, - queue: Promise.resolve(), - activeCalls: 0, - queuedCalls: 0, - lastUsedAt: Date.now(), - evictCount: 0, - pendingToolsChanged: false, - } - sessions.set(serverKey, session) - return session -} - -async function ensureConnected( - session: ManagedSession, -): Promise { - if (session.connection) return session.connection - if (!session.connecting) { - session.connecting = connectMcpClient(session.server, { - headers: session.headers, - sessionId: session.lastSessionId, - onNotification: (notification) => { - recordNotification(session, notification) - }, - }) - .catch(async (error) => { - if (session.lastSessionId && isRetainedSessionRejected(error)) { - delete session.lastSessionId - return connectMcpClient(session.server, { - headers: session.headers, - onNotification: (notification) => { - recordNotification(session, notification) - }, - }) - } - throw error - }) - .then((connection) => { - session.connection = connection - attachStderrLog(session, connection) - void logDaemon( - `started server=${session.serverKey} pid=${connection.pid() ?? 'unknown'}`, - ) - return connection - }) - .finally(() => { - delete session.connecting - }) - } - return session.connecting -} - -async function evictSession( - message: Extract, -): Promise> { - const session = sessions.get(message.serverKey) - if (!session) return { evicted: false } - if (session.server.transport === 'stdio') - return { evicted: false, reason: 'stdio' } - - const evictTask = session.queue.then(async () => { - // Real HTTP servers can bind session ids to the old bearer token. - await closeSession(session, { - retainHttpSessionId: message.reason !== 'auth-refreshed', - }) - if (message.reason === 'auth-refreshed') delete session.lastSessionId - session.evictCount += 1 - return { evicted: true } - }) - session.queue = evictTask.catch(() => {}) - - const timeout = new Promise<{ evicted: false; timedOut: true }>((resolve) => { - setTimeout( - () => resolve({ evicted: false, timedOut: true }), - EVICT_DEADLINE_MS, - ).unref() - }) - return Promise.race([evictTask, timeout]) -} - -function recordNotification( - session: ManagedSession, - notification: { method: string; params?: unknown }, -): void { - const normalized = normalizeNotification(notification) - if (session.currentBuffer) { - session.currentBuffer.add(normalized) - return +async function cleanupIdleRuntime(server: net.Server): Promise { + if (stopping || !runtimePromise) return + const runtime = await runtimePromise + await runtime.cleanupIdleSessions(SESSION_IDLE_MS) + if ( + runtime.activeSessionCount() === 0 && + runtime.activeAuthenticationFlows() === 0 && + Date.now() - lastActivityAt >= DAEMON_IDLE_MS + ) { + await stopRuntime(server) } - if (normalized.method === 'notifications/tools/list_changed') { - session.pendingToolsChanged = true - } -} - -function normalizeNotification(notification: { - method: string - params?: unknown -}): McpNotification { - const normalized: McpNotification = { method: notification.method } - if ('params' in notification) normalized.params = notification.params - return normalized -} - -async function flushNotifications( - buffer: NotificationBuffer, -): Promise { - const notifications = buffer.flush() - if (!buffer.isOversize()) return notifications - - const json = JSON.stringify(notifications, null, 2) - const hash = createHash('sha256').update(json).digest('hex').slice(0, 16) - const filePath = path.join(tmpdir(), `mcpx-notifications-${hash}.json`) - // Stable filenames make repeated oversize notification payloads dedupe naturally. - await fs.writeFile(filePath, `${json}\n`, 'utf8') - return [{ method: '$oversize', params: { savedTo: filePath } }] -} - -function attachStderrLog( - session: ManagedSession, - connection: ConnectedSession, -): void { - const stderr = connection.stderr - if (!stderr) return - - stderr.on('data', (chunk) => { - void appendLog(serverLogPath(session.serverKey), chunk.toString()) - }) } -async function cleanupIdleSessions(server: net.Server): Promise { - if (stopping) return - - const now = Date.now() - for (const session of sessions.values()) { - if (session.activeCalls > 0 || session.queuedCalls > 0) continue - if (now - session.lastUsedAt < CHILD_IDLE_TTL_MS) continue - await closeSession(session) - sessions.delete(session.serverKey) - } - - if (sessions.size === 0 && now - lastDaemonActivity >= DAEMON_IDLE_TTL_MS) { - await logDaemon('mcpxd idle timeout reached') - server.close() - } +async function getRuntime(): Promise { + runtimePromise ??= openRuntimeStores(daemonDir()).then( + (stores) => new McpRuntime(stores), + ) + return runtimePromise } -async function stopDaemon(): Promise { +async function stopRuntime(server: net.Server): Promise { stopping = true - await Promise.all( - [...sessions.values()].map((session) => session.queue.catch(() => {})), - ) - await Promise.all( - [...sessions.values()].map((session) => closeSession(session)), - ) - sessions.clear() await fs.rm(daemonSocketPath(), { force: true }).catch(() => {}) + server.close() process.exitCode = 0 - setTimeout(() => process.exit(0), 10).unref() } -async function closeSession( - session: ManagedSession, - options: { retainHttpSessionId?: boolean } = {}, -): Promise { - const retainHttpSessionId = options.retainHttpSessionId ?? true - if (session.server.transport === 'http' && retainHttpSessionId) { - session.lastSessionId = - session.connection?.sessionId() ?? session.lastSessionId - } - await session.connection?.close().catch(() => {}) - delete session.connection - delete session.connecting - await logDaemon(`stopped server=${session.serverKey}`) -} - -function status(): DaemonStatus { - const now = Date.now() - return { - pid: process.pid, - protocolVersion: DAEMON_PROTOCOL_VERSION, - version: MCPX_VERSION, - activeServers: sessions.size, - servers: [...sessions.values()].map((session) => { - const item: DaemonStatus['servers'][number] = { - serverKey: session.serverKey, - transport: session.server.transport ?? 'http', - labels: [...session.labels].sort(), - activeCalls: session.activeCalls, - queuedCalls: session.queuedCalls, - idleMs: now - session.lastUsedAt, - evictCount: session.evictCount, - hasRetainedSessionId: session.lastSessionId !== undefined, - } - if (session.server.transport === 'stdio') { - item.pid = session.connection?.pid() ?? null - } else { - item.url = redactedUrl(session.server.url) - } - if (session.lastSessionId) - item.sessionIdHash = shortHash(session.lastSessionId) - return item - }), - } -} - -function okResponse(result: unknown): DaemonMessage { - return { ok: true, result } -} - -function callResponse(result: DaemonCallResult): DaemonMessage { - const response: DaemonMessage = { ok: true, result: result.result } - if (result.notifications.length > 0) - response.notifications = result.notifications - if (result.toolsChanged) response.toolsChanged = true - return response +function isHelloMessage( + value: unknown, +): value is Extract { + return ( + !!value && + typeof value === 'object' && + (value as { op?: unknown }).op === 'hello' && + typeof (value as { protocolVersion?: unknown }).protocolVersion === + 'number' && + typeof (value as { clientVersion?: unknown }).clientVersion === 'string' + ) } function errorResponse(code: string, message: string): DaemonMessage { @@ -549,72 +203,14 @@ async function listen(server: net.Server, socketPath: string): Promise { }) } -async function logDaemon(message: string): Promise { - await appendLog(daemonLogPath(), `${new Date().toISOString()} ${message}\n`) -} - -async function appendLog(filePath: string, text: string): Promise { - await rotateLogIfNeeded(filePath, Buffer.byteLength(text)).catch(() => {}) - await fs.appendFile(filePath, text, 'utf8').catch(() => {}) -} - -async function rotateLogIfNeeded( - filePath: string, - incomingBytes: number, -): Promise { - const stat = await fs.stat(filePath).catch(() => undefined) - if (!stat || stat.size + incomingBytes <= LOG_MAX_BYTES) return - - await fs.rm(`${filePath}.2`, { force: true }).catch(() => {}) - await fs.rename(`${filePath}.1`, `${filePath}.2`).catch(() => {}) - await fs.rename(filePath, `${filePath}.1`).catch(() => {}) -} - -function isClientMessage(value: unknown): value is ClientMessage { - if (!value || typeof value !== 'object') return false - const op = (value as { op?: unknown }).op - return ( - op === 'hello' || - op === 'listTools' || - op === 'call' || - op === 'status' || - op === 'stop' || - op === 'evictSession' - ) -} - -function redactedUrl(value: string): string { - const url = new URL(value) - return `${url.host}${url.pathname}` -} - -function shortHash(value: string): string { - return createHash('sha256').update(value).digest('hex').slice(0, 16) -} - -function isRetainedSessionRejected(error: unknown): boolean { - const message = error instanceof Error ? error.message : String(error) - const normalized = message.toLowerCase() - return ( - message.includes('404') || - normalized.includes('bad session') || - // PostHog returns this when a retained session id outlives its access token. - normalized.includes('invalid api key') || - normalized.includes('invalid_api_key') - ) -} - -function isUnauthorizedError(error: unknown): boolean { - const message = error instanceof Error ? error.message : String(error) - return ( - message.includes('401') || message.toLowerCase().includes('unauthorized') - ) -} - async function isLiveSocket(socketPath: string): Promise { - const socket = await connectSocket(socketPath).catch(() => undefined) - if (!socket) return false + let socket: Socket | undefined try { + socket = await new Promise((resolve, reject) => { + const candidate = net.createConnection(socketPath) + candidate.once('connect', () => resolve(candidate)) + candidate.once('error', reject) + }) const parsed = (await requestJsonLine(socket, { op: 'hello', protocolVersion: DAEMON_PROTOCOL_VERSION, @@ -626,14 +222,6 @@ async function isLiveSocket(socketPath: string): Promise { } catch { return false } finally { - socket.destroy() + socket?.destroy() } } - -async function connectSocket(socketPath: string): Promise { - return new Promise((resolve, reject) => { - const socket = net.createConnection(socketPath) - socket.once('connect', () => resolve(socket)) - socket.once('error', reject) - }) -} diff --git a/src/discovery.ts b/src/discovery.ts index da72ffe..2a4745a 100644 --- a/src/discovery.ts +++ b/src/discovery.ts @@ -11,7 +11,6 @@ import { authFromBearerValues, describeBearerAuth } from './bearer' import { resolveProbeHeaders } from './headers' import { listMcpTools } from './mcp-client' import { assignCommandNames } from './names' -import { authenticateOAuthServer } from './oauth' export type DiscoverServerOptions = { name: string @@ -21,7 +20,6 @@ export type DiscoverServerOptions = { url: string bearer?: string | string[] headers?: Record - interactiveAuth?: boolean } | { transport: 'stdio' @@ -47,26 +45,9 @@ export async function discoverServer( const discoveredAuth = configuredAuth ?? (await discoverAuth(url, resolveProbeHeaders(seedServer))) - let auth = discoveredAuth - try { - if (options.interactiveAuth !== false && discoveredAuth.kind === 'oauth') { - auth = await authenticateOAuthServer(options.name, url, discoveredAuth) - } - } catch (error) { - const message = error instanceof Error ? error.message : String(error) - return { - server: { - ...seedServer, - auth: discoveredAuth, - discoveredAt: new Date().toISOString(), - }, - status: 'auth-required', - message, - } - } const server: ServerConfig = { ...seedServer, - auth, + auth: discoveredAuth, discoveredAt: new Date().toISOString(), } @@ -80,7 +61,7 @@ export async function discoverServer( } } catch (error) { const message = error instanceof Error ? error.message : String(error) - if (auth.kind === 'oauth' || auth.kind === 'unknown') { + if (discoveredAuth.kind === 'oauth' || discoveredAuth.kind === 'unknown') { return { server, status: 'auth-required', @@ -120,46 +101,6 @@ async function discoverStdioServer( } } -export async function refreshServer( - server: ServerConfig, - name = 'stdio', -): Promise { - const tools = await listMcpTools(server, name) - return { - ...server, - discoveredAt: new Date().toISOString(), - tools: normalizeTools(tools), - } -} - -export async function reauthenticateServer( - name: string, - server: ServerConfig, -): Promise { - if (server.transport === 'stdio') { - throw new Error( - `Server "${name}" uses stdio and does not support OAuth re-authentication.`, - ) - } - - const url = new URL(server.url) - const discoveredAuth = await discoverAuth( - url, - resolveProbeHeaders({ ...server, auth: { kind: 'none' } }), - ) - if (discoveredAuth.kind !== 'oauth') { - throw new Error( - `Server "${name}" did not advertise OAuth authentication metadata.`, - ) - } - - const auth = await authenticateOAuthServer(name, url, discoveredAuth) - return { - ...server, - auth, - } -} - export function normalizeTools(tools: McpTool[]): ToolDefinition[] { const names = tools.map((tool) => tool.name) const commandNames = assignCommandNames(names) diff --git a/src/headers.ts b/src/headers.ts index 192c9bf..c57eb5a 100644 --- a/src/headers.ts +++ b/src/headers.ts @@ -1,8 +1,6 @@ import type { HttpServerConfig } from './types' import { resolveBearerHeader, resolveBearerHeaderForProbe } from './bearer' -import { refreshOAuthToken, shouldRefreshOAuthToken } from './oauth' -import { getOAuthTokenForUpdate, putOAuthTokenInCache } from './token-cache' export type ResolvedHeaders = { headers: Record @@ -31,22 +29,12 @@ export async function resolveHeadersWithState( server: HttpServerConfig, ): Promise { const headers = baseHeaders(server) - let authRefreshed = false if (server.auth.kind === 'bearer') { headers.Authorization = await resolveBearerHeader(server.url, server.auth) } - if (server.auth.kind === 'oauth-token') { - const result = await resolveOAuthToken(server) - const token = result?.token - authRefreshed = result?.refreshed ?? false - if (token) { - headers.Authorization = `${normalizeAuthScheme(token.tokenType)} ${token.accessToken}` - } - } - - return { headers, authRefreshed } + return { headers, authRefreshed: false } } function baseHeaders(server: HttpServerConfig): Record { @@ -56,30 +44,6 @@ function baseHeaders(server: HttpServerConfig): Record { } } -async function resolveOAuthToken(server: HttpServerConfig) { - if (server.auth.kind !== 'oauth-token') return undefined - - const { cache, token } = await getOAuthTokenForUpdate(server.auth.tokenKey) - if (!token || !shouldRefreshOAuthToken(token)) - return { token, refreshed: false } - - const refreshed = await refreshOAuthToken({ - issuer: issuerFromTokenKey(server.auth.tokenKey), - resourceUrl: server.url, - token, - }) - await putOAuthTokenInCache(cache, server.auth.tokenKey, refreshed) - return { token: refreshed, refreshed: true } -} - -function issuerFromTokenKey(tokenKey: string): string { - const separator = tokenKey.indexOf(':') - if (separator === -1) { - throw new Error(`Invalid OAuth token key: ${tokenKey}`) - } - return tokenKey.slice(separator + 1) -} - export function normalizeAuthScheme(tokenType: string): string { const normalized = tokenType.toLowerCase() return normalized === 'bearer' || diff --git a/src/main.ts b/src/main.ts index d951fc1..64ce66d 100755 --- a/src/main.ts +++ b/src/main.ts @@ -1,12 +1,10 @@ #!/usr/bin/env bun +import { runDaemonServer } from './daemon-server' import { runMcpx } from './router' -import { - runSchemaRefreshWorker, - shouldRunSchemaRefreshWorker, -} from './schema-refresh' -if (shouldRunSchemaRefreshWorker()) { - await runSchemaRefreshWorker() +const argv = process.argv.slice(2) +if (argv[0] === '@daemon' && argv[1] === 'server') { + await runDaemonServer() } else { - await runMcpx(process.argv.slice(2), process.cwd(), import.meta.path) + await runMcpx(argv, process.cwd(), import.meta.path) } diff --git a/src/mcp-client.ts b/src/mcp-client.ts index 46d6ec6..39c9061 100644 --- a/src/mcp-client.ts +++ b/src/mcp-client.ts @@ -7,8 +7,6 @@ import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/ import type { McpTool, ServerConfig, ToolAnnotations } from './types' -import { callToolViaDaemon, listToolsViaDaemon } from './daemon-client' -import { shouldUseDaemon } from './daemon-protocol' import { resolveHeaders } from './headers' import { MCPX_VERSION } from './version' @@ -136,11 +134,8 @@ function stdioTransportParams(server: ServerConfig) { export async function listMcpTools( server: ServerConfig, - serverName = 'stdio', + _serverName = 'stdio', ): Promise { - if (shouldUseDaemon()) { - return listToolsViaDaemon(server, serverName) - } return withMcpClient(server, async (client) => listAllMcpTools(client)) } @@ -175,11 +170,8 @@ export async function callMcpTool( server: ServerConfig, toolName: string, input: Record, - serverName = 'stdio', + _serverName = 'stdio', ): Promise { - if (shouldUseDaemon()) { - return callToolViaDaemon(server, serverName, toolName, input) - } return withMcpClient(server, async (client) => client.callTool( { name: toolName, arguments: input }, diff --git a/src/notifications.ts b/src/notifications.ts index 593ba67..13dd70f 100644 --- a/src/notifications.ts +++ b/src/notifications.ts @@ -71,6 +71,20 @@ export function createNotificationBuffer(): NotificationBuffer { } } +export async function flushNotificationBuffer( + buffer: NotificationBuffer, +): Promise { + const notifications = buffer.flush() + if (!buffer.isOversize()) return notifications + + const json = JSON.stringify(notifications, null, 2) + const hash = createHash('sha256').update(json).digest('hex').slice(0, 16) + const filePath = path.join(tmpdir(), `mcpx-notifications-${hash}.json`) + // Stable filenames make repeated oversize notification payloads dedupe naturally. + await fs.writeFile(filePath, `${json}\n`, 'utf8') + return [{ method: '$oversize', params: { savedTo: filePath } }] +} + function isProgressNotification( notification: McpNotification, ): notification is Extract< @@ -82,3 +96,7 @@ function isProgressNotification( typeof notification.params === 'object' ) } +import { createHash } from 'node:crypto' +import fs from 'node:fs/promises' +import { tmpdir } from 'node:os' +import path from 'node:path' diff --git a/src/oauth.ts b/src/oauth.ts index 51b336c..39ad2f8 100644 --- a/src/oauth.ts +++ b/src/oauth.ts @@ -1,19 +1,24 @@ -import { cancel, confirm, isCancel, note, password, text } from '@clack/prompts' import { createHash, randomBytes } from 'node:crypto' import type { AuthDiscovery, OAuthServerMetadata, OAuthToken } from './types' -import { - getOAuthClientSecret, - putOAuthTokenWithClientSecret, -} from './token-cache' - type OAuthClientRegistration = { clientId: string clientSecret?: string clientSecretKey?: string } +export type ManualOAuthClientRequest = { + serverName: string + redirectUri: string + issuer: string + scopes: string[] +} + +export type ManualOAuthClientProvider = ( + request: ManualOAuthClientRequest, +) => Promise + type OAuthCallbackResult = { code: string state: string @@ -34,11 +39,17 @@ const CALLBACK_TIMEOUT_MS = 5 * 60 * 1000 // Manual OAuth clients require a pre-registered redirect URI; random ports break that contract. const MANUAL_CALLBACK_PORT = 65245 -export async function authenticateOAuthServer( +export async function performOAuthAuthentication( serverName: string, resourceUrl: URL, auth: DiscoveredOAuth, -): Promise { + signal?: AbortSignal, + manualClientProvider?: ManualOAuthClientProvider, +): Promise<{ + auth: AuthenticatedOAuth + token: OAuthToken + clientSecret?: string +}> { const authorizationServer = auth.authorizationServers?.[0] if (!authorizationServer) { throw new Error( @@ -46,22 +57,39 @@ export async function authenticateOAuthServer( ) } - const metadata = await fetchAuthorizationServerMetadata(authorizationServer) + const metadata = await fetchAuthorizationServerMetadata(authorizationServer, { + signal, + }) + if ( + !metadata.registrationEndpoint && + metadata.tokenEndpointAuthMethodsSupported && + !metadata.tokenEndpointAuthMethodsSupported.includes('client_secret_post') + ) { + throw new Error( + 'OAuth server requires a client authentication method mcpx cannot use yet.', + ) + } const verifier = base64Url(randomBytes(32)) const challenge = base64Url(createHash('sha256').update(verifier).digest()) const state = base64Url(randomBytes(24)) const manualClient = metadata.registrationEndpoint ? undefined - : await promptForManualOAuthClient(serverName, metadata, auth) + : await requireManualClientProvider(manualClientProvider)({ + serverName, + redirectUri: `http://${LOCALHOST}:${MANUAL_CALLBACK_PORT}${CALLBACK_PATH}`, + issuer: metadata.issuer, + scopes: auth.scopesSupported ?? [], + }) const callback = await waitForOAuthCallback( state, manualClient ? MANUAL_CALLBACK_PORT : 0, + signal, ) try { const client = manualClient ?? - (await registerOAuthClient(metadata, callback.redirectUri)) + (await registerOAuthClient(metadata, callback.redirectUri, { signal })) const scope = chooseOAuthScope( auth.scopesSupported, metadata.scopesSupported, @@ -90,6 +118,7 @@ export async function authenticateOAuthServer( resourceUrl: string code: string verifier: string + signal?: AbortSignal } = { metadata, clientId: client.clientId, @@ -99,11 +128,20 @@ export async function authenticateOAuthServer( verifier, } if (client.clientSecret) exchangeOptions.clientSecret = client.clientSecret + if (signal) exchangeOptions.signal = signal const token = await exchangeAuthorizationCode(exchangeOptions) const tokenKey = `${serverName}:${metadata.issuer}` - await putOAuthTokenWithClientSecret(tokenKey, token, client.clientSecret) - return { kind: 'oauth-token', tokenKey, confidence: 'confirmed' } + const completed: { + auth: AuthenticatedOAuth + token: OAuthToken + clientSecret?: string + } = { + auth: { kind: 'oauth-token', tokenKey, confidence: 'confirmed' }, + token, + } + if (client.clientSecret) completed.clientSecret = client.clientSecret + return completed } finally { callback.close() } @@ -111,6 +149,7 @@ export async function authenticateOAuthServer( export async function fetchAuthorizationServerMetadata( issuer: string, + options: { signal?: AbortSignal } = {}, ): Promise { const metadataUrls = authorizationServerMetadataUrls(issuer) const failures: string[] = [] @@ -118,6 +157,7 @@ export async function fetchAuthorizationServerMetadata( for (const metadataUrl of metadataUrls) { const response = await fetch(metadataUrl, { headers: { Accept: 'application/json' }, + signal: options.signal, }) if (!response.ok) { failures.push(`${metadataUrl} (${response.status})`) @@ -191,6 +231,7 @@ function parseAuthorizationServerMetadata( export async function registerOAuthClient( metadata: OAuthServerMetadata, redirectUri: string, + options: { signal?: AbortSignal } = {}, ): Promise { if (!metadata.registrationEndpoint) { throw new Error( @@ -209,6 +250,7 @@ export async function registerOAuthClient( token_endpoint_auth_method: 'none', application_type: 'native', }), + signal: options.signal, }) if (!response.ok) { @@ -252,6 +294,8 @@ export async function refreshOAuthToken(options: { issuer: string resourceUrl: string token: OAuthToken + clientSecret?: string + signal?: AbortSignal }): Promise { if (!options.token.refreshToken) { throw new Error( @@ -263,16 +307,16 @@ export async function refreshOAuthToken(options: { 'OAuth token is expired and cannot be refreshed because it was created by an older mcpx version. Run mcpx @add again.', ) } - const metadata = await fetchAuthorizationServerMetadata(options.issuer) + const metadata = await fetchAuthorizationServerMetadata(options.issuer, { + signal: options.signal, + }) const body = new URLSearchParams({ grant_type: 'refresh_token', client_id: options.token.clientId, refresh_token: options.token.refreshToken, resource: options.resourceUrl, }) - const clientSecret = options.token.clientSecretKey - ? await getOAuthClientSecret(options.token.clientSecretKey) - : undefined + const clientSecret = options.clientSecret if (options.token.clientSecretKey && !clientSecret) { throw new Error('OAuth client secret is missing. Run mcpx @add again.') } @@ -285,6 +329,7 @@ export async function refreshOAuthToken(options: { Accept: 'application/json', }, body, + signal: options.signal, }) if (!response.ok) { @@ -333,6 +378,7 @@ async function exchangeAuthorizationCode(options: { resourceUrl: string code: string verifier: string + signal?: AbortSignal }): Promise { const body = new URLSearchParams({ grant_type: 'authorization_code', @@ -351,6 +397,7 @@ async function exchangeAuthorizationCode(options: { Accept: 'application/json', }, body, + signal: options.signal, }) if (!response.ok) { @@ -402,79 +449,15 @@ export function parseOAuthTokenResponse( return token } -async function promptForManualOAuthClient( - serverName: string, - metadata: OAuthServerMetadata, - auth: DiscoveredOAuth, -): Promise { - if (!process.stdin.isTTY) { - throw new Error( - 'Manual OAuth client authentication requires an interactive terminal.', - ) - } - if ( - metadata.tokenEndpointAuthMethodsSupported && - !metadata.tokenEndpointAuthMethodsSupported.includes('client_secret_post') - ) { - throw new Error( - 'OAuth server requires a client authentication method mcpx cannot use yet.', - ) - } - - const redirectUri = `http://${LOCALHOST}:${MANUAL_CALLBACK_PORT}${CALLBACK_PATH}` - note( - [ - 'This OAuth server does not support dynamic client registration.', - 'Before continuing, open the provider app settings, add this exact Redirect URL, and save it.', - `Redirect URL: ${redirectUri}`, - 'For Slack, this is under "OAuth & Permissions" -> "Redirect URLs".', - `Authorization server: ${metadata.issuer}`, - auth.scopesSupported?.length - ? `Requested scopes: ${auth.scopesSupported.join(', ')}` - : '', - ] - .filter(Boolean) - .join('\n'), - `${serverName} OAuth client`, - ) - - const redirectConfigured = await confirm({ - message: - 'I have already added and saved this exact Redirect URL in the provider app', - initialValue: false, - }) - if (isCancel(redirectConfigured) || !redirectConfigured) { - cancel('OAuth authentication cancelled.') +function requireManualClientProvider( + provider: ManualOAuthClientProvider | undefined, +): ManualOAuthClientProvider { + if (!provider) { throw new Error( - `Add ${redirectUri} as a redirect URL, then run mcpx @add again.`, + 'Manual OAuth client authentication requires Caller Input from mcpx @refresh.', ) } - - const clientId = await text({ - message: 'OAuth client_id', - validate: (value) => - value && value.trim() ? undefined : 'client_id is required.', - }) - if (isCancel(clientId)) { - cancel('OAuth authentication cancelled.') - throw new Error('OAuth authentication cancelled.') - } - - const clientSecret = await password({ - message: 'OAuth client_secret', - validate: (value) => - value && value.trim() ? undefined : 'client_secret is required.', - }) - if (isCancel(clientSecret)) { - cancel('OAuth authentication cancelled.') - throw new Error('OAuth authentication cancelled.') - } - - return { - clientId: clientId.trim(), - clientSecret: clientSecret.trim(), - clientSecretKey: clientSecretKey(clientId.trim()), - } + return provider } function parseSuccessfulTokenPayload( @@ -510,11 +493,21 @@ function clientSecretKey(clientId: string): string { function waitForOAuthCallback( expectedState: string, port: number, + signal?: AbortSignal, ): OAuthCallbackServer { let server: Bun.Server | undefined let timer: ReturnType | undefined let settled = false + let closeWithError = () => {} const result = new Promise((finish, fail) => { + const onAbort = () => { + complete( + 'fail', + signal?.reason instanceof Error + ? signal.reason + : new Error('OAuth authentication cancelled.'), + ) + } const complete = ( outcome: 'finish' | 'fail', value: OAuthCallbackResult | Error, @@ -522,6 +515,7 @@ function waitForOAuthCallback( if (settled) return settled = true if (timer) clearTimeout(timer) + signal?.removeEventListener('abort', onAbort) server?.stop() if (outcome === 'finish') { finish(value as OAuthCallbackResult) @@ -529,10 +523,17 @@ function waitForOAuthCallback( } fail(value) } + closeWithError = () => + complete('fail', new Error('OAuth callback server closed.')) timer = setTimeout(() => { complete('fail', new Error('OAuth authentication timed out.')) }, CALLBACK_TIMEOUT_MS) + if (signal?.aborted) { + onAbort() + return + } + signal?.addEventListener('abort', onAbort, { once: true }) server = Bun.serve({ hostname: LOCALHOST, @@ -563,19 +564,21 @@ function waitForOAuthCallback( }, }) }) + // Registration can fail before anyone awaits the callback; keep close rejection handled. + void result.catch(() => {}) if (!server) { throw new Error('Failed to start local OAuth callback server.') } + // The timeout owns the authentication lifetime; a stopped Bun server must not + // strand the CLI if its native socket cleanup lags behind the rejected promise. + server.unref() return { redirectUri: `http://${LOCALHOST}:${server.port}${CALLBACK_PATH}`, result, close: () => { - if (settled) return - settled = true - if (timer) clearTimeout(timer) - server?.stop() + closeWithError() }, } } diff --git a/src/project-service.ts b/src/project-service.ts deleted file mode 100644 index 2236556..0000000 --- a/src/project-service.ts +++ /dev/null @@ -1,47 +0,0 @@ -import type { RegistryConfig, ServerConfig } from './types' - -import { readRegistryConfig, writeRegistryConfig } from './config' -import { reauthenticateServer, refreshServer } from './discovery' - -export type ProjectService = { - config: RegistryConfig - ensureServerReady: (name: string) => Promise - reauthenticateServer: (name: string) => Promise - save: () => Promise -} - -export async function loadProjectService(): Promise { - const config = await readRegistryConfig() - - return { - config, - ensureServerReady: async (name: string) => { - const server = config.servers[name] - if (!server) { - throw new Error( - `Unknown MCP server "${name}". Run "mcpx @add --name ${name} --url " first.`, - ) - } - if (server.tools && server.tools.length > 0) { - return server - } - const refreshed = await refreshServer(server, name) - config.servers[name] = refreshed - await writeRegistryConfig(config) - return refreshed - }, - reauthenticateServer: async (name: string) => { - const server = config.servers[name] - if (!server) { - throw new Error( - `Unknown MCP server "${name}". Run "mcpx @add --name ${name} --url " first.`, - ) - } - const refreshed = await reauthenticateServer(name, server) - config.servers[name] = refreshed - await writeRegistryConfig(config) - return refreshed - }, - save: () => writeRegistryConfig(config), - } -} diff --git a/src/refresh-progress.ts b/src/refresh-progress.ts deleted file mode 100644 index 4cb1e8b..0000000 --- a/src/refresh-progress.ts +++ /dev/null @@ -1,238 +0,0 @@ -import type { - RefreshProgressEvent, - ServerRefreshResult, -} from './schema-refresh' - -const ESC = String.fromCharCode(27) -const ANSI_HIDE_CURSOR = ESC + '[?25l' -const ANSI_SHOW_CURSOR = ESC + '[?25h' -const ANSI_ERASE_LINE = ESC + '[2K' -const ANSI_CR = '\r' -const ANSI_DIM = ESC + '[2m' -const ANSI_RESET = ESC + '[0m' -const ANSI_GREEN = ESC + '[32m' -const ANSI_YELLOW = ESC + '[33m' -const ANSI_RED = ESC + '[31m' -const ANSI_CYAN = ESC + '[36m' - -const SPINNER_FRAMES = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'] - -export type ProgressReporter = { - handle: (event: RefreshProgressEvent) => void - dispose: () => void -} - -export function createRefreshProgressReporter(): ProgressReporter { - const stream = process.stderr - const isTTY = Boolean((stream as NodeJS.WriteStream).isTTY) - - if (!isTTY) { - return createPlainReporter(stream) - } - return createTtyReporter(stream) -} - -function createPlainReporter(stream: NodeJS.WriteStream): ProgressReporter { - const write = (line: string) => { - stream.write(line + '\n') - } - return { - handle(event) { - switch (event.type) { - case 'start': - write(`Refreshing ${event.total} server(s)...`) - break - case 'server-done': - write( - ` [${event.completed}/${event.total}] ${event.name}: ${describeResultStatus(event.result)}`, - ) - break - case 'reauth-start': - write(`Re-authenticating ${event.name} (${event.remaining} left)...`) - break - case 'reauth-done': - write(` ${event.name}: ${describeResultStatus(event.result)}`) - break - case 'complete': - write('Refresh complete.') - break - } - }, - dispose() {}, - } -} - -function createTtyReporter(stream: NodeJS.WriteStream): ProgressReporter { - let total = 0 - let completed = 0 - let active: string[] = [] - let phase: 'idle' | 'refresh' | 'reauth' = 'idle' - let reauthName: string | undefined - let frameIndex = 0 - let interval: NodeJS.Timeout | undefined - let disposed = false - - const cols = () => stream.columns || 80 - - const writeStatusLine = () => { - if (disposed) return - const frame = SPINNER_FRAMES[frameIndex % SPINNER_FRAMES.length] - frameIndex += 1 - let line: string - if (phase === 'reauth') { - line = `${ANSI_CYAN}${frame}${ANSI_RESET} Re-authenticating ${ANSI_CYAN}${reauthName ?? ''}${ANSI_RESET}` - } else if (phase === 'refresh') { - const counter = `${ANSI_DIM}[${completed}/${total}]${ANSI_RESET}` - const list = active.length > 0 ? active.join(', ') : 'finishing...' - const max = Math.max(20, cols() - 12) - const trimmed = list.length > max ? list.slice(0, max - 1) + '…' : list - line = `${ANSI_CYAN}${frame}${ANSI_RESET} ${counter} ${trimmed}` - } else { - return - } - stream.write(ANSI_CR + ANSI_ERASE_LINE + line) - } - - const clearLine = () => { - stream.write(ANSI_CR + ANSI_ERASE_LINE) - } - - const startSpinner = () => { - if (interval) return - stream.write(ANSI_HIDE_CURSOR) - interval = setInterval(writeStatusLine, 80) - writeStatusLine() - } - - const stopSpinner = () => { - if (interval) { - clearInterval(interval) - interval = undefined - } - clearLine() - } - - const writeLine = (line: string) => { - clearLine() - stream.write(line + '\n') - writeStatusLine() - } - - return { - handle(event) { - switch (event.type) { - case 'start': { - total = event.total - completed = 0 - active = [] - phase = 'refresh' - stream.write( - `${ANSI_DIM}↻${ANSI_RESET} Refreshing ${ANSI_CYAN}${total}${ANSI_RESET} server(s)...\n`, - ) - if (total === 0) { - phase = 'idle' - return - } - startSpinner() - break - } - case 'server-start': { - active = event.active - break - } - case 'server-done': { - completed = event.completed - active = event.active - writeLine( - ` ${statusIcon(event.result)} ${event.name} ${ANSI_DIM}${describeResultStatus(event.result)}${ANSI_RESET}`, - ) - break - } - case 'reauth-start': { - phase = 'reauth' - reauthName = event.name - writeLine( - `${ANSI_YELLOW}!${ANSI_RESET} Re-auth required for ${ANSI_CYAN}${event.name}${ANSI_RESET} (${event.remaining}/${event.total})`, - ) - break - } - case 'reauth-done': { - reauthName = undefined - writeLine( - ` ${statusIcon(event.result)} ${event.name} ${ANSI_DIM}${describeResultStatus(event.result)}${ANSI_RESET}`, - ) - break - } - case 'complete': { - phase = 'idle' - stopSpinner() - const s = event.summary - const parts: string[] = [] - if (s.refreshed.length > 0) - parts.push( - `${ANSI_GREEN}${s.refreshed.length} refreshed${ANSI_RESET}`, - ) - if (s.unchanged.length > 0) - parts.push( - `${ANSI_DIM}${s.unchanged.length} unchanged${ANSI_RESET}`, - ) - if (s.authRefreshed.length > 0) - parts.push( - `${ANSI_CYAN}${s.authRefreshed.length} auth-refreshed${ANSI_RESET}`, - ) - if (s.reauthenticated.length > 0) - parts.push( - `${ANSI_CYAN}${s.reauthenticated.length} reauthenticated${ANSI_RESET}`, - ) - if (s.reauthRequired.length > 0) - parts.push( - `${ANSI_YELLOW}${s.reauthRequired.length} reauth-required${ANSI_RESET}`, - ) - if (s.unreachable.length > 0) - parts.push( - `${ANSI_RED}${s.unreachable.length} unreachable${ANSI_RESET}`, - ) - stream.write( - `${ANSI_GREEN}✓${ANSI_RESET} Done${parts.length ? ': ' + parts.join(', ') : ''}\n`, - ) - break - } - } - }, - dispose() { - disposed = true - stopSpinner() - stream.write(ANSI_SHOW_CURSOR) - }, - } -} - -function statusIcon(result: ServerRefreshResult): string { - switch (result.status) { - case 'schema-refreshed': - case 'auth-refreshed': - case 'reauthenticated': - return `${ANSI_GREEN}✓${ANSI_RESET}` - case 'reauth-required': - return `${ANSI_YELLOW}!${ANSI_RESET}` - case 'unreachable': - return `${ANSI_RED}✗${ANSI_RESET}` - } -} - -function describeResultStatus(result: ServerRefreshResult): string { - switch (result.status) { - case 'schema-refreshed': - return result.schemaChanged ? 'schema updated' : 'no changes' - case 'auth-refreshed': - return result.schemaChanged ? 'auth + schema updated' : 'auth refreshed' - case 'reauthenticated': - return result.schemaChanged - ? 'reauthenticated + schema updated' - : 'reauthenticated' - case 'reauth-required': - return 'reauth required' - case 'unreachable': - return result.message ? `unreachable: ${result.message}` : 'unreachable' - } -} diff --git a/src/router.ts b/src/router.ts index 61cc938..22d9e17 100644 --- a/src/router.ts +++ b/src/router.ts @@ -1,28 +1,29 @@ -import { cancel, isCancel, multiselect } from '@clack/prompts' +import { + cancel, + confirm, + isCancel, + multiselect, + note, + password, + text, +} from '@clack/prompts' import { toStandardJsonSchema } from '@valibot/to-json-schema' import { c, cli, createDefaultSchemaExplorer, group, type Router } from 'argc' import * as v from 'valibot' -import type { ServerConfig } from './types' +import type { RuntimeInputRequest, RuntimeIntent } from './runtime-protocol' +import type { RuntimeRegistrySnapshot } from './runtime-stores' +import type { RegistryView } from './skill-command' +import type { RegistryConfig } from './types' -import { removeServerConfig, upsertServerConfig } from './config' -import { daemonStatus, stopDaemon } from './daemon-client' -import { unwrapDaemonOutput } from './daemon-result' +import { notificationModeFromEnv } from './daemon-protocol' +import { daemonOutputEnvelope } from './daemon-result' import { runDaemonServer } from './daemon-server' -import { discoverServer, refreshServer } from './discovery' import { jsonSchemaToStandardSchema } from './json-schema-standard' -import { callMcpTool } from './mcp-client' import { assertServerName } from './names' import { printOutput, type McpxContext } from './output' -import { loadProjectService, type ProjectService } from './project-service' -import { createRefreshProgressReporter } from './refresh-progress' -import { - isReauthRequiredMessage, - refreshAllServers, - startSchemaRefreshWorkerIfNeeded, -} from './schema-refresh' +import { requestRuntime } from './runtime-client' import { runSkillCommand } from './skill-command' -import { removeOAuthToken } from './token-cache' import { MCPX_VERSION } from './version' const s = toStandardJsonSchema @@ -107,14 +108,15 @@ export async function runMcpx( cwd: string, mainPath: string, ): Promise { - const service = await loadProjectService() - const isRefreshCommand = argv.includes('@refresh') - if (!isRefreshCommand) { - await refreshMissingSchemas(service) - await startSchemaRefreshWorkerIfNeeded(service.config, mainPath) + const snapshot = (await requestRuntime( + { requestId: crypto.randomUUID(), op: 'registrySnapshot' }, + mainPath, + )) as RuntimeRegistrySnapshot + const registry: RegistryView = { + servers: snapshot.servers as unknown as RegistryConfig['servers'], } - const app = cli(buildRouter(service), { + const app = cli(buildRouter(registry), { name: 'mcpx', version: MCPX_VERSION, description: 'Global MCP registry and agent-facing command surface.', @@ -137,7 +139,7 @@ export async function runMcpx( } await app.run( - { handlers: buildHandlers(service, cwd) } as never, + { handlers: buildHandlers(registry, cwd, mainPath) } as never, normalizedArgv, ) } @@ -149,9 +151,9 @@ function normalizeArgv(argv: string[]): string[] | null { return [...argv.filter((arg) => arg !== '--raw'), '--raw'] } -function buildRouter(service: ProjectService): Router { +function buildRouter(registry: RegistryView): Router { return { - ...buildServerRouter(service), + ...buildServerRouter(registry), '@add': c .meta({ description: @@ -208,9 +210,9 @@ function buildRouter(service: ProjectService): Router { } } -function buildServerRouter(service: ProjectService): Record { +function buildServerRouter(registry: RegistryView): Record { const servers: Record = {} - for (const [serverName, server] of Object.entries(service.config.servers)) { + for (const [serverName, server] of Object.entries(registry.servers)) { const tools = server.tools ?? [] const children: Record = {} for (const tool of tools) { @@ -275,26 +277,30 @@ function toolAnnotationHints( } function buildHandlers( - service: ProjectService, + registry: RegistryView, cwd: string, + mainPath = process.argv[1] ?? '', ): Record { const handlers: Record = {} - for (const [serverName, server] of Object.entries(service.config.servers)) { + for (const [serverName, server] of Object.entries(registry.servers)) { const serverHandlers: Record = {} for (const tool of server.tools ?? []) { serverHandlers[tool.commandName] = async ( options: HandlerOptions>, ) => { - const readyServer = await service.ensureServerReady(serverName) - const result = await callToolWithReauthRetry( - service, - serverName, - readyServer, - tool.name, - options.input, + const result = await requestRuntime( + { + requestId: crypto.randomUUID(), + op: 'call', + serverName, + toolName: tool.name, + input: options.input, + notificationMode: notificationModeFromEnv(), + }, + mainPath, ) - await printOutput(result, options.context) + await printOutput(runtimeCallOutput(result), options.context) } } handlers[serverName] = serverHandlers @@ -303,84 +309,75 @@ function buildHandlers( handlers['@add'] = async (options: HandlerOptions) => { const input = options.input const name = assertServerName(input.name) - const result = await discoverServer(addDiscoverOptions(name, input)) - await upsertServerConfig(name, result.server) - await printOutput( - { - name, - transport: result.server.transport ?? 'http', - status: result.status, - auth: - result.server.transport === 'stdio' ? undefined : result.server.auth, - tools: result.server.tools?.length ?? 0, - message: result.message, - }, - options.context, - ) + const discovered = addDiscoverOptions(name, input) + const intent: Extract = { + requestId: crypto.randomUUID(), + op: 'addServer', + serverName: name, + transport: discovered.transport, + } + if (discovered.transport === 'stdio') { + intent.command = discovered.command + if (discovered.args) intent.args = discovered.args + if (discovered.env) intent.env = discovered.env + } else { + intent.url = discovered.url + const bearer = normalizeStringList(discovered.bearer) + if (bearer) intent.bearer = bearer + } + await printOutput(await requestRuntime(intent, mainPath), options.context) } handlers['@remove'] = async (options: HandlerOptions<{ name?: string }>) => { const rawName = options.input.name const names = rawName === undefined - ? await promptForServersToRemove(service) + ? await promptForServersToRemove(registry) : parseRemoveNames(rawName) - const removals: Array<{ - name: string - removed: boolean - tokenRemoved: boolean - }> = [] - for (const name of names) { - const removed = await removeServerConfig(name) - if (!removed) { - throw new Error(`Unknown MCP server "${name}".`) - } - const tokenRemoved = - removed.transport !== 'stdio' && removed.auth.kind === 'oauth-token' - ? await removeOAuthToken(removed.auth.tokenKey) - : false - removals.push({ name, removed: true, tokenRemoved }) - } - - // Preserve the single-server output shape so existing agent consumers keep - // working when only one name is supplied. - if (removals.length === 1) { - await printOutput(removals[0]!, options.context) - return - } - await printOutput({ removed: removals }, options.context) + await printOutput( + await requestRuntime( + { + requestId: crypto.randomUUID(), + op: 'removeServers', + serverNames: names, + }, + mainPath, + ), + options.context, + ) } handlers['@refresh'] = async ( options: HandlerOptions>, ) => { - const reporter = createRefreshProgressReporter() - const cleanup = () => reporter.dispose() - process.once('SIGINT', cleanup) - process.once('SIGTERM', cleanup) - try { - const summary = await refreshAllServers({ - onProgress: reporter.handle, - }) - await printOutput(summary, options.context) - } finally { - reporter.dispose() - process.off('SIGINT', cleanup) - process.off('SIGTERM', cleanup) - } + await printOutput( + await requestRuntime( + { requestId: crypto.randomUUID(), op: 'refreshServers' }, + mainPath, + { onInput: promptForRuntimeInput }, + ), + options.context, + ) } handlers['@daemon'] = { status: async (options: HandlerOptions>) => { await printOutput( - await daemonStatus(process.argv[1] ?? ''), + await requestRuntime( + { requestId: crypto.randomUUID(), op: 'status' }, + mainPath, + ), options.context, ) }, stop: async (options: HandlerOptions>) => { await printOutput( - await stopDaemon(process.argv[1] ?? ''), + await requestRuntime( + { requestId: crypto.randomUUID(), op: 'stop' }, + mainPath, + { start: false }, + ), options.context, ) }, @@ -392,49 +389,73 @@ function buildHandlers( handlers['@skill'] = async ( options: HandlerOptions<{ servers?: string; show?: string }>, ) => { - await runSkillCommand(service, cwd, options.input) + await runSkillCommand(registry, cwd, options.input) } return handlers } -async function callToolWithReauthRetry( - service: ProjectService, - serverName: string, - server: ServerConfig, - toolName: string, - input: Record, +async function promptForRuntimeInput( + request: RuntimeInputRequest, + signal: AbortSignal, ): Promise { - try { - const result = await callMcpTool(server, toolName, input, serverName) - await refreshToolsIfChanged(service, serverName, server, result) - return result - } catch (error) { - const message = error instanceof Error ? error.message : String(error) - if (server.transport === 'stdio' || !isReauthRequiredMessage(message)) - throw error - const reauthenticated = await service.reauthenticateServer(serverName) - const result = await callMcpTool( - reauthenticated, - toolName, - input, - serverName, - ) - await refreshToolsIfChanged(service, serverName, reauthenticated, result) - return result + if (request.type !== 'oauth-client') { + return { cancelled: true, reason: 'Unsupported Runtime input request.' } } + if (!process.stdin.isTTY) { + return { cancelled: true, reason: 'Interactive terminal required.' } + } + note( + [ + 'This OAuth server does not support dynamic client registration.', + 'Add and save this exact Redirect URL in the provider app settings.', + `Redirect URL: ${request.redirectUri}`, + `Authorization server: ${request.issuer}`, + request.scopes.length + ? `Requested scopes: ${request.scopes.join(', ')}` + : '', + ] + .filter(Boolean) + .join('\n'), + `${request.serverName} OAuth client`, + ) + const configured = await confirm({ + message: 'I have added and saved this exact Redirect URL', + initialValue: false, + signal, + }) + if (isCancel(configured) || !configured) return { cancelled: true } + const clientId = await text({ + message: 'OAuth client_id', + signal, + validate: (value) => (value?.trim() ? undefined : 'client_id is required.'), + }) + if (isCancel(clientId)) return { cancelled: true } + const clientSecret = await password({ + message: 'OAuth client_secret', + signal, + validate: (value) => + value?.trim() ? undefined : 'client_secret is required.', + }) + if (isCancel(clientSecret)) return { cancelled: true } + return { clientId: clientId.trim(), clientSecret: clientSecret.trim() } } -async function refreshToolsIfChanged( - service: ProjectService, - serverName: string, - server: ServerConfig, - result: unknown, -): Promise { - const daemonResult = unwrapDaemonOutput(result) - if (!daemonResult?.toolsChanged) return - service.config.servers[serverName] = await refreshServer(server, serverName) - await service.save() +function runtimeCallOutput(value: unknown): unknown { + if (!value || typeof value !== 'object' || !('result' in value)) return value + const response = value as { + result: unknown + notifications?: Parameters[0]['notifications'] + toolsChanged?: boolean + } + if (response.notifications?.length || response.toolsChanged) { + return daemonOutputEnvelope({ + result: response.result, + notifications: response.notifications ?? [], + toolsChanged: response.toolsChanged === true, + }) + } + return response.result } function addDiscoverOptions(name: string, input: AddServerInput) { @@ -484,25 +505,6 @@ function normalizeStringList( return Array.isArray(value) ? value : [value] } -async function refreshMissingSchemas(service: ProjectService): Promise { - let changed = false - - for (const [name, server] of Object.entries(service.config.servers)) { - if (server.tools && server.tools.length > 0) continue - try { - service.config.servers[name] = await refreshServer(server, name) - changed = true - } catch { - // Keep startup usable when auth is not available yet; add/discover records - // the auth state so the next run can retry after credentials are configured. - } - } - - if (changed) { - await service.save() - } -} - export const __test = { buildServerRouter, buildRouter, @@ -514,9 +516,9 @@ export const __test = { } async function promptForServersToRemove( - service: ProjectService, + registry: RegistryView, ): Promise { - const names = Object.keys(service.config.servers).sort() + const names = Object.keys(registry.servers).sort() if (names.length === 0) { throw new Error('No MCP servers are registered. Nothing to remove.') } @@ -531,7 +533,7 @@ async function promptForServersToRemove( message: 'Select MCP server(s) to remove (space to toggle, enter to confirm)', options: names.map((name) => { - const server = service.config.servers[name]! + const server = registry.servers[name]! const toolCount = server.tools?.length ?? 0 const transport = server.transport === 'stdio' ? 'stdio' : `http (${server.auth.kind})` diff --git a/src/runtime-authentication.ts b/src/runtime-authentication.ts new file mode 100644 index 0000000..a5e4298 --- /dev/null +++ b/src/runtime-authentication.ts @@ -0,0 +1,209 @@ +import type { RuntimeCaller } from './runtime-caller' +import type { DeclaredServer, RuntimeStores } from './runtime-stores' + +import { AuthenticationCoordinator } from './authentication-coordinator' +import { performOAuthAuthentication, refreshOAuthToken } from './oauth' +import { RuntimeOperationError } from './runtime-call' + +type RefreshResult = Awaited> +type AuthenticationResult = Awaited< + ReturnType +> + +export type RuntimeRefreshOutcome = + | { status: 'completed'; refreshed: string[] } + | { status: 'disconnected' } + +export class RuntimeAuthentication { + readonly #stores: RuntimeStores + readonly #coordinator: AuthenticationCoordinator + readonly #refreshToken: typeof refreshOAuthToken + readonly #authenticate: typeof performOAuthAuthentication + + constructor( + stores: RuntimeStores, + options: { + coordinator?: AuthenticationCoordinator + refreshToken?: typeof refreshOAuthToken + authenticate?: typeof performOAuthAuthentication + } = {}, + ) { + this.#stores = stores + this.#coordinator = options.coordinator ?? new AuthenticationCoordinator() + this.#refreshToken = options.refreshToken ?? refreshOAuthToken + this.#authenticate = options.authenticate ?? performOAuthAuthentication + } + + async close(): Promise { + await this.#coordinator.close() + } + + activeFlows(): number { + return this.#coordinator.activeFlows() + } + + async refreshServers( + serverNames: string[] | undefined, + caller: RuntimeCaller, + ): Promise { + const registry = await this.#stores.registry.read() + const names = serverNames ?? Object.keys(registry.servers).sort() + const refreshed: string[] = [] + + for (const name of names) { + const server = registry.servers[name] + if (!server) { + throw new RuntimeOperationError( + 'operation-failed', + `Unknown MCP server: ${name}.`, + ) + } + if (server.transport === 'stdio') continue + + const outcome = + server.auth.kind === 'oauth-token' + ? await this.#refreshExisting(name, server, caller) + : server.auth.kind === 'oauth' + ? await this.#authenticateDiscovered(name, server, caller) + : { status: 'completed' as const } + if (outcome.status === 'disconnected') return outcome + if (server.auth.kind === 'oauth-token' || server.auth.kind === 'oauth') { + refreshed.push(name) + } + } + + return { status: 'completed', refreshed } + } + + async #refreshExisting( + name: string, + server: Extract, + caller: RuntimeCaller, + ) { + if (server.auth.kind !== 'oauth-token') { + return { status: 'completed' as const } + } + const tokenKey = server.auth.tokenKey + const credentials = await this.#stores.credentials.read() + const token = credentials.oauth[tokenKey] + if (!token) throw reauthRequired(name) + const separator = tokenKey.indexOf(':') + if (separator === -1) throw reauthRequired(name) + const issuer = tokenKey.slice(separator + 1) + const clientSecret = token.clientSecretKey + ? credentials.oauthClientSecrets[token.clientSecretKey] + : undefined + + return this.#coordinator.join(`oauth:${tokenKey}`, caller, { + start: async (signal) => { + const current = (await this.#stores.credentials.read()).oauth[tokenKey] + // A caller can arrive after a shared flow persisted but before it observed + // completion; reuse that rotation instead of refreshing the stale token again. + if (current && current.accessToken !== token.accessToken) return current + return this.#refreshToken({ + issuer, + resourceUrl: server.url, + token, + clientSecret, + signal, + }) + }, + persist: async (refreshed) => { + await this.#stores.updateState((state) => { + state.credentials.oauth[tokenKey] = refreshed + }) + }, + }) + } + + async #authenticateDiscovered( + name: string, + server: Extract, + caller: RuntimeCaller, + ) { + if (server.auth.kind !== 'oauth') { + return { status: 'completed' as const } + } + const auth = server.auth + const authorizationServer = auth.authorizationServers?.[0] + if (!authorizationServer) throw reauthRequired(name) + const identity = `oauth:${name}:${authorizationServer}` + + return this.#coordinator.join(identity, caller, { + start: (signal, requestInput) => + this.#authenticate( + name, + new URL(server.url), + auth, + signal, + async (request) => + parseManualClient( + await requestInput({ + type: 'oauth-client', + ...request, + }), + ), + ), + persist: async (completed) => { + await this.#stores.updateState((state) => { + state.credentials.oauth[completed.auth.tokenKey] = completed.token + if (completed.clientSecret && completed.token.clientSecretKey) { + state.credentials.oauthClientSecrets[ + completed.token.clientSecretKey + ] = completed.clientSecret + } + const current = state.registry.servers[name] + if (current && current.transport !== 'stdio') { + state.registry.servers[name] = { ...current, auth: completed.auth } + } + }) + }, + }) + } +} + +function parseManualClient(value: unknown): { + clientId: string + clientSecret: string + clientSecretKey: string +} { + if (!value || typeof value !== 'object') + throw new RuntimeOperationError( + 'cancelled', + 'OAuth authentication cancelled.', + ) + const input = value as { + cancelled?: unknown + clientId?: unknown + clientSecret?: unknown + } + if (input.cancelled === true) + throw new RuntimeOperationError( + 'cancelled', + 'OAuth authentication cancelled.', + ) + if ( + typeof input.clientId !== 'string' || + !input.clientId.trim() || + typeof input.clientSecret !== 'string' || + !input.clientSecret.trim() + ) { + throw new RuntimeOperationError( + 'operation-failed', + 'OAuth client credentials were invalid.', + ) + } + const clientId = input.clientId.trim() + return { + clientId, + clientSecret: input.clientSecret.trim(), + clientSecretKey: `oauth-client:${clientId}`, + } +} + +function reauthRequired(serverName: string): RuntimeOperationError { + return new RuntimeOperationError( + 'reauth-required', + `Credentials for ${serverName} must be refreshed.`, + ) +} diff --git a/src/runtime-call.ts b/src/runtime-call.ts new file mode 100644 index 0000000..4a62bc8 --- /dev/null +++ b/src/runtime-call.ts @@ -0,0 +1,234 @@ +import type { RuntimeCaller } from './runtime-caller' +import type { RuntimeError, RuntimeFrame } from './runtime-protocol' + +export type RuntimeCallState = 'accepted' | 'queued' | 'active' | 'terminal' +export type CallCancellationCause = + | 'caller-disconnected' + | 'runtime-stopping' + | 'timeout' + +export class RuntimeOperationError extends Error { + readonly code: RuntimeError['code'] + + constructor(code: RuntimeError['code'], message: string) { + super(message) + this.name = 'RuntimeOperationError' + this.code = code + } +} + +type Executor = (signal: AbortSignal) => Promise +type QueueEntry = { + call: RuntimeCall + execute: Executor + unsubscribeTerminal: () => void +} + +export class RuntimeCall { + readonly #caller: RuntimeCaller + readonly #controller = new AbortController() + readonly #terminalListeners = new Set<() => void>() + #state: RuntimeCallState = 'accepted' + #cancellationCause: CallCancellationCause | undefined + #callerConnected = true + #unsubscribeDisconnect: () => void = () => {} + readonly #settledPromise: Promise + #resolveSettled = () => {} + + constructor(caller: RuntimeCaller) { + this.#caller = caller + this.#settledPromise = new Promise((resolve) => { + this.#resolveSettled = resolve + }) + const unsubscribe = caller.onDisconnect(() => { + this.#callerConnected = false + void this.#finish({ cause: 'caller-disconnected', abort: true }) + }) + // A disconnected adapter may notify synchronously during subscription. + if (this.#state === 'terminal') unsubscribe() + else this.#unsubscribeDisconnect = unsubscribe + } + + get id(): string { + return this.#caller.id + } + + get state(): RuntimeCallState { + return this.#state + } + + get signal(): AbortSignal { + return this.#controller.signal + } + + get cancellationCause(): CallCancellationCause | undefined { + return this.#cancellationCause + } + + get settled(): Promise { + return this.#settledPromise + } + + queue(): boolean { + if (this.#state !== 'accepted') return false + this.#state = 'queued' + return true + } + + onTerminal(listener: () => void): () => void { + this.#terminalListeners.add(listener) + return () => this.#terminalListeners.delete(listener) + } + + async execute(run: Executor): Promise { + if (this.#state !== 'queued') return + this.#state = 'active' + try { + const result = await run(this.#controller.signal) + await this.#finish({ + frame: { requestId: this.id, kind: 'result', result }, + }) + } catch (error) { + await this.#finish({ + frame: { + requestId: this.id, + kind: 'error', + error: operationError(error), + }, + }) + } + } + + async cancel(cause: CallCancellationCause): Promise { + const frame: RuntimeFrame | undefined = + cause === 'timeout' + ? { + requestId: this.id, + kind: 'error', + error: { code: 'timeout', message: 'Runtime Call timed out.' }, + } + : cause === 'runtime-stopping' + ? { + requestId: this.id, + kind: 'error', + error: { + code: 'cancelled', + message: 'MCP Runtime is stopping.', + }, + } + : undefined + await this.#finish({ cause, abort: true, frame }) + } + + async fail(error: RuntimeError): Promise { + await this.#finish({ + frame: { requestId: this.id, kind: 'error', error }, + }) + } + + async #finish(options: { + cause?: CallCancellationCause + abort?: boolean + frame?: RuntimeFrame + }): Promise { + if (this.#state === 'terminal') return + + this.#state = 'terminal' + this.#cancellationCause = options.cause + this.#unsubscribeDisconnect() + if (options.abort && !this.#controller.signal.aborted) { + this.#controller.abort(new RuntimeCallCancelled(options.cause)) + } + for (const listener of [...this.#terminalListeners]) listener() + this.#terminalListeners.clear() + this.#resolveSettled() + + if (options.frame && this.#callerConnected) { + await this.#caller.send(options.frame) + } + } +} + +export class CancelableFifo { + readonly #entries: QueueEntry[] = [] + #active: RuntimeCall | undefined + #draining: Promise | undefined + + enqueue(call: RuntimeCall, execute: Executor): void { + if (!call.queue()) return + const entry: QueueEntry = { + call, + execute, + unsubscribeTerminal: () => {}, + } + entry.unsubscribeTerminal = call.onTerminal(() => { + const index = this.#entries.indexOf(entry) + if (index !== -1) this.#entries.splice(index, 1) + }) + this.#entries.push(entry) + this.#ensureDraining() + } + + status(): { activeCalls: number; queuedCalls: number } { + return { + activeCalls: this.#active ? 1 : 0, + queuedCalls: this.#entries.length, + } + } + + async idle(): Promise { + while (this.#draining) await this.#draining + } + + async cancelAll(cause: CallCancellationCause): Promise { + const calls = [ + this.#active, + ...this.#entries.map((entry) => entry.call), + ].filter((call): call is RuntimeCall => call !== undefined) + await Promise.all(calls.map((call) => call.cancel(cause))) + await this.idle() + } + + #ensureDraining(): void { + if (this.#draining) return + this.#draining = this.#drain().finally(() => { + this.#draining = undefined + if (this.#entries.length > 0) this.#ensureDraining() + }) + } + + async #drain(): Promise { + while (this.#entries.length > 0) { + const entry = this.#entries.shift() + if (!entry) return + entry.unsubscribeTerminal() + if (entry.call.state === 'terminal') continue + + this.#active = entry.call + try { + // Queue ownership ends at terminal transition, not at socket delivery. + void entry.call.execute(entry.execute) + await entry.call.settled + } finally { + this.#active = undefined + } + } + } +} + +class RuntimeCallCancelled extends Error { + constructor(cause: CallCancellationCause | undefined) { + super(`Runtime Call cancelled: ${cause ?? 'unknown'}.`) + this.name = 'RuntimeCallCancelled' + } +} + +function operationError(error: unknown): RuntimeError { + if (error instanceof RuntimeOperationError) { + return { code: error.code, message: error.message } + } + return { + code: 'operation-failed', + message: error instanceof Error ? error.message : String(error), + } +} diff --git a/src/runtime-caller.ts b/src/runtime-caller.ts new file mode 100644 index 0000000..1a44302 --- /dev/null +++ b/src/runtime-caller.ts @@ -0,0 +1,176 @@ +import type { Socket } from 'node:net' + +import type { + RuntimeFrame, + RuntimeInputFrame, + RuntimeInputRequest, +} from './runtime-protocol' + +export type RuntimeCaller = { + id: string + onDisconnect: (listener: () => void) => () => void + requestInput: ( + request: Omit, + signal?: AbortSignal, + ) => Promise + send: (frame: RuntimeFrame) => Promise +} + +export type InMemoryRuntimeCaller = RuntimeCaller & { + frames: RuntimeFrame[] + disconnect: () => void +} + +export type SocketRuntimeCaller = RuntimeCaller & { + receiveInput: (frame: RuntimeInputFrame) => void +} + +type FrameWriter = (frame: RuntimeFrame) => Promise + +export function createInMemoryRuntimeCaller( + requestId: string, +): InMemoryRuntimeCaller { + const frames: RuntimeFrame[] = [] + const disconnectListeners = new Set<() => void>() + let disconnected = false + const caller = createGuardedCaller( + requestId, + async (frame) => { + frames.push(frame) + }, + disconnectListeners, + () => disconnected, + async () => { + throw new Error('In-memory caller has no input provider.') + }, + ) + + return { + ...caller, + frames, + disconnect: () => { + if (disconnected) return + disconnected = true + for (const listener of [...disconnectListeners]) listener() + }, + } +} + +export function createSocketRuntimeCaller( + requestId: string, + socket: Socket, +): SocketRuntimeCaller { + const disconnectListeners = new Set<() => void>() + const inputWaiters = new Map< + string, + { resolve: (value: unknown) => void; reject: (error: Error) => void } + >() + let disconnected = socket.destroyed + const notifyDisconnect = () => { + if (disconnected) return + disconnected = true + for (const listener of [...disconnectListeners]) listener() + for (const waiter of inputWaiters.values()) { + waiter.reject(new Error(`Runtime caller ${requestId} is disconnected.`)) + } + inputWaiters.clear() + } + socket.once('close', notifyDisconnect) + + const caller = createGuardedCaller( + requestId, + (frame) => + new Promise((resolve, reject) => { + if (disconnected || socket.destroyed) { + reject(new Error(`Runtime caller ${requestId} is disconnected.`)) + return + } + socket.write(`${JSON.stringify(frame)}\n`, (error) => { + if (error) reject(error) + else resolve() + }) + }), + disconnectListeners, + () => disconnected || socket.destroyed, + (request, signal) => { + const inputId = crypto.randomUUID() + return new Promise((resolve, reject) => { + const onAbort = () => { + inputWaiters.delete(inputId) + reject(signal?.reason ?? new Error('Runtime input cancelled.')) + } + if (signal?.aborted) { + onAbort() + return + } + inputWaiters.set(inputId, { + resolve: (value) => { + signal?.removeEventListener('abort', onAbort) + resolve(value) + }, + reject: (error) => { + signal?.removeEventListener('abort', onAbort) + reject(error) + }, + }) + signal?.addEventListener('abort', onAbort, { once: true }) + void caller + .send({ + requestId, + kind: 'event', + event: { type: 'input-required', data: { ...request, inputId } }, + }) + .catch((error) => { + inputWaiters.delete(inputId) + reject(error) + }) + }) + }, + ) + return { + ...caller, + receiveInput: (frame) => { + if (frame.requestId !== requestId) return + const waiter = inputWaiters.get(frame.inputId) + if (!waiter) return + inputWaiters.delete(frame.inputId) + waiter.resolve(frame.value) + }, + } +} + +function createGuardedCaller( + requestId: string, + write: FrameWriter, + disconnectListeners: Set<() => void>, + isDisconnected: () => boolean, + requestInput: RuntimeCaller['requestInput'], +): RuntimeCaller { + let terminal = false + return { + id: requestId, + requestInput, + onDisconnect: (listener) => { + if (isDisconnected()) { + listener() + return () => {} + } + disconnectListeners.add(listener) + return () => disconnectListeners.delete(listener) + }, + send: async (frame) => { + if (frame.requestId !== requestId) { + throw new Error( + `Runtime frame ${frame.requestId} does not match caller ${requestId}.`, + ) + } + if (terminal) { + throw new Error( + `Runtime caller ${requestId} already received a terminal frame.`, + ) + } + await write(frame) + if (frame.kind !== 'event') terminal = true + }, + } +} diff --git a/src/runtime-client.ts b/src/runtime-client.ts new file mode 100644 index 0000000..3969c59 --- /dev/null +++ b/src/runtime-client.ts @@ -0,0 +1,151 @@ +import type { + RuntimeFrame, + RuntimeInputRequest, + RuntimeIntent, +} from './runtime-protocol' + +import { connectDaemonSocket, ensureDaemon } from './daemon-client' +import { readJsonLines, requestJsonLine, writeJsonLine } from './daemon-io' +import { helloMessage } from './daemon-protocol' + +export async function requestRuntime( + intent: RuntimeIntent, + mainPath: string, + options: { + start?: boolean + onEvent?: (frame: RuntimeFrame) => void + onInput?: ( + request: RuntimeInputRequest, + signal: AbortSignal, + ) => Promise + } = {}, +): Promise { + if (options.start ?? true) await ensureDaemon(mainPath) + const socket = await connectDaemonSocket() + try { + const hello = await requestJsonLine(socket, helloMessage()) + if (!isCompatibleHello(hello)) { + throw new Error('Invalid MCP Runtime handshake.') + } + return await new Promise((resolve, reject) => { + let unsubscribe = () => {} + const inputControllers = new Map() + const cleanup = () => { + for (const controller of inputControllers.values()) controller.abort() + inputControllers.clear() + unsubscribe() + socket.off('error', onError) + socket.off('close', onClose) + socket.off('end', onClose) + } + const onError = (error: Error) => { + cleanup() + reject(error) + } + const onClose = () => { + cleanup() + reject( + new Error('MCP Runtime connection closed before a terminal frame.'), + ) + } + socket.once('error', onError) + socket.once('close', onClose) + socket.once('end', onClose) + unsubscribe = readJsonLines( + socket, + (value) => { + if (!isRuntimeFrame(value) || value.requestId !== intent.requestId) { + cleanup() + reject(new Error('Invalid MCP Runtime frame.')) + return + } + if (value.kind === 'event') { + options.onEvent?.(value) + const request = runtimeInputRequest(value) + if (request) { + const controller = new AbortController() + inputControllers.set(request.inputId, controller) + void respondToInput( + socket, + intent.requestId, + request, + options.onInput, + controller.signal, + ).finally(() => inputControllers.delete(request.inputId)) + } + return + } + cleanup() + if (value.kind === 'error') reject(runtimeClientError(value)) + else resolve(value.result) + }, + (error) => { + cleanup() + reject(error) + }, + ) + writeJsonLine(socket, intent) + }) + } finally { + socket.end() + } +} + +async function respondToInput( + socket: import('node:net').Socket, + requestId: string, + request: RuntimeInputRequest, + provider: + | ((request: RuntimeInputRequest, signal: AbortSignal) => Promise) + | undefined, + signal: AbortSignal, +): Promise { + const value = provider + ? await provider(request, signal).catch((error) => ({ + cancelled: true, + reason: error instanceof Error ? error.message : String(error), + })) + : { cancelled: true, reason: 'Interactive input is unavailable.' } + if (signal.aborted || socket.destroyed || !socket.writable) return + writeJsonLine(socket, { + kind: 'input', + requestId, + inputId: request.inputId, + value, + }) +} + +function runtimeInputRequest( + frame: Extract, +): RuntimeInputRequest | undefined { + if (frame.event.type !== 'input-required') return undefined + return frame.event.data as RuntimeInputRequest +} + +function isCompatibleHello(value: unknown): boolean { + return ( + !!value && + typeof value === 'object' && + 'ok' in value && + (value as { ok?: unknown }).ok === true + ) +} + +function isRuntimeFrame(value: unknown): value is RuntimeFrame { + if (!value || typeof value !== 'object') return false + const frame = value as { requestId?: unknown; kind?: unknown } + return ( + typeof frame.requestId === 'string' && + (frame.kind === 'event' || + frame.kind === 'result' || + frame.kind === 'error') + ) +} + +function runtimeClientError( + frame: Extract, +): Error { + const error = new Error(frame.error.message) as Error & { code?: string } + error.code = frame.error.code + return error +} diff --git a/src/runtime-protocol.ts b/src/runtime-protocol.ts new file mode 100644 index 0000000..ca37f8c --- /dev/null +++ b/src/runtime-protocol.ts @@ -0,0 +1,275 @@ +import { MCPX_VERSION } from './version' + +export const DAEMON_PROTOCOL_VERSION = 3 + +export type NotificationMode = 'buffer' | 'discard' + +export type RuntimeErrorCode = + | 'caller-disconnected' + | 'cancelled' + | 'connection-complete' + | 'handshake-required' + | 'invalid-frame' + | 'operation-failed' + | 'protocol-mismatch' + | 'reauth-required' + | 'timeout' + +export type RuntimeError = { + code: RuntimeErrorCode + message: string +} + +export type RuntimeEvent = { + type: string + message?: string + data?: unknown +} + +export type RuntimeInputRequest = { + inputId: string + type: 'oauth-client' + serverName: string + redirectUri: string + issuer: string + scopes: string[] +} + +export type RuntimeInputFrame = { + kind: 'input' + requestId: string + inputId: string + value: unknown +} + +export function isRuntimeInputFrame( + value: unknown, +): value is RuntimeInputFrame { + return ( + isRecord(value) && + hasExactly(value, ['kind', 'requestId', 'inputId', 'value']) && + value.kind === 'input' && + isNonEmptyString(value.requestId) && + isNonEmptyString(value.inputId) + ) +} + +export type RuntimeFrame = + | { requestId: string; kind: 'event'; event: RuntimeEvent } + | { requestId: string; kind: 'result'; result: unknown } + | { requestId: string; kind: 'error'; error: RuntimeError } + +export type RuntimeIntent = + | { requestId: string; op: 'registrySnapshot' } + | { + requestId: string + op: 'call' + serverName: string + toolName: string + input: Record + notificationMode?: NotificationMode + } + | { + requestId: string + op: 'addServer' + serverName: string + transport: 'http' | 'stdio' + url?: string + bearer?: string[] + command?: string + args?: string[] + env?: Record + } + | { requestId: string; op: 'removeServers'; serverNames: string[] } + | { requestId: string; op: 'refreshServers'; serverNames?: string[] } + | { requestId: string; op: 'status' } + | { requestId: string; op: 'stop' } + +export type RuntimeHello = { + kind: 'hello' + protocolVersion: number + clientVersion: string +} + +export type RuntimeConnectionInput = RuntimeHello | RuntimeIntent + +export type RuntimeConnectionAcceptance = + | { kind: 'hello' } + | { kind: 'intent'; intent: RuntimeIntent } + | { kind: 'error'; error: RuntimeError } + +export function createHello(): RuntimeHello { + return { + kind: 'hello', + protocolVersion: DAEMON_PROTOCOL_VERSION, + clientVersion: MCPX_VERSION, + } +} + +export function parseRuntimeIntent( + value: unknown, +): RuntimeIntent | RuntimeError { + if (!isRecord(value) || !isNonEmptyString(value.requestId)) { + return invalidFrame() + } + + switch (value.op) { + case 'registrySnapshot': + case 'status': + case 'stop': + return hasExactly(value, ['requestId', 'op']) + ? (value as RuntimeIntent) + : invalidFrame() + case 'call': { + const allowed = [ + 'requestId', + 'op', + 'serverName', + 'toolName', + 'input', + 'notificationMode', + ] + if ( + !hasOnly(value, allowed) || + !hasRequired(value, allowed.slice(0, 5)) || + !isNonEmptyString(value.serverName) || + !isNonEmptyString(value.toolName) || + !isRecord(value.input) || + (value.notificationMode !== undefined && + value.notificationMode !== 'buffer' && + value.notificationMode !== 'discard') + ) { + return invalidFrame() + } + return value as RuntimeIntent + } + case 'addServer': { + const allowed = [ + 'requestId', + 'op', + 'serverName', + 'transport', + 'url', + 'bearer', + 'command', + 'args', + 'env', + ] + if ( + !hasOnly(value, allowed) || + !hasRequired(value, ['requestId', 'op', 'serverName', 'transport']) || + !isNonEmptyString(value.serverName) || + (value.transport !== 'http' && value.transport !== 'stdio') || + (value.url !== undefined && !isNonEmptyString(value.url)) || + (value.command !== undefined && !isNonEmptyString(value.command)) || + (value.bearer !== undefined && !isStringArray(value.bearer)) || + (value.args !== undefined && !isStringArray(value.args)) || + (value.env !== undefined && !isStringRecord(value.env)) + ) { + return invalidFrame() + } + return value as RuntimeIntent + } + case 'removeServers': + return hasExactly(value, ['requestId', 'op', 'serverNames']) && + isStringArray(value.serverNames) + ? (value as RuntimeIntent) + : invalidFrame() + case 'refreshServers': + return hasOnly(value, ['requestId', 'op', 'serverNames']) && + hasRequired(value, ['requestId', 'op']) && + (value.serverNames === undefined || isStringArray(value.serverNames)) + ? (value as RuntimeIntent) + : invalidFrame() + default: + return invalidFrame() + } +} + +export class RuntimeConnectionState { + #state: 'awaiting-hello' | 'ready' | 'complete' = 'awaiting-hello' + + accept(value: unknown): RuntimeConnectionAcceptance { + if (this.#state === 'complete') { + return { + kind: 'error', + error: { + code: 'connection-complete', + message: 'A Runtime connection accepts exactly one operation.', + }, + } + } + + if (this.#state === 'awaiting-hello') { + if (!isRuntimeHello(value)) { + return { + kind: 'error', + error: { + code: 'handshake-required', + message: 'A Runtime connection must begin with a handshake.', + }, + } + } + if (value.protocolVersion !== DAEMON_PROTOCOL_VERSION) { + this.#state = 'complete' + return { + kind: 'error', + error: { + code: 'protocol-mismatch', + message: `Unsupported Runtime protocol ${value.protocolVersion}; expected ${DAEMON_PROTOCOL_VERSION}.`, + }, + } + } + this.#state = 'ready' + return { kind: 'hello' } + } + + const intent = parseRuntimeIntent(value) + if ('code' in intent) return { kind: 'error', error: intent } + this.#state = 'complete' + return { kind: 'intent', intent } + } +} + +function isRuntimeHello(value: unknown): value is RuntimeHello { + return ( + isRecord(value) && + hasExactly(value, ['kind', 'protocolVersion', 'clientVersion']) && + value.kind === 'hello' && + typeof value.protocolVersion === 'number' && + isNonEmptyString(value.clientVersion) + ) +} + +function invalidFrame(): RuntimeError { + return { code: 'invalid-frame', message: 'Invalid Runtime intent.' } +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +function isNonEmptyString(value: unknown): value is string { + return typeof value === 'string' && value.length > 0 +} + +function isStringArray(value: unknown): value is string[] { + return Array.isArray(value) && value.every(isNonEmptyString) +} + +function isStringRecord(value: unknown): value is Record { + return isRecord(value) && Object.values(value).every(isNonEmptyString) +} + +function hasExactly(value: Record, keys: string[]): boolean { + return hasOnly(value, keys) && hasRequired(value, keys) +} + +function hasOnly(value: Record, keys: string[]): boolean { + const allowed = new Set(keys) + return Object.keys(value).every((key) => allowed.has(key)) +} + +function hasRequired(value: Record, keys: string[]): boolean { + return keys.every((key) => Object.hasOwn(value, key)) +} diff --git a/src/runtime-session-pool.ts b/src/runtime-session-pool.ts new file mode 100644 index 0000000..5caae1e --- /dev/null +++ b/src/runtime-session-pool.ts @@ -0,0 +1,476 @@ +import { createHash } from 'node:crypto' + +import type { McpConnection } from './mcp-client' +import type { RuntimeCaller } from './runtime-caller' +import type { NotificationMode } from './runtime-protocol' +import type { DeclaredServer, RuntimeStores } from './runtime-stores' +import type { McpTool, ServerConfig, ToolDefinition } from './types' + +import { normalizeAuthScheme } from './headers' +import { + connectMcpClient, + listAllMcpTools, + toolCallRequestOptions, +} from './mcp-client' +import { assignCommandNames } from './names' +import { + createNotificationBuffer, + flushNotificationBuffer, +} from './notifications' +import { + CancelableFifo, + RuntimeCall, + RuntimeOperationError, +} from './runtime-call' + +type Connect = typeof connectMcpClient + +type ResolvedServer = { + key: string + config: ServerConfig +} + +type ManagedRuntimeSession = { + key: string + serverNames: Set + config: ServerConfig + connection?: McpConnection + connecting?: Promise + queue: CancelableFifo + lastSessionId?: string + currentBuffer?: ReturnType + pendingToolsChanged: boolean + lastUsedAt: number + evictCount: number + closing: boolean +} + +export type RuntimeCallInput = { + serverName: string + toolName: string + input: Record + notificationMode?: NotificationMode +} + +export type RuntimeSessionStatus = { + serverKey: string + labels: string[] + transport: 'stdio' | 'http' + pid?: number | null + url?: string + activeCalls: number + queuedCalls: number + idleMs: number + evictCount: number + hasRetainedSessionId: boolean +} + +export class RuntimeSessionPool { + readonly #stores: RuntimeStores + readonly #connect: Connect + readonly #sessions = new Map() + readonly #bearerCursors = new Map() + #accepting = true + + constructor(stores: RuntimeStores, options: { connect?: Connect } = {}) { + this.#stores = stores + this.#connect = options.connect ?? connectMcpClient + } + + async call(call: RuntimeCall, input: RuntimeCallInput): Promise { + if (!this.#accepting) { + throw new RuntimeOperationError('cancelled', 'MCP Runtime is stopping.') + } + const resolved = await this.#resolveServer(input.serverName) + this.#assertAccepting() + const session = this.#sessionFor(input.serverName, resolved) + session.queue.enqueue(call, async (signal) => { + const headers = await this.#resolveHeaders(input.serverName) + signal.throwIfAborted() + if (session.closing) { + throw new RuntimeOperationError( + 'cancelled', + 'MCP Runtime session is closing.', + ) + } + const buffer = createNotificationBuffer() + const requestOptions = toolCallRequestOptions() + if (input.notificationMode !== 'discard') { + requestOptions.onprogress = (progress) => { + buffer.add({ + method: 'notifications/progress', + params: { progressToken: call.id, ...progress }, + }) + } + } + const timeout = setTimeout(() => { + void call.cancel('timeout') + }, requestOptions.timeout) + timeout.unref() + session.currentBuffer = + input.notificationMode === 'discard' ? undefined : buffer + try { + const connection = await this.#ensureConnected(session, headers) + const result = await connection.client.callTool( + { name: input.toolName, arguments: input.input }, + undefined, + { ...requestOptions, signal }, + ) + const notifications = await flushNotificationBuffer(buffer) + const toolsChanged = + buffer.toolsChanged() || session.pendingToolsChanged + session.pendingToolsChanged = false + if (toolsChanged) { + const tools = await listAllMcpTools(connection.client) + await this.#stores.updateState((state) => { + state.schemas.servers[input.serverName] = { + tools: withCommandNames(tools), + discoveredAt: new Date().toISOString(), + refreshStatus: { + checkedAt: new Date().toISOString(), + status: 'ok', + }, + } + }) + } + return { result, notifications, toolsChanged } + } catch (error) { + if (isUnauthorizedError(error)) { + session.evictCount += 1 + await this.#closeSession(session, false) + throw new RuntimeOperationError( + 'reauth-required', + `Credentials for ${input.serverName} must be refreshed.`, + ) + } + throw error + } finally { + clearTimeout(timeout) + delete session.currentBuffer + session.lastUsedAt = Date.now() + } + }) + await call.settled + } + + async listTools( + serverName: string, + caller: RuntimeCaller, + ): Promise>> { + this.#assertAccepting() + const resolved = await this.#resolveServer(serverName) + this.#assertAccepting() + const session = this.#sessionFor(serverName, resolved) + let tools: Awaited> = [] + const childCaller: RuntimeCaller = { + id: `${caller.id}:schema:${serverName}`, + onDisconnect: caller.onDisconnect, + requestInput: caller.requestInput, + send: async (frame) => { + if (frame.kind === 'result') tools = frame.result as typeof tools + }, + } + const call = new RuntimeCall(childCaller) + session.queue.enqueue(call, async (signal) => { + const headers = await this.#resolveHeaders(serverName) + signal.throwIfAborted() + if (session.closing) { + throw new RuntimeOperationError( + 'cancelled', + 'MCP Runtime session is closing.', + ) + } + return listAllMcpTools( + (await this.#ensureConnected(session, headers)).client, + ) + }) + await call.settled + return tools + } + + status(): RuntimeSessionStatus[] { + const now = Date.now() + return [...this.#sessions.values()].map((session) => { + const item: RuntimeSessionStatus = { + serverKey: session.key, + labels: [...session.serverNames].sort(), + transport: session.config.transport === 'stdio' ? 'stdio' : 'http', + ...session.queue.status(), + idleMs: now - session.lastUsedAt, + evictCount: session.evictCount, + hasRetainedSessionId: session.lastSessionId !== undefined, + } + if (session.config.transport === 'stdio') { + item.pid = session.connection?.pid() ?? null + } else { + const url = new URL(session.config.url) + item.url = `${url.host}${url.pathname}` + } + return item + }) + } + + async cleanupIdle(maxIdleMs: number): Promise { + const cutoff = Date.now() - maxIdleMs + const idle = [...this.#sessions.values()].filter((session) => { + const status = session.queue.status() + return ( + status.activeCalls === 0 && + status.queuedCalls === 0 && + session.lastUsedAt <= cutoff + ) + }) + for (const session of idle) this.#sessions.delete(session.key) + await Promise.all( + idle.map((session) => this.#closeSession(session, false, true)), + ) + } + + sessionCount(): number { + return this.#sessions.size + } + + async close(): Promise { + this.#accepting = false + await Promise.all( + [...this.#sessions.values()].map((session) => + session.queue.cancelAll('runtime-stopping'), + ), + ) + const sessions = [...this.#sessions.values()] + this.#sessions.clear() + await Promise.all( + sessions.map(async (session) => { + await session.queue.idle() + await this.#closeSession(session, false, true) + }), + ) + } + + async #resolveServer(serverName: string): Promise { + const { registry, credentials } = await this.#stores.readState() + const declared = registry.servers[serverName] + if (!declared) { + throw new RuntimeOperationError( + 'operation-failed', + `Unknown MCP server: ${serverName}.`, + ) + } + + const key = stableServerKey(declared) + if (declared.transport === 'stdio') { + return { + key: stableServerKey({ name: serverName, ...declared }), + config: { ...declared, env: credentials.stdioEnv[serverName] }, + } + } + + return { key, config: { url: declared.url, auth: { kind: 'none' } } } + } + + async #resolveHeaders( + serverName: string, + ): Promise | undefined> { + const { registry, credentials } = await this.#stores.readState() + const declared = registry.servers[serverName] + if (!declared) + throw new RuntimeOperationError( + 'operation-failed', + `Unknown MCP server: ${serverName}.`, + ) + if (declared.transport === 'stdio') return undefined + + const headers: Record = { + Accept: 'application/json, text/event-stream', + ...(credentials.headers[serverName] ?? {}), + } + if (declared.auth.kind === 'bearer') { + const key = stableServerKey(declared) + const cursor = this.#bearerCursors.get(key) ?? 0 + const credential = + declared.auth.credentials[cursor % declared.auth.credentials.length] + if (!credential) { + throw reauthRequired(serverName) + } + this.#bearerCursors.set( + key, + (cursor + 1) % declared.auth.credentials.length, + ) + const value = + credential.kind === 'env' + ? process.env[credential.name] + : credentials.bearer[credential.key] + if (!value) throw reauthRequired(serverName) + headers.Authorization = value.startsWith('Bearer ') + ? value + : `Bearer ${value}` + } else if (declared.auth.kind === 'oauth-token') { + const token = credentials.oauth[declared.auth.tokenKey] + if (!token || oauthTokenIsUnusable(token.expiresAt)) { + throw reauthRequired(serverName) + } + headers.Authorization = `${normalizeAuthScheme(token.tokenType)} ${token.accessToken}` + } else if ( + declared.auth.kind === 'oauth' || + declared.auth.kind === 'unknown' + ) { + throw reauthRequired(serverName) + } + + return headers + } + + #sessionFor( + serverName: string, + resolved: ResolvedServer, + ): ManagedRuntimeSession { + const existing = this.#sessions.get(resolved.key) + if (existing) { + existing.serverNames.add(serverName) + return existing + } + + const session: ManagedRuntimeSession = { + key: resolved.key, + serverNames: new Set([serverName]), + config: resolved.config, + queue: new CancelableFifo(), + pendingToolsChanged: false, + lastUsedAt: Date.now(), + evictCount: 0, + closing: false, + } + this.#sessions.set(resolved.key, session) + return session + } + + async #ensureConnected( + session: ManagedRuntimeSession, + headers?: Record, + ): Promise { + if (session.closing) { + throw new RuntimeOperationError( + 'cancelled', + 'MCP Runtime session is closing.', + ) + } + if (session.connection) { + if (headers) session.connection.updateHeaders(headers) + return session.connection + } + if (!session.connecting) { + session.connecting = this.#connect(session.config, { + headers, + sessionId: session.lastSessionId, + onNotification: (notification) => { + if (session.currentBuffer) { + session.currentBuffer.add(notification) + } else if ( + notification.method === 'notifications/tools/list_changed' + ) { + session.pendingToolsChanged = true + } + }, + }) + .then(async (connection) => { + if (session.closing) { + await connection.close().catch(() => {}) + throw new RuntimeOperationError( + 'cancelled', + 'MCP Runtime session closed while connecting.', + ) + } + session.connection = connection + return connection + }) + .finally(() => { + delete session.connecting + }) + } + return session.connecting + } + + async #closeSession( + session: ManagedRuntimeSession, + retainSessionId: boolean, + retire = false, + ): Promise { + session.closing = true + if (retainSessionId && session.config.transport !== 'stdio') { + session.lastSessionId = + session.connection?.sessionId() ?? session.lastSessionId + } + const connection = + session.connection ?? (await session.connecting?.catch(() => undefined)) + await connection?.close().catch(() => {}) + delete session.connection + delete session.connecting + if (!retainSessionId) delete session.lastSessionId + if (!retire) session.closing = false + } + + #assertAccepting(): void { + if (!this.#accepting) { + throw new RuntimeOperationError('cancelled', 'MCP Runtime is stopping.') + } + } +} + +function stableServerKey(server: unknown): string { + return createHash('sha256') + .update(JSON.stringify(sortValue(server))) + .digest('hex') + .slice(0, 16) +} + +function sortValue(value: unknown): unknown { + if (Array.isArray(value)) return value.map(sortValue) + if (!value || typeof value !== 'object') return value + return Object.fromEntries( + Object.entries(value as Record) + .sort(([left], [right]) => left.localeCompare(right)) + .map(([key, item]) => [key, sortValue(item)]), + ) +} + +function oauthTokenIsUnusable(expiresAt: string | undefined): boolean { + if (!expiresAt) return false + const expiresAtMs = Date.parse(expiresAt) + return !Number.isFinite(expiresAtMs) || expiresAtMs <= Date.now() + 60_000 +} + +function reauthRequired(serverName: string): RuntimeOperationError { + return new RuntimeOperationError( + 'reauth-required', + `Credentials for ${serverName} must be refreshed.`, + ) +} + +function isUnauthorizedError(error: unknown): boolean { + const message = error instanceof Error ? error.message : String(error) + return ( + message.includes('401') || message.toLowerCase().includes('unauthorized') + ) +} + +function withCommandNames(tools: McpTool[]): ToolDefinition[] { + const names = assignCommandNames(tools.map((tool) => tool.name)) + return tools.map((tool) => { + const normalized: ToolDefinition = { + name: tool.name, + commandName: names.get(tool.name) ?? tool.name, + } + if (tool.title) normalized.title = tool.title + if (tool.description) normalized.description = tool.description + if ( + tool.inputSchema && + typeof tool.inputSchema === 'object' && + !Array.isArray(tool.inputSchema) + ) { + normalized.inputSchema = tool.inputSchema as Record + } + if (tool.annotations) normalized.annotations = tool.annotations + if (tool._meta) normalized._meta = tool._meta + return normalized + }) +} diff --git a/src/runtime-stores.ts b/src/runtime-stores.ts new file mode 100644 index 0000000..f576747 --- /dev/null +++ b/src/runtime-stores.ts @@ -0,0 +1,549 @@ +import { randomUUID } from 'node:crypto' +import fs from 'node:fs/promises' +import path from 'node:path' + +import type { + HttpServerConfig, + OAuthToken, + RegistryConfig, + ServerConfig, + ServerRefreshStatus, + StdioServerConfig, + ToolDefinition, +} from './types' + +import { assignCommandNames } from './names' + +const STATE_DIR = 'state-v2' +const REGISTRY_FILE = 'registry.json' +const CREDENTIALS_FILE = 'credentials.json' +const SCHEMAS_FILE = 'schema-cache.json' +const TRANSACTION_FILE = 'transaction.json' +const LEGACY_REGISTRY_FILE = 'servers.json' +const LEGACY_CREDENTIALS_FILE = 'tokens.json' +const LEGACY_REGISTRY_BACKUP = 'servers.v1.backup.json' +const LEGACY_CREDENTIALS_BACKUP = 'tokens.v1.backup.json' + +export type DeclaredBearerCredential = + | { kind: 'env'; name: string } + | { kind: 'stored'; key: string } + +type DeclaredAuth = + | Exclude + | { + kind: 'bearer' + credentials: DeclaredBearerCredential[] + strategy: 'round-robin' + confidence: 'configured' + } + +type DeclaredHttpServer = Omit< + HttpServerConfig, + 'headers' | 'tools' | 'discoveredAt' | 'refreshStatus' | 'auth' +> & { auth: DeclaredAuth } + +type DeclaredStdioServer = Omit< + StdioServerConfig, + 'env' | 'tools' | 'discoveredAt' | 'refreshStatus' +> + +export type DeclaredServer = DeclaredHttpServer | DeclaredStdioServer + +export type DeclaredRegistry = { + version: 2 + servers: Record +} + +export type CredentialState = { + version: 2 + oauth: Record + oauthClientSecrets: Record + bearer: Record + headers: Record> + stdioEnv: Record> +} + +export type ServerSchemaState = { + discoveredAt?: string + tools?: ToolDefinition[] + refreshStatus?: ServerRefreshStatus + dirty?: boolean +} + +export type SchemaCache = { + version: 1 + servers: Record +} + +export type RuntimeRegistrySnapshot = { + version: 2 + servers: Record +} + +type JsonStore = { + read: () => Promise +} + +export type RuntimeState = { + registry: DeclaredRegistry + credentials: CredentialState + schemas: SchemaCache +} + +export type RuntimeStores = { + registry: JsonStore + credentials: JsonStore + schemas: JsonStore + readState: () => Promise + updateState: ( + update: (state: RuntimeState) => T | Promise, + ) => Promise + readSnapshot: () => Promise + upsertServer: (name: string, server: ServerConfig) => Promise + removeServers: ( + names: string[], + ) => Promise> +} + +type LegacyCredentials = { + version: 1 + oauth: Record + oauthClientSecrets?: Record +} + +export async function openRuntimeStores(root: string): Promise { + const stateDir = path.join(root, STATE_DIR) + await cleanupUnpublishedMigrations(root) + if (!(await exists(stateDir))) await migrateLegacyState(root, stateDir) + await recoverTransaction(stateDir) + await archiveLegacyFiles(root) + + const paths = statePaths(stateDir) + let tail = Promise.resolve() + let poison: Error | undefined + const exclusive = (operation: () => Promise): Promise => { + const guarded = () => { + if (poison) throw poison + return operation() + } + const result = tail.then(guarded, guarded) + tail = result.then( + () => undefined, + () => undefined, + ) + return result + } + const readState = () => readRuntimeState(paths) + const commitState = async (state: RuntimeState) => { + try { + await commitRuntimeState(paths, state) + } catch (error) { + poison = new Error( + 'Runtime stores are unavailable after an unrecoverable commit failure.', + { cause: error }, + ) + throw poison + } + } + const registry = serializedJsonStore( + paths.registry, + exclusive, + ) + const credentials = serializedJsonStore( + paths.credentials, + exclusive, + ) + const schemas = serializedJsonStore(paths.schemas, exclusive) + const updateState = ( + update: (state: RuntimeState) => T | Promise, + ): Promise => + exclusive(async () => { + const state = await readState() + const result = await update(state) + await commitState(state) + return result + }) + + return { + registry, + credentials, + schemas, + readState: () => exclusive(readState), + updateState, + readSnapshot: () => + exclusive(async () => { + const { registry: declarations, schemas: cache } = await readState() + const servers: RuntimeRegistrySnapshot['servers'] = {} + for (const [name, server] of Object.entries(declarations.servers)) { + servers[name] = { ...server, ...cache.servers[name] } + } + return { version: 2, servers } + }), + upsertServer: (name, server) => + updateState((state) => { + const { + registry: declared, + credentials: secrets, + schemas: cache, + } = state + const split = splitLegacyState( + { version: 1, servers: { [name]: server } }, + { version: 1, oauth: {} }, + ) + declared.servers[name] = split.registry.servers[name]! + delete secrets.headers[name] + delete secrets.stdioEnv[name] + for (const key of Object.keys(secrets.bearer)) { + if (key.startsWith(`${name}:bearer:`)) delete secrets.bearer[key] + } + Object.assign(secrets.bearer, split.credentials.bearer) + Object.assign(secrets.headers, split.credentials.headers) + Object.assign(secrets.stdioEnv, split.credentials.stdioEnv) + if (split.schemas.servers[name]) { + cache.servers[name] = split.schemas.servers[name] + } else { + delete cache.servers[name] + } + sweepUnreferencedOAuth(state) + }), + removeServers: (names) => + updateState((state) => { + const { + registry: declared, + credentials: secrets, + schemas: cache, + } = state + const removed: Array<{ name: string; tokenRemoved: boolean }> = [] + for (const name of names) { + const server = declared.servers[name] + if (!server) continue + let tokenRemoved = false + delete declared.servers[name] + delete cache.servers[name] + delete secrets.headers[name] + delete secrets.stdioEnv[name] + for (const key of Object.keys(secrets.bearer)) { + if (key.startsWith(`${name}:bearer:`)) delete secrets.bearer[key] + } + if ( + server.transport !== 'stdio' && + server.auth.kind === 'oauth-token' + ) { + tokenRemoved = !isOAuthTokenReferenced( + declared, + server.auth.tokenKey, + ) + } + removed.push({ name, tokenRemoved }) + } + sweepUnreferencedOAuth(state) + return removed + }), + } +} + +async function cleanupUnpublishedMigrations(root: string): Promise { + // Runtime startup ownership makes every staging directory here an abandoned transaction. + const entries = await fs + .readdir(root, { withFileTypes: true }) + .catch((error) => { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return [] + throw error + }) + await Promise.all( + entries + .filter( + (entry) => + entry.isDirectory() && entry.name.startsWith(`.${STATE_DIR}.`), + ) + .map((entry) => + fs.rm(path.join(root, entry.name), { recursive: true, force: true }), + ), + ) +} + +async function migrateLegacyState( + root: string, + stateDir: string, +): Promise { + const registry = (await readOptionalJson( + path.join(root, LEGACY_REGISTRY_FILE), + )) ?? { version: 1, servers: {} } + const credentials = (await readOptionalJson( + path.join(root, LEGACY_CREDENTIALS_FILE), + )) ?? { version: 1, oauth: {} } + const migrated = splitLegacyState(registry, credentials) + const stagingDir = path.join(root, `.${STATE_DIR}.${randomUUID()}`) + + await fs.mkdir(stagingDir, { recursive: true, mode: 0o700 }) + try { + await Promise.all([ + writeJson(path.join(stagingDir, REGISTRY_FILE), migrated.registry), + writeJson(path.join(stagingDir, CREDENTIALS_FILE), migrated.credentials), + writeJson(path.join(stagingDir, SCHEMAS_FILE), migrated.schemas), + ]) + // Directory publication is the commit point, so readers never see a subset. + await fs.rename(stagingDir, stateDir) + } catch (error) { + await fs.rm(stagingDir, { recursive: true, force: true }).catch(() => {}) + throw error + } +} + +function splitLegacyState( + registry: RegistryConfig, + legacyCredentials: LegacyCredentials, +): { + registry: DeclaredRegistry + credentials: CredentialState + schemas: SchemaCache +} { + const declarations: DeclaredRegistry = { version: 2, servers: {} } + const credentials: CredentialState = { + version: 2, + oauth: legacyCredentials.oauth, + oauthClientSecrets: legacyCredentials.oauthClientSecrets ?? {}, + bearer: {}, + headers: {}, + stdioEnv: {}, + } + const schemas: SchemaCache = { version: 1, servers: {} } + + for (const [name, server] of Object.entries(registry.servers)) { + if (server.transport === 'stdio') { + const { env, tools, discoveredAt, refreshStatus, ...declared } = server + declarations.servers[name] = declared + if (env && Object.keys(env).length > 0) credentials.stdioEnv[name] = env + recordSchema(name, schemas, tools, discoveredAt, refreshStatus) + } else { + const { headers, tools, discoveredAt, refreshStatus, ...declared } = + server + declarations.servers[name] = { + ...declared, + auth: migrateAuth(name, server, credentials), + } + if (headers && Object.keys(headers).length > 0) { + credentials.headers[name] = headers + } + recordSchema(name, schemas, tools, discoveredAt, refreshStatus) + } + } + + return { registry: declarations, credentials, schemas } +} + +function migrateAuth( + name: string, + server: HttpServerConfig, + credentials: CredentialState, +): DeclaredAuth { + if (server.auth.kind !== 'bearer') return server.auth + + return { + ...server.auth, + credentials: server.auth.credentials.map((credential, index) => { + if (credential.kind === 'env') return credential + const key = `${name}:bearer:${index}` + credentials.bearer[key] = credential.value + return { kind: 'stored' as const, key } + }), + } +} + +function recordSchema( + name: string, + cache: SchemaCache, + tools: ToolDefinition[] | undefined, + discoveredAt: string | undefined, + refreshStatus: ServerRefreshStatus | undefined, +): void { + const schema: ServerSchemaState = {} + if (discoveredAt) schema.discoveredAt = discoveredAt + if (refreshStatus) schema.refreshStatus = refreshStatus + if (tools && tools.length > 0) schema.tools = normalizeTools(tools) + if (Object.keys(schema).length > 0) cache.servers[name] = schema +} + +function normalizeTools(tools: ToolDefinition[]): ToolDefinition[] { + const commandNames = assignCommandNames(tools.map((tool) => tool.name)) + return tools.map((tool) => { + const { outputSchema: _outputSchema, ...rest } = tool as ToolDefinition & { + outputSchema?: unknown + } + return { + ...rest, + commandName: commandNames.get(tool.name) ?? tool.name, + } + }) +} + +function serializedJsonStore( + filePath: string, + exclusive: (operation: () => Promise) => Promise, +): JsonStore { + return { + read: () => exclusive(() => readJson(filePath)), + } +} + +type RuntimeStatePaths = { + registry: string + credentials: string + schemas: string + transaction: string +} + +function statePaths(stateDir: string): RuntimeStatePaths { + return { + registry: path.join(stateDir, REGISTRY_FILE), + credentials: path.join(stateDir, CREDENTIALS_FILE), + schemas: path.join(stateDir, SCHEMAS_FILE), + transaction: path.join(stateDir, TRANSACTION_FILE), + } +} + +async function readRuntimeState( + paths: RuntimeStatePaths, +): Promise { + const [registry, credentials, schemas] = await Promise.all([ + readJson(paths.registry), + readJson(paths.credentials), + readJson(paths.schemas), + ]) + credentials.stdioEnv ??= {} + return { registry, credentials, schemas } +} + +async function commitRuntimeState( + paths: RuntimeStatePaths, + state: RuntimeState, + publish: StatePublisher = publishRuntimeState, +): Promise { + // The journal makes a multi-file commit recoverable after any process crash. + await atomicWriteJson(paths.transaction, state) + try { + await publish(paths, state) + } catch { + // Do not permit a mixed generation to feed the next transaction in this process. + await publish(paths, state) + } + await fs.rm(paths.transaction) +} + +type StatePublisher = ( + paths: RuntimeStatePaths, + state: RuntimeState, +) => Promise + +async function publishRuntimeState( + paths: RuntimeStatePaths, + state: RuntimeState, +): Promise { + await Promise.all([ + atomicWriteJson(paths.registry, state.registry), + atomicWriteJson(paths.credentials, state.credentials), + atomicWriteJson(paths.schemas, state.schemas), + ]) +} + +async function recoverTransaction(stateDir: string): Promise { + const paths = statePaths(stateDir) + const state = await readOptionalJson(paths.transaction) + if (!state) return + await Promise.all([ + atomicWriteJson(paths.registry, state.registry), + atomicWriteJson(paths.credentials, state.credentials), + atomicWriteJson(paths.schemas, state.schemas), + ]) + await fs.rm(paths.transaction) +} + +function isOAuthTokenReferenced( + registry: DeclaredRegistry, + tokenKey: string, +): boolean { + return Object.values(registry.servers).some( + (server) => + server.transport !== 'stdio' && + server.auth.kind === 'oauth-token' && + server.auth.tokenKey === tokenKey, + ) +} + +function sweepUnreferencedOAuth(state: RuntimeState): void { + for (const [tokenKey, token] of Object.entries(state.credentials.oauth)) { + if (isOAuthTokenReferenced(state.registry, tokenKey)) continue + delete state.credentials.oauth[tokenKey] + if ( + token.clientSecretKey && + !Object.values(state.credentials.oauth).some( + (candidate) => candidate.clientSecretKey === token.clientSecretKey, + ) + ) { + delete state.credentials.oauthClientSecrets[token.clientSecretKey] + } + } +} + +export const __test = { + commitRuntimeState, +} + +async function readJson(filePath: string): Promise { + return JSON.parse(await fs.readFile(filePath, 'utf8')) as T +} + +async function atomicWriteJson( + filePath: string, + value: unknown, +): Promise { + const tempPath = `${filePath}.${randomUUID()}.tmp` + await writeJson(tempPath, value) + await fs.rename(tempPath, filePath) +} + +async function writeJson(filePath: string, value: unknown): Promise { + await fs.writeFile(filePath, `${JSON.stringify(value, null, '\t')}\n`, { + encoding: 'utf8', + mode: 0o600, + }) +} + +async function readOptionalJson(filePath: string): Promise { + try { + return JSON.parse(await fs.readFile(filePath, 'utf8')) as T + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return undefined + throw error + } +} + +async function archiveLegacyFiles(root: string): Promise { + await Promise.all([ + archiveLegacyFile(root, LEGACY_REGISTRY_FILE, LEGACY_REGISTRY_BACKUP), + archiveLegacyFile(root, LEGACY_CREDENTIALS_FILE, LEGACY_CREDENTIALS_BACKUP), + ]) +} + +async function archiveLegacyFile( + root: string, + legacyName: string, + backupName: string, +): Promise { + const legacyPath = path.join(root, legacyName) + const backupPath = path.join(root, backupName) + if (!(await exists(legacyPath))) return + if (await exists(backupPath)) { + // Published v2 state is authoritative; keep exactly one recovery copy. + await fs.rm(legacyPath) + return + } + await fs.rename(legacyPath, backupPath) +} + +async function exists(filePath: string): Promise { + return fs + .access(filePath) + .then(() => true) + .catch(() => false) +} diff --git a/src/runtime.ts b/src/runtime.ts new file mode 100644 index 0000000..6117740 --- /dev/null +++ b/src/runtime.ts @@ -0,0 +1,222 @@ +import type { RuntimeCaller } from './runtime-caller' +import type { RuntimeIntent } from './runtime-protocol' +import type { RuntimeStores } from './runtime-stores' + +import { discoverServer, normalizeTools } from './discovery' +import { RuntimeAuthentication } from './runtime-authentication' +import { RuntimeCall, RuntimeOperationError } from './runtime-call' +import { DAEMON_PROTOCOL_VERSION } from './runtime-protocol' +import { RuntimeSessionPool } from './runtime-session-pool' +import { MCPX_VERSION } from './version' + +export class McpRuntime { + readonly #stores: RuntimeStores + readonly #sessions: RuntimeSessionPool + readonly #authentication: RuntimeAuthentication + + constructor( + stores: RuntimeStores, + options: { + sessions?: RuntimeSessionPool + authentication?: RuntimeAuthentication + } = {}, + ) { + this.#stores = stores + this.#sessions = options.sessions ?? new RuntimeSessionPool(stores) + this.#authentication = + options.authentication ?? new RuntimeAuthentication(stores) + } + + async cleanupIdleSessions(maxIdleMs: number): Promise { + await this.#sessions.cleanupIdle(maxIdleMs) + } + + activeSessionCount(): number { + return this.#sessions.sessionCount() + } + + activeAuthenticationFlows(): number { + return this.#authentication.activeFlows() + } + + async handle(intent: RuntimeIntent, caller: RuntimeCaller): Promise { + switch (intent.op) { + case 'registrySnapshot': + await caller.send({ + requestId: intent.requestId, + kind: 'result', + result: await this.#stores.readSnapshot(), + }) + return + case 'call': { + const call = new RuntimeCall(caller) + try { + await this.#sessions.call(call, intent) + } catch (error) { + await call.fail(runtimeError(error)) + } + return + } + case 'status': { + const sessions = this.#sessions.status() + await caller.send({ + requestId: intent.requestId, + kind: 'result', + result: { + pid: process.pid, + protocolVersion: DAEMON_PROTOCOL_VERSION, + version: MCPX_VERSION, + activeServers: sessions.length, + servers: sessions, + }, + }) + return + } + case 'stop': + await Promise.all([ + this.#sessions.close(), + this.#authentication.close(), + ]) + await caller.send({ + requestId: intent.requestId, + kind: 'result', + result: { stopping: true }, + }) + return + case 'refreshServers': + try { + const outcome = await this.#authentication.refreshServers( + intent.serverNames, + caller, + ) + if (outcome.status === 'disconnected') return + const snapshot = await this.#stores.readSnapshot() + const names = + intent.serverNames ?? Object.keys(snapshot.servers).sort() + for (const name of names) { + const tools = await this.#sessions.listTools(name, caller) + await this.#stores.updateState((state) => { + state.schemas.servers[name] = { + tools: normalizeTools(tools), + discoveredAt: new Date().toISOString(), + refreshStatus: { + checkedAt: new Date().toISOString(), + status: 'ok', + }, + } + }) + } + await caller.send({ + requestId: intent.requestId, + kind: 'result', + result: outcome, + }) + } catch (error) { + await caller.send({ + requestId: intent.requestId, + kind: 'error', + error: runtimeError(error), + }) + } + return + case 'addServer': { + try { + const result = await discoverServer(discoveryOptions(intent)) + await this.#stores.upsertServer(intent.serverName, result.server) + await caller.send({ + requestId: intent.requestId, + kind: 'result', + result: { + name: intent.serverName, + transport: result.server.transport ?? 'http', + status: result.status, + auth: + result.server.transport === 'stdio' + ? undefined + : result.server.auth, + tools: result.server.tools?.length ?? 0, + message: result.message, + }, + }) + } catch (error) { + await caller.send({ + requestId: intent.requestId, + kind: 'error', + error: runtimeError(error), + }) + } + return + } + case 'removeServers': { + const snapshot = await this.#stores.readSnapshot() + const missing = intent.serverNames.filter( + (name) => !snapshot.servers[name], + ) + if (missing.length > 0) { + await caller.send({ + requestId: intent.requestId, + kind: 'error', + error: { + code: 'operation-failed', + message: `Unknown MCP server(s): ${missing.join(', ')}.`, + }, + }) + return + } + const removed = await this.#stores.removeServers(intent.serverNames) + await caller.send({ + requestId: intent.requestId, + kind: 'result', + result: + removed.length === 1 + ? { ...removed[0], removed: true } + : { + removed: removed.map((item) => ({ + ...item, + removed: true, + })), + }, + }) + } + } + } +} + +function discoveryOptions(intent: Extract) { + if (intent.transport === 'stdio') { + if (!intent.command) throw new Error('Stdio MCP servers require a command.') + return { + name: intent.serverName, + transport: 'stdio' as const, + command: intent.command, + args: intent.args, + env: intent.env, + } + } + if (!intent.url) throw new Error('HTTP MCP servers require a URL.') + return { + name: intent.serverName, + transport: 'http' as const, + url: intent.url, + bearer: intent.bearer, + } +} + +function runtimeError(error: unknown): { + code: 'operation-failed' | 'reauth-required' + message: string +} { + if (error instanceof RuntimeOperationError) { + return { + code: + error.code === 'reauth-required' + ? 'reauth-required' + : 'operation-failed', + message: error.message, + } + } + return { + code: 'operation-failed', + message: error instanceof Error ? error.message : String(error), + } +} diff --git a/src/schema-refresh.ts b/src/schema-refresh.ts deleted file mode 100644 index e48c833..0000000 --- a/src/schema-refresh.ts +++ /dev/null @@ -1,631 +0,0 @@ -import fs from 'node:fs/promises' -import { homedir } from 'node:os' -import path from 'node:path' - -import type { - OAuthToken, - RegistryConfig, - ServerConfig, - ServerRefreshStatus, -} from './types' - -import { discoverAuth } from './auth-discovery' -import { readRegistryConfig, writeRegistryConfig } from './config' -import { reauthenticateServer, refreshServer } from './discovery' -import { resolveHeaders } from './headers' -import { getOAuthToken } from './token-cache' - -const REFRESH_AFTER_MS = 24 * 60 * 60 * 1000 -const STALE_LOCK_AFTER_MS = 10 * 60 * 1000 -const LOCK_PATH = path.join(homedir(), '.agents', 'mcpx', 'schema-refresh.lock') -const WORKER_ENV = 'MCPX_SCHEMA_REFRESH_WORKER' - -export type ServerRefreshResult = { - server: string - status: - | 'schema-refreshed' - | 'auth-refreshed' - | 'reauthenticated' - | 'reauth-required' - | 'unreachable' - toolsBefore: number - toolsAfter?: number - schemaChanged?: boolean - message?: string -} - -export type RefreshSummary = { - checkedAt: string - refreshed: string[] - unchanged: string[] - authRefreshed: string[] - reauthenticated: string[] - reauthRequired: string[] - unreachable: string[] - servers: ServerRefreshResult[] -} - -export function isSchemaRefreshStale( - server: ServerConfig, - now: Date = new Date(), -): boolean { - if (!server.tools || server.tools.length === 0) return false - if (!server.discoveredAt) return true - - const discoveredAt = Date.parse(server.discoveredAt) - if (!Number.isFinite(discoveredAt)) return true - - return now.getTime() - discoveredAt >= REFRESH_AFTER_MS -} - -export function hasStaleSchemas( - config: RegistryConfig, - now: Date = new Date(), -): boolean { - return Object.values(config.servers).some((server) => - isSchemaRefreshStale(server, now), - ) -} - -export async function startSchemaRefreshWorkerIfNeeded( - config: RegistryConfig, - mainPath: string, -): Promise { - if (process.env[WORKER_ENV]) return - if (!hasStaleSchemas(config)) return - if (await isLockActive()) return - - const subprocess = Bun.spawn([process.execPath, mainPath], { - env: { - ...process.env, - [WORKER_ENV]: '1', - }, - stdin: 'ignore', - stdout: 'ignore', - stderr: 'ignore', - }) - subprocess.unref() -} - -export function shouldRunSchemaRefreshWorker(): boolean { - return process.env[WORKER_ENV] === '1' -} - -export async function runSchemaRefreshWorker(): Promise { - await withRefreshLock(async () => { - const initialConfig = await readRegistryConfig() - const staleNames = Object.entries(initialConfig.servers) - .filter((entry) => isSchemaRefreshStale(entry[1])) - .map((entry) => entry[0]) - - for (const name of staleNames) { - await refreshOneServer(name, { staleOnly: true, interactiveAuth: false }) - } - }) -} - -const DEFAULT_REFRESH_CONCURRENCY = 8 - -export type RefreshProgressEvent = - | { type: 'start'; total: number; names: string[] } - | { type: 'server-start'; name: string; active: string[] } - | { - type: 'server-done' - name: string - result: ServerRefreshResult - completed: number - total: number - active: string[] - } - | { type: 'reauth-start'; name: string; remaining: number; total: number } - | { type: 'reauth-done'; name: string; result: ServerRefreshResult } - | { type: 'complete'; summary: RefreshSummary } - -export type RefreshAllOptions = { - concurrency?: number - interactiveAuth?: boolean - onProgress?: (event: RefreshProgressEvent) => void -} - -export async function refreshAllServers( - options: RefreshAllOptions = {}, -): Promise { - const config = await readRegistryConfig() - const names = Object.keys(config.servers) - const concurrency = Math.max( - 1, - options.concurrency ?? DEFAULT_REFRESH_CONCURRENCY, - ) - const interactiveAuth = options.interactiveAuth ?? true - const onProgress = options.onProgress - - onProgress?.({ type: 'start', total: names.length, names }) - - // Phase 1: refresh schemas concurrently without interactive auth so that - // browser-based OAuth flows do not race. Reauth-required servers are - // handled serially afterwards. - const results = new Map() - const active = new Set() - const queue = [...names] - - async function runOne(name: string): Promise { - active.add(name) - onProgress?.({ type: 'server-start', name, active: [...active] }) - try { - const result = await refreshOneServer(name, { - staleOnly: false, - interactiveAuth: false, - }) - results.set(name, result) - } catch (error) { - results.set(name, { - server: name, - status: 'unreachable', - toolsBefore: 0, - message: errorMessage(error), - }) - } finally { - active.delete(name) - onProgress?.({ - type: 'server-done', - name, - result: results.get(name)!, - completed: results.size, - total: names.length, - active: [...active], - }) - } - } - - const workers: Promise[] = [] - const worker = async () => { - while (queue.length > 0) { - const name = queue.shift() - if (!name) return - await runOne(name) - } - } - for (let i = 0; i < Math.min(concurrency, names.length); i++) { - workers.push(worker()) - } - await Promise.all(workers) - - // Phase 2: handle servers that require an interactive reauth one at a time. - if (interactiveAuth) { - const reauthQueue = names.filter( - (name) => results.get(name)?.status === 'reauth-required', - ) - for (let i = 0; i < reauthQueue.length; i++) { - const name = reauthQueue[i]! - onProgress?.({ - type: 'reauth-start', - name, - remaining: reauthQueue.length - i, - total: reauthQueue.length, - }) - let result: ServerRefreshResult - try { - result = await refreshOneServer(name, { - staleOnly: false, - interactiveAuth: true, - }) - } catch (error) { - result = { - server: name, - status: 'unreachable', - toolsBefore: results.get(name)?.toolsBefore ?? 0, - message: errorMessage(error), - } - } - results.set(name, result) - onProgress?.({ type: 'reauth-done', name, result }) - } - } - - const ordered = names - .map((name) => results.get(name)) - .filter((result): result is ServerRefreshResult => Boolean(result)) - const summary = buildRefreshSummary(ordered) - onProgress?.({ type: 'complete', summary }) - return summary -} - -async function refreshOneServer( - name: string, - options: { staleOnly: boolean; interactiveAuth: boolean }, -): Promise { - const beforeRefresh = await readRegistryConfig() - const server = beforeRefresh.servers[name] - if (!server) { - return { - server: name, - status: 'unreachable', - toolsBefore: 0, - message: 'Server was removed before refresh started.', - } - } - - const toolsBefore = server.tools?.length ?? 0 - if (options.staleOnly && !isSchemaRefreshStale(server)) { - return { - server: name, - status: 'schema-refreshed', - toolsBefore, - toolsAfter: toolsBefore, - schemaChanged: false, - } - } - - const authResult = await ensureAuthReady( - name, - server, - toolsBefore, - options.interactiveAuth, - ) - if (authResult.status !== 'ready') return authResult.result - - const readyServer = authResult.server - let refreshed: ServerConfig - try { - refreshed = await refreshServer(readyServer, name) - } catch (error) { - return writeFailureResult(name, readyServer, toolsBefore, error) - } - const beforeWrite = await readRegistryConfig() - const current = beforeWrite.servers[name] - if (!current || !isSameRefreshTarget(readyServer, current)) { - return { - server: name, - status: 'unreachable', - toolsBefore, - message: 'Server changed before refresh completed.', - } - } - - const toolsAfter = refreshed.tools?.length ?? 0 - const schemaChanged = - JSON.stringify(readyServer.tools ?? []) !== - JSON.stringify(refreshed.tools ?? []) - const status = - authResult.authStatus === 'none' - ? 'schema-refreshed' - : authResult.authStatus - const result: ServerRefreshResult = { - server: name, - status, - toolsBefore, - toolsAfter, - schemaChanged, - } - - // Refreshing can take longer than a foreground registry edit. Re-read before - // writing so a background worker does not resurrect a removed server. - beforeWrite.servers[name] = { - ...refreshed, - refreshStatus: refreshStatusFromResult(result), - } - await writeRegistryConfig(beforeWrite) - return result -} - -type AuthReadyResult = - | { - status: 'ready' - server: ServerConfig - authStatus: 'none' | 'auth-refreshed' | 'reauthenticated' - } - | { status: 'blocked'; result: ServerRefreshResult } - -async function ensureAuthReady( - name: string, - server: ServerConfig, - toolsBefore: number, - interactiveAuth: boolean, -): Promise { - if (server.transport === 'stdio') { - return { status: 'ready', server, authStatus: 'none' } - } - - if (server.auth.kind === 'oauth') { - if (!interactiveAuth) { - const result: ServerRefreshResult = { - server: name, - status: 'reauth-required', - toolsBefore, - message: - 'OAuth authentication is required. Run mcpx @refresh in an interactive shell.', - } - await writeRefreshStatus(name, server, refreshStatusFromResult(result)) - return { status: 'blocked', result } - } - return reauthenticateOneServer(name, server, toolsBefore) - } - - const tokenBefore = await readOAuthToken(server) - let headers: Record - try { - headers = await resolveHeaders(server) - } catch (error) { - if (interactiveAuth && isReauthRequiredMessage(errorMessage(error))) { - return reauthenticateOneServer(name, server, toolsBefore) - } - const result = await writeFailureResult(name, server, toolsBefore, error) - return { status: 'blocked', result } - } - const auth = await discoverAuth(new URL(server.url), headers) - if (auth.kind === 'oauth') { - if (interactiveAuth) - return reauthenticateOneServer(name, server, toolsBefore) - const result: ServerRefreshResult = { - server: name, - status: 'reauth-required', - toolsBefore, - message: - 'OAuth authentication is required. Run mcpx @refresh in an interactive shell.', - } - await writeRefreshStatus(name, server, refreshStatusFromResult(result)) - return { status: 'blocked', result } - } - if (auth.kind === 'unknown') { - const result: ServerRefreshResult = { - server: name, - status: 'unreachable', - toolsBefore, - message: auth.reason, - } - await writeRefreshStatus(name, server, refreshStatusFromResult(result)) - return { status: 'blocked', result } - } - const tokenAfter = await readOAuthToken(server) - const authChanged = tokenChanged(tokenBefore, tokenAfter) - const currentConfig = await readRegistryConfig() - const current = currentConfig.servers[name] - return { - status: 'ready', - server: current && isSameRefreshTarget(server, current) ? current : server, - authStatus: authChanged ? 'auth-refreshed' : 'none', - } -} - -async function reauthenticateOneServer( - name: string, - server: ServerConfig, - toolsBefore: number, -): Promise { - let reauthenticated: ServerConfig - try { - reauthenticated = await reauthenticateServer(name, server) - } catch (error) { - const message = error instanceof Error ? error.message : String(error) - const result: ServerRefreshResult = { - server: name, - status: isReauthRequiredMessage(message) - ? 'reauth-required' - : 'unreachable', - toolsBefore, - message, - } - await writeRefreshStatus(name, server, refreshStatusFromResult(result)) - return { status: 'blocked', result } - } - - const beforeWrite = await readRegistryConfig() - const current = beforeWrite.servers[name] - if (!current || !isSameRefreshTarget(server, current)) { - const result: ServerRefreshResult = { - server: name, - status: 'unreachable', - toolsBefore, - message: 'Server changed before re-authentication completed.', - } - return { status: 'blocked', result } - } - - beforeWrite.servers[name] = { - ...reauthenticated, - refreshStatus: { - checkedAt: new Date().toISOString(), - status: 'ok', - }, - } - await writeRegistryConfig(beforeWrite) - return { - status: 'ready', - server: reauthenticated, - authStatus: 'reauthenticated', - } -} - -async function writeFailureResult( - name: string, - server: ServerConfig, - toolsBefore: number, - error: unknown, -): Promise { - const message = errorMessage(error) - const result: ServerRefreshResult = { - server: name, - status: isReauthRequiredMessage(message) - ? 'reauth-required' - : 'unreachable', - toolsBefore, - message, - } - await writeRefreshStatus(name, server, refreshStatusFromResult(result)) - return result -} - -function isSameRefreshTarget(left: ServerConfig, right: ServerConfig): boolean { - if (left.transport === 'stdio' || right.transport === 'stdio') { - return ( - left.transport === 'stdio' && - right.transport === 'stdio' && - left.command === right.command && - JSON.stringify(left.args ?? []) === JSON.stringify(right.args ?? []) && - JSON.stringify(left.env ?? {}) === JSON.stringify(right.env ?? {}) - ) - } - - return ( - left.url === right.url && - JSON.stringify(left.headers ?? null) === - JSON.stringify(right.headers ?? null) && - JSON.stringify(left.auth) === JSON.stringify(right.auth) - ) -} - -async function writeRefreshStatus( - name: string, - server: ServerConfig, - status: ServerRefreshStatus, -): Promise { - const config = await readRegistryConfig() - const current = config.servers[name] - if (!current || !isSameRefreshTarget(server, current)) return - config.servers[name] = { - ...current, - refreshStatus: status, - } - await writeRegistryConfig(config) -} - -function refreshStatusFromResult( - result: ServerRefreshResult, -): ServerRefreshStatus { - const status = - result.status === 'reauth-required' - ? 'reauth-required' - : result.status === 'unreachable' - ? 'unreachable' - : 'ok' - const refreshStatus: ServerRefreshStatus = { - checkedAt: new Date().toISOString(), - status, - } - if (result.message) refreshStatus.message = result.message - return refreshStatus -} - -export function buildRefreshSummary( - results: ServerRefreshResult[], -): RefreshSummary { - return { - checkedAt: new Date().toISOString(), - refreshed: results - .filter( - (result) => - (result.status === 'schema-refreshed' || - result.status === 'auth-refreshed' || - result.status === 'reauthenticated') && - result.schemaChanged === true, - ) - .map((result) => result.server), - unchanged: results - .filter( - (result) => - (result.status === 'schema-refreshed' || - result.status === 'auth-refreshed' || - result.status === 'reauthenticated') && - result.schemaChanged === false, - ) - .map((result) => result.server), - authRefreshed: results - .filter((result) => result.status === 'auth-refreshed') - .map((result) => result.server), - reauthenticated: results - .filter((result) => result.status === 'reauthenticated') - .map((result) => result.server), - reauthRequired: results - .filter((result) => result.status === 'reauth-required') - .map((result) => result.server), - unreachable: results - .filter((result) => result.status === 'unreachable') - .map((result) => result.server), - servers: results, - } -} - -async function readOAuthToken( - server: ServerConfig, -): Promise { - if (server.transport === 'stdio') return undefined - if (server.auth.kind !== 'oauth-token') return undefined - return getOAuthToken(server.auth.tokenKey) -} - -function tokenChanged( - left: OAuthToken | undefined, - right: OAuthToken | undefined, -): boolean { - if (!left || !right) return false - return JSON.stringify(left) !== JSON.stringify(right) -} - -export function isReauthRequiredMessage(message: string): boolean { - const normalized = message.toLowerCase() - return ( - message.includes('Run mcpx @add again') || - message.includes('OAuth token refresh failed') || - message.includes('OAuth client secret is missing') || - normalized.includes('invalid_token') || - normalized.includes('unauthorized') || - normalized.includes('http 401') - ) -} - -function errorMessage(error: unknown): string { - return error instanceof Error ? error.message : String(error) -} - -async function withRefreshLock( - callback: () => Promise, -): Promise { - await clearStaleLock() - let lock: fs.FileHandle | undefined - try { - await fs.mkdir(path.dirname(LOCK_PATH), { recursive: true }) - lock = await fs.open(LOCK_PATH, 'wx') - await lock.writeFile( - JSON.stringify( - { - pid: process.pid, - startedAt: new Date().toISOString(), - }, - null, - 2, - ), - 'utf8', - ) - return await callback() - } catch (error) { - if ((error as NodeJS.ErrnoException).code === 'EEXIST') return undefined - throw error - } finally { - await lock?.close() - if (lock) { - await fs.rm(LOCK_PATH, { force: true }) - } - } -} - -async function isLockActive(): Promise { - await clearStaleLock() - try { - await fs.access(LOCK_PATH) - return true - } catch (error) { - if ((error as NodeJS.ErrnoException).code === 'ENOENT') return false - return true - } -} - -async function clearStaleLock(): Promise { - try { - const stat = await fs.stat(LOCK_PATH) - if (Date.now() - stat.mtimeMs > STALE_LOCK_AFTER_MS) { - await fs.rm(LOCK_PATH, { force: true }) - } - } catch (error) { - if ((error as NodeJS.ErrnoException).code !== 'ENOENT') throw error - } -} diff --git a/src/skill-command.ts b/src/skill-command.ts index cfa3bb7..9479acb 100644 --- a/src/skill-command.ts +++ b/src/skill-command.ts @@ -1,6 +1,6 @@ import { cancel, isCancel, multiselect } from '@clack/prompts' -import type { ProjectService } from './project-service' +import type { RegistryConfig } from './types' import { buildMcpxSkillMarkdown, @@ -13,12 +13,14 @@ export type SkillCommandInput = { show?: string } +export type RegistryView = Pick + export async function runSkillCommand( - service: ProjectService, + registry: RegistryView, cwd: string, input: SkillCommandInput, ): Promise { - const availableServers = Object.keys(service.config.servers).sort() + const availableServers = Object.keys(registry.servers).sort() if (availableServers.length === 0) { throw new Error( 'No MCP servers are registered. Run "mcpx @add --name --url " first.', diff --git a/src/token-cache.ts b/src/token-cache.ts deleted file mode 100644 index 07d4897..0000000 --- a/src/token-cache.ts +++ /dev/null @@ -1,112 +0,0 @@ -import fs from 'node:fs/promises' -import { homedir } from 'node:os' -import path from 'node:path' - -import type { OAuthToken, TokenCache } from './types' - -const TOKEN_CACHE_PATH = path.join('.agents', 'mcpx', 'tokens.json') - -export function getTokenCachePath(): string { - return path.join(homedir(), TOKEN_CACHE_PATH) -} - -export async function readTokenCache(): Promise { - const filePath = getTokenCachePath() - try { - const raw = await fs.readFile(filePath, 'utf8') - const parsed = JSON.parse(raw) as TokenCache - if ( - parsed.version !== 1 || - !parsed.oauth || - typeof parsed.oauth !== 'object' - ) { - throw new Error(`Invalid mcpx token cache at ${filePath}.`) - } - return { - version: 1, - oauth: parsed.oauth, - oauthClientSecrets: parsed.oauthClientSecrets ?? {}, - } - } catch (error) { - if ((error as NodeJS.ErrnoException).code === 'ENOENT') { - return { version: 1, oauth: {}, oauthClientSecrets: {} } - } - throw error - } -} - -export async function writeTokenCache(cache: TokenCache): Promise { - const filePath = getTokenCachePath() - await fs.mkdir(path.dirname(filePath), { recursive: true, mode: 0o700 }) - await fs.writeFile(filePath, `${JSON.stringify(cache, null, 2)}\n`, { - encoding: 'utf8', - mode: 0o600, - }) -} - -export async function getOAuthToken( - tokenKey: string, -): Promise { - const cache = await readTokenCache() - return cache.oauth[tokenKey] -} - -export async function getOAuthTokenForUpdate(tokenKey: string): Promise<{ - cache: TokenCache - token: OAuthToken | undefined -}> { - const cache = await readTokenCache() - return { cache, token: cache.oauth[tokenKey] } -} - -export async function putOAuthTokenWithClientSecret( - tokenKey: string, - token: OAuthToken, - clientSecret?: string, -): Promise { - const cache = await readTokenCache() - cache.oauth[tokenKey] = token - if (clientSecret && token.clientSecretKey) { - cache.oauthClientSecrets ??= {} - cache.oauthClientSecrets[token.clientSecretKey] = clientSecret - } - await writeTokenCache(cache) -} - -export async function getOAuthClientSecret( - secretKey: string, -): Promise { - const cache = await readTokenCache() - return cache.oauthClientSecrets?.[secretKey] -} - -export async function putOAuthTokenInCache( - cache: TokenCache, - tokenKey: string, - token: OAuthToken, -): Promise { - cache.oauth[tokenKey] = token - await writeTokenCache(cache) -} - -export async function removeOAuthToken(tokenKey: string): Promise { - const cache = await readTokenCache() - const removed = removeOAuthTokenFromCache(cache, tokenKey) - if (removed) { - await writeTokenCache(cache) - } - return removed -} - -export function removeOAuthTokenFromCache( - cache: TokenCache, - tokenKey: string, -): boolean { - const token = cache.oauth[tokenKey] - if (!token) return false - if (token.clientSecretKey && cache.oauthClientSecrets) { - delete cache.oauthClientSecrets[token.clientSecretKey] - } - delete cache.oauth[tokenKey] - return true -} diff --git a/tests/authentication-coordinator.test.ts b/tests/authentication-coordinator.test.ts new file mode 100644 index 0000000..710fc79 --- /dev/null +++ b/tests/authentication-coordinator.test.ts @@ -0,0 +1,172 @@ +import { describe, expect, it } from 'bun:test' + +import { AuthenticationCoordinator } from '../src/authentication-coordinator' +import { createInMemoryRuntimeCaller } from '../src/runtime-caller' + +describe('Runtime Authentication Flow coordinator', () => { + it('shares one flow and one persistence write across five callers', async () => { + const coordinator = new AuthenticationCoordinator() + const callers = Array.from({ length: 5 }, (_, index) => + createInMemoryRuntimeCaller(`refresh-${index}`), + ) + let starts = 0 + let writes = 0 + + const outcomes = await Promise.all( + callers.map((caller) => + coordinator.join('oauth:fixture', caller, { + start: async () => { + starts += 1 + await Bun.sleep(5) + return { accessToken: 'internal-secret' } + }, + persist: async () => { + writes += 1 + }, + }), + ), + ) + + expect(starts).toBe(1) + expect(writes).toBe(1) + expect(outcomes).toEqual(callers.map(() => ({ status: 'completed' }))) + expect(callers.flatMap((caller) => caller.frames)).toEqual([]) + }) + + it('keeps the shared flow alive while one waiter remains', async () => { + const coordinator = new AuthenticationCoordinator() + const first = createInMemoryRuntimeCaller('first') + const second = createInMemoryRuntimeCaller('second') + let flowSignal: AbortSignal | undefined + let finish = (_value: string) => {} + const result = new Promise((resolve) => { + finish = resolve + }) + const flow = { + start: async (signal: AbortSignal) => { + flowSignal = signal + return result + }, + persist: async () => {}, + } + const firstRefresh = coordinator.join('oauth:fixture', first, flow) + const secondRefresh = coordinator.join('oauth:fixture', second, flow) + await waitFor(() => flowSignal !== undefined) + + first.disconnect() + expect(flowSignal?.aborted).toBe(false) + finish('rotated') + expect(await firstRefresh).toEqual({ status: 'disconnected' }) + expect(await secondRefresh).toEqual({ status: 'completed' }) + + expect(first.frames).toEqual([]) + expect(second.frames).toEqual([]) + }) + + it('aborts and settles the flow when its final waiter disconnects', async () => { + const coordinator = new AuthenticationCoordinator() + const caller = createInMemoryRuntimeCaller('only-waiter') + let callbackOpen = false + const refresh = coordinator.join('oauth:fixture', caller, { + start: (signal) => + new Promise((_resolve, reject) => { + callbackOpen = true + signal.addEventListener( + 'abort', + () => { + callbackOpen = false + reject(signal.reason) + }, + { once: true }, + ) + }), + persist: async () => {}, + }) + await waitFor(() => callbackOpen) + + caller.disconnect() + expect(await refresh).toEqual({ status: 'disconnected' }) + + expect(callbackOpen).toBe(false) + expect(caller.frames).toEqual([]) + expect(coordinator.activeFlows()).toBe(0) + }) + + it('clears owned timeout resources after timeout and rejection', async () => { + const coordinator = new AuthenticationCoordinator({ timeoutMs: 5 }) + const timedOut = createInMemoryRuntimeCaller('timed-out') + let timeoutResourceOpen = false + const timeout = coordinator.join('oauth:timeout', timedOut, { + start: (signal) => + new Promise((_resolve, reject) => { + timeoutResourceOpen = true + signal.addEventListener( + 'abort', + () => { + timeoutResourceOpen = false + reject(signal.reason) + }, + { once: true }, + ) + }), + persist: async () => {}, + }) + expect(timeout).rejects.toMatchObject({ code: 'timeout' }) + await timeout.catch(() => {}) + + const rejected = createInMemoryRuntimeCaller('rejected') + const rejection = coordinator.join('oauth:rejected', rejected, { + start: () => { + throw new Error('provider rejected') + }, + persist: async () => {}, + }) + expect(rejection).rejects.toMatchObject({ + code: 'operation-failed', + message: 'provider rejected', + }) + await rejection.catch(() => {}) + + expect(timeoutResourceOpen).toBe(false) + expect(timedOut.frames).toEqual([]) + expect(rejected.frames).toEqual([]) + expect(coordinator.activeFlows()).toBe(0) + }) + + it('cancels and awaits every flow before refusing shutdown-era admission', async () => { + const coordinator = new AuthenticationCoordinator() + const caller = createInMemoryRuntimeCaller('shutdown') + let started = false + const waiting = coordinator.join('credential', caller, { + start: async (signal) => { + started = true + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + }, + persist: async () => {}, + }) + await waitFor(() => started) + + await coordinator.close() + + await expect(waiting).rejects.toThrow('stopping') + expect(coordinator.activeFlows()).toBe(0) + await expect( + coordinator.join('late', createInMemoryRuntimeCaller('late'), { + start: async () => 'late', + persist: async () => {}, + }), + ).rejects.toThrow('stopping') + }) +}) + +async function waitFor(predicate: () => boolean): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (predicate()) return + await Bun.sleep(1) + } + throw new Error('Condition was not met.') +} diff --git a/tests/config.test.ts b/tests/config.test.ts deleted file mode 100644 index 2fc1d48..0000000 --- a/tests/config.test.ts +++ /dev/null @@ -1,114 +0,0 @@ -import { describe, expect, it } from 'bun:test' - -import type { RegistryConfig } from '../src/types' - -import { normalizeRegistryConfig, removeServerFromConfig } from '../src/config' - -describe('registry config', () => { - it('derives command names for cached tool schemas', () => { - const config = { - version: 1, - servers: { - posthog: { - url: 'https://mcp.posthog.com/mcp', - auth: { kind: 'none' }, - tools: [ - { - name: 'alert.create', - description: 'Create an alert', - }, - ], - }, - }, - } as unknown as RegistryConfig - - expect( - normalizeRegistryConfig(config).servers.posthog?.tools?.[0]?.commandName, - ).toBe('alert-create') - }) - - it('preserves cached MCP tool metadata', () => { - const config = { - version: 1, - servers: { - browser: { - url: 'https://browser.example/mcp', - auth: { kind: 'none' }, - tools: [ - { - name: 'close_page', - commandName: 'stale', - title: 'Close Page', - annotations: { destructiveHint: true }, - _meta: { source: 'server' }, - }, - ], - }, - }, - } as unknown as RegistryConfig - - expect( - normalizeRegistryConfig(config).servers.browser?.tools?.[0], - ).toMatchObject({ - commandName: 'close_page', - title: 'Close Page', - annotations: { destructiveHint: true }, - _meta: { source: 'server' }, - }) - }) - - it('drops stale cached output schemas', () => { - const config = { - version: 1, - servers: { - browser: { - url: 'https://browser.example/mcp', - auth: { kind: 'none' }, - tools: [ - { - name: 'list_pages', - outputSchema: { type: 'object' }, - }, - ], - }, - }, - } as unknown as RegistryConfig - - expect( - normalizeRegistryConfig(config).servers.browser?.tools?.[0], - ).not.toHaveProperty('outputSchema') - }) - - it('removes a server from registry config', () => { - const config = { - version: 1, - servers: { - posthog: { - url: 'https://mcp.posthog.com/mcp', - auth: { - kind: 'oauth-token', - tokenKey: 'posthog', - confidence: 'confirmed', - }, - }, - }, - } as RegistryConfig - - const removed = removeServerFromConfig(config, 'posthog') - expect(removed?.transport).not.toBe('stdio') - expect( - removed && removed.transport !== 'stdio' ? removed.url : undefined, - ).toBe('https://mcp.posthog.com/mcp') - expect(config.servers.posthog).toBeUndefined() - }) - - it('keeps registry config unchanged when removing an unknown server', () => { - const config = { - version: 1, - servers: {}, - } as RegistryConfig - - expect(removeServerFromConfig(config, 'posthog')).toBeUndefined() - expect(config.servers).toEqual({}) - }) -}) diff --git a/tests/daemon-client.test.ts b/tests/daemon-client.test.ts index 3ece6d2..d0f1a6e 100644 --- a/tests/daemon-client.test.ts +++ b/tests/daemon-client.test.ts @@ -1,457 +1,250 @@ import { afterEach, beforeEach, describe, expect, it } from 'bun:test' import fs from 'node:fs/promises' -import http from 'node:http' import net from 'node:net' import { tmpdir } from 'node:os' import path from 'node:path' -import type { StdioServerConfig } from '../src/types' - -import { daemonStatus, stopDaemon } from '../src/daemon-client' +import { connectDaemonSocket } from '../src/daemon-client' import { requestJsonLine, writeJsonLine } from '../src/daemon-io' import { daemonSocketPath } from '../src/daemon-paths' -import { buildServerKey, helloMessage } from '../src/daemon-protocol' -import { callMcpTool, listMcpTools } from '../src/mcp-client' +import { helloMessage } from '../src/daemon-protocol' +import { requestRuntime } from '../src/runtime-client' const mainPath = path.join(import.meta.dir, '..', 'src', 'main.ts') -let previousHome: string | undefined -let previousMcpxHome: string | undefined -let previousDisableDaemon: string | undefined -let previousArgvOne: string | undefined -let home: string -let daemon: ReturnType | undefined -let fakeDaemon: net.Server | undefined -let httpFixture: http.Server | undefined +describe('MCP Runtime process client', () => { + let previousHome: string | undefined + let previousMcpxHome: string | undefined + let home: string + let fakeDaemon: net.Server | undefined -describe('mcpxd daemon client', () => { beforeEach(async () => { previousHome = process.env.HOME previousMcpxHome = process.env.MCPX_HOME - previousDisableDaemon = process.env.MCPX_DISABLE_DAEMON - previousArgvOne = process.argv[1] - home = await fs.mkdtemp(path.join(tmpdir(), 'mcpxd-test-')) + home = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-client-')) process.env.HOME = home process.env.MCPX_HOME = home - delete process.env.MCPX_DISABLE_DAEMON }) afterEach(async () => { - await stopDaemon(mainPath).catch(() => {}) + await requestRuntime( + { requestId: crypto.randomUUID(), op: 'stop' }, + mainPath, + { start: false }, + ).catch(() => {}) await stopFakeDaemon() - await stopHttpFixture() - daemon?.kill() - daemon = undefined await fs.rm(home, { recursive: true, force: true }) - if (previousHome === undefined) { - delete process.env.HOME - } else { - process.env.HOME = previousHome - } - if (previousDisableDaemon === undefined) { - delete process.env.MCPX_DISABLE_DAEMON - } else { - process.env.MCPX_DISABLE_DAEMON = previousDisableDaemon - } - if (previousMcpxHome === undefined) { - delete process.env.MCPX_HOME - } else { - process.env.MCPX_HOME = previousMcpxHome - } - if (previousArgvOne === undefined) { - process.argv.splice(1, 1) - } else { - process.argv[1] = previousArgvOne - } + if (previousHome === undefined) delete process.env.HOME + else process.env.HOME = previousHome + if (previousMcpxHome === undefined) delete process.env.MCPX_HOME + else process.env.MCPX_HOME = previousMcpxHome }) - it('starts mcpxd on demand for stdio calls', async () => { - process.argv[1] = mainPath - const server = fixtureServer() + it('starts the Runtime on demand with a private socket', async () => { + const status = (await requestRuntime( + { requestId: crypto.randomUUID(), op: 'status' }, + mainPath, + )) as { protocolVersion: number; activeServers: number } - expect( - (await listMcpTools(server, 'fixture')).map((tool) => tool.name), - ).toContain('echo') - expect((await daemonStatus(mainPath)).activeServers).toBe(1) + expect(status.protocolVersion).toBe(3) + expect(status.activeServers).toBe(0) expect((await fs.stat(daemonSocketPath())).mode & 0o777).toBe(0o600) }) - it('reuses a stdio server process across separate calls', async () => { - await startDaemon() - const server = fixtureServer() - - expect( - (await listMcpTools(server, 'fixture')).map((tool) => tool.name), - ).toContain('pid') - const firstPid = text(await callMcpTool(server, 'pid', {}, 'fixture')) - expect(text(await callMcpTool(server, 'increment', {}, 'fixture'))).toBe( - '1', - ) - expect(text(await callMcpTool(server, 'increment', {}, 'fixture'))).toBe( - '2', - ) - const secondPid = text(await callMcpTool(server, 'pid', {}, 'fixture')) - - expect(secondPid).toBe(firstPid) - expect((await daemonStatus(mainPath)).activeServers).toBe(1) - }) - - it('passes cwd to managed stdio processes', async () => { - await startDaemon() - const cwd = await fs.mkdtemp(path.join(home, 'cwd-')) - const server = fixtureServer({ cwd }) - - expect( - await fs.realpath(text(await callMcpTool(server, 'cwd', {}, 'fixture'))), - ).toBe(await fs.realpath(cwd)) - }) - - it('routes HTTP servers through mcpxd and preserves session ids', async () => { - await startDaemon() - const fixture = await startHttpFixture() - const server = { - url: fixture.url, - auth: { kind: 'none' as const }, - } + it('stops an older protocol daemon before starting v3', async () => { + await startFakeDaemon({ protocolVersion: 2, version: '0.0.0' }) - expect( - (await listMcpTools(server, 'http-fixture')).map((tool) => tool.name), - ).toEqual(['echo']) - expect(text(await callMcpTool(server, 'echo', {}, 'http-fixture'))).toBe( - 'http-ok', - ) + const status = (await requestRuntime( + { requestId: crypto.randomUUID(), op: 'status' }, + mainPath, + )) as { protocolVersion: number } - expect( - fixture.sessions.slice(1).every((session) => session === 'session-1'), - ).toBe(true) - expect((await daemonStatus(mainPath)).servers[0]).toMatchObject({ - transport: 'http', - url: expect.stringContaining('/mcp'), - }) + expect(status.protocolVersion).toBe(3) + expect(fakeDaemon).toBeUndefined() }) - it('retries HTTP session creation after an initial connection failure', async () => { - await startDaemon() - const port = await freePort() - const server = { - url: `http://127.0.0.1:${port}/mcp`, - auth: { kind: 'none' as const }, - } + it('stops an older same-protocol Runtime before restart', async () => { + await startFakeDaemon({ protocolVersion: 3, version: '0.0.0' }) - await expect(listMcpTools(server, 'http-flaky')).rejects.toThrow() + const status = (await requestRuntime( + { requestId: crypto.randomUUID(), op: 'status' }, + mainPath, + )) as { version: string } - const fixture = await startHttpFixture(port) - expect( - (await listMcpTools(server, 'http-flaky')).map((tool) => tool.name), - ).toEqual(['echo']) - expect(fixture.sessions.length).toBeGreaterThan(0) + expect(status.version).not.toBe('0.0.0') + expect(fakeDaemon).toBeUndefined() }) - it('drops retained HTTP session ids after invalid api key rejection', async () => { - await startDaemon() - const fixture = await startHttpFixture(undefined, { - rejectRetainedCallOnce: true, - }) - const server = { - transport: 'http' as const, - url: fixture.url, - auth: { kind: 'none' as const }, - } - - expect(text(await callMcpTool(server, 'echo', {}, 'http-fixture'))).toBe( - 'http-ok', - ) - await evictDaemonSession(buildServerKey(server)) - expect(text(await callMcpTool(server, 'echo', {}, 'http-fixture'))).toBe( - 'http-ok', + it('rejects a second operation on the same connection', async () => { + await requestRuntime( + { requestId: crypto.randomUUID(), op: 'status' }, + mainPath, ) - - expect(fixture.sessions).toContain('session-1') - expect(fixture.sessions).toContain('session-2') - }) - - it('drops retained HTTP session ids for auth-refreshed eviction', async () => { - await startDaemon() - const fixture = await startHttpFixture(undefined, { - rejectRetainedCallOnce: true, - }) - const server = { - transport: 'http' as const, - url: fixture.url, - auth: { kind: 'none' as const }, + const socket = await connectDaemonSocket() + try { + expect(await requestJsonLine(socket, helloMessage())).toMatchObject({ + ok: true, + protocolVersion: 3, + }) + expect( + await requestJsonLine(socket, { requestId: 'first', op: 'status' }), + ).toMatchObject({ requestId: 'first', kind: 'result' }) + expect( + await requestJsonLine(socket, { requestId: 'second', op: 'status' }), + ).toMatchObject({ + ok: false, + error: { code: 'connection-complete' }, + }) + } finally { + socket.destroy() } + }) - expect(text(await callMcpTool(server, 'echo', {}, 'http-fixture'))).toBe( - 'http-ok', + it('returns a terminal error when Runtime state loading fails', async () => { + await requestRuntime( + { requestId: crypto.randomUUID(), op: 'status' }, + mainPath, ) - const sessionsBeforeEvict = fixture.sessions.length - await evictDaemonSession(buildServerKey(server), 'auth-refreshed') - expect(text(await callMcpTool(server, 'echo', {}, 'http-fixture'))).toBe( - 'http-ok', + const registryPath = path.join( + home, + '.agents', + 'mcpx', + 'state-v2', + 'registry.json', ) - - expect(fixture.sessions.slice(sessionsBeforeEvict)).not.toContain( - 'session-1', - ) - expect(fixture.sessions.slice(sessionsBeforeEvict)).toContain('session-2') - }) - - it('stops an incompatible daemon before starting a compatible one', async () => { - await startFakeDaemon() - process.argv[1] = mainPath - - expect( - text(await callMcpTool(fixtureServer(), 'echo', {}, 'fixture')), - ).toBe('ok') - expect(fakeDaemon).toBeUndefined() - expect((await daemonStatus(mainPath)).protocolVersion).toBe(2) - }) - - it('stops an older same-protocol daemon before starting the current version', async () => { - await startFakeDaemon({ protocolVersion: 2, version: '0.0.0' }) - process.argv[1] = mainPath - - expect( - text(await callMcpTool(fixtureServer(), 'echo', {}, 'fixture')), - ).toBe('ok') - expect(fakeDaemon).toBeUndefined() - expect((await daemonStatus(mainPath)).version).not.toBe('0.0.0') - }) - - it('reports and stops the daemon', async () => { - await startDaemon() - const server = fixtureServer() - - await callMcpTool(server, 'echo', {}, 'fixture') - expect((await daemonStatus(mainPath)).servers[0]?.labels).toEqual([ - 'fixture', - ]) - expect(await stopDaemon(mainPath)).toEqual({ stopping: true }) - await waitForStopped() - }) -}) - -async function startDaemon(): Promise { - daemon = Bun.spawn([process.execPath, mainPath, '@daemon', 'server'], { - env: { - ...process.env, - MCPX_DAEMON_SERVER: '1', - }, - stdin: 'ignore', - stdout: 'ignore', - stderr: 'ignore', - }) - const deadline = Date.now() + 3_000 - while (Date.now() < deadline) { + const original = await fs.readFile(registryPath, 'utf8') + await fs.writeFile(registryPath, '{invalid') try { - await daemonStatus(mainPath) - return - } catch { - await new Promise((resolve) => setTimeout(resolve, 50)) + await expect( + requestRuntime( + { requestId: 'broken-state', op: 'registrySnapshot' }, + mainPath, + ), + ).rejects.toThrow() + } finally { + await fs.writeFile(registryPath, original) } - } - throw new Error('daemon did not start') -} - -function fixtureServer( - overrides: Partial = {}, -): StdioServerConfig { - return { - transport: 'stdio', - command: process.execPath, - args: [path.join(import.meta.dir, 'fixtures', 'stdio-server.mjs')], - ...overrides, - } -} - -function text(result: unknown): string { - const content = (result as { content?: { text?: string }[] }).content - return content?.[0]?.text ?? '' -} - -async function waitForStopped(): Promise { - const deadline = Date.now() + 3_000 - while (Date.now() < deadline) { - try { - await daemonStatus(mainPath) - } catch { - return - } - await new Promise((resolve) => setTimeout(resolve, 50)) - } - throw new Error('daemon did not stop') -} - -async function startFakeDaemon( - options: { protocolVersion?: number; version?: string } = {}, -): Promise { - const protocolVersion = options.protocolVersion ?? 0 - const version = options.version ?? 'old' - await fs.mkdir(path.dirname(daemonSocketPath()), { recursive: true }) - fakeDaemon = net.createServer((socket) => { - socket.on('data', (chunk) => { - const text = chunk.toString('utf8') - if (text.includes('"op":"hello"')) { - writeJsonLine(socket, { - ok: true, - protocolVersion, - result: { protocolVersion, version }, - }) - } - if (text.includes('"op":"stop"')) { - writeJsonLine(socket, { ok: true, result: { stopping: true } }) - void stopFakeDaemon() - } - }) }) - await new Promise((resolve, reject) => { - fakeDaemon?.once('error', reject) - fakeDaemon?.listen(daemonSocketPath(), () => { - fakeDaemon?.off('error', reject) - resolve() - }) - }) -} -async function stopFakeDaemon(): Promise { - const server = fakeDaemon - fakeDaemon = undefined - if (!server) return - await new Promise((resolve) => server.close(() => resolve())) - await fs.rm(daemonSocketPath(), { force: true }).catch(() => {}) -} + it('rejects when a daemon closes before sending a terminal frame', async () => { + await startFakeDaemon({ + protocolVersion: 3, + version: '0.9.15', + closeOnIntent: true, + }) -async function evictDaemonSession( - serverKey: string, - reason: 'auth-refreshed' | 'unauthorized' | 'manual' = 'manual', -): Promise { - const socket = await new Promise((resolve, reject) => { - const connected = net.createConnection(daemonSocketPath()) - connected.once('connect', () => resolve(connected)) - connected.once('error', reject) + await expect( + requestRuntime({ requestId: 'closed', op: 'status' }, mainPath, { + start: false, + }), + ).rejects.toThrow('closed before a terminal frame') }) - try { - await requestJsonLine(socket, helloMessage()) - await requestJsonLine(socket, { op: 'evictSession', serverKey, reason }) - } finally { - socket.end() - } -} - -async function startHttpFixture( - port?: number, - options: { rejectRetainedCallOnce?: boolean } = {}, -): Promise<{ url: string; sessions: (string | null)[] }> { - const sessions: (string | null)[] = [] - let nextSession = 1 - let rejectedRetainedCall = false - let toolCalls = 0 - httpFixture = http.createServer(async (request, response) => { - if (request.method !== 'POST') { - response.writeHead(405).end() - return - } - const body = await readRequestBody(request) - const message = JSON.parse(body) as { id?: string | number; method: string } - const sessionId = request.headers['mcp-session-id']?.toString() ?? null - sessions.push(sessionId) - - if ( - options.rejectRetainedCallOnce && - !rejectedRetainedCall && - toolCalls > 0 && - message.method === 'tools/call' && - sessionId === 'session-1' - ) { - rejectedRetainedCall = true - response - .writeHead(200, { - 'content-type': 'application/json', - 'mcp-session-id': 'session-1', - }) - .end( - JSON.stringify({ - jsonrpc: '2.0', - id: message.id, - error: { code: -32000, message: 'INVALID_API_KEY' }, + it('aborts an active CLI input provider after the terminal frame', async () => { + await startFakeDaemon({ + protocolVersion: 3, + version: '0.9.15', + inputThenTerminal: true, + }) + let inputAborted = false + + await requestRuntime( + { requestId: 'input-cleanup', op: 'status' }, + mainPath, + { + start: false, + onInput: async (_request, signal) => + new Promise((resolve) => { + signal.addEventListener( + 'abort', + () => { + inputAborted = true + resolve({ cancelled: true }) + }, + { once: true }, + ) }), - ) - return - } - if (message.method === 'tools/call') toolCalls += 1 - - if (message.method === 'notifications/initialized') { - response - .writeHead(202, { 'mcp-session-id': sessionId ?? 'session-1' }) - .end() - return - } - - const responseSession = - message.method === 'initialize' - ? `session-${nextSession++}` - : (sessionId ?? 'session-1') + }, + ) - const result = - message.method === 'initialize' - ? { - protocolVersion: '2025-11-25', - capabilities: { tools: {} }, - serverInfo: { name: 'fixture', version: '1.0.0' }, + expect(inputAborted).toBe(true) + }) + + async function startFakeDaemon(options: { + protocolVersion: number + version: string + closeOnIntent?: boolean + inputThenTerminal?: boolean + }): Promise { + await fs.mkdir(path.dirname(daemonSocketPath()), { recursive: true }) + fakeDaemon = net.createServer((socket) => { + let buffer = '' + socket.on('data', (chunk) => { + buffer += chunk.toString('utf8') + while (buffer.includes('\n')) { + const newline = buffer.indexOf('\n') + const message = JSON.parse(buffer.slice(0, newline)) as { + op?: string + requestId?: string } - : message.method === 'tools/list' - ? { tools: [{ name: 'echo', inputSchema: { type: 'object' } }] } - : { content: [{ type: 'text', text: 'http-ok' }] } - - response - .writeHead(200, { - 'content-type': 'application/json', - 'mcp-session-id': responseSession, + buffer = buffer.slice(newline + 1) + if (message.op === 'hello') { + writeJsonLine(socket, { + ok: true, + protocolVersion: options.protocolVersion, + result: options, + }) + } else if (message.op === 'stop') { + writeJsonLine( + socket, + message.requestId + ? { + requestId: message.requestId, + kind: 'result', + result: { stopping: true }, + } + : { ok: true, result: { stopping: true } }, + ) + void stopFakeDaemon() + } else if (options.closeOnIntent) { + socket.destroy() + } else if (options.inputThenTerminal && message.requestId) { + writeJsonLine(socket, { + requestId: message.requestId, + kind: 'event', + event: { + type: 'input-required', + data: { + inputId: 'input-1', + type: 'oauth-client', + serverName: 'fixture', + redirectUri: 'http://127.0.0.1/callback', + issuer: 'http://127.0.0.1', + scopes: [], + }, + }, + }) + writeJsonLine(socket, { + requestId: message.requestId, + kind: 'result', + result: { done: true }, + }) + } + } }) - .end(JSON.stringify({ jsonrpc: '2.0', id: message.id, result })) - }) - - await new Promise((resolve, reject) => { - httpFixture?.once('error', reject) - httpFixture?.listen(port ?? 0, '127.0.0.1', () => { - httpFixture?.off('error', reject) - resolve() }) - }) - const address = httpFixture.address() - if (!address || typeof address === 'string') - throw new Error('HTTP fixture did not bind.') - return { url: `http://127.0.0.1:${address.port}/mcp`, sessions } -} - -async function freePort(): Promise { - const server = http.createServer() - await new Promise((resolve, reject) => { - server.once('error', reject) - server.listen(0, '127.0.0.1', () => { - server.off('error', reject) - resolve() + await new Promise((resolve, reject) => { + fakeDaemon?.once('error', reject) + fakeDaemon?.listen(daemonSocketPath(), resolve) }) - }) - const address = server.address() - await new Promise((resolve) => server.close(() => resolve())) - if (!address || typeof address === 'string') - throw new Error('Free port lookup failed.') - return address.port -} - -async function stopHttpFixture(): Promise { - const server = httpFixture - httpFixture = undefined - if (!server) return - await new Promise((resolve) => server.close(() => resolve())) -} + } -async function readRequestBody(request: http.IncomingMessage): Promise { - const chunks: Buffer[] = [] - for await (const chunk of request) { - chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)) + async function stopFakeDaemon(): Promise { + const server = fakeDaemon + fakeDaemon = undefined + if (server) { + await new Promise((resolve) => server.close(() => resolve())) + } + await fs.rm(daemonSocketPath(), { force: true }).catch(() => {}) } - return Buffer.concat(chunks).toString('utf8') -} +}) diff --git a/tests/daemon-protocol.test.ts b/tests/daemon-protocol.test.ts index c7439d5..6528832 100644 --- a/tests/daemon-protocol.test.ts +++ b/tests/daemon-protocol.test.ts @@ -1,102 +1,21 @@ import { describe, expect, it } from 'bun:test' -import type { HttpServerConfig } from '../src/types' - import { NOTIFICATION_MODE_ENV, - buildServerKey, notificationModeFromEnv, } from '../src/daemon-protocol' -describe('daemon protocol', () => { - it('keeps HTTP server keys stable across resolved token values', () => { - const server: HttpServerConfig = { - url: 'https://mcp.example.com/mcp', - headers: { Authorization: 'Bearer old' }, - auth: { - kind: 'bearer', - credentials: [{ kind: 'env', name: 'MCP_TOKEN' }], - strategy: 'round-robin', - confidence: 'configured', - }, - } - const rotated: HttpServerConfig = { - ...server, - headers: { Authorization: 'Bearer new' }, - } - - expect(buildServerKey(rotated)).toBe(buildServerKey(server)) - }) - - it('isolates HTTP server keys by auth reference', () => { - const first: HttpServerConfig = { - url: 'https://mcp.example.com/mcp', - auth: { - kind: 'bearer', - credentials: [{ kind: 'env', name: 'FIRST_TOKEN' }], - strategy: 'round-robin', - confidence: 'configured', - }, - } - const second: HttpServerConfig = { - ...first, - auth: { - kind: 'bearer', - credentials: [{ kind: 'env', name: 'SECOND_TOKEN' }], - strategy: 'round-robin', - confidence: 'configured', - }, - } - - expect(buildServerKey(second)).not.toBe(buildServerKey(first)) - }) - - it('isolates HTTP server keys by bearer credential list', () => { - const first: HttpServerConfig = { - url: 'https://mcp.example.com/mcp', - auth: { - kind: 'bearer', - credentials: [ - { kind: 'env', name: 'FIRST_TOKEN' }, - { kind: 'env', name: 'SECOND_TOKEN' }, - ], - strategy: 'round-robin', - confidence: 'configured', - }, - } - const second: HttpServerConfig = { - ...first, - auth: { - kind: 'bearer', - credentials: [ - { kind: 'env', name: 'FIRST_TOKEN' }, - { kind: 'env', name: 'THIRD_TOKEN' }, - ], - strategy: 'round-robin', - confidence: 'configured', - }, - } - - expect(buildServerKey(second)).not.toBe(buildServerKey(first)) - }) - +describe('daemon protocol adapter', () => { it('defaults notification buffering unless explicitly discarded', () => { const previous = process.env[NOTIFICATION_MODE_ENV] try { delete process.env[NOTIFICATION_MODE_ENV] expect(notificationModeFromEnv()).toBe('buffer') - - process.env[NOTIFICATION_MODE_ENV] = 'buffer' - expect(notificationModeFromEnv()).toBe('buffer') - process.env[NOTIFICATION_MODE_ENV] = 'discard' expect(notificationModeFromEnv()).toBe('discard') } finally { - if (previous === undefined) { - delete process.env[NOTIFICATION_MODE_ENV] - } else { - process.env[NOTIFICATION_MODE_ENV] = previous - } + if (previous === undefined) delete process.env[NOTIFICATION_MODE_ENV] + else process.env[NOTIFICATION_MODE_ENV] = previous } }) @@ -108,11 +27,8 @@ describe('daemon protocol', () => { 'Invalid MCPX_NOTIFICATION_MODE value "off". Expected "buffer" or "discard".', ) } finally { - if (previous === undefined) { - delete process.env[NOTIFICATION_MODE_ENV] - } else { - process.env[NOTIFICATION_MODE_ENV] = previous - } + if (previous === undefined) delete process.env[NOTIFICATION_MODE_ENV] + else process.env[NOTIFICATION_MODE_ENV] = previous } }) }) diff --git a/tests/headers.test.ts b/tests/headers.test.ts index 0773151..0ff1adc 100644 --- a/tests/headers.test.ts +++ b/tests/headers.test.ts @@ -1,16 +1,9 @@ import { describe, expect, it } from 'bun:test' -import fs from 'node:fs/promises' -import { tmpdir } from 'node:os' -import path from 'node:path' import type { HttpServerConfig } from '../src/types' import { authFromBearerValues } from '../src/bearer' -import { - normalizeAuthScheme, - resolveHeaders, - resolveProbeHeaders, -} from '../src/headers' +import { normalizeAuthScheme, resolveProbeHeaders } from '../src/headers' describe('headers', () => { it('canonicalizes bearer token auth scheme', () => { @@ -61,50 +54,6 @@ describe('headers', () => { } } }) - - it('round-robins bearer credentials across CLI invocations', async () => { - const previousHome = process.env.MCPX_HOME - const previousFirst = process.env.FIRST_MCPX_TEST_TOKEN - const previousSecond = process.env.SECOND_MCPX_TEST_TOKEN - const home = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-bearer-test-')) - process.env.MCPX_HOME = home - process.env.FIRST_MCPX_TEST_TOKEN = 'first' - process.env.SECOND_MCPX_TEST_TOKEN = 'Bearer second' - - try { - const server = bearerServer([ - 'env:FIRST_MCPX_TEST_TOKEN', - 'env:SECOND_MCPX_TEST_TOKEN', - ]) - - await expect(resolveHeaders(server)).resolves.toMatchObject({ - Authorization: 'Bearer first', - }) - await expect(resolveHeaders(server)).resolves.toMatchObject({ - Authorization: 'Bearer second', - }) - await expect(resolveHeaders(server)).resolves.toMatchObject({ - Authorization: 'Bearer first', - }) - } finally { - await fs.rm(home, { recursive: true, force: true }) - if (previousHome === undefined) { - delete process.env.MCPX_HOME - } else { - process.env.MCPX_HOME = previousHome - } - if (previousFirst === undefined) { - delete process.env.FIRST_MCPX_TEST_TOKEN - } else { - process.env.FIRST_MCPX_TEST_TOKEN = previousFirst - } - if (previousSecond === undefined) { - delete process.env.SECOND_MCPX_TEST_TOKEN - } else { - process.env.SECOND_MCPX_TEST_TOKEN = previousSecond - } - } - }) }) function bearerServer(values: string[]): HttpServerConfig { diff --git a/tests/notification-dogfood.test.ts b/tests/notification-dogfood.test.ts index d121775..1f86429 100644 --- a/tests/notification-dogfood.test.ts +++ b/tests/notification-dogfood.test.ts @@ -193,6 +193,23 @@ describe('notification fixture dogfood', () => { hasRetainedSessionId: false, }), ) + + const removed = JSON.parse( + (await runMcpx(['@remove', '--name', 'notification-fixture', '--raw'])) + .stdout, + ) + expect(removed).toEqual({ + name: 'notification-fixture', + removed: true, + tokenRemoved: false, + }) + const registry = JSON.parse( + await fs.readFile( + path.join(home, '.agents', 'mcpx', 'state-v2', 'registry.json'), + 'utf8', + ), + ) + expect(registry.servers).toEqual({}) }) }) @@ -254,7 +271,7 @@ async function runMcpx( async function readRegisteredServerDiscoveredAt(): Promise { const raw = await fs.readFile( - path.join(home, '.agents', 'mcpx', 'servers.json'), + path.join(home, '.agents', 'mcpx', 'state-v2', 'schema-cache.json'), 'utf8', ) return JSON.parse(raw).servers['notification-fixture'].discoveredAt diff --git a/tests/router.test.ts b/tests/router.test.ts index 09b6a8c..032ab92 100644 --- a/tests/router.test.ts +++ b/tests/router.test.ts @@ -5,14 +5,7 @@ import { __test } from '../src/router' describe('router', () => { it('keeps mcpx control commands under the @ namespace', () => { const router = __test.buildRouter({ - config: { version: 1, servers: {} }, - ensureServerReady: async () => { - throw new Error('not used') - }, - reauthenticateServer: async () => { - throw new Error('not used') - }, - save: async () => {}, + servers: {}, }) expect(Object.keys(router).sort()).toEqual([ diff --git a/tests/runtime-authentication.test.ts b/tests/runtime-authentication.test.ts new file mode 100644 index 0000000..08862ec --- /dev/null +++ b/tests/runtime-authentication.test.ts @@ -0,0 +1,307 @@ +import { afterEach, describe, expect, it } from 'bun:test' +import fs from 'node:fs/promises' +import { tmpdir } from 'node:os' +import path from 'node:path' + +import { McpRuntime } from '../src/runtime' +import { RuntimeAuthentication } from '../src/runtime-authentication' +import { createInMemoryRuntimeCaller } from '../src/runtime-caller' +import { openRuntimeStores } from '../src/runtime-stores' + +describe('Runtime explicit authentication', () => { + const roots: string[] = [] + const servers: Bun.Server[] = [] + + afterEach(async () => { + for (const server of servers.splice(0)) server.stop(true) + await Promise.all( + roots.splice(0).map((root) => fs.rm(root, { recursive: true })), + ) + }) + + it('single-flights five explicit refresh callers through one local token request', async () => { + let refreshRequests = 0 + let issuer = '' + const fixture = Bun.serve({ + hostname: '127.0.0.1', + port: 0, + async fetch(request): Promise { + if (request.url.endsWith('/.well-known/oauth-authorization-server')) { + return Response.json({ + issuer, + authorization_endpoint: `${issuer}/authorize`, + token_endpoint: `${issuer}/token`, + }) + } + if (request.url.endsWith('/token')) { + refreshRequests += 1 + await Bun.sleep(20) + return Response.json({ + access_token: 'rotated-access', + refresh_token: 'rotated-refresh', + token_type: 'bearer', + expires_in: 3600, + }) + } + return new Response(null, { status: 404 }) + }, + }) + servers.push(fixture) + issuer = `http://127.0.0.1:${fixture.port}` + const tokenKey = `fixture:${issuer}` + const stores = await createStores( + { + url: `${issuer}/mcp`, + auth: { kind: 'oauth-token', tokenKey, confidence: 'confirmed' }, + }, + { + [tokenKey]: { + accessToken: 'expired', + refreshToken: 'refresh-1', + clientId: 'fixture-client', + tokenType: 'bearer', + expiresAt: '2000-01-01T00:00:00.000Z', + }, + }, + ) + const runtime = new McpRuntime(stores) + const callers = Array.from({ length: 5 }, (_, index) => + createInMemoryRuntimeCaller(`refresh-${index}`), + ) + + await Promise.all( + callers.map((caller) => + runtime.handle( + { + requestId: caller.id, + op: 'refreshServers', + serverNames: ['fixture'], + }, + caller, + ), + ), + ) + + expect(refreshRequests).toBe(1) + expect((await stores.credentials.read()).oauth[tokenKey]?.accessToken).toBe( + 'rotated-access', + ) + expect(callers.every((caller) => caller.frames.length === 1)).toBe(true) + }) + + it('keeps an interactive flow for remaining waiters and aborts after the final disconnect', async () => { + const stores = await createStores({ + url: 'http://127.0.0.1:1/mcp', + auth: { + kind: 'oauth', + confidence: 'confirmed', + authorizationServers: ['http://127.0.0.1:1'], + }, + }) + let signal: AbortSignal | undefined + let callbackOpen = false + let starts = 0 + const authentication = new RuntimeAuthentication(stores, { + authenticate: async (_name, _url, _auth, flowSignal) => { + starts += 1 + signal = flowSignal + callbackOpen = true + return new Promise((_resolve, reject) => { + flowSignal?.addEventListener( + 'abort', + () => { + callbackOpen = false + reject(flowSignal.reason) + }, + { once: true }, + ) + }) + }, + }) + const runtime = new McpRuntime(stores, { authentication }) + const first = createInMemoryRuntimeCaller('first') + const secondBase = createInMemoryRuntimeCaller('second') + let secondJoined = false + const second = { + ...secondBase, + onDisconnect: (listener: () => void) => { + // Opening the callback only proves the first waiter joined; wait for the + // second subscription so slower CI runners cannot disconnect too early. + secondJoined = true + return secondBase.onDisconnect(listener) + }, + } + const firstRun = runtime.handle( + { requestId: 'first', op: 'refreshServers', serverNames: ['fixture'] }, + first, + ) + const secondRun = runtime.handle( + { requestId: 'second', op: 'refreshServers', serverNames: ['fixture'] }, + second, + ) + await waitFor(() => callbackOpen && secondJoined) + + first.disconnect() + expect(signal?.aborted).toBe(false) + second.disconnect() + await Promise.all([firstRun, secondRun]) + + expect(starts).toBe(1) + expect(callbackOpen).toBe(false) + expect(first.frames).toEqual([]) + expect(second.frames).toEqual([]) + }) + + it('requests manual OAuth client input from the CLI caller and persists it in Runtime state', async () => { + const stores = await createStores({ + url: 'http://127.0.0.1:1/mcp', + auth: { + kind: 'oauth', + confidence: 'confirmed', + authorizationServers: ['http://127.0.0.1:1'], + }, + }) + let inputRequests = 0 + const authentication = new RuntimeAuthentication(stores, { + authenticate: async (_name, _url, _auth, _signal, manualClient) => { + const client = await manualClient?.({ + serverName: 'fixture', + redirectUri: 'http://127.0.0.1:65245/callback', + issuer: 'http://127.0.0.1:1', + scopes: ['scope:read'], + }) + if (!client) throw new Error('Missing manual client.') + return { + auth: { + kind: 'oauth-token', + tokenKey: 'fixture:issuer', + confidence: 'confirmed', + }, + token: { + accessToken: 'local-access', + tokenType: 'bearer', + clientId: client.clientId, + clientSecretKey: client.clientSecretKey, + }, + clientSecret: client.clientSecret, + } + }, + }) + const runtime = new McpRuntime(stores, { authentication }) + const caller = { + ...createInMemoryRuntimeCaller('manual'), + requestInput: async (request: { type: string }) => { + inputRequests += 1 + expect(request.type).toBe('oauth-client') + return { clientId: 'local-client', clientSecret: 'local-secret' } + }, + } + + await runtime.handle( + { requestId: 'manual', op: 'refreshServers', serverNames: ['fixture'] }, + caller, + ) + + expect(inputRequests).toBe(1) + const state = await stores.readState() + expect(state.credentials.oauth['fixture:issuer']?.clientId).toBe( + 'local-client', + ) + expect( + state.credentials.oauthClientSecrets['oauth-client:local-client'], + ).toBe('local-secret') + expect(state.registry.servers.fixture).toMatchObject({ + auth: { kind: 'oauth-token' }, + }) + }) + + it('moves manual OAuth input to a surviving waiter after the first caller disconnects', async () => { + const stores = await createStores({ + url: 'http://127.0.0.1:1/mcp', + auth: { + kind: 'oauth', + confidence: 'confirmed', + authorizationServers: ['http://127.0.0.1:1'], + }, + }) + let firstPrompted = false + let secondPrompted = false + const authentication = new RuntimeAuthentication(stores, { + authenticate: async (_name, _url, _auth, _signal, manualClient) => { + const client = await manualClient?.({ + serverName: 'fixture', + redirectUri: 'http://127.0.0.1:65245/callback', + issuer: 'http://127.0.0.1:1', + scopes: [], + }) + if (!client) throw new Error('Missing manual client.') + return { + auth: { + kind: 'oauth-token', + tokenKey: 'fixture:issuer', + confidence: 'confirmed', + }, + token: { + accessToken: 'local-access', + tokenType: 'bearer', + clientId: client.clientId, + }, + } + }, + }) + const firstBase = createInMemoryRuntimeCaller('first-input') + const first = { + ...firstBase, + requestInput: () => { + firstPrompted = true + return new Promise((_resolve, reject) => { + firstBase.onDisconnect(() => reject(new Error('caller disconnected'))) + }) + }, + } + const second = { + ...createInMemoryRuntimeCaller('second-input'), + requestInput: async () => { + secondPrompted = true + return { clientId: 'survivor', clientSecret: 'local-secret' } + }, + } + const firstRun = authentication.refreshServers(['fixture'], first) + const secondRun = authentication.refreshServers(['fixture'], second) + await waitFor(() => firstPrompted) + + firstBase.disconnect() + const outcomes = await Promise.all([firstRun, secondRun]) + + expect(secondPrompted).toBe(true) + expect(outcomes).toEqual([ + { status: 'disconnected' }, + { status: 'completed', refreshed: ['fixture'] }, + ]) + }) + + async function createStores( + server: Record, + oauth: Record = {}, + ) { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-auth-')) + roots.push(root) + await fs.writeFile( + path.join(root, 'servers.json'), + JSON.stringify({ version: 1, servers: { fixture: server } }), + ) + await fs.writeFile( + path.join(root, 'tokens.json'), + JSON.stringify({ version: 1, oauth, oauthClientSecrets: {} }), + ) + return openRuntimeStores(root) + } +}) + +async function waitFor(predicate: () => boolean): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (predicate()) return + await Bun.sleep(1) + } + throw new Error('Condition was not met.') +} diff --git a/tests/runtime-call.test.ts b/tests/runtime-call.test.ts new file mode 100644 index 0000000..80786d1 --- /dev/null +++ b/tests/runtime-call.test.ts @@ -0,0 +1,200 @@ +import { describe, expect, it } from 'bun:test' + +import type { RuntimeCaller } from '../src/runtime-caller' + +import { CancelableFifo, RuntimeCall } from '../src/runtime-call' +import { createInMemoryRuntimeCaller } from '../src/runtime-caller' + +describe('caller-owned Runtime Call lifecycle', () => { + it('removes a disconnected queued Call before activation', async () => { + const fifo = new CancelableFifo() + const firstCaller = createInMemoryRuntimeCaller('first') + const queuedCaller = createInMemoryRuntimeCaller('queued') + const first = new RuntimeCall(firstCaller) + const queued = new RuntimeCall(queuedCaller) + const reached: string[] = [] + let releaseFirst = () => {} + const firstBlocked = new Promise((resolve) => { + releaseFirst = resolve + }) + + fifo.enqueue(first, async () => { + reached.push('first') + await firstBlocked + return 'first-result' + }) + fifo.enqueue(queued, async () => { + reached.push('queued') + return 'queued-result' + }) + await waitFor(() => first.state === 'active') + + queuedCaller.disconnect() + releaseFirst() + await fifo.idle() + + expect(reached).toEqual(['first']) + expect(queued.state).toBe('terminal') + expect(queued.cancellationCause).toBe('caller-disconnected') + expect(queuedCaller.frames).toEqual([]) + expect(fifo.status()).toEqual({ activeCalls: 0, queuedCalls: 0 }) + }) + + it('aborts an active Call exactly once and never writes to its dead caller', async () => { + const fifo = new CancelableFifo() + const caller = createInMemoryRuntimeCaller('active') + const call = new RuntimeCall(caller) + let abortEvents = 0 + + fifo.enqueue(call, (signal) => { + return new Promise((_resolve, reject) => { + signal.addEventListener( + 'abort', + () => { + abortEvents += 1 + reject(signal.reason) + }, + { once: true }, + ) + }) + }) + await waitFor(() => call.state === 'active') + + caller.disconnect() + caller.disconnect() + await fifo.idle() + + expect(abortEvents).toBe(1) + expect(call.cancellationCause).toBe('caller-disconnected') + expect(caller.frames).toEqual([]) + }) + + it('removes the disconnect listener before normal completion', async () => { + const fifo = new CancelableFifo() + const caller = createInMemoryRuntimeCaller('completed') + const call = new RuntimeCall(caller) + + fifo.enqueue(call, async () => 'done') + await fifo.idle() + caller.disconnect() + + expect(call.signal.aborted).toBe(false) + expect(call.cancellationCause).toBeUndefined() + expect(caller.frames).toEqual([ + { requestId: 'completed', kind: 'result', result: 'done' }, + ]) + }) + + it('keeps timeout distinct from caller disconnect in owned state', async () => { + const fifo = new CancelableFifo() + const caller = createInMemoryRuntimeCaller('timed-out') + const call = new RuntimeCall(caller) + + fifo.enqueue(call, (signal) => { + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + }) + await waitFor(() => call.state === 'active') + await call.cancel('timeout') + await fifo.idle() + + expect(call.cancellationCause).toBe('timeout') + expect(caller.frames).toEqual([ + { + requestId: 'timed-out', + kind: 'error', + error: { code: 'timeout', message: 'Runtime Call timed out.' }, + }, + ]) + }) + + it('continues FIFO execution after failure, timeout, and cancellation', async () => { + const fifo = new CancelableFifo() + const order: string[] = [] + const failed = new RuntimeCall(createInMemoryRuntimeCaller('failed')) + const timedOut = new RuntimeCall(createInMemoryRuntimeCaller('timeout')) + const cancelledCaller = createInMemoryRuntimeCaller('cancelled') + const cancelled = new RuntimeCall(cancelledCaller) + const succeeded = new RuntimeCall(createInMemoryRuntimeCaller('succeeded')) + + fifo.enqueue(failed, async () => { + order.push('failed') + throw new Error('fixture failure') + }) + fifo.enqueue(timedOut, async (signal) => { + order.push('timeout') + await timedOut.cancel('timeout') + throw signal.reason + }) + fifo.enqueue(cancelled, async () => { + order.push('cancelled') + return null + }) + cancelledCaller.disconnect() + fifo.enqueue(succeeded, async () => { + order.push('succeeded') + return 'ok' + }) + + await fifo.idle() + + expect(order).toEqual(['failed', 'timeout', 'succeeded']) + expect(fifo.status()).toEqual({ activeCalls: 0, queuedCalls: 0 }) + }) + + it('never activates a Call whose caller disconnected before construction', async () => { + const caller = createInMemoryRuntimeCaller('already-gone') + caller.disconnect() + const call = new RuntimeCall(caller) + const fifo = new CancelableFifo() + let activated = false + + fifo.enqueue(call, async () => { + activated = true + }) + await fifo.idle() + + expect(activated).toBe(false) + expect(call.cancellationCause).toBe('caller-disconnected') + }) + + it('releases FIFO ownership before terminal socket delivery completes', async () => { + let releaseSend = () => {} + const blockedSend = new Promise((resolve) => { + releaseSend = resolve + }) + const blockedCaller: RuntimeCaller = { + id: 'blocked', + onDisconnect: () => () => {}, + requestInput: async () => undefined, + send: () => blockedSend, + } + const fifo = new CancelableFifo() + let secondRan = false + + fifo.enqueue(new RuntimeCall(blockedCaller), async () => 'first') + fifo.enqueue( + new RuntimeCall(createInMemoryRuntimeCaller('second')), + async () => { + secondRan = true + return 'second' + }, + ) + await fifo.idle() + + expect(secondRan).toBe(true) + expect(fifo.status()).toEqual({ activeCalls: 0, queuedCalls: 0 }) + releaseSend() + }) +}) + +async function waitFor(predicate: () => boolean): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (predicate()) return + await Bun.sleep(1) + } + throw new Error('Condition was not met.') +} diff --git a/tests/runtime-caller.test.ts b/tests/runtime-caller.test.ts new file mode 100644 index 0000000..cdd89c9 --- /dev/null +++ b/tests/runtime-caller.test.ts @@ -0,0 +1,82 @@ +import { describe, expect, it } from 'bun:test' + +import { createInMemoryRuntimeCaller } from '../src/runtime-caller' + +describe('RuntimeCaller', () => { + it('keeps every frame correlated to its request', async () => { + const caller = createInMemoryRuntimeCaller('request-1') + + await caller.send({ + requestId: 'request-1', + kind: 'event', + event: { type: 'progress', message: 'working' }, + }) + await caller.send({ + requestId: 'request-1', + kind: 'result', + result: { ok: true }, + }) + + expect(caller.frames).toEqual([ + { + requestId: 'request-1', + kind: 'event', + event: { type: 'progress', message: 'working' }, + }, + { + requestId: 'request-1', + kind: 'result', + result: { ok: true }, + }, + ]) + }) + + it('rejects mismatched request ids and every frame after a terminal', async () => { + const caller = createInMemoryRuntimeCaller('request-1') + + expect( + caller.send({ + requestId: 'request-2', + kind: 'result', + result: null, + }), + ).rejects.toThrow('does not match caller request-1') + + await caller.send({ + requestId: 'request-1', + kind: 'error', + error: { code: 'operation-failed', message: 'failed' }, + }) + expect( + caller.send({ + requestId: 'request-1', + kind: 'result', + result: null, + }), + ).rejects.toThrow('already received a terminal frame') + }) + + it('removes disconnect subscriptions independently', () => { + const caller = createInMemoryRuntimeCaller('request-1') + let first = 0 + let second = 0 + const unsubscribe = caller.onDisconnect(() => first++) + caller.onDisconnect(() => second++) + + unsubscribe() + caller.disconnect() + + expect(first).toBe(0) + expect(second).toBe(1) + }) + + it('immediately notifies subscribers that arrive after disconnect', () => { + const caller = createInMemoryRuntimeCaller('request-1') + caller.disconnect() + let notifications = 0 + + caller.onDisconnect(() => notifications++) + + expect(notifications).toBe(1) + }) +}) diff --git a/tests/runtime-cli.test.ts b/tests/runtime-cli.test.ts new file mode 100644 index 0000000..f47c39b --- /dev/null +++ b/tests/runtime-cli.test.ts @@ -0,0 +1,356 @@ +import { afterEach, beforeEach, describe, expect, it } from 'bun:test' +import fs from 'node:fs/promises' +import { tmpdir } from 'node:os' +import path from 'node:path' + +const mainPath = path.join(import.meta.dir, '..', 'src', 'main.ts') +const fixtureServers: Bun.Server[] = [] +describe('Runtime CLI adapter lifecycle', () => { + let home: string + + beforeEach(async () => { + home = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-cli-')) + }) + + afterEach(async () => { + await runCli(['@daemon', 'stop', '--raw']).catch(() => {}) + for (const server of fixtureServers.splice(0)) server.stop(true) + await fs.rm(home, { recursive: true, force: true }) + }) + + it('cancels the originating active Call when the CLI process exits and reuses the session', async () => { + const fixture = startHttpFixture() + await addFixture(fixture.url) + + const active = Array.from({ length: 5 }, () => + spawnCli([ + 'controlled', + 'controlled', + '--scenario', + 'acknowledge', + '--raw', + ]), + ) + await waitFor(async () => { + const status = await runCli(['@daemon', 'status', '--raw']) + if (status.exitCode !== 0) return false + const session = JSON.parse(status.stdout).servers[0] + return session?.activeCalls === 1 && session?.queuedCalls === 4 + }) + + for (const child of active) child.kill() + await Promise.all(active.map((child) => child.exited)) + await waitFor(async () => { + const status = await runCli(['@daemon', 'status', '--raw']) + return ( + status.exitCode === 0 && + JSON.parse(status.stdout).servers[0]?.activeCalls === 0 + ) + }) + + const reused = await runCli(['controlled', 'echo', '--raw']) + expect(reused).toEqual({ + exitCode: 0, + stdout: 'echo-ok\n', + stderr: '', + }) + }, 10_000) + + it('exits five concurrent success and error CLIs without surviving PIDs', async () => { + const fixture = startHttpFixture() + await addFixture(fixture.url) + const successes = Array.from({ length: 5 }, () => + spawnObserved(['controlled', 'echo', '--raw']), + ) + const successResults = await Promise.all(successes.map(waitForExit)) + expect(successResults).toEqual( + Array.from({ length: 5 }, () => ({ + exitCode: 0, + stdout: 'echo-ok\n', + stderr: '', + })), + ) + for (const child of successes) expect(isAlive(child.pid)).toBe(false) + + const failures = Array.from({ length: 5 }, () => + spawnObserved(['controlled', 'fail', '--raw']), + ) + const failureResults = await Promise.all(failures.map(waitForExit)) + expect(failureResults.every((result) => result.exitCode === 1)).toBe(true) + expect( + failureResults.every((result) => + result.stderr.includes('fixture failure'), + ), + ).toBe(true) + for (const child of failures) expect(isAlive(child.pid)).toBe(false) + }, 10_000) + + it('treats a closed stdout pipe as a bounded terminal path', async () => { + const fixture = startHttpFixture() + await addFixture(fixture.url) + const children = Array.from({ length: 5 }, () => + spawnCli(['controlled', 'echo', '--raw']), + ) + await Promise.all(children.map((child) => child.stdout.cancel())) + const exitCodes = await Promise.all( + children.map((child) => + Promise.race([child.exited, Bun.sleep(3_000).then(() => undefined)]), + ), + ) + expect(exitCodes.every((code) => code !== undefined)).toBe(true) + for (const child of children) expect(isAlive(child.pid)).toBe(false) + }, 10_000) + + it('single-flights five explicit refresh CLI processes through one local token request', async () => { + const fixture = startHttpFixture() + await seedExpiredOAuth(fixture.url, fixture.issuer) + const refreshes = Array.from({ length: 5 }, () => + spawnObserved(['@refresh', '--raw']), + ) + const results = await Promise.all(refreshes.map(waitForExit)) + + expect(results.every((result) => result.exitCode === 0)).toBe(true) + expect(fixture.tokenRequests()).toBe(1) + for (const child of refreshes) expect(isAlive(child.pid)).toBe(false) + }, 10_000) + + it('cancels an active refresh flow before Runtime stop completes', async () => { + const fixture = startHttpFixture({ holdToken: true }) + await seedExpiredOAuth(fixture.url, fixture.issuer) + const refresh = spawnObserved(['@refresh', '--raw']) + await waitFor(async () => fixture.tokenRequests() === 1) + + const stopped = await runCli(['@daemon', 'stop', '--raw']) + const refreshResult = await waitForExit(refresh) + + expect(stopped.exitCode).toBe(0) + expect(refreshResult.exitCode).toBe(1) + expect(isAlive(refresh.pid)).toBe(false) + }, 10_000) + + async function addFixture(url: string): Promise { + const added = await runCli([ + '@add', + '--name', + 'controlled', + '--url', + url, + '--raw', + ]) + expect(added.exitCode).toBe(0) + } + + async function seedExpiredOAuth(url: string, issuer: string): Promise { + const root = path.join(home, '.agents', 'mcpx') + await fs.mkdir(root, { recursive: true }) + const tokenKey = `controlled:${issuer}` + await fs.writeFile( + path.join(root, 'servers.json'), + JSON.stringify({ + version: 1, + servers: { + controlled: { + url, + auth: { kind: 'oauth-token', tokenKey, confidence: 'confirmed' }, + }, + }, + }), + ) + await fs.writeFile( + path.join(root, 'tokens.json'), + JSON.stringify({ + version: 1, + oauth: { + [tokenKey]: { + accessToken: 'expired', + refreshToken: 'local-refresh', + clientId: 'local-client', + tokenType: 'bearer', + expiresAt: '2000-01-01T00:00:00.000Z', + }, + }, + }), + ) + } + + function spawnCli(args: string[]) { + return Bun.spawn([process.execPath, mainPath, ...args], { + env: { ...process.env, HOME: home, MCPX_HOME: home }, + stdout: 'pipe', + stderr: 'pipe', + }) + } + + function spawnObserved(args: string[]) { + const proc = spawnCli(args) + return { + proc, + pid: proc.pid, + stdout: new Response(proc.stdout).text(), + stderr: new Response(proc.stderr).text(), + } + } + + async function waitForExit(child: ReturnType) { + const exitCode = await Promise.race([ + child.proc.exited, + Bun.sleep(3_000).then(() => undefined), + ]) + if (exitCode === undefined) { + child.proc.kill() + throw new Error(`CLI ${child.pid} did not exit before the deadline.`) + } + return { + exitCode, + stdout: await child.stdout, + stderr: await child.stderr, + } + } + + async function runCli(args: string[]): Promise<{ + exitCode: number + stdout: string + stderr: string + }> { + const proc = spawnCli(args) + const stdout = new Response(proc.stdout).text() + const stderr = new Response(proc.stderr).text() + return { + exitCode: await proc.exited, + stdout: await stdout, + stderr: await stderr, + } + } +}) + +function isAlive(pid: number): boolean { + try { + process.kill(pid, 0) + return true + } catch (error) { + return (error as NodeJS.ErrnoException).code !== 'ESRCH' + } +} + +function startHttpFixture(options: { holdToken?: boolean } = {}): { + url: string + issuer: string + tokenRequests: () => number +} { + let tokenRequests = 0 + let issuer = '' + const pending = new Map void>() + const server = Bun.serve({ + hostname: '127.0.0.1', + port: 0, + async fetch(request) { + const requestUrl = new URL(request.url) + if (requestUrl.pathname === '/.well-known/oauth-authorization-server') { + return Response.json({ + issuer, + authorization_endpoint: `${issuer}/authorize`, + token_endpoint: `${issuer}/token`, + }) + } + if (requestUrl.pathname === '/token') { + tokenRequests += 1 + if (options.holdToken) await new Promise(() => {}) + await Bun.sleep(25) + return Response.json({ + access_token: 'local-access', + refresh_token: 'local-refresh-2', + token_type: 'bearer', + expires_in: 3600, + }) + } + if (requestUrl.pathname !== '/mcp' || request.method !== 'POST') { + return new Response(null, { status: 404 }) + } + const message = (await request.json()) as { + id?: string | number + method: string + params?: Record + } + if (message.method === 'initialize') { + return rpcResponse(message.id, { + protocolVersion: message.params?.protocolVersion, + capabilities: { tools: {} }, + serverInfo: { name: 'local-cli-fixture', version: '1.0.0' }, + }) + } + if (message.method === 'notifications/initialized') + return acceptedResponse() + if (message.method === 'tools/list') { + return rpcResponse(message.id, { + tools: [ + { name: 'echo', inputSchema: { type: 'object' } }, + { name: 'fail', inputSchema: { type: 'object' } }, + { + name: 'controlled', + inputSchema: { + type: 'object', + properties: { scenario: { type: 'string' } }, + required: ['scenario'], + }, + }, + ], + }) + } + if (message.method === 'notifications/cancelled') { + const id = message.params?.requestId as string | number + pending.get(id)?.(rpcResponse(id, toolResult('cancelled'))) + pending.delete(id) + return acceptedResponse() + } + if (message.method !== 'tools/call') return acceptedResponse() + if (message.params?.name === 'echo') + return rpcResponse(message.id, toolResult('echo-ok')) + if (message.params?.name === 'fail') { + return rpcError(message.id, -32_000, 'fixture failure') + } + return new Promise((resolve) => { + if (message.id !== undefined) pending.set(message.id, resolve) + }) + }, + }) + fixtureServers.push(server) + issuer = `http://127.0.0.1:${server.port}` + return { + url: `${issuer}/mcp`, + issuer, + tokenRequests: () => tokenRequests, + } +} + +function acceptedResponse(): Response { + return new Response(null, { + status: 202, + headers: { 'mcp-session-id': 'local-session' }, + }) +} + +function rpcResponse(id: unknown, result: unknown): Response { + return Response.json( + { jsonrpc: '2.0', id, result }, + { headers: { 'mcp-session-id': 'local-session' } }, + ) +} + +function rpcError(id: unknown, code: number, message: string): Response { + return Response.json( + { jsonrpc: '2.0', id, error: { code, message } }, + { headers: { 'mcp-session-id': 'local-session' } }, + ) +} + +function toolResult(text: string): Record { + return { content: [{ type: 'text', text }] } +} + +async function waitFor(predicate: () => Promise): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (await predicate()) return + await Bun.sleep(20) + } + throw new Error('Condition was not met.') +} diff --git a/tests/runtime-protocol.test.ts b/tests/runtime-protocol.test.ts new file mode 100644 index 0000000..1372205 --- /dev/null +++ b/tests/runtime-protocol.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, it } from 'bun:test' + +import { + DAEMON_PROTOCOL_VERSION, + RuntimeConnectionState, + createHello, + parseRuntimeIntent, +} from '../src/runtime-protocol' + +describe('Runtime protocol v3', () => { + it('accepts one handshake followed by one operation', () => { + const connection = new RuntimeConnectionState() + + expect(connection.accept(createHello())).toEqual({ kind: 'hello' }) + expect( + connection.accept({ + requestId: 'request-1', + op: 'call', + serverName: 'fixture', + toolName: 'search', + input: { query: 'runtime' }, + }), + ).toEqual({ + kind: 'intent', + intent: { + requestId: 'request-1', + op: 'call', + serverName: 'fixture', + toolName: 'search', + input: { query: 'runtime' }, + }, + }) + expect(connection.accept({ requestId: 'request-2', op: 'status' })).toEqual( + { + kind: 'error', + error: { + code: 'connection-complete', + message: 'A Runtime connection accepts exactly one operation.', + }, + }, + ) + }) + + it('rejects protocol mismatch and operations before the handshake', () => { + const beforeHandshake = new RuntimeConnectionState() + expect( + beforeHandshake.accept({ requestId: 'request-1', op: 'status' }), + ).toMatchObject({ kind: 'error', error: { code: 'handshake-required' } }) + + const mismatch = new RuntimeConnectionState() + expect( + mismatch.accept({ + kind: 'hello', + protocolVersion: DAEMON_PROTOCOL_VERSION - 1, + clientVersion: '0.0.0', + }), + ).toMatchObject({ kind: 'error', error: { code: 'protocol-mismatch' } }) + }) + + it('parses only the final intent shapes without credential material', () => { + expect( + parseRuntimeIntent({ + requestId: 'request-1', + op: 'refreshServers', + serverNames: ['fixture'], + }), + ).toEqual({ + requestId: 'request-1', + op: 'refreshServers', + serverNames: ['fixture'], + }) + + expect( + parseRuntimeIntent({ + requestId: 'request-2', + op: 'call', + serverName: 'fixture', + toolName: 'search', + input: {}, + headers: { Authorization: 'secret' }, + }), + ).toEqual({ + code: 'invalid-frame', + message: 'Invalid Runtime intent.', + }) + }) +}) diff --git a/tests/runtime-session-pool.test.ts b/tests/runtime-session-pool.test.ts new file mode 100644 index 0000000..06e6981 --- /dev/null +++ b/tests/runtime-session-pool.test.ts @@ -0,0 +1,713 @@ +import { afterEach, describe, expect, it } from 'bun:test' +import fs from 'node:fs/promises' +import { tmpdir } from 'node:os' +import path from 'node:path' + +import type { McpConnection } from '../src/mcp-client' +import type { ServerConfig } from '../src/types' + +import { RuntimeCall } from '../src/runtime-call' +import { createInMemoryRuntimeCaller } from '../src/runtime-caller' +import { RuntimeSessionPool } from '../src/runtime-session-pool' +import { openRuntimeStores } from '../src/runtime-stores' + +type CancellationScenario = + | 'acknowledge' + | 'ignore' + | 'race' + | 'complete-before-cancel' + +const cancellationScenarios: CancellationScenario[] = [ + 'acknowledge', + 'ignore', + 'race', + 'complete-before-cancel', +] + +describe('Runtime session pool', () => { + const roots: string[] = [] + + afterEach(async () => { + await Promise.all( + roots.splice(0).map((root) => fs.rm(root, { recursive: true })), + ) + }) + + it('resolves declarations and credentials inside the Runtime', async () => { + const stores = await createStores({ + url: 'http://127.0.0.1:1/mcp', + headers: { 'x-api-key': 'header-secret' }, + auth: { + kind: 'bearer', + strategy: 'round-robin', + confidence: 'configured', + credentials: [{ kind: 'literal', value: 'bearer-secret' }], + }, + }) + let connectedServer: ServerConfig | undefined + let connectedHeaders: Record | undefined + let receivedSignal: AbortSignal | undefined + const pool = new RuntimeSessionPool(stores, { + connect: async (server, options) => { + connectedServer = server + connectedHeaders = options?.headers + return fakeConnection(async (_params, _schema, requestOptions) => { + receivedSignal = requestOptions?.signal + return 'ok' + }) + }, + }) + const caller = createInMemoryRuntimeCaller('request-1') + const call = new RuntimeCall(caller) + + await pool.call(call, { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }) + + expect(connectedServer).toEqual({ + url: 'http://127.0.0.1:1/mcp', + auth: { kind: 'none' }, + }) + expect(connectedHeaders).toEqual({ + Accept: 'application/json, text/event-stream', + Authorization: 'Bearer bearer-secret', + 'x-api-key': 'header-secret', + }) + expect(receivedSignal).toBe(call.signal) + expect(JSON.stringify(caller.frames)).not.toContain('secret') + }) + + it('removes queued calls and reuses the session after active cancellation', async () => { + const stores = await createStores({ + transport: 'stdio', + command: process.execPath, + auth: undefined, + }) + let connectCount = 0 + const upstreamCalls: string[] = [] + const pool = new RuntimeSessionPool(stores, { + connect: async () => { + connectCount += 1 + return fakeConnection(async (params, _schema, options) => { + upstreamCalls.push(params.name) + if (params.name === 'hold') { + return new Promise((_resolve, reject) => { + options?.signal?.addEventListener( + 'abort', + () => reject(options.signal?.reason), + { once: true }, + ) + }) + } + return 'echo-ok' + }) + }, + }) + const activeCaller = createInMemoryRuntimeCaller('active') + const queuedCaller = createInMemoryRuntimeCaller('queued') + const active = pool.call(new RuntimeCall(activeCaller), { + serverName: 'fixture', + toolName: 'hold', + input: {}, + }) + await waitFor(() => upstreamCalls.length === 1) + const queued = pool.call(new RuntimeCall(queuedCaller), { + serverName: 'fixture', + toolName: 'never', + input: {}, + }) + + queuedCaller.disconnect() + activeCaller.disconnect() + await Promise.all([active, queued]) + + const reuseCaller = createInMemoryRuntimeCaller('reuse') + await pool.call(new RuntimeCall(reuseCaller), { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }) + + expect(connectCount).toBe(1) + expect(upstreamCalls).toEqual(['hold', 'echo']) + expect(activeCaller.frames).toEqual([]) + expect(queuedCaller.frames).toEqual([]) + expect(reuseCaller.frames[0]).toMatchObject({ + kind: 'result', + result: { result: 'echo-ok' }, + }) + }) + + it('maps 401 to reauth-required without starting authentication', async () => { + const stores = await createStores({ + url: 'http://127.0.0.1:1/mcp', + auth: { kind: 'none' }, + }) + let closed = 0 + const pool = new RuntimeSessionPool(stores, { + connect: async () => + fakeConnection( + async () => { + throw new Error('HTTP 401 Unauthorized') + }, + () => closed++, + ), + }) + const caller = createInMemoryRuntimeCaller('request-1') + + await pool.call(new RuntimeCall(caller), { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }) + + expect(caller.frames).toEqual([ + { + requestId: 'request-1', + kind: 'error', + error: { + code: 'reauth-required', + message: 'Credentials for fixture must be refreshed.', + }, + }, + ]) + expect(closed).toBe(1) + }) + + it('rotates bearer credentials only at queued-to-active ownership', async () => { + const stores = await createStores({ + url: 'http://127.0.0.1:1/mcp', + auth: { + kind: 'bearer', + strategy: 'round-robin', + confidence: 'configured', + credentials: [ + { kind: 'literal', value: 'token-a' }, + { kind: 'literal', value: 'token-b' }, + ], + }, + }) + let authorization = '' + let releaseFirst = () => {} + const firstBlocked = new Promise((resolve) => { + releaseFirst = resolve + }) + const observed: string[] = [] + const pool = new RuntimeSessionPool(stores, { + connect: async (_server, options) => { + authorization = options?.headers?.Authorization ?? '' + return { + ...fakeConnection(async (params) => { + observed.push(`${params.name}:${authorization}`) + if (params.name === 'first') await firstBlocked + return 'ok' + }), + updateHeaders: (headers) => { + authorization = headers.Authorization ?? '' + }, + } + }, + }) + const first = pool.call( + new RuntimeCall(createInMemoryRuntimeCaller('first')), + { + serverName: 'fixture', + toolName: 'first', + input: {}, + }, + ) + await waitFor(() => observed.length === 1) + const second = pool.call( + new RuntimeCall(createInMemoryRuntimeCaller('second')), + { + serverName: 'fixture', + toolName: 'second', + input: {}, + }, + ) + await Bun.sleep(5) + expect(observed).toEqual(['first:Bearer token-a']) + releaseFirst() + await Promise.all([first, second]) + expect(observed).toEqual(['first:Bearer token-a', 'second:Bearer token-b']) + }) + + it('owns the timeout cause instead of inferring it from the SDK error', async () => { + const previous = process.env.MCPX_TOOL_CALL_TIMEOUT_MS + process.env.MCPX_TOOL_CALL_TIMEOUT_MS = '5' + try { + const stores = await createStores({ + transport: 'stdio', + command: process.execPath, + }) + const pool = new RuntimeSessionPool(stores, { + connect: async () => + fakeConnection( + async (_params, _schema, options) => + new Promise((_resolve, reject) => { + options?.signal?.addEventListener( + 'abort', + () => reject(new Error('SDK -32001')), + { once: true }, + ) + }), + ), + }) + const caller = createInMemoryRuntimeCaller('timeout-owned') + const call = new RuntimeCall(caller) + + const settlement = await Promise.race([ + pool + .call(call, { + serverName: 'fixture', + toolName: 'hold', + input: {}, + }) + .then(() => 'settled'), + Bun.sleep(50).then(() => 'deadline'), + ]) + + expect(settlement).toBe('settled') + expect(call.cancellationCause).toBe('timeout') + expect(caller.frames[0]).toMatchObject({ + kind: 'error', + error: { code: 'timeout' }, + }) + } finally { + if (previous === undefined) delete process.env.MCPX_TOOL_CALL_TIMEOUT_MS + else process.env.MCPX_TOOL_CALL_TIMEOUT_MS = previous + } + }) + + it('preserves one stdio SDK connection across cancellation races and late responses', async () => { + const stores = await createStores({ + transport: 'stdio', + command: process.execPath, + args: [ + path.join( + import.meta.dir, + '..', + 'prototypes', + 'mcp-cancellation', + 'stdio-fixture.mjs', + ), + ], + }) + const pool = new RuntimeSessionPool(stores) + try { + await runCancellationMatrix(pool) + } finally { + await pool.close() + } + }) + + it('preserves one Streamable HTTP session across cancellation races and late responses', async () => { + const fixture = startCancellationHttpFixture() + const stores = await createStores({ + url: fixture.url, + auth: { kind: 'none' }, + }) + const pool = new RuntimeSessionPool(stores) + try { + await runCancellationMatrix(pool) + const initializedSessions = fixture.sessionIds.filter(Boolean) + expect(new Set(initializedSessions)).toEqual(new Set(['local-session'])) + } finally { + await pool.close() + fixture.stop() + } + }) + + it('closes idle sessions and rejects admission after shutdown', async () => { + const stores = await createStores({ + transport: 'stdio', + command: process.execPath, + }) + let closes = 0 + const pool = new RuntimeSessionPool(stores, { + connect: async () => + fakeConnection( + async () => 'ok', + () => closes++, + ), + }) + await pool.call(new RuntimeCall(createInMemoryRuntimeCaller('first')), { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }) + + await pool.cleanupIdle(0) + expect(pool.sessionCount()).toBe(0) + expect(closes).toBe(1) + await pool.close() + const rejected = createInMemoryRuntimeCaller('rejected') + await expect( + pool.call(new RuntimeCall(rejected), { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }), + ).rejects.toThrow('stopping') + }) + + it('cancels active and queued Calls before shutdown releases the pool', async () => { + const stores = await createStores({ + transport: 'stdio', + command: process.execPath, + }) + const pool = new RuntimeSessionPool(stores, { + connect: async () => + fakeConnection( + async (_params, _schema, options) => + new Promise((_resolve, reject) => { + options?.signal?.addEventListener( + 'abort', + () => reject(options.signal?.reason), + { once: true }, + ) + }), + ), + }) + const activeCaller = createInMemoryRuntimeCaller('active-stop') + const queuedCaller = createInMemoryRuntimeCaller('queued-stop') + const active = pool.call(new RuntimeCall(activeCaller), { + serverName: 'fixture', + toolName: 'hold', + input: {}, + }) + await waitFor(() => pool.status()[0]?.activeCalls === 1) + const queued = pool.call(new RuntimeCall(queuedCaller), { + serverName: 'fixture', + toolName: 'queued', + input: {}, + }) + await waitFor(() => pool.status()[0]?.queuedCalls === 1) + + await pool.close() + await Promise.all([active, queued]) + + expect(activeCaller.frames[0]).toMatchObject({ + kind: 'error', + error: { code: 'cancelled' }, + }) + expect(queuedCaller.frames[0]).toMatchObject({ + kind: 'error', + error: { code: 'cancelled' }, + }) + expect(pool.status()).toEqual([]) + }) + + it('rechecks admission after a delayed store read before creating a session', async () => { + const stores = await createStores({ + transport: 'stdio', + command: process.execPath, + }) + const originalReadState = stores.readState + let readStarted = false + let releaseRead = () => {} + const readBlocked = new Promise((resolve) => { + releaseRead = resolve + }) + const delayedStores = { + ...stores, + readState: async () => { + readStarted = true + await readBlocked + return originalReadState() + }, + } + let connects = 0 + const pool = new RuntimeSessionPool(delayedStores, { + connect: async () => { + connects += 1 + return fakeConnection(async () => 'ok') + }, + }) + const run = pool.call( + new RuntimeCall(createInMemoryRuntimeCaller('late-admission')), + { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }, + ) + await waitFor(() => readStarted) + + await pool.close() + releaseRead() + + await expect(run).rejects.toThrow('stopping') + expect(connects).toBe(0) + expect(pool.sessionCount()).toBe(0) + }) + + it('closes a connection that resolves after shutdown has begun', async () => { + const stores = await createStores({ + transport: 'stdio', + command: process.execPath, + }) + let releaseConnect = () => {} + let connectStarted = false + let closes = 0 + const connection = fakeConnection( + async () => 'ok', + () => closes++, + ) + const pool = new RuntimeSessionPool(stores, { + connect: async () => { + connectStarted = true + await new Promise((resolve) => { + releaseConnect = resolve + }) + return connection + }, + }) + const run = pool.call( + new RuntimeCall(createInMemoryRuntimeCaller('late-connect')), + { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }, + ) + await waitFor(() => connectStarted) + const closing = pool.close() + await Bun.sleep(2) + + releaseConnect() + await Promise.all([run, closing]) + + expect(closes).toBe(1) + expect(pool.sessionCount()).toBe(0) + }) + + it('never starts a connection after shutdown completes during header resolution', async () => { + const stores = await createStores({ + transport: 'stdio', + command: process.execPath, + }) + const originalReadState = stores.readState + let reads = 0 + let headerReadStarted = false + let releaseHeaderRead = () => {} + const headerReadBlocked = new Promise((resolve) => { + releaseHeaderRead = resolve + }) + const delayedStores = { + ...stores, + readState: async () => { + reads += 1 + if (reads === 2) { + headerReadStarted = true + await headerReadBlocked + } + return originalReadState() + }, + } + let connects = 0 + const pool = new RuntimeSessionPool(delayedStores, { + connect: async () => { + connects += 1 + return fakeConnection(async () => 'ok') + }, + }) + const run = pool.call( + new RuntimeCall(createInMemoryRuntimeCaller('header-stop')), + { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }, + ) + await waitFor(() => headerReadStarted) + + await pool.close() + releaseHeaderRead() + await run + await Bun.sleep(2) + + expect(connects).toBe(0) + expect(pool.sessionCount()).toBe(0) + }) + + async function createStores(server: Record) { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-session-pool-')) + roots.push(root) + const normalized = { ...server } + if (normalized.transport === 'stdio') delete normalized.auth + await fs.writeFile( + path.join(root, 'servers.json'), + JSON.stringify({ version: 1, servers: { fixture: normalized } }), + ) + return openRuntimeStores(root) + } +}) + +async function runCancellationMatrix(pool: RuntimeSessionPool): Promise { + for (const scenario of cancellationScenarios) { + const caller = createInMemoryRuntimeCaller(`controlled-${scenario}`) + const call = new RuntimeCall(caller) + const execution = pool.call(call, { + serverName: 'fixture', + toolName: 'controlled', + input: { scenario }, + }) + + if (scenario === 'complete-before-cancel') { + await execution + caller.disconnect() + expect(call.signal.aborted).toBe(false) + } else { + await Bun.sleep(15) + caller.disconnect() + await execution + expect(call.cancellationCause).toBe('caller-disconnected') + expect(caller.frames).toEqual([]) + } + + expect(await echoThroughPool(pool, `${scenario}-before-late`)).toBe( + 'echo-ok', + ) + await Bun.sleep(scenario === 'ignore' ? 110 : 35) + expect(await echoThroughPool(pool, `${scenario}-after-late`)).toBe( + 'echo-ok', + ) + } +} + +async function echoThroughPool( + pool: RuntimeSessionPool, + requestId: string, +): Promise { + const caller = createInMemoryRuntimeCaller(requestId) + await pool.call(new RuntimeCall(caller), { + serverName: 'fixture', + toolName: 'echo', + input: {}, + }) + const frame = caller.frames[0] + if (!frame || frame.kind !== 'result') throw new Error('Missing echo result.') + const callResult = frame.result as { result?: unknown } + const mcpResult = callResult.result as { + content?: Array<{ type?: string; text?: string }> + } + return mcpResult.content?.[0]?.text ?? String(callResult.result) +} + +function startCancellationHttpFixture(): { + url: string + sessionIds: Array + stop: () => void +} { + type Pending = { + scenario: CancellationScenario + resolve: (response: Response) => void + timer: Timer + } + const pending = new Map() + const sessionIds: Array = [] + const server = Bun.serve({ + hostname: '127.0.0.1', + port: 0, + async fetch(request) { + if (request.method !== 'POST') return new Response(null, { status: 405 }) + const message = (await request.json()) as { + id?: string | number + method: string + params?: Record + } + sessionIds.push(request.headers.get('mcp-session-id')) + if (message.method === 'initialize') { + return rpcResponse(message.id, { + protocolVersion: message.params?.protocolVersion, + capabilities: { tools: {} }, + serverInfo: { name: 'local-http-fixture', version: '1.0.0' }, + }) + } + if (message.method === 'notifications/initialized') { + return acceptedResponse() + } + if (message.method === 'notifications/cancelled') { + const requestId = message.params?.requestId as string | number + const entry = pending.get(requestId) + if (entry?.scenario === 'acknowledge' || entry?.scenario === 'race') { + clearTimeout(entry.timer) + pending.delete(requestId) + entry.resolve(rpcResponse(requestId, toolResult(entry.scenario))) + } + return acceptedResponse() + } + if (message.method !== 'tools/call') return acceptedResponse() + if (message.params?.name === 'echo') { + return rpcResponse(message.id, toolResult('echo-ok')) + } + + const id = message.id + if (id === undefined) return acceptedResponse() + const scenario = message.params?.arguments + ?.scenario as CancellationScenario + if (scenario === 'complete-before-cancel') { + await Bun.sleep(5) + return rpcResponse(id, toolResult(scenario)) + } + return new Promise((resolve) => { + const timer = setTimeout( + () => { + pending.delete(id) + resolve(rpcResponse(id, toolResult(scenario))) + }, + scenario === 'ignore' ? 80 : 2_000, + ) + pending.set(id, { scenario, resolve, timer }) + }) + }, + }) + return { + url: `http://127.0.0.1:${server.port}/mcp`, + sessionIds, + stop: () => server.stop(true), + } +} + +function acceptedResponse(): Response { + return new Response(null, { + status: 202, + headers: { 'mcp-session-id': 'local-session' }, + }) +} + +function rpcResponse(id: unknown, result: unknown): Response { + return Response.json( + { jsonrpc: '2.0', id, result }, + { headers: { 'mcp-session-id': 'local-session' } }, + ) +} + +function toolResult(text: string): Record { + return { content: [{ type: 'text', text }] } +} + +function fakeConnection( + callTool: (...args: any[]) => Promise, + onClose: () => void = () => {}, +): McpConnection { + return { + client: { callTool } as McpConnection['client'], + close: async () => onClose(), + pid: () => null, + stderr: null, + sessionId: () => 'fixture-session', + updateHeaders: () => {}, + } +} + +async function waitFor(predicate: () => boolean): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (predicate()) return + await Bun.sleep(1) + } + throw new Error('Condition was not met.') +} diff --git a/tests/runtime-stores.test.ts b/tests/runtime-stores.test.ts new file mode 100644 index 0000000..38055d7 --- /dev/null +++ b/tests/runtime-stores.test.ts @@ -0,0 +1,334 @@ +import { afterEach, describe, expect, it } from 'bun:test' +import fs from 'node:fs/promises' +import { tmpdir } from 'node:os' +import path from 'node:path' + +import { __test, openRuntimeStores } from '../src/runtime-stores' + +describe('Runtime state stores', () => { + const roots: string[] = [] + + afterEach(async () => { + await Promise.all( + roots.splice(0).map((root) => fs.rm(root, { recursive: true })), + ) + }) + + it('migrates declarations, credentials, and cached schemas into separate stores', async () => { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-stores-')) + roots.push(root) + await fs.writeFile( + path.join(root, 'servers.json'), + JSON.stringify({ + version: 1, + servers: { + fixture: { + url: 'https://fixture.example/mcp', + headers: { 'x-api-key': 'header-secret' }, + auth: { + kind: 'bearer', + strategy: 'round-robin', + confidence: 'configured', + credentials: [ + { kind: 'literal', value: 'bearer-secret' }, + { kind: 'env', name: 'FIXTURE_TOKEN' }, + ], + }, + discoveredAt: '2026-08-01T00:00:00.000Z', + refreshStatus: { + checkedAt: '2026-08-02T00:00:00.000Z', + status: 'ok', + }, + tools: [{ name: 'search.items', description: 'Search items' }], + }, + stdio: { + transport: 'stdio', + command: 'fixture-command', + env: { FIXTURE_SECRET: 'stdio-secret' }, + }, + }, + }), + ) + await fs.writeFile( + path.join(root, 'tokens.json'), + JSON.stringify({ + version: 1, + oauth: { + 'fixture:issuer': { + accessToken: 'oauth-secret', + tokenType: 'bearer', + }, + }, + oauthClientSecrets: { fixture: 'client-secret' }, + }), + ) + + const stores = await openRuntimeStores(root) + + expect(await stores.registry.read()).toEqual({ + version: 2, + servers: { + fixture: { + url: 'https://fixture.example/mcp', + auth: { + kind: 'bearer', + strategy: 'round-robin', + confidence: 'configured', + credentials: [ + { kind: 'stored', key: 'fixture:bearer:0' }, + { kind: 'env', name: 'FIXTURE_TOKEN' }, + ], + }, + }, + stdio: { transport: 'stdio', command: 'fixture-command' }, + }, + }) + expect(await stores.credentials.read()).toEqual({ + version: 2, + oauth: { + 'fixture:issuer': { + accessToken: 'oauth-secret', + tokenType: 'bearer', + }, + }, + oauthClientSecrets: { fixture: 'client-secret' }, + bearer: { 'fixture:bearer:0': 'bearer-secret' }, + headers: { fixture: { 'x-api-key': 'header-secret' } }, + stdioEnv: { stdio: { FIXTURE_SECRET: 'stdio-secret' } }, + }) + expect(await stores.schemas.read()).toEqual({ + version: 1, + servers: { + fixture: { + discoveredAt: '2026-08-01T00:00:00.000Z', + refreshStatus: { + checkedAt: '2026-08-02T00:00:00.000Z', + status: 'ok', + }, + tools: [ + { + name: 'search.items', + commandName: 'search-items', + description: 'Search items', + }, + ], + }, + }, + }) + + const snapshot = await stores.readSnapshot() + expect(snapshot.servers.fixture?.tools?.[0]?.commandName).toBe( + 'search-items', + ) + expect(JSON.stringify(snapshot)).not.toContain('secret') + }) + + it('serializes concurrent state updates without losing identities', async () => { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-stores-')) + roots.push(root) + const stores = await openRuntimeStores(root) + + await Promise.all( + Array.from({ length: 12 }, (_, index) => + stores.updateState(async (state) => { + await Bun.sleep(index % 3) + state.credentials.bearer[`identity-${index}`] = `value-${index}` + }), + ), + ) + + expect(Object.keys((await stores.credentials.read()).bearer)).toHaveLength( + 12, + ) + }) + + it('retains shared OAuth material until its final declaration is removed', async () => { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-stores-')) + roots.push(root) + const stores = await openRuntimeStores(root) + await stores.updateState((state) => { + state.registry.servers.first = { + url: 'http://127.0.0.1:1/mcp', + auth: { + kind: 'oauth-token', + tokenKey: 'shared', + confidence: 'confirmed', + }, + } + state.registry.servers.second = { + url: 'http://127.0.0.1:1/mcp', + auth: { + kind: 'oauth-token', + tokenKey: 'shared', + confidence: 'confirmed', + }, + } + state.credentials.oauth.shared = { + accessToken: 'shared-secret', + tokenType: 'bearer', + clientId: 'client', + clientSecretKey: 'client-secret-key', + } + state.credentials.oauthClientSecrets['client-secret-key'] = + 'client-secret' + }) + + await stores.removeServers(['first']) + expect((await stores.credentials.read()).oauth.shared).toBeDefined() + await stores.removeServers(['second']) + const credentials = await stores.credentials.read() + expect(credentials.oauth.shared).toBeUndefined() + expect(credentials.oauthClientSecrets['client-secret-key']).toBeUndefined() + }) + + it('reopens the published state without reviving a legacy registry', async () => { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-stores-')) + roots.push(root) + await fs.writeFile( + path.join(root, 'servers.json'), + JSON.stringify({ + version: 1, + servers: { + original: { + url: 'https://original.example/mcp', + auth: { kind: 'none' }, + }, + }, + }), + ) + await openRuntimeStores(root) + await fs.writeFile( + path.join(root, 'servers.json'), + JSON.stringify({ + version: 1, + servers: { + stale: { + url: 'https://stale.example/mcp', + auth: { kind: 'none' }, + }, + }, + }), + ) + + const reopened = await openRuntimeStores(root) + + expect(Object.keys((await reopened.registry.read()).servers)).toEqual([ + 'original', + ]) + expect(await fileExists(path.join(root, 'servers.json'))).toBe(false) + expect(await fileExists(path.join(root, 'servers.v1.backup.json'))).toBe( + true, + ) + }) + + it('recovers from an unpublished migration staging directory', async () => { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-stores-')) + roots.push(root) + await fs.writeFile( + path.join(root, 'servers.json'), + JSON.stringify({ + version: 1, + servers: { + fixture: { + url: 'https://fixture.example/mcp', + auth: { kind: 'none' }, + }, + }, + }), + ) + const interrupted = path.join(root, '.state-v2.interrupted') + await fs.mkdir(interrupted) + await fs.writeFile( + path.join(interrupted, 'registry.json'), + '{"partial":true}', + ) + + const stores = await openRuntimeStores(root) + + expect(Object.keys((await stores.registry.read()).servers)).toEqual([ + 'fixture', + ]) + expect(await fileExists(interrupted)).toBe(false) + }) + + it('finishes a journaled multi-store commit after an interrupted publication', async () => { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-stores-')) + roots.push(root) + await openRuntimeStores(root) + const stateDir = path.join(root, 'state-v2') + await fs.writeFile( + path.join(stateDir, 'transaction.json'), + JSON.stringify({ + registry: { + version: 2, + servers: { recovered: { transport: 'stdio', command: 'fixture' } }, + }, + credentials: { + version: 2, + oauth: {}, + oauthClientSecrets: {}, + bearer: {}, + headers: {}, + stdioEnv: { recovered: { LOCAL_SECRET: 'secret' } }, + }, + schemas: { version: 1, servers: {} }, + }), + ) + await fs.writeFile( + path.join(stateDir, 'registry.json'), + JSON.stringify({ version: 2, servers: {} }), + ) + + const recovered = await openRuntimeStores(root) + + expect((await recovered.registry.read()).servers.recovered).toBeDefined() + expect((await recovered.credentials.read()).stdioEnv.recovered).toEqual({ + LOCAL_SECRET: 'secret', + }) + expect(await fileExists(path.join(stateDir, 'transaction.json'))).toBe( + false, + ) + }) + + it('replays a publication failure before allowing another state generation', async () => { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-stores-')) + roots.push(root) + const stores = await openRuntimeStores(root) + const state = await stores.readState() + state.registry.servers.replayed = { transport: 'stdio', command: 'fixture' } + state.credentials.stdioEnv.replayed = { LOCAL_SECRET: 'secret' } + const stateDir = path.join(root, 'state-v2') + const paths = { + registry: path.join(stateDir, 'registry.json'), + credentials: path.join(stateDir, 'credentials.json'), + schemas: path.join(stateDir, 'schema-cache.json'), + transaction: path.join(stateDir, 'transaction.json'), + } + let attempts = 0 + + await __test.commitRuntimeState(paths, state, async (target, value) => { + attempts += 1 + await fs.writeFile(target.registry, JSON.stringify(value.registry)) + if (attempts === 1) throw new Error('injected publication failure') + await Promise.all([ + fs.writeFile(target.credentials, JSON.stringify(value.credentials)), + fs.writeFile(target.schemas, JSON.stringify(value.schemas)), + ]) + }) + + expect(attempts).toBe(2) + const reopened = await openRuntimeStores(root) + expect((await reopened.registry.read()).servers.replayed).toBeDefined() + expect((await reopened.credentials.read()).stdioEnv.replayed).toEqual({ + LOCAL_SECRET: 'secret', + }) + expect(await fileExists(paths.transaction)).toBe(false) + }) +}) + +async function fileExists(filePath: string): Promise { + return fs + .access(filePath) + .then(() => true) + .catch(() => false) +} diff --git a/tests/runtime.test.ts b/tests/runtime.test.ts new file mode 100644 index 0000000..8d909d8 --- /dev/null +++ b/tests/runtime.test.ts @@ -0,0 +1,86 @@ +import { afterEach, describe, expect, it } from 'bun:test' +import fs from 'node:fs/promises' +import { tmpdir } from 'node:os' +import path from 'node:path' + +import { McpRuntime } from '../src/runtime' +import { createInMemoryRuntimeCaller } from '../src/runtime-caller' +import { openRuntimeStores } from '../src/runtime-stores' + +describe('MCP Runtime', () => { + const roots: string[] = [] + + afterEach(async () => { + await Promise.all( + roots.splice(0).map((root) => fs.rm(root, { recursive: true })), + ) + }) + + it('returns a secret-free registry snapshot through the caller seam', async () => { + const runtime = new McpRuntime( + await storesFor({ + url: 'http://127.0.0.1:1/mcp', + headers: { 'x-api-key': 'secret' }, + auth: { kind: 'none' }, + }), + ) + const caller = createInMemoryRuntimeCaller('snapshot') + + await runtime.handle( + { requestId: 'snapshot', op: 'registrySnapshot' }, + caller, + ) + + expect(caller.frames[0]).toMatchObject({ + requestId: 'snapshot', + kind: 'result', + }) + expect(JSON.stringify(caller.frames)).not.toContain('secret') + }) + + it('returns reauth-required without attempting an ordinary OAuth flow', async () => { + const runtime = new McpRuntime( + await storesFor({ + url: 'http://127.0.0.1:1/mcp', + auth: { + kind: 'oauth-token', + tokenKey: 'fixture:http://127.0.0.1:1', + confidence: 'confirmed', + }, + }), + ) + const caller = createInMemoryRuntimeCaller('call') + + await runtime.handle( + { + requestId: 'call', + op: 'call', + serverName: 'fixture', + toolName: 'echo', + input: {}, + }, + caller, + ) + + expect(caller.frames).toEqual([ + { + requestId: 'call', + kind: 'error', + error: { + code: 'reauth-required', + message: 'Credentials for fixture must be refreshed.', + }, + }, + ]) + }) + + async function storesFor(server: Record) { + const root = await fs.mkdtemp(path.join(tmpdir(), 'mcpx-runtime-')) + roots.push(root) + await fs.writeFile( + path.join(root, 'servers.json'), + JSON.stringify({ version: 1, servers: { fixture: server } }), + ) + return openRuntimeStores(root) + } +}) diff --git a/tests/schema-refresh.test.ts b/tests/schema-refresh.test.ts deleted file mode 100644 index a7d283b..0000000 --- a/tests/schema-refresh.test.ts +++ /dev/null @@ -1,114 +0,0 @@ -import { describe, expect, it } from 'bun:test' - -import { - buildRefreshSummary, - hasStaleSchemas, - isSchemaRefreshStale, -} from '../src/schema-refresh' - -describe('schema refresh', () => { - const now = new Date('2026-05-15T00:00:00.000Z') - - it('marks cached schemas stale after one day', () => { - expect( - isSchemaRefreshStale( - { - url: 'https://mcp.example.com/mcp', - auth: { kind: 'none' }, - discoveredAt: '2026-05-13T23:59:59.000Z', - tools: [{ name: 'search', commandName: 'search' }], - }, - now, - ), - ).toBe(true) - }) - - it('does not refresh missing schemas through the background stale path', () => { - expect( - isSchemaRefreshStale( - { - url: 'https://mcp.example.com/mcp', - auth: { kind: 'none' }, - }, - now, - ), - ).toBe(false) - }) - - it('detects a registry with at least one stale server', () => { - expect( - hasStaleSchemas( - { - version: 1, - servers: { - fresh: { - url: 'https://fresh.example.com/mcp', - auth: { kind: 'none' }, - discoveredAt: '2026-05-14T12:00:00.000Z', - tools: [{ name: 'search', commandName: 'search' }], - }, - stale: { - url: 'https://stale.example.com/mcp', - auth: { kind: 'none' }, - discoveredAt: '2026-05-13T00:00:00.000Z', - tools: [{ name: 'search', commandName: 'search' }], - }, - }, - }, - now, - ), - ).toBe(true) - }) - - it('summarizes refresh status as server name lists', () => { - expect( - buildRefreshSummary([ - { - server: 'changed', - status: 'schema-refreshed', - toolsBefore: 1, - toolsAfter: 2, - schemaChanged: true, - }, - { - server: 'same', - status: 'schema-refreshed', - toolsBefore: 2, - toolsAfter: 2, - schemaChanged: false, - }, - { - server: 'token', - status: 'auth-refreshed', - toolsBefore: 3, - toolsAfter: 3, - schemaChanged: false, - }, - { - server: 'reauth', - status: 'reauthenticated', - toolsBefore: 3, - toolsAfter: 4, - schemaChanged: true, - }, - { - server: 'expired', - status: 'reauth-required', - toolsBefore: 4, - }, - { - server: 'down', - status: 'unreachable', - toolsBefore: 5, - }, - ]), - ).toMatchObject({ - refreshed: ['changed', 'reauth'], - unchanged: ['same', 'token'], - authRefreshed: ['token'], - reauthenticated: ['reauth'], - reauthRequired: ['expired'], - unreachable: ['down'], - }) - }) -}) diff --git a/tests/skill-command.test.ts b/tests/skill-command.test.ts index 96f7e6a..3e5b952 100644 --- a/tests/skill-command.test.ts +++ b/tests/skill-command.test.ts @@ -3,29 +3,19 @@ import { mkdtemp, readFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' -import type { ProjectService } from '../src/project-service' +import type { RegistryView } from '../src/skill-command' import { runSkillCommand } from '../src/skill-command' -function fixtureService(): ProjectService { +function fixtureService(): RegistryView { return { - config: { - version: 1, - servers: { - slack: { - transport: 'stdio', - command: 'slack-mcp', - tools: [], - }, + servers: { + slack: { + transport: 'stdio', + command: 'slack-mcp', + tools: [], }, }, - ensureServerReady: async () => { - throw new Error('not used') - }, - reauthenticateServer: async () => { - throw new Error('not used') - }, - save: async () => {}, } } diff --git a/tests/token-cache.test.ts b/tests/token-cache.test.ts deleted file mode 100644 index 49c49b4..0000000 --- a/tests/token-cache.test.ts +++ /dev/null @@ -1,37 +0,0 @@ -import { describe, expect, it } from 'bun:test' - -import type { TokenCache } from '../src/types' - -import { removeOAuthTokenFromCache } from '../src/token-cache' - -describe('token cache', () => { - it('removes an oauth token from cache', () => { - const cache: TokenCache = { - version: 1, - oauth: { - posthog: { - accessToken: 'secret', - tokenType: 'Bearer', - clientSecretKey: 'oauth-client:posthog', - }, - }, - oauthClientSecrets: { - 'oauth-client:posthog': 'client-secret', - }, - } - - expect(removeOAuthTokenFromCache(cache, 'posthog')).toBe(true) - expect(cache.oauth.posthog).toBeUndefined() - expect(cache.oauthClientSecrets?.['oauth-client:posthog']).toBeUndefined() - }) - - it('keeps cache unchanged when removing an unknown token', () => { - const cache: TokenCache = { - version: 1, - oauth: {}, - } - - expect(removeOAuthTokenFromCache(cache, 'posthog')).toBe(false) - expect(cache.oauth).toEqual({}) - }) -})