diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index ce45661a..b852cba2 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -933,7 +933,7 @@ The harness layers several capabilities over the bare tools + loop: Proxy / TypedArrays / MapSet) runs it against `tools..(input)` Promises. Built-in coding tools stay direct. Nested calls honor `--exclude-tools` / `--deny-tool`. Budgets: 30s wall clock, 64 nested calls, 256 KiB result, **4 MiB JS heap**, **128 KiB JS stack**, and **one QuickJS - runtime per process** (a `serve` daemon never holds two heaps at once). An allow-list of built-ins + runtime per process** (a `serve` daemon never holds two heaps at once). QuickJS evaluates synchronously, so evaluation runs on the blocking pool under a local executor, never on the runtime thread other sessions share: a spinning script holds a pool thread until its interrupt (deadline, cancellation, or the call being dropped — a guard trips it) fires, and the one-heap slot is released when the evaluation really ends. Nested host calls are spawned back onto the session's runtime, under the script's deadline (which so holds across a nested await, not only while JS runs), and aborted by the same guard however the run ends — cancelled or dropped — so no host tool runs on after it. An allow-list of built-ins (`--tools read,bash`) still keeps `execute` so the deferred catalog is reachable; `--exclude-tools execute` (or `--no-tools`) drops it. Subagents inherit the parent's flag so a child cannot re-advertise the catalog the parent deferred. QuickJS is a cargo feature (`code-mode`) because idle @@ -2104,7 +2104,7 @@ the model reads. Body/time caps are ax's: 20 MB, 30 s. **Egress safety is the load-bearing part** (`web/ssrf.rs`). ax is a local CLI a human runs; this is a server fetching URLs the _model_ chose. Because the tool is read-only (in `READ_ONLY_TOOLS`, so `--approve all` never prompts on it — the operator's choice), the SSRF filter is the **sole** boundary -between the model and an internal fetch. It is deny-by-default, in two layers: +between the model and an internal fetch. So it **fails closed**: if the hardened client (resolver, timeout, no redirects) cannot be built — the builder errs or panics — the call fails and nothing is sent; there is no fallback to a default client (`a_client_that_cannot_be_built_fails_the_call_and_sends_nothing`). The MCP Events direct-HTTP client fails closed the same way. It is deny-by-default, in two layers: - **A validating `reqwest::dns::Resolve`.** reqwest connects to exactly what the resolver returns, so checking every resolved `SocketAddr` validates the _actual_ connection IP — on the initial request and @@ -2487,7 +2487,7 @@ standing in for `serve_session` so the test decides exactly when an exiting task `#[tokio::main]` builds one, and that's what this binary used to inherit. Session tasks now share that process runtime (`serve_ws` `tokio::spawn`s each `serve_session`; the future is `Send`). CPU-bound tool work (`grep`/`find`/image resize) is `spawn_blocking`, which a current-thread -runtime still has a blocking pool for. Extra workers cost per-thread stacks and mimalloc heaps, on +runtime still has a blocking pool for. So is anything that blocks the thread without using CPU: `web` builds its client and spawns and waits on its isolated parser child there, `bash`'s stale-spill sweep of the system temp directory runs there, `memory` merges and renders large search results there, and Code Mode evaluates there. `tests/tool_reactor_stall.rs` holds every tool that needs no external service to it with two probes on a `current_thread` runtime — the CPU the runtime thread spends inside the tool's polls, and the worst gap of a ticker task on the same runtime with run-queue wait (`/proc/thread-self/schedstat`) taken out, so wall-clock blocking (a sleep, blocking disk or NFS I/O, a contended `std` mutex) is caught and host load is not. Extra workers cost per-thread stacks and mimalloc heaps, on a binary whose idle RSS in a 768 MB guest is already the density budget. **The default now splits by mode** (`main.rs::build_runtime`): `current_thread` for a local daemon, @@ -4837,7 +4837,7 @@ mcp_events_subscribe (any session) ──► owned by that session session-bound HTTP server goes through rmcp's own connection, which carries its `Mcp-Session-Id`. A server→client request (an elicitation, a sampling request) raised during an `events/*` request is attributed to the session that owns it over rmcp — the same rule as during a `tools/call` - (`track_call`) — and, over direct HTTP, where nothing could route it to a session, answered at once + (`track_call`); a session starts its hub only after its elicitation and sampling gates are installed, so one raised during the very first `events/*` request reaches its client rather than the default gate's decline — and, over direct HTTP, where nothing could route it to a session, answered at once with an error rather than left unanswered (`tests/mcp_events_nested.rs`). That answer is POSTed with the same `Accept` (`application/json, text/event-stream`) as every events request — a streamable-HTTP server refuses an answer without it `406` and keeps waiting (the fixture enforces diff --git a/crates/agent/src/bin/mcp_fixture_events_server.rs b/crates/agent/src/bin/mcp_fixture_events_server.rs index 65434e61..699cfbc1 100644 --- a/crates/agent/src/bin/mcp_fixture_events_server.rs +++ b/crates/agent/src/bin/mcp_fixture_events_server.rs @@ -111,6 +111,10 @@ struct State { /// `MCP_FIXTURE_NESTED_DURING=poll|stream`: the client's answers to the nested /// `elicitation/create` raised during that `events/*` request. nested_answers: Vec, + /// `MCP_FIXTURE_NESTED_ON_SIGNAL=1`: the nested request is held until `POST + /// /control/raise_nested` sets this — so a test raises it once its client is attached, rather + /// than guessing a delay. + nested_released: bool, /// `MCP_FIXTURE_FORBID_FIRST=`: the first `n` `events/poll`/`events/subscribe` requests are /// refused with `-32012` (forbidden) — credentials that are refreshed and then work. forbid_left: u64, @@ -1058,6 +1062,10 @@ async fn control(state: &Shared, method: &str, path: &str, body: &[u8]) -> (u16, None => (404, json!({"error": "no such event"})), } } + ("POST", "/control/raise_nested") => { + state.lock().unwrap().nested_released = true; + (200, json!({ "ok": true })) + } ("POST", "/control/jwks") => { let mut st = state.lock().unwrap(); st.jwks_status = req["status"].as_u64().map(|s| s as u16); @@ -1593,6 +1601,11 @@ async fn nested_elicitation(state: &Shared, out: &Stdout, waiters: &Waiters, dur if delay > 0 { tokio::time::sleep(Duration::from_millis(delay)).await; } + if env_flag("MCP_FIXTURE_NESTED_ON_SIGNAL") { + while !state.lock().unwrap().nested_released { + tokio::time::sleep(Duration::from_millis(10)).await; + } + } let id = format!("nested-{during}"); let (tx, rx) = tokio::sync::oneshot::channel(); waiters.lock().unwrap().insert(json!(id).to_string(), tx); diff --git a/crates/agent/src/exec_endpoint.rs b/crates/agent/src/exec_endpoint.rs index d824080d..85602da3 100644 --- a/crates/agent/src/exec_endpoint.rs +++ b/crates/agent/src/exec_endpoint.rs @@ -889,7 +889,7 @@ impl crate::tools::fs::FsBackend for CellFs { async fn write_if_unchanged( &self, path: &std::path::Path, - bytes: &[u8], + bytes: Vec, expected: Option, ) -> Result { self.inner()? diff --git a/crates/agent/src/tools/code_mode.rs b/crates/agent/src/tools/code_mode.rs index e4f3883c..5e983dc7 100644 --- a/crates/agent/src/tools/code_mode.rs +++ b/crates/agent/src/tools/code_mode.rs @@ -192,6 +192,47 @@ mod runtime { } } + /// Sets the evaluation's cancel flag when dropped, so its interrupt handler stops the script at + /// its next check however `execute` ends (returned, cancelled, or dropped mid-await). + /// …and aborts every nested host call still in flight, so a dropped `execute` (an abort that + /// drops the tool future without cancelling its token) leaves no host tool running on, and the + /// evaluation — and the JS slot it holds — ends at once rather than when that call would have. + struct InterruptOnDrop { + cancelled: Arc, + nested: NestedCalls, + } + + impl Drop for InterruptOnDrop { + fn drop(&mut self) { + self.cancelled.store(true, Ordering::Relaxed); + for call in lock_nested(&self.nested).drain(..) { + call.abort(); + } + } + } + + /// The nested host calls of one evaluation (finished ones included; aborting those is a no-op). + type NestedCalls = Arc>>; + + fn lock_nested( + nested: &NestedCalls, + ) -> std::sync::MutexGuard<'_, Vec> { + nested + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// How a script's host calls reach the session: its runtime (where host tools' connections + /// and timers live), the calls in flight, and the script's deadline — which holds across a + /// nested call too, not only while JS runs (the interrupt handler only fires then). + #[derive(Clone)] + struct Bridge { + host: tokio::runtime::Handle, + nested: NestedCalls, + deadline: std::time::Instant, + timed_out: Arc, + } + /// The model-facing Code Mode tool. pub struct Execute { tools: Arc>>, @@ -248,19 +289,53 @@ mod runtime { let code = code.to_string(); let progress = progress.clone(); - let run = run_js( - tools, - catalog, - code, - calls, - max_calls, - memory_limit, - stack_size, + // Take the process-wide slot *before* allocating a runtime, so two `serve` sessions never + // hold two QuickJS heaps at once. `acquire` is cancel-safe; a cancelled `execute` drops + // the wait without starving the holder. The permit then travels with the evaluation and + // is released when it really ends — after an interrupt, not merely when `execute` returns. + let permit = JS_SLOT + .acquire() + .await + .map_err(|_| ToolError::Execution("code mode runtime slot closed".into()))?; + // However `execute` ends — cancelled, or its future dropped by an abort — the evaluation + // is interrupted at its next check rather than left running to the deadline. + let nested: NestedCalls = Arc::default(); + let _interrupt = InterruptOnDrop { + cancelled: cancelled.clone(), + nested: nested.clone(), + }; + + // QuickJS evaluates synchronously: a busy script (`while (true) {}`) holds whatever thread + // runs it until the interrupt fires. So it runs on the blocking pool, driven by a local + // executor, never on the runtime thread this session shares with others (their progress, + // their aborts). Host calls go back to this runtime (`install_bridge`). + let bridge = Bridge { + host: tokio::runtime::Handle::current(), + nested, deadline, - timed_out.clone(), - cancelled.clone(), - progress.clone(), - ); + timed_out: timed_out.clone(), + }; + let run = tokio::task::spawn_blocking({ + let (timed_out, cancelled, progress) = + (timed_out.clone(), cancelled.clone(), progress.clone()); + move || { + let _permit = permit; + futures::executor::block_on(run_js( + tools, + catalog, + code, + calls, + max_calls, + memory_limit, + stack_size, + deadline, + timed_out, + cancelled, + progress, + bridge, + )) + } + }); tokio::select! { biased; @@ -269,6 +344,9 @@ mod runtime { Err(ToolError::Execution("code mode cancelled".into())) } result = run => { + let result = result.unwrap_or_else(|e| { + Err(ToolError::Execution(format!("code mode evaluation failed: {e}"))) + }); if timed_out.load(Ordering::Relaxed) { Err(ToolError::Execution(format!( "code mode timed out after {}ms", @@ -340,15 +418,8 @@ mod runtime { timed_out: Arc, cancelled: Arc, progress: ToolProgress, + bridge: Bridge, ) -> Result { - // Take the process-wide slot *before* allocating a runtime, so two `serve` sessions never - // hold two QuickJS heaps at once. `acquire` is cancel-safe; a cancelled `execute` drops the - // wait without starving the holder. - let _permit = JS_SLOT - .acquire() - .await - .map_err(|_| ToolError::Execution("code mode runtime slot closed".into()))?; - let catalog_json = serde_json::to_string(catalog.as_ref()).map_err(|e| { ToolError::Execution(format!("code mode catalog serialize failed: {e}")) })?; @@ -388,7 +459,7 @@ mod runtime { .map_err(js_err)?; ctx.async_with(async move |ctx| -> Result { - install_bridge(&ctx, tools, calls, max_calls, progress)?; + install_bridge(&ctx, tools, calls, max_calls, progress, bridge)?; ctx.globals() .set("__catalogJson", catalog_json) .map_err(js_err)?; @@ -408,6 +479,7 @@ mod runtime { calls: Arc, max_calls: usize, progress: ToolProgress, + bridge: Bridge, ) -> Result<(), ToolError> { let func = Function::new( ctx.clone(), @@ -415,18 +487,46 @@ mod runtime { let tools = tools.clone(); let calls = calls.clone(); let progress = progress.clone(); + let bridge = bridge.clone(); async move { - Ok::( - invoke_nested( - tools, - calls, - max_calls, - name, - input.0.unwrap_or_else(|| "{}".into()), - progress, - ) - .await, - ) + // The evaluation runs on a blocking-pool thread; the host tool runs on the + // session's runtime, where its connections and timers live — under the + // script's deadline (past it the call is dropped and the script times out), + // registered so a dropped `execute` aborts it, and raced against cancellation + // so an abort does not wait on it. + let call = invoke_nested( + tools, + calls, + max_calls, + name, + input.0.unwrap_or_else(|| "{}".into()), + progress.clone(), + ); + let (deadline, timed_out) = (bridge.deadline, bridge.timed_out.clone()); + let nested = bridge.host.spawn(async move { + match tokio::time::timeout_at(deadline.into(), call).await { + Ok(answer) => answer, + Err(_) => { + timed_out.store(true, Ordering::Relaxed); + envelope("code mode timed out during a nested tool call", true) + } + } + }); + // Aborted by `execute`'s `InterruptOnDrop` however the run ends (cancelled, + // dropped): the one place a nested call is stopped. + lock_nested(&bridge.nested).push(nested.abort_handle()); + let answer = + match futures::future::select(nested, Box::pin(progress.cancelled())).await + { + futures::future::Either::Left((Ok(answer), _)) => answer, + futures::future::Either::Left((Err(e), _)) => { + envelope(format!("nested tool call failed: {e}"), true) + } + futures::future::Either::Right(_) => { + envelope("code mode cancelled", true) + } + }; + Ok::(answer) } }), ) @@ -661,6 +761,120 @@ mod runtime { ) } + /// A host tool that takes `delay` and then leaves a mark — the side effect a nested call + /// must not have once its run was cancelled, dropped, or past its deadline. + struct Slow { + delay: Duration, + done: Arc, + } + + #[async_trait] + impl Tool for Slow { + fn name(&self) -> &str { + "mcp__slow__work" + } + fn description(&self) -> &str { + "slow" + } + fn input_schema(&self) -> Value { + json!({ "type": "object", "properties": {} }) + } + async fn run(&self, _input: Value) -> Result { + tokio::time::sleep(self.delay).await; + self.done.store(true, Ordering::SeqCst); + Ok("done".into()) + } + } + + fn slow(delay: Duration, timeout: Duration) -> (Execute, Arc) { + let done = Arc::new(AtomicBool::new(false)); + let tool = Execute::with_limits( + vec![Arc::new(Slow { + delay, + done: done.clone(), + })], + Limits { + timeout, + max_calls: 16, + max_output: 16 * 1024, + memory_limit: 4 * 1024 * 1024, + stack_size: 128 * 1024, + }, + ); + (tool, done) + } + + const CALL_SLOW: &str = "return await tools.slow.work({});"; + + /// Cancelling a run while its script awaits a nested call aborts that call: the host tool + /// does not go on to finish (and leave its side effect) after the cancel. + #[tokio::test] + async fn cancelling_a_run_aborts_its_nested_call() { + let (tool, done) = slow(Duration::from_millis(500), Duration::from_secs(30)); + let token = CancellationToken::new(); + let (tx, _rx) = futures::channel::mpsc::unbounded(); + let progress = ToolProgress::new(tx, "x".into(), NAME.into(), token.clone()); + let cancel = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(100)).await; + token.cancel(); + }); + let err = tool + .run_streaming(json!({ "code": CALL_SLOW }), &progress) + .await + .unwrap_err(); + assert!(err.to_string().contains("cancelled"), "{err}"); + cancel.await.unwrap(); + tokio::time::sleep(Duration::from_millis(900)).await; + assert!( + !done.load(Ordering::SeqCst), + "the nested call ran on after the cancel" + ); + } + + /// Dropping `execute` without cancelling its token (an abort that drops the tool future) + /// aborts its nested call too, and frees the JS slot at once rather than when the call + /// would have finished. + #[tokio::test] + async fn dropping_execute_aborts_its_nested_call_and_frees_the_slot() { + let (tool, done) = slow(Duration::from_secs(5), Duration::from_secs(30)); + let dropped = + tokio::time::timeout(Duration::from_millis(150), run_code(&tool, CALL_SLOW)).await; + assert!( + dropped.is_err(), + "still awaiting the nested call when dropped" + ); + let quick = + tokio::time::timeout(Duration::from_secs(2), run_code(&tool, "return 'next';")) + .await + .expect("the slot is free: the dropped run's nested call was aborted") + .unwrap(); + assert_eq!(quick, "next"); + assert!( + !done.load(Ordering::SeqCst), + "the nested call ran on after the drop" + ); + } + + /// The script's deadline holds while it awaits a nested call, not only while JS runs: the + /// call is dropped at the deadline and the run fails as timed out. + #[tokio::test] + async fn the_deadline_holds_across_a_nested_call() { + let (tool, done) = slow(Duration::from_secs(5), Duration::from_millis(300)); + let started = std::time::Instant::now(); + let err = run_code(&tool, CALL_SLOW).await.unwrap_err(); + assert!(err.to_string().contains("timed out"), "{err}"); + assert!( + started.elapsed() < Duration::from_secs(3), + "{:?}", + started.elapsed() + ); + tokio::time::sleep(Duration::from_millis(300)).await; + assert!( + !done.load(Ordering::SeqCst), + "the nested call ran past the deadline" + ); + } + async fn run_code(tool: &Execute, code: &str) -> Result { Ok(tool.run(json!({ "code": code })).await?.text) } diff --git a/crates/agent/src/tools/edit.rs b/crates/agent/src/tools/edit.rs index 11a21205..e8526b64 100644 --- a/crates/agent/src/tools/edit.rs +++ b/crates/agent/src/tools/edit.rs @@ -176,15 +176,17 @@ impl Edit { let initial_mtime = meta.as_ref().and_then(|m| m.mtime); let raw_bytes = self.backend.read_bytes(p, 0, MAX_EDIT_BYTES).await?; - let raw = String::from_utf8(raw_bytes) - .map_err(|_| ToolError::Execution(format!("read {path}: file is not valid UTF-8")))?; // The match/splice work is pure CPU over an in-memory string and can run for tens of // milliseconds on a multi-MB file, so it stays on a blocking worker — inline it would pin a // worker of `serve_ws`'s shared runtime (and a `current_thread` test runtime), stalling that // session's event pump and abort/steer loop. See `tests/tool_reactor_stall.rs`. + // So is validating the bytes as UTF-8 — a pass over every byte of the file. let plan_path = path.clone(); let (content, ok) = match tokio::task::spawn_blocking(move || { + let raw = String::from_utf8(raw_bytes).map_err(|_| { + ToolError::Execution(format!("read {plan_path}: file is not valid UTF-8")) + })?; plan_edits(&plan_path, &raw, edits, replace_all) }) .await @@ -210,7 +212,7 @@ impl Edit { let handle = tokio::spawn(async move { let _write_lock = write_lock; backend - .write_if_unchanged(std::path::Path::new(&write_path), &content, initial_mtime) + .write_if_unchanged(std::path::Path::new(&write_path), content, initial_mtime) .await }); let unchanged = match handle.await { @@ -1472,7 +1474,11 @@ mod tests { .unwrap(); let wrote = backend - .write_if_unchanged(std::path::Path::new(p), restored.as_bytes(), initial_mtime) + .write_if_unchanged( + std::path::Path::new(p), + restored.into_bytes(), + initial_mtime, + ) .await .unwrap(); assert!( @@ -1509,7 +1515,7 @@ mod tests { let wrote = backend .write_if_unchanged( std::path::Path::new(p), - raw.replace("quick", "slow").as_bytes(), + raw.replace("quick", "slow").into_bytes(), initial_mtime, ) .await diff --git a/crates/agent/src/tools/fs/local.rs b/crates/agent/src/tools/fs/local.rs index b9a58ac1..51e4015b 100644 --- a/crates/agent/src/tools/fs/local.rs +++ b/crates/agent/src/tools/fs/local.rs @@ -121,11 +121,10 @@ impl FsBackend for LocalFs { async fn write_if_unchanged( &self, path: &Path, - bytes: &[u8], + bytes: Vec, expected: Option, ) -> Result { let path = path.to_path_buf(); - let bytes = bytes.to_vec(); blocking("write", move || { let current = std::fs::metadata(&path).and_then(|m| m.modified()).ok(); if current != expected { diff --git a/crates/agent/src/tools/fs/mod.rs b/crates/agent/src/tools/fs/mod.rs index a99dc3f6..c263eac4 100644 --- a/crates/agent/src/tools/fs/mod.rs +++ b/crates/agent/src/tools/fs/mod.rs @@ -286,10 +286,13 @@ pub trait FsBackend: Send + Sync { /// This is one method rather than a stat-then-write by the caller because the check and the write /// have to be as close together as the backend can make them; splitting it across the trait would /// put a guaranteed round trip inside the race window it exists to narrow. + /// + /// Owned, like [`write_bytes`](Self::write_bytes)'s: a backend that writes from another thread + /// takes the buffer as it is, rather than copying the whole file on the caller's (runtime) thread. async fn write_if_unchanged( &self, path: &Path, - bytes: &[u8], + bytes: Vec, expected: Option, ) -> Result; diff --git a/crates/agent/src/tools/fs/shell.rs b/crates/agent/src/tools/fs/shell.rs index 18d45498..00ffe193 100644 --- a/crates/agent/src/tools/fs/shell.rs +++ b/crates/agent/src/tools/fs/shell.rs @@ -362,7 +362,7 @@ impl FsBackend for ShellFs { async fn write_if_unchanged( &self, path: &Path, - bytes: &[u8], + bytes: Vec, expected: Option, ) -> Result { // Two round trips, and therefore a genuinely wider race window than `LocalFs`'s single @@ -374,7 +374,7 @@ impl FsBackend for ShellFs { if !mtimes_match(current, expected) { return Ok(false); } - self.write_slice(path, bytes).await?; + self.write_slice(path, &bytes).await?; Ok(true) } diff --git a/crates/agent/src/tools/mcp_events/mod.rs b/crates/agent/src/tools/mcp_events/mod.rs index 96a04799..587adcb3 100644 --- a/crates/agent/src/tools/mcp_events/mod.rs +++ b/crates/agent/src/tools/mcp_events/mod.rs @@ -508,6 +508,17 @@ struct Hub { dropped_seen: Mutex>, } +/// The direct-HTTP events client, or why it could not be built. **Fails closed**: a client that +/// cannot be built with redirects off is an error for the request, never a fallback to a default +/// client (which would follow redirects itself, past the URL that was checked). +fn events_http_client( + build: impl FnOnce() -> Result, +) -> Result { + build().map_err(|e| { + format!("cannot build the MCP Events HTTP client, so no request was made: {e}") + }) +} + /// A session's MCP Events client: its subscriptions, its coalescer, and the injection path into /// the session's own command channel. Owned by `serve_session`; dropping it cancels everything. pub struct McpEventsHub { @@ -1103,18 +1114,20 @@ impl Hub { async fn conn(&self, server: &str) -> Result { let peer = self.catalog.events_peer(server).await?; let http = match &peer { - crate::tools::mcp::EventsPeer::Http { .. } => Some( - self.http - .get_or_init(|| { + crate::tools::mcp::EventsPeer::Http { .. } => Some(match self.http.get() { + Some(client) => client.clone(), + None => { + let built = events_http_client(|| { // Before any builder call: `rustls-no-provider` panics without one. agent_core::ensure_provider(); reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) .build() - .unwrap_or_default() - }) - .clone(), - ), + .map_err(|e| e.to_string()) + })?; + self.http.get_or_init(|| built).clone() + } + }), crate::tools::mcp::EventsPeer::Rmcp { .. } => None, }; Ok(Conn { peer, http }) @@ -2574,6 +2587,14 @@ mod tests { assert!(!render_injection(1, &[gap(false)]).contains("possibly not an event")); } + /// The direct-HTTP events client fails closed: one that cannot be built is the request's error, + /// never a default client (which would follow redirects past the URL that was checked). + #[test] + fn an_events_http_client_that_cannot_be_built_fails_closed() { + let e = events_http_client(|| Err("injected build failure".into())).unwrap_err(); + assert!(e.contains("no request was made"), "{e}"); + } + /// Every injection's text names its batch on its first line — for anyone reading the /// transcript; delivery itself is tracked by tag (`AgentEvent::Steered`), not by text. #[test] diff --git a/crates/agent/src/tools/memory.rs b/crates/agent/src/tools/memory.rs index cd00200b..6b56d64d 100644 --- a/crates/agent/src/tools/memory.rs +++ b/crates/agent/src/tools/memory.rs @@ -157,6 +157,10 @@ fn render_hits(query: &str, hits: &[Hit]) -> ToolOutput { for h in hits { // Infallible: writing into a String never errors. let _ = writeln!(out, "{}:{}: {}", h.path, h.line, h.text); + // Past the cap the rest would only be cut off: stop (one line over is what marks it cut). + if out.len() > crate::tools::output::MAX_LISTING_BYTES { + break; + } } cap_listing_bytes(&mut out, "narrow the query"); out.into() @@ -312,11 +316,19 @@ impl Tool for Memory { let query = str_field(&input, "query")?; // Fan across every mount; each backend labels its hits with its own root, so the merged // listing is unambiguous. Durable first (mount order), then session. - let mut hits: Vec = Vec::new(); + let mut per_mount: Vec> = Vec::with_capacity(self.mounts.len()); for mount in &self.mounts { - hits.extend(mount.backend.search(query).await.map_err(map_err)?); + per_mount.push(mount.backend.search(query).await.map_err(map_err)?); } - Ok(render_hits(query, &hits)) + // Merging, rendering and freeing a large hit list is work for the blocking pool, + // not the runtime thread the session's other tasks are polled on. + let query = query.to_owned(); + tokio::task::spawn_blocking(move || { + let hits: Vec = per_mount.into_iter().flatten().collect(); + render_hits(&query, &hits) + }) + .await + .map_err(|e| ToolError::Execution(format!("memory search: {e}"))) } other => Err(ToolError::InvalidInput(format!( "unknown memory command `{other}` (expected view/create/str_replace/insert/delete/rename/search)" diff --git a/crates/agent/src/tools/output.rs b/crates/agent/src/tools/output.rs index 471fc535..2beb0802 100644 --- a/crates/agent/src/tools/output.rs +++ b/crates/agent/src/tools/output.rs @@ -236,6 +236,13 @@ const STALE_TEMP_FILE_AGE: Duration = Duration::from_secs(24 * 60 * 60); /// best-effort, never fails the caller" convention (see `mark_spill_broken`) — since a missed sweep /// just means the next one catches it. fn sweep_stale_temp_files(prefix: &str) { + #[cfg(test)] + if let Some(record) = SWEPT_ON.get() { + record + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push((prefix.to_owned(), std::thread::current().id())); + } let Ok(entries) = std::fs::read_dir(std::env::temp_dir()) else { return; }; @@ -263,6 +270,11 @@ fn sweep_stale_temp_files(prefix: &str) { } } +/// Which thread each sweep ran on, by prefix (tests only). +#[cfg(test)] +static SWEPT_ON: std::sync::OnceLock>> = + std::sync::OnceLock::new(); + /// Incrementally accumulates streaming output with bounded memory. /// /// Feed raw bytes with [`append`](Self::append); call [`finish`](Self::finish) when the stream ends; @@ -590,8 +602,17 @@ impl OutputAccumulator { return; } // Sweep abandoned spill files from earlier commands before adding a new one — see - // `sweep_stale_temp_files`'s doc comment for why nothing else ever cleans these up. - sweep_stale_temp_files(&self.temp_prefix); + // `sweep_stale_temp_files`'s doc comment for why nothing else ever cleans these up. It + // lists and stats the whole system temp directory (thousands of entries on a busy host, + // over NFS on some), so on a runtime it runs on the blocking pool, not on the runtime + // thread this command's output is being read on: best-effort, nothing waits for it. + match tokio::runtime::Handle::try_current() { + Ok(runtime) => { + let prefix = self.temp_prefix.clone(); + drop(runtime.spawn_blocking(move || sweep_stale_temp_files(&prefix))); + } + Err(_) => sweep_stale_temp_files(&self.temp_prefix), + } let path = std::env::temp_dir().join(format!("{}-{}.log", self.temp_prefix, random_hex16())); let mut opts = std::fs::OpenOptions::new(); @@ -1045,6 +1066,52 @@ mod tests { } } + /// On a runtime the stale-spill sweep — a listing and stat of the whole system temp directory — + /// runs on the blocking pool, not on the thread reading the command's output (a runtime thread + /// other sessions share); with no runtime it runs inline. + #[test] + fn the_stale_sweep_runs_off_the_runtime_thread() { + let swept = SWEPT_ON.get_or_init(Default::default); + let spill = |prefix: &str| { + let mut acc = OutputAccumulator::with_prefix(prefix); + for _ in 0..300u32 { + acc.append(&[b'x'; 1024]); + } + acc.finish(); + if let Some(path) = acc.snapshot(true).full_output_path { + let _ = std::fs::remove_file(path); + } + }; + let thread_of = |prefix: &str| { + swept + .lock() + .unwrap() + .iter() + .find(|(p, _)| p == prefix) + .map(|(_, t)| *t) + }; + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let runtime_thread = std::thread::current().id(); + rt.block_on(async { spill("test-sweep-on-runtime") }); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + let on = loop { + if let Some(t) = thread_of("test-sweep-on-runtime") { + break t; + } + assert!(std::time::Instant::now() < deadline, "the sweep never ran"); + std::thread::sleep(Duration::from_millis(5)); + }; + assert_ne!(on, runtime_thread, "swept on the runtime thread"); + drop(rt); + spill("test-sweep-inline"); + assert_eq!( + thread_of("test-sweep-inline"), + Some(std::thread::current().id()) + ); + } + #[test] fn ensure_temp_file_sweeps_stale_siblings_but_leaves_fresh_ones() { // Regression test for the safety-audit finding: nothing in the codebase ever deleted a diff --git a/crates/agent/src/tools/web/isolate.rs b/crates/agent/src/tools/web/isolate.rs index eae1b94a..e7357bbf 100644 --- a/crates/agent/src/tools/web/isolate.rs +++ b/crates/agent/src/tools/web/isolate.rs @@ -89,8 +89,20 @@ enum Response { /// exactly as they would in-process. A child that crashes, is killed by its seccomp filter, or times /// out is reported as an execution error naming the cause — never silently retried in-process, which /// would defeat the point of the isolation. -pub(super) fn parse(mode: Mode, input: &Value, html: &str) -> Result { - let exe = parser_binary()?; +/// +/// Blocking — it spawns the child and waits on it (up to [`PARSE_TIMEOUT`]) — so it is called on the +/// blocking pool, never on a runtime thread. `exe` overrides the parser binary (else +/// `BEYOND_AI_AGENT_WEB_PARSER`, else this executable). +pub(super) fn parse( + exe: Option<&std::path::Path>, + mode: Mode, + input: &Value, + html: &str, +) -> Result { + let exe = match exe { + Some(exe) => exe.to_path_buf(), + None => parser_binary()?, + }; let mut cmd = Command::new(exe); cmd.arg(SUBCOMMAND) diff --git a/crates/agent/src/tools/web/mod.rs b/crates/agent/src/tools/web/mod.rs index e1714ae8..06ff450a 100644 --- a/crates/agent/src/tools/web/mod.rs +++ b/crates/agent/src/tools/web/mod.rs @@ -52,9 +52,20 @@ pub struct Web { /// startup, for a tool a given session may never call. `OnceLock` keeps `new` trivial and builds the /// client exactly once, on demand. client: std::sync::OnceLock, + /// The binary that answers the isolated parse (`None`: the usual lookup, see `isolate::parse`). + parser: Option, + /// Builds the hardened client ([`build_client`]; a test swaps in one that fails). + build: fn(Duration, EgressPolicy) -> Result, } impl Web { + /// Builder-style: parse in this binary's `__web-parse` child instead of the usual lookup — for an + /// embedder (or a test) whose own executable is not the agent. + pub fn with_parser_binary(mut self, path: impl Into) -> Self { + self.parser = Some(path.into()); + self + } + /// Configure the tool. `allow_private`/`allow_hosts` come from `--web-allow-private`/ /// `--web-allow-host`; `timeout_ms` from `--web-timeout-ms` (else 30 s). The client is built lazily. pub fn new(allow_private: bool, allow_hosts: &[String], timeout_ms: Option) -> Self { @@ -64,31 +75,62 @@ impl Web { .map(Duration::from_millis) .unwrap_or(DEFAULT_TIMEOUT), client: std::sync::OnceLock::new(), + parser: None, + build: build_client, } } /// The owned client, built on first use. `redirect::Policy::none()` — we follow redirects ourselves /// so the SSRF URL check runs on each hop — plus the [`SsrfResolver`], so every connection (initial /// and any hop) validates the *resolved* IP. - fn client(&self) -> &reqwest::Client { - self.client.get_or_init(|| { - // Before *any* builder call: with reqwest's `rustls-no-provider`, a missing process-wide - // provider panics inside `build()` rather than returning `Err`, so the `unwrap_or_else` - // fallback below could never catch it (and its `Client::new()` would panic too). - agent_core::ensure_provider(); - reqwest::Client::builder() - .user_agent(USER_AGENT) - .connect_timeout(Duration::from_secs(10)) - .timeout(self.timeout) - .redirect(reqwest::redirect::Policy::none()) - .dns_resolver(SsrfResolver::new(self.policy.clone())) - .build() - // The builder only fails if the TLS backend can't initialize — fatal and identical for - // every call, so a fallback default client that will error consistently on use beats - // panicking. - .unwrap_or_else(|_| reqwest::Client::new()) - }) + /// + /// Built on the blocking pool: the first build initializes the TLS provider and loads the root + /// certificates (tens of milliseconds of CPU, and file reads on some hosts), which must not hold + /// the runtime thread this session shares with others. + /// + /// **Fails closed.** If the hardened client cannot be built (the builder errs, or panics), the + /// call fails: there is no fallback to a default client, which would have no SSRF resolver, no + /// timeout and would follow redirects itself. Nothing is cached, so a later call tries again. + async fn client(&self) -> Result<&reqwest::Client, ToolError> { + if let Some(client) = self.client.get() { + return Ok(client); + } + let (timeout, policy, build) = (self.timeout, self.policy.clone(), self.build); + let built = tokio::task::spawn_blocking(move || build(timeout, policy)) + .await + .map_err(|e| format!("building it panicked: {e}")) + .and_then(|built| built) + .map_err(|e| { + ToolError::Execution(format!( + "web: cannot build the hardened HTTP client, so no request was made: {e}" + )) + })?; + Ok(self.client.get_or_init(|| built)) } + + /// Builder-style: build the client with `build` instead (a test injects a failing one). + #[cfg(test)] + fn with_client_builder( + mut self, + build: fn(Duration, EgressPolicy) -> Result, + ) -> Self { + self.build = build; + self + } +} + +fn build_client(timeout: Duration, policy: EgressPolicy) -> Result { + // Before *any* builder call: with reqwest's `rustls-no-provider`, a missing process-wide + // provider panics inside `build()` rather than returning `Err`. + agent_core::ensure_provider(); + reqwest::Client::builder() + .user_agent(USER_AGENT) + .connect_timeout(Duration::from_secs(10)) + .timeout(timeout) + .redirect(reqwest::redirect::Policy::none()) + .dns_resolver(SsrfResolver::new(policy)) + .build() + .map_err(|e| e.to_string()) } /// One request/response round trip's outcome, before mode-specific rendering. @@ -119,7 +161,10 @@ impl Web { let start = Instant::now(); let mut redirects = 0; loop { - let mut rb = self.client().request(req.method.clone(), url.clone()); + let mut rb = self + .client() + .await? + .request(req.method.clone(), url.clone()); for (k, v) in &req.headers { rb = rb.header(k, v); } @@ -407,9 +452,31 @@ impl Tool for Web { let req = parse_request(&input)?; let fetched = self.fetch(&req).await?; - let text = match req.mode { + // Rendering is work over a body of up to megabytes, and every parsing mode spawns and waits on + // its isolated child (up to its timeout): all of it on the blocking pool, so the runtime + // thread this session shares with others is never held for it. + let mode = req.mode; + let input = input.clone(); + let parser = self.parser.clone(); + let text = + tokio::task::spawn_blocking(move || render(mode, &input, &fetched, parser.as_deref())) + .await + .map_err(|e| ToolError::Execution(format!("web: render task failed: {e}")))??; + Ok(text.into()) + } +} + +/// The tool's output for one fetched response. Blocking (see [`isolate::parse`]). +fn render( + mode: Mode, + input: &Value, + fetched: &Fetched, + parser: Option<&std::path::Path>, +) -> Result { + { + let text = match mode { // `fetch` always reports the status itself. - Mode::Fetch => render_fetch(&fetched), + Mode::Fetch => render_fetch(fetched), // The parsing modes work on the body regardless of status (an error page is still HTML), but // a model asking to `extract` from a 500 would otherwise get empty rows with no clue why — // so prepend the status when it isn't a success. @@ -425,7 +492,7 @@ impl Tool for Web { // Every parsing mode runs in a locked-down child: this is the one tool whose // input is bytes an arbitrary web server chose, and `scraper`/`htmd` reach a large // `unsafe` surface with them. See `isolate`'s module comment. - let parsed = isolate::parse(other, &input, &fetched.text())?; + let parsed = isolate::parse(parser, other, input, &fetched.text())?; let rendered = match other { Mode::Markdown => finalize(parsed, fetched.truncated), _ => parsed, @@ -434,7 +501,7 @@ impl Tool for Web { out } }; - Ok(text.into()) + Ok(text) } } @@ -443,6 +510,37 @@ mod tests { use super::*; use std::io::{Read, Write}; + /// Fails closed: when the hardened client cannot be built — the builder errs or panics — the + /// call fails and no request goes out at all (no fallback to a default client, which would have + /// no SSRF resolver, no timeout, and would follow redirects itself). + #[tokio::test] + async fn a_client_that_cannot_be_built_fails_the_call_and_sends_nothing() { + fn errs(_: Duration, _: EgressPolicy) -> Result { + Err("injected build failure".into()) + } + fn panics(_: Duration, _: EgressPolicy) -> Result { + panic!("injected build panic") + } + for build in [errs as fn(_, _) -> _, panics] { + let listener = beyond_ai_test_support::ports::listener(); + listener.set_nonblocking(true).unwrap(); + let url = format!( + "http://127.0.0.1:{}/", + listener.local_addr().unwrap().port() + ); + let tool = Web::new(true, &[], None).with_client_builder(build); + let e = tool + .run(serde_json::json!({ "url": url })) + .await + .unwrap_err(); + assert!(e.to_string().contains("no request was made"), "{e}"); + assert!( + matches!(listener.accept(), Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock), + "a request went out through an unhardened client" + ); + } + } + /// A one-shot loopback HTTP/1.1 server that replies with `response` to the first connection, and /// records the raw request line + headers it received. Returns `(port, join_handle_for_request)`. fn one_shot(response: &'static str) -> (u16, std::thread::JoinHandle) { diff --git a/crates/agent/tests/fs_backend_cost.rs b/crates/agent/tests/fs_backend_cost.rs index 8319211a..5bd35018 100644 --- a/crates/agent/tests/fs_backend_cost.rs +++ b/crates/agent/tests/fs_backend_cost.rs @@ -86,7 +86,7 @@ impl FsBackend for Counting { async fn write_if_unchanged( &self, path: &Path, - bytes: &[u8], + bytes: Vec, expected: Option, ) -> Result { self.tick(); diff --git a/crates/agent/tests/mcp_events_nested.rs b/crates/agent/tests/mcp_events_nested.rs index fc1391d4..eed02bd2 100644 --- a/crates/agent/tests/mcp_events_nested.rs +++ b/crates/agent/tests/mcp_events_nested.rs @@ -100,6 +100,10 @@ fn a_nested_elicitation_during_an_events_stream_reaches_the_owning_session() { /// In a daemon the server connection is shared by every session, so "this session" is not a given: /// the nested request raised during the events session's `events/poll` must reach a client attached /// to *that* session — not the process-wide host, and not another session's client. +/// +/// The server holds the request until the test says both clients are attached (`raise_nested`), +/// rather than for a guessed delay: under load an attach can take longer than any delay, and an +/// elicitation with no client attached is (rightly) declined at once. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn in_a_daemon_a_nested_elicitation_during_an_events_poll_reaches_the_events_session() { let home = tempfile::tempdir().unwrap(); @@ -109,7 +113,7 @@ async fn in_a_daemon_a_nested_elicitation_during_an_events_poll_reaches_the_even json!([stdio_server( "tickets", &control_file, - json!({ "MCP_FIXTURE_NESTED_DURING": "poll", "MCP_FIXTURE_NESTED_DELAY_MS": "1500" }), + json!({ "MCP_FIXTURE_NESTED_DURING": "poll", "MCP_FIXTURE_NESTED_ON_SIGNAL": "1" }), json!([{ "name": "ticket.updated", "delivery": "poll", "action": "notify" }]), )]), ); @@ -124,6 +128,20 @@ async fn in_a_daemon_a_nested_elicitation_during_an_events_poll_reaches_the_even // Another session's client is attached too; it must not be the one asked. let mut other = common::ws_connect(port, Some("bystander")).await; let mut ws = common::ws_connect(port, Some(common::mcp_events_fixture::EVENTS_SESSION)).await; + // Both attached for certain: each has answered a command on its own session. + for (client, id) in [(&mut other, "o"), (&mut ws, "e")] { + common::ws_send(client, json!({ "type": "get_state", "id": id })).await; + common::mcp_events_fixture::ws_next(client, Duration::from_secs(30), "attached", |f| { + f["type"] == "response" && f["id"] == id + }) + .await; + } + common::mcp_events_fixture::control( + &control, + "POST", + "/control/raise_nested", + Some(&json!({})), + ); let ask = common::mcp_events_fixture::ws_next( &mut ws, Duration::from_secs(30), diff --git a/crates/agent/tests/mcp_events_routing.rs b/crates/agent/tests/mcp_events_routing.rs index e92da6c9..4bdacb91 100644 --- a/crates/agent/tests/mcp_events_routing.rs +++ b/crates/agent/tests/mcp_events_routing.rs @@ -291,20 +291,27 @@ async fn a_permanently_refused_configured_subscription_is_reported_and_does_not_ let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); let (_d, port) = spawn_daemon(home.path(), &base, &["--session-idle-timeout", "1"]); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; - let status = ws_next(&mut ws, Duration::from_secs(20), "the refusal", |f| { - f["type"] == "mcp_event_status" && f["kind"] == "refused" - }) - .await; + // The refusal as the session records it, polled for: its `refused` frame goes out once, when + // the events session starts at boot — under load, before this client has attached. + let deadline = std::time::Instant::now() + Duration::from_secs(30); + let mut n = 0; + let l = loop { + n += 1; + let l = ws_list(&mut ws, &format!("l{n}")).await; + if l["data"]["unestablished"][0]["state"] == "refused" { + break l; + } + assert!(std::time::Instant::now() < deadline, "never refused: {l:#}"); + tokio::time::sleep(Duration::from_millis(100)).await; + }; + assert_eq!(l["data"]["unestablished"][0]["name"], "no.such.event"); assert!( - status["error"] + l["data"]["unestablished"][0]["last_error"] .as_str() .unwrap() .contains("offers no event"), - "{status:#}" + "{l:#}" ); - let l = ws_list(&mut ws, "l").await; - assert_eq!(l["data"]["unestablished"][0]["state"], "refused", "{l:#}"); - assert_eq!(l["data"]["unestablished"][0]["name"], "no.such.event"); let lists = || { state(&fixture)["methods"] .as_array() diff --git a/crates/agent/tests/tool_reactor_stall.rs b/crates/agent/tests/tool_reactor_stall.rs index 7ba79cc1..9ab4881b 100644 --- a/crates/agent/tests/tool_reactor_stall.rs +++ b/crates/agent/tests/tool_reactor_stall.rs @@ -1,7 +1,7 @@ // Test target: `.unwrap()` asserts preconditions; that's the point. #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] -//! A tool must not pin the async runtime for the duration of its file I/O and CPU work. +//! A tool must not pin the async runtime: not with its CPU work, and not by blocking the thread. //! //! `serve_ws` runs every session as a task on one **shared** runtime, so a tool that does its work //! inline — rather than handing it to `spawn_blocking` — stalls a worker that other sessions also @@ -9,27 +9,155 @@ //! stdin/WebSocket command loop that carries `abort` and `steer`. `serve.rs`'s `persist_blocking` //! already moved `sync_all` off the reactor for exactly this reason, and its doc comment says so. //! -//! **The probe measures CPU, not wall time.** Each tool call is polled on a `current_thread` -//! runtime, and the probe adds up the *thread CPU time* the runtime thread spends inside the tool -//! future's `poll`s: the time the tool holds the executor doing its own work. Work handed to -//! `spawn_blocking` is charged to a pool thread, not to this one. The bar is relative to a baseline -//! taken in the same run: the CPU the runtime thread spends doing the same work inline. An inline -//! tool holds the executor for about the whole baseline; a yielding one for a small fraction of it. +//! Two probes, each on a `current_thread` runtime, check two different things: //! -//! Thread CPU time does not advance while the thread is preempted. So a loaded host can make a run -//! slower, but cannot make a yielding tool look like a blocking one. A wall-clock probe (a ticker's -//! worst gap) could, and did: it counts every preemption anywhere in the call. -//! `the_probe_catches_a_tool_that_works_inline` keeps the probe's teeth honest. +//! - **CPU held** ([`held`]): the *thread CPU time* the runtime thread spends inside the tool +//! future's `poll`s, against a baseline of the same work done inline in the same run. It catches +//! a tool doing its CPU work on the executor. It cannot see a tool that blocks the thread without +//! using CPU — `std::thread::sleep`, blocking disk or NFS I/O, a contended `std` mutex — since a +//! blocked thread accrues no CPU time. +//! - **Responsiveness** ([`stalls`]): a ticker task on the same runtime, waking every millisecond +//! while the tool runs; the worst gap between its ticks is how long the runtime could not poll +//! anything else, by CPU work or by blocking alike. Each gap has the time the thread spent +//! *waiting on a run queue* (preempted: `/proc/thread-self/schedstat`) taken out, so a loaded host +//! cannot make a yielding tool look like a blocking one — the trap a plain wall-clock ticker fell +//! into before. What is left is time the thread ran, or blocked, without returning to the +//! executor; it must stay under [`MAX_UNRESPONSIVE`]. +//! +//! Every tool in `beyond_ai_agent::tools` that runs without an external service is put through the +//! responsiveness probe with a workload large enough that inline work would show +//! (`*_keeps_the_runtime_responsive`); the CPU probe covers the file tools. MCP tools are +//! network clients over rmcp's async transports, and `subagent` runs a nested agent loop on the same +//! runtime: neither is probed here. `the_probe_catches_a_tool_that_works_inline` and +//! `the_responsiveness_probe_catches_a_tool_that_blocks_without_cpu` keep both probes' teeth +//! honest. use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use std::time::Duration; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + use agent_core::Tool; -use beyond_ai_agent::tools::{edit, ls, write}; +use beyond_ai_agent::tools::{bash, edit, find, grep, ls, memory, read, todo, web, write}; use serde_json::json; +/// The longest the runtime may go without polling anything else while a tool runs, preemption +/// excluded. A tool's own poll (parsing its input, handing work off, assembling the result) takes +/// well under a millisecond; every workload below would block for far longer if done inline. +const MAX_UNRESPONSIVE: Duration = Duration::from_millis(20); + +/// How long this thread has waited on a run queue (preempted, runnable but not running) so far — +/// the second field of `/proc/thread-self/schedstat`. `None` where the kernel does not report it +/// (off Linux, or without scheduler statistics): the responsiveness probe is then **skipped**, not +/// run on plain wall clock, which a loaded host would make flake. +fn run_queue_wait() -> Option { + std::fs::read_to_string("/proc/thread-self/schedstat") + .ok() + .and_then(|s| s.split_whitespace().nth(1)?.parse::().ok()) + .map(Duration::from_nanos) +} + +/// One ticker sample: when, and the thread's run-queue wait so far (if the kernel reports it). +type Sample = (Instant, Option); + +/// The worst gap between consecutive ticker samples with the run-queue wait between them taken out; +/// `None` when any sample had no run-queue figure (see [`run_queue_wait`]). +fn worst_stretch(samples: &[Sample]) -> Option { + samples + .windows(2) + .map(|w| { + let (waited_then, waited_now) = (w[0].1?, w[1].1?); + Some((w[1].0 - w[0].0).saturating_sub(waited_now.saturating_sub(waited_then))) + }) + .try_fold(Duration::ZERO, |worst, gap| Some(worst.max(gap?))) +} + +/// Why a responsiveness probe did not run, said once per probe. +fn skip_without_schedstat(tool: &str) { + eprintln!( + "SKIP {tool}: /proc/thread-self/schedstat is unavailable here, so the responsiveness probe \ + cannot take preemption out of its measurement and is not run (the CPU probe still is)" + ); +} + +/// Held for each test's whole run, setup included. Under `cargo test` the tests share one process, +/// and one test's large allocations and frees (building a tree of thousands of files, a 4 MB source) +/// make another's runtime thread wait in the kernel on the shared address space — time that is +/// neither preemption nor the tool's, and would read as a stall. (`nextest`, which CI uses, runs each +/// test in its own process anyway.) +static ONE_TEST_AT_A_TIME: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +async fn one_at_a_time() -> tokio::sync::MutexGuard<'static, ()> { + ONE_TEST_AT_A_TIME.lock().await +} + +/// The worst stretch, preemption excluded, that the runtime thread went without polling a ticker +/// task while `work` ran on the same `current_thread` runtime; `None` without schedstat. +async fn stalls(work: F) -> (F::Output, Option) { + let samples: Arc>> = Arc::default(); + let stop = Arc::new(AtomicBool::new(false)); + let ticker = tokio::spawn({ + let samples = samples.clone(); + let stop = stop.clone(); + async move { + while !stop.load(Ordering::Relaxed) { + samples + .lock() + .unwrap() + .push((Instant::now(), run_queue_wait())); + tokio::time::sleep(Duration::from_millis(1)).await; + } + } + }); + // Let the ticker take its first sample before the work starts. + tokio::task::yield_now().await; + let out = work.await; + samples + .lock() + .unwrap() + .push((Instant::now(), run_queue_wait())); + stop.store(true, Ordering::Relaxed); + ticker.abort(); + let worst = worst_stretch(&samples.lock().unwrap()); + (out, worst) +} + +fn assert_responsive(tool: &str, worst: Duration, why: &str) { + eprintln!("PROBE {tool}: worst unresponsive stretch {worst:?}"); + assert!( + worst < MAX_UNRESPONSIVE, + "`{tool}` kept the current_thread runtime from polling anything else for {worst:?} \ + (preemption excluded): every other task on this worker — other sessions, the abort/steer \ + command loop — waited that long. {why}" + ); +} + +/// [`assert_responsive`] on the least of up to three runs of `$attempt` (an expression evaluated +/// afresh for each run, so any per-run setup in it happens before the measurement starts). A tool +/// that works or blocks inline stalls the runtime on every run; the host's own one-off stalls of +/// the runtime thread (a page fault served from disk on a busy machine) do not repeat. +macro_rules! assert_responsive_least { + ($tool:expr, $why:expr, $attempt:expr) => {{ + let mut least = Duration::MAX; + for _ in 0..3 { + let (_, worst) = stalls($attempt).await; + let Some(worst) = worst else { + skip_without_schedstat($tool); + return; + }; + least = least.min(worst); + if least < MAX_UNRESPONSIVE { + break; + } + } + assert_responsive($tool, least, $why); + }}; +} + /// CPU time consumed by the calling thread so far. fn thread_cpu() -> Duration { let ts = rustix::time::clock_gettime(rustix::time::ClockId::ThreadCPUTime); @@ -136,6 +264,7 @@ const NEW: &str = " let x_40001 = compute(i, 40001) + adjust(); // edited"; #[tokio::test(flavor = "current_thread")] async fn edit_does_not_stall_the_runtime() { + let _serial = one_at_a_time().await; let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("subject.rs"); // An inline `edit` spends tens of ms here (read + normalize + match + splice + write). @@ -171,6 +300,7 @@ async fn edit_does_not_stall_the_runtime() { #[tokio::test(flavor = "current_thread")] async fn write_does_not_stall_the_runtime() { + let _serial = one_at_a_time().await; let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("out.rs"); let src = big_ascii_source(SUBJECT_LINES); @@ -210,6 +340,7 @@ fn wide_dir() -> tempfile::TempDir { #[tokio::test(flavor = "current_thread")] async fn ls_does_not_stall_the_runtime() { + let _serial = one_at_a_time().await; let dir = wide_dir(); let baseline = inline_cpu(|| assert_eq!(scan(dir.path()), 5_000)); @@ -234,6 +365,7 @@ async fn ls_does_not_stall_the_runtime() { /// whole baseline, and never yields, so `assert_yields` would fail it. #[tokio::test(flavor = "current_thread")] async fn the_probe_catches_a_tool_that_works_inline() { + let _serial = one_at_a_time().await; let dir = wide_dir(); let baseline = inline_cpu(|| assert_eq!(scan(dir.path()), 5_000)); let held = held(async { scan(dir.path()) }).await; @@ -244,3 +376,438 @@ async fn the_probe_catches_a_tool_that_works_inline() { held.cpu ); } + +/// The responsiveness probe has teeth where the CPU probe has none: a tool that yields once and +/// then blocks the thread without using CPU (here `std::thread::sleep`, standing in for blocking +/// disk or NFS I/O or a contended `std` mutex) holds almost no CPU, so the CPU probe passes it — +/// and the runtime still could not poll anything for the whole sleep. +#[tokio::test(flavor = "current_thread")] +async fn the_responsiveness_probe_catches_a_tool_that_blocks_without_cpu() { + let _serial = one_at_a_time().await; + let blocking = || async { + tokio::task::yield_now().await; + std::thread::sleep(Duration::from_millis(150)); + }; + let cpu = held(blocking()).await; + assert!(cpu.pendings > 0); + assert!( + cpu.cpu < Duration::from_millis(20), + "the CPU probe cannot see it: {:?}", + cpu.cpu + ); + let ((), worst) = stalls(blocking()).await; + let Some(worst) = worst else { + return skip_without_schedstat("the probe's own teeth"); + }; + assert!( + worst >= Duration::from_millis(100), + "the responsiveness probe must: worst stretch {worst:?}" + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn edit_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("subject.rs"); + let src = big_ascii_source(80_000); + let tool = edit::Edit::new(dir.path()); + let p = path.to_str().unwrap().to_string(); + assert_responsive_least!("edit", "Its file I/O belongs on `spawn_blocking`.", { + // Fresh for each run (an edit changes the file), and before the measurement starts. + std::fs::write(&path, &src).unwrap(); + async { + tool.run(json!({ "path": p, "old_string": OLD, "new_string": NEW })) + .await + .unwrap() + } + }); +} + +#[tokio::test(flavor = "current_thread")] +async fn write_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("out.rs"); + let input = json!({ "path": path.to_str().unwrap(), "content": big_ascii_source(80_000) }); + let tool = write::Write::new(dir.path()); + assert_responsive_least!("write", "Its file I/O belongs on `spawn_blocking`.", { + let input = input.clone(); + async { tool.run(input).await.unwrap() } + }); +} + +#[tokio::test(flavor = "current_thread")] +async fn ls_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let dir = wide_dir(); + let d = dir.path().to_str().unwrap().to_string(); + assert_responsive_least!( + "ls", + "Its directory scan belongs on `spawn_blocking`.", + async { + ls::Ls::default() + .run(json!({ "path": d, "limit": 5000 })) + .await + .unwrap() + } + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn read_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("big.rs"); + std::fs::write(&path, big_ascii_source(80_000)).unwrap(); + let p = path.to_str().unwrap().to_string(); + let tool = read::Read::new(dir.path()); + assert_responsive_least!("read", "Its file I/O belongs on `spawn_blocking`.", async { + tool.run(json!({ "path": p, "offset": 70_000, "limit": 2000 })) + .await + .unwrap() + }); +} + +/// A tree of `files` source files, for the search tools. +fn source_tree(files: usize) -> tempfile::TempDir { + let dir = tempfile::tempdir().unwrap(); + for i in 0..files { + let sub = dir.path().join(format!("d{}", i % 50)); + std::fs::create_dir_all(&sub).unwrap(); + std::fs::write(sub.join(format!("f{i}.rs")), big_ascii_source(200)).unwrap(); + } + dir +} + +#[tokio::test(flavor = "current_thread")] +async fn grep_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let dir = source_tree(2_000); + let tool = grep::Grep::new(dir.path()); + assert_responsive_least!( + "grep", + "Its walk and match belong on `spawn_blocking`.", + async { + tool.run(json!({ "pattern": "x_199 = compute", "limit": 5000 })) + .await + .unwrap() + } + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn find_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let dir = source_tree(2_000); + let tool = find::Find::new(dir.path()); + assert_responsive_least!("find", "Its walk belongs on `spawn_blocking`.", async { + tool.run(json!({ "pattern": "**/*.rs", "limit": 5000 })) + .await + .unwrap() + }); +} + +#[tokio::test(flavor = "current_thread")] +async fn bash_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let dir = tempfile::tempdir().unwrap(); + let tool = bash::Bash::real().with_root(dir.path()); + assert_responsive_least!( + "bash", + "Waiting on the child, reading its output and spilling it must not block the runtime.", + async { + let out = tool + .run(json!({ "command": "sleep 0.3; seq 1 20000" })) + .await + .unwrap(); + assert!(out.text.contains("20000"), "{}", out.text); + } + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn todo_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let tool = todo::Todo::new(); + let todos: Vec<_> = (0..500) + .map(|i| json!({ "content": format!("task {i}"), "activeForm": format!("doing task {i}"), "status": "pending" })) + .collect(); + assert_responsive_least!( + "todo", + "It is in-memory; nothing here should block.", + async { tool.run(json!({ "todos": todos })).await.unwrap() } + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn memory_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let dir = tempfile::tempdir().unwrap(); + // Every line a hit: a search returning hundreds of thousands of lines, whose rendering (and + // freeing) inline would hold the runtime for tens of milliseconds. + for i in 0..2_000 { + std::fs::write( + dir.path().join(format!("note_{i}.md")), + format!( + "# note {i}\n{}", + "some remembered fact about the project\n".repeat(200) + ), + ) + .unwrap(); + } + let backend = Arc::new(beyond_ai_agent::memory::file::FileBackend::at( + dir.path().to_path_buf(), + )); + std::fs::write(dir.path().join("big.md"), big_ascii_source(80_000)).unwrap(); + let tool = memory::Memory::new(backend); + assert_responsive_least!( + "memory search", + "Its file I/O, and rendering a large hit list, belong on `spawn_blocking`.", + async { + tool.run(json!({ "command": "search", "query": "remembered fact" })) + .await + .unwrap() + } + ); + assert_responsive_least!( + "memory view", + "Its file I/O, and rendering a large document, belong on `spawn_blocking`.", + async { + tool.run(json!({ "command": "view", "path": "/memories/big.md" })) + .await + .unwrap() + } + ); +} + +/// A local HTTP server answering every request with `body` as HTML. +async fn html_server(body: String) -> String { + let listener = beyond_ai_test_support::ports::listener(); + listener.set_nonblocking(true).unwrap(); + let addr = listener.local_addr().unwrap(); + std::thread::spawn(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + rt.block_on(async move { + let listener = tokio::net::TcpListener::from_std(listener).unwrap(); + while let Ok((mut stream, _)) = listener.accept().await { + let body = body.clone(); + tokio::spawn(async move { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf).await; + let head = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + let _ = stream.write_all(head.as_bytes()).await; + let _ = stream.write_all(body.as_bytes()).await; + }); + } + }); + }); + format!("http://{addr}/") +} + +/// One `markdown` fetch of `url`, which must have rendered the page. +async fn web_markdown(tool: &web::Web, url: &str) { + let out = tool + .run(json!({ "url": url, "mode": "markdown" })) + .await + .unwrap(); + assert!( + out.text.contains("paragraph"), + "{}", + &out.text[..out.text.len().min(300)] + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn web_keeps_the_runtime_responsive() { + let _serial = one_at_a_time().await; + let mut html = String::from(""); + for i in 0..20_000 { + html.push_str(&format!( + "

paragraph {i} with a link

" + )); + } + html.push_str(""); + let url = html_server(html).await; + // The isolated parse runs in the agent binary's `__web-parse` child, not in this test binary. + let fresh = + || web::Web::new(true, &[], None).with_parser_binary(env!("CARGO_BIN_EXE_beyond-ai-agent")); + let why = "Fetching is async; building the client, and parsing in (and waiting on) the \ + isolated child, belong on the blocking pool."; + // A first call builds the client (TLS provider, root certificates): a fresh tool each run. + let tools: Vec = (0..3).map(|_| fresh()).collect(); + let mut cold = tools.iter(); + assert_responsive_least!( + "web (first call)", + why, + web_markdown(cold.next().unwrap(), &url) + ); + // Every call parses the page in an isolated child and waits on it. + let warm = fresh(); + web_markdown(&warm, &url).await; + assert_responsive_least!("web", why, web_markdown(&warm, &url)); +} + +/// `execute` (Code Mode) runs QuickJS's synchronous evaluation: a busy script must hold a +/// blocking-pool thread, never the runtime thread — whatever else this worker runs (other sessions, +/// their progress and aborts) keeps being polled while it spins. +#[cfg(feature = "code-mode")] +#[tokio::test(flavor = "current_thread")] +async fn code_mode_keeps_the_runtime_responsive_while_a_script_spins() { + let _serial = one_at_a_time().await; + let tool = beyond_ai_agent::tools::code_mode::Execute::new(vec![]); + let spin_of = |n: u64| json!({ "code": format!("let x = 0; for (let i = 0; i < {n}; i++) {{ x += i % 7; }} return x;") }); + // Long enough to show, whatever this build's QuickJS speed: doubled until a run takes 200 ms. + let mut n = 100_000u64; + let ran = loop { + let started = Instant::now(); + tool.run(spin_of(n)).await.unwrap(); + let ran = started.elapsed(); + if ran >= Duration::from_millis(200) || n >= 1 << 30 { + break ran; + } + n *= 2; + }; + eprintln!("PROBE execute spins {n} iterations in {ran:?}"); + let spin = spin_of(n); + assert_responsive_least!( + "execute", + "QuickJS evaluates synchronously; it belongs on the blocking pool.", + async { tool.run(spin.clone()).await.unwrap() } + ); + let held = held(async { tool.run(spin).await.unwrap() }).await; + assert!( + held.cpu < Duration::from_millis(20), + "`execute` spent {:?} of the runtime thread's CPU on a script that spins for {ran:?}", + held.cpu + ); +} + +/// An abort still stops a spinning script promptly — by its cancellation token, or by the call +/// being dropped — and frees the process's one JS slot, so the next `execute` runs at once rather +/// than after the 30 s deadline. +#[cfg(feature = "code-mode")] +#[tokio::test(flavor = "current_thread")] +async fn code_mode_abort_interrupts_a_spinning_script_promptly() { + use agent_core::tool::ToolProgress; + let _serial = one_at_a_time().await; + let tool = beyond_ai_agent::tools::code_mode::Execute::new(vec![]); + let forever = json!({ "code": "while (true) {}" }); + let quick = json!({ "code": "return 'next';" }); + + // Cancelled through the call's token, from another task on this runtime. + let token = tokio_util::sync::CancellationToken::new(); + let (tx, _rx) = futures::channel::mpsc::unbounded(); + let progress = ToolProgress::new(tx, "x".into(), "execute".into(), token.clone()); + let canceller = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(200)).await; + token.cancel(); + }); + let started = Instant::now(); + let err = tool + .run_streaming(forever.clone(), &progress) + .await + .unwrap_err(); + assert!(err.to_string().contains("cancelled"), "{err}"); + assert!( + started.elapsed() < Duration::from_secs(5), + "{:?}", + started.elapsed() + ); + canceller.await.unwrap(); + let next = tokio::time::timeout(Duration::from_secs(5), tool.run(quick.clone())) + .await + .expect("the slot is free again: the cancelled script was interrupted") + .unwrap(); + assert_eq!(next.text, "next"); + + // Dropped mid-run, as an abort that drops the tool future does. + let dropped = tokio::time::timeout(Duration::from_millis(200), tool.run(forever)).await; + assert!(dropped.is_err(), "still spinning when dropped"); + let next = tokio::time::timeout(Duration::from_secs(5), tool.run(quick)) + .await + .expect("the slot is free again: the dropped script was interrupted") + .unwrap(); + assert_eq!(next.text, "next"); +} + +/// All the CPU this (`current_thread`) runtime's thread spends while `work` runs — in `work`'s own +/// polls and in any task it spawns onto this runtime — so per-byte work moved into a spawned task +/// is seen too. +async fn runtime_thread_cpu(work: F) -> Duration { + let start = thread_cpu(); + work.await; + thread_cpu() - start +} + +/// `edit`'s cost on the runtime thread does not grow with the file: validating it as UTF-8 and +/// handing its new contents to the writer are per-byte work, done on the blocking pool (the writer +/// takes the buffer, it does not copy it here). +#[tokio::test(flavor = "current_thread")] +async fn edit_does_no_per_byte_work_on_the_runtime_thread() { + let _serial = one_at_a_time().await; + let dir = tempfile::tempdir().unwrap(); + let tool = edit::Edit::new(dir.path()); + let mut cost = Vec::new(); + for lines in [80_000usize, 640_000] { + let path = dir.path().join(format!("subject_{lines}.rs")); + let src = big_ascii_source(lines); + let p = path.to_str().unwrap().to_string(); + // The least of three, each on a fresh copy: per-byte work shows on every run. + let mut least = Duration::MAX; + for _ in 0..3 { + std::fs::write(&path, &src).unwrap(); + let spent = runtime_thread_cpu(async { + tool.run(json!({ "path": p, "old_string": "let x_40001 = compute(i, 40001)", "new_string": "let x_40001 = compute(i, 40001) /* edited */" })) + .await + .unwrap(); + }) + .await; + least = least.min(spent); + } + eprintln!( + "PROBE edit {} MB: runtime thread {least:?}", + src.len() >> 20 + ); + cost.push(least); + } + assert!( + cost[1] < cost[0] + Duration::from_micros(600), + "an 8x larger file cost the runtime thread {:?} against {:?}: per-byte work is on it", + cost[1], + cost[0] + ); +} + +/// Without scheduler statistics the responsiveness probe reports "unavailable" (and its tests +/// skip, saying so) rather than measuring plain wall clock, which counts every preemption and would +/// make a loaded host's run flake. +#[test] +fn the_responsiveness_probe_is_unavailable_without_schedstat_not_wall_clock() { + let t = Instant::now(); + let later = t + Duration::from_millis(50); + assert_eq!( + worst_stretch(&[(t, None), (later, None)]), + None, + "no run-queue figure: no measurement" + ); + assert_eq!( + worst_stretch(&[(t, Some(Duration::ZERO)), (later, None)]), + None + ); + // With it, the run-queue wait is taken out. + assert_eq!( + worst_stretch(&[ + (t, Some(Duration::from_millis(1))), + (later, Some(Duration::from_millis(31))) + ]), + Some(Duration::from_millis(20)) + ); +}