diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 84aeb5da..9238aac1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -64,12 +64,13 @@ jobs: # groups are stable, reproducible locally with the same `-E` expression, and a red shard names # the area outright. # - # agent-serve is deliberately not in this matrix. On this PR every harness for - # `binary(~serve)` compiled in ~2–4m then froze the runner: the test step stayed - # `in_progress` past job/step timeouts, logs never uploaded, and `needs.test` - # became `cancelled` so Check went red. The same 32 binaries pass under - # sequential `cargo test` locally in ~4 minutes. Leaving them in the required - # fan-in only reproduces the 30:00 death. + # agent-serve was left out of this matrix from #39 on, after its harnesses froze the runner (the + # test step stuck `in_progress` past every timeout). Out of CI they rotted unseen: by 2026-10 48 + # of them failed or hung on main, almost all because the session-title request `serve` began + # making after a first run took a reply from the in-order mock model's script, shifting every + # later turn by one. With that fixed in the harness, and the runner-freezing causes since + # addressed (systemd-scoped test step, nextest terminate-after and global-timeout below), they + # run again — as their own shard, so a hang there names itself instead of hiding another's. name: test (${{ matrix.name }}) runs-on: ubuntu-latest timeout-minutes: 25 @@ -103,6 +104,9 @@ jobs: - name: agent-rest filter: package(beyond-ai-agent) and kind(test) and not binary(~serve) and not binary(exec_endpoint_live) extra: "" + - name: agent-serve + filter: package(beyond-ai-agent) and kind(test) and binary(~serve) + extra: "" - name: core-providers filter: package(beyond-ai-agent-core) or package(beyond-ai-providers) extra: "" diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 446ad4bf..ac261866 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -492,8 +492,7 @@ The harness layers several capabilities over the bare tools + loop: empty/whitespace-only) `query` is a true no-op — the existing recency-sorted, unfiltered behavior is unchanged `list` already does, and resuming without an explicit session id picks the newest session matching the - **current cwd**, not just the globally newest. `list`/`list_all` sort by `updated_at` (most-recently- - active first), not `created_at` — matching pi's own session list (sorted by `modified`); a session + **current cwd**, not just the globally newest. `list`/`list_all` sort by `updated_at` (most-recently- active first; a same-second tie by id, descending — a generated id leads with its creation time, so the newest still comes first, and `offset` paging sees one order), not `created_at` — matching pi's own session list (sorted by `modified`); a session created long ago but touched again today must still surface at the top, not sink behind sessions born more recently but since abandoned. These four fields are `#[serde(skip)]` on `SessionMeta` itself (so a stale scan value can never leak into the on-disk header, which reuses the same @@ -3602,8 +3601,7 @@ context. Only host-authored records count: a skill-tool result (matched to its ` design the working group did not adopt, and is honoured too). Never on a timer, and only against a server that is up: starting a reaped one just to re-list would undo the reaping. A held entry past its `ttlMs` is re-fetched with `skills/get` before a load. A listing marked `cacheScope: "private"` - is never written to the on-disk manifest (that file outlives the authorization context); the server - is forgotten there instead, so the next boot asks it again. A re-list that fails keeps the last good + is never written to the on-disk manifest (that file outlives the authorization context); the server is forgotten there instead, so the next boot asks it again. A listing re-fetched after connect is written back to that manifest at once (`mcp_manifest::store_skills`, amending only the skills fields of the invocation's recorded manifest) — or, if it is now private, the server is forgotten — so a restart answering from the cache does not advertise the listing the server replaced. A re-list that fails keeps the last good listing; a notification is a generation count the listing catches up with only once a re-list succeeds, so one cut short by the refresh deadline (30s, `BEYOND_AI_AGENT_MCP_SKILLS_REFRESH_TIMEOUT_MS`) leaves the notice pending. @@ -3625,9 +3623,9 @@ context. Only host-authored records count: a skill-tool result (matched to its ` code execution without the user's explicit, per-skill approval: - _Activation._ A skill the model chose is loaded only once the user approves it — before `SKILL.md` is fetched, bound to the entry's manifest fingerprint (every `{uri, digest}`), so a - changed skill is asked about again. A `/skill:` the user typed is that consent. A nested skill (its + changed skill is asked about again. A `/skill:` the user typed is that consent — for that invocation only, never remembered, so a later load the model chooses is asked about. A nested skill (its directory inside another skill's) is no exception: its approval is its own, never covered by the - enclosing skill's, and the question names the enclosing skill. + enclosing skill's, and the question names the enclosing skill. The question names the agent that chose the skill (`origin`: `"main"`, or the subagent's `{agent, spawn_id}`): a subagent's registry binds the loading tools in its own name (`filter_by_enabled_as`), the same provenance its code-execution questions carry. - _Code execution._ While the session is acting on any MCP skill (from load to the end of the session, and again whenever its transcript still holds the skill — the spec allows longer, never shorter), every `bash`/`execute` call needs approval, asked per call and naming the active skills; @@ -3644,7 +3642,8 @@ context. Only host-authored records count: a skill-tool result (matched to its ` before it is fetched; `{purpose: "execute", active_skills}`; or `{purpose: "cross_origin_read", from_servers, to_server}`), answered by the usual `approve` command. `run` has no one to ask and denies; `--approve-mcp-skills` (both commands) approves - everything in advance. + everything in advance. Without `--approve`, the only questions a `serve` session raises are MCP skills', so an `approve` naming none of its pending questions is refused (pointing at `--approve`) rather than acknowledged. + - _Not persisted._ Remembered ("session"-scoped) decisions live in memory for the session's life in this process and are never written down or read back. Any approval on disk would be forgeable: the session file is the user's, and a model with file-write tools (`write`/`edit`) can append to it — and it knows the fingerprint a key embeds, which is in the loaded skill's tag — so a restored approval could be one the model planted, and an `execute` approval would then run code unasked after a restart. A resumed, switched-to or restarted session is therefore asked again. The `mcp_skill_approval` kind stays reserved (`session_store::HOST_CUSTOM_KINDS`, with the MCP task journal's kinds): a client's `append_custom` of a kind the agent writes itself is refused, so it cannot forge one (`tests/mcp_skills_resume.rs`, `tests/serve_session_tree.rs`). - `allowed-tools` and other declarative frontmatter are not honoured for any skill, so nothing an MCP skill declares can widen permissions. - **Names are per origin**: `:`, or `:` for two entries of one @@ -3672,7 +3671,7 @@ context. Only host-authored records count: a skill-tool result (matched to its ` its `skill://` resources are ordinary resources. - **Manifest cache.** `mcp_manifest` (version 3) records the entries (`None` = extension not declared) and their diagnostics, so a reaped or never-started server still lists its skills; the - first load dials it (and re-fetches the entry, whose freshness a cache cannot vouch for). + first load dials it (and re-fetches the entry, whose freshness a cache cannot vouch for). Every write is a read-modify-write of one file shared by all connections (a connect, a forget, a refreshed listing), so each runs under a lockfile (`mcp_manifest::FileLock`, the auth and settings stores' pattern) and renames a temporary of its own into place; concurrent refreshes lose nobody's update. - **Nested-request attribution.** Every skills request (`skills/list`, `skills/get`, `resources/read`) goes through `mcp::request_tracked`, registered with `track_call` under the calling session's host for as long as it is outstanding, so a nested request the server raises @@ -4091,7 +4090,7 @@ backend an operator points it at. | `src/bin/mcp_skills_fixture_server.rs` | Hand-rolled SEP-2640 server for `tests/mcp_skills*.rs` (stdio, or `--http`): an ordinary skill with a supporting file and a nested skill, a same-name pair, a `disable-model-invocation` skill, a tampered one, a dynamic one, an unlisted one, two over the per-skill limits, and a late one published by its `publish_late` tool (with `notifications/resources/list_changed`); caching hints from env; `_meta` on every result like a `2026-07-28` server; logs every request so tests assert what was (not) fetched | | `benches/search.rs` | Criterion macro-bench: `grep` (1 vs auto threads) and `find` (sequential) over a 5,000-file tree | | `benches/serve_runtime.rs` | Process-level Tokio runtime A/B for `serve --listen`: spawns the real binary under `current_thread` vs multi-thread (`BEYOND_AI_AGENT_TOKIO_WORKER_THREADS`) and drives concurrent WebSocket sessions, reporting RSS, throughput, p95 latency, and CPU. Custom harness, not Criterion/divan — the cost is thread stacks and per-thread heaps, which an in-process alloc profiler never sees. | -| `tests/common/mod.rs` | Shared test harness: mock Anthropic-SSE model server, `serve`/`run` command builders, `read_until_response`, gateway binary locator, port/connection helpers, process-lifetime JetStream for managed-gateway tests — used by every file below | +| `tests/common/mod.rs` | Shared test harness: mock Anthropic-SSE model server (an in-order server answers `serve`'s session-title request itself, off the conversation script and out of its record, counting them for `title_calls`; a test about titles matches `SESSION_TITLE_MARKER` on a routed one), ports with no window for another process to take them (`spawn_listening`: `serve --listen 127.0.0.1:0`, the bound port read from its announcement; `HeldPort`: a listener the test binds and hands over by socket activation, for a port named in `serve`'s own arguments; `DeadPort`: bound, never listening — `free_port` is left only for the gateway and `nats-server`), `frames_with_deadline` (a stdout reader that fails instead of hanging), `serve`/`run` command builders, `read_until_response`, gateway binary locator, port/connection helpers, process-lifetime JetStream for managed-gateway tests — used by every file below | | `tests/run_*.rs` | `run` binary against a mock model server (no gateway in the loop), split by domain: `run_core_flow` (tool round trips, refusal exit codes, text/json mode, stdin/@file input), `run_skills_prompts` (skill/prompt-template expansion, trust gating), `run_session_management` (`--session`/`--continue`/`--name`/`--fork`, same- and cross-project), `run_cli_flags` (`--version`/`--help`/`list-models`/`export`, flags reaching the wire request, including `models.json` header/api_key overrides, `--idle-timeout-ms`, `--block-images`), `run_stdout_robustness` (Task #10: a closed stdout pipe, simulated by dropping the child's stdout handle mid-stream, must exit 0 rather than panic on `EPIPE`, in both text and `--json` mode) | | `tests/serve_*.rs` | `serve` binary NDJSON protocol round-trip, split by domain: `serve_session_lifecycle` (startup/resume/crash recovery/jsonl framing/export), `serve_state_reporting` (`get_state`/`get_session_stats`/`cwd_stale`), `serve_session_tree` (branches/forks/`get_tree`), `serve_models_thinking` (model/thinking-level/`--models`), `serve_compaction_retry` (compaction + whole-run auto-retry), `serve_trust_prompts` (trust gating + system prompt), `serve_tools_bash` (tool exclusion + host `bash`), `serve_prompt_flow` (`prompt`/`steer`/`follow_up` queuing, busy semantics, refusal, abort), `serve_lifecycle` (host-emitted run lifecycle POSTs: started/progress/terminal) — each a separate Cargo test binary, ~400-1200 lines apiece rather than one ~6,200-line file | | `tests/gateway_e2e.rs` | `run` binary → real gateway binary → mock upstream (proves key-swap + the virtual key never reaches upstream). Boots a real nats-server so the fail-closed allowance watcher can seed | diff --git a/crates/agent/src/bin/mcp_skills_fixture_server.rs b/crates/agent/src/bin/mcp_skills_fixture_server.rs index 6cc9cd80..03c08296 100644 --- a/crates/agent/src/bin/mcp_skills_fixture_server.rs +++ b/crates/agent/src/bin/mcp_skills_fixture_server.rs @@ -41,6 +41,8 @@ //! reached. //! - `MCP_SKILLS_FIXTURE_TTL_MS` (default `0`) and `MCP_SKILLS_FIXTURE_CACHE_SCOPE` (default //! `public`): the caching hints on `skills/list` / `skills/get`. +//! - `MCP_SKILLS_FIXTURE_PRIVATE_FLAG=`: while that file exists, the scope is `private`, so a +//! test can turn a listing private after connect. //! //! Also served: `skill://git-workflow/nested-helper/SKILL.md`, a skill nested inside `git-workflow` //! (whose manifest lists its file, as completeness requires); and two skills over the spec's @@ -139,8 +141,14 @@ fn cache_hints() -> (u64, String) { .ok() .and_then(|v| v.parse().ok()) .unwrap_or(0); - let scope = - std::env::var("MCP_SKILLS_FIXTURE_CACHE_SCOPE").unwrap_or_else(|_| "public".to_string()); + let private_now = std::env::var("MCP_SKILLS_FIXTURE_PRIVATE_FLAG") + .is_ok_and(|p| std::path::Path::new(&p).exists()); + let scope = match private_now { + true => "private".to_string(), + false => { + std::env::var("MCP_SKILLS_FIXTURE_CACHE_SCOPE").unwrap_or_else(|_| "public".to_string()) + } + }; (ttl, scope) } diff --git a/crates/agent/src/lib.rs b/crates/agent/src/lib.rs index d7d28792..45b951c4 100644 --- a/crates/agent/src/lib.rs +++ b/crates/agent/src/lib.rs @@ -33,6 +33,8 @@ pub mod service; pub mod session_store; pub mod settings; pub mod skills; +#[cfg(test)] +pub(crate) mod test_support; pub mod timing; pub mod tools; #[cfg(test)] diff --git a/crates/agent/src/resources.rs b/crates/agent/src/resources.rs index 85a4fada..081f309d 100644 --- a/crates/agent/src/resources.rs +++ b/crates/agent/src/resources.rs @@ -1832,7 +1832,7 @@ mod tests { #[tokio::test] async fn load_context_files_via_walks_home_then_every_ancestor_nearest_last() { - let tmp = tempfile::tempdir().unwrap(); + let tmp = crate::test_support::isolated_tempdir(); let home = tmp.path().join("home"); let repo = tmp.path().join("box/repo"); let ws = repo.join("app"); @@ -1853,7 +1853,7 @@ mod tests { #[tokio::test] async fn load_context_files_via_treats_an_empty_file_as_absent() { - let tmp = tempfile::tempdir().unwrap(); + let tmp = crate::test_support::isolated_tempdir(); let ws = tmp.path().join("ws"); fs::create_dir_all(&ws).unwrap(); fs::write(ws.join("AGENTS.md"), " \n").unwrap(); diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index ebcbe6f1..8857a05e 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -148,7 +148,9 @@ //! active tip and advance the tip to it → `data: {id}` (the new entry's id); `data` defaults to `{}` //! when omitted, `kind` identifies the shape of `data` to whatever produced it (this crate never //! interprets either) — `SessionStore::append_custom`, fully built and tested but previously -//! unreachable from any RPC command; same persistence-required error as `set_label` +//! unreachable from any RPC command; same persistence-required error as `set_label`. A kind the +//! agent writes itself (`session_store::HOST_CUSTOM_KINDS`) is refused, so a client cannot forge +//! one //! - `{type:"export_html", output_path?}` render the active session's transcript as a single //! self-contained HTML file → `data: {path}`. `output_path` defaults to a timestamped //! `session-.html` in the current directory; parent directories are created as @@ -2380,7 +2382,7 @@ pub(crate) async fn scan_session_dirs( .flat_map(|dir| crate::session_store::scan_session_dir_in(dir, &layout)) .collect(); let mut metas = crate::session_store::scan_listings_in(paths, &layout, &on_progress); - metas.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); + metas.sort_by(crate::session_store::by_recency); metas }) .await @@ -5065,7 +5067,7 @@ pub(crate) async fn serve_session( // must be answerable mid-run. `accepted:false` tells the losing // client in a multi-attach race that its answer arrived too late. "approve" => { - let _ = out_tx.send(handle_approve(cid, &c, &pending_approvals)); + let _ = out_tx.send(handle_approve(cid, &c, &pending_approvals, approval.is_some())); } "elicit" => { let _ = out_tx.send(handle_elicit(cid, &c, Some(&pending_elicitations))); @@ -6533,6 +6535,18 @@ pub(crate) async fn serve_session( )), }, "append_custom" => match cmd.get("kind").and_then(Value::as_str) { + // The host's own journal kinds: an entry of one is trusted when the session is + // replayed (a task to resume, a result to deliver), so a client must not be able to + // write one. + Some(kind) if crate::session_store::is_host_custom_kind(kind) => emit!(response( + id, + "append_custom", + false, + None, + Some(&format!( + "custom entry kind `{kind}` is reserved for entries the agent writes itself" + )) + )), Some(kind) => { let data = cmd.get("data").cloned().unwrap_or_else(|| json!({})); match persistence.append_custom(kind, data) { @@ -6576,7 +6590,12 @@ pub(crate) async fn serve_session( // Reachable while idle only for a stale/duplicate answer (`accepted:false`): a real question // can only be outstanding while a run is in flight, which is the busy arm above. "approve" => { - emit!(handle_approve(id, &cmd, &pending_approvals)); + emit!(handle_approve( + id, + &cmd, + &pending_approvals, + approval.is_some() + )); } "get_todos" => { // Straight from the session while idle — no mirror needed, and no chance of one going @@ -10253,7 +10272,17 @@ fn approval_resolved_frame( /// `accepted` is `false` when the `request_id` names no outstanding question — it was already answered /// by another attached client, it timed out, or the run was aborted. That is not an error: it is the /// answer a client races and loses. -fn handle_approve(id: Option, cmd: &Value, pending: &PendingApprovals) -> OutFrame { +/// +/// `tool_gate` is whether this session asks about tool calls (`--approve`). Without one, the only +/// questions it can raise are MCP skills' — so an answer that matches none of them is refused with a +/// pointer to `--approve` rather than acknowledged, since a client sending it is most likely expecting +/// tool approvals this session will never ask for. +fn handle_approve( + id: Option, + cmd: &Value, + pending: &PendingApprovals, + tool_gate: bool, +) -> OutFrame { let Some(request_id) = cmd.get("request_id").and_then(Value::as_str) else { return response(id, "approve", false, None, Some("missing `request_id`")); }; @@ -10262,6 +10291,18 @@ fn handle_approve(id: Option, cmd: &Value, pending: &PendingApprovals) - Err(e) => return response(id, "approve", false, None, Some(&e)), }; let accepted = resolve_approval(pending, request_id, decision); + if !accepted && !tool_gate { + return response( + id, + "approve", + false, + None, + Some( + "no such question: this session asks only about MCP skills; tool approvals need \ + `--approve`", + ), + ); + } response( id, "approve", @@ -12900,7 +12941,7 @@ mod tests { async fn git_branch_reports_none_outside_a_git_repository() { // Task #25 (pi-parity fix): a lookup failure (no repo here at all) must report `None`, never // an error surfaced to the RPC caller. - let dir = tempfile::tempdir().unwrap(); + let dir = crate::test_support::isolated_tempdir(); assert_eq!(git_branch(dir.path()).await, None); } diff --git a/crates/agent/src/session_store.rs b/crates/agent/src/session_store.rs index 311f9f1e..1c6c3afe 100644 --- a/crates/agent/src/session_store.rs +++ b/crates/agent/src/session_store.rs @@ -339,13 +339,40 @@ pub fn search_sessions(sessions: Vec, query: Option<&str>) -> Vec bool { + HOST_CUSTOM_KINDS.contains(&kind) +} + +/// Most recently active first; ties (same second — `updated_at` has one-second resolution, so a +/// burst of sessions shares one) broken by id, descending, so a listing has one order however the +/// directory happened to be scanned. Without the tie-break, `offset` paging over sessions written in +/// the same second could serve one session twice and skip another. Descending, because a generated +/// id ([`new_id`]) leads with its creation time in fixed-width hex nanoseconds: within the second the +/// newest still comes first. (A client-chosen id has no such meaning, but still one stable order.) +pub(crate) fn by_recency(a: &SessionMeta, b: &SessionMeta) -> std::cmp::Ordering { + b.updated_at + .cmp(&a.updated_at) + .then_with(|| b.id.cmp(&a.id)) +} + /// One persisted line. Internally tagged on `type`. A `Message` entry flattens the wrapped /// `agent_core::Message` (`{role, content}`) alongside its tree fields, so the wire shape is /// `{"type":"message", id, parent_id, role, content}` — `agent_core::Message` itself carries no tree @@ -2877,7 +2904,7 @@ impl SessionRepo { ); } let mut metas = scan_listings(paths, &on_progress); - metas.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); + metas.sort_by(by_recency); Ok(metas) } @@ -2933,7 +2960,7 @@ impl SessionRepo { } } let mut metas = scan_listings_in(paths, self.layout(), &on_progress); - metas.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); + metas.sort_by(by_recency); Ok(metas) } @@ -8831,6 +8858,32 @@ mod tests { m } + /// Sessions active in the same second (one-second `updated_at`) still have one order, and it is + /// newest-first: a generated id leads with its creation time, so the later one sorts first. With + /// no tie-break (or an ascending one) `offset` paging over them could repeat or skip a session. + #[test] + fn same_second_sessions_list_newest_first_in_one_stable_order() { + let ids: Vec = (0..4).map(|_| new_id()).collect(); + let mut metas: Vec = ids + .iter() + .map(|id| meta_for_search(id, None, "x", 100)) + .collect(); + metas.push(meta_for_search("older", None, "x", 99)); + // Oldest first in, whatever the scan order was. + metas.sort_by(by_recency); + let got: Vec<&str> = metas.iter().map(|m| m.id.as_str()).collect(); + let mut want: Vec<&str> = ids.iter().rev().map(String::as_str).collect(); + want.push("older"); + assert_eq!(got, want); + // The same order from any starting order. + metas.reverse(); + metas.sort_by(by_recency); + assert_eq!( + metas.iter().map(|m| m.id.as_str()).collect::>(), + want + ); + } + #[test] fn search_sessions_with_no_query_returns_the_input_unchanged_and_in_order() { let sessions = vec![ diff --git a/crates/agent/src/skills.rs b/crates/agent/src/skills.rs index af295437..4cb18297 100644 --- a/crates/agent/src/skills.rs +++ b/crates/agent/src/skills.rs @@ -1513,7 +1513,7 @@ mod tests { #[test] fn find_git_repo_root_returns_none_when_no_git_repo_encloses_it() { - let tmp = tempfile::tempdir().unwrap(); + let tmp = crate::test_support::isolated_tempdir(); let deep = tmp.path().join("a/b/c"); fs::create_dir_all(&deep).unwrap(); assert_eq!(find_git_repo_root(&deep), None); @@ -1541,7 +1541,7 @@ mod tests { #[test] fn collect_ancestor_agents_skill_dirs_walks_to_the_filesystem_root_when_no_git_repo_encloses_it() { - let tmp = tempfile::tempdir().unwrap(); + let tmp = crate::test_support::isolated_tempdir(); let start = tmp.path().join("a/b"); fs::create_dir_all(&start).unwrap(); let dirs = collect_ancestor_agents_skill_dirs(&start); diff --git a/crates/agent/src/test_support.rs b/crates/agent/src/test_support.rs new file mode 100644 index 00000000..c26f0b5b --- /dev/null +++ b/crates/agent/src/test_support.rs @@ -0,0 +1,67 @@ +//! Shared helpers for this crate's unit tests. + +use std::path::Path; + +/// What makes a directory "inside a project" to the code under test: an enclosing git repo, a +/// context file, or a project skills root anywhere above it. +const PROJECT_MARKERS: [&str; 4] = [".git", "CLAUDE.md", "AGENTS.md", ".agents"]; + +/// Whether `dir` or any ancestor carries a [project marker](PROJECT_MARKERS). +fn inside_a_project(dir: &Path) -> bool { + dir.ancestors() + .any(|a| PROJECT_MARKERS.iter().any(|m| a.join(m).exists())) +} + +/// A temp directory that no ancestor walk can escape into a real project. +/// +/// Tests of "no enclosing repo" / "no context file above" behaviour walk to the filesystem root, so +/// a plain `tempfile::tempdir()` makes them depend on where `TMPDIR` points: under a git checkout or +/// a `CLAUDE.md` (a worktree's own `target/tmp`, say) they find the surrounding project and fail. +/// This picks the first candidate root with no project above it — `TMPDIR` itself when it is clean. +/// +/// # Panics +/// When every candidate root sits inside a project, naming them, rather than letting the test +/// pass or fail for a reason that has nothing to do with it. +pub(crate) fn isolated_tempdir() -> tempfile::TempDir { + let candidates = [ + std::env::temp_dir(), + "/tmp".into(), + "/var/tmp".into(), + "/dev/shm".into(), + ]; + for root in &candidates { + let Ok(root) = root.canonicalize() else { + continue; + }; + if inside_a_project(&root) { + continue; + } + if let Ok(dir) = tempfile::tempdir_in(&root) { + return dir; + } + } + panic!( + "no temp root outside a project (git repo, CLAUDE.md, AGENTS.md or .agents above it) among \ + {candidates:?}" + ); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn an_isolated_tempdir_has_no_project_above_it() { + let dir = isolated_tempdir(); + assert!(!inside_a_project(dir.path()), "{}", dir.path().display()); + } + + #[test] + fn a_marker_anywhere_above_counts() { + let dir = isolated_tempdir(); + let deep = dir.path().join("a/b"); + std::fs::create_dir_all(&deep).unwrap(); + std::fs::write(dir.path().join("CLAUDE.md"), "x").unwrap(); + assert!(inside_a_project(&deep)); + } +} diff --git a/crates/agent/src/tools/mcp.rs b/crates/agent/src/tools/mcp.rs index e1ef9ea1..bb7f8fe4 100644 --- a/crates/agent/src/tools/mcp.rs +++ b/crates/agent/src/tools/mcp.rs @@ -646,6 +646,16 @@ pub fn server_name_from_registered(tool_name: &str) -> Option<&str> { /// A session that declared it renders MCP Apps sees its apps view's tools instead of `tools` — the /// same servers, dialed with the extension advertised (see [`crate::tools::mcp_apps`]). pub fn filter_by_enabled(tools: &[Arc], enabled: &McpEnabledSet) -> Vec> { + filter_by_enabled_as(tools, enabled, &crate::approval::ApprovalOrigin::Main) +} + +/// [`filter_by_enabled`] for the agent `origin` — a subagent's registry — so a skill load it makes +/// asks the user in that subagent's name. +pub fn filter_by_enabled_as( + tools: &[Arc], + enabled: &McpEnabledSet, + origin: &crate::approval::ApprovalOrigin, +) -> Vec> { // A declared session sees the apps flavor's tools for every server whose apps connection is up, // and the plain tools for every other one — a server whose apps dial failed (or has not landed // yet) keeps working as a text tool rather than vanishing. @@ -663,7 +673,12 @@ pub fn filter_by_enabled(tools: &[Arc], enabled: &McpEnabledSet) -> Ve Some(server) => enabled.allows(server), None => true, }) - .map(|t| enabled.skills.bind(t).unwrap_or_else(|| t.clone())) + .map(|t| { + enabled + .skills + .bind_as(t, origin) + .unwrap_or_else(|| t.clone()) + }) .collect() } @@ -961,6 +976,11 @@ impl McpConnection { .map(|live| live.client.clone()) } + /// The server's configuration — what its manifest is cached under. + pub(crate) fn config(&self) -> &McpServerConfig { + &self.config + } + /// The flag a list-changed notification raises on this connection, whichever dial received it. pub(crate) fn skills_changed(&self) -> Arc { self.dial.skills_changed.clone() @@ -2114,6 +2134,7 @@ fn tools_from_manifest( dial: &Dial, manifest: crate::tools::mcp_manifest::ServerManifest, idle_reap_after: Duration, + manifest_dir: &crate::tools::mcp_manifest::ManifestDir, ) -> (Vec>, McpServerCatalog) { let conn = Arc::new(McpConnection::dormant( config.clone(), @@ -2201,6 +2222,7 @@ fn tools_from_manifest( &conn, crate::tools::mcp_skills::Listing::cached(entries, diagnostics), &mut tools, + Some(manifest_dir), ) }); let catalog = McpServerCatalog { @@ -2222,13 +2244,21 @@ async fn connect_one( manifest_dir: Option<&crate::tools::mcp_manifest::ManifestDir>, ) -> Result<(Vec>, McpServerCatalog), String> { // Cache hit: advertise from the manifest and start nothing. - if let Some(manifest) = manifest_dir.and_then(|d| crate::tools::mcp_manifest::load(d, config)) { + if let Some(dir) = manifest_dir + && let Some(manifest) = crate::tools::mcp_manifest::load(dir, config) + { tracing::debug!( server = %config.name, tools = manifest.tools.len(), "advertising MCP tools from the cached manifest; not starting the server" ); - return Ok(tools_from_manifest(config, dial, manifest, idle_reap_after)); + return Ok(tools_from_manifest( + config, + dial, + manifest, + idle_reap_after, + dir, + )); } let (client, proc) = connect_one_client(config, dial).await?; tools_from_client(config, dial, client, proc, idle_reap_after, manifest_dir).await @@ -3173,8 +3203,9 @@ async fn tools_from_client( })); } - let skills = skill_listing - .map(|listing| crate::tools::mcp_skills::attach(&config.name, &conn, listing, &mut tools)); + let skills = skill_listing.map(|listing| { + crate::tools::mcp_skills::attach(&config.name, &conn, listing, &mut tools, manifest_dir) + }); Ok(( tools, McpServerCatalog { diff --git a/crates/agent/src/tools/mcp_manifest.rs b/crates/agent/src/tools/mcp_manifest.rs index 031b0e06..1995064d 100644 --- a/crates/agent/src/tools/mcp_manifest.rs +++ b/crates/agent/src/tools/mcp_manifest.rs @@ -26,7 +26,9 @@ //! longer exists; it cannot make a call silently do the wrong thing. use std::collections::BTreeMap; -use std::path::PathBuf; +use std::io::ErrorKind; +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant, SystemTime}; use serde::{Deserialize, Serialize}; @@ -175,14 +177,53 @@ pub fn load(dir: &ManifestDir, config: &McpServerConfig) -> Option bool) { + let Ok(_lock) = FileLock::acquire(&dir.file()) else { + return; + }; + let mut all = read_store(dir); + if change(&mut all) { + write_store(dir, &all); + } +} + /// Drop `config`'s cached manifest, if any — for a server whose answer must not be cached here (a /// skills listing marked `cacheScope: "private"`). Best-effort, like [`store`]. pub fn forget(dir: &ManifestDir, config: &McpServerConfig) { - let mut all = read_store(dir); - if all.remove(&config.name).is_none() { + update_store(dir, |all| all.remove(&config.name).is_some()); +} + +/// Replace only the skills listing in `config`'s cached manifest — for a listing re-fetched after +/// connect (`ttlMs` ran out, or the server said it changed), which would otherwise reach the cache +/// only at the next live connect, leaving a restart to advertise the stale one. A `private` listing +/// (`cacheScope: "private"`) forgets the server instead, as at connect. Without a manifest already +/// recorded for this exact invocation there is nothing to amend: the next live connect writes a +/// whole one. Best-effort, like [`store`]. +pub fn store_skills( + dir: &ManifestDir, + config: &McpServerConfig, + skills: Vec, + skill_diagnostics: Vec, + private: bool, +) { + if private { + forget(dir, config); return; } - write_store(dir, &all); + let key = invocation_key(config); + update_store(dir, |all| { + let Some(manifest) = all.get_mut(&config.name).filter(|m| m.key == key) else { + return false; + }; + manifest.skills = Some(skills); + manifest.skill_diagnostics = skill_diagnostics; + true + }); } /// Record what `config`'s server advertises. Best-effort: a cache that cannot be written costs a @@ -196,19 +237,18 @@ pub fn store( skills: Option>, skill_diagnostics: Vec, ) { - let mut all = read_store(dir); - all.insert( - config.name.clone(), - ServerManifest { - key: invocation_key(config), - tools, - resources, - prompts, - skills, - skill_diagnostics, - }, - ); - write_store(dir, &all); + let manifest = ServerManifest { + key: invocation_key(config), + tools, + resources, + prompts, + skills, + skill_diagnostics, + }; + update_store(dir, |all| { + all.insert(config.name.clone(), manifest); + true + }); } fn write_store(dir: &ManifestDir, all: &Store) { @@ -220,10 +260,72 @@ fn write_store(dir: &ManifestDir, all: &Store) { let _ = std::fs::create_dir_all(parent); } // Write-then-rename: a crash mid-write must not leave a truncated manifest that reads as a - // *different* tool set. A miss is fine; a plausible-looking wrong answer is not. - let tmp = path.with_extension("json.tmp"); - if std::fs::write(&tmp, &bytes).is_ok() { - let _ = std::fs::rename(&tmp, &path); + // *different* tool set. A miss is fine; a plausible-looking wrong answer is not. The temporary + // name is this writer's own, so a writer that lost the lock to a stale-lock break still cannot + // rename another's half-written file into place. + static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let seq = SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let tmp = path.with_extension(format!("json.{}.{seq}.tmp", std::process::id())); + if std::fs::write(&tmp, &bytes).is_ok() && std::fs::rename(&tmp, &path).is_err() { + let _ = std::fs::remove_file(&tmp); + } +} + +const LOCK_TIMEOUT: Duration = Duration::from_secs(5); +const LOCK_RETRY_INTERVAL: Duration = Duration::from_millis(5); +const STALE_LOCK_AGE: Duration = Duration::from_secs(10); + +/// A cross-process advisory lock via atomic lockfile creation — `mcp_auth_store.rs::FileLock`'s +/// pattern (itself `trust_store.rs`'s and `settings.rs`'s), duplicated rather than shared like each of +/// those. Held for one read-modify-write of the manifest; a lockfile older than [`STALE_LOCK_AGE`] is +/// a crashed holder's and is broken. +struct FileLock { + path: PathBuf, +} + +impl FileLock { + fn acquire(store_path: &Path) -> std::io::Result { + let mut os = store_path.as_os_str().to_owned(); + os.push(".lock"); + let lock_path = PathBuf::from(os); + if let Some(parent) = lock_path.parent() { + std::fs::create_dir_all(parent)?; + } + let deadline = Instant::now() + LOCK_TIMEOUT; + loop { + match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&lock_path) + { + Ok(_) => return Ok(Self { path: lock_path }), + Err(e) if e.kind() == ErrorKind::AlreadyExists => { + let stale = std::fs::metadata(&lock_path) + .and_then(|m| m.modified()) + .ok() + .and_then(|t| SystemTime::now().duration_since(t).ok()) + .is_some_and(|age| age > STALE_LOCK_AGE); + if stale { + let _ = std::fs::remove_file(&lock_path); + continue; + } + if Instant::now() >= deadline { + return Err(std::io::Error::new( + ErrorKind::TimedOut, + format!("timed out waiting for {}", lock_path.display()), + )); + } + std::thread::sleep(LOCK_RETRY_INTERVAL); + } + Err(e) => return Err(e), + } + } + } +} + +impl Drop for FileLock { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.path); } } @@ -243,6 +345,70 @@ mod tests { } } + /// Connects and skills refreshes on many servers at once — the manifest is one file they all + /// read, change and write back — lose nobody's entry and nobody's refresh. + #[test] + fn concurrent_stores_and_skills_refreshes_lose_no_update() { + const SERVERS: usize = 24; + let dir = tempfile::tempdir().unwrap(); + let manifest = ManifestDir::at(dir.path()); + let configs: Vec = (0..SERVERS) + .map(|i| stdio(&format!("s{i}"), "npx", &[&format!("server-{i}")])) + .collect(); + let race = |work: &(dyn Fn(usize) + Sync)| { + let barrier = std::sync::Barrier::new(SERVERS); + std::thread::scope(|scope| { + for i in 0..SERVERS { + let barrier = &barrier; + scope.spawn(move || { + barrier.wait(); + work(i); + }); + } + }); + }; + race(&|i| { + store( + &manifest, + &configs[i], + vec![], + vec![], + vec![], + Some(vec![]), + vec![], + ); + }); + for (i, c) in configs.iter().enumerate() { + assert!(load(&manifest, c).is_some(), "s{i}'s connect was lost"); + } + race(&|i| { + store_skills( + &manifest, + &configs[i], + vec![], + vec![format!("refreshed-{i}")], + false, + ); + }); + for (i, c) in configs.iter().enumerate() { + let m = load(&manifest, c).unwrap_or_else(|| panic!("s{i}'s entry was lost")); + assert_eq!( + m.skill_diagnostics, + [format!("refreshed-{i}")], + "s{i}'s refresh was lost" + ); + } + let leftovers: Vec<_> = std::fs::read_dir(dir.path()) + .unwrap() + .map(|e| e.unwrap().file_name().to_string_lossy().into_owned()) + .filter(|n| n != FILE) + .collect(); + assert!( + leftovers.is_empty(), + "no lock or temporary file is left: {leftovers:?}" + ); + } + #[test] fn the_key_changes_when_the_pinned_version_does() { // The realistic staleness case: bumping `@playwright/mcp@0.0.78` to `@0.0.79` in settings.json diff --git a/crates/agent/src/tools/mcp_skills.rs b/crates/agent/src/tools/mcp_skills.rs index bfb252e3..b088343a 100644 --- a/crates/agent/src/tools/mcp_skills.rs +++ b/crates/agent/src/tools/mcp_skills.rs @@ -655,6 +655,9 @@ pub struct ServerSkills { /// The newest entry known per skill URI — the listing's, a `skills/get` refresh of it, or a skill /// learned by URI alone — with how long it stays current. known: Mutex)>>, + /// Where this server's manifest is cached, if anywhere: a refreshed listing is written back there + /// (see [`super::mcp_manifest::store_skills`]). + manifest: Option, } /// A skill loaded and verified, ready to enter the model's context. @@ -664,12 +667,13 @@ struct Loaded { } /// Who asked for a skill to be activated. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub enum Consent { /// The user named it (`/skill::`): that is the explicit consent activation needs. User, - /// The model chose it: activation waits for the user's approval. - Model, + /// A model chose it — the session's own agent or a subagent, as the approval question says: + /// activation waits for the user's approval. + Model(ApprovalOrigin), } fn lock(m: &Mutex) -> std::sync::MutexGuard<'_, T> { @@ -683,6 +687,7 @@ impl ServerSkills { conn: Arc, listing: Listing, changed: Arc, + manifest: Option, ) -> Self { let known = listing .entries @@ -697,6 +702,7 @@ impl ServerSkills { changed, listing: Mutex::new(listing), known: Mutex::new(known), + manifest, } } @@ -739,6 +745,17 @@ impl ServerSkills { known.insert(e.uri.clone(), (e.clone(), listing.fresh_until)); } } + // The cache must not keep advertising what the server just replaced — nor keep at all a + // listing it now marks private. + if let Some(dir) = &self.manifest { + super::mcp_manifest::store_skills( + dir, + self.conn.config(), + listing.entries.clone(), + listing.diagnostics.clone(), + listing.private, + ); + } *lock(&self.listing) = listing; true } @@ -845,7 +862,7 @@ impl ServerSkills { let client = self.client().await?; let reload = session.is_active(&self.server, uri); let (entry, refreshed) = self.current_entry(&client, uri, reload).await?; - session.authorize_activation(self, &entry, consent).await?; + session.authorize_activation(self, &entry, &consent).await?; let (entry, body) = match self.fetch_skill_md(&client, &entry).await { Ok(body) => (entry, body), // Stale (or tampered) against the held entry: refresh once and retry under the current @@ -860,7 +877,7 @@ impl ServerSkills { lock(&self.known).insert(uri.to_string(), (current.clone(), until)); // A changed manifest revokes an approval bound to the old one. session - .authorize_activation(self, ¤t, consent) + .authorize_activation(self, ¤t, &consent) .await?; let body = self.fetch_skill_md(&client, ¤t).await?; (current, body) @@ -942,7 +959,12 @@ impl ServerSkills { } /// One call of the loading tool, on behalf of `session`. See [`McpSkillTool`]'s description. - async fn handle(&self, session: &SkillSession, uri: &str) -> Result { + async fn handle( + &self, + session: &SkillSession, + uri: &str, + origin: &ApprovalOrigin, + ) -> Result { let mut is_known = lock(&self.known).contains_key(uri); // A SKILL.md the session can already read as part of a loaded skill: it may be a nested skill // nobody listed. Passing a SKILL.md URI asks for a load, so ask the server; only if it says @@ -962,7 +984,9 @@ impl ServerSkills { } } if is_known { - let loaded = self.load(session, uri, Consent::Model, true).await?; + let loaded = self + .load(session, uri, Consent::Model(origin.clone()), true) + .await?; return Ok(text_output(self.render(&loaded, ""))); } if let Some((entry, listed)) = session.owner(&self.server, uri) { @@ -981,7 +1005,9 @@ impl ServerSkills { if uri.ends_with("/SKILL.md") { // A skill this host never saw listed — handed over by the user, the server's // instructions, or another skill. `skills/get` vets it; an unknown URI is an error. - let loaded = self.load(session, uri, Consent::Model, true).await?; + let loaded = self + .load(session, uri, Consent::Model(origin.clone()), true) + .await?; return Ok(text_output(self.render(&loaded, ""))); } Err(format!( @@ -1101,7 +1127,9 @@ struct Active { struct SessionState { /// Skills the session is acting on, by `(server, SKILL.md uri)`. active: HashMap<(String, String), Active>, - /// Remembered "session"-scoped decisions, by key (which embeds the manifest fingerprint). + /// Remembered "session"-scoped decisions, by key (which embeds the manifest fingerprint). Held in + /// memory for this process's life of the session and never persisted: anything on disk that + /// grants an approval is something a model with file-write tools can write itself. decisions: HashMap, } @@ -1187,6 +1215,11 @@ impl SkillSession { *lock(&self.state) = SessionState::default(); } + /// Remember a decision for the rest of the session. + fn remember(&self, key: String, allow: bool) { + lock(&self.state).decisions.insert(key, allow); + } + /// Rebuild the acting window from a transcript: every MCP skill whose `SKILL.md` is still in the /// session's context — a skill-tool result, or a `/skill:` expansion — is one the model is acting /// on, whether this process loaded it or the session was resumed, switched to, forked or reopened @@ -1268,10 +1301,18 @@ impl SkillSession { } /// The session-bound loading tool replacing `tool`, if `tool` is one of the registered servers'. - pub(crate) fn bind(self: &Arc, tool: &Arc) -> Option> { + /// Bound for the agent `origin` (the session's own, or a subagent), so the activation questions + /// its loads raise say which agent asked. + pub(crate) fn bind_as( + self: &Arc, + tool: &Arc, + origin: &ApprovalOrigin, + ) -> Option> { let sources = lock(&self.sources); let source = sources.iter().find(|s| s.tool == tool.name())?; - Some(Arc::new(McpSkillTool::new(source.clone(), self.clone()))) + Some(Arc::new( + McpSkillTool::new(source.clone(), self.clone()).with_origin(origin.clone()), + )) } fn is_active(&self, server: &str, uri: &str) -> bool { @@ -1337,7 +1378,7 @@ impl SkillSession { match gate.request(request, &CancellationToken::new()).await { Ok(decision) => { if remember && decision.scope == ApprovalScope::Session { - lock(&self.state).decisions.insert(key, decision.allow); + self.remember(key, decision.allow); } Ok(decision.allow) } @@ -1348,7 +1389,9 @@ impl SkillSession { } /// Activation needs the user's consent, bound to the entry's manifest. A user's own - /// `/skill:` invocation *is* that consent. A nested skill is no exception — its approval is its + /// `/skill:` invocation *is* that consent — for that invocation only: it is not remembered, so a + /// later load the model chooses is asked about like any other (typing `/skill:x` once is not "let + /// the model load x whenever it likes"). A nested skill is no exception — its approval is its /// own, keyed by its own URI, so approving the enclosing skill never covers it — and the /// question says it is nested, so the user knows what they are agreeing to. The question carries /// the entry's frontmatter and file manifest, so a client can show the user what they would load @@ -1357,14 +1400,14 @@ impl SkillSession { &self, source: &ServerSkills, entry: &SkillEntry, - consent: Consent, + consent: &Consent, ) -> Result<(), String> { let fingerprint = entry.fingerprint(); let key = format!("activate\0{}\0{}\0{fingerprint}", source.server, entry.uri); - if consent == Consent::User { - lock(&self.state).decisions.insert(key, true); - return Ok(()); - } + let origin = match consent { + Consent::User => return Ok(()), + Consent::Model(origin) => origin.clone(), + }; let nested_in = source.enclosing(&entry.uri); let files: Vec = entry .files @@ -1375,7 +1418,7 @@ impl SkillSession { tool: source.tool.clone(), summary: json!({ "uri": entry.uri }), scope_key: format!("mcp-skill:{}:{}@{fingerprint}", source.server, entry.uri), - origin: ApprovalOrigin::Main, + origin, context: Some(json!({ "purpose": "activate", "server": source.server, @@ -1596,6 +1639,7 @@ pub(crate) fn attach( conn: &Arc, listing: Listing, tools: &mut Vec>, + manifest: Option<&super::mcp_manifest::ManifestDir>, ) -> Arc { let skills = Arc::new(ServerSkills::new( server, @@ -1603,6 +1647,7 @@ pub(crate) fn attach( conn.clone(), listing, conn.skills_changed(), + manifest.cloned(), )); tools.push(Arc::new(McpSkillTool::new( skills.clone(), @@ -1618,6 +1663,8 @@ pub(crate) struct McpSkillTool { description: String, skills: Arc, session: Arc, + /// The agent this binding serves, named in the activation questions its loads raise. + origin: ApprovalOrigin, } impl McpSkillTool { @@ -1636,8 +1683,14 @@ impl McpSkillTool { description, skills, session, + origin: ApprovalOrigin::Main, } } + + fn with_origin(mut self, origin: ApprovalOrigin) -> Self { + self.origin = origin; + self + } } #[async_trait] @@ -1671,7 +1724,7 @@ impl Tool for McpSkillTool { .filter(|u| !u.is_empty()) .ok_or_else(|| ToolError::InvalidInput("`uri` (a string) is required".into()))?; self.skills - .handle(&self.session, uri) + .handle(&self.session, uri, &self.origin) .await .map_err(ToolError::Execution) } diff --git a/crates/agent/src/tools/mod.rs b/crates/agent/src/tools/mod.rs index 5c0edc3f..33f1f1ed 100644 --- a/crates/agent/src/tools/mod.rs +++ b/crates/agent/src/tools/mod.rs @@ -772,7 +772,7 @@ mod tests { #[test] fn root_is_inside_git_repo_is_false_with_no_git_dir_anywhere() { - let dir = tempfile::tempdir().unwrap(); + let dir = crate::test_support::isolated_tempdir(); let nested = dir.path().join("src"); std::fs::create_dir_all(&nested).unwrap(); assert!(!root_is_inside_git_repo(&nested)); diff --git a/crates/agent/src/tools/subagent.rs b/crates/agent/src/tools/subagent.rs index 13b03289..0a47f17a 100644 --- a/crates/agent/src/tools/subagent.rs +++ b/crates/agent/src/tools/subagent.rs @@ -807,7 +807,14 @@ impl Subagent { } }; - let mut registry = self.build_child_registry(def, &root); + // One id per spawned child: it disambiguates the prompt-cache key *and* names this child in an + // approval request, so a UI can say "reviewer#7 wants to run bash" rather than "something does". + let spawn_id = CHILD_SEQ.fetch_add(1, Ordering::Relaxed); + let origin = crate::approval::ApprovalOrigin::Subagent { + agent: def.name.clone(), + spawn_id: spawn_id.to_string(), + }; + let mut registry = self.build_child_registry(def, &root, &origin); // The parent asked for typed data back. Registered *after* `build_child_registry`'s own // `--tools` filter, like every other explicitly-requested tool, and *before* the system prompt is // built so `has_structured_output` picks it up and the child is actually told the contract. @@ -853,10 +860,6 @@ impl Subagent { .unwrap_or_else(|| self.ctx.parent_model.clone()); let transport = (self.ctx.factory)(&model)?; - // One id per spawned child: it disambiguates the prompt-cache key *and* names this child in an - // approval request, so a UI can say "reviewer#7 wants to run bash" rather than "something does". - let spawn_id = CHILD_SEQ.fetch_add(1, Ordering::Relaxed); - let mut agent = Agent::new(transport, model) .with_tools(registry) .with_system(system) @@ -890,10 +893,7 @@ impl Subagent { agent = agent.with_hooks(Arc::new(ChildHooks { policy: policy.clone(), approval: self.ctx.approval.clone(), - origin: crate::approval::ApprovalOrigin::Subagent { - agent: def.name.clone(), - spawn_id: spawn_id.to_string(), - }, + origin, mcp: self.ctx.mcp_enabled.clone(), })); // Deliberately no `.with_checkpoint_hook`: `Agent::new`'s `NoCheckpoint` default is correct. A @@ -1026,11 +1026,20 @@ impl Subagent { /// The child's tool set: rooted at `root`, restricted to its effective tools, and carrying a /// depth-incremented `subagent` only when the definition asked for one *and* the cap allows it. - fn build_child_registry(&self, def: &AgentDef, root: &Path) -> ToolRegistry { + fn build_child_registry( + &self, + def: &AgentDef, + root: &Path, + origin: &crate::approval::ApprovalOrigin, + ) -> ToolRegistry { // Read through the gate *now*, not when the ctx was built: `set_mcp_enabled` is a live - // session command, and a child must see the kit its parent has at this moment. - let mcp_tools = - crate::tools::mcp::filter_by_enabled(&self.ctx.mcp_tools, &self.ctx.mcp_enabled); + // session command, and a child must see the kit its parent has at this moment. Its skill + // tools are bound in its own name, so a load it makes asks the user as this child. + let mcp_tools = crate::tools::mcp::filter_by_enabled_as( + &self.ctx.mcp_tools, + &self.ctx.mcp_enabled, + origin, + ); let mut registry = super::default_registry_with_config(&super::ToolConfig { bash_timeout_ms: self.ctx.tool_cfg.bash_timeout_ms, bash_shell_path: self.ctx.tool_cfg.bash_shell_path.as_deref(), diff --git a/crates/agent/src/worktree.rs b/crates/agent/src/worktree.rs index b59ca366..c67dcd90 100644 --- a/crates/agent/src/worktree.rs +++ b/crates/agent/src/worktree.rs @@ -1108,7 +1108,7 @@ mod tests { #[tokio::test] async fn preflight_rejects_a_non_repo_and_a_repo_with_no_commits() { - let plain = tempfile::tempdir().unwrap(); + let plain = crate::test_support::isolated_tempdir(); let err = preflight(&Git::Local, plain.path()).await.unwrap_err(); assert!(err.contains("not inside one"), "{err}"); diff --git a/crates/agent/tests/common/mcp_events_fixture.rs b/crates/agent/tests/common/mcp_events_fixture.rs index 658dec37..2f23982a 100644 --- a/crates/agent/tests/common/mcp_events_fixture.rs +++ b/crates/agent/tests/common/mcp_events_fixture.rs @@ -258,24 +258,37 @@ pub fn model_requests_with(bodies: &std::sync::Mutex>, needle: &str) /// The daemon session that owns configured subscriptions by default. pub const EVENTS_SESSION: &str = "mcp-events"; -/// A `serve --listen` daemon on `port` with `--mcp-events-callback-url` pointing back at it. -pub fn spawn_daemon(home: &Path, base: &str, port: u16, extra: &[&str]) -> ChildGuard { - spawn_daemon_env(home, base, port, extra, &[]) +/// A listening `serve` daemon with `--mcp-events-callback-url` pointing back at it, and its port. The +/// callback URL names the port before the daemon starts, so the port is a [`super::HeldPort`]. +pub fn spawn_daemon(home: &Path, base: &str, extra: &[&str]) -> (ChildGuard, u16) { + spawn_daemon_env(home, base, extra, &[]) } /// [`spawn_daemon`] with extra environment. pub fn spawn_daemon_env( home: &Path, base: &str, - port: u16, + extra: &[&str], + env: &[(&str, &str)], +) -> (ChildGuard, u16) { + let held = super::HeldPort::bind(); + // `held` drops on return: the daemon's is the only copy, so if it dies the port is refused. + (spawn_daemon_on(home, base, &held, extra, env), held.port()) +} + +/// [`spawn_daemon_env`] on a port the test holds — to restart a daemon where its callback URL, and +/// so every subscription's, still points. +pub fn spawn_daemon_on( + home: &Path, + base: &str, + held: &super::HeldPort, extra: &[&str], env: &[(&str, &str)], ) -> ChildGuard { + let port = held.port(); let mut cmd = super::serve_dir_cmd(super::BIN, base, &home.join("sessions").to_string_lossy()); cmd.envs(env.iter().copied()); cmd.args([ - "--listen", - &format!("127.0.0.1:{port}"), "--mcp-events-callback-url", &format!("http://127.0.0.1:{port}"), ]) @@ -288,9 +301,8 @@ pub fn spawn_daemon_env( .stderr(Stdio::from( std::fs::File::create(home.join(format!("serve-{port}.stderr"))).unwrap(), )); - let child = cmd.spawn_guarded(); - super::wait_for_port(port); - child + held.hand_to(&mut cmd); + cmd.spawn_guarded() } /// The next WebSocket frame matching `pred`, or a panic naming `what` after `timeout`. diff --git a/crates/agent/tests/common/mod.rs b/crates/agent/tests/common/mod.rs index 75b7b5b7..33505f03 100644 --- a/crates/agent/tests/common/mod.rs +++ b/crates/agent/tests/common/mod.rs @@ -160,12 +160,16 @@ pub fn gateway_bin() -> PathBuf { // rather than re-imported at each call site: every test that already says `common::turn_text` // keeps working, and there is one implementation of the wire format rather than two. pub use beyond_ai_test_support::{ - spawn_model_server, spawn_model_server_routed, spawn_model_server_with_stalled_response, sse, - turn_refusal, turn_text, turn_text_responses, turn_tool_use, + SCRIPTED_SESSION_TITLE, SESSION_TITLE_MARKER, spawn_model_server, spawn_model_server_routed, + spawn_model_server_with_stalled_response, sse, title_calls, turn_refusal, turn_text, + turn_text_responses, turn_tool_use, }; -/// A free localhost port (bind `:0`, read it back, release). A subprocess must bind it promptly; -/// there's a small TOCTOU window, acceptable for tests. +/// A free localhost port (bind `:0`, read it back, release) — for a subprocess that can neither +/// report the port it bound nor adopt one it is handed (the gateway, `nats-server`). It has a window: +/// between the release and the child's bind any other process can take the port, and under parallel +/// suites one did, handing a test's connections to another test's server. A `serve` child never needs +/// this: use [`spawn_listening`], or [`HeldPort`] when the port must be known before it starts. pub fn free_port() -> u16 { TcpListener::bind("127.0.0.1:0") .unwrap() @@ -174,6 +178,168 @@ pub fn free_port() -> u16 { .port() } +/// The line `serve` prints on stderr once its `--listen` socket is bound (`serve_ws::serve_ws`). +const LISTENING: &str = "serve: websocket listening on "; + +/// The port in `serve`'s listening announcement, if `line` is one. +fn announced_port(line: &str) -> Option { + let addr = line.strip_prefix(LISTENING)?.split_whitespace().next()?; + addr.parse::().ok().map(|a| a.port()) +} + +/// Spawn a `serve` command listening on a port the kernel picks, and return that port: appends +/// `--listen 127.0.0.1:0` and reads the bound address back from the line `serve` announces it with. +/// The port is `serve`'s from the moment it exists, so — unlike [`free_port`], which releases the +/// port it picked and hopes the child binds it first — no other process can take it. When this +/// returns the listener is up, so there is nothing to wait for. +/// +/// Takes the child's stderr to read that line; everything it writes is forwarded line by line to the +/// test's (captured) stderr, so a child's panic or startup error is shown with a failing test instead +/// of being discarded. A child that exits without announcing fails here, with what it said. +pub fn spawn_listening(cmd: &mut Command) -> (ChildGuard, u16) { + spawn_listening_logged(cmd, None) +} + +/// [`spawn_listening`], also copying the child's stderr to `log` — for a test that reads it back. +pub fn spawn_listening_logged(cmd: &mut Command, log: Option) -> (ChildGuard, u16) { + cmd.args(["--listen", "127.0.0.1:0"]).stderr(Stdio::piped()); + let mut child = ChildGuard::spawn(cmd); + let stderr = child.stderr.take().expect("stderr is piped"); + let (tx, rx) = std::sync::mpsc::channel::>(); + // Spawned threads inherit the test's output capture, so `eprintln!` here lands in it. + thread::spawn(move || { + let mut log = log; + let mut pending = Some((tx, String::new())); + for line in std::io::BufReader::new(stderr).lines() { + let Ok(line) = line else { break }; + eprintln!("[serve] {line}"); + if let Some(f) = log.as_mut() { + let _ = writeln!(f, "{line}"); + } + if let Some((tx, said)) = pending.as_mut() { + match announced_port(&line) { + Some(port) => { + let _ = tx.send(Ok(port)); + pending = None; + } + None => { + said.push_str(&line); + said.push('\n'); + } + } + } + } + if let Some((tx, said)) = pending { + let _ = tx.send(Err(said)); + } + }); + match rx.recv_timeout(Duration::from_secs(60)) { + Ok(Ok(port)) => (child, port), + Ok(Err(said)) => panic!("serve exited without listening; it said:\n{said}"), + Err(e) => panic!("serve did not announce its listener within 60s ({e})"), + } +} + +/// A loopback port the test holds from before `serve` starts, for a test that has to name the port +/// in `serve`'s own arguments (an MCP Events callback URL pointing back at the daemon), so cannot +/// learn it afterwards. The test binds the listener and hands it to `serve` as a socket-activated +/// one (`LISTEN_FDS`, exactly as systemd does) — so the port is never released between being chosen +/// and being served, and no other process can take it. Held for as long as this value lives, so a +/// daemon restarted on the same port gets the same socket ([`HeldPort::down`] covers the gap). +/// +/// It is passed as the child's stdin (inetd's convention, `LISTEN_FDS_FIRST_FD=0`): std's `dup2` onto +/// a standard stream is the one way to give a child a descriptor without `unsafe`, which the workspace +/// forbids, and `serve` never reads stdin once it serves a listener. Drop it once the daemon is up +/// unless it will be restarted: while the test holds the socket, a daemon that died leaves it +/// listening, and a client would wait on it instead of being refused. +pub struct HeldPort { + listener: TcpListener, + port: u16, +} + +impl HeldPort { + pub fn bind() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + Self { listener, port } + } + + pub fn port(&self) -> u16 { + self.port + } + + /// Make this socket `cmd`'s listener. Call it after anything else that sets stdin, and pass no + /// `--listen`/`--listen-uds` (either turns socket activation off). + pub fn hand_to(&self, cmd: &mut Command) { + let fd = std::os::fd::OwnedFd::from(self.listener.try_clone().unwrap()); + cmd.stdin(Stdio::from(fd)) + .env("LISTEN_FDS", "1") + .env("LISTEN_FDS_FIRST_FD", "0") + .env_remove("LISTEN_PID"); + } + + /// Between a daemon's death and its restart: until the guard drops, every connection is accepted + /// and closed at once, so a client fails promptly — as against a stopped daemon — rather than + /// waiting in the backlog of a socket the test still holds. + pub fn down(&self) -> PortDown { + let listener = self.listener.try_clone().unwrap(); + listener.set_nonblocking(true).unwrap(); + let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let stopping = stop.clone(); + let thread = thread::spawn(move || { + while !stopping.load(std::sync::atomic::Ordering::Relaxed) { + match listener.accept() { + Ok((conn, _)) => drop(conn), + Err(_) => thread::sleep(Duration::from_millis(5)), + } + } + }); + PortDown { + stop, + thread: Some(thread), + } + } +} + +/// See [`HeldPort::down`]. +pub struct PortDown { + stop: Arc, + thread: Option>, +} + +impl Drop for PortDown { + fn drop(&mut self) { + self.stop.store(true, std::sync::atomic::Ordering::Relaxed); + if let Some(t) = self.thread.take() { + let _ = t.join(); + } + } +} + +/// A loopback port nothing listens on, held for as long as the value lives: bound but never +/// listening, so a connection is refused, and taken, so no other process can start listening there +/// in the meantime — which a released "dead" port from [`free_port`] cannot promise. +pub struct DeadPort { + _socket: tokio::net::TcpSocket, + port: u16, +} + +impl DeadPort { + pub fn bind() -> Self { + let socket = tokio::net::TcpSocket::new_v4().unwrap(); + socket.bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let port = socket.local_addr().unwrap().port(); + Self { + _socket: socket, + port, + } + } + + pub fn port(&self) -> u16 { + self.port + } +} + /// Block until `port` accepts a TCP connection, or panic after ~5s. pub fn wait_for_port(port: u16) { for _ in 0..500 { @@ -292,6 +458,74 @@ pub fn read_until_response(reader: &mut impl BufRead, command: &str) -> Vec, + current: Vec, + pos: usize, + deadline: Duration, +} + +/// [`Frames`] over `out`. +pub fn frames_with_deadline(out: impl Read + Send + 'static, deadline: Duration) -> Frames { + let (tx, lines) = std::sync::mpsc::channel(); + thread::spawn(move || { + let mut out = std::io::BufReader::new(out); + loop { + let mut line = String::new(); + match out.read_line(&mut line) { + Ok(0) | Err(_) => break, + Ok(_) if tx.send(line).is_err() => break, + Ok(_) => {} + } + } + }); + Frames { + lines, + current: Vec::new(), + pos: 0, + deadline, + } +} + +impl Read for Frames { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let available = self.fill_buf()?; + let n = available.len().min(buf.len()); + buf[..n].copy_from_slice(&available[..n]); + self.consume(n); + Ok(n) + } +} + +impl BufRead for Frames { + fn fill_buf(&mut self) -> std::io::Result<&[u8]> { + if self.pos >= self.current.len() { + match self.lines.recv_timeout(self.deadline) { + Ok(line) => { + self.current = line.into_bytes(); + self.pos = 0; + } + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => return Ok(&[]), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => panic!( + "serve wrote nothing for {:?}: the run is stalled (a scripted reply taken by \ + the wrong request — the session-title call, say — leaves it waiting forever)", + self.deadline + ), + } + } + Ok(&self.current[self.pos..]) + } + + fn consume(&mut self, n: usize) { + self.pos += n; + } +} + /// Read stdout frames from a `serve` child until an `event` frame whose body satisfies `matches` /// arrives; return all frames seen, that event last. /// @@ -488,13 +722,12 @@ pub fn ws_request( request } -/// A `serve --service` child: a grant verifier, one or more `--shard =` mounts, and a -/// listener. Deliberately *not* built on [`serve_cmd`] — that passes `--key` and `--session-file`, -/// both of which service mode refuses at startup. +/// A `serve --service` child: a grant verifier and one or more `--shard =` mounts. Spawn +/// it with [`spawn_listening`], which adds the listener. Deliberately *not* built on [`serve_cmd`] — +/// that passes `--key` and `--session-file`, both of which service mode refuses at startup. pub fn serve_service_cmd( bin: &str, base: &str, - port: u16, grant_key_flag: &str, seal_key: &std::path::Path, shards: &[(&str, &std::path::Path)], @@ -508,8 +741,6 @@ pub fn serve_service_cmd( base, "--model", "claude-test", - "--listen", - &format!("127.0.0.1:{port}"), "--grant-key", grant_key_flag, "--seal-key", @@ -520,8 +751,7 @@ pub fn serve_service_cmd( } c.env("HOME", ISOLATED_HOME) .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::piped()); + .stdout(Stdio::null()); c } diff --git a/crates/agent/tests/common/service.rs b/crates/agent/tests/common/service.rs index 1d6dc9a5..2cb34d45 100644 --- a/crates/agent/tests/common/service.rs +++ b/crates/agent/tests/common/service.rs @@ -18,7 +18,7 @@ use beyond_ai_agent::session_store::{Layout, RepoOptions, SessionRepo, TenantCod use super::exec_mock::ExecMock; use super::grant::{Claims, Minter, Secrets}; -use super::{ChildGuard, SpawnGuarded, free_port, serve_service_cmd, wait_for_port}; +use super::{ChildGuard, serve_service_cmd, spawn_listening, spawn_listening_logged}; /// The header value every exec request from a service session must carry — the grant's sealed /// `exec_headers`. A test asserts both that it *reaches the sandbox* and that it never appears in a @@ -39,6 +39,8 @@ pub struct Service { /// `` → directory, in `--shard` flag order. pub shards: Vec<(String, PathBuf)>, pub child: ChildGuard, + /// Everything the replica wrote to stderr (the operator's channel), as it writes it. + pub stderr_log: PathBuf, /// The model server this replica was pointed at, so a [`Peer`] can share it. pub model_base: String, /// Held so every directory above outlives the replica. @@ -87,7 +89,6 @@ impl Service { }) .collect(); let minter = Minter::new(dir.path()); - let port = free_port(); let flag_shards: Vec<(&str, &Path)> = shards .iter() .map(|(n, p)| (n.as_str(), p.as_path())) @@ -95,7 +96,6 @@ impl Service { let mut cmd = serve_service_cmd( super::BIN, model_base, - port, &minter.grant_key_flag(), minter.seal_key(), &flag_shards, @@ -107,8 +107,9 @@ impl Service { for (k, v) in &opts.env { cmd.env(k, v); } - let child = cmd.spawn_guarded(); - wait_for_port(port); + let stderr_log = dir.path().join("replica.stderr"); + let log = std::fs::File::create(&stderr_log).unwrap(); + let (child, port) = spawn_listening_logged(&mut cmd, Some(log)); Self { port, minter, @@ -117,6 +118,7 @@ impl Service { sandbox_home, shards, child, + stderr_log, model_base: model_base.to_string(), dir, } @@ -130,18 +132,15 @@ impl Service { .iter() .map(|(n, p)| (n.as_str(), p.as_path())) .collect(); - let port = free_port(); let mut cmd = serve_service_cmd( super::BIN, &self.model_base, - port, &self.minter.grant_key_flag(), self.minter.seal_key(), &flag_shards, ); cmd.args(extra_args); - let child = cmd.spawn_guarded(); - wait_for_port(port); + let (child, port) = spawn_listening(&mut cmd); Peer { port, child } } diff --git a/crates/agent/tests/mcp_apps_events.rs b/crates/agent/tests/mcp_apps_events.rs index e03b370e..ad988f7f 100644 --- a/crates/agent/tests/mcp_apps_events.rs +++ b/crates/agent/tests/mcp_apps_events.rs @@ -11,9 +11,7 @@ use std::time::Duration; use common::mcp_events_fixture::{ emit, eventually, runs_for_event, spawn_daemon, spawn_http_fixture, state, write_settings, }; -use common::{ - free_port, spawn_model_server_routed, turn_text, ws_connect, ws_read_until_response, ws_send, -}; +use common::{spawn_model_server_routed, turn_text, ws_connect, ws_read_until_response, ws_send}; use serde_json::json; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -31,8 +29,7 @@ async fn a_declared_apps_session_leaves_events_on_the_plain_connection_one_run_p }]), ); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("handled")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &[]); + let (_d, port) = spawn_daemon(home.path(), &base, &[]); // The events session subscribes on its own, over the plain connection. eventually( diff --git a/crates/agent/tests/mcp_apps_resilience.rs b/crates/agent/tests/mcp_apps_resilience.rs index 92075c8d..4966f402 100644 --- a/crates/agent/tests/mcp_apps_resilience.rs +++ b/crates/agent/tests/mcp_apps_resilience.rs @@ -365,37 +365,27 @@ async fn service_mode_dials_a_sessions_own_apps_pool_from_its_grant() { ], turn_tool_use("toolu_svc_x", "mcp__wx__show_weather", "{}"), ); - // The replica takes a TCP port picked by `free_port` — another test can grab it first, so retry - // a start whose port turned out not to be ours. - let mut attempt = 0; - let (_svc, mut ws) = loop { - attempt += 1; - let svc = Service::start_with( - &base, - &["s1"], - Options { - extra_args: vec!["--mcp-allow-private".into()], - ..Default::default() - }, - ) - .await; - let mut claims = svc.claims("tenant-a", "s1.apps", &svc.shards[0].0); - claims.mcp = [("wx".to_string(), fixture.url.clone())] - .into_iter() - .collect(); - let token = svc.minter.mint(&claims, &svc.secrets()); - match tokio_tungstenite::connect_async(common::ws_request( - svc.port, - Some("s1.apps"), - &svc.header(&token), - )) - .await - { - Ok((ws, _)) => break (svc, ws), - Err(e) if attempt < 3 => eprintln!("replica start {attempt} unusable: {e}"), - Err(e) => panic!("replica never accepted a session: {e}"), - } - }; + let svc = Service::start_with( + &base, + &["s1"], + Options { + extra_args: vec!["--mcp-allow-private".into()], + ..Default::default() + }, + ) + .await; + let mut claims = svc.claims("tenant-a", "s1.apps", &svc.shards[0].0); + claims.mcp = [("wx".to_string(), fixture.url.clone())] + .into_iter() + .collect(); + let token = svc.minter.mint(&claims, &svc.secrets()); + let (mut ws, _) = tokio_tungstenite::connect_async(common::ws_request( + svc.port, + Some("s1.apps"), + &svc.header(&token), + )) + .await + .expect("the replica accepts the session"); // The session's own connector is dialed again, with the extension advertised. assert_eq!( diff --git a/crates/agent/tests/mcp_client.rs b/crates/agent/tests/mcp_client.rs index af35fe83..f9158813 100644 --- a/crates/agent/tests/mcp_client.rs +++ b/crates/agent/tests/mcp_client.rs @@ -714,9 +714,10 @@ fn mcp_http_streamable_a_server_that_refuses_the_connection_is_skipped_fail_soft // handshake itself failing, not a spawn failure), which is otherwise completely untested. let home = tempfile::tempdir().unwrap(); let cwd = tempfile::tempdir().unwrap(); - // A free port nothing is listening on — connecting to it fails fast with connection-refused, unlike - // an unroutable address, which would hang until a timeout. - let dead_port = common::free_port(); + // A port nothing is listening on — connecting to it fails fast with connection-refused, unlike + // an unroutable address, which would hang until a timeout. Held, so nothing can start to. + let dead = common::DeadPort::bind(); + let dead_port = dead.port(); let fixture = spawn_http_mcp_fixture(); let url = fixture.url.clone(); write_global_settings( diff --git a/crates/agent/tests/mcp_events_durability.rs b/crates/agent/tests/mcp_events_durability.rs index f97d1a7f..e1a4ed55 100644 --- a/crates/agent/tests/mcp_events_durability.rs +++ b/crates/agent/tests/mcp_events_durability.rs @@ -19,8 +19,7 @@ use common::mcp_events_fixture::{ ws_next, ws_wait_active, }; use common::{ - BIN, SpawnGuarded, free_port, serve_dir_cmd, spawn_model_server_routed, turn_text, ws_connect, - ws_send, + BIN, SpawnGuarded, serve_dir_cmd, spawn_model_server_routed, turn_text, ws_connect, ws_send, }; use serde_json::{Value, json}; @@ -38,8 +37,7 @@ async fn a_follow_up_held_during_a_run_survives_sigterm_and_reaches_the_model_on ); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("handled")); - let port = free_port(); - let mut first = spawn_daemon(home.path(), &base, port, &[]); + let (mut first, port) = spawn_daemon(home.path(), &base, &[]); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; ws_wait_active(&mut ws).await; // A long run is in flight, so the follow-up is held — not yet seen by the model. @@ -85,8 +83,7 @@ async fn a_follow_up_held_during_a_run_survives_sigterm_and_reaches_the_model_on assert_eq!(saved[0]["event"]["eventId"], "durable-1", "{saved:#?}"); // Restart: the pending event is injected, once. - let port = free_port(); - let _second = spawn_daemon(home.path(), &base, port, &[]); + let (_second, _) = spawn_daemon(home.path(), &base, &[]); eventually( Duration::from_secs(30), "the model to receive the held event", diff --git a/crates/agent/tests/mcp_events_lifecycle.rs b/crates/agent/tests/mcp_events_lifecycle.rs index 7adec46b..0936d946 100644 --- a/crates/agent/tests/mcp_events_lifecycle.rs +++ b/crates/agent/tests/mcp_events_lifecycle.rs @@ -8,44 +8,21 @@ mod common; -use std::path::Path; -use std::process::{Command, Stdio}; +use std::process::Command; use std::time::{Duration, Instant}; use common::mcp_events_fixture::{ - Frames, emit, eventually, model_requests_with, send, spawn_http_fixture, state, wait_active, - write_settings, + Frames, emit, eventually, model_requests_with, send, spawn_daemon, spawn_http_fixture, state, + wait_active, write_settings, }; use common::{ - BIN, ChildGuard, SpawnGuarded, TestWs, free_port, serve_cmd, serve_dir_cmd, - spawn_model_server_routed, turn_text, wait_for_port, ws_connect, ws_next_frame, ws_send, + BIN, ChildGuard, SpawnGuarded, TestWs, serve_cmd, spawn_model_server_routed, turn_text, + ws_connect, ws_next_frame, ws_send, }; use serde_json::{Value, json}; type Bodies = std::sync::Arc>>; -fn daemon(home: &Path, base: &str, port: u16, extra: &[&str]) -> ChildGuard { - let mut cmd = serve_dir_cmd(BIN, base, &home.join("sessions").to_string_lossy()); - cmd.args([ - "--listen", - &format!("127.0.0.1:{port}"), - "--mcp-events-callback-url", - &format!("http://127.0.0.1:{port}"), - ]) - .args(extra) - .env("HOME", home) - .env("BEYOND_AI_AGENT_MCP_EVENTS_COALESCE_MS", "300") - .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::from( - std::fs::File::create(home.join(format!("serve-{port}.stderr"))).unwrap(), - )); - let child = cmd.spawn_guarded(); - wait_for_port(port); - child -} - fn hooks_server(mcp_url: &str, action: &str) -> Value { json!([{ "name": "hooks", @@ -140,10 +117,9 @@ async fn an_event_arriving_while_detached_is_delivered_and_seen_on_reattach() { let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks_server(&mcp_url, "follow_up")); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("handled while away")); - let port = free_port(); // A 1 s idle timeout: without the keep-alive this session would be reaped (and unsubscribed) // long before the event below arrives. - let _d = daemon(home.path(), &base, port, &["--session-idle-timeout", "1"]); + let (_d, port) = spawn_daemon(home.path(), &base, &["--session-idle-timeout", "1"]); let mut ws = ws_connect(port, Some("mcp-events")).await; ws_wait_active(&mut ws).await; @@ -189,11 +165,9 @@ async fn mcp_events_reapable_lets_the_reaper_end_a_detached_sessions_subscriptio let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks_server(&mcp_url, "notify")); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = daemon( + let (_d, port) = spawn_daemon( home.path(), &base, - port, &["--session-idle-timeout", "1", "--mcp-events-reapable"], ); let mut ws = ws_connect(port, Some("mcp-events")).await; @@ -222,8 +196,7 @@ async fn a_restarted_daemon_resumes_webhook_delivery_from_its_persisted_cursor_w write_settings(home.path(), hooks_server(&mcp_url, "notify")); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let mut first = daemon(home.path(), &base, port, &[]); + let (mut first, port) = spawn_daemon(home.path(), &base, &[]); let mut ws = ws_connect(port, Some("mcp-events")).await; ws_wait_active(&mut ws).await; emit(&fixture, json!({ "event_id": "r-1", "data": { "n": 1 } })); @@ -255,8 +228,7 @@ async fn a_restarted_daemon_resumes_webhook_delivery_from_its_persisted_cursor_w emit(&fixture, json!({ "event_id": "r-2", "data": { "n": 2 } })); emit(&fixture, json!({ "event_id": "r-3", "data": { "n": 3 } })); - let port = free_port(); - let _second = daemon(home.path(), &base, port, &[]); + let (_second, port) = spawn_daemon(home.path(), &base, &[]); let mut ws = ws_connect(port, Some("mcp-events")).await; let got = collect(&mut ws, Duration::from_secs(6), |f| { f["type"] == "mcp_event" diff --git a/crates/agent/tests/mcp_events_receiver.rs b/crates/agent/tests/mcp_events_receiver.rs index 708d5176..641d433d 100644 --- a/crates/agent/tests/mcp_events_receiver.rs +++ b/crates/agent/tests/mcp_events_receiver.rs @@ -7,17 +7,13 @@ mod common; -use std::process::Stdio; use std::time::{Duration, Instant}; use common::mcp_events_fixture::{ - EVENTS_SESSION, control, emit, eventually, raw_request, runs_for_event, spawn_http_fixture, - state, write_settings, ws_next, ws_wait_active, -}; -use common::{ - BIN, ChildGuard, SpawnGuarded, free_port, serve_dir_cmd, spawn_model_server_routed, turn_text, - wait_for_port, ws_connect, ws_send, + EVENTS_SESSION, control, emit, eventually, raw_request, runs_for_event, spawn_daemon_env, + spawn_daemon_on, spawn_http_fixture, state, write_settings, ws_next, ws_wait_active, }; +use common::{ChildGuard, HeldPort, spawn_model_server_routed, turn_text, ws_connect, ws_send}; use serde_json::{Value, json}; fn hooks(mcp_url: &str, action: &str) -> Value { @@ -28,26 +24,9 @@ fn hooks(mcp_url: &str, action: &str) -> Value { }]) } -fn daemon(home: &std::path::Path, base: &str, port: u16, env: &[(&str, &str)]) -> ChildGuard { - let mut cmd = serve_dir_cmd(BIN, base, &home.join("sessions").to_string_lossy()); - cmd.args([ - "--listen", - &format!("127.0.0.1:{port}"), - "--mcp-events-callback-url", - &format!("http://127.0.0.1:{port}"), - ]) - .env("HOME", home) - .env("BEYOND_AI_AGENT_MCP_EVENTS_COALESCE_MS", "300") - .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") - .envs(env.iter().copied()) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::from( - std::fs::File::create(home.join(format!("serve-{port}.stderr"))).unwrap(), - )); - let child = cmd.spawn_guarded(); - wait_for_port(port); - child +/// A daemon on `held`, so a restart keeps the callback URL every subscription points at. +fn daemon(home: &std::path::Path, base: &str, held: &HeldPort, env: &[(&str, &str)]) -> ChildGuard { + spawn_daemon_on(home, base, held, &[], env) } /// The events session's state snapshot, once it exists. @@ -79,11 +58,10 @@ async fn the_pending_queue_is_bounded_in_bytes_and_answers_503_past_it() { let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url, "follow_up")); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = daemon( + let (_d, port) = spawn_daemon_env( home.path(), &base, - port, + &[], &[("BEYOND_AI_AGENT_MCP_EVENTS_MAX_PENDING_BYTES", "8000")], ); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; @@ -111,8 +89,7 @@ async fn a_delivery_that_cannot_be_stored_gets_503_and_lands_on_retry() { let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url, "follow_up")); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = daemon(home.path(), &base, port, &[]); + let (_d, port) = spawn_daemon_env(home.path(), &base, &[], &[]); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; ws_wait_active(&mut ws).await; // The pending log cannot be written: its place is taken by a directory. @@ -147,11 +124,10 @@ async fn a_delivery_body_that_never_finishes_is_timed_out() { let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url, "notify")); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = daemon( + let (_d, port) = spawn_daemon_env( home.path(), &base, - port, + &[], &[("BEYOND_AI_AGENT_MCP_EVENTS_BODY_TIMEOUT_MS", "500")], ); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; @@ -185,8 +161,8 @@ async fn the_callback_survives_a_crash_so_retried_deliveries_still_land() { let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url, "follow_up")); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let mut first = daemon(home.path(), &base, port, &[]); + let held = HeldPort::bind(); + let mut first = daemon(home.path(), &base, &held, &[]); let url = eventually(Duration::from_secs(20), "the subscription", || { state(&fixture)["hooks"][0]["url"] .as_str() @@ -204,6 +180,8 @@ async fn the_callback_survives_a_crash_so_retried_deliveries_still_land() { }); first.kill().unwrap(); let _ = first.wait(); + // Down, as a crashed daemon is — but still held, so the restart gets the same port. + let down = held.down(); // While it is down the server's delivery fails. let r = emit( &fixture, @@ -211,7 +189,8 @@ async fn the_callback_survives_a_crash_so_retried_deliveries_still_land() { ); assert_ne!(r["deliveries"][0]["status"], 200, "{r:#}"); - let _second = daemon(home.path(), &base, port, &[]); + drop(down); + let _second = daemon(home.path(), &base, &held, &[]); eventually(Duration::from_secs(20), "the resubscribe", || { let hooks = state(&fixture)["hooks"].as_array().unwrap().clone(); (hooks.len() == 1 && hooks[0]["refreshes"].as_u64().unwrap_or(0) >= 1 @@ -256,8 +235,8 @@ async fn a_retry_that_beats_the_resubscribe_after_a_restart_is_told_to_retry_not let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url, "follow_up")); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let mut first = daemon(home.path(), &base, port, &[]); + let held = HeldPort::bind(); + let mut first = daemon(home.path(), &base, &held, &[]); eventually(Duration::from_secs(30), "the subscription", || { (state(&fixture)["hooks"].as_array().unwrap().len() == 1).then_some(()) }); @@ -272,13 +251,16 @@ async fn a_retry_that_beats_the_resubscribe_after_a_restart_is_told_to_retry_not }); first.kill().unwrap(); let _ = first.wait(); + // Down, as a crashed daemon is — but still held, so the restart gets the same port. + let down = held.down(); let r = emit( &fixture, json!({ "event_id": "early-1", "data": { "summary": "retried early" } }), ); assert_ne!(r["deliveries"][0]["status"], 200, "{r:#}"); - let _second = daemon(home.path(), &base, port, &[]); + 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 r = control( @@ -315,8 +297,9 @@ async fn after_a_restart_a_runtime_subscriptions_old_callback_is_gone_not_held() let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url, "notify")); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let mut first = daemon(home.path(), &base, port, &[]); + let held = HeldPort::bind(); + let port = held.port(); + let mut first = daemon(home.path(), &base, &held, &[]); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; ws_wait_active(&mut ws).await; ws_send( @@ -351,7 +334,10 @@ async fn after_a_restart_a_runtime_subscriptions_old_callback_is_gone_not_held() drop(ws); first.kill().unwrap(); let _ = first.wait(); - let _second = daemon(home.path(), &base, port, &[]); + // Down, as a crashed daemon is — but still held, so the restart gets the same port. + let down = held.down(); + drop(down); + let _second = daemon(home.path(), &base, &held, &[]); // The configured one comes back on its old callback. eventually( Duration::from_secs(20), diff --git a/crates/agent/tests/mcp_events_routing.rs b/crates/agent/tests/mcp_events_routing.rs index e98b96d8..e8aa2e6b 100644 --- a/crates/agent/tests/mcp_events_routing.rs +++ b/crates/agent/tests/mcp_events_routing.rs @@ -14,7 +14,7 @@ use common::mcp_events_fixture::{ EVENTS_SESSION, control, daemon_sessions, emit, eventually, runs_for_event, spawn_daemon, spawn_http_fixture, state, write_settings, ws_list, ws_next, }; -use common::{free_port, spawn_model_server_routed, turn_text, ws_connect, ws_send}; +use common::{spawn_model_server_routed, turn_text, ws_connect, ws_send}; use serde_json::json; fn hooks(mcp_url: &str) -> serde_json::Value { @@ -33,8 +33,7 @@ async fn configured_subscriptions_are_owned_by_one_session_and_one_event_runs_th let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url)); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("handled")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &[]); + let (_d, port) = spawn_daemon(home.path(), &base, &[]); // Two ordinary clients, connected before the event. let mut a = ws_connect(port, Some("plain-a")).await; @@ -98,8 +97,7 @@ async fn plain_sessions_stay_reapable_while_the_events_session_is_kept_alive() { let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url)); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &["--session-idle-timeout", "1"]); + let (_d, port) = spawn_daemon(home.path(), &base, &["--session-idle-timeout", "1"]); eventually( Duration::from_secs(20), "the events session's subscription", @@ -140,8 +138,7 @@ async fn a_session_whose_only_subscription_terminated_becomes_reapable() { servers[0]["events"] = json!([]); write_settings(home.path(), servers); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &["--session-idle-timeout", "1"]); + let (_d, port) = spawn_daemon(home.path(), &base, &["--session-idle-timeout", "1"]); let mut ws = ws_connect(port, Some("runtime")).await; ws_send( &mut ws, @@ -180,8 +177,7 @@ async fn a_configured_subscription_the_server_ends_is_resubscribed() { let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url)); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &["--session-idle-timeout", "1"]); + let (_d, port) = spawn_daemon(home.path(), &base, &["--session-idle-timeout", "1"]); let subscribes = || { state(&fixture)["methods"] .as_array() @@ -221,8 +217,7 @@ async fn configured_subscriptions_that_fail_at_boot_are_retried_and_kept_alive() let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url)); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &["--session-idle-timeout", "1"]); + let (_d, port) = spawn_daemon(home.path(), &base, &["--session-idle-timeout", "1"]); tokio::time::sleep(Duration::from_millis(4000)).await; assert!(state(&fixture)["hooks"].as_array().unwrap().is_empty()); assert_eq!( @@ -254,8 +249,7 @@ async fn unsubscribing_a_configured_subscription_while_its_server_is_down_stops_ let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url)); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &[]); + let (_d, port) = spawn_daemon(home.path(), &base, &[]); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; // Let the first attempt fail. tokio::time::sleep(Duration::from_millis(1500)).await; diff --git a/crates/agent/tests/mcp_events_v1a.rs b/crates/agent/tests/mcp_events_v1a.rs index ac005694..804e9d75 100644 --- a/crates/agent/tests/mcp_events_v1a.rs +++ b/crates/agent/tests/mcp_events_v1a.rs @@ -15,8 +15,8 @@ use common::mcp_events_fixture::{ write_settings, ws_next, ws_wait_active, }; use common::{ - BIN, SpawnGuarded, TestWs, free_port, serve_dir_cmd, spawn_model_server_routed, turn_text, - wait_for_port, ws_connect, ws_next_frame, ws_send, + BIN, HeldPort, SpawnGuarded, TestWs, serve_dir_cmd, spawn_model_server_routed, turn_text, + ws_connect, ws_next_frame, ws_send, }; use serde_json::{Value, json}; @@ -48,20 +48,19 @@ async fn deliveries_must_carry_a_valid_v1a_signature_when_the_server_publishes_a }]), ); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = serve_dir_cmd(BIN, &base, &home.path().join("s").to_string_lossy()) - .args([ - "--listen", - &format!("127.0.0.1:{port}"), - "--mcp-events-callback-url", - &format!("http://127.0.0.1:{port}"), - ]) - .env("HOME", home.path()) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded(); - wait_for_port(port); + let held = HeldPort::bind(); + let port = held.port(); + let mut cmd = serve_dir_cmd(BIN, &base, &home.path().join("s").to_string_lossy()); + cmd.args([ + "--mcp-events-callback-url", + &format!("http://127.0.0.1:{port}"), + ]) + .env("HOME", home.path()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + held.hand_to(&mut cmd); + let _d = cmd.spawn_guarded(); + drop((cmd, held)); let mut ws = ws_connect(port, Some("mcp-events")).await; // The verification challenge itself was v1a-signed and checked: the subscribe succeeds. let mut active = false; @@ -122,8 +121,7 @@ async fn once_keys_are_seen_a_server_that_stops_publishing_is_still_held_to_them let home = tempfile::tempdir().unwrap(); write_settings(home.path(), hooks(&mcp_url)); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &[]); + let (_d, port) = spawn_daemon(home.path(), &base, &[]); let mut ws = ws_connect(port, Some(EVENTS_SESSION)).await; ws_wait_active(&mut ws).await; @@ -189,8 +187,7 @@ async fn a_failing_jwks_fetch_refuses_the_first_subscribe_instead_of_enforcing_n }]), ); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = spawn_daemon(home.path(), &base, port, &[]); + let (_d, port) = spawn_daemon(home.path(), &base, &[]); let mut ws = ws_connect(port, Some("identity-refused")).await; ws_send( &mut ws, @@ -238,11 +235,9 @@ async fn a_jwks_body_that_stalls_is_timed_out_and_leaves_nothing_behind() { }]), ); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let _d = common::mcp_events_fixture::spawn_daemon_env( + let (_d, port) = common::mcp_events_fixture::spawn_daemon_env( home.path(), &base, - port, &[], &[("BEYOND_AI_AGENT_MCP_EVENTS_JWKS_TIMEOUT_MS", "300")], ); diff --git a/crates/agent/tests/mcp_events_webhook.rs b/crates/agent/tests/mcp_events_webhook.rs index 0a0099bc..8287dd27 100644 --- a/crates/agent/tests/mcp_events_webhook.rs +++ b/crates/agent/tests/mcp_events_webhook.rs @@ -18,8 +18,8 @@ use common::mcp_events_fixture::{ control, emit, eventually, model_requests_with, spawn_http_fixture, state, write_settings, }; use common::{ - BIN, ChildGuard, SpawnGuarded, TestWs, free_port, serve_dir_cmd, spawn_model_server_routed, - turn_text, wait_for_port, ws_connect, ws_next_frame, ws_send, + BIN, ChildGuard, DeadPort, HeldPort, SpawnGuarded, TestWs, serve_dir_cmd, spawn_listening, + spawn_model_server_routed, turn_text, ws_connect, ws_next_frame, ws_send, }; use serde_json::{Value, json}; @@ -35,12 +35,12 @@ fn start_daemon(mcp_servers: Value) -> Daemon { write_settings(home.path(), mcp_servers); let (base, bodies) = spawn_model_server_routed(vec![], turn_text("noted")); let sessions = home.path().join("sessions"); - let port = free_port(); + // The callback URL names the port before the daemon starts. + let held = HeldPort::bind(); + let port = held.port(); let stderr = home.path().join("serve.stderr"); let mut cmd = serve_dir_cmd(BIN, &base, &sessions.to_string_lossy()); cmd.args([ - "--listen", - &format!("127.0.0.1:{port}"), "--mcp-events-callback-url", &format!("http://127.0.0.1:{port}"), ]) @@ -50,8 +50,8 @@ fn start_daemon(mcp_servers: Value) -> Daemon { .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::from(std::fs::File::create(&stderr).unwrap())); + held.hand_to(&mut cmd); let child = cmd.spawn_guarded(); - wait_for_port(port); Daemon { child, port, @@ -319,22 +319,18 @@ async fn a_callback_url_nothing_answers_fails_the_challenge_and_the_subscribe() let home = tempfile::tempdir().unwrap(); write_settings(home.path(), json!([webhook_server(&mcp_url, json!([]))])); let (base, _bodies) = spawn_model_server_routed(vec![], turn_text("noted")); - let port = free_port(); - let dead = free_port(); + // Points somewhere nothing listens (and, held, nothing can start to): the server's verification + // POST cannot land. + let dead = DeadPort::bind(); let mut cmd = serve_dir_cmd(BIN, &base, &home.path().join("s").to_string_lossy()); cmd.args([ - "--listen", - &format!("127.0.0.1:{port}"), - // Points somewhere nothing listens: the server's verification POST cannot land. "--mcp-events-callback-url", - &format!("http://127.0.0.1:{dead}"), + &format!("http://127.0.0.1:{}", dead.port()), ]) .env("HOME", home.path()) .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()); - let _child = cmd.spawn_guarded(); - wait_for_port(port); + .stdout(Stdio::null()); + let (_child, port) = spawn_listening(&mut cmd); let mut ws = ws_connect(port, Some("dead-callback")).await; ws_send( &mut ws, diff --git a/crates/agent/tests/mcp_skills_approval.rs b/crates/agent/tests/mcp_skills_approval.rs index 4454207f..5f151c7c 100644 --- a/crates/agent/tests/mcp_skills_approval.rs +++ b/crates/agent/tests/mcp_skills_approval.rs @@ -346,3 +346,74 @@ fn serve_asks_per_call_before_a_cross_origin_read() { let bodies = turn_bodies(&bodies.lock().unwrap()); assert!(tool_result(&bodies, "NOTES-BODY").contains("NOTES-BODY")); } + +/// A skill a subagent chooses is asked about in that subagent's name — the same provenance its +/// code-execution questions carry — not as the session's own agent. +#[test] +fn a_subagents_skill_load_is_asked_about_in_its_own_name() { + let env = Env::new(json!({})); + let agents = env.home.join(".claude/agents"); + std::fs::create_dir_all(&agents).unwrap(); + std::fs::write( + agents.join("helper.md"), + "---\nname: helper\ndescription: helps\n---\nYou are HELPER-MARKER.\n", + ) + .unwrap(); + let (mut serve, _bodies) = serve_with( + &env, + vec![ + // The parent, once the child has answered. + ("CHILD-LOADED", turn_text("parent done")), + // The child, after its load (allowed or not). + ("GIT-WORKFLOW-BODY", turn_text("CHILD-LOADED")), + ("did not approve loading", turn_text("CHILD-LOADED")), + // The child's first turn: it chooses the skill. + ("HELPER-MARKER", read_skill("c1", GIT)), + ( + "delegate-load", + common::turn_tool_use( + "d1", + "subagent", + &json!({ "agent": "helper", "task": "load the git skill" }).to_string(), + ), + ), + ], + ); + serve.send(json!({ "type": "prompt", "message": "delegate-load" })); + let q = next_question(&mut serve); + assert_eq!(q["mcp_skill"]["purpose"], "activate", "{q}"); + assert_eq!( + q["origin"]["agent"], "helper", + "the question must name the subagent that chose the skill: {q}" + ); + assert!(q["origin"]["spawn_id"].is_string(), "{q}"); + serve.approve(&q, "allow", "once"); + serve.read_until(|f| is_response(f, "prompt")); + serve.finish(); +} + +/// A `/skill:` the user typed is consent for that one load, not a standing approval: when the model +/// loads the same skill later, the user is asked. +#[test] +fn a_typed_skill_invocation_does_not_approve_the_models_later_loads() { + let env = Env::new(json!({})); + let (mut serve, _bodies) = serve_with( + &env, + vec![ + ("\"tool_use_id\":\"t1\"", turn_text("done")), + ("load-it-again", read_skill("t1", GIT)), + ], + ); + let frames = serve.call( + json!({ "type": "prompt", "message": "/skill:docs:git-workflow load it" }), + "prompt", + ); + assert_eq!(frames.last().unwrap()["success"], true, "{frames:#?}"); + serve.send(json!({ "type": "prompt", "message": "load-it-again" })); + let q = next_question(&mut serve); + assert_eq!(q["mcp_skill"]["purpose"], "activate", "{q}"); + assert_eq!(q["mcp_skill"]["uri"], GIT, "{q}"); + serve.approve(&q, "deny", "once"); + serve.read_until(|f| is_response(f, "prompt")); + serve.finish(); +} diff --git a/crates/agent/tests/mcp_skills_listing.rs b/crates/agent/tests/mcp_skills_listing.rs index 8b7a81e0..3c2b64f8 100644 --- a/crates/agent/tests/mcp_skills_listing.rs +++ b/crates/agent/tests/mcp_skills_listing.rs @@ -3,6 +3,8 @@ //! - A skill published after connect gets listed: on a list-changed notification, or once the //! listing's `ttlMs` runs out — and not before (`ttlMs` is honoured, not polled past). //! - A `cacheScope: "private"` listing is never written to the on-disk manifest cache. +//! - A listing re-fetched after connect is written back to that cache at once (or, now private, +//! forgotten), so a restart does not advertise the listing the server replaced. //! - Skills over the spec's per-skill limits (512 files / 16 MiB) are declined, and every entry left //! out — over-limit, dynamic — is surfaced through the skills diagnostics (`get_commands`' //! `collisions`, `run`'s warnings), as is a name one server lists twice. @@ -184,3 +186,65 @@ fn entries_left_out_are_surfaced_in_get_commands_and_run_warnings() { out.stderr ); } + +fn cached_manifest(env: &Env) -> String { + std::fs::read_to_string(env.home.join(".claude/mcp-manifest.json")).unwrap_or_default() +} + +/// A listing re-fetched mid-session reaches the on-disk cache then, not at the next live connect: +/// a restart that answers from the cache offers the skill published meanwhile. +#[test] +fn a_refreshed_listing_is_written_back_to_the_manifest_cache() { + let env = Env::new(json!({ + "MCP_SKILLS_FIXTURE_LATE_FLAG": "LATE", + "MCP_SKILLS_FIXTURE_TTL_MS": "0", + })); + let flag = env.cwd.join("LATE"); + let (mut serve, _bodies) = routed_serve(&env, vec![]); + prompt(&mut serve, "first-turn"); + assert!(cached_manifest(&env).contains("skill://git-workflow/SKILL.md")); + assert!(!cached_manifest(&env).contains("skill://late/SKILL.md")); + std::fs::write(&flag, "published").unwrap(); + prompt(&mut serve, "second-turn"); + // Written while the session is still up, by the re-list itself. + assert!( + cached_manifest(&env).contains("skill://late/SKILL.md"), + "the refreshed listing must be written back: {}", + cached_manifest(&env) + ); + serve.finish(); + + // And a restart answering from the cache (it starts no server) offers it. + env.clear_log(); + let (mut serve, bodies) = routed_serve(&env, vec![]); + prompt(&mut serve, "after-restart"); + serve.finish(); + assert!( + !env.log().contains(&"start -".to_string()), + "the restart answered from the cache: {:?}", + env.log() + ); + let bodies = bodies.lock().unwrap().clone(); + assert!(system_of(&bodies, "after-restart").contains("docs:late")); +} + +/// A listing that turns private on a re-fetch is dropped from the cache, as one private at connect +/// is never written. +#[test] +fn a_listing_that_turns_private_on_refresh_is_forgotten_by_the_cache() { + let env = Env::new(json!({ + "MCP_SKILLS_FIXTURE_PRIVATE_FLAG": "PRIVATE", + "MCP_SKILLS_FIXTURE_TTL_MS": "0", + })); + let (mut serve, _bodies) = routed_serve(&env, vec![]); + prompt(&mut serve, "first-turn"); + assert!(cached_manifest(&env).contains("skill://git-workflow/SKILL.md")); + std::fs::write(env.cwd.join("PRIVATE"), "now").unwrap(); + prompt(&mut serve, "second-turn"); + serve.finish(); + assert!( + !cached_manifest(&env).contains("skill://"), + "a now-private listing must not stay on disk: {}", + cached_manifest(&env) + ); +} diff --git a/crates/agent/tests/mcp_skills_resume.rs b/crates/agent/tests/mcp_skills_resume.rs index a06a6afa..5e223165 100644 --- a/crates/agent/tests/mcp_skills_resume.rs +++ b/crates/agent/tests/mcp_skills_resume.rs @@ -4,12 +4,15 @@ //! got here: resumed by a later `run --session-id`, switched back to in `serve`, forked, or reopened //! after `serve` restarted. The code-execution gate must be on for it each time; the spec lets the //! window be longer than the skill's time in context, never shorter. +//! +//! The user's approvals do not come back with it: they last the session in memory and are never +//! read from disk, where a model with file-write tools could plant one. #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] mod common; -use common::skills_env::{Env, Serve, bash, is_response, last_message_text}; +use common::skills_env::{Env, Serve, bash, is_response, last_message_text, read_skill}; use common::{spawn_model_server_routed, turn_text}; use serde_json::{Value, json}; @@ -167,3 +170,102 @@ fn a_restarted_serve_reopening_the_session_is_gated() { expect_gated(&mut serve); serve.finish(); } + +/// The model's next activation question, or `None` if the prompt finished without one. +fn activation_question(serve: &mut Serve) -> Option { + let frames = serve.read_until(|f| f["type"] == "approval_request" || is_response(f, "prompt")); + let last = frames.last().unwrap().clone(); + (last["type"] == "approval_request").then_some(last) +} + +/// Deny the question and let the prompt finish. +fn deny(serve: &mut Serve, q: &Value) { + assert_eq!(q["mcp_skill"]["purpose"], "activate", "{q}"); + serve.approve(q, "deny", "once"); + serve.read_until(|f| is_response(f, "prompt")); +} + +/// Approvals last the session, in memory, and nothing on disk or from a client can grant one. The +/// session file is the user's, writable by a model with file-write tools, so an approval read back +/// from it would be one the model could have planted; and `append_custom` refuses the kind. (The +/// audit's forgery probe, adopted.) +#[test] +fn an_approval_lasts_the_session_and_cannot_be_planted_on_disk_or_by_a_client() { + let env = Env::new(json!({})); + const GIT: &str = "skill://git-workflow/SKILL.md"; + // Each prompt has the model load the skill once, under its own tool-call id. + let routes = (1..=5) + .rev() + .flat_map(|n| { + [ + (format!("\"tool_use_id\":\"t{n}\""), turn_text("done")), + (format!("load-{n}"), read_skill(&format!("t{n}"), GIT)), + ] + }) + .collect(); + let (base, _bodies) = spawn_model_server_routed(routes, turn_text("ok")); + let session_file = env.dir.path().join("approvals.jsonl"); + + // 1. Asked once and approved for the session: the same content is not asked about again in it. + let mut serve = env.serve_on(&base, &session_file, &[]); + serve.send(json!({ "type": "prompt", "message": "load-1" })); + let q = activation_question(&mut serve).expect("the first load is asked about"); + let fingerprint = q["mcp_skill"]["manifest"].as_str().unwrap().to_string(); + let session_id = session_id(&mut serve); + serve.approve(&q, "allow", "session"); + serve.read_until(|f| is_response(f, "prompt")); + serve.send(json!({ "type": "prompt", "message": "load-2" })); + assert_eq!( + activation_question(&mut serve), + None, + "approved for the session" + ); + serve.finish(); + + // 2. Nothing was written down, so a new process asks again. + let text = std::fs::read_to_string(&session_file).unwrap(); + assert!(!text.contains("mcp_skill_approval"), "{text}"); + let mut serve = env.serve_on(&base, &session_file, &[]); + serve.send(json!({ "type": "prompt", "message": "load-3" })); + let q = activation_question(&mut serve).expect("a restart asks again"); + deny(&mut serve, &q); + serve.finish(); + + // 3. An approval written into the session file (what a model's `write` could add) is ignored. + let text = std::fs::read_to_string(&session_file).unwrap(); + let last: Value = serde_json::from_str(text.lines().last().unwrap()).unwrap(); + let forged = json!({ + "type": "custom", + "id": "forged0001", + "parent_id": last["id"], + "timestamp": last["timestamp"], + "kind": "mcp_skill_approval", + "data": { + "sessionId": session_id, + "key": format!("activate\u{0}docs\u{0}{GIT}\u{0}{fingerprint}"), + "allow": true, + }, + }); + std::fs::write(&session_file, format!("{text}{forged}\n")).unwrap(); + let mut serve = env.serve_on(&base, &session_file, &[]); + serve.send(json!({ "type": "prompt", "message": "load-4" })); + let q = activation_question(&mut serve).expect("a forged approval in the file grants nothing"); + deny(&mut serve, &q); + + // 4. Nor can a client plant one: the kind is the agent's own. + let frames = serve.call( + json!({ "type": "append_custom", "kind": "mcp_skill_approval", "data": { + "sessionId": session_id, + "key": format!("activate\u{0}docs\u{0}{GIT}\u{0}{fingerprint}"), + "allow": true, + } }), + "append_custom", + ); + let r = frames.last().unwrap(); + assert_eq!(r["success"], false, "{r}"); + assert!(r["error"].as_str().unwrap().contains("reserved"), "{r}"); + serve.send(json!({ "type": "prompt", "message": "load-5" })); + let q = activation_question(&mut serve).expect("a client-appended approval grants nothing"); + deny(&mut serve, &q); + serve.finish(); +} diff --git a/crates/agent/tests/mcp_skills_sessions.rs b/crates/agent/tests/mcp_skills_sessions.rs index cac7b3a4..41cd930f 100644 --- a/crates/agent/tests/mcp_skills_sessions.rs +++ b/crates/agent/tests/mcp_skills_sessions.rs @@ -13,39 +13,37 @@ use std::process::{Command, Stdio}; use common::skills_env::{BIN, Env, bash, last_message_text, read_skill, system_text}; use common::{ - ChildGuard, SpawnGuarded, free_port, spawn_model_server, spawn_model_server_routed, turn_text, - turn_tool_use, wait_for_port, ws_connect, ws_read_until_response, ws_send, + ChildGuard, SpawnGuarded, spawn_listening, spawn_model_server, spawn_model_server_routed, + turn_text, turn_tool_use, ws_connect, ws_read_until_response, ws_send, }; use serde_json::json; const GIT: &str = "skill://git-workflow/SKILL.md"; const GUIDE: &str = "skill://git-workflow/references/GUIDE.md"; -fn daemon(env: &Env, base: &str, port: u16) -> ChildGuard { +fn daemon(env: &Env, base: &str) -> (ChildGuard, u16) { let sessions = env.dir.path().join("sessions"); - Command::new(BIN) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - sessions.to_str().unwrap(), - // Approvals are not what this proves; the loaded-skill state is. - "--approve-mcp-skills", - ]) - .env("HOME", &env.home) - .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") - .current_dir(&env.cwd) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() + spawn_listening( + Command::new(BIN) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + sessions.to_str().unwrap(), + // Approvals are not what this proves; the loaded-skill state is. + "--approve-mcp-skills", + ]) + .env("HOME", &env.home) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .current_dir(&env.cwd) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } #[tokio::test(flavor = "multi_thread")] @@ -67,9 +65,7 @@ async fn a_daemon_session_cannot_read_a_skill_another_session_loaded() { ], turn_text("fallback"), ); - let port = free_port(); - let _daemon = daemon(&env, &base, port); - wait_for_port(port); + let (_daemon, port) = daemon(&env, &base); let mut a = ws_connect(port, Some("skills-alpha")).await; let mut b = ws_connect(port, Some("skills-bravo")).await; diff --git a/crates/agent/tests/mcp_tasks_daemon.rs b/crates/agent/tests/mcp_tasks_daemon.rs index 988455da..d8452575 100644 --- a/crates/agent/tests/mcp_tasks_daemon.rs +++ b/crates/agent/tests/mcp_tasks_daemon.rs @@ -13,33 +13,31 @@ use std::process::{Command, Stdio}; use std::time::Duration; use common::{ - ChildGuard, SpawnGuarded, TestWs, free_port, spawn_model_server_routed, turn_text, - turn_tool_use, wait_for_port, ws_connect, ws_next_frame, ws_send, + ChildGuard, TestWs, spawn_listening, spawn_model_server_routed, turn_text, turn_tool_use, + ws_connect, ws_next_frame, ws_send, }; use mcp_tasks_env::Env; use serde_json::{Value, json}; -fn daemon(env: &Env, base: &str, port: u16) -> ChildGuard { - Command::new(common::BIN) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - env.dir.path().join("sessions").to_str().unwrap(), - ]) - .env("HOME", &env.home) - .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn daemon(env: &Env, base: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(common::BIN) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + env.dir.path().join("sessions").to_str().unwrap(), + ]) + .env("HOME", &env.home) + .env("BEYOND_AI_AGENT_MCP_IDLE_SECS", "0") + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// Frames until `pred` matches (that frame last), failing after 30 s. @@ -128,9 +126,7 @@ async fn daemon_routes_mcp_questions_to_the_session_whose_call_asked() { ], turn_text("fallback"), ); - let port = free_port(); - let _daemon = daemon(&env, &base, port); - wait_for_port(port); + let (_daemon, port) = daemon(&env, &base); let mut a = ws_connect(port, Some("mcp-alpha")).await; // Bravo starts second: under a process-wide hub its gates would answer for alpha too. @@ -270,9 +266,7 @@ fn concurrent_routes() -> Vec<(String, String)> { /// is still in flight then. Returns both sessions' frames for the next 3.5 s. async fn alpha_asks_while_bravo_holds(env: &Env) -> (TestWs, Vec, Vec, ChildGuard) { let (base, _requests) = spawn_model_server_routed(concurrent_routes(), turn_text("fallback")); - let port = free_port(); - let daemon = daemon(env, &base, port); - wait_for_port(port); + let (daemon, port) = daemon(env, &base); let mut a = ws_connect(port, Some("route-alpha")).await; let mut b = ws_connect(port, Some("route-bravo")).await; ws_send( diff --git a/crates/agent/tests/serve_approval.rs b/crates/agent/tests/serve_approval.rs index 2790cb00..acbd7e5c 100644 --- a/crates/agent/tests/serve_approval.rs +++ b/crates/agent/tests/serve_approval.rs @@ -13,8 +13,8 @@ use std::process::{ChildStdin, Command, Stdio}; use std::time::Duration; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, read_until_response, spawn_model_server, - turn_text, turn_tool_use, wait_for_port, ws_connect, ws_next_frame, ws_send, + ChildGuard, ISOLATED_HOME, SpawnGuarded, read_until_response, spawn_listening, + spawn_model_server, turn_text, turn_tool_use, ws_connect, ws_next_frame, ws_send, }; use serde_json::{Value, json}; @@ -587,27 +587,25 @@ fn two_gated_calls_in_one_turn_are_asked_about_one_at_a_time() { // ---- multi-attach, over the real WebSocket transport ------------------------------------------------ -fn serve_ws_approving(base: &str, session_dir: &str, port: u16, extra: &[&str]) -> ChildGuard { - Command::new(BIN) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - ]) - .args(extra) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws_approving(base: &str, session_dir: &str, extra: &[&str]) -> (ChildGuard, u16) { + spawn_listening( + Command::new(BIN) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + ]) + .args(extra) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// How long any single frame may take to arrive before the test calls the socket stuck. @@ -694,14 +692,11 @@ async fn every_attached_client_sees_the_question_and_the_first_answer_wins() { write_turn("tu_1", target.to_str().unwrap()), turn_text("wrote it"), ]); - let port = free_port(); - let child = serve_ws_approving( + let (child, port) = serve_ws_approving( &base, dir.path().to_str().unwrap(), - port, &["--approve", "writes"], ); - wait_for_port(port); // Phone and TUI on the same session — they coexist, neither evicts the other. const SID: &str = "approvalshared1"; @@ -754,15 +749,12 @@ async fn a_question_with_no_client_left_to_answer_it_is_denied_rather_than_hung( write_turn("tu_1", target.to_str().unwrap()), turn_text("could not write"), ]); - let port = free_port(); - let child = serve_ws_approving( + let (child, port) = serve_ws_approving( &base, dir.path().to_str().unwrap(), - port, // Wait forever: only the no-client check can free this. &["--approve", "writes", "--approval-timeout", "0"], ); - wait_for_port(port); const SID: &str = "approvaldetach1"; let mut ws = Client::new(ws_connect(port, Some(SID)).await); diff --git a/crates/agent/tests/serve_drain.rs b/crates/agent/tests/serve_drain.rs index 8b3d56cd..45bb1201 100644 --- a/crates/agent/tests/serve_drain.rs +++ b/crates/agent/tests/serve_drain.rs @@ -18,8 +18,8 @@ use std::process::{Command, Stdio}; use std::time::{Duration, Instant}; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, turn_text, - turn_tool_use, wait_for_port, ws_connect, ws_next_frame, + ChildGuard, ISOLATED_HOME, spawn_listening, spawn_model_server, turn_text, turn_tool_use, + ws_connect, ws_next_frame, }; use futures::SinkExt; use serde_json::{Value, json}; @@ -29,29 +29,27 @@ use tokio_tungstenite::tungstenite::Message; const BIN: &str = env!("CARGO_BIN_EXE_beyond-ai-agent"); -fn serve_ws(base: &str, session_dir: &str, port: u16, drain_grace: &str) -> ChildGuard { - Command::new(BIN) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - "--session-idle-timeout", - "0", - "--drain-grace", - drain_grace, - ]) - .env("HOME", ISOLATED_HOME) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws(base: &str, session_dir: &str, drain_grace: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(BIN) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + "--session-idle-timeout", + "0", + "--drain-grace", + drain_grace, + ]) + .env("HOME", ISOLATED_HOME) + .stdout(Stdio::null()), + ) } /// One bare HTTP GET, returning (status, body). @@ -119,9 +117,7 @@ async fn a_draining_replica_refuses_new_sessions_but_keeps_serving_the_ones_it_o ), turn_text("done"), ]); - let port = free_port(); - let child = serve_ws(&base, dir.path().to_str().unwrap(), port, "30"); - wait_for_port(port); + let (child, port) = serve_ws(&base, dir.path().to_str().unwrap(), "30"); let pid = child.id(); // A session this replica owns, mid-run. @@ -207,9 +203,7 @@ async fn a_drain_lets_an_in_flight_run_finish_and_persists_it() { ), turn_text("finished after the signal"), ]); - let port = free_port(); - let mut child = serve_ws(&base, dir.path().to_str().unwrap(), port, "30"); - wait_for_port(port); + let (mut child, port) = serve_ws(&base, dir.path().to_str().unwrap(), "30"); let pid = child.id(); let mut ws = ws_connect(port, Some("drain-inflight")).await; @@ -267,9 +261,7 @@ async fn without_a_grace_a_signal_still_shuts_down_at_once() { // Ctrl-C on a laptop daemon ends it now, not in thirty seconds. let dir = tempfile::tempdir().unwrap(); let (base, _bodies) = spawn_model_server(vec![]); - let port = free_port(); - let mut child = serve_ws(&base, dir.path().to_str().unwrap(), port, "0"); - wait_for_port(port); + let (mut child, _) = serve_ws(&base, dir.path().to_str().unwrap(), "0"); let pid = child.id(); let began = Instant::now(); diff --git a/crates/agent/tests/serve_health.rs b/crates/agent/tests/serve_health.rs index 5a1e674c..eb9ddf9d 100644 --- a/crates/agent/tests/serve_health.rs +++ b/crates/agent/tests/serve_health.rs @@ -13,8 +13,8 @@ use std::process::{Command, Stdio}; use common::service::{Options, Service}; use common::{ - BIN, ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, turn_text, - wait_for_port, ws_connect_with_headers, ws_read_until_response, ws_send, + BIN, ISOLATED_HOME, spawn_listening, spawn_model_server, turn_text, ws_connect_with_headers, + ws_read_until_response, ws_send, }; use serde_json::{Value, json}; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; @@ -302,27 +302,23 @@ async fn readyz_does_not_interfere_with_a_session_attach() { async fn health_exists_on_a_non_service_daemon() { let (base, _requests) = spawn_model_server(vec![]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let _child: ChildGuard = Command::new(BIN) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - &base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - &dir.path().to_string_lossy(), - ]) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded(); - wait_for_port(port); + let (_child, port) = spawn_listening( + Command::new(BIN) + .args([ + "serve", + "--gateway-url", + &base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + &dir.path().to_string_lossy(), + ]) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ); let (status, response) = get(port, "/livez", &[]).await; assert_eq!(status, 200, "{response}"); diff --git a/crates/agent/tests/serve_http.rs b/crates/agent/tests/serve_http.rs index c948fc82..b0934c2a 100644 --- a/crates/agent/tests/serve_http.rs +++ b/crates/agent/tests/serve_http.rs @@ -12,33 +12,31 @@ use std::process::{Command, Stdio}; use std::time::{Duration, Instant}; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, WS_PATH, free_port, spawn_model_server, turn_text, - wait_for_port, ws_connect, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, WS_PATH, spawn_listening, spawn_model_server, turn_text, ws_connect, + ws_read_until_response, ws_send, }; use serde_json::{Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; -fn serve_ws_child(base: &str, session_dir: &str, port: u16) -> ChildGuard { - Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - ]) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws_child(base: &str, session_dir: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + ]) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } struct HttpReply { @@ -159,9 +157,7 @@ async fn wait_for_assistant(port: u16, session_id: &str) -> Value { async fn http_post_get_state_mints_a_session_and_returns_the_protocol_frame() { let (base, _requests) = spawn_model_server(vec![turn_text("unused")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let reply = post_cmd(port, None, json!({ "type": "get_state", "id": "s1" })).await; assert_eq!(reply.status, 200, "{}", reply.body); @@ -195,9 +191,7 @@ async fn http_post_get_state_mints_a_session_and_returns_the_protocol_frame() { async fn http_post_prompt_is_accepted_and_the_run_outlives_the_request() { let (base, _requests) = spawn_model_server(vec![turn_text("hello from post")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let reply = post_cmd( port, @@ -237,9 +231,7 @@ async fn http_post_prompt_is_accepted_and_the_run_outlives_the_request() { async fn http_post_prompt_without_message_is_a_protocol_error_not_accepted() { let (base, _requests) = spawn_model_server(vec![turn_text("unused")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let reply = post_cmd( port, @@ -268,9 +260,7 @@ async fn http_post_prompt_without_message_is_a_protocol_error_not_accepted() { async fn http_post_list_daemon_sessions_does_not_require_a_session() { let (base, _requests) = spawn_model_server(vec![turn_text("one")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let created = post_cmd( port, @@ -317,9 +307,7 @@ async fn http_post_list_daemon_sessions_does_not_require_a_session() { async fn a_refused_request_still_gets_its_answer_with_a_large_body_in_flight() { let (base, _requests) = spawn_model_server(vec![turn_text("unused")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let body = vec![b'x'; 256 * 1024]; @@ -336,9 +324,7 @@ async fn a_refused_request_still_gets_its_answer_with_a_large_body_in_flight() { async fn http_wrong_path_is_404_and_wrong_method_is_405() { let (base, _requests) = spawn_model_server(vec![turn_text("unused")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let (status, _headers, _body) = http_exchange(port, "POST", "/nope", Some(b"{\"type\":\"get_state\"}")).await; @@ -369,9 +355,7 @@ async fn http_wrong_path_is_404_and_wrong_method_is_405() { async fn http_get_without_upgrade_is_426() { let (base, _requests) = spawn_model_server(vec![turn_text("unused")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let (status, headers, _body) = http_exchange(port, "GET", WS_PATH, None).await; assert_eq!(status, 426); diff --git a/crates/agent/tests/serve_list_sessions.rs b/crates/agent/tests/serve_list_sessions.rs index 227151b1..f676a00f 100644 --- a/crates/agent/tests/serve_list_sessions.rs +++ b/crates/agent/tests/serve_list_sessions.rs @@ -10,33 +10,31 @@ mod common; use std::process::{Command, Stdio}; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, turn_text, - wait_for_port, ws_connect, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, spawn_listening, spawn_model_server, turn_text, ws_connect, + ws_read_until_response, ws_send, }; use serde_json::{Value, json}; -/// Spawn a `serve --listen 127.0.0.1:` child against `base` (the mock gateway), persisting +/// Spawn a listening `serve` child (and its port) against `base` (the mock gateway), persisting /// per-session files under `session_dir`. -fn serve_ws_child(base: &str, session_dir: &str, port: u16) -> ChildGuard { - Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - ]) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws_child(base: &str, session_dir: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + ]) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// Prompt a fresh session `id` to completion so it persists under the daemon's `--session-dir`, then @@ -93,9 +91,7 @@ async fn list_daemon_sessions_reports_live_sessions_with_listing_fields() { // A turn apiece for the two sessions. let (base, _requests) = spawn_model_server(vec![turn_text("one"), turn_text("two")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); prompt_session(port, "aaaaaaaaaaaaaaaa").await; prompt_session(port, "bbbbbbbbbbbbbbbb").await; @@ -131,10 +127,7 @@ async fn list_daemon_sessions_reports_persisted_sessions_as_not_live_after_resta let (base, _requests) = spawn_model_server(vec![turn_text("one"), turn_text("two")]); let dir = tempfile::tempdir().unwrap(); let dir_str = dir.path().to_str().unwrap().to_string(); - let port = free_port(); - - let mut child = serve_ws_child(&base, &dir_str, port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, &dir_str); prompt_session(port, "aaaaaaaaaaaaaaaa").await; prompt_session(port, "bbbbbbbbbbbbbbbb").await; @@ -147,9 +140,7 @@ async fn list_daemon_sessions_reports_persisted_sessions_as_not_live_after_resta let _ = child.kill(); let _ = child.wait(); - let port2 = free_port(); - let mut child2 = serve_ws_child(&base, &dir_str, port2); - wait_for_port(port2); + let (mut child2, port2) = serve_ws_child(&base, &dir_str); let cold = list_daemon_sessions(port2).await; for id in ["aaaaaaaaaaaaaaaa", "bbbbbbbbbbbbbbbb"] { diff --git a/crates/agent/tests/serve_reaper.rs b/crates/agent/tests/serve_reaper.rs index 39ca6c1e..3b797d25 100644 --- a/crates/agent/tests/serve_reaper.rs +++ b/crates/agent/tests/serve_reaper.rs @@ -11,34 +11,32 @@ use std::process::{Command, Stdio}; use std::time::Duration; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, turn_text, - turn_tool_use, wait_for_port, ws_connect, ws_next_frame, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, spawn_listening, spawn_model_server, turn_text, turn_tool_use, + ws_connect, ws_next_frame, ws_read_until_response, ws_send, }; use serde_json::{Value, json}; /// Spawn a `serve --listen` child with the idle reaper armed at `idle_secs` (`0` disables reaping). -fn serve_reaper_child(base: &str, session_dir: &str, port: u16, idle_secs: u64) -> ChildGuard { - Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - "--session-idle-timeout", - &idle_secs.to_string(), - ]) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_reaper_child(base: &str, session_dir: &str, idle_secs: u64) -> (ChildGuard, u16) { + spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + "--session-idle-timeout", + &idle_secs.to_string(), + ]) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// Send `list_daemon_sessions` on a fresh connection and return the `data.sessions` array. @@ -64,9 +62,7 @@ async fn idle_session_is_reaped_but_its_turn_persists() { let (base, _requests) = spawn_model_server(vec![turn_text("committedreply"), turn_text("again")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_reaper_child(&base, dir.path().to_str().unwrap(), port, 1); - wait_for_port(port); + let (mut child, port) = serve_reaper_child(&base, dir.path().to_str().unwrap(), 1); // Connect, prompt to completion (persists), then disconnect. { @@ -119,9 +115,7 @@ async fn zero_idle_timeout_disables_the_reaper() { const SID: &str = "neverreaped001"; let (base, _requests) = spawn_model_server(vec![turn_text("pinnedreply")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_reaper_child(&base, dir.path().to_str().unwrap(), port, 0); - wait_for_port(port); + let (mut child, port) = serve_reaper_child(&base, dir.path().to_str().unwrap(), 0); { let mut ws = ws_connect(port, Some(SID)).await; @@ -159,9 +153,7 @@ async fn attached_session_past_timeout_is_not_reaped() { const SID: &str = "attachedsession1"; let (base, _requests) = spawn_model_server(vec![turn_text("stillhere")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_reaper_child(&base, dir.path().to_str().unwrap(), port, 1); - wait_for_port(port); + let (mut child, port) = serve_reaper_child(&base, dir.path().to_str().unwrap(), 1); // Attach and prompt to completion, but KEEP the connection open. let mut ws = ws_connect(port, Some(SID)).await; @@ -210,9 +202,7 @@ async fn detached_but_mid_run_session_is_not_reaped() { turn_text("finishedmidrun"), ]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_reaper_child(&base, dir.path().to_str().unwrap(), port, 1); - wait_for_port(port); + let (mut child, port) = serve_reaper_child(&base, dir.path().to_str().unwrap(), 1); // Start the run; wait until the bash tool is actually running, then drop the connection (detached, // but the run keeps going on the server). diff --git a/crates/agent/tests/serve_reattach_catchup.rs b/crates/agent/tests/serve_reattach_catchup.rs index 6797ecab..e5aab42a 100644 --- a/crates/agent/tests/serve_reattach_catchup.rs +++ b/crates/agent/tests/serve_reattach_catchup.rs @@ -21,34 +21,32 @@ mod common; use std::process::{Command, Stdio}; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, turn_text, - turn_tool_use, wait_for_port, ws_connect, ws_next_frame, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, spawn_listening, spawn_model_server, turn_text, turn_tool_use, + ws_connect, ws_next_frame, ws_read_until_response, ws_send, }; use serde_json::{Value, json}; /// Marker text the model commits in turn 1 — the history a re-attaching client must get back. const COMMITTED: &str = "COMMITTED_BEFORE_DROP"; -fn serve_ws_child(base: &str, session_dir: &str, port: u16) -> ChildGuard { - Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - ]) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws_child(base: &str, session_dir: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + ]) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// Drive a session to the state both tests need: one **committed** turn in history (containing @@ -59,16 +57,14 @@ fn serve_ws_child(base: &str, session_dir: &str, port: u16) -> ChildGuard { /// calls `bash sleep 5` (the in-flight window); its second turn ends the run. async fn session_with_history_and_a_run_in_flight( dir: &str, - port: u16, sid: &str, -) -> (ChildGuard, String) { +) -> (ChildGuard, String, u16) { let (base, _requests) = spawn_model_server(vec![ turn_text(COMMITTED), turn_tool_use("t1", "bash", &json!({ "command": "sleep 5" }).to_string()), turn_text("AFTER_THE_SLEEP"), ]); - let child = serve_ws_child(&base, dir, port); - wait_for_port(port); + let (child, port) = serve_ws_child(&base, dir); let mut ws1 = ws_connect(port, Some(sid)).await; @@ -107,7 +103,7 @@ async fn session_with_history_and_a_run_in_flight( assert!(in_flight, "the second prompt should reach the bash sleep"); drop(ws1); // the mobile network drops mid-run - (child, sid.to_string()) + (child, sid.to_string(), port) } /// A frame that carries the committed history back to the client — whatever shape it arrives in @@ -132,9 +128,8 @@ fn is_live_event(f: &Value) -> bool { async fn reattach_midrun_can_fetch_missed_history() { const SID: &str = "reattachcatchup1"; let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let (mut child, sid) = - session_with_history_and_a_run_in_flight(dir.path().to_str().unwrap(), port, SID).await; + let (mut child, sid, port) = + session_with_history_and_a_run_in_flight(dir.path().to_str().unwrap(), SID).await; let mut ws2 = ws_connect(port, Some(&sid)).await; ws_send(&mut ws2, json!({ "type": "get_messages" })).await; @@ -171,9 +166,8 @@ async fn reattach_midrun_can_fetch_missed_history() { async fn reattach_delivers_history_before_live_frames() { const SID: &str = "reattachcatchup2"; let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let (mut child, sid) = - session_with_history_and_a_run_in_flight(dir.path().to_str().unwrap(), port, SID).await; + let (mut child, sid, port) = + session_with_history_and_a_run_in_flight(dir.path().to_str().unwrap(), SID).await; // Re-attach and just *listen*, issuing no command at all. Everything that arrives, arrives because // the server sent it — which is the point: catch-up must not be a thing the client has to ask for @@ -210,9 +204,8 @@ async fn reattach_delivers_history_before_live_frames() { async fn reattach_catchup_frame_has_the_get_messages_payload_shape() { const SID: &str = "reattachcatchup3"; let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let (mut child, sid) = - session_with_history_and_a_run_in_flight(dir.path().to_str().unwrap(), port, SID).await; + let (mut child, sid, port) = + session_with_history_and_a_run_in_flight(dir.path().to_str().unwrap(), SID).await; let mut ws2 = ws_connect(port, Some(&sid)).await; let first = ws_next_frame(&mut ws2).await.expect("a first frame"); @@ -252,11 +245,10 @@ async fn reattach_catchup_frame_has_the_get_messages_payload_shape() { async fn reattach_replays_the_in_flight_turn_not_just_committed_history() { const SID: &str = "reattachcatchup4"; let dir = tempfile::tempdir().unwrap(); - let port = free_port(); // The driving connection saw `tool_start{bash}` before it dropped (that is how the helper knows the // run is in flight) — so that frame is, by construction, one this reconnecting client missed. - let (mut child, sid) = - session_with_history_and_a_run_in_flight(dir.path().to_str().unwrap(), port, SID).await; + let (mut child, sid, port) = + session_with_history_and_a_run_in_flight(dir.path().to_str().unwrap(), SID).await; let mut ws2 = ws_connect(port, Some(&sid)).await; diff --git a/crates/agent/tests/serve_service_grant.rs b/crates/agent/tests/serve_service_grant.rs index 56d2ddb4..6977688f 100644 --- a/crates/agent/tests/serve_service_grant.rs +++ b/crates/agent/tests/serve_service_grant.rs @@ -280,18 +280,12 @@ async fn the_grant_header_is_not_echoed_in_the_handshake_response() { async fn a_daemon_without_service_mode_ignores_the_grant_header() { let (base, _requests) = spawn_model_server(vec![]); let dir = tempfile::tempdir().unwrap(); - let port = common::free_port(); let mut cmd = common::serve_dir_cmd( common::BIN, &base, &dir.path().join("sessions").to_string_lossy(), ); - cmd.arg("--listen").arg(format!("127.0.0.1:{port}")); - let _child = { - use common::SpawnGuarded as _; - cmd.spawn_guarded() - }; - common::wait_for_port(port); + let (_child, port) = common::spawn_listening(&mut cmd); let mut ws = ws_connect_with_headers(port, Some("plain"), &[("x-beyond-grant", "nonsense")]) .await diff --git a/crates/agent/tests/serve_service_mode.rs b/crates/agent/tests/serve_service_mode.rs index 25f1ff55..0e970d20 100644 --- a/crates/agent/tests/serve_service_mode.rs +++ b/crates/agent/tests/serve_service_mode.rs @@ -9,13 +9,13 @@ mod common; -use std::io::Read; use std::path::Path; use std::process::{Command, Stdio}; +use std::time::Duration; use common::service::Service; use common::{ - BIN, SpawnGuarded, free_port, spawn_model_server, turn_text, ws_connect_with_headers, + BIN, SpawnGuarded, spawn_model_server, turn_text, ws_connect_with_headers, ws_read_until_response, ws_send, }; use serde_json::{Value, json}; @@ -58,7 +58,8 @@ fn keys() -> Keys { } /// The minimal valid `--service` command line, for a test to spoil one piece of. -fn service_cmd(k: &Keys, port: u16) -> Command { +/// Every use is a startup that must fail; were one to succeed, port 0 is still the kernel's to pick. +fn service_cmd(k: &Keys) -> Command { let mut c = Command::new(BIN); c.args([ "serve", @@ -68,7 +69,7 @@ fn service_cmd(k: &Keys, port: u16) -> Command { "--model", "claude-test", "--listen", - &format!("127.0.0.1:{port}"), + "127.0.0.1:0", "--grant-key", &k.flag, "--seal-key", @@ -83,7 +84,6 @@ fn service_cmd(k: &Keys, port: u16) -> Command { #[test] fn service_mode_demands_a_verifier_a_shard_and_a_listener() { let k = keys(); - let port = free_port(); let base: Vec = vec![ "serve".into(), @@ -91,7 +91,7 @@ fn service_mode_demands_a_verifier_a_shard_and_a_listener() { "--model".into(), "claude-test".into(), ]; - let listen = ["--listen".to_string(), format!("127.0.0.1:{port}")]; + let listen = ["--listen".to_string(), "127.0.0.1:0".to_string()]; let verifier = [ "--grant-key".to_string(), k.flag.clone(), @@ -124,7 +124,7 @@ fn service_mode_demands_a_verifier_a_shard_and_a_listener() { fn a_shard_must_be_a_name_and_an_absolute_path() { let k = keys(); for bad in ["a", "a=relative/path", "a.b=/tmp", "=/tmp"] { - let mut c = service_cmd(&k, free_port()); + let mut c = service_cmd(&k); // Append a second `--shard`; the bad one is what fails. c.arg("--shard").arg(bad); let err = startup_error(&mut c); @@ -159,7 +159,7 @@ fn every_refused_flag_fails_at_startup() { vec!["--bash-shell-path".into(), "/bin/sh".into()], ]; for flag in &refused { - let mut c = service_cmd(&k, free_port()); + let mut c = service_cmd(&k); c.args(flag); let err = startup_error(&mut c); assert!( @@ -174,7 +174,7 @@ fn every_refused_flag_fails_at_startup() { #[test] fn an_ambient_agent_key_refuses_startup() { let k = keys(); - let mut c = service_cmd(&k, free_port()); + let mut c = service_cmd(&k); c.env("AI_AGENT_KEY", "bai_v1.the-replicas-own-key"); let err = startup_error(&mut c); assert!(err.contains("--key") && err.contains("--service"), "{err}"); @@ -186,7 +186,7 @@ fn an_ambient_agent_key_refuses_startup() { #[cfg(feature = "code-mode")] fn code_mode_is_refused() { let k = keys(); - let mut c = service_cmd(&k, free_port()); + let mut c = service_cmd(&k); c.arg("--code-mode"); let err = startup_error(&mut c); assert!(err.contains("--code-mode"), "{err}"); @@ -409,10 +409,18 @@ async fn the_replica_announces_service_mode_and_its_shards() { .unwrap(), ); let _ = svc.child.kill(); - let mut stderr = String::new(); - if let Some(mut pipe) = svc.child.stderr.take() { - let _ = pipe.read_to_string(&mut stderr); - } + let _ = svc.child.wait(); + // Copied to the log line by line as the replica writes it; give the copy a moment to catch up. + let deadline = std::time::Instant::now() + Duration::from_secs(10); + let stderr = loop { + let stderr = std::fs::read_to_string(&svc.stderr_log).unwrap_or_default(); + let complete = + stderr.contains("service mode") && stderr.contains("s1") && stderr.contains("s2"); + if complete || std::time::Instant::now() > deadline { + break stderr; + } + tokio::time::sleep(Duration::from_millis(20)).await; + }; assert!(stderr.contains("service mode"), "{stderr}"); assert!(stderr.contains("s1") && stderr.contains("s2"), "{stderr}"); } diff --git a/crates/agent/tests/serve_session_density.rs b/crates/agent/tests/serve_session_density.rs index 5bb2f4c7..38dab661 100644 --- a/crates/agent/tests/serve_session_density.rs +++ b/crates/agent/tests/serve_session_density.rs @@ -18,36 +18,34 @@ use std::process::{Command, Stdio}; use std::time::{Duration, Instant}; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, TestWs, free_port, spawn_model_server, turn_text, - wait_for_port, ws_connect, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, TestWs, spawn_listening, spawn_model_server, turn_text, ws_connect, + ws_read_until_response, ws_send, }; use serde_json::json; const BIN: &str = env!("CARGO_BIN_EXE_beyond-ai-agent"); -fn serve_ws(base: &str, session_dir: &str, port: u16) -> ChildGuard { - Command::new(BIN) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - // Keep every session for the measurement window; we hold the sockets anyway. - "--session-idle-timeout", - "0", - ]) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws(base: &str, session_dir: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(BIN) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + // Keep every session for the measurement window; we hold the sockets anyway. + "--session-idle-timeout", + "0", + ]) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } #[derive(Clone, Copy, Debug)] @@ -143,9 +141,7 @@ async fn ping_all(sockets: &mut [TestWs], rounds: usize) -> Vec { async fn idle_sessions_do_not_take_a_thread_each() { let (base, _r) = spawn_model_server(vec![turn_text("unused")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws(&base, dir.path().to_str().unwrap()); let pid = child.id(); // One throwaway connection so the daemon has completed its first accept/handshake path. @@ -207,9 +203,7 @@ async fn idle_sessions_do_not_take_a_thread_each() { async fn density_sweep() { let (base, _r) = spawn_model_server(vec![turn_text("unused")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws(&base, dir.path().to_str().unwrap()); let pid = child.id(); let ticks_per_sec: u64 = 100; // Linux USER_HZ; RSS/thread counts are the load-bearing columns. diff --git a/crates/agent/tests/serve_session_isolation.rs b/crates/agent/tests/serve_session_isolation.rs index d954a283..3c2c150a 100644 --- a/crates/agent/tests/serve_session_isolation.rs +++ b/crates/agent/tests/serve_session_isolation.rs @@ -12,35 +12,32 @@ mod common; use std::process::{Command, Stdio}; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, - spawn_model_server_routed, turn_text, turn_tool_use, wait_for_port, ws_connect, - ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, spawn_listening, spawn_model_server, spawn_model_server_routed, + turn_text, turn_tool_use, ws_connect, ws_read_until_response, ws_send, }; use serde_json::{Value, json}; const BIN: &str = env!("CARGO_BIN_EXE_beyond-ai-agent"); -fn serve_ws(base: &str, session_dir: &str, port: u16, extra: &[&str]) -> ChildGuard { - Command::new(BIN) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - ]) - .args(extra) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws(base: &str, session_dir: &str, extra: &[&str]) -> (ChildGuard, u16) { + spawn_listening( + Command::new(BIN) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + ]) + .args(extra) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } fn dump(frames: &[Value]) -> String { @@ -82,9 +79,7 @@ async fn concurrent_sessions_do_not_leak_transcripts() { turn_text("unexpected-fallback"), ); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws(&base, dir.path().to_str().unwrap(), port, &[]); - wait_for_port(port); + let (mut child, port) = serve_ws(&base, dir.path().to_str().unwrap(), &[]); let mut a = ws_connect(port, Some("iso-alpha")).await; let mut b = ws_connect(port, Some("iso-bravo")).await; @@ -181,9 +176,7 @@ async fn concurrent_sessions_do_not_leak_exec_endpoints() { turn_text("second"), ]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws(&base, dir.path().to_str().unwrap(), port, &[]); - wait_for_port(port); + let (mut child, port) = serve_ws(&base, dir.path().to_str().unwrap(), &[]); let mut a = ws_connect(port, Some("exec-alpha")).await; let mut b = ws_connect(port, Some("exec-bravo")).await; @@ -281,9 +274,7 @@ async fn concurrent_sessions_do_not_leak_session_memory() { turn_text("bob-read"), ]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws(&base, dir.path().to_str().unwrap(), port, &[]); - wait_for_port(port); + let (mut child, port) = serve_ws(&base, dir.path().to_str().unwrap(), &[]); let mut a = ws_connect(port, Some("mem-alpha")).await; let mut b = ws_connect(port, Some("mem-bravo")).await; @@ -369,14 +360,11 @@ async fn concurrent_sessions_do_not_leak_approvals() { ), turn_text("wrote-b"), ]); - let port = free_port(); - let mut child = serve_ws( + let (mut child, port) = serve_ws( &base, dir.path().to_str().unwrap(), - port, &["--approve", "writes"], ); - wait_for_port(port); let mut a = ws_connect(port, Some("appr-alpha")).await; let mut b = ws_connect(port, Some("appr-bravo")).await; @@ -464,9 +452,7 @@ async fn concurrent_sessions_persist_to_distinct_files() { turn_text("unexpected-fallback"), ); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws(&base, dir.path().to_str().unwrap(), port, &[]); - wait_for_port(port); + let (mut child, port) = serve_ws(&base, dir.path().to_str().unwrap(), &[]); let mut a = ws_connect(port, Some("persist-alpha")).await; let mut b = ws_connect(port, Some("persist-bravo")).await; diff --git a/crates/agent/tests/serve_session_title.rs b/crates/agent/tests/serve_session_title.rs index 4166524d..d58e745d 100644 --- a/crates/agent/tests/serve_session_title.rs +++ b/crates/agent/tests/serve_session_title.rs @@ -10,6 +10,7 @@ mod common; use std::io::{BufReader, Write}; +use std::time::Duration; use common::{ SpawnGuarded, read_until_response, serve_dir_cmd, spawn_model_server, @@ -29,11 +30,15 @@ fn only_session(frames: &[Value]) -> Value { fn a_session_names_itself_from_its_opening_exchange() { let dir = tempfile::tempdir().unwrap(); let dir_str = dir.path().to_string_lossy().into_owned(); - // Two turns: the prompt, then the title call it triggers. - let (base, _bodies) = spawn_model_server(vec![ + // The prompt's turn, and the title call it triggers — matched by the title request's own + // system prompt (a scripted server answers title calls itself, off the conversation script). + let (base, _bodies) = spawn_model_server_routed( + vec![( + common::SESSION_TITLE_MARKER.to_string(), + turn_text("Fixing the parser"), + )], turn_text("here is the answer"), - turn_text("Fixing the parser"), - ]); + ); let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_dir_cmd(bin, &base, &dir_str).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); @@ -104,3 +109,66 @@ fn an_unusable_title_is_attempted_once_and_costs_the_session_nothing() { drop(stdin); child.wait().unwrap(); } + +/// The scripted servers recognise the title request by a phrase of its system prompt. If the prompt +/// is reworded, this fails here — not as a conversation test elsewhere hanging on an off-by-one. +#[test] +fn the_scripted_servers_recognise_the_real_title_request() { + assert!( + agent_core::session_title::SESSION_TITLE_SYSTEM.starts_with(common::SESSION_TITLE_MARKER) + ); +} + +/// A scripted (in-order) server is a script of the conversation. The title call `serve` makes after +/// the first run must not take the second prompt's reply — before the harness answered title calls +/// itself, this second prompt got the reply scripted for it one call late and the test hung. And it +/// is made once: a session is titled after its first run, not on every turn — counted, since the +/// harness keeps title calls out of the record. +#[test] +fn a_title_call_does_not_consume_the_conversation_script() { + let dir = tempfile::tempdir().unwrap(); + let dir_str = dir.path().to_string_lossy().into_owned(); + let (base, bodies) = spawn_model_server(vec![ + turn_text("FIRST-REPLY"), + turn_text("SECOND-REPLY"), + turn_text("THIRD-REPLY"), + ]); + let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); + let mut child = serve_dir_cmd(bin, &base, &dir_str).spawn_guarded(); + let mut stdin = child.stdin.take().unwrap(); + let mut stdout = + common::frames_with_deadline(child.stdout.take().unwrap(), Duration::from_secs(20)); + let mut texts = Vec::new(); + for message in ["one", "two", "three"] { + writeln!(stdin, "{}", json!({ "type": "prompt", "message": message })).unwrap(); + stdin.flush().unwrap(); + let frames = read_until_response(&mut stdout, "prompt"); + let done = frames.last().unwrap(); + assert_eq!(done["success"], true, "{done}"); + writeln!(stdin, "{}", json!({ "type": "get_last_assistant_text" })).unwrap(); + stdin.flush().unwrap(); + let frames = read_until_response(&mut stdout, "get_last_assistant_text"); + texts.push(frames.last().unwrap()["data"]["text"].clone()); + } + writeln!(stdin, "{}", json!({ "type": "list_sessions" })).unwrap(); + stdin.flush().unwrap(); + let session = only_session(&read_until_response(&mut stdout, "list_sessions")); + drop(stdin); + child.wait().unwrap(); + assert_eq!( + texts, + [ + json!("FIRST-REPLY"), + json!("SECOND-REPLY"), + json!("THIRD-REPLY") + ] + ); + assert_eq!(session["title"], common::SCRIPTED_SESSION_TITLE); + // The record is the conversation's: three turns, not the title call answered beside them. + assert_eq!(bodies.lock().unwrap().len(), 3); + assert_eq!( + common::title_calls(&base), + 1, + "one title call per session, after its first run" + ); +} diff --git a/crates/agent/tests/serve_session_tree.rs b/crates/agent/tests/serve_session_tree.rs index 02348297..ee1a0b0f 100644 --- a/crates/agent/tests/serve_session_tree.rs +++ b/crates/agent/tests/serve_session_tree.rs @@ -2626,3 +2626,48 @@ fn serve_switch_session_response_carries_reasoning_effort_like_its_three_sibling drop(stdin); child.wait().unwrap(); } + +/// The kinds the agent writes itself — trusted when a session is replayed (an MCP task to resume, its +/// result) — are refused from a client, and nothing lands in the tree; any other kind is accepted. +#[test] +fn serve_append_custom_refuses_the_kinds_the_agent_writes_itself() { + let dir = tempfile::tempdir().unwrap(); + let session_file = dir.path().join("s.jsonl").to_string_lossy().into_owned(); + let (base, _bodies) = spawn_model_server(vec![turn_text("hi there")]); + let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); + let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); + let mut stdin = child.stdin.take().unwrap(); + let mut stdout = BufReader::new(child.stdout.take().unwrap()); + writeln!(stdin, "{}", json!({ "type": "prompt", "message": "hi" })).unwrap(); + stdin.flush().unwrap(); + read_until_response(&mut stdout, "prompt"); + + for kind in ["mcp_task", "mcp_task_result", "mcp_skill_approval"] { + writeln!( + stdin, + "{}", + json!({ "type": "append_custom", "kind": kind, "data": {"sessionId": "x"} }) + ) + .unwrap(); + stdin.flush().unwrap(); + let frames = read_until_response(&mut stdout, "append_custom"); + let resp = frames.last().unwrap(); + assert_eq!(resp["success"], false, "{kind}: {resp:#?}"); + assert!( + resp["error"].as_str().unwrap().contains("reserved"), + "{kind}: {resp:#?}" + ); + } + writeln!(stdin, "{}", json!({ "type": "get_tree" })).unwrap(); + stdin.flush().unwrap(); + let frames = read_until_response(&mut stdout, "get_tree"); + let nodes = frames.last().unwrap()["data"]["nodes"].as_array().unwrap(); + assert!( + nodes + .iter() + .all(|n| !n["preview"].as_str().unwrap_or("").starts_with("[custom:")), + "a refused entry must not reach the tree: {nodes:#?}" + ); + drop(stdin); + child.wait().unwrap(); +} diff --git a/crates/agent/tests/serve_shared_client.rs b/crates/agent/tests/serve_shared_client.rs index 1e90f60d..6c0b5d7f 100644 --- a/crates/agent/tests/serve_shared_client.rs +++ b/crates/agent/tests/serve_shared_client.rs @@ -13,34 +13,32 @@ mod common; use std::process::{Command, Stdio}; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, turn_text, - wait_for_port, ws_connect, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, spawn_listening, spawn_model_server, turn_text, ws_connect, + ws_read_until_response, ws_send, }; use serde_json::json; /// Spawn a `serve --listen` child with a shared upstream client in `--upstream-http2 `. -fn serve_ws_child_shared(base: &str, session_dir: &str, port: u16, http2_mode: &str) -> ChildGuard { - Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--upstream-http2", - http2_mode, - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - ]) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws_child_shared(base: &str, session_dir: &str, http2_mode: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--upstream-http2", + http2_mode, + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + ]) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// Two distinct sessions, driven concurrently through the one shared `--upstream-http2 auto` client, @@ -53,9 +51,7 @@ async fn shared_client_two_concurrent_sessions_both_complete() { let (base, _requests) = spawn_model_server(vec![turn_text("alpha reply"), turn_text("bravo reply")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child_shared(&base, dir.path().to_str().unwrap(), port, "auto"); - wait_for_port(port); + let (mut child, port) = serve_ws_child_shared(&base, dir.path().to_str().unwrap(), "auto"); let mut a = ws_connect(port, Some("alpha")).await; let mut b = ws_connect(port, Some("bravo")).await; diff --git a/crates/agent/tests/serve_startup_stall.rs b/crates/agent/tests/serve_startup_stall.rs index 8cad501a..9c8d84f6 100644 --- a/crates/agent/tests/serve_startup_stall.rs +++ b/crates/agent/tests/serve_startup_stall.rs @@ -23,8 +23,8 @@ use std::process::{Command, Stdio}; use std::time::{Duration, Instant}; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, TestWs, free_port, spawn_model_server, wait_for_port, - ws_connect, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, TestWs, spawn_listening, spawn_model_server, ws_connect, + ws_read_until_response, ws_send, }; use serde_json::json; @@ -44,30 +44,28 @@ const MAX_RPC: Duration = Duration::from_millis(1_000); /// short enough that the open is still in flight when the measured RPC lands. const OPEN_IN_FLIGHT: Duration = Duration::from_millis(300); -fn serve_ws_child(base: &str, session_dir: &str, port: u16) -> ChildGuard { - Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - ]) - .env("HOME", ISOLATED_HOME) - .env( - "BEYOND_AI_AGENT_TEST_SLOW_SESSION_OPEN_MS", - SLOW_OPEN_MS.to_string(), - ) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws_child(base: &str, session_dir: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + ]) + .env("HOME", ISOLATED_HOME) + .env( + "BEYOND_AI_AGENT_TEST_SLOW_SESSION_OPEN_MS", + SLOW_OPEN_MS.to_string(), + ) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// Round-trip one command and return how long the reply took. `get_todos` is answered straight out of @@ -91,9 +89,7 @@ async fn timed_rpc(ws: &mut TestWs, id: &str) -> Duration { async fn one_session_opening_does_not_stall_another_sessions_rpc() { let (base, _requests) = spawn_model_server(vec![]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let _child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (_child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); // The observer session, driven to a reply first so its own (equally stalled) open is behind us — // what this measures is one session answering *while another opens*, not a session opening. diff --git a/crates/agent/tests/serve_state_reporting.rs b/crates/agent/tests/serve_state_reporting.rs index 16a53379..bcfbd880 100644 --- a/crates/agent/tests/serve_state_reporting.rs +++ b/crates/agent/tests/serve_state_reporting.rs @@ -3,14 +3,18 @@ mod common; -use std::io::{BufRead, BufReader, Write}; +use std::io::{BufRead, Write}; use common::{ - SpawnGuarded, read_until_event, read_until_response, serve_cmd, serve_dir_cmd, - spawn_model_server, turn_text, turn_tool_use, + SpawnGuarded, frames_with_deadline, read_until_event, read_until_response, serve_cmd, + serve_dir_cmd, spawn_model_server, turn_text, turn_tool_use, }; use serde_json::{Value, json}; +/// How long a test waits for `serve`'s next frame before failing. Every run here is a few scripted +/// turns and a sub-second `sleep`; a stall (not slowness) is what this catches. +const FRAME_DEADLINE: std::time::Duration = std::time::Duration::from_secs(20); + #[test] fn serve_get_session_stats_reports_the_same_field_set_idle_or_mid_turn() { // Fix 8 (pi-parity gap): the busy-mode arm used to send the raw `LiveStats::snapshot()` verbatim @@ -41,7 +45,7 @@ fn serve_get_session_stats_reports_the_same_field_set_idle_or_mid_turn() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); writeln!( stdin, @@ -109,7 +113,7 @@ fn serve_get_state_and_get_session_stats_answer_live_during_a_prompt() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); writeln!(stdin, "{}", json!({ "type": "prompt", "message": "go" })).unwrap(); stdin.flush().unwrap(); @@ -170,7 +174,7 @@ fn serve_get_tree_since_works_from_the_busy_loop_mid_prompt() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); writeln!(stdin, "{}", json!({ "type": "prompt", "message": "first" })).unwrap(); stdin.flush().unwrap(); @@ -244,7 +248,7 @@ fn serve_get_state_reports_pending_tool_ids_while_a_tool_is_running() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); writeln!(stdin, "{}", json!({ "type": "prompt", "message": "go" })).unwrap(); stdin.flush().unwrap(); @@ -298,7 +302,7 @@ fn serve_get_state_reports_runtime_settings_and_queue_depth() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); // Defaults, nothing queued yet. Fix 1 (pi-parity gap): `claude-test` (this test's model, via // `serve_cmd`) supports reasoning, so a fresh process with no `--reasoning-effort` now starts at @@ -361,7 +365,7 @@ fn serve_steer_and_follow_up_expose_the_actual_queued_text_not_just_a_count() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); writeln!( stdin, @@ -430,7 +434,7 @@ fn serve_reports_cwd_stale_false_for_a_freshly_created_session() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); let mut ready = String::new(); stdout.read_line(&mut ready).unwrap(); @@ -468,7 +472,7 @@ fn serve_reports_cwd_stale_true_when_the_recorded_directory_no_longer_exists() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file.to_string_lossy()).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); let mut ready = String::new(); stdout.read_line(&mut ready).unwrap(); @@ -512,7 +516,7 @@ fn serve_switch_session_reports_cwd_stale_for_the_newly_active_session() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_dir_cmd(bin, &base, &session_dir.to_string_lossy()).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); // The freshly (auto-)created active session must not be stale. let mut ready = String::new(); @@ -565,7 +569,7 @@ fn serve_switch_session_restores_the_reopened_sessions_own_model() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_dir_cmd(bin, &base, &session_dir.to_string_lossy()).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); let mut ready = String::new(); stdout.read_line(&mut ready).unwrap(); @@ -612,7 +616,7 @@ fn serve_get_session_stats_reports_context_usage_after_a_real_turn() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); writeln!(stdin, "{}", json!({ "type": "get_session_stats" })).unwrap(); stdin.flush().unwrap(); @@ -654,7 +658,7 @@ fn serve_set_session_name_and_get_last_assistant_text() { let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); // No name set yet, no assistant reply yet. writeln!(stdin, "{}", json!({ "type": "get_state" })).unwrap(); @@ -741,7 +745,7 @@ fn serve_set_session_name_strips_newlines_and_pushes_a_session_info_changed_fram let bin = env!("CARGO_BIN_EXE_beyond-ai-agent"); let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); writeln!( stdin, @@ -809,7 +813,7 @@ fn get_state_reports_session_file_and_is_streaming_idle_vs_mid_run() { // startup, before any turn ever ran). let mut child = serve_cmd(bin, &base, &session_file).spawn_guarded(); let mut stdin = child.stdin.take().unwrap(); - let mut stdout = BufReader::new(child.stdout.take().unwrap()); + let mut stdout = frames_with_deadline(child.stdout.take().unwrap(), FRAME_DEADLINE); writeln!(stdin, "{}", json!({ "type": "get_state" })).unwrap(); stdin.flush().unwrap(); let idle_frames = read_until_response(&mut stdout, "get_state"); diff --git a/crates/agent/tests/serve_uds.rs b/crates/agent/tests/serve_uds.rs index 056c5c6a..f7e2ebb5 100644 --- a/crates/agent/tests/serve_uds.rs +++ b/crates/agent/tests/serve_uds.rs @@ -14,15 +14,15 @@ use std::process::{Command, Stdio}; use std::time::Duration; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, turn_text, - wait_for_port, ws_connect, ws_connect_uds, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, SpawnGuarded, spawn_listening, spawn_model_server, turn_text, + ws_connect, ws_connect_uds, ws_read_until_response, ws_send, }; use serde_json::json; /// Spawn `serve --listen-uds ` (and optionally `--listen `) against `base` (the mock /// gateway), persisting per-session files under `session_dir`. In listener mode stdio is unused, so /// null it. -fn serve_uds_child(base: &str, session_dir: &str, sock: &str, tcp: Option) -> ChildGuard { +fn uds_cmd(base: &str, session_dir: &str, sock: &str) -> Command { let mut c = Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")); c.args([ "serve", @@ -37,14 +37,21 @@ fn serve_uds_child(base: &str, session_dir: &str, sock: &str, tcp: Option) "--session-dir", session_dir, ]); - if let Some(port) = tcp { - c.args(["--listen", &format!("127.0.0.1:{port}")]); - } c.env("HOME", ISOLATED_HOME) .stdin(Stdio::null()) .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() + .stderr(Stdio::null()); + c +} + +/// `serve --listen-uds ` alone. +fn serve_uds_child(base: &str, session_dir: &str, sock: &str) -> ChildGuard { + uds_cmd(base, session_dir, sock).spawn_guarded() +} + +/// `serve --listen-uds ` and a TCP listener beside it, with the TCP port. +fn serve_uds_and_tcp_child(base: &str, session_dir: &str, sock: &str) -> (ChildGuard, u16) { + spawn_listening(&mut uds_cmd(base, session_dir, sock)) } /// Block until the UDS at `path` accepts a connection, or panic after ~5s. The socket *file* appearing @@ -66,12 +73,7 @@ async fn uds_get_state_then_prompt_streams_ack_event_response() { let (base, _requests) = spawn_model_server(vec![turn_text("hello there")]); let dir = tempfile::tempdir().unwrap(); let sock = dir.path().join("agent.sock"); - let mut child = serve_uds_child( - &base, - dir.path().to_str().unwrap(), - sock.to_str().unwrap(), - None, - ); + let mut child = serve_uds_child(&base, dir.path().to_str().unwrap(), sock.to_str().unwrap()); wait_for_uds(&sock).await; let mut ws = ws_connect_uds(&sock, None).await; @@ -120,14 +122,8 @@ async fn uds_cross_transport_reattach_shares_supervisor() { let (base, _requests) = spawn_model_server(vec![turn_text("crossreply")]); let dir = tempfile::tempdir().unwrap(); let sock = dir.path().join("agent.sock"); - let port = free_port(); - let mut child = serve_uds_child( - &base, - dir.path().to_str().unwrap(), - sock.to_str().unwrap(), - Some(port), - ); - wait_for_port(port); + let (mut child, port) = + serve_uds_and_tcp_child(&base, dir.path().to_str().unwrap(), sock.to_str().unwrap()); wait_for_uds(&sock).await; // Create the session over TCP and run a prompt to completion. @@ -176,12 +172,7 @@ async fn uds_socket_has_owner_only_permissions() { let (base, _requests) = spawn_model_server(vec![]); let dir = tempfile::tempdir().unwrap(); let sock = dir.path().join("agent.sock"); - let mut child = serve_uds_child( - &base, - dir.path().to_str().unwrap(), - sock.to_str().unwrap(), - None, - ); + let mut child = serve_uds_child(&base, dir.path().to_str().unwrap(), sock.to_str().unwrap()); wait_for_uds(&sock).await; let mode = std::fs::metadata(&sock).unwrap().permissions().mode(); @@ -212,12 +203,7 @@ async fn uds_stale_socket_file_is_reclaimed() { } assert!(sock.exists(), "the stale socket node should be on disk"); - let mut child = serve_uds_child( - &base, - dir.path().to_str().unwrap(), - sock.to_str().unwrap(), - None, - ); + let mut child = serve_uds_child(&base, dir.path().to_str().unwrap(), sock.to_str().unwrap()); wait_for_uds(&sock).await; // It serves normally after reclaiming the stale node. diff --git a/crates/agent/tests/serve_websocket.rs b/crates/agent/tests/serve_websocket.rs index aeb5ce24..c8814817 100644 --- a/crates/agent/tests/serve_websocket.rs +++ b/crates/agent/tests/serve_websocket.rs @@ -12,33 +12,31 @@ use std::process::{Child, Command, Stdio}; use std::time::Duration; use common::{ - ChildGuard, ISOLATED_HOME, SpawnGuarded, free_port, spawn_model_server, turn_text, - turn_tool_use, wait_for_port, ws_connect, ws_next_frame, ws_read_until_response, ws_send, + ChildGuard, ISOLATED_HOME, spawn_listening, spawn_model_server, turn_text, turn_tool_use, + ws_connect, ws_next_frame, ws_read_until_response, ws_send, }; use serde_json::json; -/// Spawn a `serve --listen 127.0.0.1:` child against `base` (the mock gateway), persisting +/// Spawn a listening `serve` child (and its port) against `base` (the mock gateway), persisting /// per-session files under `session_dir`. In `--listen` mode stdio is unused, so null it. -fn serve_ws_child(base: &str, session_dir: &str, port: u16) -> ChildGuard { - Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{port}"), - "--gateway-url", - base, - "--key", - "bai_v1.test", - "--model", - "claude-test", - "--session-dir", - session_dir, - ]) - .env("HOME", ISOLATED_HOME) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .spawn_guarded() +fn serve_ws_child(base: &str, session_dir: &str) -> (ChildGuard, u16) { + spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--gateway-url", + base, + "--key", + "bai_v1.test", + "--model", + "claude-test", + "--session-dir", + session_dir, + ]) + .env("HOME", ISOLATED_HOME) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// (i) DoD: a full session over the socket — `get_state` answers with a `session_id`, and a `prompt` @@ -47,9 +45,7 @@ fn serve_ws_child(base: &str, session_dir: &str, port: u16) -> ChildGuard { async fn ws_get_state_then_prompt_streams_ack_event_response() { let (base, _requests) = spawn_model_server(vec![turn_text("hello there")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let mut ws = ws_connect(port, None).await; @@ -101,9 +97,7 @@ async fn ws_run_survives_dropped_connection_and_reattaches() { turn_text("alldone"), ]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); // Connection #1 starts the prompt and waits until the run is provably in flight — the `bash` // tool_start event means turn 1 finished and the 2s sleep is now running — then drops mid-run. @@ -175,9 +169,7 @@ async fn ws_two_connections_both_receive_the_stream() { let (base, _requests) = spawn_model_server(vec![turn_text("firstreply"), turn_text("secondreply")]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); // Both connect to the SAME session id — they coexist (no eviction). let mut ws1 = ws_connect(port, Some(SID)).await; @@ -271,9 +263,7 @@ async fn ws_distinct_sessions_run_concurrently_and_stay_isolated() { turn_text("done"), ]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let mut a = ws_connect(port, Some("alpha")).await; let mut b = ws_connect(port, Some("bravo")).await; @@ -357,9 +347,7 @@ async fn ws_sigterm_persists_in_flight_session_and_exits() { turn_text("done"), ]); let dir = tempfile::tempdir().unwrap(); - let port = free_port(); - let mut child = serve_ws_child(&base, dir.path().to_str().unwrap(), port); - wait_for_port(port); + let (mut child, port) = serve_ws_child(&base, dir.path().to_str().unwrap()); let pid = child.id(); let mut ws = ws_connect(port, Some(SID)).await; diff --git a/crates/agent/tests/smoke.rs b/crates/agent/tests/smoke.rs index 7538777f..c6c26570 100644 --- a/crates/agent/tests/smoke.rs +++ b/crates/agent/tests/smoke.rs @@ -24,9 +24,9 @@ use std::process::{Command, Output, Stdio}; use std::thread; use common::{ - ChildGuard, DEV_PUBKEY_B64, DEV_TOKEN, SpawnGuarded, free_port, gateway_bin, unused_nats_port, - wait_for_allowance_ready, wait_for_port, ws_connect, ws_connect_uds, ws_next_frame, - ws_read_until_response, ws_send, + ChildGuard, DEV_PUBKEY_B64, DEV_TOKEN, SpawnGuarded, free_port, gateway_bin, spawn_listening, + unused_nats_port, wait_for_allowance_ready, wait_for_port, ws_connect, ws_connect_uds, + ws_next_frame, ws_read_until_response, ws_send, }; use serde_json::{Value, json}; @@ -2097,34 +2097,26 @@ fn smoke_whole_run_auto_retry_recovers_from_a_real_dropped_connection_live() { /// Spawn a `serve --listen` child against the gateway on `gw_port`, persisting per-session files under /// `session_dir`, bound to WebSocket `ws_port`. The WebSocket transport uses stdio for nothing, so null /// it. Mirrors [`serve_child`] but for the `serve_ws` path. -fn serve_ws_child( - gw_port: u16, - cwd: &Path, - model: &str, - session_dir: &Path, - ws_port: u16, -) -> ChildGuard { - Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{ws_port}"), - "--gateway-url", - &format!("http://127.0.0.1:{gw_port}"), - "--key", - DEV_TOKEN, - "--model", - model, - "--max-steps", - "6", - "--session-dir", - session_dir.to_str().unwrap(), - ]) - .current_dir(cwd) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::inherit()) - .spawn_guarded() +fn serve_ws_child(gw_port: u16, cwd: &Path, model: &str, session_dir: &Path) -> (ChildGuard, u16) { + spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--gateway-url", + &format!("http://127.0.0.1:{gw_port}"), + "--key", + DEV_TOKEN, + "--model", + model, + "--max-steps", + "6", + "--session-dir", + session_dir.to_str().unwrap(), + ]) + .current_dir(cwd) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ) } /// (iv) Live end-to-end over the **WebSocket** transport: a real prompt drives a real tool round-trip @@ -2146,9 +2138,7 @@ async fn smoke_ws_prompt_round_trip_live() { std::fs::write(dir.path().join("marker.txt"), format!("{token}\n")).unwrap(); let (gw_port, mut gateway) = boot_gateway(dir.path(), p.pool, &key); let session_dir = dir.path().join("sessions"); - let ws_port = free_port(); - let mut child = serve_ws_child(gw_port, dir.path(), p.model, &session_dir, ws_port); - wait_for_port(ws_port); + let (mut child, ws_port) = serve_ws_child(gw_port, dir.path(), p.model, &session_dir); let mut ws = ws_connect(ws_port, None).await; ws_send( @@ -2213,9 +2203,7 @@ async fn smoke_ws_run_survives_reconnect_live() { let dir = tempfile::tempdir().unwrap(); let (gw_port, mut gateway) = boot_gateway(dir.path(), p.pool, &key); let session_dir = dir.path().join("sessions"); - let ws_port = free_port(); - let mut child = serve_ws_child(gw_port, dir.path(), p.model, &session_dir, ws_port); - wait_for_port(ws_port); + let (mut child, ws_port) = serve_ws_child(gw_port, dir.path(), p.model, &session_dir); // ws1 starts a run that runs `bash sleep 8` (a wide, real in-flight window), waits until the // tool is actually running, then drops mid-run. @@ -2410,32 +2398,28 @@ async fn smoke_h2c_shared_pool_live() { std::fs::write(dir.path().join("marker.txt"), format!("{token}\n")).unwrap(); let (gw_port, mut gateway) = boot_gateway(dir.path(), p.pool, &key); let session_dir = dir.path().join("sessions"); - let ws_port = free_port(); // serve with the shared h2c client. The gateway (this branch) accepts downstream h2c. - let mut child = Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) - .args([ - "serve", - "--listen", - &format!("127.0.0.1:{ws_port}"), - "--upstream-http2", - "h2c", - "--gateway-url", - &format!("http://127.0.0.1:{gw_port}"), - "--key", - DEV_TOKEN, - "--model", - p.model, - "--max-steps", - "6", - "--session-dir", - session_dir.to_str().unwrap(), - ]) - .current_dir(dir.path()) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::inherit()) - .spawn_guarded(); - wait_for_port(ws_port); + let (mut child, ws_port) = spawn_listening( + Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) + .args([ + "serve", + "--upstream-http2", + "h2c", + "--gateway-url", + &format!("http://127.0.0.1:{gw_port}"), + "--key", + DEV_TOKEN, + "--model", + p.model, + "--max-steps", + "6", + "--session-dir", + session_dir.to_str().unwrap(), + ]) + .current_dir(dir.path()) + .stdin(Stdio::null()) + .stdout(Stdio::null()), + ); // Two concurrent sessions, both reading the marker through the ONE shared h2c client. let mut a = ws_connect(ws_port, Some("h2csessa")).await; @@ -2491,9 +2475,7 @@ async fn smoke_multi_attach_both_see_live_stream() { std::fs::write(dir.path().join("marker.txt"), format!("{token}\n")).unwrap(); let (gw_port, mut gateway) = boot_gateway(dir.path(), p.pool, &key); let session_dir = dir.path().join("sessions"); - let ws_port = free_port(); - let mut child = serve_ws_child(gw_port, dir.path(), p.model, &session_dir, ws_port); - wait_for_port(ws_port); + let (mut child, ws_port) = serve_ws_child(gw_port, dir.path(), p.model, &session_dir); const SID: &str = "multiattachlive1"; // Two devices on ONE session, at the same time. diff --git a/crates/test-support/src/lib.rs b/crates/test-support/src/lib.rs index 5bbcc42a..de619e4b 100644 --- a/crates/test-support/src/lib.rs +++ b/crates/test-support/src/lib.rs @@ -158,8 +158,100 @@ fn read_http_request(stream: &mut TcpStream) -> String { } } +/// A phrase only the agent's session-title request carries (the opening of +/// `agent_core::session_title::SESSION_TITLE_SYSTEM`); `session_title_request_is_recognised` in the +/// agent's tests keeps it in step. +pub const SESSION_TITLE_MARKER: &str = "You write short titles for saved conversations"; + +/// The title a scripted server gives every session-title request. +pub const SCRIPTED_SESSION_TITLE: &str = "Scripted test session"; + +/// Whether a raw request is the session-title call `serve` makes after a session's first run. +pub fn is_session_title_request(req: &str) -> bool { + req.contains(SESSION_TITLE_MARKER) +} + +fn reply(stream: &mut TcpStream, resp: &str) { + let http = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n{resp}" + ); + let _ = stream.write_all(http.as_bytes()); + let _ = stream.flush(); +} + +/// The next *conversation* request on a scripted server, with its raw text. +/// +/// A scripted (in-order) server is a script of the conversation's turns. The session-title request +/// is not one of them: `serve` makes it once, after a session's first successful run, in the +/// background of whatever the test does next. Let it take a scripted reply and every later turn is +/// off by one — the next prompt gets the reply meant for the title, and a test written against the +/// conversation hangs waiting for a turn that already went to the title (or fails on a reply meant +/// for a different turn). So title requests are answered here, with [`SCRIPTED_SESSION_TITLE`], and +/// neither consume the script nor appear in the record, which holds the conversation's requests +/// exactly. A test about titles scripts and observes them on a routed server, matching +/// [`SESSION_TITLE_MARKER`]. +fn accept_scripted( + listener: &TcpListener, + recorder: Option<&Mutex>>, +) -> Option<(TcpStream, String)> { + let base = listener + .local_addr() + .ok() + .map(|a| format!("http://{a}")) + .unwrap_or_default(); + loop { + let (mut stream, _) = listener.accept().ok()?; + // Drain the *whole* request before answering. Closing a socket with unread data in its + // receive buffer makes the kernel send an RST instead of a FIN, which discards whatever the + // peer had not yet read — including the response just written. + let req = read_http_request(&mut stream); + if is_session_title_request(&req) { + *title_counts() + .lock() + .unwrap() + .entry(base.clone()) + .or_default() += 1; + reply(&mut stream, &turn_text(SCRIPTED_SESSION_TITLE)); + continue; + } + if let Some(recorder) = recorder { + recorder.lock().unwrap().push(req.clone()); + } + return Some((stream, req)); + } +} + +/// Title requests answered so far, per scripted server (by base URL). +fn title_counts() -> &'static Mutex> { + static COUNTS: std::sync::OnceLock>> = + std::sync::OnceLock::new(); + COUNTS.get_or_init(Default::default) +} + +/// How many session-title requests the scripted server at `base` has answered itself. They are kept +/// out of the conversation's record, so this is where a change in how often `serve` asks for a +/// title — once per session, after its first successful run — shows. +pub fn title_calls(base: &str) -> usize { + title_counts() + .lock() + .unwrap() + .get(base) + .copied() + .unwrap_or(0) +} + +/// After a script runs out: keep answering title requests (one can arrive after the last turn), +/// and refuse anything else by closing the connection unanswered, as a closed listener would. +fn serve_titles_only(listener: TcpListener, recorder: Option>>>) { + while let Some((stream, _)) = accept_scripted(&listener, recorder.as_deref()) { + drop(stream); + } +} + /// Spawn a model server answering `responses` in order, recording each full raw request (headers + /// body). Returns the base URL and the shared record of requests. +/// +/// Session-title requests do not consume a response (see [`accept_scripted`]). pub fn spawn_model_server(responses: Vec) -> (String, Arc>>) { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); @@ -167,16 +259,12 @@ pub fn spawn_model_server(responses: Vec) -> (String, Arc