diff --git a/crates/agent-core/ARCHITECTURE.md b/crates/agent-core/ARCHITECTURE.md index e5f77d99..cdb7efff 100644 --- a/crates/agent-core/ARCHITECTURE.md +++ b/crates/agent-core/ARCHITECTURE.md @@ -62,7 +62,12 @@ builder on `Agent` or `ModelRequest`, and each is exercised by unit tests): A queued message is a `SteeringMessage` (text plus optional image attachments), not a bare `String` — the same shape a fresh `prompt` accepts, so `steer`/`follow_up` aren't a lesser channel that silently drops attachments a `prompt` would have kept. `From<&str>`/`From` build a text-only message, - so every plain-text `push`/`push_steer` call site keeps compiling unchanged. On drain, a follow-up + so every plain-text `push`/`push_steer` call site keeps compiling unchanged. A message may carry an + opaque `tag` (`SteeringMessage::with_tag`); the `Steered` event that reports its injection lists the + tags injected (`AgentEvent::Steered::tags`, not serialized), emitted after the transcript holding + them is checkpointed on both paths — so a host learns exactly which of its queued messages reached + the model, whatever later happens to the transcript (`serve` records MCP Events delivery this way). + On drain, a follow-up with images becomes a real `Message::user_with_images` turn instead of a plain-text one; a mid-run steer's images are appended as `ContentBlock::Image` blocks after its text block, onto the same tool-results turn its text rides on. diff --git a/crates/agent-core/src/agent.rs b/crates/agent-core/src/agent.rs index ed0d4059..e4ce1fd8 100644 --- a/crates/agent-core/src/agent.rs +++ b/crates/agent-core/src/agent.rs @@ -66,8 +66,15 @@ pub enum AgentEvent { }, /// One model turn completed. TurnEnd { stop_reason: StopReason, step: u32 }, - /// Queued steering/follow-up messages were injected at a stop boundary; the run continues. - Steered { messages: usize }, + /// Queued steering/follow-up messages were injected (folded into a tool-results turn, or at a + /// stop boundary) and the transcript holding them checkpointed; the run continues. `tags` are + /// the [`crate::steering::SteeringMessage::tag`]s of the messages injected, for the host that + /// queued them (not serialized). + Steered { + messages: usize, + #[serde(skip)] + tags: Vec, + }, /// The run finished normally (the model ended its turn and no steering was queued). AgentEnd { steps: u32 }, /// A compaction round has begun — [`Agent::compact`] confirmed a worthwhile prefix exists and is @@ -1250,6 +1257,7 @@ impl Agent { return Ok(()); } let count = injected.len(); + let tags: Vec = injected.iter().filter_map(|m| m.tag).collect(); for msg in injected { if msg.images.is_empty() { session.user(msg.text); @@ -1265,10 +1273,16 @@ impl Agent { if let Some(context) = steering.take_turn_context() { session.attach_request_block(context); } - sink(AgentEvent::Steered { messages: count }); // A plain user message ends the visible history here — a valid, resumable checkpoint - // (see `CheckpointHook`) before the next model call. + // (see `CheckpointHook`) before the next model call. Checkpointed *before* the + // `Steered` event, as on the tool-results path below: a host that records delivery + // on that event (by its tags) must find the messages already in the persisted + // transcript. self.checkpoint_guarded(session).await; + sink(AgentEvent::Steered { + messages: count, + tags, + }); continue; } @@ -1811,6 +1825,7 @@ impl Agent { // lane, injected only at the stop boundary below. let steered = steering.drain_steer(); let steered_count = steered.len(); + let steered_tags: Vec = steered.iter().filter_map(|m| m.tag).collect(); for msg in steered { result_blocks.push(ContentBlock::text(msg.text)); for source in msg.images { @@ -1828,6 +1843,14 @@ impl Agent { // both committed now, so this is a valid, resumable checkpoint (see `CheckpointHook`) — the // one mid-run point a crash between here and the run's eventual end would otherwise lose. self.checkpoint_guarded(session).await; + // Reported before any way out of the loop below: the steered messages are in the + // checkpointed transcript now, whether or not the run goes on. + if steered_count > 0 { + sink(AgentEvent::Steered { + messages: steered_count, + tags: steered_tags, + }); + } if terminate { // A tool requested completion (e.g. an `attempt_completion`/`exit` tool) and the whole // batch agreed. The results are already recorded; end the run as if the model had @@ -1838,11 +1861,6 @@ impl Agent { }); return Ok(()); } - if steered_count > 0 { - sink(AgentEvent::Steered { - messages: steered_count, - }); - } // A graceful-stop request is honored here too, after this turn's tool results (and any // folded-in steer text) are already committed — the same turn-boundary contract as the // tool-less branch above. Checked *after* the `Steered` event so a client sees its steer @@ -4898,6 +4916,57 @@ mod tests { )); } + /// A steer folded into the tool-results turn of a batch that ends the run (`terminate`) did + /// reach the model's transcript: its `Steered` event, with its tag, is still reported. + #[tokio::test] + async fn a_steer_folded_into_a_terminating_batch_is_still_reported() { + struct ExitTool; + #[async_trait] + impl Tool for ExitTool { + fn name(&self) -> &str { + "exit" + } + fn description(&self) -> &str { + "End the run." + } + fn input_schema(&self) -> Value { + serde_json::json!({ "type": "object" }) + } + async fn run( + &self, + _: Value, + ) -> std::result::Result { + Ok(crate::tool::ToolOutput::text("done").with_terminate(true)) + } + } + let mut tools = ToolRegistry::new(); + tools.register(Arc::new(ExitTool)); + let (agent, _mock) = agent_with(vec![turn::tool_call("tu_1", "exit", "{}")], tools); + let mut session = Session::new(); + session.user("finish up"); + let steering = Steering::new(); + steering.push_steer(SteeringMessage::new("one more thing", Vec::new()).with_tag(42)); + let mut tags = Vec::new(); + agent + .run_events_steered( + &mut session, + |ev| { + if let AgentEvent::Steered { tags: t, .. } = ev { + tags.extend(t); + } + }, + CancellationToken::new(), + steering, + ) + .await + .unwrap(); + assert_eq!( + tags, + vec![42], + "the steer reached the transcript and is reported" + ); + } + #[tokio::test] async fn a_terminate_request_only_wins_when_every_call_in_the_batch_agrees() { // pi-parity coverage (`packages/agent/test/agent-loop.test.ts`, "should continue after @@ -9584,6 +9653,43 @@ mod tests { ); } + /// A queued message's tag comes back in the `Steered` event that reports it reached the model — + /// mid-run (folded into a tool-results turn) and at a stop boundary alike — so the host that + /// queued it knows exactly what was received, whatever later happens to the transcript. + #[tokio::test] + async fn a_steered_message_reports_its_tag_when_it_reaches_the_model() { + let mut tools = ToolRegistry::new(); + tools.register(Arc::new(EchoTool)); + let (agent, _mock) = agent_with( + vec![ + turn::tool_call("tu_1", "echo", r#"{"text":"pong"}"#), + turn::text("acknowledged"), + turn::text("and again"), + ], + tools, + ); + let mut session = Session::new(); + session.user("start"); + let steering = Steering::new(); + steering.push_steer(SteeringMessage::new("mid-run", Vec::new()).with_tag(7)); + steering.push(SteeringMessage::new("at the stop", Vec::new()).with_tag(8)); + let mut tags = Vec::new(); + agent + .run_events_steered( + &mut session, + |ev| { + if let AgentEvent::Steered { tags: t, .. } = ev { + tags.extend(t); + } + }, + CancellationToken::new(), + steering, + ) + .await + .unwrap(); + assert_eq!(tags, vec![7, 8]); + } + #[tokio::test] async fn steering_is_injected_mid_run_between_tool_turns() { // A steering message queued while the agent is mid-tool-call must be folded into the *same* diff --git a/crates/agent-core/src/steering.rs b/crates/agent-core/src/steering.rs index faeb28a2..022fae9a 100644 --- a/crates/agent-core/src/steering.rs +++ b/crates/agent-core/src/steering.rs @@ -80,6 +80,10 @@ use crate::tool::ToolRegistry; pub struct SteeringMessage { pub text: String, pub images: Vec, + /// An opaque id the queuer chose, reported back in [`crate::AgentEvent::Steered`] at the moment + /// the message reaches the model — so a host can tell exactly which of its queued messages the + /// model received, whatever later happens to the transcript (a compaction, say). + pub tag: Option, } impl SteeringMessage { @@ -88,8 +92,15 @@ impl SteeringMessage { Self { text: text.into(), images, + tag: None, } } + + /// The same message, tagged (see [`Self::tag`]). + pub fn with_tag(mut self, tag: u64) -> Self { + self.tag = Some(tag); + self + } } impl From<&str> for SteeringMessage { @@ -97,6 +108,7 @@ impl From<&str> for SteeringMessage { Self { text: text.to_string(), images: Vec::new(), + tag: None, } } } @@ -106,6 +118,7 @@ impl From for SteeringMessage { Self { text, images: Vec::new(), + tag: None, } } } diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 911b0ced..8031a60f 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -2194,7 +2194,7 @@ the HTML arrives, so a slow tool renders in progress; `mcp_app_result` carries t `CallToolResult`. The model gets the same text it always did (`tool_output_from_result` reads only `content`), so `structuredContent` and `_meta` reach the view and never model context. **The model's result never waits for its view**: when the call finishes first it returns at once, and the view -(if it loads within 10 s) follows as `mcp_app_open` then `mcp_app_result`. A view over 4 MiB (`MAX_VIEW_BYTES`) — refused **unread** when the server advertised its `size` in `resources/list` (recorded in both manifests), refused at the transport when it did not, without the whole message ever being held: the stdio pump (`mcp_stdio::read_capped`) caps every inbound message at 64 MiB (`BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES`) and the answer to a `ui://` `resources/read` it wrote at `MAX_VIEW_BYTES`, keyed by that request's id. The id is read by walking the top-level object's members in the message's first 4 KiB (`mcp_stdio::scan_head` — never a text search, which a nested `"id":` would match), and an over-cap message is answered with a JSON-RPC error only when that prefix proves it is a response to a host request (a top-level `id` before a top-level `result`/`error`, and no `method`); anything else over the cap — a server→client request or notification (its id is the server's and can equal a host request's), or a message whose id comes after its body — is dropped with a warning, and its request, if any, times out; over streamable HTTP a view read goes through `mcp_wire::HttpClient::post_bounded` — the request path `skills/*` takes, with rmcp's own status handling (401/403 `WWW-Authenticate` → `AuthRequired`/`InsufficientScope`, a 404 on a session → `SessionExpired`) and an SSE response handed to rmcp as a stream, so every other event on it is routed — refused on its `Content-Length` before the body, or as soon as the streamed body or one SSE event passes the cap (the event replaced by the refusal and the stream ended); other HTTP requests keep rmcp's client — with the wrong mime type, or missing costs only its frames — the spec's own text +(if it loads within 10 s) follows as `mcp_app_open` then `mcp_app_result`. A view over 4 MiB (`MAX_VIEW_BYTES`) — refused **unread** when the server advertised its `size` in `resources/list` (recorded in both manifests), refused at the transport when it did not, without the whole message ever being held: the stdio pump (`mcp_stdio::read_capped`) caps every inbound message at 64 MiB (`BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES`) and the answer to a `ui://` `resources/read` it wrote at `MAX_VIEW_BYTES`, keyed by that request's id. The id is read by walking the top-level object's members in the message's first 4 KiB (`mcp_stdio::scan_head` — never a text search, which a nested `"id":` would match), and an over-cap message is answered with a JSON-RPC error only when that prefix proves it is a response to a host request (a top-level `id` before a top-level `result`/`error`, and no `method`); anything else over the cap — a server→client request or notification (its id is the server's and can equal a host request's), or a message whose id comes after its body — is dropped with a warning, and its request, if any, times out; over streamable HTTP a view read goes through `mcp_wire::HttpClient::post_bounded` — the request path `skills/*` takes, with rmcp's own status handling (401/403 `WWW-Authenticate` → `AuthRequired`/`InsufficientScope`, a 404 on a session → `SessionExpired`) and an SSE response handed to rmcp as a stream, so every other event on it is routed — refused on its `Content-Length` before the body, or as soon as the streamed body or one SSE event passes the cap (the event replaced by the refusal and the stream ended); every other streamable-HTTP POST goes through `post_bounded` too, under the general per-message cap (rmcp's own `post_message` is no longer used — see `mcp_wire.rs` below) — with the wrong mime type, or missing costs only its frames — the spec's own text fallback. The CSP of every view that opens is logged (spec SHOULD). An aborted call sends `mcp_app_result{cancelled}`; its reason goes through service mode's path redaction. @@ -4126,7 +4126,7 @@ backend an operator points it at. | `src/bin/mcp_fixture_events_server.rs` | Hand-rolled MCP Events **server** fixture (stdio or streamable HTTP) with a control API for tests: emits events, signs webhook deliveries (Standard Webhooks, dual-signing after a rotation, tamper modes), runs the verification challenge, grants TTLs, terminates subscriptions, and records what it observed | | `src/tools/mcp_skills.rs` | SEP-2640 Skills extension, client side: entry validation (incl. the 512-file / 16 MiB per-skill limits), `ServerSkills` (per connection: listing with `ttlMs`/list-changed refresh, learned entries), `SkillSession` (per session, on `McpEnabledSet`: loaded skills, content-bound approvals, the activation and code-execution approval gate, `--approve-mcp-skills`), the per-server loading tool `mcp____skill__read` rebound to each session, `RunHooks`, and `McpSkills` — the session's gated view: ``, `get_commands` entries, diagnostics, `/skill::` expansion. See "MCP skills (SEP-2640)" | | `src/tools/mcp_view_http.rs` | The streamable-HTTP client an apps connection dials through: `ViewCappedHttp` sends a `ui://` `resources/read` through `mcp_wire::HttpClient::post_bounded` (rmcp's auth and session status handling, SSE handed to rmcp as a stream) and refuses its response over `MAX_VIEW_BYTES` by `Content-Length` before the body, or as soon as the streamed body or one SSE event passes it; everything else goes to `mcp_wire::HttpClient` | -| `src/tools/mcp_wire.rs` | Works around rust-sdk#1197 over streamable HTTP: `HttpClient` (a `StreamableHttpClient` that answers `skills/*` POSTs itself — and MCP App view reads, for `mcp_view_http` — through `post_bounded`, applying `mcp_stdio::rescue` to the JSON body or to each SSE event, delegating the rest to `reqwest`). A unit test pins the upstream bug so the module can be deleted when it is fixed | +| `src/tools/mcp_wire.rs` | Works around rust-sdk#1197 over streamable HTTP, and caps every streamable-HTTP message: `HttpClient` (a `StreamableHttpClient` that answers **every** POST itself through `post_bounded` (rmcp's own `reqwest` `post_message` is never called: the request headers it sends — bearer, `Mcp-Session-Id`, `Accept`, protocol version, custom headers — and the response shapes it acts on — `401`/`403` `WWW-Authenticate`, a `404` on a session, `202`, JSON with its session id, SSE — are pinned against rmcp's client by the differential tests `an_ordinary_request_is_sent_as_rmcps_own_client_would` / `…_answered_as_rmcps_own_client_would`) — applying `mcp_stdio::rescue` to the JSON body or to each SSE event, and holding each to the one per-message cap, `mcp_stdio::max_message_bytes`: a request's over-cap answer becomes that request's JSON-RPC error; MCP App view reads, for `mcp_view_http`, get the tighter view cap; the GET stream is capped through `rmcp`'s SSE-event limit). A unit test pins the upstream bug | | `src/bin/mcp_skills_fixture_server.rs` | Hand-rolled SEP-2640 server for `tests/mcp_skills*.rs` (stdio, or `--http`): an ordinary skill with a supporting file and a nested skill, a same-name pair, a `disable-model-invocation` skill, a tampered one, a dynamic one, an unlisted one, two over the per-skill limits, and a late one published by its `publish_late` tool (with `notifications/resources/list_changed`); caching hints from env; `_meta` on every result like a `2026-07-28` server; logs every request so tests assert what was (not) fetched | | `benches/search.rs` | Criterion macro-bench: `grep` (1 vs auto threads) and `find` (sequential) over a 5,000-file tree | | `benches/serve_runtime.rs` | Process-level Tokio runtime A/B for `serve --listen`: spawns the real binary under `current_thread` vs multi-thread (`BEYOND_AI_AGENT_TOKIO_WORKER_THREADS`) and drives concurrent WebSocket sessions, reporting RSS, throughput, p95 latency, and CPU. Custom harness, not Criterion/divan — the cost is thread stacks and per-thread heaps, which an in-process alloc profiler never sees. | @@ -4178,6 +4178,8 @@ backend an operator points it at. | `--mcp-events-reapable` / `AI_AGENT_MCP_EVENTS_REAPABLE` (`serve` daemon) | off (a session with a live MCP Events subscription is exempt from the idle reaper) | Let `--session-idle-timeout` reap a detached session even while it holds live subscriptions (its teardown unsubscribes) | | `--mcp-events-session` / `AI_AGENT_MCP_EVENTS_SESSION` (`serve` daemon) | `mcp-events` | MCP Events: the daemon session that owns the subscriptions configured in `mcp_servers[].events`; started at boot, the only session configured events are delivered to | | `BEYOND_AI_AGENT_MCP_EVENTS_COALESCE_MS` / `_POLL_FLOOR_MS` / `_STREAM_IDLE_MS` / `_WEBHOOK_TTL_MS` / `_MAX_PENDING` / `_MAX_PENDING_BYTES` / `_STREAM_BUFFER` / `_JWKS_TIMEOUT_MS` / `_BODY_TIMEOUT_MS` / `_HEALTHY_MS` | `2000` / `1000` / `70000` / `3600000` / `1000` / `16777216` / `256` / `5000` / `10000` / `600000` | MCP Events tuning: minimum gap between event injections into a session; floor on `nextPollMs`; push-stream silence before a reconnect; the webhook `ttlMs` suggested (the server's grant wins); most undelivered events a session holds, by count and by stored bytes, before refusing new ones; notifications queued per push stream before it overflows and resumes; the deadline on one JWKS fetch (headers and body); the deadline on one webhook delivery's body; how long a runtime subscription must stay up for its re-discovery budget to be restored | +| `BEYOND_AI_AGENT_MCP_EVENTS_REFUSED_RETRY_MS` | `3600000` | MCP Events: how long a subscription the server refused for good (a permanent error) waits before it is tried again; it does not keep the session alive meanwhile | +| `BEYOND_AI_AGENT_MCP_EVENTS_RUNTIME_TTL_MS` / `_MAX_RESTORED_SESSIONS` | `604800000` / `32` | MCP Events runtime subscriptions: how long after its session's last client command one is still restored after a restart; how many such sessions a daemon starts at boot | | `--mcp-allow-private` / `AI_AGENT_MCP_ALLOW_PRIVATE` (`serve --service`) | off | Let a grant's MCP connectors reach loopback/private/link-local addresses. Off by default: a connector URL is refused unless it resolves to a public one, so a tenant cannot point an "MCP server" at the instance metadata service, a storage mount target or a peer replica. For a dev replica whose connectors run beside it (and this repo's own tests); no grant can ask for it | | `RUST_LOG` (`tracing_subscriber::EnvFilter::from_default_env`) | unset (no logs) | Verbosity of `tracing` spans/events emitted by the binary's subscriber | | `AI_AGENT_TIMING` | unset (no timing output) | `=1` prints a startup-timing breakdown (resource discovery, system-prompt build, session open, agent construction) to stderr just before the first turn/`ready` frame — pi's own `PI_TIMING=1`; every checkpoint is a zero-cost no-op when unset | @@ -4351,21 +4353,20 @@ transport makes goes through `OAuthHttp`, which attaches the server's _current_ carries no `auth_header` of its own); direct `events/*` POSTs (`mcp_events::wire`'s `Conn::send`) likewise read the current token per request, not the one their headers were built with. **One rule on both paths:** any **401** from a server with a login — with or without a `WWW-Authenticate` -challenge, whatever its body (for such a server every POST is answered by -`mcp_wire::HttpClient::post_bounded`, which decides from the status; rmcp's own client reads a 401 -carrying a JSON-RPC error body as an ordinary error response) — makes `ServerAuth::after_rejection` force a refresh through the same `AuthorizationManager` +challenge, whatever its body (every POST, for every server, is answered by +`mcp_wire::HttpClient::post_bounded`, which decides from the status — a deliberate divergence from rmcp's own client (see below), which reads a 401 carrying a JSON-RPC error body as an ordinary error +response; pinned in `an_ordinary_request_is_answered_as_rmcps_own_client_would`) — makes `ServerAuth::after_rejection` force a refresh through the same `AuthorizationManager` and `McpAuthStore` `mcp-login` uses (so the new token is persisted), and the request is retried **once**; a second 401 is returned as the server's answer. That covers `tools/call`, `resources/*`, `prompts/*`, `skills/*`, MCP App view reads, the handshake, the standalone stream and `events/*`. A 403 (`InsufficientScope`) never refreshes. -**`post_bounded` answers what rmcp's client would, with one deliberate difference.** Since an OAuth -server's POSTs no longer reach rmcp's client, `post_bounded` carries its semantics: `Mcp-Session-Id` +**`post_bounded` answers what rmcp's client would, with two deliberate differences.** Since no server's POSTs reach rmcp's client any more, `post_bounded` carries its semantics: `Mcp-Session-Id` in and out, a 404 on a session as `SessionExpired` (rmcp re-initializes and retries), an SSE response handed back as a stream, 202/204 — or an empty 200 to a notification or reply — as accepted, a JSON 200 that is not a JSON-RPC message, for a notification or reply, as accepted, a non-JSON success as an error (rmcp's unexpected content type), and a 4xx to `server/discover` as the legacy server's cue to -`initialize`. **The difference:** a JSON 200 that is not a JSON-RPC message, answering a _request_, +`initialize`. **The differences:** any 401 is `AuthRequired` (above; rmcp reads one with a JSON-RPC body as an error response — pinned in `an_ordinary_request_is_answered_as_rmcps_own_client_would`); and a JSON 200 that is not a JSON-RPC message, answering a _request_, is an error here. rmcp calls it accepted and then waits for a response that cannot come (a request is answered on its own POST, as JSON or an SSE stream carrying it) until the request's timeout; failing at once with the body in the error is the honest answer. Pinned end to end against the OAuth fixture @@ -4447,10 +4448,11 @@ mcp_events_subscribe (any session) ──► owned by that session their route. The token and the last confirmed secret persist in the state file, so after a restart the same callback URL is registered again before anything is sent — a delivery the server retries while the daemon was down (a crash leaves its subscription in place) still lands and verifies. - Persisted tokens of _configured_ subscriptions are held from the moment the state is read (up to - 10 min): a retry that arrives before its subscription has re-registered gets `503` (retry), never - `410` (which the draft treats as "stop"). A runtime subscription is not resubscribed after a - restart, so its token is not held and its old callback answers `410`. + Persisted tokens of every subscription the session will subscribe again — configured ones, and + runtime ones being restored (below) — are held from the moment the state is read (up to 10 min): a + retry that arrives before its subscription has re-registered gets `503` (retry), never `410` (which + the draft treats as "stop"). A runtime subscription the server ended is not restored, so its token + is not held and its old callback answers `410`. After a _graceful_ shutdown the subscription was unsubscribed, so an event type the server cannot replay from a cursor loses what it emitted while the daemon was down; one with replay resumes from the persisted cursor. Refresh runs at ¾ of the way to `refreshBefore` (an unparseable one refreshes after @@ -4472,19 +4474,21 @@ mcp_events_subscribe (any session) ──► owned by that session The coalescer commits before it injects. A batch leaves the pending queue only when the **model has received it** and the transcript holding it has been persisted (`McpEventsHub::finish_run`, from the `prompt` arm): the run's own prompt counts once the run did not fail (an abort still leaves the - prompt in the transcript the next run reads); a batch steered in mid-run counts only if it is in the - transcript, found by the batch id its text starts with (`[MCP events · batch N]`), so two batches - with the same content are told apart — an abort clears the steer lane - (`Steering::clear_run_scoped`) before the model sees what is queued there. A batch whose run failed, whose transcript did not persist, or whose + prompt in the transcript the next run reads); a batch steered in mid-run counts exactly when the + model receives it: the steer is queued with its batch as the message's tag + (`SteeringMessage::tag`), and the run's `AgentEvent::Steered` reports the tags it injected _after_ + checkpointing the transcript holding them — `mcp_events::Receipts::received` records the batch as + delivered (durably) right then. Nothing is re-derived from the transcript later, so a compaction + that summarizes the steered turn away does not make it look undelivered; an abort clears the steer + lane (`Steering::clear_run_scoped`) before the model sees what is queued there, so no `Steered` + reports it. A batch whose run failed, whose transcript did not persist, or whose steer was dropped goes back to pending and is injected again; after 3 such attempts an event is dropped and reported (`mcp_event_status` `kind: "dropped"`), so a poison event cannot re-run the model forever. A restart re-injects whatever is still pending. **Guarantee:** every acknowledged event reaches the model at least once (or is reported dropped after 3 attempts that never reached - it). It reaches it exactly once except in four cases, each of which re-injects the batch: a run - fails or is aborted before the model saw a steered batch (by design); a compaction during the run - summarizes a steered batch's turn away, so it is no longer found in the transcript (the model did - see it — this one is a duplicate); the process dies, or the `done` record fails to be written, - between a run's transcript persist and that record. A graceful exit writes the record first, so + it). It reaches it exactly once except in three cases, each of which re-injects the batch: a run + fails or is aborted before the model saw a steered batch (by design); the process dies, or the + `done` record fails to be written, between a turn's checkpoint and that record. A graceful exit writes the record first, so SIGTERM, stdin EOF and the reaper never duplicate (`tests/mcp_events_durability.rs`). - **Into the session.** Each subscription dedups by `eventId` (newest 1024, persisted), then every accepted occurrence is broadcast as `mcp_event`. Unless `action: "notify"`, it is queued for the @@ -4518,6 +4522,33 @@ mcp_events_subscribe (any session) ──► owned by that session `events/list` re-fetched). While one is not up it counts as live for the keep-alive: a daemon whose servers are all briefly down keeps its events session. An explicit `mcp_events_unsubscribe` stops keeping it up, including while a retry is waiting (membership is re-checked before every attempt). + A **permanent refusal** — the server is unknown (gone from settings), does not offer the event, + does not speak the extension (`-32601`), rejects the request as invalid (`-32602`, `-32011`, + `-32014` other than `schema_changed`), or no delivery mode fits — is different: it is reported as + `refused` (an `mcp_event_status` frame, and `mcp_events_list`'s `unestablished` list, which also + shows what is `starting` or `retrying`), retried only every + `BEYOND_AI_AGENT_MCP_EVENTS_REFUSED_RETRY_MS` (1 h), and does **not** count as live for the + keep-alive. **Forbidden** (`-32012`) is not permanent at first — an expired credential looks like + this: it is retried on a short backoff (1 s doubling to 5 s), each attempt dialing afresh so + credentials are re-resolved, and becomes a refusal only after 5 in a row. +- **Runtime subscriptions survive a restart.** A successful `mcp_events_subscribe` stores its whole + spec in the session's events state (`PersistedSub::runtime`). When the session starts again it + subscribes each one again (retried like a configured one until it is back, or until it is + explicitly unsubscribed); a daemon starts every session holding such a subscription at boot + (`mcp_events::sessions_with_runtime_subscriptions`, pinned and unpinned like the events session), + so its triggers resume with no client attached. An explicit unsubscribe forgets the spec, and so + does the server ending the subscription (unless re-discovery brings it back), and so does a restore + the server refuses for good (its server renamed or removed, say) — so a dead subscription cannot + pin its session or bring it back at every boot. Forgetting one forgets all of it — spec, cursor, + and webhook callback (token and secret) — so no credential outlives the subscription in the + snapshot. Runtime subscriptions also **expire**: a session's last client command is recorded + (coarsely, in the events state; creating a runtime subscription counts as one, so a subscription + made by a single command still has a start to its time to live), and once it is older than + `BEYOND_AI_AGENT_MCP_EVENTS_RUNTIME_TTL_MS` (7 days) its runtime subscriptions are forgotten + instead of restored, and a daemon no longer starts it at boot. A state with runtime subscriptions + and no record at all is not trusted: it is not restored. Boot restoration is bounded too: at most + `BEYOND_AI_AGENT_MCP_EVENTS_MAX_RESTORED_SESSIONS` (32) sessions are started — the ones a client + used most recently — and the rest come back when their clients do. - **Schema changes.** For a runtime subscription, a termination saying the event type was removed (`-32011`, `kind: event`) or changed in place (`-32014`, `reason: schema_changed`) re-fetches `events/list` and resubscribes — the draft's SHOULD — up to 5 times, with backoff; the budget is @@ -4603,7 +4634,24 @@ mcp_events_subscribe (any session) ──► owned by that session are up), and `main` returning — or a panic unwinding out of it — passes the `mcp_stdio::ExitSweep` guard's `Drop` (`tests/mcp_stdio_transport.rs`: `serve` on EOF and SIGTERM, `run` finishing, cancelled mid-turn and refused for bad configuration; the panic path by - `mcp_stdio::tests::a_panic_unwinding_past_the_exit_guard_still_sweeps`). The only exits that skip it + `mcp_stdio::tests::a_panic_unwinding_past_the_exit_guard_still_sweeps`). A panic out of the main + future does not reach that guard by dropping the runtime — that would wait for every blocking + task, and `serve`'s stdin reader is one, parked on a read that may never return, so the process + hung. `main` catches the unwind instead, abandons the runtime's threads + (`Runtime::shutdown_background`), sweeps, and exits `101` + (`a_panic_inside_serve_exits_promptly_and_still_sweeps`, through a debug-build-only `__test_panic` + command). In a daemon a session is a task, not the main future, so its panic is **contained**: + `serve_ws` catches the session future's unwind (in-flight runs and MCP calls are dropped, its + events hub cancelled and its state written by their `Drop`s), sends its clients an `error` frame, + ends the session, and keeps serving every other one — sound because sessions share only state + behind plain replace-the-value locks, and MCP server processes belong to the process, swept on its + exit (`a_panic_in_a_daemon_session_ends_that_session_only`). The one session the daemon depends on + — the MCP Events session, which runs the configured subscriptions — is started again after a + panic, with backoff (1 s doubling to a minute, reset after ten minutes up, per session), so + configured triggers do not stop until the next restart + (`the_events_session_is_restarted_after_a_panic`); so is any session holding runtime + subscriptions, which also run with no client attached + (`a_session_with_runtime_subscriptions_is_restarted_after_a_panic`). The only exits that skip it are argument-parsing failures, before anything is spawned. **The cost, by design:** a process that ran a stdio server exits a little later — measured at about 166 ms instead of 44 ms for a server that exits on EOF — because each server is given the chance to exit cleanly and is then swept, @@ -4618,10 +4666,20 @@ mcp_events_subscribe (any session) ──► owned by that session passes through the same `rescue`). Stateless (`2026-07-28`) streamable-HTTP servers get `events/*` directly over HTTP (`MCP-Protocol-Version`, `Mcp-Method`, per-request `_meta`, the server's resolved headers and OAuth bearer), with bodies - bounded (responses 4 MiB, SSE events 1 MiB, JWKS 64 KiB) and an SSE reader whose scan and drain + bounded by the same per-message cap as every MCP transport (`mcp_stdio::max_message_bytes`, `BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES`: a unary JSON body, a unary SSE answer's event, an `events/stream` event — an over-cap answer fails its request, never read whole). An over-cap **event notification** (on a push stream, or from a stdio server) is skipped, not reconnected into: its bounded head is read structurally (`mcp_stdio::oversized_stand_in`, the same top-level member walk as `scan_head`) for its routing, cursor and id, the rest is skipped unread, and a small `$oversized` stand-in takes its place — the subscription keeps that cursor (or the next heartbeat's), records an `oversized` gap (a `gap` frame, and a notice to the model that an event was dropped unread) and carries on (`tests/mcp_message_cap.rs`) — JWKS documents at 64 KiB — and an SSE reader whose scan and drain are both linear (events are parsed in place; the consumed prefix is dropped once per read, not once per event); an older, session-bound HTTP server goes through rmcp's own connection, which carries its `Mcp-Session-Id`. + A server→client request (an elicitation, a sampling request) raised during an `events/*` request + is attributed to the session that owns it over rmcp — the same rule as during a `tools/call` + (`track_call`) — and, over direct HTTP, where nothing could route it to a session, answered at once + with an error rather than left unanswered (`tests/mcp_events_nested.rs`). That answer is POSTed + with the same `Accept` (`application/json, text/event-stream`) as every events request — a + streamable-HTTP server refuses an answer without it `406` and keeps waiting (the fixture enforces + it) — and through the same `Conn::send`, so with an OAuth login it carries the server's current + token and a 401 refreshes and resends once (no `Mcp-Method`: an answer has no method) — and on an `events/stream` it is answered inline, one at a time, so a server flooding + requests onto the stream cannot make the client spawn without limit + (`a_flood_of_server_requests_on_an_events_stream_is_answered_one_at_a_time`). - **`run` and service mode.** One-shot `run` does not subscribe: it ends when the model stops, and an event is a trigger for a _future_ run. Service mode refuses the `mcp_events_*` commands outright (`service::refused_command`), grant connectors carry no `events`, and the webhook and diff --git a/crates/agent/src/bin/mcp_fixture_events_server.rs b/crates/agent/src/bin/mcp_fixture_events_server.rs index 8f21319c..e243c1b0 100644 --- a/crates/agent/src/bin/mcp_fixture_events_server.rs +++ b/crates/agent/src/bin/mcp_fixture_events_server.rs @@ -108,6 +108,62 @@ struct State { jwks_stall: bool, /// How many JWKS requests are stalled right now (they hold their connection open). jwks_stalled: u64, + /// `MCP_FIXTURE_NESTED_DURING=poll|stream`: the client's answers to the nested + /// `elicitation/create` raised during that `events/*` request. + nested_answers: Vec, + /// `MCP_FIXTURE_FORBID_FIRST=`: the first `n` `events/poll`/`events/subscribe` requests are + /// refused with `-32012` (forbidden) — credentials that are refreshed and then work. + forbid_left: u64, + /// Answers the client POSTed to requests this server raised over HTTP (`id` and no `method`). + client_answers: Vec, + /// Answer POSTs being handled right now, and the most at once (`MCP_FIXTURE_ANSWER_DELAY_MS` + /// holds each a while, so concurrent ones overlap). + answers_in_flight: u64, + answers_max_in_flight: u64, + /// Whether the HTTP nested request (`MCP_FIXTURE_NESTED_DURING`) has been raised. + http_nested_done: bool, +} + +/// The nested `elicitation/create` this server raises toward the client during an `events/*` +/// request (`MCP_FIXTURE_NESTED_DURING`). +fn nested_request(during: &str) -> Value { + json!({ + "jsonrpc": "2.0", + "id": format!("nested-{during}"), + "method": "elicitation/create", + "params": { + "mode": "form", + "message": format!("Approve during events/{during}?"), + "requestedSchema": { "type": "object", "properties": { "ok": { "type": "boolean" } } }, + }, + }) +} + +/// Over HTTP: wait for the client's answer to the nested request, and record it (or a timeout). +async fn await_http_answer(state: &Shared, during: &str) { + let id = json!(format!("nested-{during}")); + let deadline = std::time::Instant::now() + Duration::from_secs(20); + let answer = loop { + let found = state + .lock() + .unwrap() + .client_answers + .iter() + .find(|a| a["id"] == id) + .cloned(); + if let Some(a) = found { + break a; + } + if std::time::Instant::now() > deadline { + break json!({ "timeout": true }); + } + tokio::time::sleep(Duration::from_millis(50)).await; + }; + state + .lock() + .unwrap() + .nested_answers + .push(json!({ "during": during, "answer": answer })); } fn env_flag(name: &str) -> bool { @@ -444,16 +500,20 @@ async fn rpc( params: Value, principal: Option, ) -> Result { - let events_down = { + let (events_down, forbidden) = { let mut st = state.lock().unwrap(); st.methods.push(method.to_owned()); + let forbidden = matches!(method, "events/poll" | "events/subscribe") && st.forbid_left > 0; + if forbidden { + st.forbid_left -= 1; + } if params .pointer("/_meta/io.modelcontextprotocol~1clientCapabilities/extensions/io.modelcontextprotocol~1ui") .is_some() { st.ui_methods.push(method.to_owned()); } - st.events_down + (st.events_down, forbidden) }; let no_events = env_flag("MCP_FIXTURE_NO_EVENTS"); let capabilities = if no_events { @@ -479,7 +539,18 @@ async fn rpc( "name": "echo", "description": "Echoes back its `text` argument.", "inputSchema": { "type": "object", "properties": { "text": { "type": "string" } } } + }, { + "name": "blob", + "description": "Returns a text result of `bytes` bytes.", + "inputSchema": { "type": "object", "properties": { "bytes": { "type": "integer" } } } }] })), + "tools/call" if params["name"] == "blob" => { + let n = params + .pointer("/arguments/bytes") + .and_then(Value::as_u64) + .unwrap_or(0) as usize; + Ok(json!({ "content": [{ "type": "text", "text": "x".repeat(n) }], "isError": false })) + } "tools/call" => { let text = params .pointer("/arguments/text") @@ -493,12 +564,19 @@ async fn rpc( m if m.starts_with("events/") && events_down => { Err(err(-32603, &format!("temporarily unavailable: {m}"))) } + m if forbidden => Err(err(-32012, &format!("forbidden: {m}"))), "events/list" => { // `MCP_FIXTURE_LIST_DELAY_MS`: a slow server, for concurrency tests. let delay = env_u64("MCP_FIXTURE_LIST_DELAY_MS", 0); if delay > 0 { tokio::time::sleep(Duration::from_millis(delay)).await; } + // `MCP_FIXTURE_LIST_PAD_BYTES`: a discovery answer far larger than any client should + // read whole. + let pad = env_u64("MCP_FIXTURE_LIST_PAD_BYTES", 0) as usize; + if pad > 0 { + return Ok(json!({ "events": event_types(), "pad": "x".repeat(pad) })); + } Ok(json!({ "events": event_types() })) } "events/poll" => { @@ -928,7 +1006,8 @@ async fn control(state: &Shared, method: &str, path: &str, body: &[u8]) -> (u16, "cancelled": st.cancelled, "verifications": st.verifications, "deliveries": st.deliveries, "streams": st.streams.len(), "log": st.log.len(), "requests": st.requests, "sessionless_rejections": st.sessionless_rejections, - "jwks_stalled": st.jwks_stalled, + "jwks_stalled": st.jwks_stalled, "nested_answers": st.nested_answers, + "client_answers": st.client_answers.len(), "answers_max_in_flight": st.answers_max_in_flight, }), ) } @@ -1308,6 +1387,27 @@ async fn handle_http(state: Shared, mut stream: TcpStream) { } return respond(&mut stream, 202, "text/plain", b"").await; }; + // An answer to a request this server raised: kept for whoever is waiting on it. Like the + // official Python SDK, a POST that does not accept both JSON and SSE is refused `406` (and so + // never reaches the waiter). + if msg.get("method").is_none() { + let accept = req.headers.get("accept").map(String::as_str).unwrap_or(""); + if !(accept.contains("application/json") && accept.contains("text/event-stream")) { + return respond(&mut stream, 406, "text/plain", b"Not Acceptable").await; + } + { + let mut st = state.lock().unwrap(); + st.client_answers.push(msg.clone()); + st.answers_in_flight += 1; + st.answers_max_in_flight = st.answers_max_in_flight.max(st.answers_in_flight); + } + let delay = env_u64("MCP_FIXTURE_ANSWER_DELAY_MS", 0); + if delay > 0 { + tokio::time::sleep(Duration::from_millis(delay)).await; + } + state.lock().unwrap().answers_in_flight -= 1; + return respond(&mut stream, 202, "text/plain", b"").await; + } // Stateless streamable HTTP (2026-07-28) requires `Mcp-Method` to match the body, as the // official Python SDK enforces — so a client that forgets it fails here, not in production. if !legacy && req.headers.get("mcp-method").map(String::as_str) != Some(method.as_str()) { @@ -1325,7 +1425,47 @@ async fn handle_http(state: Shared, mut stream: TcpStream) { .get("authorization") .and_then(|a| a.strip_prefix("Bearer ")) .map(str::to_owned); + let nested_during = std::env::var("MCP_FIXTURE_NESTED_DURING").unwrap_or_default(); + let raise_nested = |during: &str| { + nested_during == during + && !std::mem::replace(&mut state.lock().unwrap().http_nested_done, true) + }; + // `MCP_FIXTURE_NESTED_DURING=poll` over HTTP: the first `events/poll` is answered as an SSE + // stream that raises the nested request first, and carries the result only after the client + // has answered it (or the wait has timed out). + // `MCP_FIXTURE_POLL_SSE_PAD_BYTES`: `events/poll` answered as an SSE stream whose one event is + // padded far past any sane message size. + let poll_pad = env_u64("MCP_FIXTURE_POLL_SSE_PAD_BYTES", 0) as usize; + if method == "events/poll" && poll_pad > 0 { + let head = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n"; + let _ = stream.write_all(head.as_bytes()).await; + let body = json!({"jsonrpc": "2.0", "id": id, "result": { "events": [], "cursor": "0", "pad": "x".repeat(poll_pad) }}); + let _ = stream + .write_all(format!("data: {body}\n\n").as_bytes()) + .await; + return; + } + if method == "events/poll" && raise_nested("poll") { + let head = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n"; + let _ = stream.write_all(head.as_bytes()).await; + let _ = stream + .write_all(format!("data: {}\n\n", nested_request("poll")).as_bytes()) + .await; + let _ = stream.flush().await; + await_http_answer(&state, "poll").await; + let body = match rpc(&state, &method, params, principal).await { + Ok(result) => json!({"jsonrpc": "2.0", "id": id, "result": like_a_2026_server(result)}), + Err((code, message, data)) => { + json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message, "data": data}}) + } + }; + let _ = stream + .write_all(format!("data: {body}\n\n").as_bytes()) + .await; + return; + } if method == "events/stream" { + let nested = raise_nested("stream"); match open_stream(&state, ¶ms, id.clone()) { Ok((key, mut rx)) => { let head = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n"; @@ -1333,6 +1473,24 @@ async fn handle_http(state: Shared, mut stream: TcpStream) { state.lock().unwrap().streams.remove(&key); return; } + if nested { + let _ = stream + .write_all(format!("data: {}\n\n", nested_request("stream")).as_bytes()) + .await; + let _ = stream.flush().await; + let state = state.clone(); + tokio::spawn(async move { await_http_answer(&state, "stream").await }); + } + // `MCP_FIXTURE_NESTED_FLOOD=`: a hostile server raising `n` requests on the + // stream at once. + for i in 0..env_u64("MCP_FIXTURE_NESTED_FLOOD", 0) { + let mut request = nested_request("flood"); + request["id"] = json!(format!("flood-{i}")); + let _ = stream + .write_all(format!("data: {request}\n\n").as_bytes()) + .await; + } + let _ = stream.flush().await; let terminated_end = loop { let Some(frame) = rx.recv().await else { break false; @@ -1399,8 +1557,52 @@ async fn write_line(out: &Stdout, v: &Value) -> bool { out.write_all(&line).await.is_ok() && out.flush().await.is_ok() } +/// Raise a nested `elicitation/create` toward the client and record its answer (or an error) in +/// `nested_answers`. Answers come back through `waiters`, routed by `run_stdio`. +async fn nested_elicitation(state: &Shared, out: &Stdout, waiters: &Waiters, during: &str) { + // `MCP_FIXTURE_NESTED_DELAY_MS`: hold the request (and so the `events/*` call it rides) a + // while first, so a client can attach before it is raised. + let delay = env_u64("MCP_FIXTURE_NESTED_DELAY_MS", 0); + if delay > 0 { + tokio::time::sleep(Duration::from_millis(delay)).await; + } + let id = format!("nested-{during}"); + let (tx, rx) = tokio::sync::oneshot::channel(); + waiters.lock().unwrap().insert(json!(id).to_string(), tx); + let request = json!({ + "jsonrpc": "2.0", + "id": id, + "method": "elicitation/create", + "params": { + "mode": "form", + "message": format!("Approve during events/{during}?"), + "requestedSchema": { + "type": "object", + "properties": { "ok": { "type": "boolean" } }, + }, + }, + }); + if !write_line(out, &request).await { + return; + } + let answer = match tokio::time::timeout(Duration::from_secs(30), rx).await { + Ok(Ok(msg)) => msg, + _ => json!({ "timeout": true }), + }; + state + .lock() + .unwrap() + .nested_answers + .push(json!({ "during": during, "answer": answer })); +} + +type Waiters = Arc>>>; + async fn run_stdio(state: Shared) { let out: Stdout = Arc::new(tokio::sync::Mutex::new(tokio::io::stdout())); + let waiters: Waiters = Arc::default(); + let nested_during = std::env::var("MCP_FIXTURE_NESTED_DURING").unwrap_or_default(); + let nested_done = Arc::new(std::sync::atomic::AtomicBool::new(false)); // `MCP_FIXTURE_GARBAGE_STDOUT=1`: a server that prints a line that is not UTF-8 before it // speaks MCP (a stray banner from a native library, say). A client must skip it, not die. if env_flag("MCP_FIXTURE_GARBAGE_STDOUT") { @@ -1422,9 +1624,44 @@ async fn run_stdio(state: Shared) { continue; }; if msg.get("method").is_none() { - continue; // a response to something we never asked + // An answer to a request this server raised, if it is one we are waiting for. + if let Some(tx) = waiters.lock().unwrap().remove(&id.to_string()) { + let _ = tx.send(msg.clone()); + } + continue; + } + // `MCP_FIXTURE_NESTED_DURING=poll`: the first `events/poll` is answered only after the + // client has answered a nested elicitation raised while it is in flight. + if method == "events/poll" + && nested_during == "poll" + && !nested_done.swap(true, std::sync::atomic::Ordering::SeqCst) + { + let (state, out, waiters) = (state.clone(), out.clone(), waiters.clone()); + tokio::spawn(async move { + nested_elicitation(&state, &out, &waiters, "poll").await; + let reply = match rpc(&state, "events/poll", params, None).await { + Ok(result) => { + json!({"jsonrpc": "2.0", "id": id, "result": like_a_2026_server(result)}) + } + Err((code, message, data)) => { + json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message, "data": data}}) + } + }; + write_line(&out, &reply).await; + }); + continue; } if method == "events/stream" { + if nested_during == "stream" + && !nested_done.swap(true, std::sync::atomic::Ordering::SeqCst) + { + let (state, out, waiters) = (state.clone(), out.clone(), waiters.clone()); + tokio::spawn(async move { + // Raised while the stream request is open (after its first notifications). + tokio::time::sleep(Duration::from_millis(300)).await; + nested_elicitation(&state, &out, &waiters, "stream").await; + }); + } match open_stream(&state, ¶ms, id.clone()) { Ok((key, mut rx)) => { let out = out.clone(); @@ -1468,6 +1705,7 @@ async fn main() { let stdio = std::env::args().any(|a| a == "--stdio"); let state: Shared = Arc::new(Mutex::new(State { events_down: env_flag("MCP_FIXTURE_EVENTS_DOWN"), + forbid_left: env_u64("MCP_FIXTURE_FORBID_FIRST", 0), ..State::default() })); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/crates/agent/src/main.rs b/crates/agent/src/main.rs index 2168b30e..4495c9de 100644 --- a/crates/agent/src/main.rs +++ b/crates/agent/src/main.rs @@ -1662,7 +1662,20 @@ fn main() -> Result<(), Box> { // The ordinary way out (`run` finishing, every subcommand that returns) and a panic unwinding // out of `run` both pass through this guard's `Drop`: the same cleanup as `exit_process`. let _sweep = tools::mcp_stdio::ExitSweep; - build_runtime(is_service_mode())?.block_on(run()) + let runtime = build_runtime(is_service_mode())?; + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| runtime.block_on(run()))) { + Ok(result) => result, + // A panic out of `run` (the panic hook has already reported it). Dropping the runtime + // here would wait for every blocking task — and `serve`'s stdin reader is one, parked on a + // read that may never return — so the process would hang instead of exiting. Abandon the + // runtime's threads, sweep (the same cleanup as every exit), and go, with Rust's own + // panic status. + Err(_) => { + runtime.shutdown_background(); + tools::mcp_stdio::sweep_before_exit(); + std::process::exit(101); + } + } } /// What every way out of the process does first: retire every stdio MCP server still running diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index ace7f955..e24d3fcf 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -4617,7 +4617,6 @@ pub(crate) async fn serve_session( // the run is over — see `mcp_events_finish` below. let own_injection = crate::tools::mcp_events::injection_batch(&cmd); let mut steered_injections: Vec = Vec::new(); - let messages_before_run = session.messages.len(); let result = 'retry: loop { tokens_before.store(0, Ordering::Relaxed); refused.store(false, Ordering::Relaxed); @@ -4671,7 +4670,19 @@ pub(crate) async fn serve_session( } let life_obs = life.clone(); let attempt_result = { + // A steered MCP Events batch is delivered the moment the model receives it: the + // run's `Steered` event reports its tag after the transcript holding it is + // checkpointed, and it is recorded — durably — right then, never re-derived + // from the transcript afterwards (a compaction may have summarized it away). + let receipts_sink = mcp_events.as_ref().map(|hub| hub.receipts()); let mut sink = move |ev: AgentEvent| { + if let AgentEvent::Steered { tags, .. } = &ev + && let Some(receipts) = &receipts_sink + { + for tag in tags { + receipts.received(*tag); + } + } // Set on `CompactionStart`, cleared on literally anything else — see // `is_compacting`'s own declaration above for why that's exact, not a // conservative approximation. @@ -4972,7 +4983,8 @@ pub(crate) async fn serve_session( if let Some(batch) = crate::tools::mcp_events::injection_batch(&c) { let text = crate::skills::expand_if_skill_invocation(m, &skills); let text = crate::prompts::expand_if_slash(&text, &prompt_templates); - let m = agent_core::SteeringMessage::new(text, parse_images(c.get("images"))); + let m = agent_core::SteeringMessage::new(text, parse_images(c.get("images"))) + .with_tag(batch); let queued = steering.push_steer(m); if !queued { deferred_events.push_back(l.to_string()); @@ -5593,10 +5605,10 @@ pub(crate) async fn serve_session( // A batch is delivered only once the model has it in a transcript that persisted: // the run's own prompt when the run did not fail (an abort still leaves the prompt - // in the transcript the next run reads); a steered one only if it actually made it - // into the transcript (found by its batch id) — an abort clears the steer lane - // (`clear_run_scoped`) before the model sees what is queued there. Everything else - // goes back to pending, to be injected again. + // in the transcript the next run reads); a steered one exactly when the run's + // `Steered` event reported it (already recorded then — see `steered_received`) — an + // abort clears the steer lane (`clear_run_scoped`) before the model sees what is + // queued there. Everything else goes back to pending, to be injected again. if let Some(hub) = &mcp_events && (own_injection.is_some() || !steered_injections.is_empty()) { @@ -5610,24 +5622,10 @@ pub(crate) async fn serve_session( returned.push(b) } } - // Compaction can shrink the transcript mid-run; then look at all of it. - let since = if session.messages.len() >= messages_before_run { - messages_before_run - } else { - 0 - }; - for b in steered_injections.drain(..) { - if run_ok - && crate::tools::mcp_events::transcript_has_injection( - &session.messages[since..], - b, - ) - { - delivered.push(b); - } else { - returned.push(b); - } - } + // A steered batch the model received has already left the pending queue (see + // `receipts_sink`); returning it is a no-op. One it never received — an abort + // cleared the steer lane — goes back to pending. + returned.append(&mut steered_injections); hub.finish_run(&delivered, &returned).await; } if running.swap(false, Ordering::Relaxed) @@ -6110,6 +6108,15 @@ pub(crate) async fn serve_session( )), } } + // Test-only, in debug builds and only with `BEYOND_AI_AGENT_TEST_PANICS` set: panic right + // here, inside the session loop, so the way a panic leaves `serve` can be tested. + #[cfg(debug_assertions)] + "__test_panic" if std::env::var_os("BEYOND_AI_AGENT_TEST_PANICS").is_some() => { + #[allow(clippy::panic)] + { + panic!("test panic requested inside serve"); + } + } "get_state" => { let mut data = session_stats(&session, ¤t_model); if let Value::Object(m) = &mut data { diff --git a/crates/agent/src/serve_ws.rs b/crates/agent/src/serve_ws.rs index 4689fa66..e5acf641 100644 --- a/crates/agent/src/serve_ws.rs +++ b/crates/agent/src/serve_ws.rs @@ -409,6 +409,9 @@ struct Supervisor { /// transition already flows through. metrics: Option>, body: SessionBody, + /// Told the id of every session that ended in a panic (see `TryPin` in `start`), so the daemon + /// can bring back one it depends on — the MCP Events session. + panicked: tokio::sync::mpsc::UnboundedSender, } /// What service mode needs at the *supervisor* level, as opposed to per session: the keyring every @@ -962,6 +965,8 @@ impl Supervisor { // For the failure paths below: whoever attached to the `Starting` slot while the lock was // being taken is told why nothing started. let starting_conn = out_conn.clone(); + let panic_conn = out_conn.clone(); + let panicked = self.panicked.clone(); let metrics_for_lock = self.metrics.clone(); tokio::spawn(async move { let mut lock = None; @@ -1013,7 +1018,31 @@ impl Supervisor { } let _ = started_tx.send(Ok(())); if exit.go_live() { - body.await; + // A panic in one session is contained to that session: its future unwinds — every + // in-flight MCP call and run is dropped, its events hub cancelled and its state + // written by their `Drop`s — its clients get an `error` frame and their connections + // end with it, and the daemon goes on serving every other session. Sound here + // because a session shares nothing with others that a half-finished update could + // leave broken: shared state is behind `lock_ignoring_poison`'s plain + // replace-the-value locks, and MCP server processes are the process's (swept on its + // exit), not the session's. + use futures::FutureExt as _; + if let Err(payload) = std::panic::AssertUnwindSafe(body).catch_unwind().await { + let message = payload + .downcast_ref::<&str>() + .map(|s| (*s).to_owned()) + .or_else(|| payload.downcast_ref::().cloned()) + .unwrap_or_else(|| "a panic".to_owned()); + eprintln!( + "serve: session {session_id} panicked ({message}); it has ended, the daemon carries on" + ); + lock_ignoring_poison(&panic_conn).broadcast(OutFrame::Value(json!({ + "type": "error", + "session_id": session_id, + "error": format!("the session ended after an internal error: {message}"), + }))); + let _ = panicked.send(session_id.clone()); + } } else { drop(body); } @@ -1753,6 +1782,7 @@ pub async fn serve_ws( // mode. let events_session = (!cfg.service_mode && !cfg.mcp_catalog.event_subscriptions().is_empty()) .then(|| cfg.mcp_events_session.clone()); + let (panicked_tx, mut panicked_rx) = tokio::sync::mpsc::unbounded_channel::(); let supervisor = Arc::new(Supervisor { table: Arc::default(), // Service mode lists per tenant, from the shards, never from one process-wide directory. @@ -1762,9 +1792,11 @@ pub async fn serve_ws( service, metrics: cfg.metrics.clone(), body: serve_session_body(cfg), + panicked: panicked_tx, }); let mut shutdown = crate::serve::ShutdownSignal::new()?; + let events_session_id = events_session.clone(); if let Some(id) = events_session { // Pinned and at once unpinned: the session starts exactly as if a client had connected and // left. It stays because its subscriptions keep it alive, not because anything holds it. @@ -1780,6 +1812,100 @@ pub async fn serve_ws( } } + // A session that ends in a panic is started again when it holds MCP Events subscriptions that + // must keep running with no client: the events session (configured subscriptions), or one with + // runtime subscriptions (`mcp_events_subscribe`) — with backoff per session (1 s doubling to a + // minute, reset once it has stayed up ten minutes), so a session that panics on every start + // cannot spin the daemon. Any other session that panics simply ends: its clients were told, and + // it starts again when one comes back. + { + let supervisor = supervisor.clone(); + let events_id = events_session_id.clone(); + let session_dir = supervisor.session_dir.clone(); + tokio::spawn(async move { + let mut backoff: std::collections::HashMap = + std::collections::HashMap::new(); + while let Some(id) = panicked_rx.recv().await { + let holds_runtime = match &session_dir { + Some(dir) => { + let dir = dir.clone(); + let probe = id.clone(); + tokio::task::spawn_blocking(move || { + crate::tools::mcp_events::sessions_with_runtime_subscriptions( + std::path::Path::new(&dir), + usize::MAX, + ) + .contains(&probe) + }) + .await + .unwrap_or(false) + } + None => false, + }; + if Some(&id) != events_id.as_ref() && !holds_runtime { + continue; + } + let entry = backoff + .entry(id.clone()) + .or_insert((0, std::time::Instant::now())); + if entry.1.elapsed() > std::time::Duration::from_secs(600) { + entry.0 = 0; + } + let delay = std::time::Duration::from_secs(1u64 << entry.0.min(6)) + .min(std::time::Duration::from_secs(60)); + entry.0 += 1; + entry.1 = std::time::Instant::now(); + let supervisor = supervisor.clone(); + tokio::spawn(async move { + tokio::time::sleep(delay).await; + match supervisor.pin(Some(id.clone()), None).await { + Ok(p) => { + supervisor.unpin(&p.id, p.incarnation); + eprintln!( + "serve: restarted session `{id}` after a panic, for its MCP Events subscriptions" + ); + } + Err(e) => eprintln!( + "serve: could not restart session `{id}` after a panic: {}", + HttpError::from(e) + ), + } + }); + } + }); + } + + // Sessions holding runtime MCP Events subscriptions (`mcp_events_subscribe`) are started too, + // the same way, so those subscriptions are restored after a restart without waiting for their + // client to come back. Never in service mode. + if let Some(dir) = supervisor.session_dir.clone() { + let events_id = events_session_id.clone(); + let ids = tokio::task::spawn_blocking(move || { + // Bounded: a pile of old sessions cannot all be woken at every boot. + crate::tools::mcp_events::sessions_with_runtime_subscriptions( + std::path::Path::new(&dir), + crate::tools::mcp_events::max_restored_sessions(), + ) + }) + .await + .unwrap_or_default(); + for id in ids + .into_iter() + .filter(|id| Some(id.as_str()) != events_id.as_deref()) + { + match supervisor.pin(Some(id.clone()), None).await { + Ok(p) => { + supervisor.unpin(&p.id, p.incarnation); + eprintln!("serve: restored session `{id}` for its MCP Events subscriptions"); + } + Err(e) => eprintln!( + "serve: could not restore session `{id}` for its MCP Events subscriptions: {}", + HttpError::from(e) + ), + } + } + } + // The idle reaper (on unless `--session-idle-timeout 0` turned it off). A background ticker that // stops dead and detached-idle-not-mid-run sessions — the same `Stopping` transition (drop the // retained input → the session persists and exits) shutdown uses. Its handle is aborted on @@ -3048,6 +3174,7 @@ mod tests { table: Arc::default(), session_dir: None, service: None, + panicked: tokio::sync::mpsc::unbounded_channel().0, body: Box::new( move |_id, _service, mut input_rx, _out, _running, _keep_alive| { let probe = probe.clone(); diff --git a/crates/agent/src/tools/mcp.rs b/crates/agent/src/tools/mcp.rs index 408bd3e9..8f9bdf81 100644 --- a/crates/agent/src/tools/mcp.rs +++ b/crates/agent/src/tools/mcp.rs @@ -2415,8 +2415,6 @@ async fn connect_http( crate::tools::mcp_oauth::OAuthHttp::new( crate::tools::mcp_view_http::ViewCappedHttp::new(crate::tools::mcp_wire::HttpClient { client, - // With a login, every POST is answered where its 401 status is still visible. - oauth: auth.is_some(), }), auth, ), diff --git a/crates/agent/src/tools/mcp_events/mod.rs b/crates/agent/src/tools/mcp_events/mod.rs index 16c6ff7e..22249cab 100644 --- a/crates/agent/src/tools/mcp_events/mod.rs +++ b/crates/agent/src/tools/mcp_events/mod.rs @@ -42,7 +42,7 @@ pub mod webhook; mod wire; use std::collections::{HashMap, HashSet, VecDeque}; -use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; @@ -166,6 +166,16 @@ fn healthy_for() -> Duration { /// Longest wait between attempts to (re)establish a configured subscription. const CONFIGURED_RETRY_CAP: Duration = Duration::from_secs(60); +/// How many times in a row a server may answer "forbidden" (`-32012`) — each attempt with +/// credentials resolved afresh — before the subscription is treated as refused for good. +const MAX_FORBIDDEN_RETRIES: u32 = 5; + +/// How long a subscription the server refused for good waits before it is tried again — rarely, in +/// case the server's configuration changed (`BEYOND_AI_AGENT_MCP_EVENTS_REFUSED_RETRY_MS`). +fn refused_retry() -> Duration { + env_ms("BEYOND_AI_AGENT_MCP_EVENTS_REFUSED_RETRY_MS", 3_600_000) +} + fn now_unix() -> i64 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -362,6 +372,7 @@ impl SubState { cursor: self.cursor(), recent, webhook: None, + runtime: None, }, ); } @@ -402,6 +413,16 @@ impl Active { } } +/// A subscription this session means to have that is not up right now (see +/// `Hub::unestablished`). +struct Unestablished { + spec: SubSpec, + last_error: Option, + /// The server refused for good (see [`SubError`]): retried only every + /// [`refused_retry`], and not counted as live. + permanent: bool, +} + /// What [`Hub::deliver`] did with one occurrence. #[derive(Debug, PartialEq, Eq)] enum Delivery { @@ -463,10 +484,14 @@ struct Hub { /// For direct HTTP events requests; built on first use, never for a stdio-only session. http: std::sync::OnceLock, keep_alive: Option>, - /// Configured subscriptions not currently up — starting, or retrying after a failure or a - /// termination. They count as live for the keep-alive: a daemon whose servers are all briefly - /// down must not have its events session reaped. - starting: AtomicUsize, + /// Subscriptions this session means to have that are not up right now — configured ones + /// starting or retrying, runtime ones being restored after a restart — and why. Those still + /// expected to come up count as live for the keep-alive (a daemon whose servers are all briefly + /// down must not have its events session reaped); a permanent refusal does not. + unestablished: Mutex>, + /// Keys of runtime subscriptions being restored after a restart (until each is up again or + /// explicitly unsubscribed). + restoring: Mutex>, store: StateStore, /// `mcp_events_*` commands in flight — spawned so none ever blocks the session's command loop. command_tasks: Mutex>, @@ -497,9 +522,14 @@ impl McpEventsHub { ) -> (mpsc::Receiver, Self) { let (tx, rx) = mpsc::channel::(crate::serve::IN_CHANNEL_BOUND); let weak = tx.downgrade(); + let store = StateStore::open(cfg.state_path, max_pending(), max_pending_bytes()); + let client_store = store.clone(); tokio::spawn(async move { let mut input_rx = input_rx; while let Some(line) = input_rx.recv().await { + // Only a client's own commands pass here (injections go around), so this is what + // "a client was here" means for runtime subscriptions' time to live. + client_store.touch_client(now_unix_ms()); if tx.send(line).await.is_err() { break; } @@ -518,8 +548,9 @@ impl McpEventsHub { shutdown: CancellationToken::new(), http: std::sync::OnceLock::new(), keep_alive: cfg.keep_alive, - starting: AtomicUsize::new(0), - store: StateStore::open(cfg.state_path, max_pending(), max_pending_bytes()), + unestablished: Mutex::new(HashMap::new()), + restoring: Mutex::new(HashSet::new()), + store, command_tasks: Mutex::new(tokio::task::JoinSet::new()), owns_configured: cfg.owns_configured, }); @@ -529,12 +560,12 @@ impl McpEventsHub { hub.refresh_keep_alive(); } }); - // Persisted webhook callbacks of *configured* subscriptions — the ones this session will - // subscribe again — are held from the moment the state is read, so a delivery the server - // retried across a restart gets `503` (retry), not `410` (stop), until the subscription has - // re-registered. A runtime subscription is not resubscribed after a restart, so its token - // is not held: `410` is the truth for it. - if hub.callback_url.is_some() { + // Once the state is read: hold the persisted webhook callbacks of every subscription this + // session will subscribe again — configured ones (re-created from settings) and runtime + // ones (restored from the state) — so a delivery the server retried across a restart gets + // `503` (retry), not `410` (stop), until the subscription has re-registered; then restore + // the runtime ones. (`start_configured` runs right after `attach`, before this does.) + { let weak_hub = Arc::downgrade(&hub); tokio::spawn(async move { let Some(store) = weak_hub.upgrade().map(|h| h.store.clone()) else { @@ -544,16 +575,35 @@ impl McpEventsHub { let Some(hub) = weak_hub.upgrade() else { return; }; - let configured = hub - .configured - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone(); - for (key, token) in store.webhook_tokens() { - if configured.contains(&key) { - webhook::reserve(&token); + // Runtime subscriptions whose session has heard from no client for longer than + // their time to live are forgotten — spec, position and webhook callback (token + // and secret) alike — not restored. + if !state::runtime_still_wanted(store.last_client_ms(), now_unix_ms()) { + for (key, _) in store.runtime_specs() { + store.forget_sub(&key); } } + let runtime: Vec<(String, SubSpec)> = store + .runtime_specs() + .into_iter() + .filter(|(key, _)| !hub.is_configured(key)) + .filter_map(|(key, cmd)| parse_spec(&cmd).ok().map(|spec| (key, spec))) + .filter(|(key, spec)| spec.key() == *key) + .collect(); + if hub.callback_url.is_some() { + for (key, token) in store.webhook_tokens() { + if hub.is_configured(&key) || runtime.iter().any(|(k, _)| *k == key) { + webhook::reserve(&token); + } + } + } + for (key, spec) in runtime { + hub.restoring + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(key); + Hub::keep_subscribed(&hub, spec, Duration::ZERO, true); + } }); } let coalescer = tokio::spawn(coalesce( @@ -592,7 +642,7 @@ impl McpEventsHub { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .insert(spec.key()); - Hub::keep_subscribed(&self.hub, spec, Duration::ZERO); + Hub::keep_subscribed(&self.hub, spec, Duration::ZERO, false); } } } @@ -635,6 +685,14 @@ impl McpEventsHub { } } + /// The handle a run's event sink records steered batches through, the moment the model receives + /// them (see [`Receipts::received`]). + pub fn receipts(&self) -> Receipts { + Receipts { + hub: self.hub.clone(), + } + } + /// The session moved to another transcript file (`new_session`, `switch_session`, …): the /// events state moves with it now, merged with whatever that transcript already had. pub async fn relocate(&self, session_file: Option<&std::path::Path>) { @@ -697,6 +755,81 @@ impl McpEventsHub { } } +/// The ids of (at most `max`) the sessions in `dir` whose events state holds runtime subscriptions +/// still worth restoring (see [`state::runtime_still_wanted`]) — for a daemon to start at boot, so +/// those subscriptions come back without waiting for a client to reattach. Reads only the small +/// snapshot files (`_.mcp-events.json`); blocking. +pub fn sessions_with_runtime_subscriptions(dir: &std::path::Path, max: usize) -> Vec { + let Ok(entries) = std::fs::read_dir(dir) else { + return Vec::new(); + }; + let now = now_unix_ms(); + let mut ranked: Vec<(i64, String)> = Vec::new(); + for entry in entries.flatten() { + let name = entry.file_name(); + let Some(stem) = name + .to_str() + .and_then(|n| n.strip_suffix(".mcp-events.json")) + else { + continue; + }; + let Some((_, id)) = stem.split_once('_') else { + continue; + }; + if let Some(last) = std::fs::read(entry.path()) + .ok() + .and_then(|b| state::snapshot_wants_restore(&b, now)) + { + ranked.push((last, id.to_owned())); + } + } + // The most recently used first: past the cap, the ones a client touched longest ago wait. + ranked.sort_by(|a, b| b.0.cmp(&a.0).then_with(|| a.1.cmp(&b.1))); + let mut ids: Vec = ranked.into_iter().map(|(_, id)| id).collect(); + if ids.len() > max { + eprintln!( + "serve: {} sessions hold runtime MCP Events subscriptions; restoring {max} of them at boot \ + (the rest come back when their clients do)", + ids.len() + ); + ids.truncate(max); + } + ids +} + +/// The most sessions a daemon starts at boot to restore runtime MCP Events subscriptions +/// (`BEYOND_AI_AGENT_MCP_EVENTS_MAX_RESTORED_SESSIONS`, default 32). +pub fn max_restored_sessions() -> usize { + env_u64("BEYOND_AI_AGENT_MCP_EVENTS_MAX_RESTORED_SESSIONS", 32) as usize +} + +/// Records steered injection batches as delivered from inside a run's (synchronous) event sink. +#[derive(Clone)] +pub struct Receipts { + hub: Arc, +} + +impl Receipts { + /// The model received `batch`: the run's `Steered` event reported it, after the transcript + /// holding it was checkpointed. Its events leave the pending queue now, and the record of that + /// is written at once — whatever a later compaction does to the transcript, the batch is never + /// injected again. + pub fn received(&self, batch: u64) { + if self.hub.store.delivered(batch) { + let store = self.hub.store.clone(); + let hub = Arc::downgrade(&self.hub); + tokio::spawn(async move { + if let Err(e) = store.commit().await { + tracing::warn!(error = %e, "could not record a received MCP events batch yet"); + } + if let Some(hub) = hub.upgrade() { + hub.refresh_keep_alive(); + } + }); + } + } +} + /// Where a session's events state lives: beside its transcript. pub fn state_path_for(session_file: &std::path::Path) -> std::path::PathBuf { session_file.with_extension("mcp-events.json") @@ -773,6 +906,20 @@ fn parse_key(cmd: &Value) -> Result { }) } +/// A spec as the `mcp_events_subscribe` command that would create it — what a runtime +/// subscription is remembered as (see `PersistedSub::runtime`), and parsed back by [`parse_spec`]. +fn spec_command(spec: &SubSpec) -> Value { + json!({ + "server": spec.server, + "name": spec.sub.name, + "arguments": spec.arguments(), + "delivery": spec.sub.delivery.map(mode_str), + "action": action_str(spec.sub.action), + "instructions": spec.sub.instructions, + "max_age_ms": spec.sub.max_age_ms, + }) +} + fn parse_spec(cmd: &Value) -> Result { let mut spec = parse_key(cmd)?; spec.sub.delivery = match cmd.get("delivery") { @@ -801,7 +948,18 @@ impl Hub { let result = match ctype { "mcp_events_list" => Ok(hub.list(cmd.get("server").and_then(Value::as_str)).await), "mcp_events_subscribe" => match parse_spec(cmd) { - Ok(spec) => Hub::subscribe(hub, spec).await, + Ok(spec) => { + let result = Hub::subscribe(hub, spec.clone()) + .await + .map_err(|e| e.message); + // A runtime subscription is remembered, so the session subscribes it again + // after a restart. (A configured one is re-created from settings.) + if result.is_ok() && !hub.is_configured(&spec.key()) { + hub.store + .set_runtime(&spec.key(), Some(spec_command(&spec))); + } + result + } Err(e) => Err(e), }, "mcp_events_unsubscribe" => match parse_key(cmd) { @@ -853,15 +1011,39 @@ impl Hub { Err(e) => available.push(json!({ "server": server, "supported": false, - "error": e, + "error": e.message, })), } } + let unestablished: Vec = { + let u = self + .unestablished + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut v: Vec<(&String, &Unestablished)> = u.iter().collect(); + v.sort_by(|a, b| a.0.cmp(b.0)); + v.into_iter() + .map(|(_, u)| { + json!({ + "server": u.spec.server, + "name": u.spec.sub.name, + "arguments": u.spec.arguments(), + "state": match (&u.last_error, u.permanent) { + (_, true) => "refused", + (Some(_), false) => "retrying", + (None, false) => "starting", + }, + "last_error": u.last_error, + }) + }) + .collect() + }; json!({ "spec_commit": SPEC_COMMIT, "webhook": self.callback_url.is_some(), "owns_configured": self.owns_configured, "subscriptions": subscriptions, + "unestablished": unestablished, "available": available, "pending": self.store.pending_len(), }) @@ -904,7 +1086,7 @@ impl Hub { } /// `events/list`, all pages, cached until the server says `list_changed`. - async fn discover(&self, server: &str) -> Result>, String> { + async fn discover(&self, server: &str) -> Result>, SubError> { let conn = self.conn(server).await?; let generation = conn.generation(); if let Some((g, events)) = self @@ -927,10 +1109,10 @@ impl Hub { .call("events/list", params, RPC_TIMEOUT) .await .map_err(|e| match e.code { - Some(-32601) => format!( + Some(-32601) => SubError::permanent(format!( "server `{server}` does not support the MCP Events extension (events/list: method not found)" - ), - _ => format!("events/list on `{server}` failed: {e}"), + )), + _ => SubError::from_rpc(&e).context(&format!("events/list on `{server}` failed")), })?; if let Some(list) = page.get("events").and_then(Value::as_array) { events.extend(list.iter().cloned()); @@ -954,7 +1136,7 @@ impl Hub { /// delivery/action/instructions replaces the old subscription **only once the new one is /// confirmed**: a replacement that fails leaves the old one running, untouched. The two share /// one cursor and dedup window, so the brief overlap cannot double-deliver. - async fn subscribe(hub: &Arc, spec: SubSpec) -> Result { + async fn subscribe(hub: &Arc, spec: SubSpec) -> Result { let key = spec.key(); let lock = hub.key_lock(&key); let _held = lock.lock().await; @@ -969,6 +1151,13 @@ impl Hub { } }; + // A server that is not (or no longer) configured will not appear by retrying. + if !hub.catalog.snapshot().iter().any(|s| s.name == spec.server) { + return Err(SubError::permanent(format!( + "unknown MCP server `{}`", + spec.server + ))); + } let events = hub.discover(&spec.server).await?; let descriptor = events .iter() @@ -978,12 +1167,12 @@ impl Hub { .iter() .filter_map(|e| e.get("name").and_then(Value::as_str)) .collect(); - format!( + SubError::permanent(format!( "server `{}` offers no event `{}` (it offers: {})", spec.server, spec.sub.name, names.join(", ") - ) + )) })?; let offered: Vec = descriptor .get("delivery") @@ -994,7 +1183,8 @@ impl Hub { .collect() }) .unwrap_or_default(); - let mode = choose_mode(spec.sub.delivery, &offered, hub.callback_url.is_some())?; + let mode = choose_mode(spec.sub.delivery, &offered, hub.callback_url.is_some()) + .map_err(SubError::permanent)?; hub.store.loaded().await; let state = match &reuse { @@ -1003,7 +1193,7 @@ impl Hub { }; let prior_state = state.with(|s| std::mem::replace(&mut s.state, "starting")); let cancel = hub.shutdown.child_token(); - let (ready_tx, ready_rx) = oneshot::channel::>(); + let (ready_tx, ready_rx) = oneshot::channel::>(); let spec_arc = Arc::new(spec.clone()); let task = { let hub = hub.clone(); @@ -1026,11 +1216,13 @@ impl Hub { let outcome = match ready { Ok(Ok(Ok(()))) => Ok(()), Ok(Ok(Err(e))) => Err(e), - Ok(Err(_)) => Err("the subscription task ended before it was ready".to_owned()), - Err(_) => Err(format!( + Ok(Err(_)) => Err(SubError::from( + "the subscription task ended before it was ready", + )), + Err(_) => Err(SubError::from(format!( "no confirmation within {}s", READY_TIMEOUT.as_secs() - )), + ))), }; if let Err(e) = outcome { cancel.cancel(); @@ -1038,12 +1230,12 @@ impl Hub { if reuse.is_some() { state.with(|s| s.state = prior_state); } - return Err(format!( - "subscribing to `{}` on `{}` ({}) failed: {e}", + return Err(e.context(&format!( + "subscribing to `{}` on `{}` ({}) failed", spec.sub.name, spec.server, mode_str(mode) - )); + ))); } let replaced = hub.lock_subs().remove(&key); if let Some(old) = replaced { @@ -1088,6 +1280,10 @@ impl Hub { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .remove(key); + self.restoring + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(key); } let Some(active) = self.lock_subs().remove(key) else { return false; @@ -1102,7 +1298,12 @@ impl Hub { /// trigger needs: a live subscription, configured ones still starting, or undelivered events. fn refresh_keep_alive(&self) { if let Some(k) = &self.keep_alive { - let live = self.starting.load(Ordering::Acquire) > 0 + let live = self + .unestablished + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .values() + .any(|u| !u.permanent) || self.store.pending_len() > 0 || self.lock_subs().values().any(Active::live); k.store(live, Ordering::Release); @@ -1169,6 +1370,29 @@ impl Hub { /// The server says events may have been skipped (`truncated`, a `gap` envelope): adopt its /// fresh cursor, tell attached clients, and — unless `notify` — tell the model too, so it can /// re-check authoritative state. + /// An event too large to accept was skipped (see the `$oversized` stand-in): note its id so a + /// replay of it is not reported twice, keep its cursor if known, and record a gap that says so. + fn oversized(&self, spec: &SubSpec, state: &SubState, stand_in: &Value) { + let event_id = stand_in.get("eventId").and_then(Value::as_str); + if let Some(id) = event_id + && let Ok(mut d) = state.dedup.lock() + && !d.first_sighting(id) + { + return; + } + tracing::warn!( + server = %spec.server, + event = %spec.sub.name, + event_id, + "skipped an MCP event over the message-size cap" + ); + let mut carrier = json!({ "reason": "oversized", "eventId": event_id }); + if let Some(cursor) = stand_in.get("cursor") { + carrier["cursor"] = cursor.clone(); + } + self.gap(spec, state, &carrier); + } + fn gap(&self, spec: &SubSpec, state: &SubState, carrier: &Value) { if carrier .as_object() @@ -1176,7 +1400,13 @@ impl Hub { { state.set_cursor_from(carrier); } - self.status_event(spec, "gap", json!({ "cursor": state.cursor() })); + let reason = carrier.get("reason").and_then(Value::as_str); + let event_id = carrier.get("eventId").filter(|v| !v.is_null()).cloned(); + self.status_event( + spec, + "gap", + json!({ "cursor": state.cursor(), "reason": reason, "event_id": event_id }), + ); if spec.sub.action != McpEventAction::Notify { let queued = self.store.push_pending(PendingEvent::new( spec.sub.action, @@ -1184,7 +1414,7 @@ impl Hub { spec.sub.name.clone(), spec.arguments(), spec.sub.instructions.clone(), - json!({ "gap": true, "cursor": state.cursor() }), + json!({ "gap": true, "cursor": state.cursor(), "reason": reason, "eventId": event_id }), )); if !queued { tracing::warn!("pending queue full; a gap notice was not queued for the model"); @@ -1206,32 +1436,60 @@ impl Hub { (self.emit)(frame); } - /// Keep a configured subscription up: subscribe after `delay`, and on failure retry with - /// capped backoff — forever, until the session ends. It counts as live (for the keep-alive) - /// the whole time it is not up. - fn keep_subscribed(hub: &Arc, spec: SubSpec, delay: Duration) { - hub.starting.fetch_add(1, Ordering::AcqRel); + fn is_configured(&self, key: &str) -> bool { + self.configured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .contains(key) + } + + fn is_restoring(&self, key: &str) -> bool { + self.restoring + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .contains(key) + } + + /// Keep a subscription up: subscribe after `delay`, and on failure try again — a configured one + /// forever (until the session ends or it is explicitly unsubscribed), a `restoring` runtime one + /// until it is back. A transient failure retries with backoff capped at + /// [`CONFIGURED_RETRY_CAP`] and keeps the session alive meanwhile; a **permanent** refusal (see + /// [`SubError`]) is reported as `refused`, retried only every [`refused_retry`], and does not + /// keep the session alive. + fn keep_subscribed(hub: &Arc, spec: SubSpec, delay: Duration, restoring: bool) { + let key = spec.key(); + hub.unestablished + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert( + key.clone(), + Unestablished { + spec: spec.clone(), + last_error: None, + permanent: false, + }, + ); hub.refresh_keep_alive(); let weak = Arc::downgrade(hub); let shutdown = hub.shutdown.clone(); - let resumed = !delay.is_zero(); + let resumed = !delay.is_zero() && !restoring; tokio::spawn(async move { let mut delay = delay; let mut failures = 0u32; + let mut forbidden_streak = 0u32; loop { tokio::select! { () = tokio::time::sleep(delay) => {} () = shutdown.cancelled() => break, } let Some(hub) = weak.upgrade() else { return }; - // An explicit unsubscribe while this was waiting ends the effort: it is no longer - // configured to be kept up. - if !hub - .configured - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .contains(&spec.key()) - { + // An explicit unsubscribe while this was waiting ends the effort. + let wanted = if restoring { + hub.is_restoring(&key) + } else { + hub.is_configured(&key) + }; + if !wanted { break; } if resumed { @@ -1247,26 +1505,75 @@ impl Hub { } break; } - Err(e) => { - failures += 1; - delay = backoff(failures, CONFIGURED_RETRY_CAP); + Err(mut e) => { + // Forbidden: every attempt dials afresh (credentials re-resolved); only a + // refusal that keeps coming back is taken as final. + if e.forbidden { + forbidden_streak += 1; + e.permanent |= forbidden_streak >= MAX_FORBIDDEN_RETRIES; + } else { + forbidden_streak = 0; + } + let kind = if e.permanent { + delay = refused_retry(); + "refused" + } else if e.forbidden { + // A short backoff: a refreshed credential should work at once. + delay = backoff(forbidden_streak - 1, Duration::from_secs(5)); + "error" + } else { + failures += 1; + delay = backoff(failures, CONFIGURED_RETRY_CAP); + "error" + }; + if let Some(u) = hub + .unestablished + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get_mut(&key) + { + u.last_error = Some(e.message.clone()); + u.permanent = e.permanent; + } + hub.refresh_keep_alive(); eprintln!( - "warning: session {}: mcp events: `{}` on `{}`: {e} (retrying in {}s)", + "warning: session {}: mcp events: `{}` on `{}`: {e} ({}; retrying in {}s)", hub.session_id, spec.sub.name, spec.server, + if e.permanent { + "refused" + } else { + "unavailable" + }, delay.as_secs() ); hub.status_event( &spec, - "error", - json!({ "error": e, "retry_in_ms": delay.as_millis() as u64 }), + kind, + json!({ "error": e.message, "retry_in_ms": delay.as_millis() as u64 }), ); + // A runtime subscription being restored that the server now refuses for + // good (its server gone from settings, say) is forgotten, its webhook + // callback with it: not retried, not restored again, keeping nothing. + if restoring && e.permanent { + hub.store.forget_sub(&key); + break; + } } } } if let Some(hub) = weak.upgrade() { - hub.starting.fetch_sub(1, Ordering::AcqRel); + hub.unestablished + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&key); + if restoring { + hub.restoring + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&key); + } hub.refresh_keep_alive(); } }); @@ -1302,10 +1609,14 @@ impl Hub { .unwrap_or_else(std::sync::PoisonError::into_inner) .contains(&spec.key()); if configured { - Hub::keep_subscribed(self, spec.clone(), Duration::from_secs(1)); + Hub::keep_subscribed(self, spec.clone(), Duration::from_secs(1), false); return; } self.refresh_keep_alive(); + // A runtime subscription the server ended is over: forgotten — spec, position, webhook + // callback (token and secret) — so it is not subscribed again after a restart (unless + // re-discovery brings it back, below, as a new subscription). + self.store.forget_sub(&spec.key()); if !error.wants_rediscovery() { return; } @@ -1337,8 +1648,12 @@ impl Hub { d.remove(&spec.server); } match Hub::subscribe(&hub, spec.clone()).await { - Ok(_) => hub.resubscribed(&spec), - Err(e) => hub.status_event(&spec, "error", json!({ "error": e })), + Ok(_) => { + hub.store + .set_runtime(&spec.key(), Some(spec_command(&spec))); + hub.resubscribed(&spec); + } + Err(e) => hub.status_event(&spec, "error", json!({ "error": e.message })), } }); } @@ -1394,12 +1709,74 @@ async fn stop_task(task: tokio::task::JoinHandle<()>) { } } +/// Why a subscribe failed, and whether trying again could help. A **permanent** refusal (the server +/// does not offer the event, does not speak the extension, rejects the request as invalid, or no +/// delivery mode fits) is retried only rarely, and does not keep a session alive. +#[derive(Debug, Clone)] +struct SubError { + pub(super) message: String, + pub(super) permanent: bool, + /// Access refused (`-32012`): retried, with credentials re-resolved, until it has failed + /// [`MAX_FORBIDDEN_RETRIES`] times in a row. + pub(super) forbidden: bool, +} + +impl SubError { + fn permanent(message: impl Into) -> Self { + Self { + message: message.into(), + permanent: true, + forbidden: false, + } + } + + fn from_rpc(e: &RpcError) -> Self { + Self { + message: e.to_string(), + permanent: e.is_permanent(), + forbidden: e.is_forbidden(), + } + } + + fn context(self, prefix: &str) -> Self { + Self { + message: format!("{prefix}: {}", self.message), + ..self + } + } +} + +impl From for SubError { + fn from(message: String) -> Self { + Self { + message, + permanent: false, + forbidden: false, + } + } +} + +impl From<&str> for SubError { + fn from(message: &str) -> Self { + message.to_owned().into() + } +} + +impl std::fmt::Display for SubError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.message) + } +} + +/// How a subscription task reports its first success or failure to `subscribe`. +type ReadyTx = oneshot::Sender>; + /// Exponential backoff, 1 s doubling to a cap. fn backoff(attempt: u32, cap: Duration) -> Duration { Duration::from_secs(1u64 << attempt.min(6)).min(cap) } -fn ready_ok(ready: &mut Option>>) { +fn ready_ok(ready: &mut Option) { if let Some(tx) = ready.take() { let _ = tx.send(Ok(())); } @@ -1414,7 +1791,7 @@ async fn run_poll( spec: Arc, state: Arc, cancel: CancellationToken, - ready: oneshot::Sender>, + ready: ReadyTx, ) { let mut ready = Some(ready); let mut failures = 0u32; @@ -1508,7 +1885,7 @@ async fn run_poll( } Err(e) => { if let Some(tx) = ready.take() { - let _ = tx.send(Err(e.to_string())); + let _ = tx.send(Err(SubError::from_rpc(&e))); return; } if e.is_terminal() { @@ -1552,7 +1929,7 @@ async fn run_push( spec: Arc, state: Arc, cancel: CancellationToken, - ready: oneshot::Sender>, + ready: ReadyTx, ) { let mut ready = Some(ready); let mut failures = 0u32; @@ -1584,7 +1961,7 @@ async fn run_push( } StreamEnd::Reconnect(why) => { if let Some(tx) = ready.take() { - let _ = tx.send(Err(why)); + let _ = tx.send(Err(why.into())); return; } failures += 1; @@ -1607,11 +1984,11 @@ fn stream_error( hub: &Arc, spec: &SubSpec, state: &SubState, - ready: &mut Option>>, + ready: &mut Option, e: RpcError, ) -> StreamEnd { if let Some(tx) = ready.take() { - let _ = tx.send(Err(e.to_string())); + let _ = tx.send(Err(SubError::from_rpc(&e))); return StreamEnd::Terminated; } if e.is_terminal() { @@ -1627,7 +2004,7 @@ fn on_stream_msg( hub: &Arc, spec: &SubSpec, state: &SubState, - ready: &mut Option>>, + ready: &mut Option, failures: &mut u32, msg: StreamMsg, ) -> Option { @@ -1645,6 +2022,13 @@ fn on_stream_msg( state.set_cursor_from(&msg.params); } } + // An event over the message-size cap, skipped by the transport: only its stand-in arrived + // (`mcp_stdio::oversized_stand_in`). It is not reconnected into — the stream goes on — its + // cursor is kept when its head carried one (else the next heartbeat's applies), and the + // model is told an event was dropped. + "notifications/events/event" if msg.params.get("$oversized") == Some(&json!(true)) => { + hub.oversized(spec, state, &msg.params); + } "notifications/events/event" => { if hub.deliver(spec, McpEventDelivery::Push, state, msg.params) == Delivery::Full { // Nothing of it was recorded: reconnecting from the cursor brings it back. @@ -1660,7 +2044,9 @@ fn on_stream_msg( "notifications/events/terminated" => { let error = msg.params.get("error").cloned().unwrap_or(Value::Null); if let Some(tx) = ready.take() { - let _ = tx.send(Err(format!("terminated: {error}"))); + let _ = tx.send(Err( + SubError::from_rpc(&RpcError::from_json(&error)).context("terminated") + )); } else { hub.ended(spec, state, RpcError::from_json(&error)); } @@ -1681,7 +2067,7 @@ async fn push_once( state: &SubState, cancel: &CancellationToken, conn: &Conn, - ready: &mut Option>>, + ready: &mut Option, failures: &mut u32, ) -> StreamEnd { let mut stream = match conn.open_stream(spec.params(state.cursor())).await { @@ -1874,19 +2260,6 @@ fn injection_marker(batch: u64) -> String { format!("[MCP events · batch {batch}]") } -/// Whether the injection carrying `batch` is in `messages` as a user turn — how a steered batch is -/// known to have reached the model. Found by its batch id, not its content. (A compaction during -/// the run that summarized the steered turn away hides it: the batch is then re-injected.) -pub fn transcript_has_injection(messages: &[agent_core::Message], batch: u64) -> bool { - let marker = injection_marker(batch); - messages.iter().rev().any(|m| { - m.role == agent_core::Role::User - && m.content.iter().any(|b| { - matches!(b, agent_core::ContentBlock::Text { text, .. } if text.starts_with(&marker)) - }) - }) -} - fn render_injection(batch: u64, events: &[PendingEvent]) -> String { let occurrences = events .iter() @@ -1911,6 +2284,20 @@ fn render_injection(batch: u64, events: &[PendingEvent]) -> String { } for e in events { if e.event.get("gap") == Some(&json!(true)) { + if e.event.get("reason") == Some(&json!("oversized")) { + out.push_str(&format!( + "\n[gap] An event for `{}` on `{}`{} was larger than the message-size limit and \ + was dropped unread. If it matters, re-check the authoritative state with tools.\n", + e.name, + e.server, + e.event + .get("eventId") + .and_then(Value::as_str) + .map(|id| format!(" (event id `{}`)", id.replace('`', ""))) + .unwrap_or_default() + )); + continue; + } out.push_str(&format!( "\n[gap] The server reported that events for `{}` on `{}` may have been missed \ (its history did not reach back far enough). If it matters, re-check the \ @@ -2095,20 +2482,39 @@ mod tests { let _ = task.await; } - /// A steered batch is recognised in the transcript by its id. Two injections with identical - /// content (two gap notices) are told apart: only the one that reached the model counts. + /// Every injection's text names its batch on its first line — for anyone reading the + /// transcript; delivery itself is tracked by tag (`AgentEvent::Steered`), not by text. #[test] - fn a_steered_batch_is_found_by_its_id_not_its_content() { + fn an_injection_names_its_batch() { let gap = || pending(McpEventAction::Steer, json!({"gap": true, "cursor": "9"})); - let seven = render_injection(7, &[gap()]); - let eight = render_injection(8, &[gap()]); - let in_transcript = vec![agent_core::Message::user(seven)]; - assert!(transcript_has_injection(&in_transcript, 7)); - assert!( - !transcript_has_injection(&in_transcript, 8), - "an identical batch that never reached the model is not delivered" + assert!(render_injection(7, &[gap()]).starts_with(&injection_marker(7))); + assert_ne!(render_injection(7, &[gap()]), render_injection(8, &[gap()])); + } + + /// Past the boot cap, the sessions restored are the ones a client used most recently — not + /// the first ones alphabetically. + #[test] + fn the_boot_cap_keeps_the_most_recently_used_sessions() { + let dir = tempfile::tempdir().unwrap(); + let now = now_unix_ms(); + for (id, ago) in [ + ("a-oldest", 3_000), + ("b-newest", 1_000), + ("c-middle", 2_000), + ] { + std::fs::write( + dir.path().join(format!("100_{id}.mcp-events.json")), + json!({ + "subscriptions": { "k": { "cursor": null, "runtime": { "server": "s", "name": "n" } } }, + "last_client_ms": now - ago, + }) + .to_string(), + ) + .unwrap(); + } + assert_eq!( + sessions_with_runtime_subscriptions(dir.path(), 2), + ["b-newest", "c-middle"] ); - assert!(!transcript_has_injection(&in_transcript, 70)); - let _ = eight; } } diff --git a/crates/agent/src/tools/mcp_events/state.rs b/crates/agent/src/tools/mcp_events/state.rs index 3cdc0797..9c8bed53 100644 --- a/crates/agent/src/tools/mcp_events/state.rs +++ b/crates/agent/src/tools/mcp_events/state.rs @@ -61,6 +61,10 @@ pub(super) struct PersistedSub { pub(super) recent: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub(super) webhook: Option, + /// A runtime subscription's whole spec (`mcp_events_subscribe`'s command), so the session + /// subscribes it again after a restart. `None` for a configured one (settings re-create it). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) runtime: Option, } /// An event a server has been told was received (its cursor advanced, its webhook acked) that the @@ -123,6 +127,45 @@ struct Snapshot { /// in `webhook`). #[serde(default)] server_keys: BTreeMap>, + /// When a client last sent this session a command (unix ms; coarse — updated at most once a + /// minute). Runtime subscriptions are not restored once it is older than their time to live. + #[serde(default, skip_serializing_if = "Option::is_none")] + last_client_ms: Option, +} + +/// How long a runtime subscription outlives its session's last client command before it is no +/// longer restored after a restart (`BEYOND_AI_AGENT_MCP_EVENTS_RUNTIME_TTL_MS`, default 7 days). +pub(super) fn runtime_ttl_ms() -> i64 { + std::env::var("BEYOND_AI_AGENT_MCP_EVENTS_RUNTIME_TTL_MS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(7 * 24 * 3_600_000) +} + +/// Whether a snapshot's runtime subscriptions are still worth restoring: its session heard from a +/// client within [`runtime_ttl_ms`]. A subscription's own creation counts as hearing from one +/// ([`StateStore::set_runtime`] records it), so a state with runtime subscriptions and no record +/// at all is not trusted: not wanted. +pub(super) fn runtime_still_wanted(last_client_ms: Option, now_ms: i64) -> bool { + last_client_ms.is_some_and(|t| now_ms.saturating_sub(t) <= runtime_ttl_ms()) +} + +fn now_ms() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0) +} + +/// What a daemon reads at boot to decide whether to start a session: does its events state hold +/// runtime subscriptions still worth restoring? (Snapshot only; blocking.) +/// Returns when the session last heard from a client, when it does — for ranking. +pub(super) fn snapshot_wants_restore(bytes: &[u8], now_ms: i64) -> Option { + let snap = serde_json::from_slice::(bytes).ok()?; + (runtime_still_wanted(snap.last_client_ms, now_ms) + && snap.subscriptions.values().any(|s| s.runtime.is_some())) + .then_some(snap.last_client_ms) + .flatten() } #[derive(serde::Serialize, serde::Deserialize)] @@ -156,6 +199,7 @@ struct Data { snapshot_dirty: bool, /// Bytes in the log file on disk. log_len: u64, + last_client_ms: Option, } impl Data { @@ -245,6 +289,9 @@ impl StateStore { if entry.webhook.is_none() { entry.webhook = d.subs.get(key).and_then(|s| s.webhook.clone()); } + if entry.runtime.is_none() { + entry.runtime = d.subs.get(key).and_then(|s| s.runtime.clone()); + } if d.subs.get(key) == Some(&entry) { return; } @@ -277,6 +324,61 @@ impl StateStore { .collect() } + /// Record (or with `None`, drop) a runtime subscription's spec, to subscribe again after a + /// restart. + pub(super) fn set_runtime(&self, key: &str, spec: Option) { + { + let mut d = lock(&self.inner.data); + if spec.is_none() && !d.subs.contains_key(key) { + return; + } + let created = spec.is_some(); + let entry = d.subs.entry(key.to_owned()).or_default(); + if entry.runtime == spec { + return; + } + entry.runtime = spec; + // Creating a runtime subscription is a client's own command: its time to live starts + // now, whatever the command channel recorded before there was runtime state. + if created { + d.last_client_ms = Some(now_ms()); + } + d.snapshot_dirty = true; + } + self.dirty(); + } + + /// A client sent this session a command just now. Recorded at most once a minute, and only + /// while there is runtime state the record matters for. + pub(super) fn touch_client(&self, now_ms: i64) { + { + let mut d = lock(&self.inner.data); + if !d.subs.values().any(|s| s.runtime.is_some()) + || d.last_client_ms + .is_some_and(|t| now_ms.saturating_sub(t) < 60_000) + { + return; + } + d.last_client_ms = Some(now_ms); + d.snapshot_dirty = true; + } + self.dirty(); + } + + /// When a client last sent this session a command (see [`Self::touch_client`]). + pub(super) fn last_client_ms(&self) -> Option { + lock(&self.inner.data).last_client_ms + } + + /// Every persisted runtime subscription spec, as `(subscription key, spec)`. + pub(super) fn runtime_specs(&self) -> Vec<(String, Value)> { + lock(&self.inner.data) + .subs + .iter() + .filter_map(|(k, s)| s.runtime.clone().map(|r| (k.clone(), r))) + .collect() + } + /// An explicit unsubscribe forgets the position; a session ending keeps it, to resume from. pub(super) fn forget_sub(&self, key: &str) { { @@ -559,6 +661,7 @@ async fn writer( d.pending_bytes = l.pending.iter().map(|p| p.size).sum(); d.subs = l.snapshot.subscriptions; d.server_keys = l.snapshot.server_keys; + d.last_client_ms = l.snapshot.last_client_ms; d.pending = l.pending; d.log_len = l.log_len; } @@ -667,6 +770,7 @@ async fn flush(data: &Arc>) -> Result<(), String> { spec_commit: SPEC_COMMIT.to_owned(), subscriptions: d.subs.clone(), server_keys: d.server_keys.clone(), + last_client_ms: d.last_client_ms, } }); (path, log_write, snapshot) @@ -777,6 +881,7 @@ async fn relocate(data: &Arc>, target: Option) { spec_commit: SPEC_COMMIT.to_owned(), subscriptions: d.subs.clone(), server_keys: d.server_keys.clone(), + last_client_ms: d.last_client_ms, }), Some(log), ) @@ -907,6 +1012,7 @@ mod tests { cursor: Some("7".into()), recent: vec!["a".into()], webhook: None, + runtime: None, }, ); let writer = store.take_writer().unwrap(); @@ -1248,4 +1354,60 @@ mod tests { "the event pushed before the failure is on disk" ); } + + /// Runtime subscriptions expire: once their session has heard from no client for longer than + /// the time to live, they are not restored, and a daemon does not start that session at boot. + #[test] + fn runtime_subscriptions_expire_without_a_client() { + let day = 24 * 3_600_000; + let now = 100 * day; + assert!(runtime_still_wanted(Some(now - day), now)); + assert!( + !runtime_still_wanted(Some(now - 8 * day), now), + "a week is the default" + ); + assert!( + !runtime_still_wanted(None, now), + "no record of a client: not trusted" + ); + let snap = |last: i64| { + serde_json::json!({ + "subscriptions": { "k": { "cursor": null, "runtime": { "server": "s", "name": "n" } } }, + "last_client_ms": last, + }) + .to_string() + }; + assert_eq!( + snapshot_wants_restore(snap(now - day).as_bytes(), now), + Some(now - day) + ); + assert_eq!( + snapshot_wants_restore(snap(now - 8 * day).as_bytes(), now), + None + ); + let no_runtime = + serde_json::json!({ "subscriptions": { "k": { "cursor": null } } }).to_string(); + assert_eq!(snapshot_wants_restore(no_runtime.as_bytes(), now), None); + } + + /// A client's commands are recorded (coarsely) only while there is runtime state to expire. + #[tokio::test] + async fn client_activity_is_recorded_only_while_it_matters() { + let s = StateStore::open(None, 10, u64::MAX); + s.loaded().await; + s.touch_client(1_000_000); + assert_eq!(s.last_client_ms(), None, "nothing to expire yet"); + s.set_runtime("k", Some(serde_json::json!({ "server": "s" }))); + assert!( + s.last_client_ms().is_some_and(|t| t > 1_000_000), + "creating a runtime subscription is client activity" + ); + lock(&s.inner.data).last_client_ms = Some(1_000_000 - 120_000); + s.touch_client(1_000_000); + assert_eq!(s.last_client_ms(), Some(1_000_000)); + s.touch_client(1_030_000); + assert_eq!(s.last_client_ms(), Some(1_000_000), "at most once a minute"); + s.touch_client(1_070_000); + assert_eq!(s.last_client_ms(), Some(1_070_000)); + } } diff --git a/crates/agent/src/tools/mcp_events/webhook.rs b/crates/agent/src/tools/mcp_events/webhook.rs index 1e9d0d15..d6bbfd2f 100644 --- a/crates/agent/src/tools/mcp_events/webhook.rs +++ b/crates/agent/src/tools/mcp_events/webhook.rs @@ -32,8 +32,8 @@ use crate::settings::McpEventDelivery; use super::wire::{KeyFetch, RpcError, parse_rfc3339_ms}; use super::{ - Delivery, Hub, RPC_TIMEOUT, SubSpec, SubState, WEBHOOK_PATH_PREFIX, backoff, now_unix, - now_unix_ms, ready_ok, webhook_ttl, + Delivery, Hub, RPC_TIMEOUT, ReadyTx, SubError, SubSpec, SubState, WEBHOOK_PATH_PREFIX, backoff, + now_unix, now_unix_ms, ready_ok, webhook_ttl, }; /// Standard Webhooks' replay window, in seconds, either side of now. @@ -486,12 +486,14 @@ pub(super) async fn run_webhook( spec: Arc, state: Arc, cancel: CancellationToken, - ready: oneshot::Sender>, + ready: ReadyTx, ) { let mut ready = Some(ready); let Some(base) = hub.callback_url.clone() else { if let Some(tx) = ready.take() { - let _ = tx.send(Err("webhook delivery is not configured".into())); + let _ = tx.send(Err(SubError::permanent( + "webhook delivery is not configured", + ))); } return; }; @@ -519,7 +521,7 @@ pub(super) async fn run_webhook( Ok(v) => v, Err(e) => { if let Some(tx) = ready.take() { - let _ = tx.send(Err(e)); + let _ = tx.send(Err(e.into())); } return; } @@ -640,7 +642,7 @@ pub(super) async fn run_webhook( }; if let Err(e) = keys { if let Some(tx) = ready.take() { - let _ = tx.send(Err(e)); + let _ = tx.send(Err(e.into())); } return; } @@ -724,7 +726,7 @@ pub(super) async fn run_webhook( } Err(e) => { if let Some(tx) = ready.take() { - let _ = tx.send(Err(e.to_string())); + let _ = tx.send(Err(SubError::from_rpc(&e))); return; } if e.is_terminal() && e.code != Some(-32015) { diff --git a/crates/agent/src/tools/mcp_events/wire.rs b/crates/agent/src/tools/mcp_events/wire.rs index 2d05a7b0..c2fd5888 100644 --- a/crates/agent/src/tools/mcp_events/wire.rs +++ b/crates/agent/src/tools/mcp_events/wire.rs @@ -18,12 +18,14 @@ use tokio::sync::mpsc; use super::env_ms; -/// Most bytes one SSE event (one JSON-RPC message on a push stream) may be. The draft keeps -/// delivery bodies at or under 256 KiB; anything far past that is a broken or hostile server. -pub(super) const MAX_SSE_EVENT_BYTES: usize = 1024 * 1024; -/// Most bytes a unary `events/*` response body may be. -pub(super) const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024; -/// Most bytes a server's webhook JWKS document may be. +/// Most bytes one message on the direct-HTTP events wire — a unary response body, or one SSE event +/// (a JSON-RPC message on a push stream or a unary SSE answer) — may be: the same per-message cap +/// every MCP transport applies ([`crate::tools::mcp_stdio::max_message_bytes`], +/// `BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES`). +fn max_message_bytes() -> usize { + crate::tools::mcp_stdio::max_message_bytes() +} +/// Most bytes a server's webhook JWKS document may be (a key document, not an MCP message). pub(super) const MAX_JWKS_BYTES: usize = 64 * 1024; /// How many notifications one push stream may have queued before its consumer catches up. Past it @@ -167,6 +169,20 @@ impl RpcError { ) } + /// A refusal retrying will not change: the request is wrong (`-32602`), the server does not + /// speak the extension (`-32601`), the event does not exist (`-32011`), or the mode is + /// unsupported (`-32014`) — unless it is the draft's "re-discover" signal. `-32012` (forbidden) + /// is not: credentials expire and are refreshed, so it is retried (see [`Self::is_forbidden`]). + pub(super) fn is_permanent(&self) -> bool { + matches!(self.code, Some(-32601 | -32602 | -32011 | -32014)) && !self.wants_rediscovery() + } + + /// Access refused (`-32012`): retried with fresh credentials on a short backoff, and treated as + /// a permanent refusal only after repeated failures. + pub(super) fn is_forbidden(&self) -> bool { + self.code == Some(-32012) + } + /// The draft's "re-discover and resubscribe" signals: the event type was removed /// (`-32011`, `data.kind: "event"`) or its schema changed in place (`-32014`, /// `data.reason: "schema_changed"`) — not an authorization failure. @@ -294,6 +310,8 @@ impl Conn { json!({ "jsonrpc": "2.0", "id": format!("beyond-events-{id}"), "method": method, "params": params }) } + /// One direct POST. `method` is empty for an answer (a JSON-RPC response has no method, so it + /// carries no `Mcp-Method`). fn http_request( http: &reqwest::Client, url: &str, @@ -306,8 +324,10 @@ impl Conn { .post(url) .header("Content-Type", "application/json") .header("Accept", "application/json, text/event-stream") - .header("MCP-Protocol-Version", protocol_version) - .header("Mcp-Method", method); + .header("MCP-Protocol-Version", protocol_version); + if !method.is_empty() { + req = req.header("Mcp-Method", method); + } for (k, v) in headers { req = req.header(k, v); } @@ -457,7 +477,7 @@ impl Conn { let sse = is_sse(&resp); let status = resp.status(); if !sse { - let bytes = read_capped(resp, MAX_RESPONSE_BYTES) + let bytes = read_capped(resp, max_message_bytes()) .await .map_err(|e| RpcError::local(format!("{method}: {e}")))?; let msg: Value = serde_json::from_slice(&bytes).map_err(|_| { @@ -470,6 +490,17 @@ impl Conn { if msg.get("id") == Some(&want) { return rpc_outcome(msg); } + if is_server_request(&msg) { + refuse_server_request( + http, + url, + headers, + auth.as_deref(), + protocol_version, + &msg, + ) + .await; + } } Err(RpcError::local(format!( "{method}: the stream ended without a response" @@ -555,7 +586,7 @@ impl Conn { .await?; if !is_sse(&resp) { // A JSON answer to a stream request is an immediate error (or a result). - let bytes = read_capped(resp, MAX_RESPONSE_BYTES) + let bytes = read_capped(resp, max_message_bytes()) .await .map_err(|e| RpcError::local(format!("events/stream: {e}")))?; let msg: Value = serde_json::from_slice(&bytes) @@ -565,9 +596,31 @@ impl Conn { } // Backpressure, not loss: the reader awaits channel space, so TCP flow control // slows the server rather than anything being dropped. + let refuse_with = ( + http.clone(), + url.clone(), + headers.clone(), + auth.clone(), + protocol_version.clone(), + ); let reader = tokio::spawn(async move { let mut events = SseReader::new(resp); while let Some(msg) = events.next().await { + // Answered inline — one at a time, each bounded — so a server that floods + // requests onto the stream cannot make the reader spawn without limit. + if is_server_request(&msg) { + let (http, url, headers, auth, version) = &refuse_with; + refuse_server_request( + http, + url, + headers, + auth.as_deref(), + version, + &msg, + ) + .await; + continue; + } if msg.get("id") == Some(&want) { let fin = match rpc_outcome(msg) { Ok(v) => StreamMsg::final_ok(v), @@ -618,6 +671,43 @@ impl Conn { } } +/// A server→client *request* (it has a `method` and an `id`) arriving on an `events/*` response. +fn is_server_request(msg: &Value) -> bool { + msg.get("method").is_some() && msg.get("id").is_some_and(|id| !id.is_null()) +} + +/// Refuse a server→client request that arrived on a direct-HTTP `events/*` response. Over direct +/// HTTP nothing routes it to a session (rmcp never sees the exchange), so rather than leave the +/// server waiting forever on an answer that will never come, it is answered at once with an error. +/// (Over rmcp the same request is attributed to the owning session, like one raised during a +/// `tools/call`.) +async fn refuse_server_request( + http: &reqwest::Client, + url: &str, + headers: &[(http::HeaderName, http::HeaderValue)], + auth: Option<&crate::tools::mcp_oauth::ServerAuth>, + protocol_version: &str, + msg: &Value, +) { + let method = msg.get("method").and_then(Value::as_str).unwrap_or("?"); + let body = json!({ + "jsonrpc": "2.0", + "id": msg["id"], + "error": { + "code": -32601, + "message": format!("`{method}` is not served during an events request over direct HTTP"), + }, + }); + // Sent as every events POST is (`Conn::send`): the same `Accept` — a streamable-HTTP server + // (the official Python SDK among them) answers a POST without it `406`, and then still waits + // for the answer — and, with an OAuth login, the server's current token, refreshed and resent + // once on a 401, so a refresh since the stream opened cannot leave the server waiting either. + let send = Conn::send(http, url, headers, auth, protocol_version, "", &body); + if let Ok(Err(e)) = tokio::time::timeout(Duration::from_secs(5), send).await { + tracing::debug!(error = %e, method, "could not refuse a server request"); + } +} + fn is_sse(resp: &reqwest::Response) -> bool { resp.headers() .get(reqwest::header::CONTENT_TYPE) @@ -635,7 +725,7 @@ fn rpc_outcome(msg: Value) -> Result { /// Minimal `text/event-stream` reader: yields each event's `data:` payload parsed as JSON. /// -/// Bounded and linear: an event larger than [`MAX_SSE_EVENT_BYTES`] ends the stream (with +/// Bounded and linear: an event larger than the per-message cap ends the stream (with /// [`Self::error`] saying why) instead of growing the buffer, and the separator scan resumes where /// it stopped rather than rescanning the whole buffer on every chunk. pub(super) struct SseReader { @@ -647,16 +737,23 @@ pub(super) struct SseReader { start: usize, /// How far into `buf` no separator can start — the next scan begins here. scanned: usize, + /// The largest event read (see [`max_message_bytes`]). + max: usize, pub(super) error: Option, } impl SseReader { pub(super) fn new(resp: reqwest::Response) -> Self { + Self::with_max(resp, max_message_bytes()) + } + + pub(super) fn with_max(resp: reqwest::Response, max: usize) -> Self { Self { resp, buf: Vec::new(), start: 0, scanned: 0, + max, error: None, } } @@ -664,7 +761,20 @@ impl SseReader { pub(super) async fn next(&mut self) -> Option { loop { if let Some((end, sep_end)) = find_event_end(&self.buf, self.scanned.max(self.start)) { - let parsed = parse_sse_event(&self.buf[self.start..end]); + let event = &self.buf[self.start..end]; + // One read can bring a whole over-cap event at once: it is held to the cap all the + // same — stood in for if it is an events notification, else the stream ends. + let parsed = if event.len() > self.max { + match crate::tools::mcp_stdio::oversized_stand_in(&sse_data_head(event, 4096)) { + Some(stand_in) => Some(stand_in), + None => { + self.error = Some(format!("an SSE event exceeded {} bytes", self.max)); + return None; + } + } + } else { + parse_sse_event(event) + }; self.start = sep_end; self.scanned = sep_end; match parsed { @@ -680,8 +790,18 @@ impl SseReader { } // A separator is at most four bytes, so one could still start in the last three. self.scanned = self.scanned.max(self.buf.len().saturating_sub(3)); - if self.buf.len() > MAX_SSE_EVENT_BYTES { - self.error = Some(format!("an SSE event exceeded {MAX_SSE_EVENT_BYTES} bytes")); + if self.buf.len() > self.max { + // An over-cap MCP Events notification is skipped, not read: its bounded head names + // it (and, before the payload, its cursor), and a small stand-in takes its place — + // see `mcp_stdio::oversized_stand_in`. Anything else ends the stream with an error. + let head = sse_data_head(&self.buf, 4096); + if let Some(stand_in) = crate::tools::mcp_stdio::oversized_stand_in(&head) { + if !self.skip_rest_of_event().await { + return None; + } + return Some(stand_in); + } + self.error = Some(format!("an SSE event exceeded {} bytes", self.max)); return None; } match self.resp.chunk().await { @@ -690,6 +810,44 @@ impl SseReader { } } } + + /// Discard the rest of the event in progress (the buffer holds only its start), reading on to + /// its end without keeping it. `false` if the stream ended first. + async fn skip_rest_of_event(&mut self) -> bool { + // Only the last three bytes can begin a separator that the next chunk completes. + let mut tail: Vec = self.buf[self.buf.len().saturating_sub(3)..].to_vec(); + self.buf.clear(); + self.start = 0; + self.scanned = 0; + loop { + let Ok(Some(chunk)) = self.resp.chunk().await else { + return false; + }; + tail.extend_from_slice(&chunk); + if let Some((_, sep_end)) = find_event_end(&tail, 0) { + self.buf.extend_from_slice(&tail[sep_end..]); + return true; + } + let keep = tail.len().saturating_sub(3); + tail.drain(..keep); + } + } +} + +/// The start of an SSE event's first `data:` payload, at most `max` bytes of it. +fn sse_data_head(event: &[u8], max: usize) -> Vec { + let at = event + .windows(5) + .position(|w| w == b"data:") + .map_or(event.len(), |i| i + 5); + let rest = &event[at..]; + let rest = rest.strip_prefix(b" ").unwrap_or(rest); + let end = rest + .iter() + .position(|&b| b == b'\n' || b == b'\r') + .unwrap_or(rest.len()) + .min(max); + rest[..end].to_vec() } /// One SSE event's `data:` lines, joined and parsed as JSON. `None` for an event with no data or @@ -970,9 +1128,9 @@ mod tests { #[tokio::test] async fn an_oversized_sse_event_ends_the_stream_instead_of_buffering_it() { - let body = format!("data: {}", "x".repeat(MAX_SSE_EVENT_BYTES + 10)); + let body = format!("data: {}", "x".repeat(4096 + 10)); let resp = reqwest::Response::from(http::Response::new(body)); - let mut reader = SseReader::new(resp); + let mut reader = SseReader::with_max(resp, 4096); assert!(reader.next().await.is_none()); assert!(reader.error.unwrap().contains("exceeded")); // A normal event still parses. @@ -1007,6 +1165,29 @@ mod tests { ); } + /// An over-cap events notification in an SSE stream is skipped to its end without being held, + /// and stands in as a small `$oversized` message; the stream goes on to the next event. + #[tokio::test] + async fn an_oversized_sse_event_is_skipped_and_the_stream_goes_on() { + let big = format!( + r#"data: {{"jsonrpc":"2.0","method":"notifications/events/event","params":{{"cursor":"7","data":"{}"}}}}"#, + "x".repeat(50_000) + ); + let body = format!( + "{big}\n\ndata: {}\n\n", + r#"{"jsonrpc":"2.0","method":"notifications/events/event","params":{"cursor":"8"}}"# + ); + let resp = reqwest::Response::from(http::Response::new(body)); + let mut reader = SseReader::with_max(resp, 4096); + let first = reader.next().await.unwrap(); + assert_eq!(first["params"]["$oversized"], true); + assert_eq!(first["params"]["cursor"], "7"); + assert!(first["params"].get("data").is_none()); + let second = reader.next().await.unwrap(); + assert_eq!(second["params"]["cursor"], "8"); + assert!(reader.error.is_none()); + } + #[test] fn a_full_stream_is_marked_overflowed_and_forwards_nothing_after() { let router = NotificationRouter::default(); @@ -1127,4 +1308,38 @@ mod tests { assert_eq!(resp.status(), 200); assert_eq!(*seen.lock().unwrap(), ["bearer fresh"]); } + + #[tokio::test] + async fn a_refused_server_request_is_answered_with_the_current_token_and_refreshed_on_401() { + // The refusal is sent the way every events POST is: the current shared token, not the + // one the stream's headers were built with, and a 401 refreshes and resends once — a + // refusal the server rejects leaves it waiting on an answer that never comes. + agent_core::ensure_provider(); + let request = json!({ "jsonrpc": "2.0", "id": 3, "method": "elicitation/create" }); + let (url, seen) = bearer_server().await; + let auth = crate::tools::mcp_oauth::ServerAuth::fake("fresh", || Ok("unused".into())); + refuse_server_request( + &reqwest::Client::new(), + &url, + &bearer("stale"), + Some(&auth), + "2026-07-28", + &request, + ) + .await; + assert_eq!(*seen.lock().unwrap(), ["bearer fresh"]); + + let (url, seen) = bearer_server().await; + let auth = crate::tools::mcp_oauth::ServerAuth::fake("stale", || Ok("fresh".into())); + refuse_server_request( + &reqwest::Client::new(), + &url, + &bearer("stale"), + Some(&auth), + "2026-07-28", + &request, + ) + .await; + assert_eq!(*seen.lock().unwrap(), ["bearer stale", "bearer fresh"]); + } } diff --git a/crates/agent/src/tools/mcp_oauth.rs b/crates/agent/src/tools/mcp_oauth.rs index 67dcb50f..03aab109 100644 --- a/crates/agent/src/tools/mcp_oauth.rs +++ b/crates/agent/src/tools/mcp_oauth.rs @@ -908,7 +908,6 @@ mod tests { OAuthHttp::new( crate::tools::mcp_view_http::ViewCappedHttp::new(crate::tools::mcp_wire::HttpClient { client: reqwest::Client::new(), - oauth: true, }), Some(auth), ) diff --git a/crates/agent/src/tools/mcp_stdio.rs b/crates/agent/src/tools/mcp_stdio.rs index aaf8ffa4..167a75f5 100644 --- a/crates/agent/src/tools/mcp_stdio.rs +++ b/crates/agent/src/tools/mcp_stdio.rs @@ -215,7 +215,9 @@ pub(crate) fn rescue(line: &[u8]) -> Option> { /// server response cannot take the host's memory with it. pub const DEFAULT_MAX_MESSAGE_BYTES: usize = 64 * 1024 * 1024; -fn max_message_bytes() -> usize { +/// The per-message cap in force — the one cap every MCP transport applies: stdio lines here, every +/// streamable-HTTP body and SSE event (`mcp_wire::HttpClient`), and the MCP Events direct-HTTP wire. +pub(crate) fn max_message_bytes() -> usize { std::env::var("BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES") .ok() .and_then(|v| v.parse().ok()) @@ -357,6 +359,90 @@ pub(crate) fn scan_head(head: &[u8]) -> Head { } } +/// A small stand-in for an over-cap **MCP Events notification** (`notifications/events/*`), built +/// from the bounded head of the message — the oversized one itself is never held. The stand-in is +/// the same notification with `params` cut down to the members that appear, whole, before the +/// head runs out (the event's `_meta` routing, `cursor`, `eventId`, `name`, … — whatever precedes the +/// payload) plus `"$oversized": true`, so the events client can skip that event, keep its cursor and +/// tell the model it was dropped, rather than reconnecting into the same giant event forever. +/// `None` for anything else (it is dropped as before). +pub(crate) fn oversized_stand_in(head: &[u8]) -> Option { + let ws = |mut i: usize| { + while head.get(i).is_some_and(u8::is_ascii_whitespace) { + i += 1; + } + i + }; + let mut i = ws(0); + if head.get(i) != Some(&b'{') { + return None; + } + i += 1; + let mut method: Option = None; + loop { + i = ws(i); + if head.get(i) != Some(&b'"') { + return None; + } + let key_end = string_end(head, i)?; + let key = &head[i + 1..key_end - 1]; + i = ws(key_end); + if head.get(i) != Some(&b':') { + return None; + } + i = ws(i + 1); + if key == b"params" { + let method = method.filter(|m| m.starts_with("notifications/events/"))?; + let mut params = serde_json::Map::new(); + if head.get(i) == Some(&b'{') { + let mut j = i + 1; + loop { + j = ws(j); + if head.get(j) != Some(&b'"') { + break; + } + let Some(k_end) = string_end(head, j) else { + break; + }; + let Ok(k) = serde_json::from_slice::(&head[j..k_end]) else { + break; + }; + j = ws(k_end); + if head.get(j) != Some(&b':') { + break; + } + j = ws(j + 1); + let Some(v_end) = value_end(head, j) else { + break; + }; + let value = &head[j..v_end]; + // Scalars, and the small `_meta` (routing); a whole nested payload is skipped. + let keep = !matches!(value.first(), Some(b'{' | b'[')) || k == "_meta"; + if keep && let Ok(v) = serde_json::from_slice::(value) { + params.insert(k, v); + } + j = ws(v_end); + if head.get(j) != Some(&b',') { + break; + } + j += 1; + } + } + params.insert("$oversized".into(), Value::Bool(true)); + return Some(json!({ "jsonrpc": "2.0", "method": method, "params": params })); + } + let end = value_end(head, i)?; + if key == b"method" { + method = serde_json::from_slice::(&head[i..end]).ok(); + } + i = ws(end); + if head.get(i) != Some(&b',') { + return None; + } + i += 1; + } +} + /// The index just past the JSON string opening at `buf[at]`, if it closes within `buf`. fn string_end(buf: &[u8], at: usize) -> Option { let mut j = at + 1; @@ -417,6 +503,8 @@ pub(crate) enum Inbound { reply_to: Option, method: Option, cap: usize, + /// For an over-cap MCP Events notification, its stand-in ([`oversized_stand_in`]). + stand_in: Option, }, Eof, } @@ -437,19 +525,22 @@ pub(crate) async fn read_capped( let mut head: Option = None; // Set once the line is over its cap. The rest is discarded unread. let mut discarding: Option = None; - let refused = |head: Head, cap: usize| { + // An over-cap events notification's stand-in, built from its head before that is let go. + let mut stand_in: Option = None; + let refused = |head: Head, cap: usize, stand_in: Option| { caps.take(head.reply_to()); Inbound::Refused { reply_to: head.reply_to().map(str::to_owned), method: head.method, cap, + stand_in, } }; loop { let avail = reader.fill_buf().await?; if avail.is_empty() { return Ok(match discarding { - Some(head) => refused(head, cap), + Some(head) => refused(head, cap, stand_in.take()), None if buf.is_empty() => Inbound::Eof, None => Inbound::Line(buf), }); @@ -470,7 +561,11 @@ pub(crate) async fn read_capped( } // Over the window means scanned above, so `head` is always `Some` here. if buf.len() > cap.max(ID_WINDOW) || (done && buf.len() > cap) { - discarding = Some(head.take().unwrap_or_default()); + let scanned = head.take().unwrap_or_default(); + if scanned.reply_to().is_none() { + stand_in = oversized_stand_in(&buf[..buf.len().min(ID_WINDOW)]); + } + discarding = Some(scanned); buf = Vec::new(); } } @@ -478,7 +573,7 @@ pub(crate) async fn read_capped( reader.consume(n); if done { return Ok(match discarding { - Some(head) => refused(head, cap), + Some(head) => refused(head, cap, stand_in.take()), None => { caps.take(head.as_ref().and_then(Head::reply_to)); Inbound::Line(buf) @@ -597,7 +692,18 @@ pub(crate) fn stdio_transport( reply_to, method, cap, + stand_in, }) => { + // An over-cap MCP Events notification goes on as its small stand-in, so the + // subscription can skip the event and report it dropped. + if let Some(stand_in) = stand_in { + let mut line = stand_in.to_string().into_bytes(); + line.push(b'\n'); + if pump_side.write_all(&line).await.is_err() { + break; + } + continue; + } // Answer in the server's place only for a response that provably answers a host // request. A server→client request or notification is dropped (its id is the // server's, and an error under it could fail an unrelated host request), as is a @@ -710,7 +816,8 @@ mod tests { Inbound::Refused { reply_to: Some("7".into()), method: None, - cap: view_cap + cap: view_cap, + stand_in: None, } ); assert!( @@ -746,7 +853,8 @@ mod tests { Inbound::Refused { reply_to: None, method: None, - cap: 100_000 + cap: 100_000, + stand_in: None, } ); assert!(peak <= 100_000 + 8192, "held {peak}"); @@ -763,7 +871,8 @@ mod tests { Inbound::Refused { reply_to: Some("\"abc\"".into()), method: None, - cap: 100_000 + cap: 100_000, + stand_in: None, } ); // A plain response under the cap passes untouched, and an untracked large one is not @@ -800,7 +909,8 @@ mod tests { Inbound::Refused { reply_to: None, method: Some("notifications/progress".into()), - cap: 100_000 + cap: 100_000, + stand_in: None, } ); // Nor when the nested id comes before the top-level members, in a response with no id. @@ -835,7 +945,8 @@ mod tests { Inbound::Refused { reply_to: None, method: Some("sampling/createMessage".into()), - cap: 100_000 + cap: 100_000, + stand_in: None, } ); // And it never steals a tight cap registered for the host's own request with that id. @@ -1050,4 +1161,38 @@ mod tests { "the panic surfaced with the server's grandchild {orphan} still running" ); } + + /// An over-cap events notification's stand-in keeps what its head shows whole — routing, cursor, + /// id — and none of the payload; anything that is not an events notification gets none. + #[test] + fn an_oversized_events_notification_gets_a_stand_in_from_its_head() { + let head = br#"{"jsonrpc":"2.0","method":"notifications/events/event","params":{"_meta":{"io.modelcontextprotocol/subscriptionId":"s1"},"cursor":"41","eventId":"e9","data":{"blob":"xxxxxxxxxx"#; + let stand_in = oversized_stand_in(head).unwrap(); + assert_eq!(stand_in["method"], "notifications/events/event"); + assert_eq!(stand_in["params"]["cursor"], "41"); + assert_eq!(stand_in["params"]["eventId"], "e9"); + assert_eq!(stand_in["params"]["$oversized"], true); + assert_eq!( + stand_in["params"]["_meta"]["io.modelcontextprotocol/subscriptionId"], + "s1" + ); + assert!(stand_in["params"].get("data").is_none(), "no payload"); + // Members past the head are simply unknown. + let late = br#"{"jsonrpc":"2.0","method":"notifications/events/event","params":{"data":{"blob":"xxxx"#; + let stand_in = oversized_stand_in(late).unwrap(); + assert!(stand_in["params"].get("cursor").is_none()); + for other in [ + &br#"{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1"#[..], + br#"{"jsonrpc":"2.0","id":3,"result":{"x":"yyyy"#, + br#"{"jsonrpc":"2.0","params":{"cursor":"1"},"method":"notifications/events/event"}"#, + b"not json", + ] { + assert_eq!( + oversized_stand_in(other), + None, + "{}", + String::from_utf8_lossy(other) + ); + } + } } diff --git a/crates/agent/src/tools/mcp_wire.rs b/crates/agent/src/tools/mcp_wire.rs index 615f84af..dccfcac4 100644 --- a/crates/agent/src/tools/mcp_wire.rs +++ b/crates/agent/src/tools/mcp_wire.rs @@ -7,10 +7,14 @@ //! [`crate::tools::mcp_stdio::rescue`] / [`crate::tools::mcp_stdio::unwrap_rescued`], applied to the //! raw bytes before `rmcp` parses them. For stdio that is the one stdio transport //! ([`crate::tools::mcp_stdio`]); for streamable HTTP it is [`HttpClient`] here, which wraps the -//! `reqwest` client `rmcp` drives and answers `skills/*` POSTs itself (with the exact URI, session, -//! auth and standard headers the transport handed it), delegating everything else untouched. The -//! same request path ([`HttpClient::post_bounded`]) serves `mcp_view_http`'s capped MCP App view -//! reads. +//! `reqwest` client `rmcp` drives and answers every POST itself (with the exact URI, session, auth +//! and standard headers the transport handed it). The same request path +//! ([`HttpClient::post_bounded`]) serves `mcp_view_http`'s capped MCP App view reads. +//! +//! It is also where every streamable-HTTP message is held to the one per-message cap +//! ([`crate::tools::mcp_stdio::max_message_bytes`], `BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES`): a JSON +//! body or SSE event over it answers its request with an error instead of being buffered whole, and +//! the GET stream is capped through `rmcp`'s own SSE-event limit. //! //! Delete this module when the upstream fix ships. @@ -26,46 +30,39 @@ use rmcp::transport::streamable_http_client::{ }; use serde_json::Value; -/// The `reqwest` client `rmcp`'s streamable-HTTP transport drives, with `skills/*` answered here so -/// their results can be [rescued](crate::tools::mcp_stdio::rescue) before `rmcp` parses them. +/// The `reqwest` client `rmcp`'s streamable-HTTP transport drives, answering **every** POST itself +/// ([`HttpClient::post_bounded`]; rmcp's own `post_message` is never called) so results are +/// [rescued](crate::tools::mcp_stdio::rescue) and size-capped before `rmcp` parses them, and so a +/// 401's status is always visible: rmcp turns a 401 whose body is a JSON-RPC error into an ordinary +/// error *response*, the status lost before `mcp_oauth` could see it and refresh. Here any 401 is +/// `AuthRequired`, for every server. #[derive(Clone)] pub(crate) struct HttpClient { pub(crate) client: reqwest::Client, - /// The server has an `agent mcp-login`: **every** POST is answered here - /// ([`HttpClient::post_bounded`]) rather than by rmcp's client, because rmcp turns a 401 whose - /// body is a JSON-RPC error into an ordinary error *response* — the status lost before - /// `mcp_oauth` could see it and refresh. Here any 401 is `AuthRequired`. - pub(crate) oauth: bool, } impl HttpClient { #[cfg(test)] pub(crate) fn new(client: reqwest::Client) -> Self { - Self { - client, - oauth: false, - } + Self { client } } } -/// The largest JSON body (or SSE event) a POST answered here for an OAuth server may be: the -/// host's per-message cap, the same as the stdio transport's. -const OAUTH_MAX_MESSAGE_BYTES: usize = crate::tools::mcp_stdio::DEFAULT_MAX_MESSAGE_BYTES; - -fn is_skills_request(message: &ClientJsonRpcMessage) -> bool { - // The method is the one thing needed; serializing is how to read it without matching every - // request variant. - matches!(message, ClientJsonRpcMessage::Request(_)) - && serde_json::to_value(message) - .ok() - .and_then(|v| v.get("method").and_then(Value::as_str).map(str::to_string)) - .is_some_and(|m| m.starts_with("skills/")) +/// The limit for an ordinary POST: the per-message cap (no larger than `max`, when the transport +/// names one); over it, a request is answered with a JSON-RPC error (its caller sees an ordinary +/// failed request), anything else fails as an undeliverable message does. +fn ordinary_limit(message: &ClientJsonRpcMessage, max: usize) -> Limit { + // Only the id is needed — read straight off the request, not by serializing the message. + let id = match message { + ClientJsonRpcMessage::Request(request) => serde_json::to_value(&request.id).ok(), + _ => None, + }; + Limit { + max: max.min(crate::tools::mcp_stdio::max_message_bytes()), + over: id.map_or(OverLimit::Fail, OverLimit::Refuse), + } } -/// The largest SSE event a `skills/*` response may carry when the transport names no limit: one -/// listing page or entry of a skill at the spec's per-skill limits, with room to spare. -const DEFAULT_MAX_SKILLS_EVENT: usize = 32 * 1024 * 1024; - /// Headers a configured custom header may not override (`rmcp`'s own reserved set). const RESERVED_HEADERS: [&str; 3] = ["accept", "mcp-session-id", "last-event-id"]; @@ -438,27 +435,8 @@ impl StreamableHttpClient for HttpClient { auth_header: Option, custom_headers: HashMap, ) -> Result> { - if self.oauth || is_skills_request(&message) { - return self - .post_bounded( - uri, - message, - session_id, - auth_header, - custom_headers, - Limit { - max: if self.oauth { - OAUTH_MAX_MESSAGE_BYTES - } else { - DEFAULT_MAX_SKILLS_EVENT - }, - over: OverLimit::Fail, - }, - ) - .await; - } - self.client - .post_message(uri, message, session_id, auth_header, custom_headers) + let limit = ordinary_limit(&message, usize::MAX); + self.post_bounded(uri, message, session_id, auth_header, custom_headers, limit) .await } @@ -471,34 +449,8 @@ impl StreamableHttpClient for HttpClient { custom_headers: HashMap, max_sse_event_size: usize, ) -> Result> { - if self.oauth || is_skills_request(&message) { - return self - .post_bounded( - uri, - message, - session_id, - auth_header, - custom_headers, - Limit { - max: if self.oauth { - OAUTH_MAX_MESSAGE_BYTES.max(max_sse_event_size) - } else { - max_sse_event_size - }, - over: OverLimit::Fail, - }, - ) - .await; - } - self.client - .post_message_with_max_sse_event_size( - uri, - message, - session_id, - auth_header, - custom_headers, - max_sse_event_size, - ) + let limit = ordinary_limit(&message, max_sse_event_size); + self.post_bounded(uri, message, session_id, auth_header, custom_headers, limit) .await } @@ -526,7 +478,14 @@ impl StreamableHttpClient for HttpClient { StreamableHttpError, > { self.client - .get_stream(uri, session_id, last_event_id, auth_header, custom_headers) + .get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_header, + custom_headers, + crate::tools::mcp_stdio::max_message_bytes(), + ) .await } @@ -549,7 +508,7 @@ impl StreamableHttpClient for HttpClient { last_event_id, auth_header, custom_headers, - max_sse_event_size, + max_sse_event_size.min(crate::tools::mcp_stdio::max_message_bytes()), ) .await } @@ -715,6 +674,229 @@ mod tests { )); } + /// Every POST now goes through [`HttpClient::post_bounded`] instead of `rmcp`'s own + /// `reqwest` client — so for an ordinary request (a tool call) the two must agree on every + /// response shape `rmcp` acts on: auth required, insufficient scope, an expired session, + /// accepted, a JSON answer (with its `Mcp-Session-Id`), and an SSE stream. Of the two deliberate + /// divergences, this pins the 401 one last (a 401 with no challenge is still `AuthRequired` + /// here); the other — a non-JSON-RPC success answering a request — is pinned in + /// `a_json_success_that_is_not_json_rpc_is_accepted_for_a_notification_but_fails_a_request`. + #[tokio::test] + async fn an_ordinary_request_is_answered_as_rmcps_own_client_would() { + use futures::StreamExt as _; + agent_core::ensure_provider(); + let call = || -> ClientJsonRpcMessage { + serde_json::from_value(serde_json::json!({ + "jsonrpc": "2.0", "id": 7, "method": "tools/call", "params": { "name": "x" }, + })) + .unwrap() + }; + let shape = + |r: &Result>| match r { + Ok(StreamableHttpPostResponse::Accepted) => "accepted".to_owned(), + Ok(StreamableHttpPostResponse::Json(m, session)) => format!( + "json {} session={session:?}", + serde_json::to_value(m).unwrap()["id"] + ), + Ok(StreamableHttpPostResponse::Sse(_, session)) => { + format!("sse session={session:?}") + } + Err(StreamableHttpError::AuthRequired(_)) => "auth-required".to_owned(), + Err(StreamableHttpError::InsufficientScope(_)) => "insufficient-scope".to_owned(), + Err(StreamableHttpError::SessionExpired) => "session-expired".to_owned(), + Err(e) => format!("other: {e}"), + Ok(_) => "other ok".to_owned(), + }; + let cases: [(&[u8], Option<&str>); 6] = [ + (b"HTTP/1.1 401 Unauthorized\r\nWWW-Authenticate: Bearer realm=\"x\"\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", None), + (b"HTTP/1.1 403 Forbidden\r\nWWW-Authenticate: Bearer scope=\"tools\"\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", None), + (b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", Some("s-1")), + (b"HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", None), + (b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nMcp-Session-Id: s-9\r\nContent-Length: 48\r\nConnection: close\r\n\r\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"content\":[]}}", None), + (b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nMcp-Session-Id: s-9\r\nConnection: close\r\n\r\ndata: {\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"content\":[]}}\n\n", None), + ]; + for (raw, session) in cases { + let ours = HttpClient::new(reqwest::Client::new()) + .post_message( + canned(raw, false).await.into(), + call(), + session.map(Into::into), + None, + HashMap::new(), + ) + .await; + let theirs = reqwest::Client::new() + .post_message( + canned(raw, false).await.into(), + call(), + session.map(Into::into), + None, + HashMap::new(), + ) + .await; + assert_eq!( + shape(&ours), + shape(&theirs), + "for {}", + String::from_utf8_lossy(&raw[..raw.len().min(40)]) + ); + // An SSE answer streams the same response through both. + if let ( + Ok(StreamableHttpPostResponse::Sse(mut a, _)), + Ok(StreamableHttpPostResponse::Sse(mut b, _)), + ) = (ours, theirs) + { + assert_eq!( + a.next().await.unwrap().unwrap().data, + b.next().await.unwrap().unwrap().data + ); + } + } + // The deliberate divergence: a 401 with no `WWW-Authenticate` whose body is a JSON-RPC + // error. rmcp hands that back as an ordinary error *response*, losing the status; here any + // 401 is `AuthRequired`, so `mcp_oauth` can refresh (and any server's caller sees why). + let raw: &[u8] = b"HTTP/1.1 401 Unauthorized\r\nContent-Type: application/json\r\nContent-Length: 69\r\nConnection: close\r\n\r\n{\"jsonrpc\":\"2.0\",\"id\":7,\"error\":{\"code\":-32001,\"message\":\"no token\"}}"; + let ours = HttpClient::new(reqwest::Client::new()) + .post_message( + canned(raw, false).await.into(), + call(), + None, + None, + HashMap::new(), + ) + .await; + let theirs = reqwest::Client::new() + .post_message( + canned(raw, false).await.into(), + call(), + None, + None, + HashMap::new(), + ) + .await; + assert_eq!(shape(&ours), "auth-required"); + assert_ne!( + shape(&theirs), + "auth-required", + "rmcp's own client: {}", + shape(&theirs) + ); + } + + /// …and they send the same request: the bearer token, the session id, `Accept`, the protocol + /// version and custom headers, and the body — so a server's auth and session handling sees + /// no difference. + #[tokio::test] + async fn an_ordinary_request_is_sent_as_rmcps_own_client_would() { + agent_core::ensure_provider(); + async fn capture() -> (String, tokio::sync::oneshot::Receiver) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut got = Vec::new(); + let mut buf = [0u8; 8192]; + loop { + let n = tokio::io::AsyncReadExt::read(&mut stream, &mut buf) + .await + .unwrap(); + got.extend_from_slice(&buf[..n]); + let text = String::from_utf8_lossy(&got).into_owned(); + if let Some((head, body)) = text.split_once("\r\n\r\n") { + let len: usize = head + .lines() + .find_map(|l| { + l.to_ascii_lowercase() + .strip_prefix("content-length:") + .map(|v| v.trim().parse().unwrap()) + }) + .unwrap_or(0); + if body.len() >= len || n == 0 { + let _ = tx.send(text); + break; + } + } + } + let _ = tokio::io::AsyncWriteExt::write_all( + &mut stream, + b"HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await; + }); + (format!("http://{addr}/mcp"), rx) + } + // The request line, the headers a server acts on (names lowercased, sorted), and the body. + let normalize = |raw: String| { + let (head, body) = raw.split_once("\r\n\r\n").unwrap(); + let mut lines = head.lines(); + let request_line = lines.next().unwrap().to_owned(); + let mut headers: Vec = lines + .filter_map(|l| l.split_once(':')) + .map(|(k, v)| format!("{}: {}", k.to_ascii_lowercase(), v.trim())) + .filter(|l| !l.starts_with("host:")) + .collect(); + headers.sort(); + (request_line, headers, body.to_owned()) + }; + let message = || -> ClientJsonRpcMessage { + serde_json::from_value(json!({ + "jsonrpc": "2.0", "id": 7, "method": "tools/call", "params": { "name": "x" }, + })) + .unwrap() + }; + let headers = || { + HashMap::from([ + ( + HeaderName::from_static("mcp-protocol-version"), + HeaderValue::from_static("2025-11-25"), + ), + ( + HeaderName::from_static("x-tenant"), + HeaderValue::from_static("t-1"), + ), + ]) + }; + let (url, ours) = capture().await; + HttpClient::new(reqwest::Client::new()) + .post_message( + url.into(), + message(), + Some("s-1".into()), + Some("tok".into()), + headers(), + ) + .await + .unwrap(); + let (url, theirs) = capture().await; + reqwest::Client::new() + .post_message( + url.into(), + message(), + Some("s-1".into()), + Some("tok".into()), + headers(), + ) + .await + .unwrap(); + let (ours, theirs) = ( + normalize(ours.await.unwrap()), + normalize(theirs.await.unwrap()), + ); + assert_eq!(ours, theirs); + for wanted in [ + "authorization: Bearer tok", + "mcp-session-id: s-1", + "x-tenant: t-1", + ] { + assert!( + ours.1.iter().any(|h| h == wanted), + "{wanted} in {:?}", + ours.1 + ); + } + } + #[tokio::test] async fn a_reserved_custom_header_is_refused() { agent_core::ensure_provider(); @@ -772,10 +954,7 @@ mod tests { async fn a_json_success_that_is_not_json_rpc_is_accepted_for_a_notification_but_fails_a_request() { agent_core::ensure_provider(); - let oauth = HttpClient { - client: reqwest::Client::new(), - oauth: true, - }; + let oauth = HttpClient::new(reqwest::Client::new()); let not_json_rpc: &'static [u8] = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 11\r\nConnection: close\r\n\r\n{\"ok\":true}"; // A notification: accepted, as rmcp's client does. let notification: ClientJsonRpcMessage = serde_json::from_value( @@ -821,14 +1000,14 @@ mod tests { } #[tokio::test] - async fn a_json_body_over_the_cap_is_refused() { + async fn a_json_body_over_the_cap_answers_the_request_with_an_error() { agent_core::ensure_provider(); let url = canned( b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 64\r\nConnection: close\r\n\r\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"skills\":[]},\"pad\":\"xxxxxxxx\"}", false, ) .await; - let e = HttpClient::new(reqwest::Client::new()) + let answer = HttpClient::new(reqwest::Client::new()) .post_message_with_max_sse_event_size( url.into(), skills_list(), @@ -838,8 +1017,36 @@ mod tests { 16, ) .await - .unwrap_err(); - assert!(format!("{e}").contains("maximum size of 16 bytes"), "{e}"); + .unwrap(); + let StreamableHttpPostResponse::Json(message, _) = answer else { + panic!("an over-cap answer to a request is that request's error"); + }; + let text = serde_json::to_string(&message).unwrap(); + assert!(text.contains("over 16 bytes refused"), "{text}"); + } + + /// Every POST — not only `skills/*` — is held to the one per-message cap: a tool call's over-cap + /// answer comes back as that call's error, not read whole. + #[tokio::test] + async fn any_request_is_held_to_the_cap() { + agent_core::ensure_provider(); + let url = canned( + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 64\r\nConnection: close\r\n\r\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"content\":[]},\"pad\":\"xxxxxxx\"}", + false, + ) + .await; + let call: ClientJsonRpcMessage = serde_json::from_value(serde_json::json!({ + "jsonrpc": "2.0", "id": 7, "method": "tools/call", "params": { "name": "x" }, + })) + .unwrap(); + let answer = HttpClient::new(reqwest::Client::new()) + .post_message_with_max_sse_event_size(url.into(), call, None, None, HashMap::new(), 16) + .await + .unwrap(); + let StreamableHttpPostResponse::Json(message, _) = answer else { + panic!("an over-cap answer to a tool call is that call's error"); + }; + assert!(serde_json::to_string(&message).unwrap().contains("refused")); } #[tokio::test] diff --git a/crates/agent/tests/mcp_events_durability.rs b/crates/agent/tests/mcp_events_durability.rs index e91d1e13..aff292da 100644 --- a/crates/agent/tests/mcp_events_durability.rs +++ b/crates/agent/tests/mcp_events_durability.rs @@ -342,3 +342,119 @@ fn an_event_whose_run_fails_is_injected_again_up_to_a_bound() { "a failed run's event is injected again, up to the bound" ); } + +/// A steered batch the model received is delivered, even when a compaction later in the same run +/// summarizes the turn that carried it out of the transcript: delivery is recorded when the model +/// receives it (the run's `Steered` event), not re-derived from the transcript at the run's end — +/// so it is not injected a second time. +#[test] +fn a_steered_batch_summarized_away_by_a_mid_run_compaction_is_not_injected_again() { + let (_fx, mcp_url, fixture) = spawn_http_fixture(&[]); + let home = tempfile::tempdir().unwrap(); + write_settings( + home.path(), + json!([{ + "name": "tickets", "transport": "http", "url": mcp_url, + "events": [{ "name": "ticket.updated", "delivery": "poll", "action": "steer" }], + }]), + ); + let (base, bodies) = spawn_model_server_routed( + vec![ + // The second event has arrived: finish. + ("second, kept".into(), turn_text("done")), + // The first has: answer, reporting a context near the window — so once the second is + // steered in at that stop boundary, the run compacts before its next turn, keeping only + // the newest user turn (the second's) and summarizing the first's away. + ( + "first, summarized away".into(), + turn_text("ok2").replace("\"input_tokens\":12", "\"input_tokens\":900"), + ), + ], + turn_text("ok"), + ); + let sessions = home.path().join("sessions"); + let mut cmd = serve_dir_cmd(BIN, &base, &sessions.to_string_lossy()); + cmd.env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_EVENTS_POLL_FLOOR_MS", "100") + .env("BEYOND_AI_AGENT_MCP_EVENTS_COALESCE_MS", "100") + // A tiny window: the big event pushes the transcript over the threshold, so the run + // compacts after the steered turn and keeps only the newest part. + .args([ + "--context-window", + "1000", + "--compaction-reserve-tokens", + "700", + "--compaction-keep-recent-tokens", + "1", + "--no-session-memory", + ]); + let mut child = cmd.spawn_guarded(); + let mut stdin = child.stdin.take().unwrap(); + let mut frames = Frames::new(&mut child, None); + wait_active(&mut stdin, &mut frames, 1); + // Each turn stalls this long: the window in which the second event must be polled, coalesced + // and steered in before the turn carrying the first ends — wide, so full-suite load on the + // host cannot close it (2.5 s did, once). + send( + &mut stdin, + json!({ "type": "prompt", "id": "long", "message": beyond_ai_test_support::stall_prompt(6000) }), + ); + frames.wait(Duration::from_secs(10), "the run's ack", |f| { + f["type"] == "ack" && f["id"] == "long" + }); + // The first event: steered in, it reaches the model at the stop boundary after the first + // (stalled) turn — as its own user turn. + emit( + &fixture, + json!({ "event_id": "compacted-1", "data": { "summary": "first, summarized away", "detail": "d".repeat(3000) } }), + ); + eventually(Duration::from_secs(20), "the turn that carries it", || { + bodies + .lock() + .unwrap() + .iter() + .any(|b| b.contains("first, summarized away")) + .then_some(()) + }); + // The second, while that turn stalls: steered in at the stop boundary after it, as the newest + // user turn. The compaction that follows keeps only that turn. + emit( + &fixture, + json!({ "event_id": "compacted-2", "data": { "summary": "second, kept" } }), + ); + // Several stalled turns and a compaction: more than `response`'s 30 s under load. + let done = frames.wait(Duration::from_secs(90), "the run's response", |f| { + f["type"] == "response" && f["id"] == "long" + }); + assert_eq!(done["success"], true, "{done:#}"); + assert!( + frames + .seen + .iter() + .any(|f| f["type"] == "event" && f["event"]["kind"] == "compaction_start"), + "the run compacted after the steered turns" + ); + std::thread::sleep(Duration::from_millis(3000)); + assert_eq!( + runs_for_event(&bodies, "first, summarized away"), + 1, + "the model received the first event once, and it was not injected again" + ); + // And it is recorded as delivered, durably: nothing is pending on disk, and the log holds the + // `done` records that say so. + let state_file = std::fs::read_dir(&sessions) + .unwrap() + .map(|e| e.unwrap().path()) + .find(|p| p.to_string_lossy().ends_with(".mcp-events.json")) + .expect("the session's events state"); + eventually( + Duration::from_secs(10), + "the pending queue on disk to drain", + || pending_on_disk(&state_file).is_empty().then_some(()), + ); + let log = std::fs::read_to_string(state_file.with_extension("log")).unwrap(); + assert!( + log.lines().filter(|l| l.contains("\"done\"")).count() >= 2, + "both events have durable `done` records: {log}" + ); +} diff --git a/crates/agent/tests/mcp_events_nested.rs b/crates/agent/tests/mcp_events_nested.rs new file mode 100644 index 00000000..6693a0c7 --- /dev/null +++ b/crates/agent/tests/mcp_events_nested.rs @@ -0,0 +1,220 @@ +//! A server may raise a nested request (here `elicitation/create`) while one of its `events/*` +//! requests is in flight — a unary `events/poll`, or an open `events/stream`. It goes to the session +//! that owns that request, exactly as one raised during a `tools/call` does (the `track_call` rule): +//! that session's client is asked and the server gets the client's answer. Real `serve`, a real +//! stdio fixture, a mock model. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +mod common; + +use std::process::Stdio; +use std::time::Duration; + +use common::mcp_events_fixture::{ + Frames, eventually, send, state, stdio_server, wait_control_file, write_settings, +}; +use common::{BIN, SpawnGuarded, serve_cmd, spawn_model_server_routed, turn_text}; +use serde_json::json; + +fn nested_during(during: &str, delivery: &str) { + let home = tempfile::tempdir().unwrap(); + let control_file = home.path().join("control"); + write_settings( + home.path(), + json!([stdio_server( + "tickets", + &control_file, + json!({ "MCP_FIXTURE_NESTED_DURING": during, "MCP_FIXTURE_HEARTBEAT_MS": "200" }), + json!([{ "name": "ticket.updated", "delivery": delivery, "action": "notify" }]), + )]), + ); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let mut cmd = serve_cmd(BIN, &base, &home.path().join("s.jsonl").to_string_lossy()); + cmd.env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .env("BEYOND_AI_AGENT_MCP_EVENTS_POLL_FLOOR_MS", "100") + .stderr(Stdio::null()); + let mut child = cmd.spawn_guarded(); + let mut stdin = child.stdin.take().unwrap(); + let mut frames = Frames::new(&mut child, None); + let control = wait_control_file(&control_file); + let ask = frames.wait( + Duration::from_secs(30), + "the nested elicitation, asked of this session's client", + |f| f["type"] == "elicitation_request", + ); + assert_eq!(ask["server"], "tickets", "{ask:#}"); + assert!( + ask.to_string().contains(&format!("events/{during}")), + "{ask:#}" + ); + send( + &mut stdin, + json!({ + "type": "elicit", + "request_id": ask["request_id"], + "action": "accept", + "content": { "ok": true }, + }), + ); + let answers = eventually(Duration::from_secs(20), "the server's answer", || { + let st = state(&control); + let a = st["nested_answers"].as_array()?.clone(); + (!a.is_empty()).then_some(a) + }); + assert_eq!(answers[0]["during"], during); + assert_eq!( + answers[0]["answer"]["result"]["action"], "accept", + "the owning session's client answered: {answers:#?}" + ); + assert_eq!(answers[0]["answer"]["result"]["content"]["ok"], true); +} + +#[test] +fn a_nested_elicitation_during_an_events_poll_reaches_the_owning_session() { + nested_during("poll", "poll"); +} + +#[test] +fn a_nested_elicitation_during_an_events_stream_reaches_the_owning_session() { + nested_during("stream", "push"); +} + +/// In a daemon the server connection is shared by every session, so "this session" is not a given: +/// the nested request raised during the events session's `events/poll` must reach a client attached +/// to *that* session — not the process-wide host, and not another session's client. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn in_a_daemon_a_nested_elicitation_during_an_events_poll_reaches_the_events_session() { + let home = tempfile::tempdir().unwrap(); + let control_file = home.path().join("control"); + write_settings( + home.path(), + json!([stdio_server( + "tickets", + &control_file, + json!({ "MCP_FIXTURE_NESTED_DURING": "poll", "MCP_FIXTURE_NESTED_DELAY_MS": "1500" }), + json!([{ "name": "ticket.updated", "delivery": "poll", "action": "notify" }]), + )]), + ); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let (_d, port) = common::mcp_events_fixture::spawn_daemon_env( + home.path(), + &base, + &[], + &[("BEYOND_AI_AGENT_MCP_EVENTS_POLL_FLOOR_MS", "100")], + ); + let control = wait_control_file(&control_file); + // Another session's client is attached too; it must not be the one asked. + let mut other = common::ws_connect(port, Some("bystander")).await; + let mut ws = common::ws_connect(port, Some(common::mcp_events_fixture::EVENTS_SESSION)).await; + let ask = common::mcp_events_fixture::ws_next( + &mut ws, + Duration::from_secs(30), + "the nested elicitation on the events session", + |f| f["type"] == "elicitation_request", + ) + .await; + common::ws_send( + &mut ws, + json!({ "type": "elicit", "request_id": ask["request_id"], "action": "accept", "content": { "ok": true } }), + ) + .await; + let answers = eventually(Duration::from_secs(20), "the server's answer", || { + let a = state(&control)["nested_answers"].as_array()?.clone(); + (!a.is_empty()).then_some(a) + }); + assert_eq!( + answers[0]["answer"]["result"]["action"], "accept", + "{answers:#?}" + ); + // The bystander saw no elicitation. + let stray = tokio::time::timeout(Duration::from_millis(500), async { + while let Some(f) = common::ws_next_frame(&mut other).await { + if f["type"] == "elicitation_request" { + return true; + } + } + false + }) + .await + .unwrap_or(false); + assert!(!stray, "another session's client was asked"); +} + +/// Over direct HTTP (a stateless `2026-07-28` server) nothing routes a server→client request to a +/// session, so it is answered at once with an error — never left unanswered, holding the server. +fn nested_over_direct_http(during: &str, delivery: &str) { + let (_fx, mcp_url, control) = common::mcp_events_fixture::spawn_http_fixture(&[ + ("MCP_FIXTURE_NESTED_DURING", during), + ("MCP_FIXTURE_HEARTBEAT_MS", "200"), + ]); + let home = tempfile::tempdir().unwrap(); + write_settings( + home.path(), + json!([{ + "name": "tickets", "transport": "http", "url": mcp_url, + "events": [{ "name": "ticket.updated", "delivery": delivery, "action": "notify" }], + }]), + ); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let mut cmd = serve_cmd(BIN, &base, &home.path().join("s.jsonl").to_string_lossy()); + cmd.env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .env("BEYOND_AI_AGENT_MCP_EVENTS_POLL_FLOOR_MS", "100") + .stderr(Stdio::null()); + let _child = cmd.spawn_guarded(); + let answers = eventually(Duration::from_secs(15), "the server's answer", || { + let a = state(&control)["nested_answers"].as_array()?.clone(); + (!a.is_empty()).then_some(a) + }); + assert_eq!(answers[0]["during"], during); + assert!( + answers[0]["answer"]["error"].is_object(), + "refused at once, not left to time out: {answers:#?}" + ); +} + +#[test] +fn over_direct_http_a_nested_request_during_an_events_poll_is_refused_not_ignored() { + nested_over_direct_http("poll", "poll"); +} + +#[test] +fn over_direct_http_a_nested_request_during_an_events_stream_is_refused_not_ignored() { + nested_over_direct_http("stream", "push"); +} + +/// A server that floods requests onto an `events/stream` gets every one answered, but one at a +/// time: the client never runs more than one answer at once per stream, however many arrive. +#[test] +fn a_flood_of_server_requests_on_an_events_stream_is_answered_one_at_a_time() { + let (_fx, mcp_url, control) = common::mcp_events_fixture::spawn_http_fixture(&[ + ("MCP_FIXTURE_NESTED_FLOOD", "20"), + ("MCP_FIXTURE_ANSWER_DELAY_MS", "50"), + ("MCP_FIXTURE_HEARTBEAT_MS", "200"), + ]); + let home = tempfile::tempdir().unwrap(); + write_settings( + home.path(), + json!([{ + "name": "tickets", "transport": "http", "url": mcp_url, + "events": [{ "name": "ticket.updated", "delivery": "push", "action": "notify" }], + }]), + ); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let mut cmd = serve_cmd(BIN, &base, &home.path().join("s.jsonl").to_string_lossy()); + cmd.env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .stderr(Stdio::null()); + let _child = cmd.spawn_guarded(); + eventually( + Duration::from_secs(30), + "every flooded request answered", + || (state(&control)["client_answers"].as_u64()? >= 20).then_some(()), + ); + assert_eq!( + state(&control)["answers_max_in_flight"], + 1, + "answered one at a time, not a task per request" + ); +} diff --git a/crates/agent/tests/mcp_events_poll.rs b/crates/agent/tests/mcp_events_poll.rs index 797efa69..36430887 100644 --- a/crates/agent/tests/mcp_events_poll.rs +++ b/crates/agent/tests/mcp_events_poll.rs @@ -249,7 +249,8 @@ fn a_server_without_the_extension_is_reported_and_its_tools_are_unchanged() { let status = s .frames .wait(Duration::from_secs(20), "the subscribe failure", |f| { - f["type"] == "mcp_event_status" && f["kind"] == "error" + // Permanent: the server does not speak the extension. + f["type"] == "mcp_event_status" && f["kind"] == "refused" }); assert!( status["error"] @@ -518,3 +519,52 @@ fn a_slow_server_does_not_serialize_subscriptions_to_others() { assert_eq!(first["success"], true); assert_eq!(s.frames.response("slow")["success"], true); } + +/// "Forbidden" (`-32012`) is how a server answers an expired credential: retried, with credentials +/// resolved afresh, on a short backoff — and the subscription comes up once they work. +#[test] +fn a_forbidden_subscribe_is_retried_and_comes_up_once_credentials_work() { + let home = tempfile::tempdir().unwrap(); + let control_file = home.path().join("control"); + let server = stdio_server( + "tickets", + &control_file, + json!({ "MCP_FIXTURE_FORBID_FIRST": "2" }), + json!([{ "name": "ticket.updated", "delivery": "poll", "action": "notify" }]), + ); + let mut s = start(json!([server]), home); + wait_active(&mut s.stdin, &mut s.frames, 1); + assert!( + !s.frames + .seen + .iter() + .any(|f| f["type"] == "mcp_event_status" && f["kind"] == "refused"), + "a forbidden answer is not a permanent refusal" + ); +} + +/// A server that keeps answering "forbidden" is taken as refusing for good — after a few attempts, +/// not on the first. +#[test] +fn a_subscribe_forbidden_over_and_over_is_eventually_refused() { + let home = tempfile::tempdir().unwrap(); + let control_file = home.path().join("control"); + let server = stdio_server( + "tickets", + &control_file, + json!({ "MCP_FIXTURE_FORBID_FIRST": "1000" }), + json!([{ "name": "ticket.updated", "delivery": "poll", "action": "notify" }]), + ); + let mut s = start(json!([server]), home); + let control = wait_control_file(&control_file); + s.frames.wait(Duration::from_secs(40), "the refusal", |f| { + f["type"] == "mcp_event_status" && f["kind"] == "refused" + }); + let polls = state(&control)["methods"] + .as_array() + .unwrap() + .iter() + .filter(|m| *m == "events/poll") + .count(); + assert!(polls >= 5, "tried {polls} times before giving up"); +} diff --git a/crates/agent/tests/mcp_events_receiver.rs b/crates/agent/tests/mcp_events_receiver.rs index 641d433d..7d98257d 100644 --- a/crates/agent/tests/mcp_events_receiver.rs +++ b/crates/agent/tests/mcp_events_receiver.rs @@ -10,8 +10,9 @@ mod common; use std::time::{Duration, Instant}; use common::mcp_events_fixture::{ - EVENTS_SESSION, control, emit, eventually, raw_request, runs_for_event, spawn_daemon_env, - spawn_daemon_on, spawn_http_fixture, state, write_settings, ws_next, ws_wait_active, + EVENTS_SESSION, control, daemon_sessions, emit, eventually, raw_request, runs_for_event, + spawn_daemon_env, spawn_daemon_on, spawn_http_fixture, state, write_settings, ws_next, + ws_wait_active, }; use common::{ChildGuard, HeldPort, spawn_model_server_routed, turn_text, ws_connect, ws_send}; use serde_json::{Value, json}; @@ -132,10 +133,11 @@ async fn a_delivery_body_that_never_finishes_is_timed_out() { ); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; ws_wait_active(&mut ws).await; - let url = state(&fixture)["hooks"][0]["url"] - .as_str() - .unwrap() - .to_owned(); + let url = eventually(Duration::from_secs(10), "the hook", || { + state(&fixture)["hooks"][0]["url"] + .as_str() + .map(str::to_owned) + }); let path = url.split_once(&format!("127.0.0.1:{port}")).unwrap().1; let started = Instant::now(); let status = raw_request( @@ -288,24 +290,26 @@ async fn a_retry_that_beats_the_resubscribe_after_a_restart_is_told_to_retry_not }); } -/// Only callbacks that will come back are held after a restart: a configured subscription is -/// subscribed again, a runtime one is not — so a retried delivery to the runtime one's old callback -/// gets `410` (stop) at once, rather than `503` (retry) for minutes on end. +/// A runtime subscription (`mcp_events_subscribe`) survives a restart like a configured one: the +/// daemon starts the session that holds it at boot — no client has to come back — and subscribes +/// it again on its old callback, so events reach that session's model as before. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn after_a_restart_a_runtime_subscriptions_old_callback_is_gone_not_held() { +async fn a_runtime_subscription_is_restored_after_a_restart_without_its_client() { let (_fx, mcp_url, fixture) = spawn_http_fixture(&[("MCP_FIXTURE_ALLOW_HTTP_CALLBACK", "1")]); let home = tempfile::tempdir().unwrap(); - write_settings(home.path(), hooks(&mcp_url, "notify")); - let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + // The server is configured, its events are not: nothing here but the runtime subscription. + let mut servers = hooks(&mcp_url, "notify"); + servers[0]["events"] = json!([]); + write_settings(home.path(), servers); + let (base, bodies) = spawn_model_server_routed(vec![], turn_text("noted")); let held = HeldPort::bind(); let port = held.port(); let mut first = daemon(home.path(), &base, &held, &[]); - let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; - ws_wait_active(&mut ws).await; + let mut ws = ws_connect(port, Some("runtime-owner")).await; ws_send( &mut ws, json!({ "type": "mcp_events_subscribe", "id": "s", "server": "hooks", "name": "ticket.updated", - "arguments": { "project": "alpha" }, "delivery": "webhook", "action": "notify" }), + "delivery": "webhook", "action": "follow_up" }), ) .await; let r = ws_next(&mut ws, Duration::from_secs(20), "the subscribe", |f| { @@ -313,58 +317,285 @@ async fn after_a_restart_a_runtime_subscriptions_old_callback_is_gone_not_held() }) .await; assert_eq!(r["success"], true, "{r:#}"); - let runtime_url = eventually(Duration::from_secs(10), "the runtime hook", || { - state(&fixture)["hooks"] - .as_array()? - .iter() - .find(|h| h["project"] == "alpha") - .and_then(|h| h["url"].as_str().map(str::to_owned)) + let url = eventually(Duration::from_secs(10), "the hook", || { + state(&fixture)["hooks"][0]["url"] + .as_str() + .map(str::to_owned) }); - eventually(Duration::from_secs(10), "both callbacks persisted", || { - let saved: Value = - serde_json::from_slice(&std::fs::read(state_file(home.path())).ok()?).ok()?; - (saved["subscriptions"] - .as_object()? - .values() - .filter(|s| s["webhook"]["token"].is_string()) + let subscribes = || { + state(&fixture)["methods"] + .as_array() + .unwrap() + .iter() + .filter(|m| *m == "events/subscribe") .count() - == 2) + }; + eventually(Duration::from_secs(10), "the runtime spec on disk", || { + std::fs::read_dir(home.path().join("sessions")) + .ok()? + .map(|e| e.unwrap().path()) + .filter(|p| p.to_string_lossy().ends_with(".mcp-events.json")) + .any(|p| { + std::fs::read_to_string(p) + .unwrap_or_default() + .contains("\"runtime\"") + }) .then_some(()) }); drop(ws); first.kill().unwrap(); let _ = first.wait(); + let before = subscribes(); // Down, as a crashed daemon is — but still held, so the restart gets the same port. let down = held.down(); drop(down); let _second = daemon(home.path(), &base, &held, &[]); - // The configured one comes back on its old callback. eventually( Duration::from_secs(20), - "the configured resubscribe", - || { - let n = state(&fixture)["methods"] - .as_array() - .unwrap() - .iter() - .filter(|m| *m == "events/subscribe") - .count(); - (n >= 3).then_some(()) - }, + "the subscription, restored with no client", + || (subscribes() > before).then_some(()), + ); + assert_eq!( + state(&fixture)["hooks"][0]["url"].as_str().unwrap(), + url, + "on its old callback" + ); + let r = emit( + &fixture, + json!({ "event_id": "after-restart-1", "data": { "summary": "restored runtime" } }), + ); + assert_eq!(r["deliveries"][0]["status"], 200, "{r:#}"); + eventually( + Duration::from_secs(20), + "the owning session's model run", + || (runs_for_event(&bodies, "restored runtime") >= 1).then_some(()), ); - let path = runtime_url - .split_once(&format!("127.0.0.1:{port}")) +} + +/// A runtime subscription whose server is gone from settings does not pin its session forever: the +/// restore after a restart is refused for good (an unknown server), the subscription is forgotten, +/// the session is reaped like any idle one, and the next restart does not bring it back. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_runtime_subscription_to_a_removed_server_is_forgotten_not_resurrected() { + let (_fx, mcp_url, _fixture) = spawn_http_fixture(&[("MCP_FIXTURE_ALLOW_HTTP_CALLBACK", "1")]); + let home = tempfile::tempdir().unwrap(); + let mut servers = hooks(&mcp_url, "notify"); + servers[0]["events"] = json!([]); + write_settings(home.path(), servers.clone()); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let held = HeldPort::bind(); + let port = held.port(); + let mut first = daemon(home.path(), &base, &held, &[]); + let mut ws = ws_connect(port, Some("orphaned")).await; + ws_send( + &mut ws, + json!({ "type": "mcp_events_subscribe", "id": "s", "server": "hooks", "name": "ticket.updated", + "delivery": "webhook", "action": "notify" }), + ) + .await; + let r = ws_next(&mut ws, Duration::from_secs(20), "the subscribe", |f| { + f["type"] == "response" && f["id"] == "s" + }) + .await; + assert_eq!(r["success"], true, "{r:#}"); + let has_runtime = |home: &std::path::Path| { + std::fs::read_dir(home.join("sessions")) + .unwrap() + .map(|e| e.unwrap().path()) + .filter(|p| p.to_string_lossy().ends_with(".mcp-events.json")) + .any(|p| { + std::fs::read_to_string(p) + .unwrap_or_default() + .contains("\"runtime\"") + }) + }; + eventually(Duration::from_secs(10), "the runtime spec on disk", || { + has_runtime(home.path()).then_some(()) + }); + drop(ws); + common::mcp_events_fixture::sigterm_and_wait(&mut first); + + // The server is renamed: nothing called `hooks` exists any more. + servers[0]["name"] = json!("hooks-renamed"); + write_settings(home.path(), servers); + let (mut second, port) = common::mcp_events_fixture::spawn_daemon( + home.path(), + &base, + &["--session-idle-timeout", "1"], + ); + let mut reaped = false; + for _ in 0..100 { + if daemon_sessions(port).await.get("orphaned") == Some(&false) { + reaped = true; + break; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + assert!(reaped, "nothing keeps the session alive: it is reaped"); + assert!( + !has_runtime(home.path()), + "the dead subscription is forgotten" + ); + // …and its webhook callback with it: no token or secret is left in the snapshot. + let leftover = std::fs::read_dir(home.path().join("sessions")) .unwrap() - .1 - .to_owned(); - let status = raw_request( - port, - &format!("POST {path} HTTP/1.1\r\nHost: x\r\nContent-Length: 2\r\n\r\n"), - b"{}", - Duration::from_secs(5), + .map(|e| e.unwrap().path()) + .filter(|p| p.to_string_lossy().ends_with(".mcp-events.json")) + .any(|p| { + let s = std::fs::read_to_string(p).unwrap_or_default(); + s.contains("\"webhook\"") || s.contains("whsec_") + }); + assert!( + !leftover, + "the forgotten subscription's token and secret are purged" ); - assert_eq!( - status, 410, - "nothing will resubscribe a runtime subscription: its callback is gone" + common::mcp_events_fixture::sigterm_and_wait(&mut second); + + // And a further restart does not resurrect it. + let (_third, port) = common::mcp_events_fixture::spawn_daemon( + home.path(), + &base, + &["--session-idle-timeout", "1"], + ); + for _ in 0..15 { + assert_ne!( + daemon_sessions(port).await.get("orphaned"), + Some(&true), + "not started again at boot" + ); + tokio::time::sleep(Duration::from_millis(200)).await; + } +} + +/// A runtime subscription a single client command created (nothing else from that client) still +/// expires: its creation starts the time to live, so once that has passed with no client it is +/// not restored after a restart. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_runtime_subscription_made_by_one_command_expires_with_no_client() { + let (_fx, mcp_url, fixture) = spawn_http_fixture(&[("MCP_FIXTURE_ALLOW_HTTP_CALLBACK", "1")]); + let home = tempfile::tempdir().unwrap(); + let mut servers = hooks(&mcp_url, "notify"); + servers[0]["events"] = json!([]); + write_settings(home.path(), servers); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let held = HeldPort::bind(); + let port = held.port(); + let mut first = daemon(home.path(), &base, &held, &[]); + let mut ws = ws_connect(port, Some("one-shot")).await; + ws_send( + &mut ws, + json!({ "type": "mcp_events_subscribe", "id": "s", "server": "hooks", "name": "ticket.updated", + "delivery": "webhook", "action": "notify" }), + ) + .await; + let r = ws_next(&mut ws, Duration::from_secs(20), "the subscribe", |f| { + f["type"] == "response" && f["id"] == "s" + }) + .await; + assert_eq!(r["success"], true, "{r:#}"); + let subscribes = || { + state(&fixture)["methods"] + .as_array() + .unwrap() + .iter() + .filter(|m| *m == "events/subscribe") + .count() + }; + eventually(Duration::from_secs(10), "the runtime spec on disk", || { + std::fs::read_dir(home.path().join("sessions")) + .ok()? + .map(|e| e.unwrap().path()) + .filter(|p| p.to_string_lossy().ends_with(".mcp-events.json")) + .any(|p| { + std::fs::read_to_string(p) + .unwrap_or_default() + .contains("\"runtime\"") + }) + .then_some(()) + }); + drop(ws); + first.kill().unwrap(); + let _ = first.wait(); + let before = subscribes(); + tokio::time::sleep(Duration::from_millis(50)).await; + let down = held.down(); + drop(down); + let _second = daemon( + home.path(), + &base, + &held, + &[("BEYOND_AI_AGENT_MCP_EVENTS_RUNTIME_TTL_MS", "1")], + ); + tokio::time::sleep(Duration::from_millis(3000)).await; + assert_eq!(subscribes(), before, "expired: not restored at boot"); + assert_ne!( + daemon_sessions(port).await.get("one-shot"), + Some(&true), + "and its session is not started" + ); +} + +/// A session that holds runtime subscriptions and ends in a panic is started again, so they keep +/// running with no client attached. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_session_with_runtime_subscriptions_is_restarted_after_a_panic() { + let (_fx, mcp_url, fixture) = spawn_http_fixture(&[("MCP_FIXTURE_ALLOW_HTTP_CALLBACK", "1")]); + let home = tempfile::tempdir().unwrap(); + let mut servers = hooks(&mcp_url, "notify"); + servers[0]["events"] = json!([]); + write_settings(home.path(), servers); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let held = HeldPort::bind(); + let port = held.port(); + let _d = daemon( + home.path(), + &base, + &held, + &[("BEYOND_AI_AGENT_TEST_PANICS", "1")], + ); + let mut ws = ws_connect(port, Some("watcher")).await; + ws_send( + &mut ws, + json!({ "type": "mcp_events_subscribe", "id": "s", "server": "hooks", "name": "ticket.updated", + "delivery": "webhook", "action": "notify" }), + ) + .await; + let r = ws_next(&mut ws, Duration::from_secs(20), "the subscribe", |f| { + f["type"] == "response" && f["id"] == "s" + }) + .await; + assert_eq!(r["success"], true, "{r:#}"); + let subscribes = || { + state(&fixture)["methods"] + .as_array() + .unwrap() + .iter() + .filter(|m| *m == "events/subscribe") + .count() + }; + eventually(Duration::from_secs(10), "the runtime spec on disk", || { + std::fs::read_dir(home.path().join("sessions")) + .ok()? + .map(|e| e.unwrap().path()) + .filter(|p| p.to_string_lossy().ends_with(".mcp-events.json")) + .any(|p| { + std::fs::read_to_string(p) + .unwrap_or_default() + .contains("\"runtime\"") + }) + .then_some(()) + }); + let before = subscribes(); + ws_send(&mut ws, json!({ "type": "__test_panic", "id": "boom" })).await; + ws_next(&mut ws, Duration::from_secs(20), "the error frame", |f| { + f["type"] == "error" + }) + .await; + drop(ws); + eventually( + Duration::from_secs(20), + "the session back, resubscribed", + || (subscribes() > before).then_some(()), ); + assert_eq!(daemon_sessions(port).await.get("watcher"), Some(&true)); } diff --git a/crates/agent/tests/mcp_events_routing.rs b/crates/agent/tests/mcp_events_routing.rs index e8aa2e6b..e92da6c9 100644 --- a/crates/agent/tests/mcp_events_routing.rs +++ b/crates/agent/tests/mcp_events_routing.rs @@ -277,3 +277,104 @@ async fn unsubscribing_a_configured_subscription_while_its_server_is_down_stops_ state(&fixture) ); } + +/// A configured subscription the server refuses for good (here: an event it does not offer) is +/// not hammered every minute forever: it is reported as `refused` — in `mcp_events_list` and as an +/// `mcp_event_status` frame — retried only rarely, and it does not keep the events session alive. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_permanently_refused_configured_subscription_is_reported_and_does_not_keep_the_session() { + let (_fx, mcp_url, fixture) = spawn_http_fixture(&[("MCP_FIXTURE_ALLOW_HTTP_CALLBACK", "1")]); + let home = tempfile::tempdir().unwrap(); + let mut servers = hooks(&mcp_url); + servers[0]["events"] = json!([{ "name": "no.such.event", "delivery": "webhook" }]); + write_settings(home.path(), servers); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let (_d, port) = spawn_daemon(home.path(), &base, &["--session-idle-timeout", "1"]); + let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; + let status = ws_next(&mut ws, Duration::from_secs(20), "the refusal", |f| { + f["type"] == "mcp_event_status" && f["kind"] == "refused" + }) + .await; + assert!( + status["error"] + .as_str() + .unwrap() + .contains("offers no event"), + "{status:#}" + ); + let l = ws_list(&mut ws, "l").await; + assert_eq!(l["data"]["unestablished"][0]["state"], "refused", "{l:#}"); + assert_eq!(l["data"]["unestablished"][0]["name"], "no.such.event"); + let lists = || { + state(&fixture)["methods"] + .as_array() + .unwrap() + .iter() + .filter(|m| *m == "events/list") + .count() + }; + let before = lists(); + tokio::time::sleep(Duration::from_millis(5000)).await; + assert!( + lists() <= before + 1, + "a permanent refusal is not retried on the transient backoff" + ); + // Nothing else keeps the events session: once detached, the reaper takes it. + drop(ws); + let mut reaped = false; + for _ in 0..100 { + if daemon_sessions(port).await.get(EVENTS_SESSION) == Some(&false) { + reaped = true; + break; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + assert!( + reaped, + "a refused subscription does not keep its session alive" + ); +} + +/// The session that owns the configured subscriptions is restarted after a panic, so configured +/// triggers keep working without waiting for a daemon restart. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn the_events_session_is_restarted_after_a_panic() { + let (_fx, mcp_url, fixture) = spawn_http_fixture(&[]); + let home = tempfile::tempdir().unwrap(); + write_settings( + home.path(), + json!([{ + "name": "tickets", "transport": "http", "url": mcp_url, + "events": [{ "name": "ticket.updated", "delivery": "poll", "action": "follow_up" }], + }]), + ); + let (base, bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let (_d, port) = common::mcp_events_fixture::spawn_daemon_env( + home.path(), + &base, + &[], + &[ + ("BEYOND_AI_AGENT_TEST_PANICS", "1"), + ("BEYOND_AI_AGENT_MCP_EVENTS_POLL_FLOOR_MS", "100"), + ], + ); + let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; + common::ws_send(&mut ws, json!({ "type": "__test_panic", "id": "boom" })).await; + ws_next(&mut ws, Duration::from_secs(20), "the error frame", |f| { + f["type"] == "error" + }) + .await; + drop(ws); + // Back on its own: a configured event emitted now reaches the model. + tokio::time::sleep(Duration::from_millis(2500)).await; + emit( + &fixture, + json!({ "event_id": "after-panic-1", "data": { "summary": "after the panic" } }), + ); + eventually( + Duration::from_secs(20), + "the event reaching the model", + || (runs_for_event(&bodies, "after the panic") >= 1).then_some(()), + ); + assert_eq!(daemon_sessions(port).await.get(EVENTS_SESSION), Some(&true)); +} diff --git a/crates/agent/tests/mcp_message_cap.rs b/crates/agent/tests/mcp_message_cap.rs new file mode 100644 index 00000000..fbabe284 --- /dev/null +++ b/crates/agent/tests/mcp_message_cap.rs @@ -0,0 +1,265 @@ +//! Every MCP message is held to the one size cap (`BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES`), on every +//! path: a tool result over it — stdio, or streamable HTTP through `mcp_wire::HttpClient` — fails +//! its call with an error (never read whole) and the connection keeps working; and the MCP Events +//! direct-HTTP wire refuses an over-cap unary JSON answer, an over-cap unary SSE answer and an +//! over-cap `events/stream` event. Real `serve`, a real fixture server, a mock model. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +mod common; + +use std::process::Stdio; + +use common::mcp_events_fixture::{ + Frames, fast_knobs, send, spawn_http_fixture, stdio_server, write_settings, +}; +use common::{BIN, SpawnGuarded, serve_cmd, spawn_model_server_routed, turn_text, turn_tool_use}; +use serde_json::{Value, json}; +use std::time::Duration; + +const LIMIT: &str = "1048576"; + +/// A `serve` whose model asks for `mcp__big__blob` with `bytes` when the prompt names a size, +/// and answers "done" once it has the tool result. +fn run_blobs(home: &std::path::Path, server: Value) -> Vec { + write_settings(home, json!([server])); + let (base, bodies) = spawn_model_server_routed( + vec![ + ("tool_result".into(), turn_text("done")), + ( + "blob of 2000000".into(), + turn_tool_use("t_big", "mcp__big__blob", r#"{"bytes":2000000}"#), + ), + ( + "blob of 10".into(), + turn_tool_use("t_small", "mcp__big__blob", r#"{"bytes":10}"#), + ), + ], + turn_text("done"), + ); + let mut cmd = serve_cmd(BIN, &base, &home.join("s.jsonl").to_string_lossy()); + cmd.env("HOME", home) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .env("BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES", LIMIT) + .stderr(Stdio::null()); + let mut child = cmd.spawn_guarded(); + let mut stdin = child.stdin.take().unwrap(); + let mut frames = Frames::new(&mut child, None); + for (id, msg) in [ + ("p1", "make a blob of 2000000"), + ("p2", "make a blob of 10"), + ] { + send( + &mut stdin, + json!({ "type": "prompt", "id": id, "message": msg }), + ); + let r = frames.response(id); + assert_eq!(r["success"], true, "{r:#}"); + // A fresh transcript, so the next prompt's request carries nothing of this one. + send( + &mut stdin, + json!({ "type": "new_session", "id": format!("ns-{id}") }), + ); + frames.response(&format!("ns-{id}")); + } + let requests = bodies.lock().unwrap().clone(); + drop(stdin); + requests +} + +fn assert_capped(requests: &[String]) { + let big = requests + .iter() + .find(|r| r.contains("t_big") && r.contains("tool_result")) + .expect("the model got the big call's result"); + assert!( + big.contains("refused by the host") && !big.contains(&"x".repeat(100_000)), + "the over-limit result failed the call instead of arriving" + ); + let small = requests + .iter() + .find(|r| r.contains("t_small") && r.contains("tool_result")) + .expect("the model got the small call's result"); + assert!( + small.contains("xxxxxxxxxx"), + "the connection still works after the discarded message" + ); +} + +#[test] +fn an_over_limit_message_from_a_stdio_server_fails_its_call_and_the_connection_survives() { + let home = tempfile::tempdir().unwrap(); + let server = stdio_server("big", &home.path().join("control"), json!({}), json!([])); + let requests = run_blobs(home.path(), server); + assert_capped(&requests); +} + +#[test] +fn an_over_limit_message_from_an_http_server_fails_its_call_and_the_connection_survives() { + let (_fx, mcp_url, _control) = spawn_http_fixture(&[]); + let home = tempfile::tempdir().unwrap(); + let server = json!({ "name": "big", "transport": "http", "url": mcp_url }); + let requests = run_blobs(home.path(), server); + assert_capped(&requests); +} + +/// A `serve` (stdio) on one streamable-HTTP fixture server called `tickets`, with a small cap. +fn events_serve( + fixture_env: &[(&str, &str)], + events: Value, +) -> ( + common::ChildGuard, + std::process::ChildStdin, + Frames, + String, + tempfile::TempDir, + common::ChildGuard, +) { + let (fx, mcp_url, control) = spawn_http_fixture(fixture_env); + let home = tempfile::tempdir().unwrap(); + write_settings( + home.path(), + json!([{ "name": "tickets", "transport": "http", "url": mcp_url, "events": events }]), + ); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let mut cmd = serve_cmd(BIN, &base, &home.path().join("s.jsonl").to_string_lossy()); + fast_knobs(&mut cmd) + .env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES", "65536") + .stderr(Stdio::null()); + let mut child = cmd.spawn_guarded(); + let stdin = child.stdin.take().unwrap(); + let frames = Frames::new(&mut child, None); + (child, stdin, frames, control, home, fx) +} + +#[test] +fn an_over_cap_unary_json_answer_on_the_events_wire_is_refused() { + let (_c, mut stdin, mut frames, _control, _home, _fx) = + events_serve(&[("MCP_FIXTURE_LIST_PAD_BYTES", "200000")], json!([])); + send(&mut stdin, json!({ "type": "mcp_events_list", "id": "l" })); + let l = frames.response("l"); + let available = &l["data"]["available"][0]; + assert_eq!(available["supported"], false, "{l:#}"); + assert!( + available["error"] + .as_str() + .unwrap() + .contains("larger than 65536"), + "{l:#}" + ); +} + +#[test] +fn an_over_cap_unary_sse_answer_on_the_events_wire_is_refused() { + let (_c, mut stdin, mut frames, _control, _home, _fx) = + events_serve(&[("MCP_FIXTURE_POLL_SSE_PAD_BYTES", "200000")], json!([])); + send( + &mut stdin, + json!({ "type": "mcp_events_subscribe", "id": "s", "server": "tickets", + "name": "ticket.updated", "delivery": "poll", "action": "notify" }), + ); + let r = frames.response("s"); + assert_eq!( + r["success"], false, + "the over-cap poll answer is not read: {r:#}" + ); +} + +/// An `events/stream` event over the cap is skipped, not read and not reconnected into: its bounded +/// head names it (and its cursor), it is reported as a gap — to the client and to the model — and the +/// next event on the same stream is delivered. Over streamable HTTP (SSE) and over stdio. +fn an_over_cap_push_event_is_skipped_and_reported(stdio: bool) { + let events = json!([{ "name": "ticket.updated", "delivery": "push", "action": "follow_up" }]); + let home = tempfile::tempdir().unwrap(); + let control_file = home.path().join("control"); + let http_fixture = (!stdio).then(|| spawn_http_fixture(&[("MCP_FIXTURE_HEARTBEAT_MS", "200")])); + let server = match &http_fixture { + Some((_fx, mcp_url, _control)) => { + json!({ "name": "tickets", "transport": "http", "url": mcp_url, "events": events }) + } + None => stdio_server( + "tickets", + &control_file, + json!({ "MCP_FIXTURE_HEARTBEAT_MS": "200" }), + events, + ), + }; + write_settings(home.path(), json!([server])); + let (base, bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let mut cmd = serve_cmd(BIN, &base, &home.path().join("s.jsonl").to_string_lossy()); + fast_knobs(&mut cmd) + .env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES", "65536") + .stderr(Stdio::null()); + let mut child = cmd.spawn_guarded(); + let mut stdin = child.stdin.take().unwrap(); + let mut frames = Frames::new(&mut child, None); + let control = match &http_fixture { + Some((_, _, control)) => control.clone(), + None => common::mcp_events_fixture::wait_control_file(&control_file), + }; + common::mcp_events_fixture::wait_active(&mut stdin, &mut frames, 1); + let streams = |control: &str| { + common::mcp_events_fixture::state(control)["methods"] + .as_array() + .unwrap() + .iter() + .filter(|m| *m == "events/stream") + .count() + }; + let streams_before = streams(&control); + common::mcp_events_fixture::emit( + &control, + json!({ "event_id": "huge-1", "data": { "blob": "x".repeat(200_000) } }), + ); + common::mcp_events_fixture::emit( + &control, + json!({ "event_id": "small-1", "data": { "ok": true } }), + ); + let gap = frames.wait( + Duration::from_secs(20), + "the gap for the dropped event", + |f| f["type"] == "mcp_event_status" && f["kind"] == "gap", + ); + assert_eq!(gap["reason"], "oversized", "{gap:#}"); + // The fixture serializes `cursor` ahead of the payload (and `eventId` after it): the cursor is + // kept, the id is not known. + assert!( + gap["cursor"].is_string(), + "the skipped event's cursor is kept: {gap:#}" + ); + let f = frames.wait(Duration::from_secs(20), "the next event", |f| { + f["type"] == "mcp_event" && f["event"]["eventId"] == "small-1" + }); + assert_eq!(f["event"]["data"]["ok"], true); + assert!( + !frames + .seen + .iter() + .any(|f| f["type"] == "mcp_event" && f["event"]["eventId"] == "huge-1"), + "the over-cap event was never read" + ); + common::mcp_events_fixture::eventually(Duration::from_secs(20), "the model told", || { + bodies + .lock() + .unwrap() + .iter() + .any(|b| b.contains("larger than the message-size limit")) + .then_some(()) + }); + assert_eq!( + streams(&control), + streams_before, + "skipped in place: the stream was not reconnected into the same event" + ); +} + +#[test] +fn an_over_cap_push_event_over_http_is_skipped_and_reported_and_the_next_is_delivered() { + an_over_cap_push_event_is_skipped_and_reported(false); +} + +#[test] +fn an_over_cap_push_event_over_stdio_is_skipped_and_reported_and_the_next_is_delivered() { + an_over_cap_push_event_is_skipped_and_reported(true); +} diff --git a/crates/agent/tests/mcp_oauth_routing.rs b/crates/agent/tests/mcp_oauth_routing.rs index dc3eaf45..66b067ba 100644 --- a/crates/agent/tests/mcp_oauth_routing.rs +++ b/crates/agent/tests/mcp_oauth_routing.rs @@ -1,6 +1,6 @@ //! An MCP server with an `agent mcp-login` is spoken to over this crate's own POST path -//! (`mcp_wire::HttpClient::post_bounded`, chosen in `connect_http` by `oauth: auth.is_some()`) instead -//! of rmcp's client — so that path has to carry everything a streamable-HTTP server does: sessions +//! (`mcp_wire::HttpClient::post_bounded`, which answers every streamable-HTTP POST) instead of +//! rmcp's client — so that path has to carry everything a streamable-HTTP server does: sessions //! (`Mcp-Session-Id`, a 404 for an expired one), SSE responses, the standalone `GET` stream, the //! legacy handshake's cues. Each test here drives a real dial (`agent run`) against the OAuth fixture //! speaking that feature, and pins one piece of the routing. @@ -25,8 +25,8 @@ fn result_of(bodies: &[String], request: usize, id: &str) -> (String, bool) { #[test] fn a_401_carrying_a_json_rpc_error_body_refreshes_through_a_real_dial() { - // rmcp's own client reads this 401 as an ordinary error *response*; only the OAuth routing - // (`oauth: auth.is_some()` in `connect_http`) sees the status and refreshes. + // rmcp's own client reads this 401 as an ordinary error *response*; only `post_bounded` + // (every POST's path) sees the status, so the OAuth layer refreshes. let (home, fixture) = logged_in(); fixture.reject_with_json_body.store(true, Ordering::SeqCst); fixture.revoke_after_calls.store(1, Ordering::SeqCst); diff --git a/crates/agent/tests/mcp_stdio_transport.rs b/crates/agent/tests/mcp_stdio_transport.rs index f2fdcc15..2069c5d0 100644 --- a/crates/agent/tests/mcp_stdio_transport.rs +++ b/crates/agent/tests/mcp_stdio_transport.rs @@ -326,3 +326,132 @@ fn a_run_refused_for_bad_configuration_after_connecting_still_sweeps() { "`run` refused its configuration and left its stdio server's grandchild {orphan} running" ); } + +/// A panic inside `serve` (here: in the session loop, while the stdin reader is parked on a read +/// that never returns) must not hang the process in runtime teardown: it exits promptly, and its +/// stdio servers are still swept. +#[test] +fn a_panic_inside_serve_exits_promptly_and_still_sweeps() { + let home = tempfile::tempdir().unwrap(); + let pidfile = home.path().join("orphan.pid"); + let (_child_guard, mut stdin, mut frames) = { + write_settings( + home.path(), + json!([stdio_server( + "tools", + &home.path().join("control"), + json!({ "MCP_FIXTURE_ORPHAN_PIDFILE": pidfile.to_string_lossy() }), + json!([]) + )]), + ); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let mut cmd = serve_cmd(BIN, &base, &home.path().join("s.jsonl").to_string_lossy()); + cmd.env("HOME", home.path()) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .env("BEYOND_AI_AGENT_TEST_PANICS", "1"); + let mut child = cmd.spawn_guarded(); + let stdin = child.stdin.take().unwrap(); + let frames = Frames::new(&mut child, None); + (child, stdin, frames) + }; + let mut child = _child_guard; + send(&mut stdin, json!({ "type": "get_mcp", "id": "m" })); + frames.response("m"); + let orphan: u32 = eventually(Duration::from_secs(10), "the orphan's pid", || { + std::fs::read_to_string(&pidfile).ok()?.trim().parse().ok() + }); + // stdin stays open: the reader thread stays parked on its read. + send(&mut stdin, json!({ "type": "__test_panic", "id": "boom" })); + let deadline = std::time::Instant::now() + Duration::from_secs(15); + let status = loop { + if let Some(status) = child.try_wait().unwrap() { + break status; + } + assert!( + std::time::Instant::now() < deadline, + "serve hung after a panic instead of exiting" + ); + std::thread::sleep(Duration::from_millis(50)); + }; + assert!(!status.success(), "a panic is a failure exit"); + assert!( + !alive(orphan), + "serve panicked and left its stdio server's grandchild {orphan} running" + ); + drop(stdin); +} + +/// In a daemon, a panic inside one session is contained to it: that session's clients get an +/// `error` frame and the session ends, the daemon keeps serving the others, and a later SIGTERM +/// still exits promptly with the stdio servers swept. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_panic_in_a_daemon_session_ends_that_session_only() { + let home = tempfile::tempdir().unwrap(); + let pidfile = home.path().join("orphan.pid"); + write_settings( + home.path(), + json!([stdio_server( + "tools", + &home.path().join("control"), + json!({ "MCP_FIXTURE_ORPHAN_PIDFILE": pidfile.to_string_lossy() }), + json!([]) + )]), + ); + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let (mut daemon, port) = common::mcp_events_fixture::spawn_daemon_env( + home.path(), + &base, + &[], + &[ + ("BEYOND_AI_AGENT_TEST_PANICS", "1"), + ("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0"), + ], + ); + let mut victim = common::ws_connect(port, Some("victim")).await; + let mut survivor = common::ws_connect(port, Some("survivor")).await; + for ws in [&mut victim, &mut survivor] { + common::ws_send(ws, json!({ "type": "get_mcp", "id": "m" })).await; + common::mcp_events_fixture::ws_next(ws, Duration::from_secs(20), "get_mcp", |f| { + f["type"] == "response" && f["id"] == "m" + }) + .await; + } + let orphan: u32 = eventually(Duration::from_secs(10), "the orphan's pid", || { + std::fs::read_to_string(&pidfile).ok()?.trim().parse().ok() + }); + + common::ws_send(&mut victim, json!({ "type": "__test_panic", "id": "boom" })).await; + let err = common::mcp_events_fixture::ws_next( + &mut victim, + Duration::from_secs(20), + "the panicking session's error frame", + |f| f["type"] == "error", + ) + .await; + assert_eq!(err["session_id"], "victim", "{err:#}"); + assert!( + err["error"].as_str().unwrap().contains("internal error"), + "{err:#}" + ); + + // The daemon, and every other session, carry on. + common::ws_send(&mut survivor, json!({ "type": "get_state", "id": "s" })).await; + let state = common::mcp_events_fixture::ws_next( + &mut survivor, + Duration::from_secs(20), + "the other session still answering", + |f| f["type"] == "response" && f["id"] == "s", + ) + .await; + assert_eq!(state["success"], true, "{state:#}"); + assert!( + daemon.try_wait().unwrap().is_none(), + "the daemon is still up" + ); + + common::mcp_events_fixture::sigterm_and_wait(&mut daemon); + assert!( + !alive(orphan), + "the daemon exited and left its stdio server's grandchild {orphan} running" + ); +}