agent: load-only flakes fixed at the cause (two real races); #137 follow-ups

crates/agent/ARCHITECTURE.md

@@ -3685,9 +3685,12 @@ the task, not the connection or the process, as the unit of work.
     session's directory), made by the session's first journal write, so a session that never
     journals a task gets no extra file. Every first writer (threads, other processes) agrees on one
     key: it is made under an exclusive lock on the sidecar and re-read once the lock is held, and
-    every journal write adopts the key on disk. A key under the earlier `with_extension` name
-    (`<session>.mcp-task-journal.json`) is read, and carried over by the next journal write. It is one of the session's sidecars, so it moves,
-    trashes and restores with it. On replay an entry without a valid `mac` is ignored: a line appended to the
+    every journal write adopts the key on disk. The lock is an OFD lock on the sidecar itself
+    (`file_lock::Target::Itself`), so another descriptor's open and close cannot release it, NFS
+    included; the holder still re-reads and writes through its locked descriptor, so "re-read, then
+    write" is one step. A key under the earlier `with_extension` name
+    (`<session>.mcp-task-journal.json`) is read, and carried over by the next journal write. Both
+    names are session sidecars, so the key moves, trashes and restores with its session. On replay an entry without a valid `mac` is ignored: a line appended to the
     `.jsonl` by a model with `write`/`edit` neither causes a poll nor is delivered as a result.
   - A session without a key (from before per-session keys, or that has not journaled yet) is read
     as before: its entries are accepted, and opening it writes nothing. Its next journal write makes

crates/agent/src/serve.rs

@@ -1296,6 +1296,21 @@ fn simulated_slow_open() {
 #[cfg(not(debug_assertions))]
 fn simulated_slow_open() {}
 
+/// Latency before a session installs its elicitation and sampling gates, simulated (debug builds
+/// only, like [`simulated_slow_open`]): holds open the window in which a server could ask a question
+/// before the session can take it, so a test can show nothing is asked in it.
+#[cfg(debug_assertions)]
+async fn simulated_slow_gate_install() {
+    if let Ok(ms) = std::env::var("BEYOND_AI_AGENT_TEST_SLOW_GATE_INSTALL_MS")
+        && let Ok(ms) = ms.parse::<u64>()
+    {
+        tokio::time::sleep(std::time::Duration::from_millis(ms)).await;
+    }
+}
+
+#[cfg(not(debug_assertions))]
+async fn simulated_slow_gate_install() {}
+
 impl Persistence {
     /// Open persistence and select the active session: in repo mode per [`SessionSelect`], in file mode
     /// by opening the named file (or creating it).
@@ -2868,7 +2883,9 @@ pub(crate) async fn serve_session(
             },
         );
         input_rx = rx;
-        hub.start_configured();
+        // Started below, once this session's elicitation and sampling gates are installed: a
+        // server may raise a nested request during the very first `events/poll`, and asked before
+        // then it would be declined as having no client.
         mcp_events = Some(hub);
     }
     // SEP-2640 skills from every connected server, seen through this session's MCP gate.
@@ -3139,6 +3156,7 @@ pub(crate) async fn serve_session(
         }};
     }
 
