diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 8262c414..614a79e3 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/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 - (`.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 + (`.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 diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index 7ed5b825..998b355a 100644 --- a/crates/agent/src/serve.rs +++ b/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::() + { + 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!(); diff --git a/crates/agent/src/serve_ws.rs b/crates/agent/src/serve_ws.rs index 5e2c920d..98e9945f 100644 --- a/crates/agent/src/serve_ws.rs +++ b/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::(); 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. diff --git a/crates/agent/src/session_store.rs b/crates/agent/src/session_store.rs index e5205e2d..997a9c12 100644 --- a/crates/agent/src/session_store.rs +++ b/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] diff --git a/crates/agent/src/tools/mcp_events/mod.rs b/crates/agent/src/tools/mcp_events/mod.rs index d911f9e7..24ed1935 100644 --- a/crates/agent/src/tools/mcp_events/mod.rs +++ b/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::().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() diff --git a/crates/agent/src/tools/mcp_events/webhook.rs b/crates/agent/src/tools/mcp_events/webhook.rs index d6bbfd2f..8cdbc252 100644 --- a/crates/agent/src/tools/mcp_events/webhook.rs +++ b/crates/agent/src/tools/mcp_events/webhook.rs @@ -105,6 +105,41 @@ fn reserved() -> std::sync::MutexGuard<'static, HashMap> = Mutex::new(None); + +fn restoring_until() -> std::sync::MutexGuard<'static, Option> { + 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", diff --git a/crates/agent/src/tools/mcp_events/wire.rs b/crates/agent/src/tools/mcp_events/wire.rs index 03ae96bf..556e9987 100644 --- a/crates/agent/src/tools/mcp_events/wire.rs +++ b/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, + /// 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 { 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 ); } diff --git a/crates/agent/src/tools/mcp_stdio.rs b/crates/agent/src/tools/mcp_stdio.rs index 487d90e6..1bd6d9ad 100644 --- a/crates/agent/src/tools/mcp_stdio.rs +++ b/crates/agent/src/tools/mcp_stdio.rs @@ -115,15 +115,32 @@ pub(crate) fn retire(server: &std::sync::Arc) { /// Every streamable-HTTP MCP session still open is ended too (`mcp_http_exit`): its `DELETE`s go out /// first, in parallel and bounded by [`CLOSE_DEADLINE`](crate::tools::mcp_http_exit::CLOSE_DEADLINE), /// overlapping the stdio grace. +/// +/// The wait is a bound, so a loaded host can stall a sweep thread past it. Whatever the threads +/// managed by then, every server still running is killed here, on the exiting thread, before it +/// returns: the guarantee that nothing is left running rests on that, not on the threads' timing. pub fn sweep_before_exit() { let http_closed = crate::tools::mcp_http_exit::begin_close_all(crate::tools::mcp_http_exit::CLOSE_DEADLINE); - retire_all(); + let live = retire_all(); http_closed(); #[cfg(unix)] - crate::tools::exec::wait_for_pending_group_kills( - SHUTDOWN_GRACE + std::time::Duration::from_secs(2), - ); + { + crate::tools::exec::wait_for_pending_group_kills( + SHUTDOWN_GRACE + std::time::Duration::from_secs(2), + ); + for server in &live { + if !server.exited() { + // Not reaped, so the pid is still this server's. + let _ = lock(&server.child).start_kill(); + } + if let Some(pgid) = server.pgid { + crate::tools::exec::kill_group_members(pgid); + } + } + } + #[cfg(not(unix))] + drop(live); } /// Runs [`sweep_before_exit`] when dropped — at the end of `main`, and while a panic unwinds out of @@ -139,7 +156,7 @@ impl Drop for ExitSweep { /// Retire every stdio server still running — for the way out of the process, where a connection /// some owner still holds would otherwise never be dropped. Follow it with /// [`crate::tools::exec::wait_for_pending_group_kills`] ([`sweep_before_exit`] does both). -pub fn retire_all() { +pub fn retire_all() -> Vec> { let live: Vec<_> = lock(&LIVE) .iter() .filter_map(std::sync::Weak::upgrade) @@ -147,6 +164,7 @@ pub fn retire_all() { for server in &live { retire(server); } + live } /// The key a rescued result is wrapped under on its way through rmcp; see [`rescue`]. @@ -1224,11 +1242,26 @@ mod tests { } tokio::time::sleep(std::time::Duration::from_millis(20)).await; }; + // Running, and not yet sent SIGKILL. A process that has been sent SIGKILL may not have been + // scheduled to die yet on a loaded host, so "killed" is read from the kernel: gone, a + // zombie, or SIGKILL pending (bit 8 of its pending masks). That is the sweep's guarantee; + // when the corpse is reaped is the scheduler's business. let alive = |pid: u32| { - std::fs::read_to_string(format!("/proc/{pid}/stat")).is_ok_and(|s| { - s.rsplit_once(") ") - .is_some_and(|(_, rest)| !rest.starts_with('Z')) - }) + let Ok(status) = std::fs::read_to_string(format!("/proc/{pid}/status")) else { + return false; + }; + let field = |name: &str| { + status + .lines() + .find_map(|l| l.strip_prefix(name)) + .map(|v| v.trim().to_owned()) + .unwrap_or_default() + }; + let kill_pending = ["SigPnd:", "ShdPnd:"] + .iter() + .any(|f| u64::from_str_radix(&field(f), 16).is_ok_and(|mask| mask & (1 << 8) != 0)); + let state = field("State:"); + !kill_pending && !state.starts_with('Z') && !state.starts_with('X') }; assert!(alive(orphan)); let held = (server, transport); diff --git a/crates/agent/tests/mcp_events_nested.rs b/crates/agent/tests/mcp_events_nested.rs index 6693a0c7..7ddf3afa 100644 --- a/crates/agent/tests/mcp_events_nested.rs +++ b/crates/agent/tests/mcp_events_nested.rs @@ -17,6 +17,10 @@ use common::{BIN, SpawnGuarded, serve_cmd, spawn_model_server_routed, turn_text} use serde_json::json; fn nested_during(during: &str, delivery: &str) { + nested_during_with(during, delivery, &[]); +} + +fn nested_during_with(during: &str, delivery: &str, env: &[(&str, &str)]) { let home = tempfile::tempdir().unwrap(); let control_file = home.path().join("control"); write_settings( @@ -33,6 +37,7 @@ fn nested_during(during: &str, delivery: &str) { cmd.env("HOME", home.path()) .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") .env("BEYOND_AI_AGENT_MCP_EVENTS_POLL_FLOOR_MS", "100") + .envs(env.iter().copied()) .stderr(Stdio::null()); let mut child = cmd.spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); @@ -75,6 +80,18 @@ fn a_nested_elicitation_during_an_events_poll_reaches_the_owning_session() { nested_during("poll", "poll"); } +/// The events hub starts only once the session can take a question. Held open on demand (a debug +/// seam delays the gates' install), the window in which the first `events/poll` could raise one +/// before then would have it declined as having no client. A loaded host used to open it by chance. +#[test] +fn a_nested_elicitation_on_the_first_poll_waits_for_the_session_to_take_questions() { + nested_during_with( + "poll", + "poll", + &[("BEYOND_AI_AGENT_TEST_SLOW_GATE_INSTALL_MS", "1500")], + ); +} + #[test] fn a_nested_elicitation_during_an_events_stream_reaches_the_owning_session() { nested_during("stream", "push"); diff --git a/crates/agent/tests/mcp_events_receiver.rs b/crates/agent/tests/mcp_events_receiver.rs index 7d98257d..ae5203c0 100644 --- a/crates/agent/tests/mcp_events_receiver.rs +++ b/crates/agent/tests/mcp_events_receiver.rs @@ -226,8 +226,10 @@ async fn the_callback_survives_a_crash_so_retried_deliveries_still_land() { } /// After a restart, a delivery the server retries can arrive before the subscription has -/// re-registered its route (here discovery is slow). It must get `503` — retry — not `410`, which -/// tells the server to give up; once the route is back, the retry lands. +/// re-registered its route (here discovery is slow), and even before the daemon has read the state +/// that names its callback (held off on demand by a debug seam; a loaded host used to open that +/// window by chance). It must get `503` — retry — not `410`, which tells the server to give up; +/// once the route is back, the retry lands. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn a_retry_that_beats_the_resubscribe_after_a_restart_is_told_to_retry_not_to_stop() { let (_fx, mcp_url, fixture) = spawn_http_fixture(&[ @@ -262,9 +264,15 @@ async fn a_retry_that_beats_the_resubscribe_after_a_restart_is_told_to_retry_not assert_ne!(r["deliveries"][0]["status"], 200, "{r:#}"); drop(down); - let _second = daemon(home.path(), &base, &held, &[]); - // The state is read at once; discovery takes 3 s, so the route is not back yet. - tokio::time::sleep(Duration::from_millis(800)).await; + let _second = daemon( + home.path(), + &base, + &held, + &[("BEYOND_AI_AGENT_TEST_SLOW_EVENTS_RESTORE_MS", "1500")], + ); + // At once: the state is not read for 1.5 s, and discovery takes 3 s after that, so neither + // the token's reservation nor its route exists yet. The port is held, so the delivery waits + // in its backlog until the daemon accepts it. let r = control( &fixture, "POST", diff --git a/crates/agent/tests/mcp_skills_integrity.rs b/crates/agent/tests/mcp_skills_integrity.rs index 7800cf00..6ee65431 100644 --- a/crates/agent/tests/mcp_skills_integrity.rs +++ b/crates/agent/tests/mcp_skills_integrity.rs @@ -170,7 +170,14 @@ fn no_skills_suppresses_mcp_skills_too() { #[test] fn a_listing_that_never_ends_is_cut_off_and_said_so() { - let env = Env::new(json!({ "MCP_SKILLS_FIXTURE_ENDLESS": "1" })); + // The cut-off is the page budget (`MAX_LIST_PAGES`), counted, not timed. The listing stays + // fresh for the whole run (`ttlMs`), so the turn does not page through all 64 again under the + // refresh's wall-clock bound: on a loaded host that second walk could outlast the bound, and + // this test is about the budget, not about refresh timing. + let env = Env::new(json!({ + "MCP_SKILLS_FIXTURE_ENDLESS": "1", + "MCP_SKILLS_FIXTURE_TTL_MS": "3600000", + })); let out = env.run_with("hello", vec![turn_text("ok")], &[]); assert!( out.stderr.contains("cut off after 64 pages"), @@ -178,6 +185,11 @@ fn a_listing_that_never_ends_is_cut_off_and_said_so() { out.stderr ); assert!(system_text(&out.bodies[0]).contains("docs:git-workflow")); + assert_eq!( + env.log_of("skills/list").len(), + 64, + "exactly the page budget was walked, once" + ); } #[test] diff --git a/crates/agent/tests/serve_health.rs b/crates/agent/tests/serve_health.rs index eb9ddf9d..96f2afe3 100644 --- a/crates/agent/tests/serve_health.rs +++ b/crates/agent/tests/serve_health.rs @@ -59,6 +59,23 @@ async fn get(port: u16, path: &str, headers: &[(&str, &str)]) -> (u16, String) { exchange(stream, "GET", path, headers).await } +/// `GET /readyz`, once the answer has settled. The shard probe is a real `create` and `unlink` with a +/// one-second deadline, and on a loaded host a probe can run past it: that caller is answered +/// "shard probe timed out", while the probe goes on and its answer lands in the memo for the next +/// caller. That answer says nothing about what a test checks, so it is waited out. Any other answer, +/// a 503 for any other reason included, is returned at once. +async fn readyz(port: u16, headers: &[(&str, &str)]) -> (u16, String) { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60); + loop { + let (status, response) = get(port, "/readyz", headers).await; + let timed_out = status == 503 && body_of(&response)["reason"] == "shard probe timed out"; + if !timed_out || std::time::Instant::now() > deadline { + return (status, response); + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } +} + fn body_of(response: &str) -> Value { let body = response .split_once("\r\n\r\n") @@ -95,13 +112,13 @@ async fn readyz_is_200_when_every_shard_is_writable() { let (base, _requests) = spawn_model_server(vec![]); let svc = Service::start(&base, &["s1", "s2"]).await; - let (status, response) = get(svc.port, "/readyz", &[]).await; + let (status, response) = readyz(svc.port, &[]).await; assert_eq!(status, 200, "{response}"); assert_eq!(body_of(&response)["status"], "ready"); // Repeated probes are served from the memo and stay correct. for _ in 0..3 { - let (status, _) = get(svc.port, "/readyz", &[]).await; + let (status, _) = readyz(svc.port, &[]).await; assert_eq!(status, 200); } @@ -125,7 +142,7 @@ async fn readyz_is_503_and_names_the_missing_shard() { let gone = svc.shard("s2").to_path_buf(); std::fs::remove_dir_all(&gone).unwrap(); - let (status, response) = get(svc.port, "/readyz", &[]).await; + let (status, response) = readyz(svc.port, &[]).await; assert_eq!(status, 503, "{response}"); let reason = body_of(&response)["reason"] .as_str() @@ -160,7 +177,7 @@ async fn readyz_is_503_when_a_shard_is_read_only() { let ro = svc.shard("s2").to_path_buf(); std::fs::set_permissions(&ro, std::fs::Permissions::from_mode(0o555)).unwrap(); - let (status, response) = get(svc.port, "/readyz", &[]).await; + let (status, response) = readyz(svc.port, &[]).await; // Restore before asserting, so a failure still lets the temp dir clean itself up. std::fs::set_permissions(&ro, std::fs::Permissions::from_mode(0o755)).unwrap(); @@ -227,7 +244,7 @@ async fn health_answers_over_a_unix_socket() { assert_eq!(body_of(&response)["status"], "ready"); // Same process, same shards: the TCP listener agrees. - let (status, _) = get(svc.port, "/readyz", &[]).await; + let (status, _) = readyz(svc.port, &[]).await; assert_eq!(status, 200); } @@ -254,7 +271,7 @@ async fn a_grant_is_never_required_and_never_echoed() { // A grant far past the header cap is dropped unread on this path too — still a 200, not a 401. let huge = "A".repeat(13 * 1024); - let (status, response) = get(svc.port, "/readyz", &[("x-beyond-grant", &huge)]).await; + let (status, response) = readyz(svc.port, &[("x-beyond-grant", &huge)]).await; assert_eq!(status, 200, "{response}"); assert!(!response.contains("AAAA"), "{response}"); } @@ -266,7 +283,7 @@ async fn readyz_does_not_interfere_with_a_session_attach() { let (base, _requests) = spawn_model_server(vec![turn_text("hi")]); let svc = Service::start(&base, &["s1"]).await; - let (status, _) = get(svc.port, "/readyz", &[]).await; + let (status, _) = readyz(svc.port, &[]).await; assert_eq!(status, 200); let token = svc.token("t1", "s1.alpha"); @@ -274,7 +291,7 @@ async fn readyz_does_not_interfere_with_a_session_attach() { .await .expect("attach with a valid grant"); - let (status, _) = get(svc.port, "/readyz", &[]).await; + let (status, _) = readyz(svc.port, &[]).await; assert_eq!(status, 200, "a probe mid-session must not fail"); ws_send( @@ -291,7 +308,7 @@ async fn readyz_does_not_interfere_with_a_session_attach() { // And the probe still only ever reports on shards — no session directory is walked, so a live // session leaves no trace in the answer. - let (status, response) = get(svc.port, "/readyz", &[]).await; + let (status, response) = readyz(svc.port, &[]).await; assert_eq!(status, 200); assert!(!response.contains("s1.alpha"), "{response}"); } @@ -324,7 +341,7 @@ async fn health_exists_on_a_non_service_daemon() { assert_eq!(status, 200, "{response}"); assert_eq!(body_of(&response)["status"], "alive"); - let (status, response) = get(port, "/readyz", &[]).await; + let (status, response) = readyz(port, &[]).await; assert_eq!(status, 200, "{response}"); assert_eq!(body_of(&response)["status"], "ready"); diff --git a/crates/agent/tests/serve_lifecycle.rs b/crates/agent/tests/serve_lifecycle.rs index 0671e230..4916dc44 100644 --- a/crates/agent/tests/serve_lifecycle.rs +++ b/crates/agent/tests/serve_lifecycle.rs @@ -34,33 +34,71 @@ const WAIT: Duration = Duration::from_secs(60); struct Collector { url: String, posts: Arc>>, + /// While set, every POST is recorded but not answered (see [`Collector::spawn_held`]). + hold: Arc, _guard: Option, } +/// Holds a collector's answers until released, and counts the ones it has sent. +#[derive(Default)] +struct Hold { + held: Mutex, + released: std::sync::Condvar, + answered: std::sync::atomic::AtomicUsize, +} + +impl Hold { + fn wait(&self) { + let mut held = self.held.lock().unwrap(); + while *held { + held = self.released.wait(held).unwrap(); + } + } +} + impl Collector { fn spawn() -> Self { - Self::spawn_with_delay(Duration::ZERO) + Self::spawn_tcp(false) + } + + /// A collector that records each POST and then answers nothing until [`Collector::release`]. + fn spawn_held() -> Self { + Self::spawn_tcp(true) } - fn spawn_with_delay(delay: Duration) -> Self { + fn spawn_tcp(held: bool) -> Self { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); let posts = Arc::new(Mutex::new(Vec::new())); - let recorder = posts.clone(); + let hold = Arc::new(Hold { + held: Mutex::new(held), + ..Hold::default() + }); + let (recorder, gate) = (posts.clone(), hold.clone()); thread::spawn(move || { for conn in listener.incoming() { let Ok(stream) = conn else { break }; - let recorder = recorder.clone(); - thread::spawn(move || handle_post(stream, delay, recorder)); + let (recorder, gate) = (recorder.clone(), gate.clone()); + thread::spawn(move || handle_post(stream, &gate, recorder)); } }); Self { url: format!("http://{addr}/lifecycle"), posts, + hold, _guard: None, } } + fn release(&self) { + *self.hold.held.lock().unwrap() = false; + self.hold.released.notify_all(); + } + + fn answered(&self) -> usize { + self.hold.answered.load(std::sync::atomic::Ordering::SeqCst) + } + fn bodies(&self) -> Vec { self.posts .lock() @@ -89,17 +127,19 @@ impl Collector { let sock = dir.path().join("lifecycle.sock"); let listener = std::os::unix::net::UnixListener::bind(&sock).unwrap(); let posts = Arc::new(Mutex::new(Vec::new())); - let recorder = posts.clone(); + let hold = Arc::new(Hold::default()); + let (recorder, gate) = (posts.clone(), hold.clone()); thread::spawn(move || { for conn in listener.incoming() { let Ok(stream) = conn else { break }; - let recorder = recorder.clone(); - thread::spawn(move || handle_post(stream, Duration::ZERO, recorder)); + let (recorder, gate) = (recorder.clone(), gate.clone()); + thread::spawn(move || handle_post(stream, &gate, recorder)); } }); Self { url: format!("unix://{}", sock.display()), posts, + hold, _guard: Some(dir), } } @@ -117,7 +157,7 @@ impl Collector { fn handle_post( mut stream: impl Read + Write, - delay: Duration, + hold: &Hold, posts: Arc>>, ) { let mut buf = Vec::new(); @@ -156,12 +196,12 @@ fn handle_post( let headers = String::from_utf8_lossy(&buf[..split]).into_owned(); let body = String::from_utf8_lossy(&buf[split.min(buf.len())..]).into_owned(); posts.lock().unwrap().push((headers, body)); - if !delay.is_zero() { - thread::sleep(delay); - } + hold.wait(); let _ = stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"); let _ = stream.flush(); + hold.answered + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); } fn of_type<'a>(bodies: &'a [Value], kind: &str) -> Vec<&'a Value> { @@ -517,10 +557,10 @@ fn serve_refusal_is_succeeded_with_refused() { fn a_slow_lifecycle_consumer_does_not_delay_the_prompt_response() { let dir = tempfile::tempdir().unwrap(); let session_file = dir.path().join("s.jsonl").to_string_lossy().into_owned(); - let collector = Collector::spawn_with_delay(Duration::from_secs(2)); - // Stall the *model* as well so we can send the prompt and then measure only the control-plane - // round trip — but a stalled lifecycle POST must not hold the response. Instant model, slow - // collector: if `emit` awaited the network, the prompt response would take ~2s. + // The collector answers no lifecycle POST until the test releases it, so the prompt response + // can only arrive if `emit` does not wait on the network: proved by order, not by a clock (a + // time bound flipped under host load). If it waited, the read below would stall and fail. + let collector = Collector::spawn_held(); let (base, _bodies) = spawn_model_server(vec![turn_text("quick")]); let mut child = serve_life(&base, &session_file, &collector.url).spawn_guarded(); @@ -531,14 +571,16 @@ fn a_slow_lifecycle_consumer_does_not_delay_the_prompt_response() { &mut stdin, json!({ "type": "prompt", "id": "fast", "message": "hi" }), ); - let start = Instant::now(); let frames = read_until_response(&mut stdout, "prompt"); - let elapsed = start.elapsed(); assert_eq!(frames.last().unwrap()["success"], true, "{frames:#?}"); - assert!( - elapsed < Duration::from_millis(1500), - "prompt response waited on the lifecycle consumer ({elapsed:?})" + assert_eq!( + collector.answered(), + 0, + "the prompt response came only after the lifecycle consumer answered" ); + collector.release(); + // The events were really posted (and held), not skipped. + collector.wait_until(|b| of_type(b, "succeeded").len() == 1); drop(stdin); child.wait().unwrap(); diff --git a/crates/agent/tests/serve_uds.rs b/crates/agent/tests/serve_uds.rs index f7e2ebb5..3f0cc944 100644 --- a/crates/agent/tests/serve_uds.rs +++ b/crates/agent/tests/serve_uds.rs @@ -195,10 +195,14 @@ async fn uds_stale_socket_file_is_reclaimed() { let dir = tempfile::tempdir().unwrap(); let sock = dir.path().join("agent.sock"); - // Pre-create a *dead* socket node: bind a UnixListener, then drop it. The path now exists on disk - // (a socket file) but nothing is listening — exactly the crashed-daemon leftover we must reclaim. + // Pre-create a *dead* socket node: the path exists on disk (a socket file) but nothing listens + // on it — exactly the crashed-daemon leftover we must reclaim. Bound as a datagram socket, never + // a listening one: a listener bound and dropped here can live on for a moment in a child another + // test forks in that instant (the descriptor closes only at its `exec`), and a stream connect + // landing then succeeds, so the daemon would see a "live" socket and refuse to start. Nothing + // can make a stream connect to a datagram node succeed. { - let stale = std::os::unix::net::UnixListener::bind(&sock).unwrap(); + let stale = std::os::unix::net::UnixDatagram::bind(&sock).unwrap(); drop(stale); } assert!(sock.exists(), "the stale socket node should be on disk"); diff --git a/crates/agent/tests/tool_reactor_stall.rs b/crates/agent/tests/tool_reactor_stall.rs index f5a659d7..7ba79cc1 100644 --- a/crates/agent/tests/tool_reactor_stall.rs +++ b/crates/agent/tests/tool_reactor_stall.rs @@ -95,8 +95,16 @@ fn inline_cpu(mut work: impl FnMut()) -> Duration { /// A yielding tool may hold the executor for at most this fraction of the inline baseline (its own /// argument parsing, the hand-off, and assembling the result). An inline one holds about all of it. +/// +/// That per-call overhead is fixed (a few hundred microseconds in a debug build) while the baseline +/// scales with the input, so the inputs are sized for a baseline of several milliseconds even on a +/// fast machine: a 4 MB `write` was under a millisecond on a CI runner, and the fixed overhead alone +/// came to a third of it. const MAX_HELD_FRACTION: f64 = 0.25; +/// Lines of [`big_ascii_source`] for the `edit` and `write` subjects: ~36 MB. +const SUBJECT_LINES: usize = 640_000; + fn assert_yields(tool: &str, held: &Held, baseline: Duration, why: &str) { eprintln!( "PROBE {tool}: held {:?} of {baseline:?}, {} pendings", @@ -130,8 +138,8 @@ const NEW: &str = " let x_40001 = compute(i, 40001) + adjust(); // edited"; async fn edit_does_not_stall_the_runtime() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("subject.rs"); - // ~4 MB: an inline `edit` spends tens of ms here (read + normalize + match + splice + write). - let src = big_ascii_source(80_000); + // An inline `edit` spends tens of ms here (read + normalize + match + splice + write). + let src = big_ascii_source(SUBJECT_LINES); std::fs::write(&path, &src).unwrap(); let scratch = dir.path().join("baseline.rs"); let baseline = inline_cpu(|| { @@ -165,7 +173,7 @@ async fn edit_does_not_stall_the_runtime() { async fn write_does_not_stall_the_runtime() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("out.rs"); - let src = big_ascii_source(80_000); + let src = big_ascii_source(SUBJECT_LINES); let scratch = dir.path().join("baseline.rs"); let baseline = inline_cpu(|| std::fs::write(&scratch, &src).unwrap());