diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 720b0e42..421180c6 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -4589,8 +4589,15 @@ mcp_events_subscribe (any session) ──► owned by that session 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`. + the draft treats as "stop"). Before any state is read, a daemon with a callback URL answers every + unknown token `503` from its start until each session it starts to restore callbacks (the events + session, and every session holding runtime subscriptions) has read its state and held its tokens + (`webhook::restore_pending` / `restored`, also bounded at 10 min); after that an unknown token is + `410` at once. A runtime subscription the server ended is not restored, so its token is not held + and its old callback answers `410`. No subscription makes a request — configured, restored after a + restart or a panic, or resubscribing — until the session can take the server's questions (its + elicitation and sampling gates are installed: `McpEventsHub::start`), since a server may ask during + the very first `events/poll`. 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 diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index bb00b84b..973c4ad2 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -3204,7 +3204,7 @@ pub(crate) async fn serve_session( ); mcp_host.sampling.install(sampling_gate); if let Some(hub) = &mcp_events { - hub.start_configured(); + hub.start(); } // Where an MCP App view's `mcp_app_open`/`mcp_app_result` frames go: this session's fanout, // held weakly like the elicitation gate's, so a view-bearing tool cannot keep it alive. diff --git a/crates/agent/src/serve_ws.rs b/crates/agent/src/serve_ws.rs index b092e190..bb54e75e 100644 --- a/crates/agent/src/serve_ws.rs +++ b/crates/agent/src/serve_ws.rs @@ -1855,10 +1855,36 @@ pub async fn serve_ws( let mut shutdown = crate::serve::ShutdownSignal::new()?; let events_session_id = events_session.clone(); - if events_session.is_some() && cfg_has_callback { - // Until the events session has read its state, a delivery for a callback it persisted - // must be told to retry, not to stop (see `mcp_events::restore_pending`). - crate::tools::mcp_events::restore_pending(); + // Sessions holding runtime MCP Events subscriptions (`mcp_events_subscribe`), started below + // after the events session. Listed now so the restore window covers them too. + let runtime_sessions: Vec = match supervisor.session_dir.clone() { + Some(dir) => { + let events_id = events_session_id.clone(); + tokio::task::spawn_blocking(move || { + // Bounded: a pile of old sessions cannot all be woken at every boot. + crate::tools::mcp_events::sessions_with_runtime_subscriptions( + std::path::Path::new(&dir), + crate::tools::mcp_events::max_restored_sessions(), + ) + }) + .await + .unwrap_or_default() + .into_iter() + .filter(|id| Some(id.as_str()) != events_id.as_deref()) + .collect() + } + None => Vec::new(), + }; + if cfg_has_callback { + // Until every session started to restore persisted callbacks has read its state, a + // delivery for one of them must be told to retry, not to stop (see + // `mcp_events::restore_pending`). Configured subscriptions and runtime ones alike. + crate::tools::mcp_events::restore_pending( + events_session + .iter() + .cloned() + .chain(runtime_sessions.iter().cloned()), + ); } if let Some(id) = events_session { // Pinned and at once unpinned: the session starts exactly as if a client had connected and @@ -1941,21 +1967,8 @@ 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 || { - // Bounded: a pile of old sessions cannot all be woken at every boot. - crate::tools::mcp_events::sessions_with_runtime_subscriptions( - std::path::Path::new(&dir), - crate::tools::mcp_events::max_restored_sessions(), - ) - }) - .await - .unwrap_or_default(); - for id in ids - .into_iter() - .filter(|id| Some(id.as_str()) != events_id.as_deref()) - { + { + for id in runtime_sessions { match supervisor.pin(Some(id.clone()), None).await { Ok(p) => { supervisor.unpin(&p.id, p.incarnation); diff --git a/crates/agent/src/tools/mcp_events/mod.rs b/crates/agent/src/tools/mcp_events/mod.rs index 24ed1935..96a04799 100644 --- a/crates/agent/src/tools/mcp_events/mod.rs +++ b/crates/agent/src/tools/mcp_events/mod.rs @@ -496,6 +496,12 @@ struct Hub { /// `mcp_events_*` commands in flight — spawned so none ever blocks the session's command loop. command_tasks: Mutex>, owns_configured: bool, + /// Cancelled once the session can take a server's questions (its elicitation and sampling gates + /// are installed; see [`McpEventsHub::start`]). Every subscription waits for it before its first + /// request — configured ones, runtime ones restored after a restart or a panic, resubscriptions — + /// since a server may ask during the very first `events/poll`, and asked before then the + /// question would be declined as having no client. + questions_ready: CancellationToken, /// The last few over-cap messages a stdio connection reported to every push stream at once /// (`$dropped_id`): each is told to the model once, however many of this session's /// subscriptions it reached. @@ -557,6 +563,7 @@ impl McpEventsHub { store, command_tasks: Mutex::new(tokio::task::JoinSet::new()), owns_configured: cfg.owns_configured, + questions_ready: CancellationToken::new(), dropped_seen: Mutex::new(std::collections::VecDeque::new()), }); let weak_hub = Arc::downgrade(&hub); @@ -569,22 +576,14 @@ impl McpEventsHub { // 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.) + // the runtime ones, which make no request before `start` (see `Hub::questions_ready`). { let weak_hub = Arc::downgrade(&hub); tokio::spawn(async move { let Some(store) = weak_hub.upgrade().map(|h| h.store.clone()) else { return; }; - // Debug builds: a test can hold the state read off, to show what a delivery that - // arrives before it is told (retry, not stop). - #[cfg(debug_assertions)] - if let Some(ms) = std::env::var("BEYOND_AI_AGENT_TEST_SLOW_EVENTS_RESTORE_MS") - .ok() - .and_then(|v| v.parse::().ok()) - { - tokio::time::sleep(Duration::from_millis(ms)).await; - } + simulated_slow_restore().await; store.loaded().await; let Some(hub) = weak_hub.upgrade() else { return; @@ -611,11 +610,9 @@ impl McpEventsHub { } } } - // The daemon held every unknown token at `503` until this session's tokens were - // reserved (`webhook::restore_pending`); they are now. - if hub.owns_configured { - webhook::restored(); - } + // The daemon may be holding every unknown token at `503` until this session's + // tokens are reserved (`webhook::restore_pending`); they are now. + webhook::restored(&hub.session_id); for (key, spec) in runtime { hub.restoring .lock() @@ -646,7 +643,11 @@ impl McpEventsHub { /// a failed subscribe is retried with capped backoff for as long as the session runs, and so is /// a subscription the server later ends. Failures surface as `mcp_event_status` frames and on /// stderr, and stay visible in `mcp_events_list`. - pub fn start_configured(&self) { + /// The session can take a server's questions now: let every subscription make its first + /// request (see `Hub::questions_ready`), and start the configured ones this session owns. Call + /// once the session's elicitation and sampling gates are installed. + pub fn start(&self) { + self.hub.questions_ready.cancel(); if !self.hub.owns_configured { return; } @@ -774,6 +775,21 @@ impl McpEventsHub { } } +/// Debug builds: a test can hold a session's state read off, to show what a delivery that arrives +/// before it is told (retry, not stop), and that a restored subscription still asks nothing early. +#[cfg(debug_assertions)] +async fn simulated_slow_restore() { + if let Some(ms) = std::env::var("BEYOND_AI_AGENT_TEST_SLOW_EVENTS_RESTORE_MS") + .ok() + .and_then(|v| v.parse::().ok()) + { + tokio::time::sleep(Duration::from_millis(ms)).await; + } +} + +#[cfg(not(debug_assertions))] +async fn simulated_slow_restore() {} + /// 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 @@ -1523,8 +1539,14 @@ impl Hub { hub.refresh_keep_alive(); let weak = Arc::downgrade(hub); let shutdown = hub.shutdown.clone(); + let ready = hub.questions_ready.clone(); let resumed = !delay.is_zero() && !restoring; tokio::spawn(async move { + // Not a request before the session can take the server's questions. + tokio::select! { + () = ready.cancelled() => {} + () = shutdown.cancelled() => return, + } let mut delay = delay; let mut failures = 0u32; let mut forbidden_streak = 0u32; diff --git a/crates/agent/src/tools/mcp_events/webhook.rs b/crates/agent/src/tools/mcp_events/webhook.rs index 8cdbc252..f3a5a2b3 100644 --- a/crates/agent/src/tools/mcp_events/webhook.rs +++ b/crates/agent/src/tools/mcp_events/webhook.rs @@ -105,35 +105,55 @@ fn reserved() -> std::sync::MutexGuard<'static, HashMap> = Mutex::new(None); +/// While set, *every* unknown token is answered `503` (retry) rather than `410` (stop): from a +/// daemon's start until each session it starts to restore persisted callbacks (the events session, +/// and every session holding runtime subscriptions) has read its state and reserved the tokens in it +/// ([`restore_pending`], [`restored`]), and for at most [`RESERVE_FOR`]. Without it, a delivery the +/// server retried across a restart that lands before the state is read would be told to stop for +/// good. +static RESTORING: Mutex> = Mutex::new(None); -fn restoring_until() -> std::sync::MutexGuard<'static, Option> { - RESTORING_UNTIL +struct Restoring { + until: std::time::Instant, + /// The sessions still to reserve their tokens. + sessions: std::collections::HashSet, +} + +fn restoring_state() -> std::sync::MutexGuard<'static, Option> { + RESTORING .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) } -/// A daemon whose events session will restore persisted callbacks is starting: answer unknown -/// tokens `503` until it has ([`restored`]), or for at most [`RESERVE_FOR`]. -pub fn restore_pending() { - *restoring_until() = Some(std::time::Instant::now() + RESERVE_FOR); +/// A daemon is starting `sessions` to restore their persisted callbacks: answer unknown tokens +/// `503` until each has reserved its own ([`restored`]), or for at most [`RESERVE_FOR`]. Nothing to +/// restore, nothing held. +pub fn restore_pending(sessions: impl IntoIterator) { + let sessions: std::collections::HashSet = sessions.into_iter().collect(); + *restoring_state() = (!sessions.is_empty()).then(|| Restoring { + until: std::time::Instant::now() + RESERVE_FOR, + sessions, + }); } -/// The events session has reserved every token it persisted: unknown tokens are `410` again. -pub(super) fn restored() { - *restoring_until() = None; +/// `session` has reserved every token it persisted. Once every session restore was waiting on has, +/// unknown tokens are `410` again. +pub(super) fn restored(session: &str) { + let mut state = restoring_state(); + if let Some(r) = state.as_mut() { + r.sessions.remove(session); + if r.sessions.is_empty() { + *state = None; + } + } } fn restoring() -> bool { - let mut until = restoring_until(); - match *until { - Some(t) if t > std::time::Instant::now() => true, + let mut state = restoring_state(); + match state.as_ref() { + Some(r) if r.until > std::time::Instant::now() => true, Some(_) => { - *until = None; + *state = None; false } None => false, diff --git a/crates/agent/tests/mcp_events_nested.rs b/crates/agent/tests/mcp_events_nested.rs index 7ddf3afa..fc1391d4 100644 --- a/crates/agent/tests/mcp_events_nested.rs +++ b/crates/agent/tests/mcp_events_nested.rs @@ -235,3 +235,88 @@ fn a_flood_of_server_requests_on_an_events_stream_is_answered_one_at_a_time() { "answered one at a time, not a task per request" ); } + +/// A runtime subscription restored when its session starts again (a restart here; a panic restart +/// takes the same path) waits for the session to take questions too, like a configured one: its +/// first `events/poll` raises a nested elicitation while a debug seam holds the session's gates off, +/// and the question still reaches the session's client rather than being declined. +#[test] +fn a_restored_runtime_subscription_asks_only_once_the_session_takes_questions() { + let home = tempfile::tempdir().unwrap(); + let control_file = home.path().join("control"); + let session = home.path().join("s.jsonl"); + let settings = |nested: bool| { + let mut env = json!({ "MCP_FIXTURE_HEARTBEAT_MS": "200" }); + if nested { + env["MCP_FIXTURE_NESTED_DURING"] = json!("poll"); + } + write_settings( + home.path(), + json!([stdio_server("tickets", &control_file, env, json!([]))]), + ); + }; + let spawn = |extra: &[(&str, &str)]| { + let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); + let mut cmd = serve_cmd(BIN, &base, &session.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") + .envs(extra.iter().copied()) + .stderr(Stdio::null()); + cmd.spawn_guarded() + }; + + // First run: a runtime subscription, polled, with no question asked. + settings(false); + let mut child = spawn(&[]); + let mut stdin = child.stdin.take().unwrap(); + let mut frames = Frames::new(&mut child, None); + send( + &mut stdin, + json!({ "type": "mcp_events_subscribe", "id": "s", "server": "tickets", + "name": "ticket.updated", "delivery": "poll", "action": "notify" }), + ); + assert_eq!(frames.response("s")["success"], true); + let control = wait_control_file(&control_file); + eventually( + Duration::from_secs(20), + "the runtime subscription polled", + || { + let polls = state(&control)["methods"] + .as_array()? + .iter() + .filter(|m| *m == "events/poll") + .count(); + (polls > 0).then_some(()) + }, + ); + drop(stdin); + child.wait().unwrap(); + let _ = std::fs::remove_file(&control_file); + + // Second run: restored from the state, and its server asks during the first poll, while the + // session is still installing the gates that would take the question. + settings(true); + let mut child = spawn(&[("BEYOND_AI_AGENT_TEST_SLOW_GATE_INSTALL_MS", "1500")]); + let mut stdin = child.stdin.take().unwrap(); + let mut frames = Frames::new(&mut child, None); + let control = wait_control_file(&control_file); + let ask = frames.wait( + Duration::from_secs(30), + "the restored subscription's nested elicitation", + |f| f["type"] == "elicitation_request", + ); + 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 a = state(&control)["nested_answers"].as_array()?.clone(); + (!a.is_empty()).then_some(a) + }); + assert_eq!( + answers[0]["answer"]["result"]["action"], "accept", + "the session's client answered, not a decline: {answers:#?}" + ); +} diff --git a/crates/agent/tests/mcp_events_receiver.rs b/crates/agent/tests/mcp_events_receiver.rs index ae5203c0..0eaa4c18 100644 --- a/crates/agent/tests/mcp_events_receiver.rs +++ b/crates/agent/tests/mcp_events_receiver.rs @@ -296,6 +296,99 @@ async fn a_retry_that_beats_the_resubscribe_after_a_restart_is_told_to_retry_not eventually(Duration::from_secs(20), "the model run", || { (runs_for_event(&bodies, "retried early") >= 1).then_some(()) }); + // Restored: the window is closed, and a token nothing holds is told to stop, at once. + assert_eq!(dead_token_status(held.port()), 410); +} + +/// A delivery to a callback token no subscription holds: the status the daemon answers. +fn dead_token_status(port: u16) -> u16 { + raw_request( + port, + "POST /_beyond/mcp-events/0000000000000000000000000000dead HTTP/1.1\r\nHost: x\r\n\ + Content-Length: 2\r\nConnection: close\r\n\r\n", + b"{}", + Duration::from_secs(10), + ) +} + +/// A daemon whose only subscriptions are runtime ones (`mcp_events_subscribe`) holds the restart +/// window too: a retried delivery that lands before the session holding the subscription has read +/// its state (held off by a debug seam) gets `503`, not `410`; once restored, the retry lands, and +/// a token nothing holds is `410` again at once. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn with_only_runtime_subscriptions_an_early_retry_is_told_to_retry_not_to_stop() { + 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("runtime-owner")).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:#}"); + eventually( + Duration::from_secs(10), + "the runtime spec and its callback 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| { + let text = std::fs::read_to_string(p).unwrap_or_default(); + text.contains("\"runtime\"") && text.contains("\"token\"") + }) + .then_some(()) + }, + ); + drop(ws); + first.kill().unwrap(); + let _ = first.wait(); + let down = held.down(); + let r = emit( + &fixture, + json!({ "event_id": "early-1", "data": { "summary": "retried early" } }), + ); + assert_ne!(r["deliveries"][0]["status"], 200, "{r:#}"); + drop(down); + let _second = daemon( + home.path(), + &base, + &held, + &[("BEYOND_AI_AGENT_TEST_SLOW_EVENTS_RESTORE_MS", "1500")], + ); + let r = control( + &fixture, + "POST", + "/control/redeliver", + Some(&json!({ "event_id": "early-1" })), + ); + assert_eq!( + r["deliveries"][0]["status"], 503, + "resuming: retry, not stop: {r:#}" + ); + eventually(Duration::from_secs(30), "the retry to land", || { + let r = control( + &fixture, + "POST", + "/control/redeliver", + Some(&json!({ "event_id": "early-1" })), + ); + (r["deliveries"][0]["status"] == 200).then_some(()) + }); + assert_eq!(dead_token_status(port), 410); } /// A runtime subscription (`mcp_events_subscribe`) survives a restart like a configured one: the diff --git a/crates/agent/tests/release_seams.rs b/crates/agent/tests/release_seams.rs index 4fe9b8f1..2ef3c6f2 100644 --- a/crates/agent/tests/release_seams.rs +++ b/crates/agent/tests/release_seams.rs @@ -18,10 +18,24 @@ fn ungated_seam_names(path: &str, source: &str) -> Vec { if !code.contains("BEYOND_AI_AGENT_TEST_") { continue; } - // A gated match arm: the attribute sits right above the arm. - let arm_gated = lines[i.saturating_sub(3)..i] - .iter() - .any(|l| l.trim() == "#[cfg(debug_assertions)]"); + // A gated match arm or statement: the attribute sits directly above the start of the arm or + // statement this line belongs to (walking back over its own continuation lines), not merely + // somewhere close above — an attribute on the *previous* statement gates that one only. + let mut start = i; + while start > 0 { + let prev = lines[start - 1].split("//").next().unwrap_or("").trim_end(); + let ends_item = prev.is_empty() + || prev.ends_with(';') + || prev.ends_with('{') + || prev.ends_with('}') + || prev.ends_with(',') + || prev.trim_start().starts_with("#["); + if ends_item { + break; + } + start -= 1; + } + let arm_gated = start > 0 && lines[start - 1].trim() == "#[cfg(debug_assertions)]"; // A gated function: the nearest enclosing `fn` (less indented than this line) carries it. let indent = line.len() - line.trim_start().len(); let fn_gated = (0..i) @@ -112,4 +126,23 @@ fn dispatch(cmd: &str) { } "; assert!(ungated_seam_names("f.rs", gated_arm).is_empty()); + // A gated statement, spanning lines, inside an ungated function. + let gated_statement = "\ +async fn body() { + #[cfg(debug_assertions)] + if let Some(ms) = std::env::var( + \"BEYOND_AI_AGENT_TEST_SLOW_X_MS\", + ).ok() {} +} +"; + assert!(ungated_seam_names("f.rs", gated_statement).is_empty()); + // The attribute gates the statement under it, not the call site after that. + let gated_neighbour = "\ +async fn body() { + #[cfg(debug_assertions)] + let x = 1; + delay(\"BEYOND_AI_AGENT_TEST_SLOW_X_MS\").await; +} +"; + assert_eq!(ungated_seam_names("f.rs", gated_neighbour).len(), 1); }