+    simulated_slow_gate_install().await;
     let (elicit_gate, pending_elicitations) = ServeElicitationGate::new(
         out_conn.clone(),
         cfg.approval_timeout,
@@ -3151,6 +3169,9 @@ pub(crate) async fn serve_session(
         persistence.session_id(),
     );
     mcp_host.sampling.install(sampling_gate);
+    if let Some(hub) = &mcp_events {
+        hub.start_configured();
+    }
     // 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.
     let app_sink: crate::tools::mcp_apps::AppSink = {
@@ -4556,6 +4577,14 @@ pub(crate) async fn serve_session(
                 // turn's model call (see `crate::mcp_resume`), by the first attempt only.
                 // A resumer that finished with a task's server unreachable left that task pending:
                 // a fresh one tries it again now (its finished results carry over).
+                //
+                // Journal writes still queued in `mcp_task_rx` are not drained here, and nothing
+                // below depends on them. A finished task's result: a restarted resumer is handed the
+                // old one's results (`mcp_carried`), so it never re-polls. A placeholder mark queued
+                // as this prompt arrives: its placeholder is not on the path yet (this prompt's
+                // `resume_into_turn` splices it, adding the resumer's own placeholder ids to the
+                // journal's), and once a run has put one there, the run-end drain journals its mark
+                // before the next prompt reads the journal.
                 let mut mcp_carried = Vec::new();
                 if mcp_resumer.session_id() != persistence.session_id() {
                     mcp_resumer = start_mcp_resumer!();

crates/agent/src/serve_ws.rs

@@ -1824,6 +1824,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 cfg_has_callback = cfg.mcp_events_callback_url.is_some();
     let (panicked_tx, mut panicked_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
     let supervisor = Arc::new(Supervisor {
         table: Arc::default(),
@@ -1839,6 +1840,11 @@ 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();
+    }
     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.

crates/agent/src/session_store.rs

@@ -877,7 +877,11 @@ fn make_journal_key(
             }
         }
     };
-    if let Some(theirs) = read_journal_key(&path) {
+    // Re-read through the locked descriptor: a second open and close would drop the lock on NFS.
+    let mut on_disk = Vec::new();
+    if lock.file().read_to_end(&mut on_disk).is_ok()
+        && let Some(theirs) = JournalAuth::from_sidecar(&on_disk)
+    {
         return Some(theirs);
     }
     let (auth, bytes) = match log.legacy_journal_key_path().and_then(|p| {
@@ -2889,7 +2893,12 @@ fn move_sibling_memory(session_jsonl: &Path, dst_dir: &Path) {
     let sidecars = SESSION_SIDECARS
         .iter()
         .map(|ext| session_jsonl.with_extension(ext))
-        .chain([journal_key_beside(session_jsonl)]);
+        .chain([
+            journal_key_beside(session_jsonl),
+            // A key under the old name not yet carried over: without it a restored session would
+            // read as keyless.
+            session_jsonl.with_extension(crate::mcp_resume::JournalAuth::SIDECAR),
+        ]);
     for sidecar in sidecars {
         if let Some(name) = sidecar.file_name()
             && sidecar.is_file()
@@ -2938,6 +2947,7 @@ fn remove_sibling_memory(session_jsonl: &Path) {
         let _ = fs::remove_file(session_jsonl.with_extension(ext));
     }
     let _ = fs::remove_file(journal_key_beside(session_jsonl));
+    let _ = fs::remove_file(session_jsonl.with_extension(crate::mcp_resume::JournalAuth::SIDECAR));
 }
 
 /// A directory of session files. `Clone` is just a `PathBuf` copy — cheap, and lets a caller move an
@@ -6761,6 +6771,44 @@ mod tests {
         assert_eq!(store.active_journal(TASK_ENTRY_KIND).len(), 1);
     }
 
+    /// A key still under the old name (not yet carried over) is trashed and restored with its
+    /// session too: left behind, the restored session would read as keyless.
+    #[test]
+    fn a_trashed_session_takes_an_old_name_key_along() {
+        use crate::mcp_resume::{JournalAuth, TASK_ENTRY_KIND};
+        let dir = tmpdir();
+        let repo = SessionRepo::open(dir.path()).unwrap();
+        let mut store = repo.create(SessionMeta::new("/w", "m")).unwrap();
+        store
+            .append_journal(TASK_ENTRY_KIND, serde_json::json!({ "taskId": "t" }))
+            .unwrap();
+        let id = store.meta().id.clone();
+        let path = store.path().to_path_buf();
+        drop(store);
+        let legacy = path.with_extension(JournalAuth::SIDECAR);
+        fs::rename(journal_key_beside(&path), &legacy).unwrap();
+
+        repo.delete(&id).unwrap();
+        assert!(
+            !legacy.exists(),
+            "the old-name key left the session directory"
+        );
+        assert!(
+            dir.path()
+                .join(".trash")
+                .join(legacy.file_name().unwrap())
+                .is_file()
+        );
+        assert!(repo.restore_session(&id).unwrap());
+        assert!(legacy.is_file(), "the old-name key came back");
+        let (store, _) = SessionStore::open(path).unwrap();
+        assert_eq!(store.active_journal(TASK_ENTRY_KIND).len(), 1);
+        assert!(
+            matches!(store.journal, JournalAuth::Keyed { .. }),
+            "the restored session is keyed, not keyless"
+        );
+    }
+
     /// Sessions whose names differ only after their last dot each get their own key: the key's
     /// name is a suffix on the whole file name, not a replaced extension.
     #[test]

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

@@ -54,7 +54,7 @@ use crate::settings::{McpEventAction, McpEventDelivery, McpEventSubscription};
 use crate::tools::mcp::McpCatalog;
 
 use state::{PendingEvent, PersistedSub, StateStore};
-pub use webhook::{MAX_WEBHOOK_BODY, WebhookReply, receive_webhook, route_exists};
+pub use webhook::{MAX_WEBHOOK_BODY, WebhookReply, receive_webhook, restore_pending, route_exists};
 pub use wire::NotificationRouter;
 use wire::{Conn, RpcError, StreamMsg};
 
@@ -576,6 +576,15 @@ impl McpEventsHub {
                 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;
+                }
                 store.loaded().await;
                 let Some(hub) = weak_hub.upgrade() else {
                     return;
@@ -602,6 +611,11 @@ 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();
+                }
                 for (key, spec) in runtime {
                     hub.restoring
                         .lock()

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

@@ -105,6 +105,41 @@ 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);
+
+fn restoring_until() -> std::sync::MutexGuard<'static, Option<std::time::Instant>> {
+    RESTORING_UNTIL
+        .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);
+}
+
+/// The events session has reserved every token it persisted: unknown tokens are `410` again.
+pub(super) fn restored() {
+    *restoring_until() = None;
+}
+
+fn restoring() -> bool {
+    let mut until = restoring_until();
+    match *until {
+        Some(t) if t > std::time::Instant::now() => true,
+        Some(_) => {
+            *until = None;
+            false
+        }
+        None => false,
+    }
+}
+
 /// Hold a persisted callback token until its subscription re-registers (see [`RESERVED`]).
 pub(super) fn reserve(token: &str) {
     if routes().contains_key(token) {
@@ -133,7 +168,7 @@ fn is_reserved(token: &str) -> bool {
 /// Whether `token` names a live subscription, or one expected back after a restart — the
 /// listener's cheap check before it reads a body.
 pub fn route_exists(token: &str) -> bool {
-    routes().contains_key(token) || is_reserved(token)
+    routes().contains_key(token) || is_reserved(token) || restoring()
 }
 
 /// CSPRNG bytes. A failure is an error, never a fallback to something predictable.
@@ -296,7 +331,7 @@ pub async fn receive_webhook(
     body: &[u8],
 ) -> WebhookReply {
     let Some(route) = routes().get(token).cloned() else {
-        if is_reserved(token) {
+        if is_reserved(token) || restoring() {
             return WebhookReply::error(
                 503,
                 "Service Unavailable",

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

@@ -795,6 +795,10 @@ pub(super) struct SseReader {
     /// The largest event read (see [`max_message_bytes`]).
     max: usize,
     pub(super) error: Option<String>,
+    /// Bytes scanned for a separator plus bytes moved by dropping a consumed prefix: the work that
+    /// must stay linear in what was read (see `many_events_in_one_chunk_drain_in_linear_time`).
+    #[cfg(test)]
+    work: usize,
 }
 
 impl SseReader {
@@ -810,12 +814,21 @@ impl SseReader {
             scanned: 0,
             max,
             error: None,
+            #[cfg(test)]
+            work: 0,
         }
     }
 
     pub(super) async fn next(&mut self) -> Option<Value> {
         loop {
-            if let Some((end, sep_end)) = find_event_end(&self.buf, self.scanned.max(self.start)) {
+            let from = self.scanned.max(self.start);
+            let found = find_event_end(&self.buf, from);
+            #[cfg(test)]
+            {
+                self.work +=
+                    found.map_or(self.buf.len(), |(_, sep_end)| sep_end) - from.min(self.buf.len());
+            }
+            if let Some((end, sep_end)) = found {
                 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.
@@ -839,6 +852,10 @@ impl SseReader {
             }
             // Need more bytes: drop what has been consumed (at most a partial event remains).
             if self.start > 0 {
+                #[cfg(test)]
+                {
+                    self.work += self.buf.len() - self.start;
+                }
                 self.buf.drain(..self.start);
                 self.scanned = self.scanned.saturating_sub(self.start);
                 self.start = 0;
@@ -1205,18 +1222,21 @@ mod tests {
         // All of it already buffered — one large read, as a fast server on a fast link delivers.
         let resp = reqwest::Response::from(http::Response::new(String::new()));
         let mut reader = SseReader::new(resp);
+        let len = body.len();
         reader.buf = body.into_bytes();
-        let started = std::time::Instant::now();
         let mut seen = 0;
         while let Some(v) = reader.next().await {
             assert_eq!(v["i"], seen);
             seen += 1;
         }
         assert_eq!(seen, n);
+        // Counted, not timed: a clock bound flips under host load. Every byte is scanned about
+        // once and the consumed prefix is dropped once per chunk; dropping it per event, or
+        // rescanning from the event's start, makes this quadratic (~n^2 / 2 bytes).
         assert!(
-            started.elapsed() < Duration::from_secs(3),
-            "draining {n} events took {:?}",
-            started.elapsed()
+            reader.work <= 2 * len,
+            "draining {n} events ({len} bytes) did {} bytes of scanning and moving",
+            reader.work
         );
     }