From 29cbcfe0c7778e339db7547ae5affe13fca01961 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 19:47:48 -0700 Subject: [PATCH 01/11] =?UTF-8?q?agent:=20MCP=20Events=20follow-ups=20?= =?UTF-8?q?=E2=80=94=20panic=20exit,=20delivery=20by=20tag,=20permanent=20?= =?UTF-8?q?refusals,=20runtime=20restore?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - A panic out of `run`/`serve` no longer hangs in runtime drop on the parked stdin reader: `main` catches the unwind, abandons the runtime's threads, sweeps stdio servers and exits 101. - Steered MCP Events batches are delivered by tag: `SteeringMessage::tag` comes back in `AgentEvent::Steered` (emitted after the checkpoint on both paths), and the batch is recorded as delivered right then — a later compaction can no longer make it look undelivered. - Permanent refusals (event not offered, extension unsupported, invalid/denied, no mode) are classified, reported as `refused` (frame + `mcp_events_list.unestablished`), retried hourly, and no longer keep the events session alive. - Runtime subscriptions are persisted and restored after a restart; a daemon boots the sessions that hold them, and their webhook tokens are held like configured ones. - Tests for nested elicitation during `events/poll` and `events/stream` (stdio and daemon). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent-core/ARCHITECTURE.md | 7 +- crates/agent-core/src/agent.rs | 61 ++- crates/agent-core/src/steering.rs | 13 + crates/agent/ARCHITECTURE.md | 52 +- .../src/bin/mcp_fixture_events_server.rs | 97 +++- crates/agent/src/main.rs | 15 +- crates/agent/src/serve.rs | 54 +- crates/agent/src/serve_ws.rs | 30 ++ crates/agent/src/tools/mcp_events/mod.rs | 463 ++++++++++++++---- crates/agent/src/tools/mcp_events/state.rs | 35 ++ crates/agent/src/tools/mcp_events/webhook.rs | 16 +- crates/agent/src/tools/mcp_events/wire.rs | 8 + crates/agent/tests/mcp_events_durability.rs | 93 ++++ crates/agent/tests/mcp_events_nested.rs | 144 ++++++ crates/agent/tests/mcp_events_poll.rs | 3 +- crates/agent/tests/mcp_events_receiver.rs | 95 ++-- crates/agent/tests/mcp_events_routing.rs | 58 +++ crates/agent/tests/mcp_stdio_transport.rs | 54 ++ 18 files changed, 1103 insertions(+), 195 deletions(-) create mode 100644 crates/agent/tests/mcp_events_nested.rs diff --git a/crates/agent-core/ARCHITECTURE.md b/crates/agent-core/ARCHITECTURE.md index e5f77d9..cdb7eff 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 ed0d405..959fc23 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 { @@ -1841,6 +1856,7 @@ impl Agent { if steered_count > 0 { sink(AgentEvent::Steered { messages: steered_count, + tags: steered_tags, }); } // A graceful-stop request is honored here too, after this turn's tool results (and any @@ -9584,6 +9600,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 faeb28a..022fae9 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 911b0ce..786fb5b 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -4178,6 +4178,7 @@ 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 | | `--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 | @@ -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,20 @@ 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 does not offer the event, does not speak the extension + (`-32601`), rejects the request as invalid or the access as denied (`-32602`, `-32011`, `-32012`, + `-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. +- **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). - **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 +4621,13 @@ 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). 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, diff --git a/crates/agent/src/bin/mcp_fixture_events_server.rs b/crates/agent/src/bin/mcp_fixture_events_server.rs index 8f21319..f54ea77 100644 --- a/crates/agent/src/bin/mcp_fixture_events_server.rs +++ b/crates/agent/src/bin/mcp_fixture_events_server.rs @@ -108,6 +108,9 @@ 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, } fn env_flag(name: &str) -> bool { @@ -479,7 +482,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") @@ -928,7 +942,7 @@ 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, }), ) } @@ -1399,8 +1413,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 +1480,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(); diff --git a/crates/agent/src/main.rs b/crates/agent/src/main.rs index 2168b30..4495c9d 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 ace7f95..b6f45af 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -4617,7 +4617,11 @@ 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(); + // The steered batches the model has received: recorded — durably, in the events + // state — from the run's `Steered` event, the moment the transcript holding them is + // checkpointed, never re-derived from the transcript afterwards (a compaction may + // have summarized them away by then). + let steered_received: Arc>> = Arc::default(); let result = 'retry: loop { tokens_before.store(0, Ordering::Relaxed); refused.store(false, Ordering::Relaxed); @@ -4671,7 +4675,19 @@ pub(crate) async fn serve_session( } let life_obs = life.clone(); let attempt_result = { + let received_sink = steered_received.clone(); + let receipts_sink = mcp_events.as_ref().map(|hub| hub.receipts()); let mut sink = move |ev: AgentEvent| { + if let AgentEvent::Steered { tags, .. } = &ev + && !tags.is_empty() + { + if let Some(receipts) = &receipts_sink { + for tag in tags { + receipts.received(*tag); + } + } + lock_ignoring_poison(&received_sink).extend(tags.iter().copied()); + } // 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 +4988,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 +5610,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,21 +5627,9 @@ 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 - }; + let received = std::mem::take(&mut *lock_ignoring_poison(&steered_received)); for b in steered_injections.drain(..) { - if run_ok - && crate::tools::mcp_events::transcript_has_injection( - &session.messages[since..], - b, - ) - { - delivered.push(b); - } else { + if !received.contains(&b) { returned.push(b); } } @@ -6110,6 +6115,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 4689fa6..13f2a2a 100644 --- a/crates/agent/src/serve_ws.rs +++ b/crates/agent/src/serve_ws.rs @@ -1765,6 +1765,7 @@ pub async fn serve_ws( }); 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 +1781,35 @@ pub async fn serve_ws( } } + // 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 || { + crate::tools::mcp_events::sessions_with_runtime_subscriptions(std::path::Path::new( + &dir, + )) + }) + .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 diff --git a/crates/agent/src/tools/mcp_events/mod.rs b/crates/agent/src/tools/mcp_events/mod.rs index 16c6ff7..070facb 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,12 @@ fn healthy_for() -> Duration { /// Longest wait between attempts to (re)establish a configured subscription. const CONFIGURED_RETRY_CAP: Duration = Duration::from_secs(60); +/// 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 +368,7 @@ impl SubState { cursor: self.cursor(), recent, webhook: None, + runtime: None, }, ); } @@ -402,6 +409,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 +480,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>, @@ -518,7 +539,8 @@ impl McpEventsHub { shutdown: CancellationToken::new(), http: std::sync::OnceLock::new(), keep_alive: cfg.keep_alive, - starting: AtomicUsize::new(0), + unestablished: Mutex::new(HashMap::new()), + restoring: Mutex::new(HashSet::new()), store: StateStore::open(cfg.state_path, max_pending(), max_pending_bytes()), command_tasks: Mutex::new(tokio::task::JoinSet::new()), owns_configured: cfg.owns_configured, @@ -529,12 +551,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 +566,27 @@ 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); + 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 +625,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 +668,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 +738,70 @@ impl McpEventsHub { } } +/// The ids of the sessions in `dir` whose events state holds runtime subscriptions to restore — 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) -> Vec { + let Ok(entries) = std::fs::read_dir(dir) else { + return Vec::new(); + }; + let mut ids = 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; + }; + let has_runtime = std::fs::read(entry.path()) + .ok() + .and_then(|b| serde_json::from_slice::(&b).ok()) + .and_then(|v| { + v.get("subscriptions")?.as_object().map(|subs| { + subs.values() + .any(|s| s.get("runtime").is_some_and(|r| !r.is_null())) + }) + }) + .unwrap_or(false); + if has_runtime { + ids.push(id.to_owned()); + } + } + ids.sort(); + ids +} + +/// 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 +878,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 +920,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 +983,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 +1058,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 +1081,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 +1108,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; @@ -978,12 +1132,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 +1148,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 +1158,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 +1181,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 +1195,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 +1245,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 +1263,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); @@ -1206,15 +1372,43 @@ 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; @@ -1224,14 +1418,13 @@ impl Hub { () = 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 { @@ -1248,25 +1441,55 @@ impl Hub { break; } Err(e) => { - failures += 1; - delay = backoff(failures, CONFIGURED_RETRY_CAP); + let kind = if e.permanent { + delay = refused_retry(); + "refused" + } 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 }), ); } } } 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 +1525,13 @@ 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: it is not subscribed again after a + // restart either (unless re-discovery brings it back, below). + self.store.set_runtime(&spec.key(), None); if !error.wants_rediscovery() { return; } @@ -1337,8 +1563,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 +1624,68 @@ 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, +} + +impl SubError { + fn permanent(message: impl Into) -> Self { + Self { + message: message.into(), + permanent: true, + } + } + + fn from_rpc(e: &RpcError) -> Self { + Self { + message: e.to_string(), + permanent: e.is_permanent(), + } + } + + 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, + } + } +} + +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 +1700,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 +1794,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 +1838,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 +1870,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 +1893,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 +1913,7 @@ fn on_stream_msg( hub: &Arc, spec: &SubSpec, state: &SubState, - ready: &mut Option>>, + ready: &mut Option, failures: &mut u32, msg: StreamMsg, ) -> Option { @@ -1660,7 +1946,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 +1969,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 +2162,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() @@ -2095,20 +2370,12 @@ 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!(!transcript_has_injection(&in_transcript, 70)); - let _ = eight; + assert!(render_injection(7, &[gap()]).starts_with(&injection_marker(7))); + assert_ne!(render_injection(7, &[gap()]), render_injection(8, &[gap()])); } } diff --git a/crates/agent/src/tools/mcp_events/state.rs b/crates/agent/src/tools/mcp_events/state.rs index 3cdc079..9a8b6d7 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 @@ -245,6 +249,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 +284,33 @@ 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 entry = d.subs.entry(key.to_owned()).or_default(); + if entry.runtime == spec { + return; + } + entry.runtime = spec; + d.snapshot_dirty = true; + } + self.dirty(); + } + + /// 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) { { @@ -907,6 +941,7 @@ mod tests { cursor: Some("7".into()), recent: vec!["a".into()], webhook: None, + runtime: None, }, ); let writer = store.take_writer().unwrap(); diff --git a/crates/agent/src/tools/mcp_events/webhook.rs b/crates/agent/src/tools/mcp_events/webhook.rs index 1e9d0d1..d6bbfd2 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 2d05a7b..f88c3a7 100644 --- a/crates/agent/src/tools/mcp_events/wire.rs +++ b/crates/agent/src/tools/mcp_events/wire.rs @@ -167,6 +167,14 @@ impl RpcError { ) } + /// A refusal retrying will not change: the request is wrong (`-32602`), the server does not + /// speak the extension (`-32601`), the event or the access does not exist (`-32011`, `-32012`), + /// or the mode is unsupported (`-32014`) — unless it is the draft's "re-discover" signal. + pub(super) fn is_permanent(&self) -> bool { + matches!(self.code, Some(-32601 | -32602 | -32011 | -32012 | -32014)) + && !self.wants_rediscovery() + } + /// 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. diff --git a/crates/agent/tests/mcp_events_durability.rs b/crates/agent/tests/mcp_events_durability.rs index e91d1e1..0fd8e55 100644 --- a/crates/agent/tests/mcp_events_durability.rs +++ b/crates/agent/tests/mcp_events_durability.rs @@ -342,3 +342,96 @@ 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(child.stdout.take().unwrap(), None); + wait_active(&mut stdin, &mut frames, 1); + send( + &mut stdin, + json!({ "type": "prompt", "id": "long", "message": beyond_ai_test_support::stall_prompt(2500) }), + ); + 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" } }), + ); + let done = frames.response("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" + ); +} diff --git a/crates/agent/tests/mcp_events_nested.rs b/crates/agent/tests/mcp_events_nested.rs new file mode 100644 index 0000000..23a7953 --- /dev/null +++ b/crates/agent/tests/mcp_events_nested.rs @@ -0,0 +1,144 @@ +//! 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(child.stdout.take().unwrap(), 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 port = common::free_port(); + let _d = common::mcp_events_fixture::spawn_daemon_env( + home.path(), + &base, + port, + &[], + &[("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"); +} diff --git a/crates/agent/tests/mcp_events_poll.rs b/crates/agent/tests/mcp_events_poll.rs index 797efa6..03111b5 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"] diff --git a/crates/agent/tests/mcp_events_receiver.rs b/crates/agent/tests/mcp_events_receiver.rs index 641d433..dda9bec 100644 --- a/crates/agent/tests/mcp_events_receiver.rs +++ b/crates/agent/tests/mcp_events_receiver.rs @@ -288,24 +288,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 +315,57 @@ 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(()) - }, - ); - let path = runtime_url - .split_once(&format!("127.0.0.1:{port}")) - .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), + "the subscription, restored with no client", + || (subscribes() > before).then_some(()), ); assert_eq!( - status, 410, - "nothing will resubscribe a runtime subscription: its callback is gone" + 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(()), ); } diff --git a/crates/agent/tests/mcp_events_routing.rs b/crates/agent/tests/mcp_events_routing.rs index e8aa2e6..97c39b9 100644 --- a/crates/agent/tests/mcp_events_routing.rs +++ b/crates/agent/tests/mcp_events_routing.rs @@ -277,3 +277,61 @@ 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 port = free_port(); + let _d = spawn_daemon(home.path(), &base, port, &["--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" + ); +} diff --git a/crates/agent/tests/mcp_stdio_transport.rs b/crates/agent/tests/mcp_stdio_transport.rs index f2fdcc1..039b18a 100644 --- a/crates/agent/tests/mcp_stdio_transport.rs +++ b/crates/agent/tests/mcp_stdio_transport.rs @@ -326,3 +326,57 @@ 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(child.stdout.take().unwrap(), 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); +} From a8983829d9ce0fe3b01d7a3d7d004452599e6711 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 20:33:34 -0700 Subject: [PATCH 02/11] =?UTF-8?q?agent:=20MCP=20Events=20follow-ups=20?= =?UTF-8?q?=E2=80=94=20audit=20fixes=20(F1=E2=80=93F6),=20session=20panic?= =?UTF-8?q?=20containment,=20events-session=20restart?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Unknown/removed servers are permanent refusals; a restore refused for good forgets the runtime subscription; runtime subscriptions expire after their session's last client command (RUNTIME_TTL_MS, 7 d); boot restoration is bounded (MAX_RESTORED_SESSIONS, 32). - Steered-batch delivery relies only on the Steered tag (the redundant received list is gone); the compaction test now asserts the durable done records and an empty pending queue. - agent-core reports Steered before ending a terminating batch. - Direct-HTTP events requests refuse server->client requests with an error instead of ignoring them (unary SSE and events/stream). - -32012 is retried with fresh credentials on a short backoff, refused only after 5 in a row. - The unused blob fixture tool is gone. - A panicking daemon session is contained (error frame, session ends, daemon carries on); the MCP Events session is restarted after a panic, with backoff. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent-core/src/agent.rs | 65 +++++++++- crates/agent/ARCHITECTURE.md | 32 ++++- .../src/bin/mcp_fixture_events_server.rs | 111 ++++++++++++++++-- crates/agent/src/serve.rs | 29 ++--- crates/agent/src/serve_ws.rs | 78 +++++++++++- crates/agent/src/tools/mcp_events/mod.rs | 89 +++++++++++--- crates/agent/src/tools/mcp_events/state.rs | 101 ++++++++++++++++ crates/agent/src/tools/mcp_events/wire.rs | 70 ++++++++++- crates/agent/tests/mcp_events_durability.rs | 17 +++ crates/agent/tests/mcp_events_nested.rs | 43 +++++++ crates/agent/tests/mcp_events_poll.rs | 49 ++++++++ crates/agent/tests/mcp_events_receiver.rs | 90 +++++++++++++- crates/agent/tests/mcp_events_routing.rs | 46 ++++++++ crates/agent/tests/mcp_stdio_transport.rs | 77 ++++++++++++ 14 files changed, 828 insertions(+), 69 deletions(-) diff --git a/crates/agent-core/src/agent.rs b/crates/agent-core/src/agent.rs index 959fc23..e4ce1fd 100644 --- a/crates/agent-core/src/agent.rs +++ b/crates/agent-core/src/agent.rs @@ -1843,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 @@ -1853,12 +1861,6 @@ impl Agent { }); return Ok(()); } - if steered_count > 0 { - sink(AgentEvent::Steered { - messages: steered_count, - tags: steered_tags, - }); - } // 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 @@ -4914,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 diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 786fb5b..63f0f87 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -4179,6 +4179,7 @@ backend an operator points it at. | `--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 | @@ -4522,20 +4523,29 @@ 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 does not offer the event, does not speak the extension - (`-32601`), rejects the request as invalid or the access as denied (`-32602`, `-32011`, `-32012`, + 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. + 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). + 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. Runtime subscriptions also **expire**: a session's + last client command is recorded (coarsely, in the events state), 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. Boot restoration is bounded too: + at most `BEYOND_AI_AGENT_MCP_EVENTS_MAX_RESTORED_SESSIONS` (32) sessions are started; 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 @@ -4627,7 +4637,15 @@ mcp_events_subscribe (any session) ──► owned by that session 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). The only exits that skip it + 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), so configured triggers + do not stop until the next restart (`the_events_session_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, @@ -4646,6 +4664,10 @@ mcp_events_subscribe (any session) ──► owned by that session 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`). - **`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 f54ea77..0ac435b 100644 --- a/crates/agent/src/bin/mcp_fixture_events_server.rs +++ b/crates/agent/src/bin/mcp_fixture_events_server.rs @@ -111,6 +111,55 @@ struct State { /// `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, + /// 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 { @@ -447,16 +496,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 { @@ -482,18 +535,7 @@ 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") @@ -507,6 +549,7 @@ 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); @@ -1322,6 +1365,11 @@ 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. + if msg.get("method").is_none() { + state.lock().unwrap().client_answers.push(msg.clone()); + 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()) { @@ -1339,7 +1387,35 @@ 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). + 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"; @@ -1347,6 +1423,14 @@ 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 }); + } let terminated_end = loop { let Some(frame) = rx.recv().await else { break false; @@ -1561,6 +1645,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/serve.rs b/crates/agent/src/serve.rs index b6f45af..e24d3fc 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -4617,11 +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(); - // The steered batches the model has received: recorded — durably, in the events - // state — from the run's `Steered` event, the moment the transcript holding them is - // checkpointed, never re-derived from the transcript afterwards (a compaction may - // have summarized them away by then). - let steered_received: Arc>> = Arc::default(); let result = 'retry: loop { tokens_before.store(0, Ordering::Relaxed); refused.store(false, Ordering::Relaxed); @@ -4675,18 +4670,18 @@ pub(crate) async fn serve_session( } let life_obs = life.clone(); let attempt_result = { - let received_sink = steered_received.clone(); + // 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 - && !tags.is_empty() + && let Some(receipts) = &receipts_sink { - if let Some(receipts) = &receipts_sink { - for tag in tags { - receipts.received(*tag); - } + for tag in tags { + receipts.received(*tag); } - lock_ignoring_poison(&received_sink).extend(tags.iter().copied()); } // Set on `CompactionStart`, cleared on literally anything else — see // `is_compacting`'s own declaration above for why that's exact, not a @@ -5627,12 +5622,10 @@ pub(crate) async fn serve_session( returned.push(b) } } - let received = std::mem::take(&mut *lock_ignoring_poison(&steered_received)); - for b in steered_injections.drain(..) { - if !received.contains(&b) { - 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) diff --git a/crates/agent/src/serve_ws.rs b/crates/agent/src/serve_ws.rs index 13f2a2a..7485634 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,6 +1792,7 @@ 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()?; @@ -1781,15 +1812,53 @@ pub async fn serve_ws( } } + // The events session is what keeps configured subscriptions running; if it ends in a panic it + // is started again — with backoff (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(); + tokio::spawn(async move { + let mut attempt = 0u32; + let mut last = std::time::Instant::now(); + while let Some(id) = panicked_rx.recv().await { + if Some(&id) != events_id.as_ref() { + continue; + } + if last.elapsed() > std::time::Duration::from_secs(600) { + attempt = 0; + } + let delay = std::time::Duration::from_secs(1u64 << attempt.min(6)) + .min(std::time::Duration::from_secs(60)); + attempt += 1; + tokio::time::sleep(delay).await; + last = std::time::Instant::now(); + match supervisor.pin(Some(id.clone()), None).await { + Ok(p) => { + supervisor.unpin(&p.id, p.incarnation); + eprintln!("serve: restarted the MCP Events session `{id}` after a panic"); + } + Err(e) => eprintln!( + "serve: could not restart the MCP Events session `{id}`: {}", + 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 || { - crate::tools::mcp_events::sessions_with_runtime_subscriptions(std::path::Path::new( - &dir, - )) + // 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(); @@ -3078,6 +3147,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_events/mod.rs b/crates/agent/src/tools/mcp_events/mod.rs index 070facb..07d7ab9 100644 --- a/crates/agent/src/tools/mcp_events/mod.rs +++ b/crates/agent/src/tools/mcp_events/mod.rs @@ -166,6 +166,10 @@ 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 { @@ -518,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; } @@ -541,7 +550,7 @@ impl McpEventsHub { keep_alive: cfg.keep_alive, unestablished: Mutex::new(HashMap::new()), restoring: Mutex::new(HashSet::new()), - store: StateStore::open(cfg.state_path, max_pending(), max_pending_bytes()), + store, command_tasks: Mutex::new(tokio::task::JoinSet::new()), owns_configured: cfg.owns_configured, }); @@ -566,6 +575,13 @@ impl McpEventsHub { let Some(hub) = weak_hub.upgrade() else { return; }; + // Runtime subscriptions whose session has heard from no client for longer than + // their time to live are forgotten, not restored. + if !state::runtime_still_wanted(store.last_client_ms(), now_unix_ms()) { + for (key, _) in store.runtime_specs() { + store.set_runtime(&key, None); + } + } let runtime: Vec<(String, SubSpec)> = store .runtime_specs() .into_iter() @@ -738,13 +754,15 @@ impl McpEventsHub { } } -/// The ids of the sessions in `dir` whose events state holds runtime subscriptions to restore — 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) -> Vec { +/// 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 ids = Vec::new(); for entry in entries.flatten() { let name = entry.file_name(); @@ -757,24 +775,28 @@ pub fn sessions_with_runtime_subscriptions(dir: &std::path::Path) -> Vec let Some((_, id)) = stem.split_once('_') else { continue; }; - let has_runtime = std::fs::read(entry.path()) - .ok() - .and_then(|b| serde_json::from_slice::(&b).ok()) - .and_then(|v| { - v.get("subscriptions")?.as_object().map(|subs| { - subs.values() - .any(|s| s.get("runtime").is_some_and(|r| !r.is_null())) - }) - }) - .unwrap_or(false); - if has_runtime { + if std::fs::read(entry.path()).is_ok_and(|b| state::snapshot_wants_restore(&b, now)) { ids.push(id.to_owned()); } } ids.sort(); + 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 { @@ -1123,6 +1145,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() @@ -1412,6 +1441,7 @@ impl Hub { tokio::spawn(async move { let mut delay = delay; let mut failures = 0u32; + let mut forbidden_streak = 0u32; loop { tokio::select! { () = tokio::time::sleep(delay) => {} @@ -1440,10 +1470,22 @@ impl Hub { } break; } - Err(e) => { + 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); @@ -1476,6 +1518,13 @@ impl Hub { 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: it is not + // retried, not restored again, and keeps nothing alive. + if restoring && e.permanent { + hub.store.set_runtime(&key, None); + break; + } } } } @@ -1631,6 +1680,9 @@ async fn stop_task(task: tokio::task::JoinHandle<()>) { 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 { @@ -1638,6 +1690,7 @@ impl SubError { Self { message: message.into(), permanent: true, + forbidden: false, } } @@ -1645,6 +1698,7 @@ impl SubError { Self { message: e.to_string(), permanent: e.is_permanent(), + forbidden: e.is_forbidden(), } } @@ -1661,6 +1715,7 @@ impl From for SubError { Self { message, permanent: false, + forbidden: false, } } } diff --git a/crates/agent/src/tools/mcp_events/state.rs b/crates/agent/src/tools/mcp_events/state.rs index 9a8b6d7..19de3ac 100644 --- a/crates/agent/src/tools/mcp_events/state.rs +++ b/crates/agent/src/tools/mcp_events/state.rs @@ -127,6 +127,36 @@ 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`] (or the state predates the record, and gets the benefit of the +/// doubt once — the next client command records it). +pub(super) fn runtime_still_wanted(last_client_ms: Option, now_ms: i64) -> bool { + last_client_ms.is_none_or(|t| now_ms.saturating_sub(t) <= runtime_ttl_ms()) +} + +/// 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.) +pub(super) fn snapshot_wants_restore(bytes: &[u8], now_ms: i64) -> bool { + let Ok(snap) = serde_json::from_slice::(bytes) else { + return false; + }; + runtime_still_wanted(snap.last_client_ms, now_ms) + && snap.subscriptions.values().any(|s| s.runtime.is_some()) } #[derive(serde::Serialize, serde::Deserialize)] @@ -160,6 +190,7 @@ struct Data { snapshot_dirty: bool, /// Bytes in the log file on disk. log_len: u64, + last_client_ms: Option, } impl Data { @@ -302,6 +333,28 @@ impl StateStore { 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) @@ -593,6 +646,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; } @@ -701,6 +755,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) @@ -811,6 +866,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), ) @@ -1283,4 +1339,49 @@ 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), + "older state gets one chance" + ); + let snap = |last: i64| { + serde_json::json!({ + "subscriptions": { "k": { "cursor": null, "runtime": { "server": "s", "name": "n" } } }, + "last_client_ms": last, + }) + .to_string() + }; + assert!(snapshot_wants_restore(snap(now - day).as_bytes(), now)); + assert!(!snapshot_wants_restore(snap(now - 8 * day).as_bytes(), now)); + let no_runtime = + serde_json::json!({ "subscriptions": { "k": { "cursor": null } } }).to_string(); + assert!(!snapshot_wants_restore(no_runtime.as_bytes(), now)); + } + + /// 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" }))); + 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/wire.rs b/crates/agent/src/tools/mcp_events/wire.rs index f88c3a7..05b9838 100644 --- a/crates/agent/src/tools/mcp_events/wire.rs +++ b/crates/agent/src/tools/mcp_events/wire.rs @@ -168,11 +168,17 @@ impl RpcError { } /// A refusal retrying will not change: the request is wrong (`-32602`), the server does not - /// speak the extension (`-32601`), the event or the access does not exist (`-32011`, `-32012`), - /// or the mode is unsupported (`-32014`) — unless it is the draft's "re-discover" signal. + /// 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 | -32012 | -32014)) - && !self.wants_rediscovery() + 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 @@ -478,6 +484,9 @@ impl Conn { if msg.get("id") == Some(&want) { return rpc_outcome(msg); } + if is_server_request(&msg) { + refuse_server_request(http, url, headers, protocol_version, &msg).await; + } } Err(RpcError::local(format!( "{method}: the stream ended without a response" @@ -573,9 +582,22 @@ 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(), + protocol_version.clone(), + ); let reader = tokio::spawn(async move { let mut events = SseReader::new(resp); while let Some(msg) = events.next().await { + if is_server_request(&msg) { + let (http, url, headers, version) = refuse_with.clone(); + tokio::spawn(async move { + refuse_server_request(&http, &url, &headers, &version, &msg).await; + }); + continue; + } if msg.get("id") == Some(&want) { let fin = match rpc_outcome(msg) { Ok(v) => StreamMsg::final_ok(v), @@ -626,6 +648,46 @@ 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)], + 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"), + }, + }); + let mut req = http + .post(url) + .header("Content-Type", "application/json") + .header("MCP-Protocol-Version", protocol_version); + for (k, v) in headers { + req = req.header(k, v); + } + if let Ok(Err(e)) = + tokio::time::timeout(Duration::from_secs(5), req.body(body.to_string()).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) diff --git a/crates/agent/tests/mcp_events_durability.rs b/crates/agent/tests/mcp_events_durability.rs index 0fd8e55..4d11f40 100644 --- a/crates/agent/tests/mcp_events_durability.rs +++ b/crates/agent/tests/mcp_events_durability.rs @@ -434,4 +434,21 @@ fn a_steered_batch_summarized_away_by_a_mid_run_compaction_is_not_injected_again 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 index 23a7953..799c614 100644 --- a/crates/agent/tests/mcp_events_nested.rs +++ b/crates/agent/tests/mcp_events_nested.rs @@ -142,3 +142,46 @@ async fn in_a_daemon_a_nested_elicitation_during_an_events_poll_reaches_the_even .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"); +} diff --git a/crates/agent/tests/mcp_events_poll.rs b/crates/agent/tests/mcp_events_poll.rs index 03111b5..3643088 100644 --- a/crates/agent/tests/mcp_events_poll.rs +++ b/crates/agent/tests/mcp_events_poll.rs @@ -519,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 dda9bec..82e2f44 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}; @@ -369,3 +370,88 @@ async fn a_runtime_subscription_is_restored_after_a_restart_without_its_client() || (runs_for_event(&bodies, "restored runtime") >= 1).then_some(()), ); } + +/// 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 port = free_port(); + let mut first = daemon(home.path(), &base, port, &[]); + 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 port = free_port(); + let mut second = common::mcp_events_fixture::spawn_daemon( + home.path(), + &base, + port, + &["--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" + ); + common::mcp_events_fixture::sigterm_and_wait(&mut second); + + // And a further restart does not resurrect it. + let port = free_port(); + let _third = common::mcp_events_fixture::spawn_daemon( + home.path(), + &base, + port, + &["--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; + } +} diff --git a/crates/agent/tests/mcp_events_routing.rs b/crates/agent/tests/mcp_events_routing.rs index 97c39b9..2f2e5c3 100644 --- a/crates/agent/tests/mcp_events_routing.rs +++ b/crates/agent/tests/mcp_events_routing.rs @@ -335,3 +335,49 @@ async fn a_permanently_refused_configured_subscription_is_reported_and_does_not_ "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 port = free_port(); + let _d = common::mcp_events_fixture::spawn_daemon_env( + home.path(), + &base, + port, + &[], + &[ + ("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_stdio_transport.rs b/crates/agent/tests/mcp_stdio_transport.rs index 039b18a..4265a20 100644 --- a/crates/agent/tests/mcp_stdio_transport.rs +++ b/crates/agent/tests/mcp_stdio_transport.rs @@ -380,3 +380,80 @@ fn a_panic_inside_serve_exits_promptly_and_still_sweeps() { ); 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 port = common::free_port(); + let mut daemon = common::mcp_events_fixture::spawn_daemon_env( + home.path(), + &base, + port, + &[], + &[ + ("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" + ); +} From 8fddf7ea4ad97610967fd5eef7160b26259875a2 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 20:47:56 -0700 Subject: [PATCH 03/11] =?UTF-8?q?agent:=20one=20MCP=20message=20cap=20on?= =?UTF-8?q?=20every=20path=20=E2=80=94=20the=20events=20HTTP=20wire=20and?= =?UTF-8?q?=20every=20HttpClient=20POST?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends #134's single per-message cap (mcp_stdio::max_message_bytes, BEYOND_AI_AGENT_MCP_MAX_MESSAGE_BYTES) to the paths it did not cover: the MCP Events direct-HTTP wire (unary JSON, unary SSE, events/stream — its own 4 MiB / 1 MiB caps are gone) and every mcp_wire::HttpClient POST (not only skills/*; a request's over-cap answer becomes its JSON-RPC error), plus the GET stream through rmcp's SSE-event limit. The skills-only 32 MiB constant is gone. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/ARCHITECTURE.md | 4 +- .../src/bin/mcp_fixture_events_server.rs | 29 +++ crates/agent/src/tools/mcp.rs | 2 - crates/agent/src/tools/mcp_events/wire.rs | 35 ++-- crates/agent/src/tools/mcp_oauth.rs | 1 - crates/agent/src/tools/mcp_stdio.rs | 4 +- crates/agent/src/tools/mcp_wire.rs | 163 +++++++------- crates/agent/tests/mcp_events_receiver.rs | 9 +- crates/agent/tests/mcp_message_cap.rs | 198 ++++++++++++++++++ 9 files changed, 332 insertions(+), 113 deletions(-) create mode 100644 crates/agent/tests/mcp_message_cap.rs diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 63f0f87..7c17a43 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -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` — 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. | @@ -4660,7 +4660,7 @@ 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 one fails its request or ends its stream, never read whole; a server that keeps resending a giant stream event keeps failing, on backoff) — 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`. diff --git a/crates/agent/src/bin/mcp_fixture_events_server.rs b/crates/agent/src/bin/mcp_fixture_events_server.rs index 0ac435b..986b9d0 100644 --- a/crates/agent/src/bin/mcp_fixture_events_server.rs +++ b/crates/agent/src/bin/mcp_fixture_events_server.rs @@ -535,7 +535,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") @@ -556,6 +567,12 @@ async fn rpc( 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" => { @@ -1395,6 +1412,18 @@ async fn handle_http(state: Shared, mut stream: TcpStream) { // `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; diff --git a/crates/agent/src/tools/mcp.rs b/crates/agent/src/tools/mcp.rs index 408bd3e..8f9bdf8 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/wire.rs b/crates/agent/src/tools/mcp_events/wire.rs index 05b9838..fc4ce3e 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 @@ -471,7 +473,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(|_| { @@ -572,7 +574,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) @@ -705,7 +707,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 { @@ -717,16 +719,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, } } @@ -750,8 +759,8 @@ 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 { + self.error = Some(format!("an SSE event exceeded {} bytes", self.max)); return None; } match self.resp.chunk().await { @@ -1040,9 +1049,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. diff --git a/crates/agent/src/tools/mcp_oauth.rs b/crates/agent/src/tools/mcp_oauth.rs index 67dcb50..03aab10 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 aaf8ffa..9d2a4ef 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()) diff --git a/crates/agent/src/tools/mcp_wire.rs b/crates/agent/src/tools/mcp_wire.rs index 615f84a..2e66c25 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,38 @@ 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 { + let id = matches!(message, ClientJsonRpcMessage::Request(_)) + .then(|| serde_json::to_value(message).ok()) + .flatten() + .and_then(|v| v.get("id").cloned()); + 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 +434,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 +448,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 +477,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 +507,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 } @@ -772,10 +730,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 +776,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 +793,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_receiver.rs b/crates/agent/tests/mcp_events_receiver.rs index 82e2f44..9335bb8 100644 --- a/crates/agent/tests/mcp_events_receiver.rs +++ b/crates/agent/tests/mcp_events_receiver.rs @@ -133,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( diff --git a/crates/agent/tests/mcp_message_cap.rs b/crates/agent/tests/mcp_message_cap.rs new file mode 100644 index 0000000..36ecfbb --- /dev/null +++ b/crates/agent/tests/mcp_message_cap.rs @@ -0,0 +1,198 @@ +//! 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(child.stdout.take().unwrap(), 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(child.stdout.take().unwrap(), 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 ends the stream with an error rather than being read +/// whole; the subscription retries (with backoff — a server that keeps resending the same giant event +/// keeps failing, it does not take the host's memory), and the session stays responsive. +#[test] +fn an_over_cap_event_on_an_events_stream_is_refused_not_read_whole() { + let (_c, mut stdin, mut frames, control, _home, _fx) = events_serve( + &[("MCP_FIXTURE_HEARTBEAT_MS", "200")], + json!([{ "name": "ticket.updated", "delivery": "push", "action": "notify" }]), + ); + common::mcp_events_fixture::wait_active(&mut stdin, &mut frames, 1); + common::mcp_events_fixture::emit( + &control, + json!({ "event_id": "huge-1", "data": { "blob": "x".repeat(200_000) } }), + ); + let status = frames.wait(Duration::from_secs(20), "the stream's refusal", |f| { + f["type"] == "mcp_event_status" + && f["kind"] == "error" + && f.to_string().contains("exceeded") + }); + assert!(status.to_string().contains("65536"), "{status:#}"); + send(&mut stdin, json!({ "type": "mcp_events_list", "id": "l" })); + let l = frames.response("l"); + assert_eq!(l["success"], true, "the session stays responsive: {l:#}"); + assert!( + !frames + .seen + .iter() + .any(|f| f["type"] == "mcp_event" && f["event"]["eventId"] == "huge-1"), + "the over-cap event was never delivered" + ); +} From 6d125b6bea64fa5d7d7dc943ee36ac95d2373f18 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 21:25:26 -0700 Subject: [PATCH 04/11] agent: an over-cap MCP Events push event is skipped and reported, not reconnected into forever MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The over-cap event's bounded head is read structurally (mcp_stdio::oversized_stand_in — the same member walk as scan_head) for its routing, cursor and id; a small $oversized stand-in replaces it on both transports (the stdio pump, and the events SSE reader, which skips the rest of the event unread). The subscription keeps the cursor (else the next heartbeat's), records an 'oversized' gap (frame + model notice) and carries on — no reconnect into the same giant event. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/src/tools/mcp_events/mod.rs | 54 +++++++- crates/agent/src/tools/mcp_events/wire.rs | 86 +++++++++++- crates/agent/src/tools/mcp_stdio.rs | 161 ++++++++++++++++++++-- crates/agent/tests/mcp_message_cap.rs | 103 +++++++++++--- 4 files changed, 374 insertions(+), 30 deletions(-) diff --git a/crates/agent/src/tools/mcp_events/mod.rs b/crates/agent/src/tools/mcp_events/mod.rs index 07d7ab9..24a71bd 100644 --- a/crates/agent/src/tools/mcp_events/mod.rs +++ b/crates/agent/src/tools/mcp_events/mod.rs @@ -1364,6 +1364,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() @@ -1371,7 +1394,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, @@ -1379,7 +1408,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"); @@ -1986,6 +2015,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. @@ -2241,6 +2277,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 \ diff --git a/crates/agent/src/tools/mcp_events/wire.rs b/crates/agent/src/tools/mcp_events/wire.rs index fc4ce3e..224f0ef 100644 --- a/crates/agent/src/tools/mcp_events/wire.rs +++ b/crates/agent/src/tools/mcp_events/wire.rs @@ -743,7 +743,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 { @@ -760,6 +773,16 @@ 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() > 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; } @@ -769,6 +792,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 @@ -1086,6 +1147,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(); diff --git a/crates/agent/src/tools/mcp_stdio.rs b/crates/agent/src/tools/mcp_stdio.rs index 9d2a4ef..167a75f 100644 --- a/crates/agent/src/tools/mcp_stdio.rs +++ b/crates/agent/src/tools/mcp_stdio.rs @@ -359,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; @@ -419,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, } @@ -439,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), }); @@ -472,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(); } } @@ -480,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) @@ -599,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 @@ -712,7 +816,8 @@ mod tests { Inbound::Refused { reply_to: Some("7".into()), method: None, - cap: view_cap + cap: view_cap, + stand_in: None, } ); assert!( @@ -748,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}"); @@ -765,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 @@ -802,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. @@ -837,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. @@ -1052,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/tests/mcp_message_cap.rs b/crates/agent/tests/mcp_message_cap.rs index 36ecfbb..6f09199 100644 --- a/crates/agent/tests/mcp_message_cap.rs +++ b/crates/agent/tests/mcp_message_cap.rs @@ -165,34 +165,101 @@ fn an_over_cap_unary_sse_answer_on_the_events_wire_is_refused() { ); } -/// An `events/stream` event over the cap ends the stream with an error rather than being read -/// whole; the subscription retries (with backoff — a server that keeps resending the same giant event -/// keeps failing, it does not take the host's memory), and the session stays responsive. -#[test] -fn an_over_cap_event_on_an_events_stream_is_refused_not_read_whole() { - let (_c, mut stdin, mut frames, control, _home, _fx) = events_serve( - &[("MCP_FIXTURE_HEARTBEAT_MS", "200")], - json!([{ "name": "ticket.updated", "delivery": "push", "action": "notify" }]), - ); +/// 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(child.stdout.take().unwrap(), 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) } }), ); - let status = frames.wait(Duration::from_secs(20), "the stream's refusal", |f| { - f["type"] == "mcp_event_status" - && f["kind"] == "error" - && f.to_string().contains("exceeded") + 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!(status.to_string().contains("65536"), "{status:#}"); - send(&mut stdin, json!({ "type": "mcp_events_list", "id": "l" })); - let l = frames.response("l"); - assert_eq!(l["success"], true, "the session stays responsive: {l:#}"); + 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 delivered" + "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); +} From 4d703cbcad017f9169d3fa66950593522716ddec Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 21:38:28 -0700 Subject: [PATCH 05/11] agent tests: follow-up tests on #135's race-free daemon ports Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/tests/mcp_events_nested.rs | 4 +--- crates/agent/tests/mcp_events_receiver.rs | 13 +++++-------- crates/agent/tests/mcp_events_routing.rs | 7 ++----- crates/agent/tests/mcp_stdio_transport.rs | 4 +--- 4 files changed, 9 insertions(+), 19 deletions(-) diff --git a/crates/agent/tests/mcp_events_nested.rs b/crates/agent/tests/mcp_events_nested.rs index 799c614..8cee6e4 100644 --- a/crates/agent/tests/mcp_events_nested.rs +++ b/crates/agent/tests/mcp_events_nested.rs @@ -97,11 +97,9 @@ async fn in_a_daemon_a_nested_elicitation_during_an_events_poll_reaches_the_even )]), ); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = common::free_port(); - let _d = common::mcp_events_fixture::spawn_daemon_env( + let (_d, port) = common::mcp_events_fixture::spawn_daemon_env( home.path(), &base, - port, &[], &[("BEYOND_AI_AGENT_MCP_EVENTS_POLL_FLOOR_MS", "100")], ); diff --git a/crates/agent/tests/mcp_events_receiver.rs b/crates/agent/tests/mcp_events_receiver.rs index 9335bb8..1a215b5 100644 --- a/crates/agent/tests/mcp_events_receiver.rs +++ b/crates/agent/tests/mcp_events_receiver.rs @@ -383,8 +383,9 @@ async fn a_runtime_subscription_to_a_removed_server_is_forgotten_not_resurrected servers[0]["events"] = json!([]); write_settings(home.path(), servers.clone()); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let mut first = daemon(home.path(), &base, port, &[]); + 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, @@ -417,11 +418,9 @@ async fn a_runtime_subscription_to_a_removed_server_is_forgotten_not_resurrected // The server is renamed: nothing called `hooks` exists any more. servers[0]["name"] = json!("hooks-renamed"); write_settings(home.path(), servers); - let port = free_port(); - let mut second = common::mcp_events_fixture::spawn_daemon( + let (mut second, port) = common::mcp_events_fixture::spawn_daemon( home.path(), &base, - port, &["--session-idle-timeout", "1"], ); let mut reaped = false; @@ -440,11 +439,9 @@ async fn a_runtime_subscription_to_a_removed_server_is_forgotten_not_resurrected common::mcp_events_fixture::sigterm_and_wait(&mut second); // And a further restart does not resurrect it. - let port = free_port(); - let _third = common::mcp_events_fixture::spawn_daemon( + let (_third, port) = common::mcp_events_fixture::spawn_daemon( home.path(), &base, - port, &["--session-idle-timeout", "1"], ); for _ in 0..15 { diff --git a/crates/agent/tests/mcp_events_routing.rs b/crates/agent/tests/mcp_events_routing.rs index 2f2e5c3..e92da6c 100644 --- a/crates/agent/tests/mcp_events_routing.rs +++ b/crates/agent/tests/mcp_events_routing.rs @@ -289,8 +289,7 @@ async fn a_permanently_refused_configured_subscription_is_reported_and_does_not_ 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 port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &["--session-idle-timeout", "1"]); + 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" @@ -350,11 +349,9 @@ async fn the_events_session_is_restarted_after_a_panic() { }]), ); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = common::mcp_events_fixture::spawn_daemon_env( + let (_d, port) = common::mcp_events_fixture::spawn_daemon_env( home.path(), &base, - port, &[], &[ ("BEYOND_AI_AGENT_TEST_PANICS", "1"), diff --git a/crates/agent/tests/mcp_stdio_transport.rs b/crates/agent/tests/mcp_stdio_transport.rs index 4265a20..3bfdfe4 100644 --- a/crates/agent/tests/mcp_stdio_transport.rs +++ b/crates/agent/tests/mcp_stdio_transport.rs @@ -398,11 +398,9 @@ async fn a_panic_in_a_daemon_session_ends_that_session_only() { )]), ); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = common::free_port(); - let mut daemon = common::mcp_events_fixture::spawn_daemon_env( + let (mut daemon, port) = common::mcp_events_fixture::spawn_daemon_env( home.path(), &base, - port, &[], &[ ("BEYOND_AI_AGENT_TEST_PANICS", "1"), From 93fbb4a05596bd1ebcd41b666611ce6587a5ffde Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 22:17:22 -0700 Subject: [PATCH 06/11] =?UTF-8?q?agent:=20MCP=20Events=20re-audit=20fixes?= =?UTF-8?q?=20(N1=E2=80=93N7)=20=E2=80=94=20runtime=20TTL=20from=20creatio?= =?UTF-8?q?n,=20refusal=20Accept,=20inline=20refusals,=20boot=20ranking,?= =?UTF-8?q?=20purge,=20restart,=20one=20POST=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - N1: creating a runtime subscription records client activity (set_runtime), and a state with runtime subscriptions but no record is no longer trusted forever: one made by a single command expires one TTL after creation. (a_runtime_subscription_made_by_one_command_expires_with_no_client) - N2: a direct-HTTP server->client request's refusal POST carries the same Accept as every events POST; the fixture now 406s an answer without it, as the Python SDK would. (over_direct_http_*) - N3: refusals on an events/stream are answered inline, one at a time — a flooding server cannot make the reader spawn without limit. (a_flood_of_server_requests_on_an_events_stream_is_answered_one_at_a_time) - N4: the boot cap keeps the sessions a client used most recently, not the first alphabetically. (the_boot_cap_keeps_the_most_recently_used_sessions) - N5: a forgotten runtime subscription (TTL expiry, permanent refusal on restore, server ended) is forgotten whole: its webhook token and secret leave the snapshot. (a_runtime_subscription_to_a_removed_server_is_forgotten_not_resurrected) - N6: a panicked session holding runtime subscriptions is restarted with per-session backoff, like the events session. (a_session_with_runtime_subscriptions_is_restarted_after_a_panic) - N7: ordinary_limit reads the request id directly instead of serializing the whole message; every streamable-HTTP POST goes through post_bounded (rmcp's post_message is unused) — pinned against rmcp's own client on request headers and response shapes (an_ordinary_request_is_sent_as_rmcps_own_client_would / ..._answered_as_...). - ARCHITECTURE.md: the above, and the over-cap push-event skip from the previous commit. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/ARCHITECTURE.md | 34 ++- .../src/bin/mcp_fixture_events_server.rs | 35 ++- crates/agent/src/serve_ws.rs | 71 +++++-- crates/agent/src/tools/mcp_events/mod.rs | 58 +++-- crates/agent/src/tools/mcp_events/state.rs | 54 +++-- crates/agent/src/tools/mcp_events/wire.rs | 11 +- crates/agent/src/tools/mcp_wire.rs | 200 +++++++++++++++++- crates/agent/tests/mcp_events_nested.rs | 35 +++ crates/agent/tests/mcp_events_receiver.rs | 146 +++++++++++++ 9 files changed, 575 insertions(+), 69 deletions(-) diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 7c17a43..9dd9103 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, and caps every streamable-HTTP message: `HttpClient` (a `StreamableHttpClient` that answers **every** POST itself through `post_bounded` — 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/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. | @@ -4540,12 +4540,16 @@ mcp_events_subscribe (any session) ──► owned by that 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. Runtime subscriptions also **expire**: a session's - last client command is recorded (coarsely, in the events state), and once it is older than + 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. Boot restoration is bounded too: - at most `BEYOND_AI_AGENT_MCP_EVENTS_MAX_RESTORED_SESSIONS` (32) sessions are started; the rest come - back when their clients do. + 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 @@ -4644,8 +4648,11 @@ mcp_events_subscribe (any session) ──► owned by that session 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), so configured triggers - do not stop until the next restart (`the_events_session_is_restarted_after_a_panic`). The only exits that skip it + 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, @@ -4660,14 +4667,19 @@ 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 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 one fails its request or ends its stream, never read whole; a server that keeps resending a giant stream event keeps failing, on backoff) — JWKS documents at 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`). + 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 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 986b9d0..e243c1b 100644 --- a/crates/agent/src/bin/mcp_fixture_events_server.rs +++ b/crates/agent/src/bin/mcp_fixture_events_server.rs @@ -116,6 +116,10 @@ struct State { 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, } @@ -1003,6 +1007,7 @@ async fn control(state: &Shared, method: &str, path: &str, body: &[u8]) -> (u16, "deliveries": st.deliveries, "streams": st.streams.len(), "log": st.log.len(), "requests": st.requests, "sessionless_rejections": st.sessionless_rejections, "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, }), ) } @@ -1382,9 +1387,25 @@ 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. + // 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() { - state.lock().unwrap().client_answers.push(msg.clone()); + 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 @@ -1460,6 +1481,16 @@ async fn handle_http(state: Shared, mut stream: TcpStream) { 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; diff --git a/crates/agent/src/serve_ws.rs b/crates/agent/src/serve_ws.rs index 7485634..e5acf64 100644 --- a/crates/agent/src/serve_ws.rs +++ b/crates/agent/src/serve_ws.rs @@ -1812,38 +1812,65 @@ pub async fn serve_ws( } } - // The events session is what keeps configured subscriptions running; if it ends in a panic it - // is started again — with backoff (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. + // 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 attempt = 0u32; - let mut last = std::time::Instant::now(); + let mut backoff: std::collections::HashMap = + std::collections::HashMap::new(); while let Some(id) = panicked_rx.recv().await { - if Some(&id) != events_id.as_ref() { + 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; } - if last.elapsed() > std::time::Duration::from_secs(600) { - attempt = 0; + 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 << attempt.min(6)) + let delay = std::time::Duration::from_secs(1u64 << entry.0.min(6)) .min(std::time::Duration::from_secs(60)); - attempt += 1; - tokio::time::sleep(delay).await; - last = std::time::Instant::now(); - match supervisor.pin(Some(id.clone()), None).await { - Ok(p) => { - supervisor.unpin(&p.id, p.incarnation); - eprintln!("serve: restarted the MCP Events session `{id}` after a panic"); + 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) + ), } - Err(e) => eprintln!( - "serve: could not restart the MCP Events session `{id}`: {}", - HttpError::from(e) - ), - } + }); } }); } diff --git a/crates/agent/src/tools/mcp_events/mod.rs b/crates/agent/src/tools/mcp_events/mod.rs index 24a71bd..22249ca 100644 --- a/crates/agent/src/tools/mcp_events/mod.rs +++ b/crates/agent/src/tools/mcp_events/mod.rs @@ -576,10 +576,11 @@ impl McpEventsHub { return; }; // Runtime subscriptions whose session has heard from no client for longer than - // their time to live are forgotten, not restored. + // 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.set_runtime(&key, None); + store.forget_sub(&key); } } let runtime: Vec<(String, SubSpec)> = store @@ -763,7 +764,7 @@ pub fn sessions_with_runtime_subscriptions(dir: &std::path::Path, max: usize) -> return Vec::new(); }; let now = now_unix_ms(); - let mut ids = Vec::new(); + let mut ranked: Vec<(i64, String)> = Vec::new(); for entry in entries.flatten() { let name = entry.file_name(); let Some(stem) = name @@ -775,11 +776,16 @@ pub fn sessions_with_runtime_subscriptions(dir: &std::path::Path, max: usize) -> let Some((_, id)) = stem.split_once('_') else { continue; }; - if std::fs::read(entry.path()).is_ok_and(|b| state::snapshot_wants_restore(&b, now)) { - ids.push(id.to_owned()); + 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())); } } - ids.sort(); + // 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 \ @@ -1548,10 +1554,10 @@ impl Hub { 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: it is not - // retried, not restored again, and keeps nothing alive. + // 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.set_runtime(&key, None); + hub.store.forget_sub(&key); break; } } @@ -1607,9 +1613,10 @@ impl Hub { return; } self.refresh_keep_alive(); - // A runtime subscription the server ended is over: it is not subscribed again after a - // restart either (unless re-discovery brings it back, below). - self.store.set_runtime(&spec.key(), None); + // 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; } @@ -2483,4 +2490,31 @@ mod tests { 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"] + ); + } } diff --git a/crates/agent/src/tools/mcp_events/state.rs b/crates/agent/src/tools/mcp_events/state.rs index 19de3ac..9c8bed5 100644 --- a/crates/agent/src/tools/mcp_events/state.rs +++ b/crates/agent/src/tools/mcp_events/state.rs @@ -143,20 +143,29 @@ pub(super) fn runtime_ttl_ms() -> i64 { } /// Whether a snapshot's runtime subscriptions are still worth restoring: its session heard from a -/// client within [`runtime_ttl_ms`] (or the state predates the record, and gets the benefit of the -/// doubt once — the next client command records it). +/// 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_none_or(|t| now_ms.saturating_sub(t) <= runtime_ttl_ms()) + 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.) -pub(super) fn snapshot_wants_restore(bytes: &[u8], now_ms: i64) -> bool { - let Ok(snap) = serde_json::from_slice::(bytes) else { - return false; - }; - runtime_still_wanted(snap.last_client_ms, now_ms) - && snap.subscriptions.values().any(|s| s.runtime.is_some()) +/// 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)] @@ -323,11 +332,17 @@ impl StateStore { 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(); @@ -1352,8 +1367,8 @@ mod tests { "a week is the default" ); assert!( - runtime_still_wanted(None, now), - "older state gets one chance" + !runtime_still_wanted(None, now), + "no record of a client: not trusted" ); let snap = |last: i64| { serde_json::json!({ @@ -1362,11 +1377,17 @@ mod tests { }) .to_string() }; - assert!(snapshot_wants_restore(snap(now - day).as_bytes(), now)); - assert!(!snapshot_wants_restore(snap(now - 8 * day).as_bytes(), now)); + 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!(!snapshot_wants_restore(no_runtime.as_bytes(), now)); + 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. @@ -1377,6 +1398,11 @@ mod tests { 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); diff --git a/crates/agent/src/tools/mcp_events/wire.rs b/crates/agent/src/tools/mcp_events/wire.rs index 224f0ef..7fd3a62 100644 --- a/crates/agent/src/tools/mcp_events/wire.rs +++ b/crates/agent/src/tools/mcp_events/wire.rs @@ -593,11 +593,11 @@ impl Conn { 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, version) = refuse_with.clone(); - tokio::spawn(async move { - refuse_server_request(&http, &url, &headers, &version, &msg).await; - }); + let (http, url, headers, version) = &refuse_with; + refuse_server_request(http, url, headers, version, &msg).await; continue; } if msg.get("id") == Some(&want) { @@ -676,9 +676,12 @@ async fn refuse_server_request( "message": format!("`{method}` is not served during an events request over direct HTTP"), }, }); + // The same `Accept` as every events POST: a streamable-HTTP server (the official Python SDK + // among them) answers a POST without it `406` — and then still waits for the answer. let mut req = http .post(url) .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") .header("MCP-Protocol-Version", protocol_version); for (k, v) in headers { req = req.header(k, v); diff --git a/crates/agent/src/tools/mcp_wire.rs b/crates/agent/src/tools/mcp_wire.rs index 2e66c25..50c02b5 100644 --- a/crates/agent/src/tools/mcp_wire.rs +++ b/crates/agent/src/tools/mcp_wire.rs @@ -52,10 +52,11 @@ impl HttpClient { /// 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 { - let id = matches!(message, ClientJsonRpcMessage::Request(_)) - .then(|| serde_json::to_value(message).ok()) - .flatten() - .and_then(|v| v.get("id").cloned()); + // 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), @@ -673,6 +674,197 @@ 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. + #[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(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 + ); + } + } + } + + /// …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(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(); diff --git a/crates/agent/tests/mcp_events_nested.rs b/crates/agent/tests/mcp_events_nested.rs index 8cee6e4..183c025 100644 --- a/crates/agent/tests/mcp_events_nested.rs +++ b/crates/agent/tests/mcp_events_nested.rs @@ -183,3 +183,38 @@ fn over_direct_http_a_nested_request_during_an_events_poll_is_refused_not_ignore 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_receiver.rs b/crates/agent/tests/mcp_events_receiver.rs index 1a215b5..7d98257 100644 --- a/crates/agent/tests/mcp_events_receiver.rs +++ b/crates/agent/tests/mcp_events_receiver.rs @@ -436,6 +436,19 @@ async fn a_runtime_subscription_to_a_removed_server_is_forgotten_not_resurrected !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() + .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" + ); common::mcp_events_fixture::sigterm_and_wait(&mut second); // And a further restart does not resurrect it. @@ -453,3 +466,136 @@ async fn a_runtime_subscription_to_a_removed_server_is_forgotten_not_resurrected 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)); +} From 8e50de358f42d8941261aae904d50ecc2ca7ce6c Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 22:20:19 -0700 Subject: [PATCH 07/11] agent: a refused server request is answered with the current OAuth token (after #138) #138 made every direct events/* POST carry the server's current shared token and refresh-and-resend once on a 401; the refusal POST for a server->client request still sent the headers built at dial, so after a refresh it spent a 401 and the server kept waiting. It now goes through the same Conn::send (Accept included; no Mcp-Method, since an answer has no method). (a_refused_server_request_is_answered_with_the_current_token_and_refreshed_on_401) Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/ARCHITECTURE.md | 3 +- crates/agent/src/tools/mcp_events/wire.rs | 85 ++++++++++++++++++----- 2 files changed, 69 insertions(+), 19 deletions(-) diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 9dd9103..1977e38 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -4677,7 +4677,8 @@ mcp_events_subscribe (any session) ──► owned by that session 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 on an `events/stream` it is answered inline, one at a time, so a server flooding + 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 diff --git a/crates/agent/src/tools/mcp_events/wire.rs b/crates/agent/src/tools/mcp_events/wire.rs index 7fd3a62..c2fd588 100644 --- a/crates/agent/src/tools/mcp_events/wire.rs +++ b/crates/agent/src/tools/mcp_events/wire.rs @@ -310,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, @@ -322,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); } @@ -487,7 +491,15 @@ impl Conn { return rpc_outcome(msg); } if is_server_request(&msg) { - refuse_server_request(http, url, headers, protocol_version, &msg).await; + refuse_server_request( + http, + url, + headers, + auth.as_deref(), + protocol_version, + &msg, + ) + .await; } } Err(RpcError::local(format!( @@ -588,6 +600,7 @@ impl Conn { http.clone(), url.clone(), headers.clone(), + auth.clone(), protocol_version.clone(), ); let reader = tokio::spawn(async move { @@ -596,8 +609,16 @@ impl Conn { // 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, version) = &refuse_with; - refuse_server_request(http, url, headers, version, &msg).await; + 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) { @@ -664,6 +685,7 @@ 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, ) { @@ -676,19 +698,12 @@ async fn refuse_server_request( "message": format!("`{method}` is not served during an events request over direct HTTP"), }, }); - // The same `Accept` as every events POST: a streamable-HTTP server (the official Python SDK - // among them) answers a POST without it `406` — and then still waits for the answer. - let mut req = http - .post(url) - .header("Content-Type", "application/json") - .header("Accept", "application/json, text/event-stream") - .header("MCP-Protocol-Version", protocol_version); - for (k, v) in headers { - req = req.header(k, v); - } - if let Ok(Err(e)) = - tokio::time::timeout(Duration::from_secs(5), req.body(body.to_string()).send()).await - { + // 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"); } } @@ -1293,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"]); + } } From 15acdaedd0bad12be86a3bf980001e03ab1cf620 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 22:45:01 -0700 Subject: [PATCH 08/11] =?UTF-8?q?agent:=20reconcile=20the=20one=20POST=20p?= =?UTF-8?q?ath=20with=20#140=20=E2=80=94=20drop=20HttpClient's=20oauth=20f?= =?UTF-8?q?lag,=20pin=20the=20401=20divergence?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #140 routed an OAuth server's POSTs through post_bounded (HttpClient { oauth }) so a 401 stays visible; this branch already routes every server's POSTs there, so the flag selected nothing: it is removed, with OAUTH_MAX_MESSAGE_BYTES (the one per-message cap, mcp_stdio::max_message_bytes, applies). The rmcp-parity test now pins the one deliberate divergence: a 401 with no challenge and a JSON-RPC error body is AuthRequired here, an error response from rmcp's own client. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/ARCHITECTURE.md | 7 +++--- crates/agent/src/tools/mcp_wire.rs | 36 +++++++++++++++++++++++++++--- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 1977e38..5bd6718 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -4353,9 +4353,10 @@ 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 — the one deliberate divergence +from rmcp's own client, 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 diff --git a/crates/agent/src/tools/mcp_wire.rs b/crates/agent/src/tools/mcp_wire.rs index 50c02b5..f583347 100644 --- a/crates/agent/src/tools/mcp_wire.rs +++ b/crates/agent/src/tools/mcp_wire.rs @@ -677,7 +677,8 @@ 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. + /// accepted, a JSON answer (with its `Mcp-Session-Id`), and an SSE stream. One divergence is + /// deliberate, and pinned last: a 401 with no challenge is still `AuthRequired` here. #[tokio::test] async fn an_ordinary_request_is_answered_as_rmcps_own_client_would() { use futures::StreamExt as _; @@ -713,7 +714,7 @@ mod tests { (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(reqwest::Client::new()) + let ours = HttpClient::new(reqwest::Client::new()) .post_message( canned(raw, false).await.into(), call(), @@ -749,6 +750,35 @@ mod tests { ); } } + // 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 @@ -826,7 +856,7 @@ mod tests { ]) }; let (url, ours) = capture().await; - HttpClient(reqwest::Client::new()) + HttpClient::new(reqwest::Client::new()) .post_message( url.into(), message(), From 34fa881c69ef3ecc150b2e18274c3a6655e17ecf Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 23:24:06 -0700 Subject: [PATCH 09/11] agent tests: widen the steered-compaction test's window so full-suite load cannot close it Its second event must be polled, coalesced and steered within one stalled turn; 2.5 s was missed once under full-suite load on the host. 6 s per turn, and a 90 s bound on the run's response. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/tests/mcp_events_durability.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/crates/agent/tests/mcp_events_durability.rs b/crates/agent/tests/mcp_events_durability.rs index 4d11f40..9e3d833 100644 --- a/crates/agent/tests/mcp_events_durability.rs +++ b/crates/agent/tests/mcp_events_durability.rs @@ -392,9 +392,12 @@ fn a_steered_batch_summarized_away_by_a_mid_run_compaction_is_not_injected_again let mut stdin = child.stdin.take().unwrap(); let mut frames = Frames::new(child.stdout.take().unwrap(), 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(2500) }), + 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" @@ -419,7 +422,10 @@ fn a_steered_batch_summarized_away_by_a_mid_run_compaction_is_not_injected_again &fixture, json!({ "event_id": "compacted-2", "data": { "summary": "second, kept" } }), ); - let done = frames.response("long"); + // 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 From 0ddecd3ec4486d64a0b50c105738d71fa5fb632c Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 23:49:53 -0700 Subject: [PATCH 10/11] =?UTF-8?q?agent:=20reconcile=20the=20one=20POST=20p?= =?UTF-8?q?ath=20with=20#142=20=E2=80=94=20every=20server,=20two=20pinned?= =?UTF-8?q?=20divergences=20from=20rmcp?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #142's routing docs and tests described post_bounded as the OAuth servers' path, chosen by the removed oauth flag; it answers every server's POSTs. ARCHITECTURE.md now names both deliberate differences from rmcp's client (any 401 is AuthRequired; a non-JSON-RPC success answering a request is an error), each with its pinning test. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/ARCHITECTURE.md | 8 +++----- crates/agent/src/tools/mcp_wire.rs | 6 ++++-- crates/agent/tests/mcp_oauth_routing.rs | 8 ++++---- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 5bd6718..8031a60 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -4354,21 +4354,19 @@ carries no `auth_header` of its own); direct `events/*` POSTs (`mcp_events::wire 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 (every POST, for every server, is answered by -`mcp_wire::HttpClient::post_bounded`, which decides from the status — the one deliberate divergence -from rmcp's own client, which reads a 401 carrying a JSON-RPC error body as an ordinary error +`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 diff --git a/crates/agent/src/tools/mcp_wire.rs b/crates/agent/src/tools/mcp_wire.rs index f583347..dccfcac 100644 --- a/crates/agent/src/tools/mcp_wire.rs +++ b/crates/agent/src/tools/mcp_wire.rs @@ -677,8 +677,10 @@ 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. One divergence is - /// deliberate, and pinned last: a 401 with no challenge is still `AuthRequired` here. + /// 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 _; diff --git a/crates/agent/tests/mcp_oauth_routing.rs b/crates/agent/tests/mcp_oauth_routing.rs index dc3eaf4..66b067b 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); From b50dd9bad878c66ade2d9b0e991139c9a1c882a4 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Tue, 6 Oct 2026 23:56:55 -0700 Subject: [PATCH 11/11] agent tests: read the MCP Events suites' stdout through #141's guarded Frames #141's ChildGuard takes the stdout pipe at spawn and its lint refuses any other route to it; the follow-up tests built Frames from child.stdout.take(). They now pass the guard, like every suite. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/tests/mcp_events_durability.rs | 2 +- crates/agent/tests/mcp_events_nested.rs | 2 +- crates/agent/tests/mcp_message_cap.rs | 6 +++--- crates/agent/tests/mcp_stdio_transport.rs | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/agent/tests/mcp_events_durability.rs b/crates/agent/tests/mcp_events_durability.rs index 9e3d833..aff292d 100644 --- a/crates/agent/tests/mcp_events_durability.rs +++ b/crates/agent/tests/mcp_events_durability.rs @@ -390,7 +390,7 @@ fn a_steered_batch_summarized_away_by_a_mid_run_compaction_is_not_injected_again ]); let mut child = cmd.spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut frames = Frames::new(child.stdout.take().unwrap(), None); + 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 diff --git a/crates/agent/tests/mcp_events_nested.rs b/crates/agent/tests/mcp_events_nested.rs index 183c025..6693a0c 100644 --- a/crates/agent/tests/mcp_events_nested.rs +++ b/crates/agent/tests/mcp_events_nested.rs @@ -36,7 +36,7 @@ fn nested_during(during: &str, delivery: &str) { .stderr(Stdio::null()); let mut child = cmd.spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut frames = Frames::new(child.stdout.take().unwrap(), None); + let mut frames = Frames::new(&mut child, None); let control = wait_control_file(&control_file); let ask = frames.wait( Duration::from_secs(30), diff --git a/crates/agent/tests/mcp_message_cap.rs b/crates/agent/tests/mcp_message_cap.rs index 6f09199..fbabe28 100644 --- a/crates/agent/tests/mcp_message_cap.rs +++ b/crates/agent/tests/mcp_message_cap.rs @@ -43,7 +43,7 @@ fn run_blobs(home: &std::path::Path, server: Value) -> Vec { .stderr(Stdio::null()); let mut child = cmd.spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut frames = Frames::new(child.stdout.take().unwrap(), None); + let mut frames = Frames::new(&mut child, None); for (id, msg) in [ ("p1", "make a blob of 2000000"), ("p2", "make a blob of 10"), @@ -128,7 +128,7 @@ fn events_serve( .stderr(Stdio::null()); let mut child = cmd.spawn_guarded(); let stdin = child.stdin.take().unwrap(); - let frames = Frames::new(child.stdout.take().unwrap(), None); + let frames = Frames::new(&mut child, None); (child, stdin, frames, control, home, fx) } @@ -193,7 +193,7 @@ fn an_over_cap_push_event_is_skipped_and_reported(stdio: bool) { .stderr(Stdio::null()); let mut child = cmd.spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut frames = Frames::new(child.stdout.take().unwrap(), None); + 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), diff --git a/crates/agent/tests/mcp_stdio_transport.rs b/crates/agent/tests/mcp_stdio_transport.rs index 3bfdfe4..2069c5d 100644 --- a/crates/agent/tests/mcp_stdio_transport.rs +++ b/crates/agent/tests/mcp_stdio_transport.rs @@ -351,7 +351,7 @@ fn a_panic_inside_serve_exits_promptly_and_still_sweeps() { .env("BEYOND_AI_AGENT_TEST_PANICS", "1"); let mut child = cmd.spawn_guarded(); let stdin = child.stdin.take().unwrap(); - let frames = Frames::new(child.stdout.take().unwrap(), None); + let frames = Frames::new(&mut child, None); (child, stdin, frames) }; let mut child = _child_guard;