agent: every MCP Events start path waits for the session's gates; restore window covers runtime-only daemons

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

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.

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<String> = 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);

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<tokio::task::JoinSet<()>>,
     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::<u64>().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::<u64>().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;

crates/agent/src/tools/mcp_events/webhook.rs

@@ -105,35 +105,55 @@ fn reserved() -> std::sync::MutexGuard<'static, HashMap<String, std::time::Insta
         .unwrap_or_else(std::sync::PoisonError::into_inner)
 }
 
-/// Until when *every* unknown token is answered `503` (retry) rather than `410` (stop): from a
-/// daemon's start until its MCP Events session has read its persisted state and reserved the tokens
-/// in it ([`restore_pending`], [`restored`]). 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_UNTIL: Mutex<Option<std::time::Instant>> = 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<Option<Restoring>> = Mutex::new(None);
 
-fn restoring_until() -> std::sync::MutexGuard<'static, Option<std::time::Instant>> {
-    RESTORING_UNTIL
+struct Restoring {
+    until: std::time::Instant,
+    /// The sessions still to reserve their tokens.
+    sessions: std::collections::HashSet<String>,
+}
+
+fn restoring_state() -> std::sync::MutexGuard<'static, Option<Restoring>> {
+    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<Item = String>) {
+    let sessions: std::collections::HashSet<String> = 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,

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:#?}"
+    );
+}