From 05b8ea5653ac43ab2ce342243e341a3561a92423 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sat, 11 Jul 2026 11:57:36 -0700 Subject: [PATCH 1/9] feat(agent): carry the model's memory writes across compaction deterministically MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recall reminder (#25) tells the model to consult `/session` after a compaction, and the pressure nudge (#26) tells it to save there before — but both still lean on the summarizing model not erasing the fact that a note *exists*. Close that: make the model's memory writes ride the existing deterministic-carry channel, so every summary lists them verbatim regardless of what the summarizer chooses to say. - `CompactionProvenance` gains `memory_notes: Vec` — the logical paths the model authored via the `memory` tool (create/str_replace/ insert), deduped, oldest-first, folded forward every round by `merge_provenance`, exactly like the read/modified-file lists. - `extract_memory_notes` keys on the `memory` tool name only (like `extract_todos` keys on `todo`) — agnostic to which root a note lives under, so agent-core never learns the host's `/session` vs `/memories` convention. - `format_memory_notes` renders a `` block appended to every summary; `Agent::compact` appends it alongside the file/todo blocks, and `previous_summary` strips it (via `CARRY_BLOCKS`) so it never accumulates across incremental rounds. - Persisted on the `Entry::Compaction` record and restored on reopen, so a serve restart past a compaction keeps the awareness. This is the guaranteed, zero-effort half of the save→carry→recall arc: after a compaction the model is *shown* what it wrote and where, not just told to go look. Proven e2e (serve_session_memory): a `/session` note the model wrote survives a compaction as a `` block in the summary. Plus unit tests for extract/format/merge/strip. ARCHITECTURE.md updated. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/agent-core/src/agent.rs | 6 +- crates/agent-core/src/compaction.rs | 138 ++++++++++++++++++++- crates/agent/ARCHITECTURE.md | 6 +- crates/agent/src/session_store.rs | 11 ++ crates/agent/tests/serve_session_memory.rs | 8 ++ 5 files changed, 165 insertions(+), 4 deletions(-) diff --git a/crates/agent-core/src/agent.rs b/crates/agent-core/src/agent.rs index f8dee65..ad6bf7a 100644 --- a/crates/agent-core/src/agent.rs +++ b/crates/agent-core/src/agent.rs @@ -2265,9 +2265,10 @@ impl Agent { // reflect every round so far, not just this one's new activity. `previous_summary` strips these // same blocks back off before the body is ever fed forward, so they never accumulate. let summary = format!( - "{summary}{}{}", + "{summary}{}{}{}", compaction::format_file_operations(&file_ops.read_files, &file_ops.modified_files), - compaction::format_todo_list(file_ops.todos.as_ref()) + compaction::format_todo_list(file_ops.todos.as_ref()), + compaction::format_memory_notes(&file_ops.memory_notes) ); compaction::apply_summary(session, first_kept, &summary, tokens_before); session.compaction = file_ops; @@ -7432,6 +7433,7 @@ mod tests { compactions: 1, last_reason: Some(CompactionReason::Threshold), modified_files: vec![], + memory_notes: vec![], }; let mock = Arc::new(MockTransport::new(vec![turn::text("turn prefix summary")])); diff --git a/crates/agent-core/src/compaction.rs b/crates/agent-core/src/compaction.rs index 01fb05c..ff4c3de 100644 --- a/crates/agent-core/src/compaction.rs +++ b/crates/agent-core/src/compaction.rs @@ -75,6 +75,14 @@ pub struct CompactionProvenance { /// older one — see [`merge_provenance`]. #[serde(default)] pub todos: Option, + /// Logical paths the model has authored via the `memory` tool (a `create`/`str_replace`/`insert`), + /// deduped, oldest-first, across every round so far. Carried for the same reason as the file lists: + /// once `apply_summary` drops the prefix, the fact that the model *wrote itself a note* — and where — + /// is gone with it, so a model reminded to consult its working memory post-compaction would have no + /// idea what it saved or that it saved anything. Rendered into the summary by [`format_memory_notes`], + /// it becomes a guaranteed, zero-effort "here is what you recorded, recall the relevant ones" pointer. + #[serde(default)] + pub memory_notes: Vec, } /// Fold `previous`'s provenance forward with whatever the current round's prefix (`messages`) adds, and @@ -104,6 +112,12 @@ pub fn merge_provenance( modified.push(path); } } + let mut memory_notes = previous.memory_notes.clone(); + for path in extract_memory_notes(messages) { + if !memory_notes.contains(&path) { + memory_notes.push(path); + } + } CompactionProvenance { read_files: read, modified_files: modified, @@ -112,6 +126,7 @@ pub fn merge_provenance( // `or_else`, not `or`: a round whose prefix cleared the list (`todos: []`) yields `Some([])` and // must win over `previous`, or a cleared plan would silently reappear at the next compaction. todos: extract_todos(messages).or_else(|| previous.todos.clone()), + memory_notes, } } @@ -656,6 +671,50 @@ pub fn extract_todos(messages: &[Message]) -> Option { }) } +/// The logical paths the model wrote to via the `memory` tool in `messages`, in first-seen order — a +/// `create`, `str_replace`, or `insert` (the write commands; `view`/`search` read, `delete`/`rename` +/// don't add content). Deduped within this scan; accumulated across rounds by [`merge_provenance`]. +/// +/// Keys on the tool name `"memory"` only, exactly as [`extract_todos`] keys on `"todo"` — it stays +/// agnostic to *which* root (`/memories` vs `/session`) a note lives under, so the lower layer never has +/// to know the host's memory-mount conventions. Both roots are worth surfacing: a note is a note. +pub fn extract_memory_notes(messages: &[Message]) -> Vec { + let mut notes = Vec::new(); + for m in messages { + for b in &m.content { + if let ContentBlock::ToolUse { name, input, .. } = b + && name == "memory" + && matches!( + input.get("command").and_then(Value::as_str), + Some("create" | "str_replace" | "insert") + ) + && let Some(path) = input.get("path").and_then(Value::as_str) + { + let path = path.to_string(); + if !notes.contains(&path) { + notes.push(path); + } + } + } + } + notes +} + +/// Render the carried `memory` note paths as a `` block appended to the summarization +/// model's *output* — the same deterministic-carry contract [`format_todo_list`] documents. Returns `""` +/// when there are none, so appending it to a summary is always safe. The model reads the block as: these +/// are files I wrote myself; recall the relevant ones with the `memory` tool's `view`. +pub fn format_memory_notes(memory_notes: &[String]) -> String { + if memory_notes.is_empty() { + String::new() + } else { + format!( + "\n\n\n{}\n", + memory_notes.join("\n") + ) + } +} + /// Render a `todos` array as the compact checklist both the `todo` tool's own result and a compaction /// summary show the model. Shared so the two can never drift. /// @@ -731,10 +790,11 @@ pub fn previous_summary(prefix: &[Message]) -> Option<&str> { /// The host-appended, deterministically-regenerated blocks a summary body carries at its tail, as /// `(open, close)` tag pairs. See [`format_file_operations`] / [`format_todo_list`]. -const CARRY_BLOCKS: [(&str, &str); 3] = [ +const CARRY_BLOCKS: [(&str, &str); 4] = [ ("", ""), ("", ""), ("", ""), + ("", ""), ]; /// Strip every trailing [`CARRY_BLOCKS`] block from a summary body, leaving only the model's own prose. @@ -1751,6 +1811,7 @@ mod tests { compactions: 1, last_reason: Some(CompactionReason::Threshold), todos: None, + memory_notes: vec![], }; let messages = vec![Message::assistant(vec![ContentBlock::tool_use( "1", @@ -1937,6 +1998,81 @@ mod tests { assert!(extract_todos(&convo()).is_none()); } + fn mem_call(id: &str, command: &str, path: &str) -> Message { + Message::assistant(vec![ContentBlock::tool_use( + id, + "memory", + json!({ "command": command, "path": path }), + )]) + } + + #[test] + fn extract_memory_notes_collects_writes_first_seen_and_ignores_the_rest() { + let messages = vec![ + mem_call("1", "create", "/session/facts.md"), + mem_call("2", "view", "/session/facts.md"), // a read — ignored + mem_call("3", "str_replace", "/session/facts.md"), // same path, deduped + mem_call("4", "insert", "/memories/notes.md"), // the other root counts too + mem_call("5", "delete", "/session/gone.md"), // a delete adds nothing + Message::assistant(vec![ContentBlock::tool_use( + "6", + "write", // a different tool that happens to carry a path — not memory + json!({ "path": "/session/decoy.md" }), + )]), + ]; + assert_eq!( + extract_memory_notes(&messages), + vec![ + "/session/facts.md".to_string(), + "/memories/notes.md".to_string() + ] + ); + } + + #[test] + fn format_memory_notes_renders_a_block_or_nothing() { + assert_eq!(format_memory_notes(&[]), ""); + let block = format_memory_notes(&["/session/a.md".into(), "/memories/b.md".into()]); + assert_eq!( + block, + "\n\n\n/session/a.md\n/memories/b.md\n" + ); + } + + #[test] + fn merge_provenance_accumulates_memory_notes_across_rounds_deduped() { + let round1 = merge_provenance( + &CompactionProvenance::default(), + &[mem_call("1", "create", "/session/a.md")], + CompactionReason::Threshold, + ); + assert_eq!(round1.memory_notes, vec!["/session/a.md".to_string()]); + let round2 = merge_provenance( + &round1, + &[ + mem_call("2", "str_replace", "/session/a.md"), // already known — deduped + mem_call("3", "create", "/session/b.md"), // new + ], + CompactionReason::Threshold, + ); + assert_eq!( + round2.memory_notes, + vec!["/session/a.md".to_string(), "/session/b.md".to_string()] + ); + } + + #[test] + fn previous_summary_strips_the_memory_notes_carry_block() { + // A summary body that ends with a memory-notes block must have it peeled off before it is fed + // forward, exactly like the file/todo blocks — otherwise every incremental round re-appends one. + let body = "The model saved some notes.\n\n\n/session/a.md\n"; + let msg = Message::user(format!("{SUMMARY_MARKER}\n\n{body}")); + assert_eq!( + previous_summary(&[msg]), + Some("The model saved some notes.") + ); + } + #[test] fn extract_todos_ignores_a_todo_call_with_a_non_array_todos_argument() { // A malformed call the loop still recorded (the tool rejected it, but the `tool_use` block is diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 6de4ff5..2b324d9 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -447,7 +447,11 @@ The harness layers several capabilities over the bare tools + loop: `/session` while it still has full detail; and a **post-compaction recall reminder** (`memory::COMPACTION_REMINDER`) fires on the cut, telling it to read `/session` back. Both are computed host-side from the `Usage` events the observer already sees (no agent-core change), gated on a session - mount, and pushed as mid-run steers. The tool is host-owned (each backend `Arc` cloned into the tool at each registry rebuild, + mount, and pushed as mid-run steers. Backing the reminder, the paths the model writes via the `memory` + tool ride the **deterministic-carry channel** (`agent_core::CompactionProvenance::memory_notes`, + parallel to the read/modified-file and `todo` lists): every summary carries a `` block + naming them, so a summarizing model that never mentions a note can't erase the model's awareness that it + saved one and where — the guaranteed, zero-effort half of the save→carry→recall arc. The tool is host-owned (each backend `Arc` cloned into the tool at each registry rebuild, like `structured_output`'s `OutputSlot`) and registered _after_ `apply_filter` so a `--tools` allow-list can't strip it. In `serve` the `/session` mount rides a shared, swappable `SessionDir` cell so a `switch_session`/`new_session`/`fork` re-points every holder with one cell write — no tool or agent diff --git a/crates/agent/src/session_store.rs b/crates/agent/src/session_store.rs index acbbcd8..dc8f757 100644 --- a/crates/agent/src/session_store.rs +++ b/crates/agent/src/session_store.rs @@ -367,6 +367,11 @@ enum Entry { /// `tool_use` block that carried it. #[serde(default, skip_serializing_if = "Option::is_none")] todos: Option, + /// The `memory`-note paths carried across this round (see `CompactionProvenance::memory_notes`), + /// restored the same way so a restart past a compaction doesn't lose the model's awareness of what + /// it wrote to its working/durable memory. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + memory_notes: Vec, }, /// A record that the active model changed, anchored to whatever message was the tip at the moment /// it did — see [`SessionStore::record_model_change`]/[`SessionStore::model_at`]. `parent_id` here @@ -945,6 +950,7 @@ impl SessionStore { modified_files, last_reason, todos, + memory_notes, .. }) => { // Each record already carries every earlier round's file-provenance folded in (via @@ -960,6 +966,7 @@ impl SessionStore { compactions: 0, // replaced by `meta.compactions` (the authoritative counter) below last_reason, todos, + memory_notes, }); compactions.insert( id, @@ -1617,6 +1624,7 @@ impl SessionStore { modified_files: meta.provenance.modified_files, last_reason: meta.provenance.last_reason, todos: meta.provenance.todos, + memory_notes: meta.provenance.memory_notes, }; // An updated header snapshot (the new `compactions`/`dropped_messages` counters) first, then the @@ -4836,6 +4844,7 @@ mod tests { compactions: 1, last_reason: Some(agent_core::compaction::CompactionReason::Threshold), todos: None, + memory_notes: vec![], }; store .rewrite_compacted( @@ -4984,6 +4993,7 @@ mod tests { compactions: 1, last_reason: Some(agent_core::compaction::CompactionReason::Threshold), todos: None, + memory_notes: vec![], }; store .rewrite_compacted( @@ -5008,6 +5018,7 @@ mod tests { compactions: 2, last_reason: Some(agent_core::compaction::CompactionReason::Manual), todos: None, + memory_notes: vec![], }; store .rewrite_compacted( diff --git a/crates/agent/tests/serve_session_memory.rs b/crates/agent/tests/serve_session_memory.rs index e5bb9ff..23971a7 100644 --- a/crates/agent/tests/serve_session_memory.rs +++ b/crates/agent/tests/serve_session_memory.rs @@ -123,6 +123,14 @@ fn a_fact_written_to_session_survives_compaction_and_is_recovered_after_the_remi "the working memory must be untouched by compaction" ); + // Deterministic carry: the summary itself lists the memory the model authored, so even a summarizer + // that never mentions it can't erase the fact that a `/session` note exists and where. + let transcript = std::fs::read_to_string(&session_file).unwrap(); + assert!( + transcript.contains("") && transcript.contains("/session/facts.md"), + "the compaction summary must carry a block naming the note the model wrote" + ); + // --- Prompt 2: the reminder rides the next request, and the model recovers the value. ------------ let before = bodies.lock().unwrap().len(); send( From 9edc0363ae78ebdcf5b5132da9174f4560f7b2b6 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sat, 11 Jul 2026 12:05:41 -0700 Subject: [PATCH 2/9] feat(agent): summarizer preserves specifics verbatim (shrink the loss at the source) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The deterministic carry guarantees a few structured artifacts survive a compaction; the summary prose still loses everything else the summarizer chooses to smooth over. Strengthen the summarization instructions so the specifics most often dropped survive the cut in the first place. `SUMMARY_SYSTEM` (shared by compaction and branch summaries) and every per-mode template (`SUMMARY_INSTRUCTION`, `UPDATE_INSTRUCTION`, `SPLIT_TURN_INSTRUCTION`, `BRANCH_SUMMARY_INSTRUCTION`) now explicitly demand that file:line references, exact literal values (numbers, ports, versions, config keys, IDs, flags), commands and their key outputs, and error messages be copied *verbatim* rather than paraphrased or rounded — "a specific you drop is gone for good." The pi-parity structural shape (section headings, checkbox/`(none)` conventions, numbered Next Steps) is unchanged; only the preserve directive is sharpened. A unit test pins that every summary-shaping prompt demands verbatim specifics and names file:line references, so this can't silently regress. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/agent-core/src/branch_summary.rs | 4 +- crates/agent-core/src/compaction.rs | 53 ++++++++++++++++++++----- crates/agent/ARCHITECTURE.md | 6 ++- 3 files changed, 52 insertions(+), 11 deletions(-) diff --git a/crates/agent-core/src/branch_summary.rs b/crates/agent-core/src/branch_summary.rs index 05b228b..4b35cda 100644 --- a/crates/agent-core/src/branch_summary.rs +++ b/crates/agent-core/src/branch_summary.rs @@ -47,7 +47,9 @@ requirements mentioned]\n- [Or \"(none)\" if none were mentioned]\n\n## Progress [Completed tasks/changes]\n\n### In Progress\n- [ ] [Work that was started but not finished]\n\n### \ Blocked\n- [Issues preventing progress, if any]\n\n## Key Decisions\n- **[Decision]**: [Brief \ rationale]\n\n## Next Steps\n1. [What should happen next to continue this work]\n\nKeep each section \ -concise. Preserve exact file paths, function names, and error messages."; +concise, but never at the cost of a specific: copy exact file paths (with line numbers), \ +function/identifier names, literal values (numbers, ports, versions, config keys, IDs, flags), commands \ +and their key outputs, and error messages verbatim rather than paraphrasing them."; /// Prefix marking a message as a branch summary, mirroring [`crate::compaction::SUMMARY_MARKER`] — an /// explicit signal (for a human or a later, nested summarization) that this text recaps an abandoned diff --git a/crates/agent-core/src/compaction.rs b/crates/agent-core/src/compaction.rs index ff4c3de..52a312f 100644 --- a/crates/agent-core/src/compaction.rs +++ b/crates/agent-core/src/compaction.rs @@ -174,9 +174,12 @@ impl Default for CompactionConfig { /// pi's "You are a context summarization assistant..." opener — functionally equivalent (both frame the /// call as compaction, not conversation), just not textually identical. pub const SUMMARY_SYSTEM: &str = "You compact a long agent transcript so the agent can keep working \ -with far fewer tokens but no loss of essential context. Be precise, concrete, and information-dense; \ -preserve file paths, identifiers, commands, and decisions exactly. Do NOT continue the conversation. \ -Do NOT respond to any questions in the conversation. ONLY output the structured summary."; +with far fewer tokens but no loss of essential context. Be precise, concrete, and information-dense. \ +Copy specifics verbatim — never paraphrase, round, or elide them: file paths (with line numbers), \ +identifiers and function names, exact literal values (numbers, ports, versions, config keys, IDs, \ +flags), commands and their key outputs, error messages, and decisions and their rationale. A specific \ +you drop is gone for good — prefer keeping an exact value over smooth prose. Do NOT continue the \ +conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary."; /// The instruction appended after the rendered transcript in the summarization call. Matches pi's own /// `SUMMARIZATION_PROMPT` (`compaction.ts`) structurally: a bracketed exemplar per section (telling the @@ -192,8 +195,11 @@ mentioned by user]\n- [Or \"(none)\" if none were mentioned]\n\n## Progress\n### tasks/changes]\n\n### In Progress\n- [ ] [Current work]\n\n### Blocked\n- [Issues preventing progress, \ if any]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale]\n\n## Next Steps\n1. [Ordered list of \ what should happen next]\n\n## Critical Context\n- [Any data, examples, or references needed to \ -continue]\n- [Or \"(none)\" if not applicable]\n\nKeep each section concise. Preserve exact file paths, \ -function names, and error messages."; +continue — copy exact literal values (numbers, ports, versions, config keys, IDs, flags), commands and \ +their key outputs, and file:line references verbatim rather than describing them]\n- [Or \"(none)\" if \ +not applicable]\n\nKeep each section concise, but never at the cost of a specific: copy exact file paths \ +(with line numbers), function/identifier names, literal values, commands and their key outputs, and \ +error messages verbatim rather than paraphrasing them."; /// The instruction used when a *previous* summary already exists — an incremental update rather than a /// from-scratch re-summarization. Without this, each compaction re-summarizes the prior summary as raw @@ -213,8 +219,10 @@ discovered]\n\n## Progress\n### Done\n- [x] [Include previously done items AND n items]\n\n### In Progress\n- [ ] [Current work - update based on progress]\n\n### Blocked\n- [Current \ blockers - remove if resolved]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale] (preserve all \ previous, add new)\n\n## Next Steps\n1. [Update based on current state]\n\n## Critical Context\n- \ -[Preserve important context, add new if needed]\n\nKeep each section concise. Preserve exact file \ -paths, function names, and error messages."; +[Preserve important context, add new if needed — keep exact literal values, commands and their key \ +outputs, and file:line references verbatim]\n\nKeep each section concise, but never at the cost of a \ +specific: copy exact file paths (with line numbers), function/identifier names, literal values, commands \ +and their key outputs, and error messages verbatim rather than paraphrasing them."; /// The instruction used when the cut point falls *inside* an in-progress turn (see [`is_split_turn`]): /// the prefix being summarized ends mid-tool-dispatch, so its concluding response — and the context @@ -226,8 +234,9 @@ pub const SPLIT_TURN_INSTRUCTION: &str = "This is the PREFIX of a turn that was The SUFFIX (recent work) is retained.\n\nSummarize the prefix to provide context for the retained \ suffix:\n\n## Original Request\n[What did the user ask for in this turn?]\n\n## Early Progress\n- \ [Key decisions and work done in the prefix]\n\n## Context for Suffix\n- [Information needed to \ -understand the retained recent work]\n\nBe concise. Focus on what's needed to understand the kept \ -suffix."; +understand the retained recent work — copy exact file paths (with line numbers), literal values, \ +commands and their key outputs, and error messages verbatim]\n\nBe concise. Focus on what's needed to \ +understand the kept suffix."; /// Prefix marking a message as a compaction summary, so a later compaction recognizes it and updates /// it incrementally instead of re-summarizing it. Created by [`apply_summary`], detected by @@ -1082,6 +1091,32 @@ mod tests { ); } + #[test] + fn every_summary_prompt_demands_verbatim_specifics() { + // The one thing a summarizer must not do is smooth a specific away — every instruction that + // shapes a summary must explicitly demand the categories most often lost (line numbers, literal + // values, commands+outputs) be copied verbatim, not paraphrased. + for (name, text) in [ + ("SUMMARY_SYSTEM", SUMMARY_SYSTEM), + ("SUMMARY_INSTRUCTION", SUMMARY_INSTRUCTION), + ("UPDATE_INSTRUCTION", UPDATE_INSTRUCTION), + ("SPLIT_TURN_INSTRUCTION", SPLIT_TURN_INSTRUCTION), + ( + "BRANCH_SUMMARY_INSTRUCTION", + crate::branch_summary::BRANCH_SUMMARY_INSTRUCTION, + ), + ] { + assert!( + text.contains("verbatim"), + "{name} must demand specifics be copied verbatim: {text}" + ); + assert!( + text.contains("line numbers"), + "{name} must name file:line references as something to preserve: {text}" + ); + } + } + #[test] fn should_compact_tracks_last_input_against_window() { let cfg = CompactionConfig { diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 2b324d9..b14342a 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -451,7 +451,11 @@ The harness layers several capabilities over the bare tools + loop: tool ride the **deterministic-carry channel** (`agent_core::CompactionProvenance::memory_notes`, parallel to the read/modified-file and `todo` lists): every summary carries a `` block naming them, so a summarizing model that never mentions a note can't erase the model's awareness that it - saved one and where — the guaranteed, zero-effort half of the save→carry→recall arc. The tool is host-owned (each backend `Arc` cloned into the tool at each registry rebuild, + saved one and where — the guaranteed, zero-effort half of the save→carry→recall arc. Complementing the + deterministic carry, the summarizer's own instructions (`agent_core::compaction::SUMMARY_SYSTEM` and the + per-mode templates) demand that specifics most often smoothed away — `file:line` references, exact + literal values, commands and their key outputs — be copied *verbatim*, shrinking the loss at the source. + The tool is host-owned (each backend `Arc` cloned into the tool at each registry rebuild, like `structured_output`'s `OutputSlot`) and registered _after_ `apply_filter` so a `--tools` allow-list can't strip it. In `serve` the `/session` mount rides a shared, swappable `SessionDir` cell so a `switch_session`/`new_session`/`fork` re-points every holder with one cell write — no tool or agent From bdd7974228684b55e6545df44687293fd25ea1f8 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sat, 11 Jul 2026 12:07:20 -0700 Subject: [PATCH 3/9] style: dprint markdown fmt (underscore emphasis in ARCHITECTURE.md) Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/agent/ARCHITECTURE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index b14342a..c3b27c7 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -454,7 +454,7 @@ The harness layers several capabilities over the bare tools + loop: saved one and where — the guaranteed, zero-effort half of the save→carry→recall arc. Complementing the deterministic carry, the summarizer's own instructions (`agent_core::compaction::SUMMARY_SYSTEM` and the per-mode templates) demand that specifics most often smoothed away — `file:line` references, exact - literal values, commands and their key outputs — be copied *verbatim*, shrinking the loss at the source. + literal values, commands and their key outputs — be copied _verbatim_, shrinking the loss at the source. The tool is host-owned (each backend `Arc` cloned into the tool at each registry rebuild, like `structured_output`'s `OutputSlot`) and registered _after_ `apply_filter` so a `--tools` allow-list can't strip it. In `serve` the `/session` mount rides a shared, swappable `SessionDir` cell so a From ba6b1f54ef765dbc35c5447f29f8e69c4484b8b6 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sat, 11 Jul 2026 18:13:34 -0700 Subject: [PATCH 4/9] fix(agent): pressure point never fires on every turn when reserve >= window MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `compaction_pressure_point` collapsed to 0 when the reserve met or exceeded the context window (a tiny-window or test config): since the check is `live_prompt > point`, a 0 point is true on every turn, firing the pre-compaction nudge (and its extra steer turn) spuriously from turn one — which, against a fixed-response test server, can exhaust or stall it. There is no proactive window to warn within when reserve >= window (compaction fires the instant any usage exists), so return a never-fires sentinel (u32::MAX) in that case instead. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/agent/src/memory/mod.rs | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/crates/agent/src/memory/mod.rs b/crates/agent/src/memory/mod.rs index da1b65a..91bf316 100644 --- a/crates/agent/src/memory/mod.rs +++ b/crates/agent/src/memory/mod.rs @@ -452,8 +452,19 @@ pub const PRESSURE_NUDGE: &str = "Your context is getting large /// (`context_window - reserve_tokens`, the point [`agent_core::compaction::should_compact`] fires at). /// This leaves runway — typically a few turns — for the model to checkpoint to `/session` before the cut. /// Integer math (÷5 ×4) so there is no float cast and no overflow even at a 1M-token window. +/// +/// Returns [`u32::MAX`] (a "never fires" sentinel — `live_prompt > MAX` is always false) when +/// `reserve_tokens >= context_window`, i.e. the threshold is 0. That is a pathological or test-only config +/// with **no proactive window at all**: compaction fires the instant any real usage exists, so there is +/// no runway to warn within. Without this guard the point would be 0, and `live_prompt > 0` is true on +/// literally every turn — the nudge would fire on turn one of any small-window session, adding a spurious +/// steer turn (and, against a fixed-response test server, potentially exhausting or stalling it). pub fn compaction_pressure_point(context_window: u32, reserve_tokens: u32) -> u32 { - context_window.saturating_sub(reserve_tokens) / 5 * 4 + let threshold = context_window.saturating_sub(reserve_tokens); + if threshold == 0 { + return u32::MAX; + } + threshold / 5 * 4 } /// The full prompt size a [`agent_core::message::TokenUsage`] report implies — uncached input plus @@ -598,9 +609,13 @@ mod tests { assert!(p < 200_000 - 16_384, "must fire before the compaction cut"); assert!(p > (200_000 - 16_384) / 2, "but not absurdly early"); assert_eq!(p, 183_616 / 5 * 4); - // Degenerate windows don't panic or overflow. - assert_eq!(compaction_pressure_point(0, 16_384), 0); - assert_eq!(compaction_pressure_point(100, 200), 0); + // A window with no proactive room (reserve >= window) disables the nudge (never-fires sentinel) + // rather than collapsing to 0 — which, as `live_prompt > 0`, would fire on every turn. + assert_eq!(compaction_pressure_point(0, 16_384), u32::MAX); + assert_eq!(compaction_pressure_point(100, 200), u32::MAX); + assert_eq!(compaction_pressure_point(200, 200), u32::MAX); + // A real (if small) window still yields a finite, sub-threshold point. + assert_eq!(compaction_pressure_point(300, 200), (300 - 200) / 5 * 4); assert_eq!(compaction_pressure_point(u32::MAX, 0), u32::MAX / 5 * 4); } From 7ec16bbf100ed41141d79fdcd5696ede65df76cd Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sun, 12 Jul 2026 15:04:10 -0700 Subject: [PATCH 5/9] fix(agent): production-safety audit remediations (7 findings) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes surfaced by the crates/agent safety audit. Each is independent and self-contained; the serve.rs hunks are only the OutSink/get_messages changes (a concurrent workstream's in-progress edits in the same file are intentionally left unstaged). 1. mcp_auth_store: create mcp_auth.json 0600 atomically + self-heal on every write. It holds live MCP OAuth access/refresh tokens; write_atomic only preserves an existing file's mode, so a brand-new file previously landed at the umask default (world-readable). Mirrors auth_store's auth.json handling. 2. serve get_messages{since}: guard the split_off with arr.len()==msg_ids.len(). msg_ids (active tree path) includes Custom entries that session.messages excludes, so after an append_custom RPC an idx past messages' end made split_off(idx+1) panic — a remotely-triggerable per-session crash. 3. worktree patch_paths: parse `git apply --numstat -z` (NUL-terminated, raw). Without -z git C-quotes non-ASCII paths ("s\303\251crets.env"), which slips past the --deny-path merge-back re-check (a **/*.env glob can't match a trailing .env"), merging a denied file back into the parent repo. 4. deps: bump quinn-proto 0.11.14->0.11.16 (RUSTSEC-2026-0185, CVSS 7.5) and crossbeam-epoch 0.9.18->0.9.20 (RUSTSEC-2026-0204). 5. session_store create_private: unlink-then-create_new instead of create(true).truncate(true) on the deterministic .tmp path, so a pre-planted file/symlink can't be opened-and-truncated (keeping a loose mode) or followed. 6. serve/serve_ws: bound each network connection's output channel (OutSink enum; 1024-frame cap, disconnect-on-full). A stalled mobile socket previously let the session buffer streamed frames without limit. stdio stays unbounded. 7. exec: prune completed receivers from PENDING_GROUP_KILLS on push, so the registry (otherwise drained only before process::exit) can't grow for the lifetime of a long-lived serve daemon. Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.lock | 24 ++++++++++---- crates/agent/src/mcp_auth_store.rs | 45 ++++++++++++++++++++++++++- crates/agent/src/serve.rs | 50 ++++++++++++++++++++++++++---- crates/agent/src/serve_ws.rs | 16 +++++++--- crates/agent/src/session_store.rs | 16 +++++++++- crates/agent/src/tools/exec.rs | 9 ++++++ crates/agent/src/worktree.rs | 14 +++++++-- 7 files changed, 153 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4c90bee..4fa2070 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -892,9 +892,9 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] @@ -1546,11 +1546,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 6.0.0", "rand_core 0.10.1", "wasip2", "wasip3", + "wasm-bindgen", ] [[package]] @@ -3398,15 +3400,16 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.14" +version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ "aws-lc-rs", "bytes", - "getrandom 0.3.4", + "getrandom 0.4.2", "lru-slab", - "rand 0.9.4", + "rand 0.10.1", + "rand_pcg", "ring", "rustc-hash", "rustls", @@ -3529,6 +3532,15 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", +] + [[package]] name = "rayon" version = "1.12.0" diff --git a/crates/agent/src/mcp_auth_store.rs b/crates/agent/src/mcp_auth_store.rs index 87343e9..ec047aa 100644 --- a/crates/agent/src/mcp_auth_store.rs +++ b/crates/agent/src/mcp_auth_store.rs @@ -132,11 +132,54 @@ fn write_store_file(path: &Path, store: &OnDisk) -> std::io::Result<()> { if let Some(parent) = path.parent() { fs::create_dir_all(parent)?; } + // This file holds live MCP OAuth access/refresh tokens. `write_atomic` only *preserves* an + // existing file's permission bits — it never imposes one on a brand-new file, so the very first + // write would otherwise land at the umask default (commonly world-readable). Create it `0600` + // atomically first, exactly as `auth_store.rs` does for `auth.json`. + if !path.exists() { + create_private(path)?; + } let body = serde_json::to_string_pretty(store).map_err(std::io::Error::other)?; let path_str = path .to_str() .ok_or_else(|| std::io::Error::other(format!("non-UTF-8 path: {}", path.display())))?; - crate::tools::write_atomic(path_str, body.as_bytes()) + crate::tools::write_atomic(path_str, body.as_bytes())?; + // Belt-and-suspenders on top of the atomic-at-creation `0600`: `write_atomic` only ever copies an + // existing file's mode forward, never re-asserts one, so a file loosened out-of-band (a stray + // `chmod`, a restore from backup) would stay that way forever. Re-assert `0600` after every write + // so it self-heals — mirrors `auth_store.rs::set_private`. + set_private(path) +} + +/// Create an empty file at `path` with `0600` permissions set atomically at creation (Unix), closing +/// the window a `set_permissions`-afterward approach would leave open. Mirrors +/// `auth_store.rs`/`session_store.rs`'s identical helper (duplicated, not shared, per this codebase's +/// small-self-contained-stores convention). +fn create_private(path: &Path) -> std::io::Result<()> { + let mut opts = fs::OpenOptions::new(); + opts.write(true).create(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + opts.mode(0o600); + } + opts.open(path)?; + Ok(()) +} + +/// Unconditionally re-assert `0600` on an already-written file (Unix only — a no-op elsewhere), so a +/// mode loosened out-of-band self-heals on the next write. Mirrors `auth_store.rs::set_private`. +fn set_private(path: &Path) -> std::io::Result<()> { + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + fs::set_permissions(path, fs::Permissions::from_mode(0o600))?; + } + #[cfg(not(unix))] + { + let _ = path; + } + Ok(()) } /// Acquire the cross-process lock, re-read the store's current on-disk state, apply `mutate`, persist diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index 81c551e..09f6666 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -1875,7 +1875,7 @@ pub async fn serve(cfg: ServeConfig) -> Result, Box { if let Value::Array(arr) = &mut messages { *arr = arr.split_off(idx + 1); @@ -7564,15 +7578,39 @@ impl From for OutFrame { /// [`crate::serve_ws`]). The supervisor [`add`](OutFanout::add)s a sink on attach and /// [`remove`](OutFanout::remove)s it on disconnect. An `std::sync::Mutex` (not tokio's) is deliberate: /// it is only ever held across non-`await` `send`s into unbounded channels. +/// One attached connection's output channel. A network transport (WebSocket/UDS) registers a +/// **bounded** sink: a client whose socket has stalled under TCP backpressure (a locked phone, a dead +/// tunnel) would otherwise let the still-running session buffer streamed frames without limit, growing +/// memory until the OS TCP layer finally times the socket out. When the bounded channel fills, +/// [`OutFanout::broadcast`] prunes the sink — a disconnect-on-full that drops the wedged connection and +/// leaves it to reconnect and replay committed state via `get_messages {since}`. The in-process +/// **stdio** transport is a single trusted local sink for the whole process life and stays +/// **unbounded**: pruning it on a transient stdout stall would sever the only client. +pub(crate) enum OutSink { + Bounded(mpsc::Sender), + Unbounded(mpsc::UnboundedSender), +} + +impl OutSink { + /// Non-blocking send. `Err` means the sink is dead (receiver gone) or — for a bounded sink — its + /// buffer is full (a stalled connection); [`OutFanout::broadcast`] prunes it either way. + fn try_send(&self, frame: OutFrame) -> Result<(), ()> { + match self { + OutSink::Bounded(tx) => tx.try_send(frame).map_err(|_| ()), + OutSink::Unbounded(tx) => tx.send(frame).map_err(|_| ()), + } + } +} + #[derive(Default)] pub(crate) struct OutFanout { next_id: u64, - sinks: Vec<(u64, mpsc::UnboundedSender)>, + sinks: Vec<(u64, OutSink)>, } impl OutFanout { /// Register a connection's sink; returns an id to [`remove`](Self::remove) it by on disconnect. - pub(crate) fn add(&mut self, tx: mpsc::UnboundedSender) -> u64 { + pub(crate) fn add(&mut self, tx: OutSink) -> u64 { let id = self.next_id; self.next_id += 1; self.sinks.push((id, tx)); @@ -7599,11 +7637,11 @@ impl OutFanout { match self.sinks.as_slice() { [] => {} [(_, tx)] => { - if tx.send(frame).is_err() { + if tx.try_send(frame).is_err() { self.sinks.clear(); } } - _ => self.sinks.retain(|(_, tx)| tx.send(frame.clone()).is_ok()), + _ => self.sinks.retain(|(_, tx)| tx.try_send(frame.clone()).is_ok()), } } } diff --git a/crates/agent/src/serve_ws.rs b/crates/agent/src/serve_ws.rs index a1bfca6..05c8fe3 100644 --- a/crates/agent/src/serve_ws.rs +++ b/crates/agent/src/serve_ws.rs @@ -52,7 +52,7 @@ use tokio_tungstenite::tungstenite::handshake::server::{ErrorResponse, Request, use tokio_util::sync::CancellationToken; use crate::serve::{ - OutFanout, OutFrame, ServeConfig, SharedOutConn, Signal, UpstreamHttp2, frame_to_line, + OutFanout, OutFrame, OutSink, ServeConfig, SharedOutConn, Signal, UpstreamHttp2, frame_to_line, lock_ignoring_poison, serve_session, }; use crate::session_store::{is_valid_session_id, new_id, scan_listings, scan_session_dir}; @@ -65,6 +65,14 @@ const WS_PATH: &str = "/_beyond/agent"; /// NAT/proxies. Also the granularity at which a wholly-dead socket is noticed (the ping send fails). const PING_INTERVAL: Duration = Duration::from_secs(30); +/// Depth of a single connection's outbound frame buffer. Bounds memory per attached socket: if a +/// client's write half stalls under TCP backpressure without erroring, the session keeps broadcasting +/// streamed frames into this channel — cap it so a wedged connection can't grow memory without limit +/// (`OutFanout::broadcast` prunes the sink when it fills, dropping the dead connection). Sized to +/// absorb a normal burst of streamed event frames while keeping the worst-case per-connection buffer +/// to a few MB. +const OUT_CHANNEL_BOUND: usize = 1024; + /// A live session, reachable by id across connections. The retained `input_tx` is what makes a /// dropped socket *not* an EOF — the session's `input_rx.recv()` pends until the next command instead /// of shutting down (see the module doc). @@ -204,12 +212,12 @@ impl Supervisor { // Register this connection's send channel as one of the session's output sinks — the session // broadcasts every frame to all registered sinks. Keep the `sink_id` to remove it on disconnect. - let (conn_tx, mut conn_rx) = mpsc::unbounded_channel::(); + let (conn_tx, mut conn_rx) = mpsc::channel::(OUT_CHANNEL_BOUND); // A direct handle to *this* connection's send channel, for supervisor-level replies (the // `list_daemon_sessions` command below) that must go back to *this* socket only, not fan out to // the other attached connections. Feeds the same `conn_rx`/send task. let reply_tx = conn_tx.clone(); - let sink_id = lock_ignoring_poison(&out_conn).add(conn_tx); + let sink_id = lock_ignoring_poison(&out_conn).add(OutSink::Bounded(conn_tx)); let (mut sink, mut stream) = ws.split(); @@ -285,7 +293,7 @@ impl Supervisor { .and_then(serde_json::Value::as_str) .map(str::to_owned); let frame = self.list_daemon_sessions(client_id).await; - let _ = reply_tx.send(frame); + let _ = reply_tx.try_send(frame); continue; } } diff --git a/crates/agent/src/session_store.rs b/crates/agent/src/session_store.rs index dc8f757..caaf235 100644 --- a/crates/agent/src/session_store.rs +++ b/crates/agent/src/session_store.rs @@ -3799,8 +3799,22 @@ fn read_capped_line(reader: &mut impl BufRead, buf: &mut Vec) -> std::io::Re /// pulled off disk — so they should never be readable by anyone but the owner, independent of the /// process umask. fn create_private(path: &Path) -> std::io::Result { + // Remove any leftover at this (deterministic `.tmp`) path first, then create it fresh with + // `create_new`. `.mode()` only takes effect when the file is actually *created*, so a plain + // `create(true).truncate(true)` on a pre-existing path would open-and-truncate a **pre-planted** + // file — keeping its existing, possibly world-readable, mode — or follow a pre-planted symlink to + // some other target, writing the transcript (which can hold secrets `read` pulled off disk) + // through it. A stale `.tmp` from a prior crash is the normal reason something is already here; + // unlink it (which removes a symlink itself, not its target) then create exclusively. If an + // attacker races a fresh plant in between, `create_new` fails outright — fail closed, never + // truncating or following it. + match fs::remove_file(path) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => return Err(e), + } let mut opts = OpenOptions::new(); - opts.write(true).create(true).truncate(true); + opts.write(true).create_new(true); #[cfg(unix)] { use std::os::unix::fs::OpenOptionsExt; diff --git a/crates/agent/src/tools/exec.rs b/crates/agent/src/tools/exec.rs index 50c6de0..aede8bb 100644 --- a/crates/agent/src/tools/exec.rs +++ b/crates/agent/src/tools/exec.rs @@ -246,6 +246,15 @@ impl Drop for GroupKillGuard { let _ = tx.send(()); }); if let Ok(mut pending) = PENDING_GROUP_KILLS.lock() { + // Opportunistically reclaim entries whose kill thread has already finished (`Ok`) or + // whose sender was dropped (`Disconnected`), so this registry — otherwise drained in + // full only just before `process::exit` — can't grow for the whole lifetime of a + // long-lived `serve` daemon that cancels one bash after another. A still-running kill + // thread's receiver (`Err(Empty)`) is kept, so `wait_for_pending_group_kills` can still + // block on it. + pending.retain(|rx| { + matches!(rx.try_recv(), Err(std::sync::mpsc::TryRecvError::Empty)) + }); pending.push(rx); } } diff --git a/crates/agent/src/worktree.rs b/crates/agent/src/worktree.rs index ee63bb0..a1c713f 100644 --- a/crates/agent/src/worktree.rs +++ b/crates/agent/src/worktree.rs @@ -405,10 +405,18 @@ pub enum ApplyOutcome { /// patch and reports `\t\t` **without modifying anything** — a hand-rolled patch /// parser here would be a second, subtly-different implementation of a format git already understands. async fn patch_paths(repo_root: &Path, patch: &[u8]) -> Result, String> { - let out = git_stdout_with_stdin(repo_root, &["apply", "--numstat", "-"], patch).await?; + // `-z` is load-bearing, not cosmetic. Without it, git **C-quotes** any path containing non-ASCII + // or special bytes: `sécrets.env` is reported as `"s\303\251crets.env"` (wrapping quotes + octal + // escapes). That quoted form then slips past the deny-glob re-check in `apply_patch` — a glob like + // `**/*.env` can't match a string ending in `.env"` — silently merging a denied file back into the + // parent repo. `-z` emits each record as `\t\t\0`, NUL-terminated and + // unquoted (renames already normalized to their destination path, same as without `-z`), so the + // deny check sees the real filename. + let out = git_stdout_with_stdin(repo_root, &["apply", "--numstat", "-z", "-"], patch).await?; Ok(String::from_utf8_lossy(&out) - .lines() - .filter_map(|line| line.rsplit('\t').next()) + .split('\0') + .filter(|record| !record.is_empty()) + .filter_map(|record| record.rsplit('\t').next()) .filter(|p| !p.is_empty()) .map(str::to_string) .collect()) From 9fbc070ef3048962250f28c6af23f673b4ebd8a8 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sun, 12 Jul 2026 15:23:39 -0700 Subject: [PATCH 6/9] fix(serve): flaky login-abort hang (the CI-blocking heisenbug) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `serve_auth` abort-login test hung intermittently under load (2-core CI runners hit it ~every run; it killed the runner at ~46min with no logs). Two independent bugs, found by reproducing under `taskset -c 0,1` and live-inspecting the stuck serve with gdb: 1. Test-ordering bug (the actual hang). `abort_login`'s own response and the aborted login's async terminal response are produced by different tasks and race — but the test read them in a fixed order. When the login terminal won the race, the abort-response read consumed and discarded it, and the following terminal-response read then hung forever on a frame already thrown away. Fixed with an order-independent `read_frames_matching_all` for both racing pairs. 2. A real serve race the fixed test then exposed. `abort_login` cleared the `pending_login` slot only asynchronously, via the detached task's `Drop` guard, so the next `login` could arrive before the slot freed and be spuriously rejected as "already in flight" — hanging a client that reasonably expects abort's success to mean the slot is free. Now `abort_login` clears the slot synchronously (`take`), and the guard is generation-tagged so a late-winding-down task can't wipe a successor's slot (which would make the next abort a silent no-op, leaving that login uncancellable). Verified: the previously-hanging test now passes 60/60 under 2-core concurrency (unfixed hung within 1-9), and the full integration suite runs clean 3x on 2 cores. New unit test pins the generation-guard behavior. This is the flake that was blocking #27's CI — pre-existing, unrelated to the memory work, surfaced because CI's slow runners lose the race almost every time. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/agent/src/serve.rs | 99 +++++++++++++++++++++++++++++--- crates/agent/tests/serve_auth.rs | 52 ++++++++++++++--- 2 files changed, 133 insertions(+), 18 deletions(-) diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index 09f6666..7db3eb2 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -2858,6 +2858,10 @@ pub(crate) async fn serve_session( // flow resolves, success or failure, so a later `login` is accepted again. let pending_login: Arc>> = Arc::new(std::sync::Mutex::new(None)); + // Monotonic tag stamped on each login that claims the slot, so a late-winding-down task's guard can + // tell its own slot from a successor's (see `PendingLoginGuard`). Only ever touched from this + // single-threaded command loop, so a plain counter suffices. + let mut login_generation: u64 = 0; loop { let line = tokio::select! { biased; @@ -5776,7 +5780,10 @@ pub(crate) async fn serve_session( // browser flow. Only a second concurrent `login` is rejected (above). let login_cancel = agent_core::CancellationToken::new(); let pending_code: PendingCodeSlot = Arc::new(std::sync::Mutex::new(None)); + login_generation += 1; + let generation = login_generation; *lock_ignoring_poison(&pending_login) = Some(PendingLogin { + generation, cancel: login_cancel.clone(), pending_code: pending_code.clone(), }); @@ -5784,10 +5791,12 @@ pub(crate) async fn serve_session( let pending_login_bg = pending_login.clone(); let login_id = id.clone(); tokio::spawn(async move { - // Always clears `pending_login` when this task ends, even on panic — see - // `PendingLoginGuard`'s own doc comment. - let _reset_pending_login_on_exit = - PendingLoginGuard(pending_login_bg.clone()); + // Always clears `pending_login` when this task ends, even on panic — but only + // if the slot still holds *this* login (see `PendingLoginGuard`'s doc comment). + let _reset_pending_login_on_exit = PendingLoginGuard { + slot: pending_login_bg.clone(), + generation, + }; let callbacks = ServeLoginCallbacks { out_tx: out_tx_bg.clone(), id: login_id.clone(), @@ -5866,7 +5875,14 @@ pub(crate) async fn serve_session( } "abort_login" => { // Idempotent no-op if none is in flight — matches `abort`'s own idle-mode convention. - if let Some(p) = lock_ignoring_poison(&pending_login).as_ref() { + // `take`, not `as_ref`: clear the slot *synchronously* here rather than leaving it for the + // detached task's `Drop` guard, so the very next `login` is accepted the instant this abort + // is acknowledged. Waiting for the guard raced the next command — a `login` arriving before + // the cancelled task wound down was spuriously rejected as "already in flight", hanging a + // client that (reasonably) expected `abort_login`'s success to mean the slot was free. The + // cancelled task still sends its own terminal `login` response; clearing the slot only frees + // the next login, and the guard's generation check keeps that task from wiping the successor. + if let Some(p) = lock_ignoring_poison(&pending_login).take() { p.cancel.cancel(); } emit!(response(id, "abort_login", true, None, None)); @@ -7196,7 +7212,14 @@ pub(crate) fn lock_ignoring_poison(m: &std::sync::Mutex) -> std::sync::Mut } /// Tracks the one `login` allowed in flight at a time — see `pending_login`'s own declaration. +/// +/// `generation` is a monotonic tag stamped when this login claimed the slot. It's what lets the +/// [`PendingLoginGuard`] tell "the slot still holds *my* login" from "a newer login has since taken it" +/// on drop, so a task winding down late can't wipe a successor's slot — the race that made a subsequent +/// `abort_login` silently no-op (the login it should have cancelled then never resolves). See +/// [`PendingLoginGuard`]. struct PendingLogin { + generation: u64, cancel: agent_core::CancellationToken, pending_code: PendingCodeSlot, } @@ -7210,11 +7233,28 @@ struct PendingLogin { /// `abort_login` can't help — it just cancels a token nothing is left alive to observe). A `Drop` /// impl runs during unwinding, unlike a plain line of cleanup code placed after the panicking call, /// so this holds regardless of where in the task body things go wrong. -struct PendingLoginGuard(Arc>>); +/// +/// **Clears only its own generation.** `abort_login` clears the slot *synchronously* (so the next +/// `login` is accepted the instant the abort is acknowledged, not whenever this detached task happens to +/// wind down), which means by the time this guard runs a *newer* login may already own the slot. Wiping +/// it unconditionally would drop that newer login's marker — the next `abort_login` would then find an +/// empty slot and no-op, leaving the newer login uncancellable and its client hung waiting for a terminal +/// response that never comes. Comparing generations makes the guard a no-op unless the slot is still this +/// login's. +struct PendingLoginGuard { + slot: Arc>>, + generation: u64, +} impl Drop for PendingLoginGuard { fn drop(&mut self) { - *lock_ignoring_poison(&self.0) = None; + let mut slot = lock_ignoring_poison(&self.slot); + if slot + .as_ref() + .is_some_and(|p| p.generation == self.generation) + { + *slot = None; + } } } @@ -7641,7 +7681,9 @@ impl OutFanout { self.sinks.clear(); } } - _ => self.sinks.retain(|(_, tx)| tx.try_send(frame.clone()).is_ok()), + _ => self + .sinks + .retain(|(_, tx)| tx.try_send(frame.clone()).is_ok()), } } } @@ -7901,12 +7943,16 @@ mod tests { // provider seam this module doesn't otherwise need. let pending_login: Arc>> = Arc::new(std::sync::Mutex::new(Some(PendingLogin { + generation: 1, cancel: CancellationToken::new(), pending_code: Arc::new(std::sync::Mutex::new(None)), }))); let slot = pending_login.clone(); let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - let _guard = PendingLoginGuard(slot.clone()); + let _guard = PendingLoginGuard { + slot: slot.clone(), + generation: 1, + }; panic!("simulated failure inside the login task"); })); assert!( @@ -7919,6 +7965,41 @@ mod tests { ); } + #[test] + fn pending_login_guard_only_clears_its_own_generation_not_a_successor() { + // The race behind the flaky `serve_auth` abort-login hang: `abort_login` clears the slot + // synchronously, so a *newer* login can claim it before the cancelled task's guard runs. That + // guard must then be a no-op — wiping the successor's marker would make the next `abort_login` + // find an empty slot and silently no-op, leaving the newer login uncancellable (its client hangs). + let slot: Arc>> = + Arc::new(std::sync::Mutex::new(None)); + // A newer login (generation 2) currently owns the slot. + let newer_cancel = CancellationToken::new(); + *lock_ignoring_poison(&slot) = Some(PendingLogin { + generation: 2, + cancel: newer_cancel.clone(), + pending_code: Arc::new(std::sync::Mutex::new(None)), + }); + // An older login's guard (generation 1) winds down now. + drop(PendingLoginGuard { + slot: slot.clone(), + generation: 1, + }); + // The successor's marker must survive, and still be the cancel `abort_login` would reach. + let held = lock_ignoring_poison(&slot); + assert!( + held.as_ref().is_some_and(|p| p.generation == 2), + "the older guard must not clear the newer login's slot" + ); + drop(held); + // And a guard whose generation *does* match clears it (the normal, own-slot case). + drop(PendingLoginGuard { + slot: slot.clone(), + generation: 2, + }); + assert!(lock_ignoring_poison(&slot).is_none()); + } + #[test] fn session_stats_reports_no_context_usage_before_any_turn_has_run() { // pi: agent-session-stats.test.ts — `contextUsage` is `null` when nothing has run yet (and, diff --git a/crates/agent/tests/serve_auth.rs b/crates/agent/tests/serve_auth.rs index 037f8bc..130fe2d 100644 --- a/crates/agent/tests/serve_auth.rs +++ b/crates/agent/tests/serve_auth.rs @@ -229,17 +229,23 @@ fn login_acks_then_a_second_concurrent_login_is_rejected_then_abort_login_cancel "{rejection:#?}" ); - // Cancel the first login; it must resolve with a failure response, not hang forever. + // Cancel the first login; it must resolve with a failure response, not hang forever. The + // `abort_login`'s own response (id 3) and the aborted login's terminal response (id 1) are produced + // by different tasks and race — accept them in either order rather than assuming one precedes the + // other (the assumption that made this test flaky under load). writeln!(stdin, "{}", json!({ "id": "3", "type": "abort_login" })).unwrap(); stdin.flush().unwrap(); - let abort_ack = - read_one_frame_matching(&mut stdout, |v| v["type"] == "response" && v["id"] == "3"); - assert_eq!(abort_ack["success"], true, "{abort_ack:#?}"); - - let login_result = read_one_frame_matching(&mut stdout, |v| { - v["type"] == "response" && v["command"] == "login" && v["id"] == "1" - }); - assert_eq!(login_result["success"], false, "{login_result:#?}"); + let framed = read_frames_matching_all( + &mut stdout, + &[&|v| v["type"] == "response" && v["id"] == "3", &|v| { + v["type"] == "response" && v["command"] == "login" && v["id"] == "1" + }], + ); + assert_eq!(framed[0]["success"], true, "abort_login: {framed:#?}"); + assert_eq!( + framed[1]["success"], false, + "aborted login must resolve failed: {framed:#?}" + ); // The slot must be cleared once the first login resolves — a fresh `login` is accepted again // immediately, not left rejected by a stale in-flight marker. Anthropic again (not a different @@ -288,3 +294,31 @@ fn read_one_frame_matching( } } } + +/// Read frames until **every** predicate in `preds` has matched at least one frame (in any order), +/// returning the matched frames in predicate order. Necessary when several independent responses race +/// and can interleave arbitrarily — e.g. an `abort_login`'s own synchronous response and the aborted +/// `login`'s async terminal response, which are produced by different tasks with no ordering guarantee +/// between them. Reading them with sequential `read_one_frame_matching` calls would discard whichever +/// arrives first while waiting for the other, then hang forever on the one already thrown away. +fn read_frames_matching_all( + reader: &mut impl std::io::BufRead, + preds: &[&dyn Fn(&Value) -> bool], +) -> Vec { + let mut found: Vec> = vec![None; preds.len()]; + let mut line = String::new(); + while found.iter().any(Option::is_none) { + line.clear(); + let n = reader.read_line(&mut line).unwrap(); + assert!(n > 0, "stream ended before all matching frames arrived"); + let Ok(v) = serde_json::from_str::(line.trim()) else { + continue; + }; + for (i, pred) in preds.iter().enumerate() { + if found[i].is_none() && pred(&v) { + found[i] = Some(v.clone()); + } + } + } + found.into_iter().map(Option::unwrap).collect() +} From 0e9d60882763356a703a749ec17a13d453ceb232 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sun, 12 Jul 2026 16:07:46 -0700 Subject: [PATCH 7/9] fix(agent-core): production-safety audit remediations (14 findings) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A safety audit of `agent-core` found the mechanical surface already clean (`unsafe_code = "forbid"`, panic lints denied workspace-wide, `overflow-checks` on in release), so every finding here is behavioral. The two that mattered: - **An un-terminated SSE event grew without bound.** `LineFramer`'s 32 MiB cap bounds a single *line* and resets on every terminator, but `SseEventBuffer` only drains on a blank line — so a relay that strips blank lines grew one `String` per `data:` line forever, under the line cap the whole way, until the process OOMed. Capped the *event*, not just the line. - **The `MaxTokens` silent-truncation retry could spin forever.** Its `overflow_recovered` guard was reset on the `Ok` path *upstream* of the check that read it, so it was dead code; the retry then decremented `steps_this_call`, so `max_steps` couldn't bound it either. Each round billed two model calls and made no progress. The flag now has a survives-exactly-one-iteration lifetime (`mem::replace` into a per-turn local), and `compact()` refuses the one cut that provably can't shrink (a re-summarized summary), so `Compacted` stays an honest signal that the caller made progress. Also: - Compaction budgets now scale with the model's real window (`CompactionConfig::for_window`). The 200k absolutes left a 32k model keeping more (20_000) than the threshold that triggers a compaction (16_384) — so it re-triggered every turn and never got back under the line. Windows >= ~64k are byte-for-byte unchanged. - `openai_responses` no longer defaults a missing `output_index` to `-1` (i.e. `usize::MAX` once cast): two tool calls collapsed onto that one synthetic index and the accumulator silently overwrote the first. Now refuses to guess, matching `anthropic::usize_at`'s posture. - Provider token counts clamp instead of truncating (`as u32` is not covered by `overflow-checks`, so `u32::MAX + 2` landed as `1` — making a full context look empty and *suppressing* the compaction that would have saved the turn). - The Codex WebSocket `send()` was the module's one unbounded await; idle cached connections are now swept rather than reaped lazily per-key. - `CheckpointHook::checkpoint` is panic-guarded like every other host seam — it's the one most likely to touch failing I/O, and a host that unwrapped on a full disk took the whole run down. - Steering lanes are bounded (they're fed straight from remote RPC); a full lane refuses the newest message and `serve` now acks that honestly instead of claiming `success: true` for a message it dropped. - Tool results are capped at the loop layer — the 50 KiB bound was a per-tool convention, and `Tool` is a public trait MCP servers implement. - A cancelled `write_lock` waiter no longer orphans its map entry; schema coercion has a work budget. Not done: `Error::Transport(String)` still flattens its source chain. `MID_STREAM_NETWORK_ERROR` is a *deliberate* transport-agnostic classification channel (a concrete `#[source]` would couple `agent-core` to `reqwest` and break the network-blind-core property), and a `Box` source means converting a tuple variant matched across two crates — diagnosability, not safety. Left for a focused change. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/agent-core/ARCHITECTURE.md | 28 +- crates/agent-core/src/agent.rs | 310 +++++++++++++++-- crates/agent-core/src/codex_websocket.rs | 320 +++++++++++++++--- crates/agent-core/src/compaction.rs | 78 +++++ crates/agent-core/src/dialect/anthropic.rs | 2 +- crates/agent-core/src/dialect/mod.rs | 81 ++++- crates/agent-core/src/dialect/openai.rs | 10 +- .../src/dialect/openai_responses.rs | 156 ++++++--- crates/agent-core/src/steering.rs | 176 +++++++++- crates/agent-core/src/validation.rs | 162 +++++++-- crates/agent-core/src/write_lock.rs | 174 ++++++++-- crates/agent/ARCHITECTURE.md | 12 +- crates/agent/src/serve.rs | 61 ++-- 13 files changed, 1376 insertions(+), 194 deletions(-) diff --git a/crates/agent-core/ARCHITECTURE.md b/crates/agent-core/ARCHITECTURE.md index 04fbd2c..5ef4252 100644 --- a/crates/agent-core/ARCHITECTURE.md +++ b/crates/agent-core/ARCHITECTURE.md @@ -210,7 +210,20 @@ builder on `Agent` or `ModelRequest`, and each is exercised by unit tests): not a fraction of the already-scaled `summary_max_tokens`, which would instead compound the two scale factors into an effective ~0.4× `reserve_tokens`), merged under a "Turn Context (split turn)" header — closer to pi's two-template - approach than summarizing the whole prefix with one minimal-context call. `summary_max_tokens` scales + approach than summarizing the whole prefix with one minimal-context call. **Both token budgets scale + with the model's real context window** (`CompactionConfig::for_window`, used by `Agent::new`): the + struct's `Default` states `reserve_tokens`/`keep_recent_tokens` as 200k-model _absolutes_ (16_384 / + 20_000), and seeding only `context_window` from the model's capabilities while leaving those two at + their 200k values breaks the invariant the pair has to satisfy between them — _what a compaction keeps + must stay comfortably under the threshold that triggers one_. The trigger fires at + `context_window - reserve_tokens`, so on the catalogue's smallest window (32_768) the un-scaled numbers + give a trigger budget of 16_384 while `keep_recent_tokens` retains 20_000: a compaction that keeps + _more_ than the line it just crossed, so every turn re-triggers, each spending a real summarization + call, and the context never gets back under the threshold. `for_window` caps the reserve at ¼ of the + window and `keep_recent_tokens` at half the resulting trigger budget. The caps are one-sided (`min`, + not a proportion) so every window at or above ~64k keeps byte-for-byte the tuning it had before — only + the genuinely small windows, the ones that were broken, see different numbers. + `summary_max_tokens` scales from the model's own `max_output` rather than a flat constant (`agent::scaled_summary_max_tokens`) — seeded by `Agent::new` and _rescaled_ by `with_compaction` whenever the incoming config still carries the struct's own flat default, so a @@ -415,9 +428,16 @@ builder on `Agent` or `ModelRequest`, and each is exercised by unit tests): while holding the same map lock every concurrent `lock()` call must also acquire to clone the `Arc` (mirrors pi's `file-mutation-queue.ts`'s identity-checked `finally`-block delete). Before this fix the map instead accumulated one entry per distinct path _ever_ locked for the registry's whole lifetime, - unbounded for a long-running `serve` process. Documented limitation: this - only serializes within one process — cross-process locking would need a filesystem advisory lock, - out of scope until a real multi-process use case needs it. + unbounded for a long-running `serve` process. The same eviction check also runs on the _acquire_ path + (a `PendingLock` drop guard, defused into the `WriteLockGuard` on success): a `lock()` future dropped + while still parked — a cancelled run — never constructs a guard, so without it the last reference to a + contended key could die with no one left to evict the entry, reintroducing exactly the leak above. + Two documented limitations: this only serializes within one process (cross-process locking would need + a filesystem advisory lock, out of scope until a real multi-process use case needs it); and + `futures::lock::Mutex` is _barging_, not FIFO — the registry guarantees mutual exclusion but no + fairness, so sustained contention on one path can starve a waiter. The only ordering that actually + matters (a turn's own same-target calls) is already guaranteed upstream, where they run serially in + call order as one group. - **Explicit reasoning/thinking disable** — `ModelCaps::reasoning_disableable` (per exact model id, not per `ThinkingShape`) drives an explicit "off" signal (Anthropic `{"type":"disabled"}`, OpenAI `{"effort":"none"}`) on a turn that isn't requesting thinking, for a model capable of one — instead of diff --git a/crates/agent-core/src/agent.rs b/crates/agent-core/src/agent.rs index ad6bf7a..4c83093 100644 --- a/crates/agent-core/src/agent.rs +++ b/crates/agent-core/src/agent.rs @@ -311,8 +311,13 @@ impl Agent { pub fn new(transport: Arc, model: impl Into) -> Self { let model = model.into(); let caps = crate::models::capabilities(&model); - let reserve_tokens = CompactionConfig::default().reserve_tokens; - let summary_max_tokens = scaled_summary_max_tokens(reserve_tokens, caps.max_output); + // Both compaction budgets scale with the model's real window (see `for_window`) — reading the + // reserve back off the *scaled* config rather than off `Default` matters here, because + // `scaled_summary_max_tokens` sizes the summarization call against the headroom that actually + // exists on this model, not against a 200k model's. + let compaction = CompactionConfig::for_window(caps.context_window); + let summary_max_tokens = + scaled_summary_max_tokens(compaction.reserve_tokens, caps.max_output); Self { transport, tools: ToolRegistry::new(), @@ -326,9 +331,8 @@ impl Agent { reasoning_effort: None, temperature: None, compaction: CompactionConfig { - context_window: caps.context_window, summary_max_tokens, - ..CompactionConfig::default() + ..compaction }, auto_retry: true, sequential_tools: false, @@ -685,10 +689,20 @@ impl Agent { // called: a crash here must not lose it, exactly the exposure `CheckpointHook`'s own doc // comment used to call out (the very first turn of a run was the one point a crash lost the // user's own submitted prompt entirely). - self.checkpoint.checkpoint(session).await; - // Set once we've already compacted to recover from a context-overflow this turn, so a second - // overflow gives up instead of looping. Reset after each turn that lands cleanly. + self.checkpoint_guarded(session).await; + // Set once we've already compacted to recover from a context-overflow *error* this turn, so a + // second overflow gives up instead of looping. Reset after each turn that lands cleanly — which + // is sound for this flag precisely because both arms that read it (`Err(e) if + // is_context_overflow(..)`) are evaluated inside the `match` on the turn's result, i.e. strictly + // *before* the `Ok` path reaches that reset. let mut overflow_recovered = false; + // The same idea for the *silent-truncation* recovery path (a `MaxTokens` stop that + // `is_hard_overflow` recognizes), which needs its own flag rather than sharing the one above: + // that check lives on the `Ok` path, *downstream* of the reset, so a shared flag is always false + // by the time it's read and the guard it's supposed to provide can never fire. Read via + // `mem::replace` into a per-turn local instead (see the take below), which gives it a + // survives-exactly-one-iteration lifetime that no reset ordering can defeat. + let mut truncation_recovered = false; // Steps taken *by this call*, distinct from `session.steps` (a lifetime total across every // call, used only for observability — the `step` field on emitted events). Checking the // ceiling against this instead means `Error::MaxSteps` is a per-call backstop a client can @@ -985,6 +999,16 @@ impl Agent { } }; overflow_recovered = false; + // Take the truncation guard for *this* turn and disarm it in one move, so it survives + // exactly one iteration: the turn immediately after a truncation recovery sees `true` (and + // therefore refuses to recover a second time), and every turn after that sees `false` again. + // Reading it into a local here — rather than testing the field directly down in the + // `MaxTokens` branch — is what makes the guard immune to the reset-ordering bug that made + // its predecessor dead code: there is no longer any path on which the flag is cleared + // between being set and being read. It also re-arms correctly through a tool round-trip, + // which returns to the top of the loop without ever reaching that branch. + let recovered_truncation_last_turn = + std::mem::replace(&mut truncation_recovered, false); last_stop_reason = turn.stop_reason; let malformed: HashMap = std::mem::take(&mut turn.malformed).into_iter().collect(); @@ -1060,7 +1084,7 @@ impl Agent { // immediately instead, leaving the queue untouched (the same persistent `Steering` handle a // later `prompt` call reads from — see `serve.rs`). if turn.stop_reason == StopReason::Refusal { - self.checkpoint.checkpoint(session).await; + self.checkpoint_guarded(session).await; sink(AgentEvent::AgentEnd { steps: session.steps, }); @@ -1092,13 +1116,13 @@ impl Agent { // hard-truncated non-answer with no indication anything went wrong. If auto-compaction is // enabled and the live prompt has genuinely reached (or this `MaxTokens` stop implies // it's about to reach) the raw context window, compact and retry this same turn instead — - // pi's own `_checkCompaction`+silent-overflow handling does the same. Guarded by the same - // `overflow_recovered` flag the error-based overflow-retry arm uses (reset at the top of - // every turn), so a second silent truncation right after this recovery doesn't loop - // forever; a genuine improvement should always show up as *some* progress within one - // retry, same rationale as the error-based path. + // pi's own `_checkCompaction`+silent-overflow handling does the same. Guarded by + // `truncation_recovered` (see its declaration) so a second silent truncation right after + // this recovery gives up and reports the truncated answer instead of looping forever; a + // genuine improvement should always show up as *some* progress within one retry, same + // rationale as the error-based path. if turn.stop_reason == StopReason::MaxTokens - && !overflow_recovered + && !recovered_truncation_last_turn && self.compaction.enabled && compaction::is_hard_overflow( session, @@ -1124,7 +1148,7 @@ impl Agent { Arc::make_mut(&mut session.messages).pop(); session.steps -= 1; steps_this_call -= 1; - overflow_recovered = true; + truncation_recovered = true; continue; } Ok(_) => { @@ -1159,7 +1183,7 @@ impl Agent { // the tool-calling half of a turn ever reached a checkpoint; a plain conversational reply // (the model's *most* common shape) never did, silently relying on a caller's own // post-run persist to ever see it recorded. - self.checkpoint.checkpoint(session).await; + self.checkpoint_guarded(session).await; // A pending graceful-stop request wins over draining follow-up/steer messages, exactly // as it wins over continuing tool-call turns below — the queue is left untouched (same // rationale as the refusal case above) so nothing queued for "next time" is lost. A @@ -1206,7 +1230,7 @@ impl Agent { sink(AgentEvent::Steered { messages: count }); // A plain user message ends the visible history here — a valid, resumable checkpoint // (see `CheckpointHook`) before the next model call. - self.checkpoint.checkpoint(session).await; + self.checkpoint_guarded(session).await; continue; } @@ -1226,7 +1250,7 @@ impl Agent { // checkpoints its own tool-less case separately) — a call requesting tools never falls // through to that other checkpoint, so this is the one and only checkpoint a `tool_use` turn // pays for here. - self.checkpoint.checkpoint(session).await; + self.checkpoint_guarded(session).await; // Run the tools and feed results back as a single user turn. A tool's own failure becomes // an error `tool_result`, not an aborted run — the model can react to it next turn. @@ -1679,7 +1703,7 @@ impl Agent { terminate &= wants_terminate; result_blocks.push(ContentBlock::ToolResult { tool_use_id: id.clone(), - content, + content: cap_tool_result(content), is_error, images, }); @@ -1702,7 +1726,7 @@ impl Agent { // A tool round-trip just landed: assistant `tool_use` and its matching `tool_result`s are // both committed now, so this is a valid, resumable checkpoint (see `CheckpointHook`) — the // one mid-run point a crash between here and the run's eventual end would otherwise lose. - self.checkpoint.checkpoint(session).await; + self.checkpoint_guarded(session).await; if terminate { // A tool requested completion (e.g. an `attempt_completion`/`exit` tool) and the whole // batch agreed. The results are already recorded; end the run as if the model had @@ -2164,6 +2188,19 @@ impl Agent { if cut.turn_start.is_none() && compaction::previous_summary(&session.messages[..1]).is_some() { + // `first_kept == 1` on a clean boundary means the prefix handed to the summarizer is + // *exactly* the prior summary and nothing else — so the "fresh" summary is a summary of a + // summary, `apply_summary` splices back a list of identical length, and the round is pure + // loss: one paid model call, zero tokens freed, and (because a summary is already bounded + // by `summary_max_tokens`) no realistic shrink either. Unlike the broader budget check + // below this holds *regardless* of how much has landed since — the content since the + // summary is all in the kept suffix, which this cut doesn't touch. Bailing here, before + // `CompactionStart` is even emitted, is what keeps `Compacted` an honest signal that the + // caller made progress; the overflow-retry paths in `run_events_steered` branch on exactly + // that and would otherwise re-prompt the identical request forever. + if cut.first_kept == 1 { + return Ok(CompactOutcome::AlreadyCompacted); + } let since_last_summary: u32 = session.messages[1..] .iter() .map(compaction::estimate_message_tokens) @@ -2284,7 +2321,7 @@ impl Agent { // very next turn immediately re-triggers the identical (already-paid-for) compaction again, and // if this was overflow-recovery, the resumed session lands right back in the same // context-overflow condition it just paid to escape. - self.checkpoint.checkpoint(session).await; + self.checkpoint_guarded(session).await; sink(AgentEvent::Compacted { messages_before: before, messages_after: session.messages.len(), @@ -2297,6 +2334,27 @@ impl Agent { Ok(CompactOutcome::Compacted) } + /// Persist the session through the host's [`CheckpointHook`], containing a panic in the host's own + /// persistence code instead of letting it destroy the run. + /// + /// Every other host-supplied seam in this loop — every `AgentHooks` method, the event `sink` — is + /// already wrapped in [`catch_tool_panic`]/`catch_sink_panic` and fails open (or, for the approval + /// gate, deliberately closed). `checkpoint` was the one that wasn't, and it is the seam *most* + /// likely to touch failing I/O: its own trait doc invites the host to do blocking work there + /// ("appending to a session file"), so a full disk or an `EACCES` in a host persistence path that + /// unwraps would unwind straight through the agent loop and take the whole run — in `serve`, the + /// whole session task — down with it. A checkpoint is an optimization (it bounds what a crash + /// loses); failing to take one must degrade to "this checkpoint didn't persist", never to "the run + /// died". Logged at `error` because a host whose persistence is panicking genuinely wants to know. + async fn checkpoint_guarded(&self, session: &Session) { + if let Err(msg) = catch_tool_panic(self.checkpoint.checkpoint(session)).await { + tracing::error!( + error = %msg, + "checkpoint hook panicked; continuing without persisting this checkpoint" + ); + } + } + /// Runs an *automatic* (proactive-threshold or hard-overflow) compaction and swallows a failure /// rather than propagating it — mirrors pi's `_runAutoCompaction`, which wraps the equivalent call /// in `try/catch` and emits `compaction_end { errorMessage }` on failure instead of letting it @@ -2528,6 +2586,43 @@ fn catch_sink_panic(f: impl FnOnce()) { /// applies uniformly to every tool, not just `bash`, which is the point: one dispatch-layer mechanism /// handles cancellation safely for all current and future tools, rather than requiring each one to /// cooperatively watch a cancellation signal and preserve its own partial state. +/// Hard ceiling, in bytes, on the tool-result content this loop will admit into a session. +/// +/// Until this existed the only bound on tool output was a *per-tool convention* — the built-ins cap +/// themselves at `tools::output::DEFAULT_MAX_BYTES` (50 KiB) — and a `Tool` is a public trait that +/// hosts and MCP servers implement. Nothing stopped an implementation from handing back 200 MB, and +/// nothing downstream would have clipped it: `compaction`'s `TOOL_RESULT_MAX_CHARS` truncates only +/// what is rendered *into the summarization prompt*, while the raw content stays in `session.messages`, +/// is `Arc`-cloned into every subsequent `ModelRequest`, and is serialized on every checkpoint. That +/// makes the memory bound a property of every tool author rather than of the system, which is the +/// wrong place for it. +/// +/// An order of magnitude above what a compliant tool emits, so this never clips one in practice — it +/// is a backstop against a rogue or buggy implementation, not a second opinion on how much a tool +/// should return. +const TOOL_RESULT_MAX_BYTES: usize = 1024 * 1024; + +/// Clamp one tool result to [`TOOL_RESULT_MAX_BYTES`], telling the model what happened rather than +/// silently handing it a truncated document it would reason over as if complete. +fn cap_tool_result(content: String) -> String { + if content.len() <= TOOL_RESULT_MAX_BYTES { + return content; + } + let dropped = content.len() - TOOL_RESULT_MAX_BYTES; + // `String` is UTF-8 and truncation is by bytes, so walk back to the nearest boundary — slicing + // mid-codepoint would panic, and this content is arbitrary tool output, not ASCII by assumption. + let mut end = TOOL_RESULT_MAX_BYTES; + while end > 0 && !content.is_char_boundary(end) { + end -= 1; + } + let mut out = content; + out.truncate(end); + out.push_str(&format!( + "\n\n[tool result truncated: {dropped} more bytes exceeded the {TOOL_RESULT_MAX_BYTES}-byte limit]" + )); + out +} + fn resolve_tool_result(result: Option) -> ToolCallResult { result.unwrap_or_else(|| { ( @@ -4247,6 +4342,179 @@ mod tests { ); } + /// Regression: a model that keeps coming back `MaxTokens`-truncated with no tool calls used to spin + /// the silent-truncation recovery arm forever. Two independent defects combined to remove every + /// backstop: (1) `overflow_recovered` was reset on the `Ok` path *upstream* of the `MaxTokens` + /// check that read it, so the "don't recover twice in a row" guard was dead code; and (2) the + /// recovery arm decremented `steps_this_call`, so `max_steps` could not bound it either. Each round + /// billed two model calls (the turn + its summarization) and made no progress. Assert the run + /// terminates, and that it does so in a bounded number of model calls rather than by luck. + #[tokio::test] + async fn a_model_that_always_truncates_terminates_instead_of_looping_forever() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct AlwaysTruncates { + calls: AtomicUsize, + } + #[async_trait] + impl ModelTransport for AlwaysTruncates { + async fn stream(&self, req: ModelRequest) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + // The summarization call is the one that carries no tools and a system prompt — answer + // it with a real (short) summary so `compact()` genuinely succeeds and the loop reaches + // the retry it used to spin on. Every *live* turn truncates, forever. + let is_summarization = req.system.is_some() && req.tools.is_empty(); + let text = if is_summarization { + "a summary" + } else { + "cut off" + }; + let stop_reason = if is_summarization { + StopReason::EndTurn + } else { + StopReason::MaxTokens + }; + let s = futures::stream::iter(vec![ + Ok(StreamEvent::MessageStart), + Ok(StreamEvent::TextDelta { + index: 0, + text: text.into(), + }), + Ok(StreamEvent::ContentBlockStop { index: 0 }), + Ok(StreamEvent::MessageStop { stop_reason }), + ]); + Ok(Box::pin(s)) + } + } + + let mut session = Session::new(); + session.messages = Arc::new(vec![ + Message::user("first request"), + Message::assistant(vec![ContentBlock::text("first done")]), + Message::user("second request"), + Message::assistant(vec![ContentBlock::text("second done")]), + Message::user("third request, the one that gets cut off"), + ]); + + let transport = Arc::new(AlwaysTruncates { + calls: AtomicUsize::new(0), + }); + let agent = + Agent::new(transport.clone(), "claude-opus-4-8").with_compaction(CompactionConfig { + // Trivially trips `is_hard_overflow` on a `MaxTokens` stop regardless of prompt size. + context_window: 10, + keep_recent_tokens: 1, + ..CompactionConfig::default() + }); + + let result = tokio::time::timeout( + Duration::from_secs(10), + agent.run_events(&mut session, |_| {}), + ) + .await + .expect("an always-truncating model must not loop forever"); + + assert!( + result.is_ok(), + "the truncated reply is the best available answer and should be reported, got: {result:?}" + ); + // Exactly one recovery is allowed: live turn + summarization + the retried live turn. The lower + // bound proves the recovery arm was actually entered (otherwise this test would pass + // vacuously); the upper bound proves the guard then refused to enter it a second time. + let calls = transport.calls.load(Ordering::SeqCst); + assert!( + (3..=4).contains(&calls), + "the silent-truncation recovery must fire exactly once — entered (>=3 calls) and then \ + refused a second time (<=4 calls); got {calls} calls" + ); + } + + #[tokio::test] + async fn a_panicking_checkpoint_hook_degrades_instead_of_killing_the_run() { + // `checkpoint` is the host seam most likely to touch failing I/O — its own trait doc invites + // blocking work like "appending to a session file" — and it was the only one not wrapped in + // `catch_tool_panic`. A host persistence path that unwraps on a full disk used to unwind + // straight through the loop and take the whole run (in `serve`, the whole session task) with + // it. Failing to persist a checkpoint must cost the checkpoint, not the run. + struct PanickingCheckpoint { + calls: Arc, + } + #[async_trait] + impl CheckpointHook for PanickingCheckpoint { + async fn checkpoint(&self, _session: &Session) { + self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + panic!("disk is full"); + } + } + + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let transport = Arc::new(MockTransport::new(vec![vec![ + StreamEvent::MessageStart, + StreamEvent::TextDelta { + index: 0, + text: "done".into(), + }, + StreamEvent::ContentBlockStop { index: 0 }, + StreamEvent::MessageStop { + stop_reason: StopReason::EndTurn, + }, + ]])); + let agent = Agent::new(transport, "claude-opus-4-8").with_checkpoint_hook(Arc::new( + PanickingCheckpoint { + calls: Arc::clone(&calls), + }, + )); + + let mut session = Session::new(); + session.push(Message::user("hi")); + let result = agent.run_events(&mut session, |_| {}).await; + + assert!( + result.is_ok(), + "a panicking checkpoint hook must not fail the run, got: {result:?}" + ); + assert!( + calls.load(std::sync::atomic::Ordering::SeqCst) > 0, + "the hook must actually have been called (otherwise this passes vacuously)" + ); + assert!( + session + .messages + .iter() + .any(|m| matches!(m.role, Role::Assistant)), + "the run must still have completed its turn normally" + ); + } + + #[test] + fn a_rogue_tool_result_is_capped_before_it_can_enter_the_session() { + // The built-in tools cap themselves at 50 KiB, but `Tool` is a public trait that hosts and MCP + // servers implement — so the bound has to live in the loop, not in every tool author's head. + // Anything under the ceiling must pass through byte-for-byte untouched. + let small = "x".repeat(1024); + assert_eq!(cap_tool_result(small.clone()), small); + + let rogue = "x".repeat(TOOL_RESULT_MAX_BYTES + 5_000); + let capped = cap_tool_result(rogue); + assert!( + capped.len() < TOOL_RESULT_MAX_BYTES + 500, + "a rogue result must be clamped to roughly the ceiling, got {} bytes", + capped.len() + ); + assert!( + capped.contains("truncated"), + "the model must be told the result was clipped rather than silently handed a partial \ + document it would reason over as if complete" + ); + + // Tool output is arbitrary bytes, not ASCII by assumption: truncating mid-codepoint would + // panic, so the cut must walk back to a char boundary. A multi-byte char straddling the limit + // is the case that would have blown up. + let multibyte = "é".repeat(TOOL_RESULT_MAX_BYTES); + let capped = cap_tool_result(multibyte); + assert!(capped.contains("truncated")); + } + #[test] fn is_retryable_mid_stream_matches_both_dialects_truncation_and_overload() { assert!(is_retryable_mid_stream(&Error::Transport( diff --git a/crates/agent-core/src/codex_websocket.rs b/crates/agent-core/src/codex_websocket.rs index 49547bf..32b6c6c 100644 --- a/crates/agent-core/src/codex_websocket.rs +++ b/crates/agent-core/src/codex_websocket.rs @@ -41,20 +41,28 @@ //! [`checkin`](CodexWebSocketCache::checkin) are both plain, synchronous, sub-millisecond critical //! sections over a [`std::sync::Mutex`]. //! -//! Unlike pi, there is no proactive idle-timeout/max-age background timer +//! Unlike pi, there is no *background* idle-timeout/max-age timer //! (`SESSION_WEBSOCKET_CACHE_TTL_MS`/`SESSION_WEBSOCKET_MAX_AGE_MS`, both enforced via `setTimeout` in -//! the real client) — a spawned-per-connection sweep task is a meaningfully bigger commitment (this -//! crate has never spawned a background task of its own before) for a purely eager-resource-reclaim -//! optimization, not a correctness requirement. Instead, [`checkout`](CodexWebSocketCache::checkout) -//! lazily evicts an entry past [`MAX_CONNECTION_AGE`] the next time it's looked up — a connection can -//! sit open somewhat longer than pi's proactive close in a quiet session, which only matters for -//! server-side resource pressure, not for this client's own correctness. A cached connection the -//! *server* silently closed between turns (idle timeout, restart) is handled reactively instead of -//! pi's `readyState`-checked-before-reuse: see [`Attempt`]'s "stale reused connection" retry, one -//! connect-with-a-fresh-socket retry before the ordinary SSE-fallback path, so an otherwise-healthy -//! session doesn't stick to SSE for its whole remaining lifetime just because one connection went -//! idle. If a real need for the proactive sweep appears (a long-lived `serve` process visibly -//! accumulating idle sockets), add it as its own follow-up. +//! the real client): a spawned-per-connection sweep task is a meaningfully bigger commitment (this +//! crate has never spawned a background task of its own before) than the reclaim it buys. What this +//! cache does instead is sweep *every* entry past [`MAX_CONNECTION_AGE`] out of the map on every +//! [`checkout`](CodexWebSocketCache::checkout) **and** [`checkin`](CodexWebSocketCache::checkin) — an +//! O(n) pass over a map that holds at most one entry per live conversation, cheap enough to run +//! unconditionally inside the same synchronous critical section, and not restricted to the one key +//! being looked up. That distinction is the whole point: an entry whose key is never checked out +//! again is exactly the entry that leaks, and a per-key lazy check can by construction never reach it. +//! A [`CachedConnection`] is not free to leave parked — it holds an open fd, its TLS session state, +//! and a [`Continuation`] carrying a full clone of that turn's entire input array — and nothing pumps +//! an idle socket, so the server's own pings go unanswered while it sits there. Since every turn of +//! every session touches this cache, any still-running client sweeps the corpses of the finished ones; +//! only a client whose sessions have *all* gone quiet holds sockets past the ceiling, and it releases +//! them when its next turn arrives or when it drops. +//! +//! A cached connection the *server* silently closed between turns (idle timeout, restart) is handled +//! reactively rather than via pi's `readyState`-checked-before-reuse: see [`Attempt`]'s "stale reused +//! connection" retry, one connect-with-a-fresh-socket retry before the ordinary SSE-fallback path, so +//! an otherwise-healthy session doesn't stick to SSE for its whole remaining lifetime just because one +//! connection went idle. //! //! # Delta diffing //! @@ -129,11 +137,30 @@ const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); /// (600s): the gateway/backend can legitimately go quiet for a long extended-thinking gap, so this /// isn't a ceiling on total turn duration, only on "has this socket gone silently dead". const IDLE_TIMEOUT: Duration = Duration::from_secs(600); -/// A cached connection older than this is treated as expired at its next checkout and closed rather -/// than reused — mirrors pi's `SESSION_WEBSOCKET_MAX_AGE_MS` (55 minutes), a defensive ceiling -/// independent of activity (some backends cap a single connection's total lifetime regardless of how -/// recently it was used). +/// Cap on getting one turn's single outbound `response.create` frame out to the OS. `SinkExt::send` is +/// feed-then-flush, so it stays pending until the bytes are actually accepted by the socket: a peer +/// that completes the upgrade and then stops reading (a TCP zero-window, a wedged load balancer, or a +/// *cached* socket whose peer went away without ever saying so) would otherwise park this await +/// forever, and this is the only await in the module that isn't already bounded. Nothing else can +/// break it — the SSE fallback is downstream of it, and the only remaining escape is the run's +/// cancellation token, which a detached background session has nobody to press. Sized like +/// [`CONNECT_TIMEOUT`] rather than [`IDLE_TIMEOUT`]: unlike a read, this +/// write isn't waiting on the model to think, so a healthy path finishes in milliseconds and has no +/// legitimate reason to outlast establishing the connection did. +const SEND_TIMEOUT: Duration = Duration::from_secs(10); +/// A cached connection older than this is swept out of the idle pool and closed rather than reused — +/// mirrors pi's `SESSION_WEBSOCKET_MAX_AGE_MS` (55 minutes), a defensive ceiling independent of +/// activity (some backends cap a single connection's total lifetime regardless of how recently it was +/// used). const MAX_CONNECTION_AGE: Duration = Duration::from_secs(55 * 60); +/// Ceiling on how many distinct `cache_key`s one cache remembers as sticky-SSE. The marker is +/// deliberately never *un*set for a key (see [`CodexWebSocketCache::mark_fallback`]), so without a +/// ceiling the set is a strictly-growing string map for the whole life of a `GatewayClient` — bounded +/// in practice by the number of conversations one client ever serves, which for a long-lived process +/// is not a bound at all. Far above any realistic live-conversation count per client, so eviction only +/// ever discards markers old enough that re-attempting the WebSocket for them is the right call +/// anyway. +const MAX_SSE_FALLBACK_KEYS: usize = 256; /// The OpenAI-Beta opt-in the real Codex backend's WebSocket upgrade requires — pi's /// `OPENAI_BETA_RESPONSES_WEBSOCKETS`. const OPENAI_BETA_RESPONSES_WEBSOCKETS: &str = "responses_websockets=2026-02-06"; @@ -276,47 +303,61 @@ struct CachedConnection { continuation: Option, } +/// Drop every entry whose age is at or past [`MAX_CONNECTION_AGE`], regardless of key. Generic over +/// the entry so the age policy itself is exercisable without a live TLS socket (a [`WsSocket`] can't +/// be constructed without a real peer to hand it). +/// +/// Callers run this with the map's guard already held: it is deliberately synchronous, and dropping a +/// swept [`CachedConnection`] closes its fd without a WebSocket close handshake — the peer will see a +/// TCP FIN on a socket it has heard nothing from for close to an hour, which is precisely what the +/// ceiling exists to tell it. +fn sweep_expired(entries: &mut HashMap, age_of: impl Fn(&T) -> Instant) { + entries.retain(|_, entry| age_of(entry).elapsed() < MAX_CONNECTION_AGE); +} + /// One per [`GatewayClient`](crate::client::GatewayClient) — see the module doc comment for the /// "present == idle" cache-key design. pub(crate) struct CodexWebSocketCache { idle: Mutex>, - sse_fallback: Mutex>, + /// Value is the instant the marker was set — read only to pick a victim when the map is at + /// [`MAX_SSE_FALLBACK_KEYS`], never to expire a marker on its own. + sse_fallback: Mutex>, } impl CodexWebSocketCache { pub(crate) fn new() -> Self { Self { idle: Mutex::new(HashMap::new()), - sse_fallback: Mutex::new(std::collections::HashSet::new()), + sse_fallback: Mutex::new(HashMap::new()), } } - /// Take ownership of `key`'s cached connection, if any is idle and not past - /// [`MAX_CONNECTION_AGE`]. A brief, synchronous, lock-held critical section only — never awaits. + /// Take ownership of `key`'s cached connection, if one is idle and still inside + /// [`MAX_CONNECTION_AGE`], sweeping every *other* expired entry out on the way through. A brief, + /// synchronous, lock-held critical section only — never awaits. fn checkout(&self, key: &str) -> Option { let mut guard = self .idle .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - let entry = guard.remove(key)?; - if entry.created_at.elapsed() >= MAX_CONNECTION_AGE { - // Dropped here: past its age ceiling, close it rather than hand it back for reuse. - None - } else { - Some(entry) - } + sweep_expired(&mut guard, |entry| entry.created_at); + guard.remove(key) } - /// Return a still-healthy connection to the idle pool after a clean turn completion. Overwrites - /// (and thus closes) any entry already present for `key` — the only way that can happen is two - /// concurrent turns on the same `cache_key` both completing and re-caching; whichever checks in - /// last wins, which is a benign, documented simplification (see the module doc comment) rather + /// Return a still-healthy connection to the idle pool after a clean turn completion, sweeping any + /// expired entries out at the same time — this is the moment a session goes quiet, so it's the + /// last chance to notice that some *other* session went quiet an hour ago and never came back. + /// + /// Overwrites (and thus closes) any entry already present for `key` — the only way that can happen + /// is two concurrent turns on the same `cache_key` both completing and re-caching; whichever checks + /// in last wins, which is a benign, documented simplification (see the module doc comment) rather /// than a correctness issue (each turn already sent/received its own complete, correct exchange). fn checkin(&self, key: String, connection: CachedConnection) { let mut guard = self .idle .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); + sweep_expired(&mut guard, |entry| entry.created_at); guard.insert(key, connection); } @@ -324,18 +365,33 @@ impl CodexWebSocketCache { /// `cache_key` skips the WebSocket attempt entirely and goes straight to HTTP/SSE, for the rest of /// this cache's lifetime (mirrors pi's `websocketSseFallbackSessions`, which is likewise never /// cleared except by an explicit debug/reset call this crate has no equivalent of). + /// + /// The set of marked keys is capped at [`MAX_SSE_FALLBACK_KEYS`], evicting the longest-marked key + /// to make room. Losing a marker is not a correctness problem — it costs the *next* turn on that + /// key one more failed WebSocket attempt before it falls back and re-marks — whereas an unbounded + /// map is a slow leak in any process whose `GatewayClient` outlives the conversations it serves. fn mark_fallback(&self, key: &str) { - self.sse_fallback + let mut guard = self + .sse_fallback .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(key.to_string()); + .unwrap_or_else(std::sync::PoisonError::into_inner); + if guard.len() >= MAX_SSE_FALLBACK_KEYS && !guard.contains_key(key) { + let oldest = guard + .iter() + .min_by_key(|(_, marked_at)| *marked_at) + .map(|(key, _)| key.clone()); + if let Some(oldest) = oldest { + guard.remove(&oldest); + } + } + guard.insert(key.to_string(), Instant::now()); } fn is_fallback_active(&self, key: &str) -> bool { self.sse_fallback .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .contains(key) + .contains_key(key) } } @@ -719,9 +775,25 @@ async fn attempt_once( }; let wire_body = build_wire_body(full_body, continuation.as_ref()); - if let Err(e) = socket.send(Message::text(wire_frame(&wire_body))).await { - tracing::debug!(error = %e, "codex websocket send failed; using SSE for this turn"); - return AttemptOnceOutcome::Fallback; + let send = socket.send(Message::text(wire_frame(&wire_body))); + match tokio::time::timeout(SEND_TIMEOUT, send).await { + // A peer that took the upgrade and then stopped reading our bytes is indistinguishable, from + // here, from one that rejected them outright — and the answer to both is the same: this turn + // goes out over HTTP/SSE. See [`SEND_TIMEOUT`] for why leaving the flush pending is not an + // option. + Err(_elapsed) => { + tracing::debug!( + timeout = ?SEND_TIMEOUT, + reused, + "codex websocket send timed out; using SSE for this turn" + ); + return AttemptOnceOutcome::Fallback; + } + Ok(Err(e)) => { + tracing::debug!(error = %e, "codex websocket send failed; using SSE for this turn"); + return AttemptOnceOutcome::Fallback; + } + Ok(Ok(())) => {} } let mut receiver = Receiver::new(socket); @@ -1150,4 +1222,174 @@ mod tests { "the marker must not leak across keys" ); } + + /// The idle pool's own values are `CachedConnection`s, which can't be built without a live peer to + /// hand a `WsSocket` — so the sweep is generic over its entry (see [`sweep_expired`]) and its age + /// policy is exercised here against bare timestamps, which is the entire input it actually reads. + fn aged_map(entries: &[(&str, Duration)]) -> HashMap { + let now = Instant::now(); + entries + .iter() + .map(|(key, age)| { + let created_at = now.checked_sub(*age).expect("test ages fit in an Instant"); + ((*key).to_string(), created_at) + }) + .collect() + } + + #[test] + fn the_sweep_drops_every_expired_entry_not_just_the_one_being_looked_up() { + // The whole point of sweeping the map rather than age-checking one key: the entry that leaks + // is precisely the one whose session ended and will never be checked out again. + let mut entries = aged_map(&[ + ( + "finished_an_hour_ago", + MAX_CONNECTION_AGE + Duration::from_secs(60), + ), + ("exactly_at_the_ceiling", MAX_CONNECTION_AGE), + ("still_live", Duration::from_secs(30)), + ]); + sweep_expired(&mut entries, |created_at| *created_at); + assert_eq!( + entries.keys().collect::>(), + vec!["still_live"], + "only the connection inside the age ceiling may survive: {entries:#?}" + ); + } + + #[test] + fn the_sweep_keeps_everything_inside_the_age_ceiling() { + let mut entries = aged_map(&[ + ("a", Duration::ZERO), + ("b", MAX_CONNECTION_AGE - Duration::from_secs(1)), + ]); + sweep_expired(&mut entries, |created_at| *created_at); + assert_eq!(entries.len(), 2, "a healthy pool must not be evicted"); + } + + /// A `WsSocket` over a real (but never spoken-to) loopback TCP connection. The sweep only ever + /// reads an entry's `created_at` and drops the rest, so no handshake or peer is needed — but the + /// type is concrete, so *some* socket is. + async fn parked_socket() -> WsSocket { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("loopback bind"); + let addr = listener.local_addr().expect("bound port"); + let client = tokio::net::TcpStream::connect(addr).await.expect("connect"); + listener.accept().await.expect("accept"); + tokio_tungstenite::WebSocketStream::from_raw_socket( + tokio_tungstenite::MaybeTlsStream::Plain(client), + tokio_tungstenite::tungstenite::protocol::Role::Client, + None, + ) + .await + } + + async fn cached(created_at: Instant) -> CachedConnection { + CachedConnection { + socket: parked_socket().await, + created_at, + continuation: None, + } + } + + #[tokio::test] + async fn checkout_and_checkin_both_sweep_expired_entries() { + // Both directions matter: checkin is the moment a session goes quiet (the last chance to + // notice that some *other* session went quiet an hour ago), checkout is every turn of every + // session. Either way the key being reaped is one nobody asked about. + let stale = Instant::now() + .checked_sub(MAX_CONNECTION_AGE + Duration::from_secs(1)) + .expect("test ages fit in an Instant"); + + for entry_point in ["checkout", "checkin"] { + let cache = CodexWebSocketCache::new(); + cache.checkin("abandoned".to_string(), cached(stale).await); + + if entry_point == "checkout" { + assert!(cache.checkout("some_other_session").is_none()); + } else { + cache.checkin( + "some_other_session".to_string(), + cached(Instant::now()).await, + ); + } + + let guard = cache.idle.lock().expect("uncontended in a unit test"); + assert!( + !guard.contains_key("abandoned"), + "{entry_point} must reap the expired entry it wasn't asked about" + ); + } + } + + #[tokio::test] + async fn checkout_returns_a_live_entry_and_refuses_an_expired_one() { + let cache = CodexWebSocketCache::new(); + cache.checkin("live".to_string(), cached(Instant::now()).await); + assert!(cache.checkout("live").is_some(), "inside the age ceiling"); + assert!( + cache.checkout("live").is_none(), + "checkout removes: a second concurrent turn on one key must not share the socket" + ); + + let stale = Instant::now() + .checked_sub(MAX_CONNECTION_AGE) + .expect("test ages fit in an Instant"); + cache.checkin("old".to_string(), cached(stale).await); + assert!( + cache.checkout("old").is_none(), + "at the age ceiling it must be closed, not handed back" + ); + } + + #[test] + fn the_sticky_sse_marker_set_is_capped() { + // Never cleared per-key by design, so the only thing standing between a long-lived + // GatewayClient and an unbounded string map is this ceiling. + let cache = CodexWebSocketCache::new(); + for n in 0..MAX_SSE_FALLBACK_KEYS { + cache.mark_fallback(&format!("session_{n}")); + } + assert_eq!( + cache + .sse_fallback + .lock() + .expect("uncontended in a unit test") + .len(), + MAX_SSE_FALLBACK_KEYS + ); + + cache.mark_fallback("one_too_many"); + let guard = cache + .sse_fallback + .lock() + .expect("uncontended in a unit test"); + assert_eq!(guard.len(), MAX_SSE_FALLBACK_KEYS, "the cap must hold"); + assert!( + guard.contains_key("one_too_many"), + "the key that just failed is the one most worth remembering" + ); + let survivors = (0..MAX_SSE_FALLBACK_KEYS) + .filter(|n| guard.contains_key(&format!("session_{n}"))) + .count(); + assert_eq!( + survivors, + MAX_SSE_FALLBACK_KEYS - 1, + "exactly one — the longest-marked — key may be evicted to make room" + ); + } + + #[test] + fn re_marking_a_key_already_at_the_cap_evicts_nothing() { + let cache = CodexWebSocketCache::new(); + for n in 0..MAX_SSE_FALLBACK_KEYS { + cache.mark_fallback(&format!("session_{n}")); + } + cache.mark_fallback("session_0"); + assert!( + (0..MAX_SSE_FALLBACK_KEYS).all(|n| cache.is_fallback_active(&format!("session_{n}"))), + "refreshing an existing marker doesn't need room made for it" + ); + } } diff --git a/crates/agent-core/src/compaction.rs b/crates/agent-core/src/compaction.rs index 52a312f..6816e7e 100644 --- a/crates/agent-core/src/compaction.rs +++ b/crates/agent-core/src/compaction.rs @@ -146,6 +146,46 @@ pub struct CompactionConfig { pub enabled: bool, } +impl CompactionConfig { + /// Defaults scaled to a model's real context window. + /// + /// [`Default`] is sized for a 200k-context Claude model and states both budgets as *absolutes* + /// (`reserve_tokens: 16_384`, `keep_recent_tokens: 20_000`). Seeding only `context_window` from the + /// model's capabilities and leaving those two at their 200k values — which is what [`Agent::new`] + /// used to do — silently breaks the one invariant the two budgets have to satisfy between them: + /// + /// > what a compaction *keeps* must be comfortably smaller than the threshold that *triggers* one. + /// + /// The trigger fires at `context_window - reserve_tokens`, so on a 32_768-token model (the smallest + /// in the catalogue) the un-scaled numbers give a trigger budget of 16_384 while `keep_recent_tokens` + /// retains 20_000 — a compaction that keeps *more* than the line it just crossed. Every turn then + /// re-triggers, each one spending a real summarization call, and the context never gets back under + /// the threshold. Scaling both budgets down with the window restores the invariant. + /// + /// The caps are deliberately one-sided (`min`, not a proportion): any window at or above ~64k gets + /// byte-for-byte the same tuning as before, since a quarter of the window already exceeds the + /// absolute reserve and half the remaining budget already exceeds `keep_recent_tokens`. Only the + /// genuinely small windows — the ones that were broken — see different numbers. + pub fn for_window(context_window: u32) -> Self { + let d = Self::default(); + // Never hand the reserve more than a quarter of the window: at the absolute 16_384 a 32k model + // would surrender half its context to headroom before the first token is sent. + let reserve_tokens = d.reserve_tokens.min(context_window / 4); + // What's actually usable once the reserve is set aside — and the exact quantity the trigger + // compares against, so `keep_recent_tokens` must stay a fraction of *this*, not of the window. + let trigger_budget = context_window.saturating_sub(reserve_tokens); + // Half the trigger budget: a compaction lands the context at roughly 50% of the threshold, so + // there's real room to work before the next one fires rather than re-triggering immediately. + let keep_recent_tokens = d.keep_recent_tokens.min(trigger_budget / 2); + Self { + context_window, + reserve_tokens, + keep_recent_tokens, + ..d + } + } +} + impl Default for CompactionConfig { fn default() -> Self { // Defaults sized for a 200k-context Claude model; override per deployment. Matches pi's own @@ -1006,6 +1046,44 @@ mod tests { use super::*; use serde_json::json; + #[test] + fn scaled_budgets_keep_less_than_the_threshold_that_triggers_a_compaction() { + // The invariant the two budgets have to satisfy between them, across every window in the + // catalogue: a compaction must *keep* less than the trigger threshold it just crossed. + // Otherwise the very next turn re-triggers, pays for another summarization, and never gets + // back under the line — which is exactly what the un-scaled 200k absolutes did on a 32k model + // (trigger budget 16_384, keep_recent 20_000). + for window in [4_096u32, 8_192, 32_768, 40_960, 128_000, 200_000, 1_000_000] { + let cfg = CompactionConfig::for_window(window); + let trigger_budget = window - cfg.reserve_tokens; + assert!( + cfg.keep_recent_tokens < trigger_budget, + "window {window}: keep_recent ({}) must stay under the trigger budget ({trigger_budget})", + cfg.keep_recent_tokens + ); + assert!( + cfg.reserve_tokens < window, + "window {window}: the reserve must not swallow the whole context" + ); + } + } + + #[test] + fn scaling_leaves_every_large_window_byte_for_byte_unchanged() { + // The caps are one-sided on purpose: they must only bite on the small windows that were + // actually broken. Any window at or above ~64k keeps precisely the tuning it had before. + let d = CompactionConfig::default(); + for window in [128_000u32, 200_000, 256_000, 1_000_000] { + let cfg = CompactionConfig::for_window(window); + assert_eq!(cfg.reserve_tokens, d.reserve_tokens, "window {window}"); + assert_eq!( + cfg.keep_recent_tokens, d.keep_recent_tokens, + "window {window}" + ); + assert_eq!(cfg.context_window, window); + } + } + fn convo() -> Vec { vec![ Message::user("the original task: refactor foo"), diff --git a/crates/agent-core/src/dialect/anthropic.rs b/crates/agent-core/src/dialect/anthropic.rs index 0c9dd15..829b4fc 100644 --- a/crates/agent-core/src/dialect/anthropic.rs +++ b/crates/agent-core/src/dialect/anthropic.rs @@ -1012,7 +1012,7 @@ fn str_at<'a>(v: Option<&'a Value>, key: &str) -> &'a str { fn u32_at(v: Option<&Value>, key: &str) -> u32 { v.and_then(|v| v.get(key)) .and_then(Value::as_u64) - .unwrap_or(0) as u32 + .map_or(0, super::saturating_u32) } /// Anthropic's own `content_block_start`/`_delta`/`_stop` all carry a real `index` — read it straight diff --git a/crates/agent-core/src/dialect/mod.rs b/crates/agent-core/src/dialect/mod.rs index f80e009..409cd8a 100644 --- a/crates/agent-core/src/dialect/mod.rs +++ b/crates/agent-core/src/dialect/mod.rs @@ -571,11 +571,29 @@ pub trait StreamDecoder: Send { #[derive(Default)] pub struct SseEventBuffer { data: Vec, + /// Running total of every buffered payload's length, so the cap below is a single add and compare + /// rather than a walk of `data` on each line. + bytes: usize, /// This not-yet-flushed event's SSE-level `event:` field, if one was sent — e.g. `"error"` for an /// Anthropic `event: error` frame. Matches pi's `state.event` (`SseDecoderState`). event: Option, } +/// Ceiling on one buffered (un-terminated) *event* — the sum of every `data:` payload accumulated +/// since the last blank line. +/// +/// [`LineFramer::extend`]'s own `MAX_BUFFERED_LINE_BYTES` bounds a single un-terminated *line*, and +/// its buffer is reset by `split_to` on every terminator — so it cannot bound this. An event is only +/// drained on a blank line ([`SseEventBuffer::take`]), so a stream of short, well-terminated `data:` +/// lines that never sends the blank-line event terminator (a spec-violating relay, or a proxy that +/// strips blank lines) grows `data` by one `String` per line forever while the line cap never trips: +/// the process OOMs instead of erroring out. The two caps are equal on purpose — a proxy that re-wraps +/// one long logical line into several `data:` lines produces an event of exactly the size the line cap +/// would have allowed unwrapped, so no stream the line cap accepts is rejected here. +/// +/// [`LineFramer::extend`]: crate::client::LineFramer::extend +const MAX_BUFFERED_EVENT_BYTES: usize = 32 * 1024 * 1024; + impl SseEventBuffer { /// A buffer with no pending data lines. pub fn new() -> Self { @@ -583,9 +601,17 @@ impl SseEventBuffer { } /// Buffer one `data:` line's payload (already stripped of the `data:` prefix and surrounding - /// whitespace). - fn push_data(&mut self, payload: &str) { + /// whitespace). Errors if this event has accumulated more than [`MAX_BUFFERED_EVENT_BYTES`] + /// without a terminating blank line — see that constant's doc comment. + fn push_data(&mut self, payload: &str) -> Result<()> { + self.bytes = self.bytes.saturating_add(payload.len()); + if self.bytes > MAX_BUFFERED_EVENT_BYTES { + return Err(Error::Transport(format!( + "SSE event exceeded {MAX_BUFFERED_EVENT_BYTES} bytes without a terminating blank line" + ))); + } self.data.push(payload.to_string()); + Ok(()) } /// Record this event's SSE-level `event:` field, overwriting any earlier value buffered for the @@ -602,6 +628,7 @@ impl SseEventBuffer { /// preserves existing behavior for a bare `event: ping`-style keepalive with no data. fn take(&mut self) -> Option<(Option, String)> { let is_error_event = self.event.as_deref() == Some("error"); + self.bytes = 0; if self.data.is_empty() && !is_error_event { self.event = None; return None; @@ -690,7 +717,7 @@ pub fn push_sse_line( if payload.is_empty() || payload == "[DONE]" { return Ok(Vec::new()); } - buf.push_data(payload); + buf.push_data(payload)?; Ok(Vec::new()) } @@ -784,6 +811,21 @@ fn repair_and_parse(decoder: &dyn StreamDecoder, payload: &str) -> Option /// Anthropic's `{"type":"error","error":{"message":…}}`, OpenAI Chat Completions' bare /// `{"error":{"message":…}}`, and the OpenAI Responses API's flat `{"type":"error","code":…, /// "message":…}` (no nested `error` object — a genuinely different shape from the other two, since +/// Narrow a wire-reported token count to the `u32` the usage types carry, clamping rather than +/// truncating. +/// +/// Every dialect reads usage out of provider JSON as a `u64` and has to land it in a `u32`. A plain +/// `as` cast is *not* covered by `overflow-checks` — it wraps silently — so a provider or proxy +/// reporting `"prompt_tokens": 4294967297` would land `1`. That value feeds `Session::record_usage` +/// → `last_input_tokens` → the compaction trigger, so a garbage-large report makes a nearly-full +/// context look empty and *suppresses* the compaction that would have saved the turn, which is a +/// strictly worse failure than over-reporting. `session`'s own accumulation is already saturating for +/// exactly this reason (provider usage numbers carry no upper-bound validation); clamping here is +/// what keeps that guarantee from being defeated one layer up. +pub(crate) fn saturating_u32(n: u64) -> u32 { + u32::try_from(n).unwrap_or(u32::MAX) +} + /// Responses streams a top-level `error` *event* rather than an in-band error field). Returns `None` /// for ordinary stream events. // `pub(crate)`: also reused by `codex_websocket`'s WebSocket receive loop, which decodes the exact @@ -1241,6 +1283,39 @@ data: {"type":"message_stop"} assert!(push_sse_line(&mut dec, &mut buf3, "").unwrap().is_empty()); } + #[test] + fn an_event_that_never_terminates_errors_instead_of_growing_without_bound() { + // `LineFramer`'s own `MAX_BUFFERED_LINE_BYTES` bounds a single un-terminated *line*, and its + // buffer resets on every terminator — so short, well-terminated `data:` lines that never see a + // blank line slip past it entirely while this buffer grows one `String` per line forever. Feed + // exactly that shape (a relay that strips blank lines) and require an error, not an OOM. + let mut buf = SseEventBuffer::new(); + let mut dec = anthropic::Decoder::default(); + let payload = "x".repeat(64 * 1024); + let line = format!("data: {payload}"); + // 32 MiB cap / 64 KiB per line = 512 lines to reach it; anything past that must error. + let mut err = None; + for _ in 0..1024 { + if let Err(e) = push_sse_line(&mut dec, &mut buf, &line) { + err = Some(e); + break; + } + } + let err = err.expect("an event with no terminating blank line must eventually error"); + assert!( + matches!(&err, Error::Transport(m) if m.contains("without a terminating blank line")), + "expected a transport error naming the cap, got: {err:?}" + ); + + // The cap is per *event*, not per stream: a normal flush resets the counter, so a long-lived + // stream of ordinary events never accumulates toward it. + let mut buf = SseEventBuffer::new(); + for _ in 0..1024 { + push_sse_line(&mut dec, &mut buf, &line).expect("each event is well under the cap"); + let _ = push_sse_line(&mut dec, &mut buf, ""); + } + } + #[test] fn repair_orphaned_tool_use_leaves_a_well_formed_list_untouched() { // pi: tool-call-without-result.test.ts's non-orphaned control case, and the general fast path: diff --git a/crates/agent-core/src/dialect/openai.rs b/crates/agent-core/src/dialect/openai.rs index d5b7926..201bfba 100644 --- a/crates/agent-core/src/dialect/openai.rs +++ b/crates/agent-core/src/dialect/openai.rs @@ -1024,18 +1024,18 @@ impl StreamDecoder for Decoder { .and_then(|d| d.get("cached_tokens")) .and_then(Value::as_u64) .or_else(|| usage.get("prompt_cache_hit_tokens").and_then(Value::as_u64)) - .unwrap_or(0) as u32; + .map_or(0, super::saturating_u32); // OpenAI itself never emits `cache_write_tokens` (no such wire concept), but // OpenRouter-compatible providers can — separate from `cached_tokens` (cache *reads*). let cache_write = usage .get("prompt_tokens_details") .and_then(|d| d.get("cache_write_tokens")) .and_then(Value::as_u64) - .unwrap_or(0) as u32; + .map_or(0, super::saturating_u32); let prompt = usage .get("prompt_tokens") .and_then(Value::as_u64) - .unwrap_or(0) as u32; + .map_or(0, super::saturating_u32); // OpenAI's `prompt_tokens` is the *whole* prompt including both cached (read) and // freshly cache-written tokens; bill only the genuinely uncached remainder as // `input_tokens` so accounting doesn't double-count either kind. Matches pi's @@ -1046,12 +1046,12 @@ impl StreamDecoder for Decoder { self.usage.output_tokens = usage .get("completion_tokens") .and_then(Value::as_u64) - .unwrap_or(0) as u32; + .map_or(0, super::saturating_u32); self.usage.reasoning_tokens = usage .get("completion_tokens_details") .and_then(|d| d.get("reasoning_tokens")) .and_then(Value::as_u64) - .unwrap_or(0) as u32; + .map_or(0, super::saturating_u32); out.push(StreamEvent::Usage(self.usage)); } diff --git a/crates/agent-core/src/dialect/openai_responses.rs b/crates/agent-core/src/dialect/openai_responses.rs index 134e654..59a2c5b 100644 --- a/crates/agent-core/src/dialect/openai_responses.rs +++ b/crates/agent-core/src/dialect/openai_responses.rs @@ -612,14 +612,22 @@ fn str_at<'a>(v: Option<&'a Value>, key: &str) -> &'a str { .unwrap_or("") } -fn i64_at(v: &Value, key: &str) -> i64 { - v.get(key).and_then(Value::as_i64).unwrap_or(-1) +/// Read a wire-provided block index. `None` on a missing, non-numeric, or negative value — never +/// defaulted, and in particular never defaulted to the `-1`-cast-to-`usize` (i.e. `usize::MAX`) this +/// used to produce. Two `output_item.added` events that both lack an `output_index` would otherwise +/// both open at that same synthetic index, and `Accumulator::apply`'s `open.insert(index, ..)` would +/// silently *overwrite* the first with the second: one tool call disappears with no error, and both +/// calls' argument fragments concatenate into a single malformed buffer. Dropping the malformed event +/// instead is the same refuse-to-guess posture `dialect::anthropic`'s own `usize_at` documents. +fn usize_at(v: &Value, key: &str) -> Option { + let n = v.get(key).and_then(Value::as_i64)?; + usize::try_from(n).ok() } fn u32_field(v: Option<&Value>, key: &str) -> u32 { v.and_then(|v| v.get(key)) .and_then(Value::as_u64) - .unwrap_or(0) as u32 + .map_or(0, super::saturating_u32) } /// Build the error message for a `response.failed` event, in the same `code: message` shape @@ -726,7 +734,7 @@ fn message_item_id_phase(item: Option<&Value>) -> (Option, Option, + open_indices: Vec, saw_terminal: bool, failed: Option, saw_tool_call: bool, @@ -750,7 +758,7 @@ impl Default for Decoder { impl Decoder { /// Emit `ev` for `index`, marking it open (a no-op if it already is). - fn emit(&mut self, out: &mut Vec, index: i64, ev: StreamEvent) { + fn emit(&mut self, out: &mut Vec, index: usize, ev: StreamEvent) { if !self.open_indices.contains(&index) { self.open_indices.push(index); } @@ -759,12 +767,10 @@ impl Decoder { /// Close `index`, if it's actually open — a `done` for an index never opened is a stale/duplicate /// event, silently ignored rather than emitting a spurious close for a block that doesn't exist. - fn close(&mut self, out: &mut Vec, index: i64) { + fn close(&mut self, out: &mut Vec, index: usize) { if let Some(pos) = self.open_indices.iter().position(|&i| i == index) { self.open_indices.remove(pos); - out.push(StreamEvent::ContentBlockStop { - index: index as usize, - }); + out.push(StreamEvent::ContentBlockStop { index }); } } @@ -789,9 +795,7 @@ impl Decoder { // event arrives — but a malformed/truncated stream must not silently drop whatever's still // open, at any index. for index in std::mem::take(&mut self.open_indices) { - out.push(StreamEvent::ContentBlockStop { - index: index as usize, - }); + out.push(StreamEvent::ContentBlockStop { index }); } let response = data.get("response"); @@ -849,7 +853,7 @@ impl Decoder { struct FastDeltaEvent { #[serde(rename = "type")] kind: FastDeltaKind, - output_index: i64, + output_index: usize, delta: String, } @@ -873,16 +877,16 @@ impl StreamDecoder for Decoder { let index = parsed.output_index; let ev = match parsed.kind { FastDeltaKind::OutputText | FastDeltaKind::Refusal => StreamEvent::TextDelta { - index: index as usize, + index, text: parsed.delta, }, FastDeltaKind::FunctionCallArguments => StreamEvent::InputJsonDelta { - index: index as usize, + index, partial_json: parsed.delta, }, FastDeltaKind::ReasoningSummaryText | FastDeltaKind::ReasoningText => { StreamEvent::ThinkingDelta { - index: index as usize, + index, text: parsed.delta, } } @@ -910,7 +914,10 @@ impl StreamDecoder for Decoder { let kind = data.get("type").and_then(Value::as_str).unwrap_or(""); match kind { "response.output_item.added" => { - let index = i64_at(data, "output_index"); + let Some(index) = usize_at(data, "output_index") else { + tracing::warn!(event = %kind, "responses event with no usable output_index; dropping"); + return out; + }; let item = data.get("item"); if item.and_then(|i| i.get("type")).and_then(Value::as_str) == Some("function_call") { @@ -926,11 +933,7 @@ impl StreamDecoder for Decoder { self.emit( &mut out, index, - StreamEvent::ToolUseStart { - index: index as usize, - id, - name, - }, + StreamEvent::ToolUseStart { index, id, name }, ); } // A message/reasoning item has no explicit "start" `StreamEvent` — it opens implicitly @@ -938,18 +941,24 @@ impl StreamDecoder for Decoder { // Nothing to emit here for those. } "response.reasoning_summary_text.delta" | "response.reasoning_text.delta" => { - let index = i64_at(data, "output_index"); + let Some(index) = usize_at(data, "output_index") else { + tracing::warn!(event = %kind, "responses event with no usable output_index; dropping"); + return out; + }; self.emit( &mut out, index, StreamEvent::ThinkingDelta { - index: index as usize, + index, text: str_at(Some(data), "delta").to_string(), }, ); } "response.reasoning_summary_part.done" => { - let index = i64_at(data, "output_index"); + let Some(index) = usize_at(data, "output_index") else { + tracing::warn!(event = %kind, "responses event with no usable output_index; dropping"); + return out; + }; // Only meaningful for an index that's already open (accrued some reasoning text) — // otherwise there's nothing to add a paragraph break to. if self.open_indices.contains(&index) { @@ -957,30 +966,36 @@ impl StreamDecoder for Decoder { &mut out, index, StreamEvent::ThinkingDelta { - index: index as usize, + index, text: "\n\n".to_string(), }, ); } } "response.output_text.delta" | "response.refusal.delta" => { - let index = i64_at(data, "output_index"); + let Some(index) = usize_at(data, "output_index") else { + tracing::warn!(event = %kind, "responses event with no usable output_index; dropping"); + return out; + }; self.emit( &mut out, index, StreamEvent::TextDelta { - index: index as usize, + index, text: str_at(Some(data), "delta").to_string(), }, ); } "response.function_call_arguments.delta" => { - let index = i64_at(data, "output_index"); + let Some(index) = usize_at(data, "output_index") else { + tracing::warn!(event = %kind, "responses event with no usable output_index; dropping"); + return out; + }; self.emit( &mut out, index, StreamEvent::InputJsonDelta { - index: index as usize, + index, partial_json: str_at(Some(data), "delta").to_string(), }, ); @@ -990,18 +1005,24 @@ impl StreamDecoder for Decoder { // the streamed deltas produced so far, so a single dropped/duplicated mid-stream delta // (a relay hiccup with no transport-level error — nothing else would ever catch it) // can't silently leave the final call's arguments corrupted. - let index = i64_at(data, "output_index"); + let Some(index) = usize_at(data, "output_index") else { + tracing::warn!(event = %kind, "responses event with no usable output_index; dropping"); + return out; + }; self.emit( &mut out, index, StreamEvent::InputJsonFinal { - index: index as usize, + index, full_json: str_at(Some(data), "arguments").to_string(), }, ); } "response.output_item.done" => { - let index = i64_at(data, "output_index"); + let Some(index) = usize_at(data, "output_index") else { + tracing::warn!(event = %kind, "responses event with no usable output_index; dropping"); + return out; + }; if !self.open_indices.contains(&index) { // A `done` for an index we never saw open — duplicate/stale event, ignored. } else { @@ -1015,7 +1036,7 @@ impl StreamDecoder for Decoder { &mut out, index, StreamEvent::SignatureDelta { - index: index as usize, + index, signature: item.map(ToString::to_string).unwrap_or_default(), }, ); @@ -1028,10 +1049,7 @@ impl StreamDecoder for Decoder { self.emit( &mut out, index, - StreamEvent::ThinkingFinal { - index: index as usize, - text, - }, + StreamEvent::ThinkingFinal { index, text }, ); } } @@ -1047,7 +1065,7 @@ impl StreamDecoder for Decoder { &mut out, index, StreamEvent::InputJsonFinal { - index: index as usize, + index, full_json: args.to_string(), }, ); @@ -1060,7 +1078,7 @@ impl StreamDecoder for Decoder { &mut out, index, StreamEvent::TextFinal { - index: index as usize, + index, text, id, phase, @@ -1143,6 +1161,64 @@ mod tests { use crate::dialect::decode_sse; use crate::message::{ImageSource, Message, ToolDef}; + #[test] + fn two_function_calls_with_no_output_index_are_dropped_not_collapsed_onto_one_index() { + // Regression: `output_index` used to default to `-1`, which every handler then cast with + // `as usize` — i.e. `usize::MAX`. Two `function_call` items lacking the field both opened at + // that one synthetic index, and the accumulator's `open.insert(index, ..)` silently overwrote + // the first call with the second: a tool call vanished with no error at all, and both calls' + // argument fragments ran together into one malformed buffer. Refuse to guess an index instead. + let sse = concat!( + "data: {\"type\":\"response.output_item.added\",\"item\":", + "{\"type\":\"function_call\",\"call_id\":\"c1\",\"id\":\"i1\",\"name\":\"first\"}}\n\n", + "data: {\"type\":\"response.output_item.added\",\"item\":", + "{\"type\":\"function_call\",\"call_id\":\"c2\",\"id\":\"i2\",\"name\":\"second\"}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\"}}\n\n", + ); + let mut dec = Decoder::default(); + let events = decode_sse(&mut dec, sse).expect("a missing index must not error the stream"); + let starts: Vec<_> = events + .iter() + .filter(|e| matches!(e, StreamEvent::ToolUseStart { .. })) + .collect(); + assert!( + starts.is_empty(), + "an event with no usable output_index must be dropped, not opened at a synthetic \ + index where a sibling call can silently overwrite it; got: {starts:?}" + ); + + // A negative index is the same class of malformed input and must be refused identically — + // it is what `-1` was being read as before, just spelled explicitly on the wire. + assert_eq!( + usize_at(&serde_json::json!({"output_index": -1}), "output_index"), + None + ); + assert_eq!(usize_at(&serde_json::json!({}), "output_index"), None); + assert_eq!( + usize_at(&serde_json::json!({"output_index": 3}), "output_index"), + Some(3) + ); + } + + #[test] + fn an_absurd_provider_token_count_clamps_instead_of_wrapping_to_a_small_number() { + // `as u32` is not covered by `overflow-checks`: it wraps silently, so `u32::MAX + 2` used to + // land as `1`. A tiny `input_tokens` is the dangerous direction — it makes a nearly-full + // context look empty and *suppresses* the compaction that would have saved the turn. + let sse = concat!( + "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",", + "\"usage\":{\"input_tokens\":4294967297,\"output_tokens\":5}}}\n\n", + ); + let mut dec = Decoder::default(); + decode_sse(&mut dec, sse).expect("decode"); + assert_eq!( + dec.usage.input_tokens, + u32::MAX, + "an out-of-range token count must clamp to u32::MAX, never wrap to a small number" + ); + assert_eq!(dec.usage.output_tokens, 5); + } + #[test] fn build_body_sends_temperature_when_set() { let req = ModelRequest::new("gpt-4o", vec![Message::user("hi")], 256).with_temperature(0.4); diff --git a/crates/agent-core/src/steering.rs b/crates/agent-core/src/steering.rs index c5328c3..adf2e02 100644 --- a/crates/agent-core/src/steering.rs +++ b/crates/agent-core/src/steering.rs @@ -122,6 +122,52 @@ impl PartialEq for SteeringMessage { type Queue = Arc>>; +/// The most messages any one lane will hold. Every push into these lanes originates, in this +/// codebase's real shape, at a *remote* RPC command (`serve.rs`'s `steer`/`follow_up`/`prompt` +/// handlers, reachable over WebSocket and UDS), and a [`SteeringMessage`] carries `images` — base64 +/// payloads that can each be megabytes. The lanes, meanwhile, drain slowly by design: the default +/// [`QueueMode::OneAtATime`] takes exactly one message per drain point, and while no run is in flight +/// the steer lane has no drain point at all. A client that pushes faster than the agent drains — or +/// pushes at all while idle — therefore pins heap in the daemon for the life of the session, and +/// nothing but [`clear`](Steering::clear) ever gives it back. A cap is what turns that from unbounded +/// into bounded. +/// +/// 64 is chosen to be far above any plausible legitimate depth and still small enough to bound the +/// worst case: these lanes hold *human-authored* redirects, drained roughly one per turn, so a lane +/// even a dozen deep already means the client is producing faster than the agent can ever consume. +/// A client that legitimately needs to hand the agent more than 64 pending instructions wants a +/// prompt, not 64 steers. +const MAX_QUEUED_PER_LANE: usize = 64; + +/// Push onto a lane unless it is already at [`MAX_QUEUED_PER_LANE`], returning whether the message was +/// queued. +/// +/// **Reject-newest, not drop-oldest.** The oldest message in a lane is the one the loop is about to +/// inject — the instruction the client has been waiting on longest. Evicting it to make room for a +/// newer one would silently reorder a client's own instruction stream (its later message lands, its +/// earlier one vanishes) and, under a sustained flood, would keep the head of the queue perpetually +/// churning so that *nothing* ever reaches the model. Refusing the newest instead keeps whatever is +/// queued a faithful, in-order prefix of what the client sent, and puts the loss at the tail — where +/// the client that caused it is still connected, still pushing, and gets `false` back to act on. This +/// is backpressure, not eviction. +/// +/// Never silent: a refusal is logged at `warn`, because a dropped steer is a user's words not reaching +/// the model. +fn push_capped(q: &Queue, lane: &'static str, message: SteeringMessage) -> bool { + let mut queue = lock(q); + if queue.len() >= MAX_QUEUED_PER_LANE { + tracing::warn!( + lane, + depth = queue.len(), + cap = MAX_QUEUED_PER_LANE, + "steering lane full; rejecting message (the queue is draining slower than it is being fed)" + ); + return false; + } + queue.push_back(message); + true +} + /// How much of a lane [`Steering::drain_steer`]/[`Steering::drain_at_stop`] consumes per call. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum QueueMode { @@ -199,14 +245,20 @@ impl Steering { /// Queue a **follow-up**: injected when the run next reaches a stop boundary. Accepts a plain /// string or a [`SteeringMessage`] (with image attachments). - pub fn push(&self, message: impl Into) { - lock(&self.follow_up).push_back(message.into()); + /// + /// Returns `false` — and queues nothing — if the follow-up lane is already at + /// [`MAX_QUEUED_PER_LANE`]; see [`push_capped`] for why the *newest* message is the one refused. + pub fn push(&self, message: impl Into) -> bool { + push_capped(&self.follow_up, "follow_up", message.into()) } /// Queue a **steer**: injected mid-run, at the next tool-results turn, to redirect a busy agent. /// Accepts a plain string or a [`SteeringMessage`] (with image attachments). - pub fn push_steer(&self, message: impl Into) { - lock(&self.steer).push_back(message.into()); + /// + /// Returns `false` — and queues nothing — if the steer lane is already full; see + /// [`push`](Self::push). + pub fn push_steer(&self, message: impl Into) -> bool { + push_capped(&self.steer, "steer", message.into()) } /// Queue a **next-turn** message: folded onto the very first request of whatever run comes next, @@ -215,8 +267,12 @@ impl Steering { /// to), though nothing here enforces that — a caller can queue one mid-run too, and it simply waits /// for the *next* fresh run to start. Accepts a plain string or a [`SteeringMessage`] (with image /// attachments). - pub fn push_next_turn(&self, message: impl Into) { - lock(&self.next_turn).push_back(message.into()); + /// + /// Returns `false` — and queues nothing — if the next-turn lane is already full; see + /// [`push`](Self::push). This lane is the one most likely to be pushed while nothing is draining + /// it (that's its whole purpose), so its cap is the one doing the most work. + pub fn push_next_turn(&self, message: impl Into) -> bool { + push_capped(&self.next_turn, "next_turn", message.into()) } /// Whether any lane currently holds any messages. @@ -338,7 +394,10 @@ impl Steering { lock(&self.next_turn).clear(); self.stop_requested.store(false, Ordering::Relaxed); *lock_opt(&self.model_switch) = None; - *lock_opt(&self.tool_switch) = None; + // Taken out and bound to a `let`, not assigned away as `*guard = None`: the discarded + // registry's destructor must not run while we hold the lock — see + // [`request_tool_set`](Self::request_tool_set). + let _discarded_tools = self.take_tool_switch(); } /// Drop the steer and follow-up lanes, plus any pending stop request, mid-run model switch, and @@ -372,7 +431,8 @@ impl Steering { lock(&self.follow_up).clear(); self.stop_requested.store(false, Ordering::Relaxed); *lock_opt(&self.model_switch) = None; - *lock_opt(&self.tool_switch) = None; + // As in [`clear`](Self::clear): take it out, drop it after the guard is gone. + let _discarded_tools = self.take_tool_switch(); } /// Request that the run stop gracefully at the next turn boundary — after the current turn's tool @@ -432,8 +492,19 @@ impl Steering { /// A second call before the first is observed replaces it outright — same "no queue of switches, /// only the most recent request matters" contract /// [`request_model_switch`](Self::request_model_switch) uses. + /// + /// The registry being *replaced* is deliberately moved out of the slot and dropped by this + /// function's own scope, rather than written over with `*guard = Some(tools)`. An assignment + /// through the guard drops the old value in place, still holding the lock — and a `ToolRegistry` + /// is `Arc`s, so that runs arbitrary *host-supplied* destructors under our mutex. A host + /// tool whose `Drop` blocks, takes another lock, or re-enters `Steering` would stall or deadlock + /// every other lane behind it; one that panics would unwind out of the assignment with the slot + /// dropped-but-not-yet-overwritten and the mutex poisoned, and [`lock_opt`]'s poison recovery + /// hands that slot straight back to the next caller. Dropping after the guard is released costs + /// nothing and makes both moot: the `let` binding outlives the temporary guard, so the old + /// registry's destructor runs on an unlocked `Steering`. pub fn request_tool_set(&self, tools: ToolRegistry) { - *lock_opt(&self.tool_switch) = Some(tools); + let _replaced_tools = lock_opt(&self.tool_switch).replace(tools); } /// Consume and return the pending tool-set switch, if any. Each request is observed at most once. @@ -971,6 +1042,93 @@ mod tests { ); } + #[test] + fn a_full_lane_rejects_the_newest_message_and_keeps_the_queued_prefix_intact() { + // The lanes are fed straight from remote RPC and drained one message per turn boundary, so a + // client that pushes faster than the agent drains must hit a wall rather than grow the queue + // (and its base64 image payloads) without bound. The wall refuses the *newest* message: what + // is already queued stays queued, in order, so the client's instruction stream is truncated at + // the tail rather than silently reordered. + let s = Steering::new(); + for i in 0..MAX_QUEUED_PER_LANE { + assert!( + s.push_steer(format!("m{i}")), + "message {i} is under the cap" + ); + } + assert_eq!(s.pending_count(), MAX_QUEUED_PER_LANE); + + assert!( + !s.push_steer("one too many"), + "the cap must refuse the push" + ); + assert_eq!( + s.pending_count(), + MAX_QUEUED_PER_LANE, + "a refused push must not grow the lane" + ); + let texts = s.steer_texts(); + assert_eq!( + texts.first().map(String::as_str), + Some("m0"), + "the oldest message is untouched" + ); + assert!( + !texts.iter().any(|t| t == "one too many"), + "the refused message must not have displaced anything" + ); + + // Draining frees exactly one slot: the cap is a live watermark, not a one-shot latch. + assert_eq!(s.drain_steer(), vec!["m0".to_string()]); + assert!( + s.push_steer("now there is room"), + "a drain makes room again" + ); + assert!(!s.push_steer("but only one slot's worth")); + } + + #[test] + fn each_lane_is_capped_independently() { + // One saturated lane must not starve the others — they have separate drain points, and a + // stalled steer lane says nothing about whether a follow-up or next-turn message can land. + let s = Steering::new(); + for i in 0..MAX_QUEUED_PER_LANE { + assert!(s.push_steer(format!("s{i}"))); + } + assert!(!s.push_steer("rejected")); + assert!(s.push("the follow-up lane has its own budget")); + assert!(s.push_next_turn("as does the next-turn lane")); + assert_eq!(s.pending_count(), MAX_QUEUED_PER_LANE + 2); + } + + #[test] + fn the_follow_up_and_next_turn_lanes_are_capped_too() { + // The next-turn lane is the one most exposed: it is pushed while the agent is idle, and + // nothing drains it until some later prompt starts a fresh run. + let s = Steering::new(); + for _ in 0..MAX_QUEUED_PER_LANE { + assert!(s.push("f")); + assert!(s.push_next_turn("n")); + } + assert!(!s.push("over"), "the follow-up lane is capped"); + assert!(!s.push_next_turn("over"), "the next-turn lane is capped"); + assert_eq!(s.pending_count(), MAX_QUEUED_PER_LANE * 2); + } + + #[test] + fn clear_releases_a_saturated_lane() { + let s = Steering::new(); + for _ in 0..MAX_QUEUED_PER_LANE { + s.push_steer("x"); + } + assert!(!s.push_steer("full")); + s.clear(); + assert!( + s.push_steer("room again"), + "clear must return the lane's whole budget, not leave it wedged" + ); + } + #[test] fn request_model_switch_leaves_thinking_level_unset() { // The plain two-arg `request_model_switch` (every existing caller) must keep meaning diff --git a/crates/agent-core/src/validation.rs b/crates/agent-core/src/validation.rs index 04071c5..22c777d 100644 --- a/crates/agent-core/src/validation.rs +++ b/crates/agent-core/src/validation.rs @@ -21,15 +21,66 @@ //! are best-effort the way pi's own composition handling is (never fails the surrounding value over a //! member mismatch — see [`apply_composition`]); everywhere else in this module, a sub-schema that //! can't be coerced fails the whole call, matching AJV's own all-or-nothing `coerceTypes`. +//! +//! That recursion runs against a schema this process does not necessarily author — an MCP tool's schema +//! is the remote server's verbatim advertisement — so the whole pass runs on a fixed work budget and +//! degrades to a no-op rather than spinning when it runs out. See [`COERCION_NODE_BUDGET`]. use serde_json::Value; +/// How many schema nodes one `coerce_tool_arguments` call may visit before it stops coercing and hands +/// the rest of the value back as-is. +/// +/// This pass runs on *every* tool dispatch against the tool's own `input_schema()`, and for an MCP tool +/// that schema is whatever the remote server advertised — verbatim, and therefore attacker-controlled if +/// that server is hostile or compromised. Composition is what makes an un-budgeted pass dangerous: +/// `allOf`/`anyOf`/`oneOf` visit *every* member with a full deep clone of the value, and each member +/// re-enters the recursion, so the pass's cost is (schema nodes visited) × (size of the value being +/// cloned at each one) — a hostile schema of a few hundred KB (thousands of composition members, nothing +/// exotic) can therefore buy itself thousands of deep clones of the whole argument value and pin the +/// dispatching task. Recursion *depth* is already safe — serde_json refuses to parse past 128 levels of +/// nesting — so this is a work blowup, not a stack overflow, and a cap on nodes visited is the smallest +/// thing that bounds it. It bounds every other shape of pathological schema at the same time. +/// +/// 10k is orders of magnitude above any real tool schema (the largest here visit a few dozen nodes), so +/// a legitimate call can never reach it, while a hostile one is capped at ~10k clones of a model-emitted +/// argument value — milliseconds, not minutes. +const COERCION_NODE_BUDGET: u32 = 10_000; + +/// Remaining nodes this coercion pass may visit. Exhaustion is not an error: coercion is a best-effort +/// convenience (see the module doc comment), so running out just means the value stops being normalized +/// and is handed back untouched — exactly what a value with no coercible schema at all would get. The +/// tool's own validation still runs afterwards and still produces its own clear error if the arguments +/// really are wrong. +struct Budget(u32); + +impl Budget { + /// Charge one visited node. `false` once the budget is gone — the caller must then stop recursing + /// and return its value as-is. + fn spend(&mut self) -> bool { + self.0 = self.0.saturating_sub(1); + self.0 > 0 + } +} + /// Coerce `input` to match `schema`'s declared type(s), recursing into `object` schemas' `properties`. /// Returns the coerced value on success, or an error string (unused by callers today beyond falling /// back to the original value — see the module doc comment) describing which type(s) it couldn't /// satisfy. +/// +/// Bounded work: at most [`COERCION_NODE_BUDGET`] schema nodes are visited, after which the remaining +/// value is returned un-coerced rather than coerced. pub fn coerce_tool_arguments(schema: &Value, input: Value) -> Result { - coerce_value(schema, input) + coerce_value(schema, input, &mut Budget(COERCION_NODE_BUDGET)) +} + +/// Coerce with an explicit budget, reporting what was left of it — the only way to observe from a test +/// that the pass really did stop early rather than merely finishing fast. +#[cfg(test)] +fn coerce_with_budget(schema: &Value, input: Value, budget: u32) -> (Result, u32) { + let mut budget = Budget(budget); + let result = coerce_value(schema, input, &mut budget); + (result, budget.0) } fn declared_types(schema: &Value) -> Option> { @@ -110,7 +161,11 @@ fn try_coerce(t: &str, input: &Value) -> Option { /// extra ones with no `additionalProperties` schema to coerce against, are left alone — /// presence/`required` enforcement stays each tool's own job (its existing /// `input.get("x").ok_or_else(...)` pattern), not this pass's. -fn coerce_object_properties(schema: &Value, input: Value) -> Result { +fn coerce_object_properties( + schema: &Value, + input: Value, + budget: &mut Budget, +) -> Result { let Value::Object(mut obj) = input else { return Ok(input); }; @@ -121,7 +176,7 @@ fn coerce_object_properties(schema: &Value, input: Value) -> Result Result Result Result { +fn coerce_array_items(schema: &Value, input: Value, budget: &mut Budget) -> Result { let Value::Array(mut items) = input else { return Ok(input); }; match schema.get("items") { Some(Value::Array(item_schemas)) => { for (item, item_schema) in items.iter_mut().zip(item_schemas) { - *item = coerce_value(item_schema, item.take())?; + *item = coerce_value(item_schema, item.take(), budget)?; } } Some(item_schema) if item_schema.is_object() => { for item in items.iter_mut() { - *item = coerce_value(item_schema, item.take())?; + *item = coerce_value(item_schema, item.take(), budget)?; } } _ => {} @@ -171,19 +226,24 @@ fn coerce_array_items(schema: &Value, input: Value) -> Result { /// member that *does* coerce, falling back to the value untouched if none do. Real, non-composition /// type mismatches are still caught downstream by `coerce_value`'s own `type` handling once composition /// hands its result off. -fn apply_composition(schema: &Value, mut value: Value) -> Value { +/// +/// This is the expensive corner of the pass — every `allOf` member is visited with a full deep clone of +/// the value, and each one re-enters the recursion — so it is also where `budget` earns its keep: see +/// [`COERCION_NODE_BUDGET`]. Once the budget is out the members are left unapplied and the value passes +/// through, which is indistinguishable from a composition whose members simply didn't coerce anything. +fn apply_composition(schema: &Value, mut value: Value, budget: &mut Budget) -> Value { if let Some(members) = schema.get("allOf").and_then(Value::as_array) { for member in members { - if let Ok(coerced) = coerce_value(member, value.clone()) { + if let Ok(coerced) = coerce_value(member, value.clone(), budget) { value = coerced; } } } if let Some(members) = schema.get("anyOf").and_then(Value::as_array) { - value = coerce_union(members, value); + value = coerce_union(members, value, budget); } if let Some(members) = schema.get("oneOf").and_then(Value::as_array) { - value = coerce_union(members, value); + value = coerce_union(members, value, budget); } value } @@ -191,9 +251,9 @@ fn apply_composition(schema: &Value, mut value: Value) -> Value { /// Try each union member in turn, keeping the first one that coerces cleanly; if none do, the value is /// handed back untouched rather than treated as a failure — same rationale as `apply_composition`'s own /// doc comment. -fn coerce_union(members: &[Value], value: Value) -> Value { +fn coerce_union(members: &[Value], value: Value, budget: &mut Budget) -> Value { for member in members { - if let Ok(coerced) = coerce_value(member, value.clone()) { + if let Ok(coerced) = coerce_value(member, value.clone(), budget) { return coerced; } } @@ -227,8 +287,15 @@ fn retag_whole_number(value: Value) -> Value { } } -fn coerce_value(schema: &Value, input: Value) -> Result { - let value = apply_composition(schema, input); +fn coerce_value(schema: &Value, input: Value, budget: &mut Budget) -> Result { + // Every re-entry into the recursion — a property, an array element, an `allOf`/`anyOf` member — + // passes through here, so charging the budget at this single point bounds the whole pass. Bailing + // returns the value un-coerced rather than an error: a hostile schema must degrade this pass into a + // no-op, never into a failed dispatch. + if !budget.spend() { + return Ok(input); + } + let value = apply_composition(schema, input, budget); let Some(types) = declared_types(schema) else { // No `type` constraint at all — an object schema with only `properties`/`additionalProperties` @@ -240,10 +307,10 @@ fn coerce_value(schema: &Value, input: Value) -> Result { .get("additionalProperties") .is_some_and(Value::is_object)) { - return coerce_object_properties(schema, value); + return coerce_object_properties(schema, value, budget); } if value.is_array() && schema.get("items").is_some() { - return coerce_array_items(schema, value); + return coerce_array_items(schema, value, budget); } return Ok(value); }; @@ -251,8 +318,8 @@ fn coerce_value(schema: &Value, input: Value) -> Result { for t in &types { if matches_type_raw(t, &value) { return match *t { - "object" => coerce_object_properties(schema, value), - "array" => coerce_array_items(schema, value), + "object" => coerce_object_properties(schema, value, budget), + "array" => coerce_array_items(schema, value, budget), "integer" | "number" => Ok(retag_whole_number(value)), _ => Ok(value), }; @@ -500,6 +567,63 @@ mod tests { ); } + #[test] + fn a_pathological_nested_all_of_schema_stops_at_the_budget_instead_of_blowing_up() { + // The shape that turns this pass into a denial of service: every level is a two-member `allOf` + // whose members nest the level below, so the number of nodes the coercion visits doubles with + // each level — and each of those visits deep-clones the value being coerced. + // + // 15 levels (~65k nodes) rather than the ~25 a real attacker would send: a JSON Schema is a + // *tree*, with no `$ref` sharing here, so the schema document itself doubles right along with + // the visit count and a depth-25 fixture is 30M+ nodes — it OOMs the test process while it's + // still being *built*, before any coercion runs. Depth is capped here only to keep the fixture + // buildable; what's being pinned down is that the coercion stops at its budget no matter how far + // the schema goes, which the exhausted-budget assertion below shows directly (a wall-clock bound + // alone would prove nothing at a depth this small). + let mut schema = json!({"type": "string"}); + let mut input = json!("innermost"); + for _ in 0..15 { + let member = json!({"type": "object", "properties": {"p": schema}}); + schema = json!({"allOf": [member.clone(), member]}); + input = json!({"p": input}); + } + + let started = std::time::Instant::now(); + let (got, remaining) = coerce_with_budget(&schema, input.clone(), COERCION_NODE_BUDGET); + let elapsed = started.elapsed(); + + assert_eq!( + remaining, 0, + "the fixture must actually exhaust the budget, or it isn't testing the bail-out at all" + ); + // Bailing out is a no-op, not an error: the value comes back intact (every coercion this schema + // asks for is an identity one) for the tool's own validation to judge. + assert_eq!( + got, + Ok(input), + "an exhausted budget must hand the value back un-coerced, never fail the dispatch" + ); + // A loose sanity bound, not a benchmark — it must not flake on a loaded CI box. + assert!( + elapsed < std::time::Duration::from_secs(5), + "coercion of a hostile schema must be bounded work, took {elapsed:?}" + ); + } + + #[test] + fn an_exhausted_budget_degrades_to_a_no_op_rather_than_an_error() { + // The bail-out has to be indistinguishable from "this schema had nothing to coerce": a hostile + // schema must be able to turn this pass off, never to turn a legitimate tool call into a failed + // dispatch. A budget of 1 is spent by the first node, so nothing is coerced. + let (got, remaining) = coerce_with_budget(&json!({"type": "integer"}), json!("42"), 1); + assert_eq!(got, Ok(json!("42"))); + assert_eq!(remaining, 0); + + // Same schema, same input, with budget to spare: coercion happens as normal. + let (got, _) = coerce_with_budget(&json!({"type": "integer"}), json!("42"), 2); + assert_eq!(got, Ok(json!(42))); + } + #[test] fn boolean_additional_properties_is_not_treated_as_a_schema() { let schema = json!({ diff --git a/crates/agent-core/src/write_lock.rs b/crates/agent-core/src/write_lock.rs index 886df56..709a272 100644 --- a/crates/agent-core/src/write_lock.rs +++ b/crates/agent-core/src/write_lock.rs @@ -12,16 +12,35 @@ //! serializes within one OS process — two separate `serve` processes against the same file would need //! a filesystem advisory lock (e.g. `flock`), which is out of scope here. //! -//! Entries are evicted once nothing references them any more: when a [`WriteLockGuard`] drops, it -//! checks whether the registry's map is the *only* remaining holder of that key's `Arc` (`strong_count -//! == 1`) — i.e. no other guard and no other in-flight `lock()` call still needs it — and if so removes -//! the map entry. The check-and-remove happens while holding the map's own lock, so a concurrent -//! `lock()` call for the same key can never observe (or create) a state where the entry is removed out -//! from under it: it either wins the race and clones the `Arc` before eviction runs (bumping the count -//! so eviction is skipped), or it runs after eviction and simply inserts a fresh mutex. This mirrors -//! pi's `file-mutation-queue.ts` `withFileMutationQueue`, which deletes its map entry in a `finally` -//! block only if the map still points at that same queue object — the identity check there and the -//! refcount check here both exist to avoid tearing down an entry a concurrent caller still needs. +//! Entries are evicted once nothing references them any more, and *whichever* reference is the last +//! one to go: the check — is the registry's map the only remaining holder of that key's `Arc` +//! (`strong_count == 1`), i.e. no other guard and no other in-flight `lock()` call still needs it? — +//! runs both when a [`WriteLockGuard`] drops and when a `lock()` future is dropped *before* it ever +//! acquires (cancellation). Only covering the guard path would leak: a waiter that clones the `Arc`, +//! parks behind a holder, and is then cancelled mid-`.await` never constructs a guard, so the holder's +//! own drop (which saw `strong_count > 1` because the waiter was still queued) would have been the last +//! chance to evict, and the entry would live in the map for the process's whole lifetime — one leaked +//! `String` + `Arc` per distinct contended-then-cancelled path in a long-lived `serve` process. +//! Making the *acquire* path carry its own drop-guard (rather than sweeping the map, or storing `Weak`s +//! and reaping dead ones later) keeps eviction exact and immediate: the reference that dies last cleans +//! up, with no scan and no window where a dead entry is still visible. +//! +//! The check-and-remove happens while holding the map's own lock, so a concurrent `lock()` call for the +//! same key can never observe (or create) a state where the entry is removed out from under it: it +//! either wins the race and clones the `Arc` before eviction runs (bumping the count so eviction is +//! skipped), or it runs after eviction and simply inserts a fresh mutex. This mirrors pi's +//! `file-mutation-queue.ts` `withFileMutationQueue`, which deletes its map entry in a `finally` block +//! only if the map still points at that same queue object — the identity check there and the refcount +//! check here both exist to avoid tearing down an entry a concurrent caller still needs. +//! +//! **Mutual exclusion, but no fairness.** `futures::lock::Mutex` is *barging*, not FIFO: a caller that +//! polls at the right moment can take a just-released lock ahead of a waiter that queued long before it. +//! So this registry guarantees that two writers to the same path never overlap — nothing more. It makes +//! no ordering or anti-starvation promise across turns or across sessions, and sustained contention on +//! one hot path can in principle starve a waiter indefinitely. That's an accepted trade: the ordering +//! that actually matters — a single turn's own same-target calls applying in the order the model asked +//! for them — is already guaranteed upstream, where `agent.rs`'s `group_runs` runs same-target calls +//! serially in call order within one group, so they never race here in the first place. use std::collections::HashMap; use std::sync::{Arc, Mutex as SyncMutex}; @@ -36,9 +55,12 @@ use futures::lock::{Mutex as AsyncMutex, OwnedMutexGuard}; type LockMap = Arc>>>>; /// A registry of per-key async mutexes. Cheap to construct and share via `Arc`. An entry exists only -/// while at least one caller holds or is waiting on the lock for its key; the last [`WriteLockGuard`] -/// to drop for a given key removes it (see the module doc comment), so the map stays bounded by the -/// number of *currently* contended paths rather than growing for the registry's whole lifetime. +/// while at least one caller holds or is waiting on the lock for its key; the last reference to drop — +/// a [`WriteLockGuard`], or a `lock()` future cancelled before it ever acquired — removes it (see the +/// module doc comment), so the map stays bounded by the number of *currently* contended paths rather +/// than growing for the registry's whole lifetime. +/// +/// Provides mutual exclusion only, not fairness — see the module doc comment. #[derive(Default)] pub struct WriteLockRegistry { locks: LockMap, @@ -51,8 +73,10 @@ impl WriteLockRegistry { /// Acquire the lock for `key`, creating its entry if this is the first call for that key. The /// returned guard is `'static` (owns its `Arc`), so it can be held across `.await` points and moved - /// into a spawned task; dropping it releases the lock and, if no one else is waiting on this key, - /// evicts the key's entry from the registry. + /// into a spawned task; dropping it releases the lock and, if no one else is holding or waiting on + /// this key, evicts the key's entry from the registry. Dropping the *future* returned here before it + /// resolves (cancellation) is equally clean: nothing is locked, and the entry is evicted just the + /// same if this call was the last thing referencing it. pub async fn lock(&self, key: &str) -> WriteLockGuard { // Scoped to a block (rather than a named `let` binding dropped explicitly) so the sync // `MutexGuard` — not `Send` — provably ends its lifetime before the `.await` below, keeping @@ -74,12 +98,19 @@ impl WriteLockRegistry { .clone(), } }; - let guard = mutex.lock_owned().await; - WriteLockGuard { - guard: Some(guard), + // Declared *before* the `lock_owned()` future below, and therefore dropped *after* it when this + // whole future is cancelled mid-`.await` (a suspended future drops its live locals in reverse + // creation order, and the awaited future is itself a later-created temporary). That order is + // what makes the strong-count check below exact: by the time this reservation's `Drop` runs, the + // queued `lock_owned()` future has already released its own clone of the `Arc`, so the count + // reflects only the map, this reservation, and any *genuinely other* holder or waiter. + let reservation = PendingLock { locks: self.locks.clone(), key: key.to_string(), - } + mutex: Some(mutex.clone()), + }; + let guard = mutex.lock_owned().await; + reservation.into_guard(guard) } #[cfg(test)] @@ -91,6 +122,61 @@ impl WriteLockRegistry { } } +/// Remove `key`'s entry if the map is the only thing still referencing its `Arc` — no other guard, and +/// no other `lock()` call still queued on it. Every caller must have released its *own* clone of the +/// `Arc` before calling this, or it will see its own reference and decline to evict. +/// +/// Checked and removed while holding the map's own lock, which every `lock()` call must also acquire to +/// clone the `Arc`, so a concurrent caller can't slip in between the check and the removal. +fn evict_if_unreferenced(locks: &LockMap, key: &str) { + let mut locks = locks + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if locks + .get(key) + .is_some_and(|arc| Arc::strong_count(arc) == 1) + { + locks.remove(key); + } +} + +/// The in-flight half of [`WriteLockRegistry::lock`]: it exists only while a caller is queued on a key's +/// mutex but has not acquired it yet. Its `Drop` is what stops a *cancelled* waiter from orphaning the +/// key's map entry — see the module doc comment. On a successful acquire it is defused by +/// [`PendingLock::into_guard`], which hands eviction duty over to the [`WriteLockGuard`]. +struct PendingLock { + locks: LockMap, + key: String, + /// `Some` while this caller is still merely *waiting*. Taken (without evicting) once the lock is + /// acquired, which is how `Drop` tells cancellation apart from a normal handoff to the guard. + mutex: Option>>, +} + +impl PendingLock { + fn into_guard(mut self, guard: OwnedMutexGuard<()>) -> WriteLockGuard { + // Acquired: the `WriteLockGuard` now owns both the lock and the eviction check, so this + // reservation must not run its own. Dropping the `Arc` here is safe — `guard` holds one. + self.mutex = None; + WriteLockGuard { + guard: Some(guard), + locks: self.locks.clone(), + key: std::mem::take(&mut self.key), + } + } +} + +impl Drop for PendingLock { + fn drop(&mut self) { + // `None` means `into_guard` already handed off; only a caller dropped *before* it acquired gets + // this far. Release our own clone of the `Arc` before the check, or we'd count ourselves. + let Some(mutex) = self.mutex.take() else { + return; + }; + drop(mutex); + evict_if_unreferenced(&self.locks, &self.key); + } +} + /// RAII guard returned by [`WriteLockRegistry::lock`]. Releases the per-key mutex on drop and, if that /// was the last outstanding reference to the key (no other guard or in-flight `lock()` call still holds /// a clone of its `Arc`), removes the key's entry from the registry so it doesn't linger forever. @@ -105,19 +191,7 @@ impl Drop for WriteLockGuard { // Drop the mutex guard first: this both unlocks the per-key mutex and releases this holder's // own strong reference to its `Arc`, so the strong-count check below reflects reality. self.guard.take(); - let mut locks = self - .locks - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(arc) = locks.get(&self.key) { - // Only the map itself still references this key's `Arc` — no other guard or queued - // `lock()` caller holds a clone — so it's safe to evict. Checked while holding `locks`' - // own lock, which every `lock()` call must also acquire to clone the `Arc`, so a - // concurrent caller can't slip in between this check and the removal. - if Arc::strong_count(arc) == 1 { - locks.remove(&self.key); - } - } + evict_if_unreferenced(&self.locks, &self.key); } } @@ -297,4 +371,40 @@ mod tests { "entry must be evicted once every concurrent holder has finished" ); } + + #[tokio::test] + async fn a_waiter_cancelled_before_it_acquires_does_not_orphan_its_map_entry() { + // The leak this pins down needs a very specific interleaving, so it's driven by hand rather + // than with tasks and sleeps: a waiter must still be queued when the holder releases (so the + // holder's own drop sees `strong_count > 1` and declines to evict), and must then die without + // ever acquiring. With eviction living only in `WriteLockGuard::drop`, nothing would be left to + // clean up after it and the entry would sit in the map for the process's whole lifetime. + let registry = WriteLockRegistry::new(); + let holder = registry.lock("cancelled.rs").await; + assert_eq!(registry.key_count(), 1); + + // One poll is enough for the waiter to clone the key's `Arc` and queue on its mutex; it can't + // acquire, because `holder` still owns the lock. + let mut waiter = Box::pin(registry.lock("cancelled.rs")); + assert!( + futures::poll!(waiter.as_mut()).is_pending(), + "sanity check: the waiter must be parked behind the holder, not acquired" + ); + + drop(holder); + assert_eq!( + registry.key_count(), + 1, + "the entry must survive the holder's release while a waiter still references it" + ); + + // The waiter's whole future is dropped mid-`lock_owned()` — cancellation — so it never + // constructs a guard. + drop(waiter); + assert_eq!( + registry.key_count(), + 0, + "a waiter cancelled before acquiring must still evict the entry it was the last to reference" + ); + } } diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index c3b27c7..da74204 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -634,7 +634,17 @@ The harness layers several capabilities over the bare tools + loop: variant — expands a `/skill:name`/`/name` invocation exactly like a fresh `prompt` does (`serve.rs`'s shared `expand_message` helper) and accepts the same optional `images` array a fresh `prompt` does (parsed via `parse_images` into an `agent_core::SteeringMessage`), so neither expansion nor image - attachments are a `prompt`-only behavior a client has to special-case. `compact`(optional `custom_instructions` steers what the summary + attachments are a `prompt`-only behavior a client has to special-case. **Each steering lane is + bounded** (`agent_core::steering`'s per-lane cap): the lanes are fed straight from remote RPC, carry + base64 image payloads, and drain at most one message per turn boundary — and the steer lane has no + drain point at all while idle — so an unbounded lane lets one client pin arbitrary heap in the daemon + for the session's life. A full lane **refuses the newest** message (rather than dropping the oldest, + which would silently reorder the client's own instruction stream and, under a sustained flood, churn + the head so nothing ever reaches the model), and the refusal is reported honestly: the ack comes back + `success: false` with a `steering queue full` error, alongside the usual `queue_content` payload so + the client can see the depth it's up against and retry once the agent drains. Acking `true` for a + message that was dropped is the one outcome a client cannot recover from, since it has no other signal + that its instruction never reached the model. `compact`(optional `custom_instructions` steers what the summary emphasizes, passed straight through to `Agent::compact` — matching pi's own `compact(customInstructions)`; response is `{compacted, reason, summary, tokens_before, tokens_after}` — all four of the latter `null` on a no-op, or a real compaction's own generated summary text and its `AgentEvent::Compacted`-carried diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index 7db3eb2..b3111d1 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -3328,17 +3328,22 @@ pub(crate) async fn serve_session( // `steer` redirects mid-run (injected at the next // tool turn); `follow_up` waits for the stop // boundary. Two separate lanes. - if cmd == "steer" { - steering.push_steer(m); + let queued = if cmd == "steer" { + steering.push_steer(m) } else { - steering.push(m); - } + steering.push(m) + }; // Fix 5 (pi-parity gap): the pushing client // learns what's actually queued from its own // ack, same round trip — see `queue_content`'s // own doc comment for why this doesn't also - // need a separate unsolicited event. - let _ = out_tx.send(response(cid, cmd, true, Some(queue_content(&steering)), None)); + // need a separate unsolicited event. A lane at + // its cap refuses the newest message, so the ack + // has to say so: acking `true` for a message that + // was dropped is the one outcome a client can't + // recover from, since it has no other signal that + // its instruction never reached the model. + let _ = out_tx.send(response(cid, cmd, queued, Some(queue_content(&steering)), (!queued).then_some(STEERING_QUEUE_FULL))); } None => { let _ = out_tx.send(response(cid, cmd, false, None, Some("missing `message`"))); @@ -3362,13 +3367,16 @@ pub(crate) async fn serve_session( m, parse_images(c.get("images")), ); - steering.push_steer(m); + let queued = steering.push_steer(m); // Fix 5 (pi-parity gap): same queue-content // visibility the dedicated `steer`/`follow_up` - // commands' own acks now carry. + // commands' own acks now carry — including a + // full lane's refusal (see the `steer` arm). let mut data = queue_content(&steering); - data["queued_as"] = json!("steer"); - let _ = out_tx.send(response(cid, "prompt", true, Some(data), None)); + if queued { + data["queued_as"] = json!("steer"); + } + let _ = out_tx.send(response(cid, "prompt", queued, Some(data), (!queued).then_some(STEERING_QUEUE_FULL))); } (Some("follow_up"), Some(m)) => { let m = expand_message(m, &skills, &prompt_templates); @@ -3376,10 +3384,12 @@ pub(crate) async fn serve_session( m, parse_images(c.get("images")), ); - steering.push(m); + let queued = steering.push(m); let mut data = queue_content(&steering); - data["queued_as"] = json!("follow_up"); - let _ = out_tx.send(response(cid, "prompt", true, Some(data), None)); + if queued { + data["queued_as"] = json!("follow_up"); + } + let _ = out_tx.send(response(cid, "prompt", queued, Some(data), (!queued).then_some(STEERING_QUEUE_FULL))); } _ => { let _ = out_tx.send(response(cid, "prompt", false, None, Some("busy: a prompt is running; only `abort`/`steer`/`follow_up`, or a `prompt` with `streaming_behavior: \"steer\"|\"follow_up\"`, are accepted"))); @@ -3799,18 +3809,19 @@ pub(crate) async fn serve_session( let m = expand_message(m, &skills, &prompt_templates); let m = agent_core::SteeringMessage::new(m, parse_images(cmd.get("images"))); - if cmd_type == "steer" { - steering.push_steer(m); + let queued = if cmd_type == "steer" { + steering.push_steer(m) } else { - steering.push(m); - } - // Fix 5 (pi-parity gap): see `queue_content`'s own doc comment. + steering.push(m) + }; + // Fix 5 (pi-parity gap): see `queue_content`'s own doc comment. A lane at its + // cap refuses the message; the ack reports that rather than claiming it queued. emit!(response( id, cmd_type, - true, + queued, Some(queue_content(&steering)), - None + (!queued).then_some(STEERING_QUEUE_FULL) )); } None => emit!(response( @@ -6379,6 +6390,16 @@ impl AgentHooks for ServeHooks { /// consulted only by [`ServeHooks::before_provider_request`], the one place this crate's own "what's /// actually sent to the model" transform lives. const HOST_BASH_EXCLUDED_LABEL: &str = "[Host bash command, excluded from model context]"; + +/// Returned on a `steer`/`follow_up` (or a `prompt` routed into one of those lanes) whose queue is at +/// its cap. The lanes are bounded because they're fed straight from remote RPC and drain at most one +/// message per turn boundary — see `agent_core::steering`'s own cap. The client is told so explicitly: +/// a lane that refuses the newest message while its ack still claims `success: true` would leave a +/// client with no signal at all that its instruction never reached the model, which is the one outcome +/// it cannot recover from on its own. The accompanying `queue_content` payload shows what *is* queued, +/// so a client can see the depth it's up against and retry once the agent drains. +const STEERING_QUEUE_FULL: &str = "steering queue full; retry once the agent drains a message"; + /// The ordinary (not excluded) counterpart to [`HOST_BASH_EXCLUDED_LABEL`] — unchanged from before Fix 9. const HOST_BASH_LABEL: &str = "[Host bash command, run outside the model's own turn]"; /// Prefix for the structured `exit_code`/`cancelled`/`truncated`/`full_output_path` status line the From bb2de06e7d02b5cb04d09c5e118f7ba2b400b782 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sun, 12 Jul 2026 17:01:23 -0700 Subject: [PATCH 8/9] fix(agent): resource-leak remediations across the daemon, tools, and export MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A leak hunt across `crates/agent` (daemon, subprocess tools, export). The two that mattered: - **An abandoned `login` leaked an entire session, permanently.** `login_cancel` was a fresh *root* token, the login ran as a detached task holding a clone of the session's `out_tx`, and teardown is `drop(out_tx); writer.await` — where the writer ends only once every sender drops. `pending_login` was never cancelled on the break path (`abort_login` is a *command*, which an abandoned login never sends). So a client that starts a login and disconnects wedged `serve_session` forever: an OS thread + its runtime, the `Agent`, the `GatewayClient` and its pool, the message history, and the OAuth callback server's *fixed* port (53692 / 1455) — which then made every subsequent login fail with `PortBindFailed`. In stdio `serve` it also made graceful shutdown hang outright. The root cause was that `ServeLoginCallbacks::prompt_text` was a bare `rx.await`. CLAUDE.md's human-in-the-loop rule says the wait "races the run's `CancellationToken` and a timeout, and fails closed" — `ApprovalGate` does exactly that; this second instance of the same seam did not. It now does, and teardown cancels a pending login unconditionally. `CallbackServer` also gained a `Drop` that releases the listener without depending on anyone to cancel first. - **`bash` output spilled to `/tmp` with no cap, and `/tmp` here is a 14 GB tmpfs** — i.e. "spill to disk" is spill to RAM. The 50 KiB cap is a display/tail cap; post-spill the writer did `write_all` unconditionally. With a 30-minute default timeout, `bash: yes` fills RAM. The spill file now has a byte budget and the reported truncation metadata stays honest about what was dropped. Also: - The daemon's idle reaper was **opt-in** (`Option`, no default), so every connection minted a session held for the daemon's life. It now defaults to 1h; `0` explicitly disables. Dead handles (a session whose loop already exited) were *excluded* from reaping by an inverted predicate and leaked forever — a closed `input_tx` is now the strongest reason to reap, not a reason to skip. - A wedged session permanently burned a tokio blocking-pool thread (the `timeout` abandoned the await, not the `j.join()`); the join now polls `is_finished()`. - The inbound command channel was unbounded while outbound was carefully capped. It's now bounded with real backpressure — *not* drop-on-full: a dropped command is a correctness bug (the client waits on a response that never comes), so a full queue stops the reader and lets TCP carry the backpressure to the client. - `git apply` wrote the whole patch to stdin before reading stdout/stderr — >64 KiB of git output deadlocked the parent, holding a child and three pipe fds. Write and drain are now concurrent. - `lcs_diff` allocated an uncapped O(n·m) table (a 25k-line rewrite is ~2.5 GB at export time); past a cell budget it falls back to a linear rendering. - `wait_for_pending_group_kills` is now called before every `process::exit` that follows a cancellation, as its own doc comment always asked — three sites were missing it, leaving backgrounded grandchildren alive. Not included: `session_store.rs`. A concurrent session is mid-refactor there (`Node::content` → `Arc`), and the `append_custom` size cap this hunt found belongs on top of that, not underneath it. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/agent/ARCHITECTURE.md | 17 +- crates/agent/src/export.rs | 71 ++++- crates/agent/src/main.rs | 26 +- crates/agent/src/oauth/callback_server.rs | 39 ++- crates/agent/src/serve.rs | 81 +++++- crates/agent/src/serve_ws.rs | 312 +++++++++++++++++++--- crates/agent/src/tools/bash.rs | 42 +++ crates/agent/src/tools/output.rs | 275 +++++++++++++++++-- crates/agent/src/worktree.rs | 102 +++++-- crates/agent/tests/serve_reaper.rs | 43 ++- 10 files changed, 904 insertions(+), 104 deletions(-) diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index da74204..7f14cf4 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -47,12 +47,17 @@ The harness layers several capabilities over the bare tools + loop: reference and the reconnect flow. Two supervisor-level daemon facilities sit above the per-session protocol: **`list_daemon_sessions`** — intercepted in the read loop (never forwarded to a session) and answered by unioning the live `session id → running session` map with an on-disk scan of - `/*.jsonl`, each entry tagged `live: true|false`; and an optional **idle reaper** - (`--session-idle-timeout `, off by default) — a background ticker that reclaims sessions with no - attached connection, idle past the timeout, and not mid-run (a per-session `running` flag `serve_session` - flips around a `prompt`), dropping the retained `input_tx` so the session persists and exits exactly as - graceful shutdown does per-entry (both share `join_handles`). A reconnect to a just-reaped id respawns - it from disk and replays via `get_messages {since}`. A third daemon facility, **shared upstream + `/*.jsonl`, each entry tagged `live: true|false`; and the **idle reaper** + (`--session-idle-timeout `, default 3600, `0` disables) — a background ticker that reclaims + sessions whose task has already ended (a closed `input_tx`: pure garbage, reaped on sight) plus those + with no attached connection, idle past the timeout, and not mid-run (a per-session `running` flag + `serve_session` flips around a `prompt`), dropping the retained `input_tx` so the session persists and + exits exactly as graceful shutdown does per-entry (both share `join_handles`, which polls + `JoinHandle::is_finished` under a grace period rather than parking a blocking join, so a wedged session + can't burn a blocking-pool thread for the daemon's life). The default has to be finite: a connection + that omits `?session_id=` mints a fresh id, and every id owns a thread, an `Agent`, and a gateway pool + until something reclaims it. A reconnect to a just-reaped id respawns it from disk and replays via + `get_messages {since}`. A third daemon facility, **shared upstream pooling** (`--upstream-http2 `, off by default): instead of each session building its own `reqwest::Client` (so N sessions ≈ N connections to the gateway on the plaintext HTTP/1.1 hop), `serve_ws` builds **one** client and injects it into every session via diff --git a/crates/agent/src/export.rs b/crates/agent/src/export.rs index 5c8bed3..31cfd62 100644 --- a/crates/agent/src/export.rs +++ b/crates/agent/src/export.rs @@ -1073,10 +1073,13 @@ fn diff_pair_html(old: &str, new: &str) -> String { /// Compute a real line-level diff between `old` and `new` (pi-parity Fix 1) — pi delegates this to the /// `diff` npm package's `diffLines` (`edit-diff.ts:385`); no equivalent crate is already a dependency of /// this workspace (checked `Cargo.lock` — nothing named `similar`/`diff`/`imara-diff`/`dissimilar` -/// exists, transitively or otherwise), and this is render-time-only, bounded content (a single `edit` -/// call's before/after text, not an arbitrarily large corpus), so a small hand-rolled diff is +/// exists, transitively or otherwise), and this is render-time-only, so a small hand-rolled diff is /// appropriate rather than reaching for a new dependency. /// +/// The content is *not* bounded: `old`/`new` are whatever the model put in an `edit` call, and nothing +/// on the export path truncates them. [`lcs_diff`] therefore caps its own quadratic table +/// ([`MAX_LCS_CELLS`]) and degrades to a plain block rendering past it. +/// /// A common prefix and suffix (lines identical between `old` and `new` at the very start/end) are /// trimmed off *before* running the actual O(n*m) LCS table in [`lcs_diff`] — the overwhelming majority /// of a real edit's old/new pair is unchanged context around one small changed region, so this keeps the @@ -1116,21 +1119,44 @@ fn diff_lines<'a>(old: &'a str, new: &'a str) -> Vec> { ops } +/// The largest LCS table [`lcs_diff`] will build, in cells. The quadratic table is the one place an +/// export can be made to allocate without limit: prefix/suffix trimming shrinks a *typical* edit to a +/// handful of changed lines, but an edit that rewrites a file wholesale leaves every line differing on +/// both sides, and `dp` is then `len(old) * len(new)` cells — a 25k-line rewrite is ~2.5 GB, a +/// single-shot OOM triggered by nothing more than exporting the session that contains it. +/// +/// 1M cells is 4 MiB of `u32` and admits any changed region up to ~1000x1000 lines, which is already far +/// past what a human reads as a diff; past it the rendering degrades but the content is all still there. +const MAX_LCS_CELLS: usize = 1_000_000; + /// A classic LCS (longest-common-subsequence) line diff over `old`/`new` — the textbook /// dynamic-programming table (`dp[i][j]` = length of the LCS of `old[i..]`/`new[j..]`), backtracked from /// `(0, 0)` to produce an edit script. `O(len(old) * len(new))` time and space, which is why /// [`diff_lines`] trims the common prefix/suffix first rather than calling this directly on the whole -/// old/new pair. +/// old/new pair, and why the table is capped at [`MAX_LCS_CELLS`]. fn lcs_diff<'a>(old: &[&'a str], new: &[&'a str]) -> Vec> { let n = old.len(); let m = new.len(); - let mut dp = vec![vec![0u32; m + 1]; n + 1]; + // Too big to diff: fall back to what this renderer did before the line-level diff existed — the whole + // old block `-`, the whole new block `+`. Linear in the input, and every line the model wrote still + // appears in the export; all that is lost is the pairing of unchanged lines across a pathologically + // large rewrite, which nobody was going to read line-by-line anyway. + if n.saturating_mul(m) > MAX_LCS_CELLS { + let mut ops = Vec::with_capacity(n + m); + ops.extend(old.iter().copied().map(LineDiffOp::Delete)); + ops.extend(new.iter().copied().map(LineDiffOp::Insert)); + return ops; + } + // One flat allocation with an explicit stride, rather than `vec![vec![]; n + 1]`: same cells, but + // `n + 1` fewer allocations and a contiguous, cache-friendly sweep. + let stride = m + 1; + let mut dp = vec![0u32; stride * (n + 1)]; for i in (0..n).rev() { for j in (0..m).rev() { - dp[i][j] = if old[i] == new[j] { - dp[i + 1][j + 1] + 1 + dp[i * stride + j] = if old[i] == new[j] { + dp[(i + 1) * stride + j + 1] + 1 } else { - dp[i + 1][j].max(dp[i][j + 1]) + dp[(i + 1) * stride + j].max(dp[i * stride + j + 1]) }; } } @@ -1142,7 +1168,7 @@ fn lcs_diff<'a>(old: &[&'a str], new: &[&'a str]) -> Vec> { ops.push(LineDiffOp::Equal(old[i])); i += 1; j += 1; - } else if dp[i + 1][j] >= dp[i][j + 1] { + } else if dp[(i + 1) * stride + j] >= dp[i * stride + j + 1] { ops.push(LineDiffOp::Delete(old[i])); i += 1; } else { @@ -3416,6 +3442,35 @@ mod tests { assert!(html.contains("line five"), "{html}"); } + #[test] + fn a_wholesale_rewrite_past_the_lcs_cap_still_exports_and_still_shows_both_sides() { + // A file rewritten wholesale: every line differs, so prefix/suffix trimming saves nothing and the + // LCS table would be 5k x 5k = 25M cells (100 MB) — and a 25k-line rewrite would be ~2.5 GB. + // Past `MAX_LCS_CELLS` the diff degrades to plain `-` old / `+` new blocks: linear, and the + // content is all still in the document. The timing assert is what catches a regression to the + // quadratic path, which for this input takes seconds and 100 MB rather than milliseconds. + let old: String = (0..5_000).map(|i| format!("old line {i}\n")).collect(); + let new: String = (0..5_000).map(|i| format!("new line {i}\n")).collect(); + let messages = vec![Message::assistant(vec![ContentBlock::ToolUse { + id: "1".into(), + name: "edit".into(), + input: serde_json::json!({ + "path": "src/lib.rs", + "old_string": old, + "new_string": new, + }), + thought_signature: None, + }])]; + let started = std::time::Instant::now(); + let html = render_html(&meta(), &messages, &[], None); + assert!( + started.elapsed() < std::time::Duration::from_secs(5), + "export of an oversized edit fell back into the quadratic diff" + ); + assert!(html.contains("class=\"diff-del\">-old line 4999"), "{html}"); + assert!(html.contains("class=\"diff-add\">+new line 4999"), "{html}"); + } + #[test] fn edit_call_with_no_overlap_still_diffs_as_a_full_delete_then_insert() { // The degenerate case (nothing in common) must still behave like the old bulk-colored diff --git a/crates/agent/src/main.rs b/crates/agent/src/main.rs index 9dd8bbe..dbc737a 100644 --- a/crates/agent/src/main.rs +++ b/crates/agent/src/main.rs @@ -686,8 +686,11 @@ enum Command { listen_uds_mode: Option, /// Daemon mode only: reap a session that has had no attached connection for this many seconds /// and isn't mid-run — dropping it so it persists and exits, exactly like a graceful shutdown - /// does per-session. Absent (the default) keeps today's behavior: a detached session lives until - /// the daemon stops. Ignored without `--listen`/`--listen-uds`. + /// does per-session. Nothing is lost: reconnecting to a reaped id respawns it and replays from + /// disk. Absent ⇒ 3600 (one hour) — long enough that a client can drop its socket and re-attach + /// to a still-running session, finite so an unattended daemon's threads and gateway pools don't + /// accumulate forever. Pass `0` to disable reaping entirely (every session then lives until the + /// daemon stops). Ignored without `--listen`/`--listen-uds`. #[arg(long, env = "AI_AGENT_SESSION_IDLE_TIMEOUT")] session_idle_timeout: Option, /// How the daemon pools its upstream (agent→gateway) connections across sessions: `off` (the @@ -2020,6 +2023,15 @@ async fn main() -> Result<(), Box> { // shutdown from a clean stdin-EOF one — previously every graceful path exited 0 // unconditionally, matching neither pi's own `rpc-mode.ts` (143/129 for SIGTERM/SIGHUP) // nor a shell's own convention for reporting which signal actually stopped a process. + // + // A signal-triggered shutdown gets here by cancelling every live run, which drops any + // in-flight bash tool future and with it its `GroupKillGuard` — whose reaping happens on a + // detached thread that is *not* among the session threads `serve`/`serve_ws` joined on the + // way out, and that `process::exit` would otherwise terminate mid-`kill`, orphaning exactly + // the backgrounded grandchildren the guard exists to reap. Bounded, so a wedged `kill`/`ps` + // shell-out can't hold the daemon's own shutdown open. + #[cfg(unix)] + tools::exec::wait_for_pending_group_kills(std::time::Duration::from_secs(2)); std::process::exit(shutdown_cause.map(serve::Signal::exit_code).unwrap_or(0)); } Command::Tools => { @@ -4163,6 +4175,12 @@ async fn run_task( // only ever captured here. persist_run_tail(&store, &session)?; if broken_pipe.load(Ordering::Relaxed) { + // Reached *because* `write_stdout_or_exit` tripped `cancel`, so the same in-flight bash tool + // future a signal would have dropped has just been dropped here too — with the same detached + // `GroupKillGuard` cleanup thread still running, and the same `process::exit` about to kill it + // mid-`kill`. See `unwrap_turn_result`, which pays this toll for the signal path. + #[cfg(unix)] + tools::exec::wait_for_pending_group_kills(std::time::Duration::from_secs(2)); std::process::exit(0); } let mut stop_reason = unwrap_turn_result(turn_result, &shutdown_cause)?; @@ -4182,6 +4200,10 @@ async fn run_task( .await; persist_run_tail(&store, &session)?; if broken_pipe.load(Ordering::Relaxed) { + // Same broken-pipe cancellation as the first turn's exit above — drain the pending + // process-group kills before tearing their threads down. + #[cfg(unix)] + tools::exec::wait_for_pending_group_kills(std::time::Duration::from_secs(2)); std::process::exit(0); } stop_reason = unwrap_turn_result(turn_result, &shutdown_cause)?; diff --git a/crates/agent/src/oauth/callback_server.rs b/crates/agent/src/oauth/callback_server.rs index a9090ed..de44d15 100644 --- a/crates/agent/src/oauth/callback_server.rs +++ b/crates/agent/src/oauth/callback_server.rs @@ -60,13 +60,17 @@ impl CallbackServer { // Nudges the blocking accept loop awake the instant cancellation fires, rather than making it // wait out its own poll tick — a latency nicety, not a correctness requirement (the loop below // also checks `is_cancelled()` on every tick regardless). - let unblocker = { + // Aborted on *every* exit from this function, including an early drop of the whole future — not + // just the normal return. It holds its own `Arc`, so a task left parked on + // `cancel.cancelled()` for a token that never fires would keep the listener (and its bound, + // fixed port) alive no matter what the accept loop below does. + let unblocker = AbortOnDrop(Some({ let server = server.clone(); tokio::spawn(async move { cancel.cancelled().await; server.unblock(); }) - }; + })); let result = tokio::task::spawn_blocking(move || { loop { @@ -95,11 +99,40 @@ impl CallbackServer { .ok() .flatten(); - unblocker.abort(); + drop(unblocker); result } } +/// Aborts a task when it goes out of scope, however that happens — a normal return, an early `?`, or +/// the enclosing future being dropped mid-`.await`. A bare `handle.abort()` at the end of a function +/// only covers the first of those. +struct AbortOnDrop(Option>); + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + if let Some(h) = self.0.take() { + h.abort(); + } + } +} + +impl Drop for CallbackServer { + /// Releases the listener — and with it the bound port — without relying on anyone to have cancelled + /// the token first. + /// + /// The accept loop runs on `spawn_blocking`, which cannot be aborted: it exits only by observing a + /// cancelled token or by `recv_timeout` failing. So if a caller merely *drops* the login future + /// rather than cancelling it, that loop keeps polling forever, holding its own `Arc` and the + /// port with it. `unblock()` makes the pending `recv_timeout` fail, the loop returns, and its `Arc` + /// goes with it. This matters more than it looks: the port is *fixed* per provider (53692 for + /// Anthropic, 1455 for Codex), not ephemeral, so a single leaked listener makes every subsequent + /// login on that provider fail outright with `PortBindFailed`. + fn drop(&mut self) { + self.server.unblock(); + } +} + /// Parse one request; `Some(code)` only when the path/state/code all check out. Deliberately /// returns `Option`, not `Result` — see [`CallbackServer::wait_for_callback`]'s doc comment on why a /// bad request must not fail the wait. diff --git a/crates/agent/src/serve.rs b/crates/agent/src/serve.rs index b3111d1..971bc60 100644 --- a/crates/agent/src/serve.rs +++ b/crates/agent/src/serve.rs @@ -1850,6 +1850,23 @@ fn resolve_startup_model_and_level( ) } +/// Depth of a session's inbound command queue, shared by both transports. +/// +/// The counterpart to [`crate::serve_ws::OUT_CHANNEL_BOUND`] on the way *in*. The consumer is a single +/// session loop that runs one command to completion — and some commands are genuinely slow (a +/// `list_sessions` walks the session directory) — so a client that pipelines faster than the loop drains +/// will always be able to get ahead of it. Left unbounded, that queue is a client-driven allocation with +/// no ceiling. +/// +/// Bounded, but explicitly **not** drop-on-full the way the outbound fanout and the `Pong` queue are: +/// those carry frames whose loss is a cosmetic degradation, whereas a *dropped command* is a +/// correctness bug — the client is left waiting on a response that will never come, with no signal that +/// anything went wrong. So a full queue makes the producer wait instead. On the WebSocket path that +/// means the read loop stops pulling from the socket and TCP's own window carries the backpressure to +/// the client; on stdio it simply pauses the stdin reader. Deep enough that ordinary pipelining never +/// touches it, shallow enough that a flood can't grow memory without limit. +pub(crate) const IN_CHANNEL_BOUND: usize = 256; + /// Run the control loop until stdin closes. /// Headless `serve` over **stdio** — the default transport. A thin wrapper over the transport- /// agnostic [`serve_session`]: pump stdin lines into its input channel (EOF drops the sender, which @@ -1858,13 +1875,17 @@ fn resolve_startup_model_and_level( /// `serve` behavior and test rides this path unchanged; [`crate::serve_ws`] is the parallel WebSocket /// entry point that reuses the same [`serve_session`] core. pub async fn serve(cfg: ServeConfig) -> Result, Box> { - let (input_tx, input_rx) = mpsc::unbounded_channel::(); + let (input_tx, input_rx) = mpsc::channel::(IN_CHANNEL_BOUND); tokio::spawn(async move { let mut lines = BufReader::new(tokio::io::stdin()).lines(); // A malformed-UTF-8 read error is treated as a clean EOF, matching the pre-refactor idle/busy // read loops (killing a long-running process over one bad stdin byte is far more disruptive). while let Ok(Some(line)) = lines.next_line().await { - if input_tx.send(line).is_err() { + // Awaits rather than dropping when the queue is full: a command the session never sees is a + // correctness bug (the client waits on a response that never comes), so the right answer to + // a producer outrunning the loop is to stop reading it, not to shed it. Here that simply + // pauses the stdin reader. + if input_tx.send(line).await.is_err() { break; // the session ended; nothing left to feed } } @@ -1909,7 +1930,7 @@ pub async fn serve(cfg: ServeConfig) -> Result, Box, + mut input_rx: mpsc::Receiver, out_conn: SharedOutConn, // Set `true` for the duration of a `prompt` run, `false` otherwise. The daemon supervisor's idle // reaper reads this to never reap a session with an in-flight background run (see @@ -5813,6 +5834,7 @@ pub(crate) async fn serve_session( id: login_id.clone(), provider: provider_id, pending_code, + cancel: login_cancel.clone(), }; let result = crate::oauth::login(provider_id, &callbacks, &login_cancel).await; @@ -5981,6 +6003,19 @@ pub(crate) async fn serve_session( } } + // A `login` still in flight owns a detached task holding its own clone of `out_tx` — and the writer + // below ends *only* once every sender is gone. Without this cancel, a client that starts a login and + // then simply disconnects (or a SIGTERM arriving mid-flow) parks that task forever, `writer.await` + // never returns, and `serve_session` never completes: the OS thread, its runtime, the `Agent`, the + // `GatewayClient` and its pool, the session history, and the OAuth callback server's bound — and + // *fixed* — port all leak for the daemon's lifetime, which also makes every subsequent login on that + // port fail outright. In stdio `serve` the same hang means the process can never exit gracefully. + // + // `abort_login` is not a substitute: it only fires when a client explicitly sends it, which is + // exactly what an abandoned login never does. Teardown is the path that has to be unconditional. + if let Some(p) = lock_ignoring_poison(&pending_login).take() { + p.cancel.cancel(); + } drop(out_tx); let _ = writer.await; Ok(shutdown_cause) @@ -7539,8 +7574,26 @@ struct ServeLoginCallbacks { id: Option, provider: crate::oauth::OAuthProviderId, pending_code: PendingCodeSlot, + /// This login's own token, so the wait below can lose a race against it. Not merely decorative: + /// two provider flows (`oauth::github_copilot`'s Enterprise-domain prompt, `oauth::openai_codex`'s + /// port-already-bound fallback) `await` `prompt_text` with *no other* cancellation observer in the + /// select, so if this wait doesn't watch the token itself, nothing does — `abort_login` becomes a + /// silent no-op and the task parks forever. + cancel: agent_core::CancellationToken, } +/// How long an unanswered `manual_code` prompt waits before the login is abandoned. +/// +/// The wait is on a human: read a browser page, copy a code, paste it back. Generous enough that a +/// distracted user doesn't lose a legitimate flow, finite because the alternative is a task parked for +/// the process's lifetime — and a parked login task is not cheap. It holds a clone of the session's +/// `out_tx`, and the session's writer ends *only* when every sender drops, so a login nobody ever +/// answers wedges `serve_session` itself: the OS thread, its runtime, the `Agent`, the `GatewayClient` +/// and its pool, the whole message history, and the OAuth callback server's bound (and fixed) port all +/// stay live until the daemon dies. Failing closed on silence is what keeps an abandoned browser tab +/// from costing a session. +const LOGIN_PROMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600); + #[async_trait::async_trait] impl crate::oauth::LoginCallbacks for ServeLoginCallbacks { async fn show_auth_url(&self, url: &str, _instructions: Option<&str>) { @@ -7598,8 +7651,26 @@ impl crate::oauth::LoginCallbacks for ServeLoginCallbacks { None, None, )); - rx.await - .map_err(|_| crate::oauth::OAuthError::LoginCancelled) + // The same shape `ServeApprovalGate::request` uses, and for the same reason: this is the second + // instance of the one human-in-the-loop seam, so it races the token and a timeout and **fails + // closed** — no answer, for any reason, abandons the login. A bare `rx.await` here (what this + // used to be) has no way out at all: the sender lives inside these callbacks, which the login + // future owns, so nothing external can even drop it to wake the wait. + // + // `biased` so a cancellation already signalled wins over a code that happened to land in the + // same poll — an aborted login must not quietly succeed on the strength of a late paste. + tokio::select! { + biased; + _ = self.cancel.cancelled() => Err(crate::oauth::OAuthError::LoginCancelled), + answer = rx => answer.map_err(|_| crate::oauth::OAuthError::LoginCancelled), + _ = tokio::time::sleep(LOGIN_PROMPT_TIMEOUT) => { + tracing::debug!( + provider = %self.provider, + "login code prompt went unanswered; abandoning the login" + ); + Err(crate::oauth::OAuthError::LoginCancelled) + } + } } async fn select( diff --git a/crates/agent/src/serve_ws.rs b/crates/agent/src/serve_ws.rs index 05c8fe3..845ed94 100644 --- a/crates/agent/src/serve_ws.rs +++ b/crates/agent/src/serve_ws.rs @@ -49,6 +49,7 @@ use tokio::time::{Duration, MissedTickBehavior}; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::handshake::server::{ErrorResponse, Request, Response}; +use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; use tokio_util::sync::CancellationToken; use crate::serve::{ @@ -71,7 +72,40 @@ const PING_INTERVAL: Duration = Duration::from_secs(30); /// (`OutFanout::broadcast` prunes the sink when it fills, dropping the dead connection). Sized to /// absorb a normal burst of streamed event frames while keeping the worst-case per-connection buffer /// to a few MB. -const OUT_CHANNEL_BOUND: usize = 1024; +pub(crate) const OUT_CHANNEL_BOUND: usize = 1024; + +/// Depth of a single connection's control-frame queue (the `Pong`s the read loop owes the client for +/// its `Ping`s). Its drain is the send task, which stalls for as long as the socket's write half is +/// backed up — so a client that floods `Ping`s while never reading would grow an unbounded queue of +/// them. Bounded, and **dropped on full**: a lost `Pong` costs nothing (the client's next `Ping` +/// re-asks, and liveness in the other direction is our own keepalive `Ping`, whose send failure is what +/// actually detects a dead socket). +const CTRL_CHANNEL_BOUND: usize = 8; + +/// Cap on a single inbound WebSocket message (tungstenite's own defaults are 64 MiB per message / 16 MiB +/// per frame — far past anything this protocol needs). Every inbound message is one command line that +/// gets queued for the session, so this is what bounds the cost of any one queued command; a client that +/// exceeds it gets a protocol error and its connection closed, never a silently truncated command. +/// Generous enough for a `prompt` carrying a large pasted body. +const MAX_INBOUND_MESSAGE_BYTES: usize = 8 * 1024 * 1024; + +/// How long a **detached** session (no attached connection) stays live before the reaper reclaims it, +/// when the operator gave no `--session-idle-timeout`. The default has to be finite: without a reaper +/// the session map only grows — a connection that omits `?session_id=` mints a fresh id, and every id +/// owns an OS thread, an `Agent`, and a gateway pool until the daemon stops. It also has to be *long*, +/// because a detached session is not a dead one: re-attaching to a still-running run is the entire point +/// of the design (see the module doc), so the window must comfortably outlast a tunnel, a locked screen, +/// or a lunch break. An hour is both. Nothing is lost when it fires — a reaped session persisted on its +/// way out, and reconnecting to its id respawns it and replays from disk. +const DEFAULT_SESSION_IDLE_TIMEOUT: Duration = Duration::from_secs(60 * 60); + +/// How long a batch of session threads gets to persist and exit before the caller stops waiting on them. +const JOIN_GRACE: Duration = Duration::from_secs(10); + +/// How often [`join_handles_within`] re-checks whether the session threads it's waiting on have exited. +/// Short enough not to add perceptible latency to a graceful shutdown (the common case: every session +/// persists in milliseconds), long enough to cost nothing while waiting. +const JOIN_POLL: Duration = Duration::from_millis(10); /// A live session, reachable by id across connections. The retained `input_tx` is what makes a /// dropped socket *not* an EOF — the session's `input_rx.recv()` pends until the next command instead @@ -80,7 +114,7 @@ struct SessionHandle { /// Feeds command lines into the session's [`serve_session`] loop. Held here (not by any socket) so /// the session outlives its connections. Every attached connection feeds this one channel, so any /// device can drive the session (its `mpsc` is multi-sender). - input_tx: mpsc::UnboundedSender, + input_tx: mpsc::Sender, /// The session's set of attached connections — its output is **broadcast** to all of them, so a /// phone and a TUI (or any N of the user's devices) on one session all see the live stream at once /// (see [`crate::serve::OutFanout`]). Each connection registers its sink on attach and removes it on @@ -160,7 +194,7 @@ impl Supervisor { } let handle = sessions.entry(id.clone()).or_insert_with(|| { - let (input_tx, input_rx) = mpsc::unbounded_channel::(); + let (input_tx, input_rx) = mpsc::channel::(crate::serve::IN_CHANNEL_BOUND); let out_conn: SharedOutConn = Arc::new(Mutex::new(OutFanout::default())); let cfg = self.session_cfg(&id); let session_out = out_conn.clone(); @@ -223,8 +257,9 @@ impl Supervisor { // The send task owns the socket's write half. Everything outbound — protocol frames, `Pong` // replies, and periodic keepalive `Ping`s — funnels through it so the single sink has one - // writer. `ctrl_rx` carries the read loop's `Pong` replies. - let (ctrl_tx, mut ctrl_rx) = mpsc::unbounded_channel::(); + // writer. `ctrl_rx` carries the read loop's `Pong` replies, bounded like every other queue a + // client can push into (see [`CTRL_CHANNEL_BOUND`]). + let (ctrl_tx, mut ctrl_rx) = mpsc::channel::(CTRL_CHANNEL_BOUND); // Stops this connection's send task at teardown (when its read loop ends on socket close). The // send task also stops on its own if `conn_rx` closes (this connection's sink removed + `reply_tx` // dropped). @@ -298,12 +333,23 @@ impl Supervisor { } } } - if input_tx.send(text.as_str().to_owned()).is_err() { + // Never dropped, never coalesced: a client that got no error believes its + // command was accepted, so discarding one leaves it waiting forever on a + // response that will never come — a dropped command is a correctness bug, + // not a shed load. So this *awaits* a full queue rather than shedding: the + // read loop stops pulling from the socket, TCP's own window closes, and the + // backpressure lands where it belongs — on the client that is outrunning the + // session loop. Any one command is separately bounded by + // [`MAX_INBOUND_MESSAGE_BYTES`], so a full queue is bounded in bytes too. + if input_tx.send(text.as_str().to_owned()).await.is_err() { break; // session gone } } Message::Ping(payload) => { - let _ = ctrl_tx.send(Message::Pong(payload)); + // Drop-on-full (see [`CTRL_CHANNEL_BOUND`]): a `Pong` is a courtesy, and a + // `Ping` flood from a client that never drains its read side must not be + // able to queue work faster than the send task can write it out. + let _ = ctrl_tx.try_send(Message::Pong(payload)); } Message::Close(_) => break, // Pong (keepalive ack), Binary (protocol is text JSON), and any future frame @@ -403,23 +449,16 @@ impl Supervisor { })) } - /// Idle reaper (only spawned when [`ServeConfig::session_idle_timeout`] is set): reclaim every - /// session that has been **detached** (`attached == 0`) for at least `timeout`, whose input channel - /// is still open, and that isn't mid-`prompt` (`running`). Removing the handle drops its retained + /// Idle reaper: reclaim every session [`is_reapable`] names. Removing the handle drops its retained /// `input_tx`, so the session observes EOF and persists+exits exactly as in [`shutdown`]; the join - /// then waits for that persist to land. A reconnect to a just-reaped id transparently respawns and - /// replays via `get_messages{since}`. + /// then waits for that persist to land (and reclaims the thread of one that had already exited). A + /// reconnect to a just-reaped id transparently respawns and replays via `get_messages{since}`. async fn reap_idle(&self, timeout: Duration) { let joins: Vec> = { let mut sessions = lock_ignoring_poison(&self.sessions); let reap: Vec = sessions .iter() - .filter(|(_, h)| { - h.attached == 0 - && h.last_detached_at.is_some_and(|d| d.elapsed() >= timeout) - && !h.input_tx.is_closed() - && !h.running.load(Ordering::Relaxed) - }) + .filter(|(_, h)| is_reapable(h, timeout)) .map(|(id, _)| id.clone()) .collect(); reap.into_iter() @@ -430,25 +469,86 @@ impl Supervisor { } } -/// Wait for a batch of session threads to finish, bounded so a wedged session can't hang the caller -/// forever. `JoinHandle::join` blocks, so it runs on the blocking pool with a total 10s cap; a -/// straggler past the cap is left behind (graceful shutdown then falls back to `process::exit`, the -/// reaper simply retries the id on its next tick). Shared by [`Supervisor::shutdown`] and +/// Whether the reaper should reclaim this session. +/// +/// A **closed `input_tx`** is the strongest reason of all: the session's loop is already gone (it +/// returned early — no credential, an unwritable session dir — or hit an internal error), so nothing +/// will ever read its input again. The entry is pure garbage: a `HashMap` slot, a `SharedOutConn`, and +/// an unreclaimed `JoinHandle` for a thread that has already exited. Reap it whatever its attach state +/// or clock says — a live connection still pinned to it is no reason to keep a corpse (that connection's +/// next command fails its `input_tx.send` and tears the socket down; a reconnect respawns the id). +/// +/// Otherwise the session is alive, and the ordinary conditions apply: **detached** (`attached == 0`) for +/// at least `timeout`, and not mid-`prompt` — a detached background run is exactly what this design +/// exists to keep alive, so `running` is never reaped out from under an in-flight turn. +fn is_reapable(h: &SessionHandle, timeout: Duration) -> bool { + if h.input_tx.is_closed() { + return true; + } + h.attached == 0 + && h.last_detached_at.is_some_and(|d| d.elapsed() >= timeout) + && !h.running.load(Ordering::Relaxed) +} + +/// Wait for a batch of session threads to finish persisting, bounded by [`JOIN_GRACE`] so a wedged +/// session can't hang the caller forever. Shared by [`Supervisor::shutdown`] and /// [`Supervisor::reap_idle`] so both persist-then-join on the same discipline. async fn join_handles(joins: Vec>) { + join_handles_within(joins, JOIN_GRACE).await; +} + +/// The waiting itself, with the grace period as a parameter so it can be tested. +/// +/// The wait **polls `is_finished`** rather than parking a blocking join behind a timeout, because a +/// timeout only abandons the *await* — the blocking join underneath it keeps running. Wedge one session +/// and a `spawn_blocking(|| handle.join())` would sit on one of tokio's blocking threads (a finite pool) +/// for the rest of the process's life; wedge enough of them and every `spawn_blocking` in the daemon +/// stalls behind an empty pool, session persistence included. `join()` is therefore only ever called on +/// a thread that has *already exited*, where it cannot block: it just reaps the thread. +/// +/// A straggler past `grace` is dropped, which detaches it — the OS reclaims its stack when it finally +/// does exit, and nothing (no caller, no pool thread) is left waiting on it. That is all the caller can +/// do: a graceful shutdown falls through to `process::exit` regardless, and the reaper has already +/// removed the id from the map, so it will never see that session again. +async fn join_handles_within(joins: Vec>, grace: Duration) { if joins.is_empty() { return; } - let waiter = tokio::task::spawn_blocking(move || { - for j in joins { - let _ = j.join(); + let deadline = Instant::now() + grace; + let mut pending = joins; + loop { + let mut still = Vec::with_capacity(pending.len()); + for j in pending { + if j.is_finished() { + let _ = j.join(); + } else { + still.push(j); + } } - }); - if tokio::time::timeout(Duration::from_secs(10), waiter) - .await - .is_err() - { - eprintln!("serve: some sessions did not persist within the join grace period"); + pending = still; + if pending.is_empty() { + return; + } + if Instant::now() >= deadline { + break; + } + tokio::time::sleep(JOIN_POLL).await; + } + eprintln!( + "serve: {} session(s) did not persist within the join grace period", + pending.len() + ); +} + +/// Resolve the idle reaper's window from [`ServeConfig::session_idle_timeout`]: unset ⇒ +/// [`DEFAULT_SESSION_IDLE_TIMEOUT`] (the reaper is *on* by default — the map has no other way to shrink, +/// see the const), and `0` ⇒ `None`, the explicit opt-out for an operator who genuinely wants every +/// session pinned for the daemon's lifetime. +fn resolve_idle_timeout(configured: Option) -> Option { + match configured { + Some(t) if t.is_zero() => None, + Some(t) => Some(t), + None => Some(DEFAULT_SESSION_IDLE_TIMEOUT), } } @@ -672,9 +772,9 @@ pub async fn serve_ws( #[cfg(not(unix))] let uds_listener: Option<()> = None; - // Read before `cfg` is moved into the supervisor: whether (and how aggressively) to reap idle - // sessions. - let idle_timeout = cfg.session_idle_timeout; + // Read before `cfg` is moved into the supervisor: how aggressively to reap idle sessions (and + // whether the operator opted out of reaping altogether). + let idle_timeout = resolve_idle_timeout(cfg.session_idle_timeout); let supervisor = Arc::new(Supervisor { sessions: Mutex::new(HashMap::new()), @@ -682,11 +782,17 @@ pub async fn serve_ws( }); let mut shutdown = crate::serve::ShutdownSignal::new()?; - // The idle reaper (off unless `--session-idle-timeout` is set). A background ticker that reclaims - // detached, idle, not-mid-run sessions — the same drop-`input_tx` → persist → join path shutdown - // uses. Its handle is aborted on shutdown so the process can exit cleanly. Tick at half the timeout - // (so a session is reaped within ~1.5× the timeout at worst), capped at 30s so a long timeout still - // ticks at a sane cadence. + // The idle reaper (on unless `--session-idle-timeout 0` turned it off). A background ticker that + // reclaims dead and detached-idle-not-mid-run sessions — the same drop-`input_tx` → persist → join + // path shutdown uses. Its handle is aborted on shutdown so the process can exit cleanly. Tick at + // half the timeout (so a session is reaped within ~1.5× the timeout at worst), capped at 30s so a + // long timeout still ticks at a sane cadence. + if idle_timeout.is_none() { + eprintln!( + "serve: idle-session reaper OFF (--session-idle-timeout 0) — every session, including one \ + no client ever re-attaches to, holds its thread and gateway pool until the daemon stops" + ); + } let reaper = idle_timeout.map(|t| { let supervisor = supervisor.clone(); eprintln!("serve: idle-session reaper on ({}s)", t.as_secs()); @@ -785,7 +891,13 @@ where // session id (parsed from the query) for use after the upgrade completes. let requested_id: Arc>> = Arc::new(Mutex::new(None)); let slot = requested_id.clone(); - let ws = tokio_tungstenite::accept_hdr_async( + // Every inbound message becomes one queued command line, so cap what a single one can cost us — + // tungstenite's defaults (64 MiB message / 16 MiB frame) let a client hand the session queue tens of + // MB at a time. Oversize is a hard protocol error (the connection closes); nothing is truncated. + let config = WebSocketConfig::default() + .max_message_size(Some(MAX_INBOUND_MESSAGE_BYTES)) + .max_frame_size(Some(MAX_INBOUND_MESSAGE_BYTES)); + let ws = tokio_tungstenite::accept_hdr_async_with_config( stream, move |req: &Request, resp: Response| -> Result { let uri = req.uri(); @@ -806,6 +918,7 @@ where } Ok(resp) }, + Some(config), ) .await?; @@ -823,3 +936,122 @@ where supervisor.attach(requested_id, ws).await; Ok(()) } + +#[cfg(test)] +#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +mod tests { + use super::*; + + /// Build a handle with no session behind it. `alive` keeps the input channel open (the receiver is + /// returned and must be held); dropping the receiver is how a test models a session whose loop ended. + fn handle( + attached: usize, + detached_ago: Option, + ) -> (SessionHandle, mpsc::Receiver) { + let (input_tx, input_rx) = mpsc::channel::(crate::serve::IN_CHANNEL_BOUND); + let h = SessionHandle { + input_tx, + out_conn: Arc::new(Mutex::new(OutFanout::default())), + join: None, + attached, + last_detached_at: detached_ago.and_then(|d| Instant::now().checked_sub(d)), + running: Arc::new(AtomicBool::new(false)), + }; + (h, input_rx) + } + + #[test] + fn idle_timeout_defaults_to_a_finite_window_and_zero_opts_out() { + assert_eq!( + resolve_idle_timeout(None), + Some(DEFAULT_SESSION_IDLE_TIMEOUT), + "no --session-idle-timeout must still reap: the map has no other way to shrink" + ); + assert_eq!( + resolve_idle_timeout(Some(Duration::ZERO)), + None, + "0 opts out" + ); + assert_eq!( + resolve_idle_timeout(Some(Duration::from_secs(5))), + Some(Duration::from_secs(5)) + ); + } + + #[test] + fn a_dead_session_is_reapable_however_it_looks_otherwise() { + // Its loop ended (input receiver gone) while a connection is *still attached* and its idle clock + // never started: every ordinary condition says "keep", and it must still be reaped — otherwise + // the entry, its fanout, and its exited thread's handle are retained for the daemon's life. + let (h, input_rx) = handle(1, None); + drop(input_rx); + assert!(is_reapable(&h, Duration::from_secs(3600))); + } + + #[test] + fn a_live_session_is_reapable_only_once_detached_past_the_timeout() { + let timeout = Duration::from_secs(60); + + let (attached, _rx) = handle(1, None); + assert!(!is_reapable(&attached, timeout), "attached: never"); + + let (fresh, _rx) = handle(0, None); + assert!( + !is_reapable(&fresh, timeout), + "detached but no clock started: never" + ); + + let (recent, _rx) = handle(0, Some(Duration::from_secs(10))); + assert!( + !is_reapable(&recent, timeout), + "detached, but not past the timeout" + ); + + let (idle, _rx) = handle(0, Some(Duration::from_secs(120))); + assert!(is_reapable(&idle, timeout), "detached past the timeout"); + + let (mid_run, _rx) = handle(0, Some(Duration::from_secs(120))); + mid_run.running.store(true, Ordering::Relaxed); + assert!( + !is_reapable(&mid_run, timeout), + "a detached background run must never be reaped out from under its turn" + ); + } + + /// A wedged session must not hold the caller (nor a blocking-pool thread) past the grace period, and + /// the threads that *did* exit must still be reaped in the same pass. + #[tokio::test] + async fn a_wedged_thread_does_not_hold_the_join_past_the_grace_period() { + let release = Arc::new(AtomicBool::new(false)); + let wedged = { + let release = release.clone(); + std::thread::spawn(move || { + while !release.load(Ordering::Relaxed) { + std::thread::sleep(Duration::from_millis(5)); + } + }) + }; + let finished = std::thread::spawn(|| {}); + + let start = Instant::now(); + join_handles_within(vec![finished, wedged], Duration::from_millis(200)).await; + let waited = start.elapsed(); + + assert!( + waited >= Duration::from_millis(200) && waited < Duration::from_secs(5), + "the wait must end at the grace period, not when the wedged session finally exits: {waited:?}" + ); + release.store(true, Ordering::Relaxed); + } + + #[tokio::test] + async fn joining_returns_as_soon_as_every_thread_has_exited() { + let threads: Vec<_> = (0..4).map(|_| std::thread::spawn(|| {})).collect(); + let start = Instant::now(); + join_handles_within(threads, Duration::from_secs(10)).await; + assert!( + start.elapsed() < Duration::from_secs(1), + "finished threads must not wait out the grace period" + ); + } +} diff --git a/crates/agent/src/tools/bash.rs b/crates/agent/src/tools/bash.rs index cb6e880..fac9e97 100644 --- a/crates/agent/src/tools/bash.rs +++ b/crates/agent/src/tools/bash.rs @@ -488,6 +488,12 @@ fn truncation_details(snap: &OutputSnapshot) -> Option { "max_bytes": snap.truncation.max_bytes, }, "full_output_path": snap.full_output_path, + // A spill file that hit its byte budget holds only a prefix of the stream (see + // `output::MAX_SPILL_BYTES`). A client rendering `full_output_path` as "the complete + // output, click to open" would otherwise present that prefix as the whole thing — the + // same dishonesty `format_output`'s marker guards against for the model's own text. + "full_output_capped": snap.full_output_capped, + "full_output_bytes": snap.full_output_bytes, }) }) } @@ -578,6 +584,8 @@ mod tests { max_bytes: DEFAULT_MAX_BYTES, }, full_output_path: Some("/tmp/pi-bash-abc.log".into()), + full_output_capped: false, + full_output_bytes: 12345, last_line_bytes: 4, }; let details = truncation_details(&snap).expect("truncated snapshot must carry details"); @@ -591,10 +599,42 @@ mod tests { assert_eq!(details["truncation"]["max_lines"], DEFAULT_MAX_LINES); assert_eq!(details["truncation"]["max_bytes"], DEFAULT_MAX_BYTES); assert_eq!(details["full_output_path"], "/tmp/pi-bash-abc.log"); + assert_eq!(details["full_output_capped"], false); + assert_eq!(details["full_output_bytes"], 12345); // `full_output_path` is a sibling of `truncation`, not nested inside it. assert!(details["truncation"].get("full_output_path").is_none()); } + #[test] + fn truncation_details_tells_a_client_when_the_spill_file_holds_only_a_prefix() { + // A spill file that hit `output::MAX_SPILL_BYTES` stops short of the stream. A client that + // renders `full_output_path` as "the complete output" would otherwise show a firehose's first + // 128 MiB as if it were the whole thing — the wire payload has to carry the same admission the + // model's own marker text does. + let snap = OutputSnapshot { + content: "tail".into(), + truncation: Truncation { + truncated: true, + truncated_by: Some(TruncatedBy::Bytes), + total_lines: 900_000, + total_bytes: 4_000_000_000, + output_lines: 2000, + output_bytes: 51_200, + last_line_partial: false, + max_lines: DEFAULT_MAX_LINES, + max_bytes: DEFAULT_MAX_BYTES, + }, + full_output_path: Some("/tmp/pi-bash-abc.log".into()), + full_output_capped: true, + full_output_bytes: 134_217_728, + last_line_bytes: 4, + }; + let details = truncation_details(&snap).expect("truncated snapshot must carry details"); + assert_eq!(details["full_output_capped"], true); + assert_eq!(details["full_output_bytes"], 134_217_728u64); + assert_eq!(details["truncation"]["total_bytes"], 4_000_000_000u64); + } + #[test] fn truncation_details_is_none_when_output_was_not_truncated() { let snap = OutputSnapshot { @@ -611,6 +651,8 @@ mod tests { max_bytes: DEFAULT_MAX_BYTES, }, full_output_path: None, + full_output_capped: false, + full_output_bytes: 0, last_line_bytes: 3, }; assert!(truncation_details(&snap).is_none()); diff --git a/crates/agent/src/tools/output.rs b/crates/agent/src/tools/output.rs index 80f37cc..2ee7b72 100644 --- a/crates/agent/src/tools/output.rs +++ b/crates/agent/src/tools/output.rs @@ -23,6 +23,25 @@ pub const DEFAULT_MAX_LINES: usize = 2000; /// Default cap on displayed bytes (50 KiB). Protects the model's context window. pub const DEFAULT_MAX_BYTES: usize = 50 * 1024; +/// Hard ceiling on how many bytes any one spill file may take (128 MiB). +/// +/// [`DEFAULT_MAX_BYTES`] bounds only what the *model is shown*; the full stream still reaches this +/// accumulator, and the spill file used to grow for as long as the command kept talking. Memory stays +/// bounded either way (the rolling tail, plus a pre-spill buffer that's flushed and freed) — the file +/// was the unbounded side, and on this crate's deployment target the file *is* memory: the system temp +/// dir is a tmpfs (see `worktree.rs`), so "spill to disk" is spill to RAM. With `bash`'s 30-minute +/// default timeout, one `yes` or `cat /dev/urandom | base64` could fill it — and the file outlives the +/// call by [`STALE_TEMP_FILE_AGE`], so the damage persists long after the command is gone. +/// +/// The number is sized against the question "what is the largest output someone would genuinely want +/// to page back through?": a fully verbose build log, a long test run, a chatty migration — real cases +/// land in the low tens of MiB at the extreme. 128 MiB clears that by roughly an order of magnitude, +/// so no legitimate command loses bytes it had a reader for, while several concurrent firehoses +/// together still stay a small fraction of a multi-GiB tmpfs. Past the budget we simply stop writing: +/// the accumulator keeps running, so the tail, the totals, and the truncation markers all stay +/// correct — the only thing dropped is the middle of a stream nobody was going to read. +pub const MAX_SPILL_BYTES: u64 = 128 * 1024 * 1024; + /// Which limit tripped truncation. Mirrors pi's `truncatedBy: "lines" | "bytes" | null`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum TruncatedBy { @@ -68,8 +87,18 @@ pub struct OutputSnapshot { pub content: String, /// Truncation metadata for `content`. pub truncation: Truncation, - /// Path to the temp file holding the *complete* output, when one was opened. + /// Path to the temp file holding the complete output, when one was opened — or, when + /// `full_output_capped` is set, only its first `full_output_bytes`. pub full_output_path: Option, + /// Whether the spill file stopped short of the complete stream because it hit + /// [`MAX_SPILL_BYTES`]. The file still holds a genuine *prefix* of the output (not a corrupt + /// one — see [`OutputAccumulator::mark_spill_broken`] for that case), so it's worth pointing the + /// model at; but a marker that names a path while the model assumes it's the whole document is + /// worse than naming no path at all, so [`format_output`] says so explicitly when this is set. + pub full_output_capped: bool, + /// How many bytes actually landed in the spill file. Equals `truncation.total_bytes` unless + /// `full_output_capped`. + pub full_output_bytes: u64, /// Full byte size of the current (last) line — used for the partial-line marker. pub last_line_bytes: u64, } @@ -253,6 +282,13 @@ pub struct OutputAccumulator { temp_path: Option, temp_writer: Option>, + // The spill file's byte budget and what's been spent of it. `spill_capped` means the budget ran + // out and the writer was closed: the file holds a prefix, and further bytes are counted (totals, + // tail) but not written anywhere. + spill_budget: u64, + spill_bytes: u64, + spill_capped: bool, + finished: bool, } @@ -288,10 +324,22 @@ impl OutputAccumulator { spill_failed: false, temp_path: None, temp_writer: None, + spill_budget: MAX_SPILL_BYTES, + spill_bytes: 0, + spill_capped: false, finished: false, } } + /// Builder-style: shrink the spill file's byte budget. Test-only — production always runs the + /// real [`MAX_SPILL_BYTES`], and a test that had to actually emit 128 MiB to reach it would be + /// writing that much to the same tmpfs the budget exists to protect. + #[cfg(test)] + fn with_spill_budget(mut self, budget: u64) -> Self { + self.spill_budget = budget; + self + } + /// Total lines over the complete stream: completed lines plus one for a trailing partial line. fn total_lines(&self) -> u64 { self.completed_lines + u64::from(self.has_open_line) @@ -310,12 +358,8 @@ impl OutputAccumulator { self.extend_tail(data); if self.spilled { - // Post-spill: stream straight to the file so memory stays bounded. - if let Some(w) = self.temp_writer.as_mut() { - if w.write_all(data).is_err() { - self.mark_spill_broken(); - } - } + // Post-spill: stream straight to the file so memory stays bounded — up to the budget. + self.write_spill(data); } else if !self.spill_failed { // Pre-spill: keep the complete bytes in memory so a later spill can flush them whole. self.raw_buffer.extend_from_slice(data); @@ -394,6 +438,11 @@ impl OutputAccumulator { .temp_path .as_ref() .map(|p| p.to_string_lossy().into_owned()), + // Only meaningful while there's still a file to describe: a spill that later *broke* has + // already dropped its path (`mark_spill_broken`), and a capped-then-broken file must read + // as "no full output", not as a prefix a reader can trust. + full_output_capped: self.spill_capped && self.temp_path.is_some(), + full_output_bytes: self.spill_bytes, last_line_bytes: self.current_line_bytes, } } @@ -506,16 +555,14 @@ impl OutputAccumulator { } match opts.open(&path) { Ok(f) => { - let mut w = BufWriter::new(f); - if w.write_all(&self.raw_buffer).is_err() { - self.spill_failed = true; - self.raw_buffer = Vec::new(); - return; - } - self.raw_buffer = Vec::new(); self.temp_path = Some(path); - self.temp_writer = Some(w); + self.temp_writer = Some(BufWriter::new(f)); self.spilled = true; + // The buffered bytes go in through the same budgeted write as every later chunk — + // one place decides what the file is allowed to hold, so the flush of everything + // seen before the spill can't slip past the budget by writing itself whole. + let buffered = std::mem::take(&mut self.raw_buffer); + self.write_spill(&buffered); } Err(_) => { self.spill_failed = true; @@ -524,6 +571,47 @@ impl OutputAccumulator { } } + /// Write `data` to the spill file, spending it against the file's byte budget and writing nothing + /// once that budget is gone — see [`MAX_SPILL_BYTES`] for why an unbounded spill file is a memory + /// leak on this deployment target, not merely a disk one. + /// + /// Reaching the budget is deliberately *not* [`mark_spill_broken`](Self::mark_spill_broken): a + /// broken file is one whose contents no longer correspond to the command's output at all, so it + /// must stop being advertised; a budget-capped file is a genuine, complete prefix, still the best + /// thing the model can go read. What it must never do is masquerade as the whole document, hence + /// `spill_capped` — the flag [`format_output`] turns into an explicit "this file stops early" + /// marker. Everything else about the accumulator carries on unchanged past the budget: the running + /// totals and rolling tail are what make the truncation metadata honest about the bytes we dropped. + fn write_spill(&mut self, data: &[u8]) { + let remaining = self.spill_budget.saturating_sub(self.spill_bytes); + let take = remaining.min(data.len() as u64) as usize; + let Some(w) = self.temp_writer.as_mut() else { + return; + }; + if take > 0 && w.write_all(&data[..take]).is_err() { + self.mark_spill_broken(); + return; + } + self.spill_bytes += take as u64; + if self.spill_bytes >= self.spill_budget { + self.close_spill_at_budget(); + } + } + + /// Retire the spill writer once its byte budget is spent: flush what's buffered so the file on + /// disk is complete up to the budget, then drop the writer so the fd is released and no later + /// chunk can reopen the tap. `temp_path` stays — the prefix is still readable, and `spill_capped` + /// is what tells the caller it stops short. + fn close_spill_at_budget(&mut self) { + if let Some(mut w) = self.temp_writer.take() { + if w.flush().is_err() { + self.mark_spill_broken(); + return; + } + } + self.spill_capped = true; + } + /// Mark the spill file as no longer trustworthy after a write/flush failure partway through /// streaming to it (e.g. the disk filled up mid-command): stop advertising it via /// `full_output_path` — the bytes on disk are now silently truncated relative to what the @@ -672,6 +760,32 @@ fn truncate_str_to_bytes_from_end(s: &str, max_bytes: usize) -> &str { &s[start..] } +/// The `Full output: ` tail of a truncation marker — pi's string verbatim in the ordinary case, +/// where the file really does hold everything the command printed. +/// +/// When the spill file hit [`MAX_SPILL_BYTES`] it holds only the first N bytes of a much larger +/// stream, and the marker has to say so: a bare `Full output: ` invites the model to read that +/// file and reason about a firehose's opening moments as if they were the whole story — a silently +/// partial document presented as complete, which is a worse failure than admitting the bytes are gone. +/// The totals in the same marker (`of {total_lines}`) already describe the *complete* stream, so this +/// segment names what's actually on disk against it. +fn full_output_segment(snapshot: &OutputSnapshot) -> String { + // pi interpolates the path directly; when there is none we render an empty string rather than the + // literal "undefined" JS would produce. + let path = snapshot.full_output_path.as_deref().unwrap_or(""); + if snapshot.full_output_capped { + // Sized from the snapshot, not from `MAX_SPILL_BYTES`: the marker must state the bytes that + // are actually in the file the model is about to read, whatever budget produced them. + format!( + "Full output too large to save — only the first {} of {} was written: {path}", + format_size(snapshot.full_output_bytes), + format_size(snapshot.truncation.total_bytes), + ) + } else { + format!("Full output: {path}") + } +} + /// Render a snapshot as the bash tool's display text — a port of `bash.ts`'s `formatOutput`. When not /// truncated, returns `content` (or `empty_text` when `content` is empty). When truncated, appends a /// blank line and one of three markers, matching pi's strings byte-for-byte. @@ -690,26 +804,23 @@ pub fn format_output(snapshot: &OutputSnapshot, empty_text: &str) -> String { .total_lines .saturating_sub(t.output_lines) .saturating_add(1); - // pi interpolates the path directly; when there is none we render an empty string rather - // than the literal "undefined" JS would produce. - let path = snapshot.full_output_path.as_deref().unwrap_or(""); + let full = full_output_segment(snapshot); let marker = if t.last_line_partial { format!( - "[Showing last {} of line {} (line is {}). Full output: {}]", + "[Showing last {} of line {} (line is {}). {full}]", format_size(t.output_bytes), end_line, format_size(snapshot.last_line_bytes), - path, ) } else if matches!(t.truncated_by, Some(TruncatedBy::Lines)) { format!( - "[Showing lines {start_line}-{end_line} of {}. Full output: {path}]", - t.total_lines, + "[Showing lines {start_line}-{end_line} of {}. {full}]", + t.total_lines ) } else { format!( - "[Showing lines {start_line}-{end_line} of {} ({} limit). Full output: {path}]", + "[Showing lines {start_line}-{end_line} of {} ({} limit). {full}]", t.total_lines, format_size(DEFAULT_MAX_BYTES as u64), ) @@ -959,6 +1070,122 @@ mod tests { assert!(acc.spill_failed); } + /// Feed `total` bytes of newline-terminated output through `acc` in 8 KiB chunks (the shape a + /// firehose actually arrives in), with the final line ending in `END` so a tail assertion can + /// prove the accumulator kept working past whatever the spill file did. + fn feed_firehose(acc: &mut OutputAccumulator, total: usize) { + let chunk = vec![b'f'; 8 * 1024]; + let mut written = 0usize; + while written + chunk.len() <= total { + acc.append(&chunk); + written += chunk.len(); + } + acc.append(b"\nEND\n"); + } + + #[test] + fn output_under_the_spill_budget_is_written_to_the_file_in_full() { + // The budget is a backstop, not a haircut: anything short of it must still land on disk + // byte-for-byte, with no "capped" admission attached to a file that is in fact complete. + let mut acc = OutputAccumulator::with_prefix("pi-bash").with_spill_budget(1024 * 1024); + let mut full = Vec::new(); + for n in 0..300u32 { + let mut line = format!("chunk-{n:04}-").into_bytes(); + line.resize(1023, b'x'); + line.push(b'\n'); + full.extend_from_slice(&line); + acc.append(&line); + } + acc.finish(); + let snap = acc.snapshot(true); + + let path = snap.full_output_path.clone().expect("must spill"); + assert!(!snap.full_output_capped); + assert_eq!(snap.full_output_bytes, full.len() as u64); + assert_eq!(read_temp(&path), full); + // The ordinary marker, unchanged — an uncapped file says nothing about budgets. + let out = format_output(&snap, "(no output)"); + assert!(out.contains(&format!("Full output: {path}]")), "got: {out}"); + + let _ = std::fs::remove_file(&path); + } + + #[test] + fn a_firehose_past_the_spill_budget_stops_growing_the_file() { + // The leak this guards: `DEFAULT_MAX_BYTES` caps only what the model is *shown*, so before the + // budget existed the spill file grew for as long as the command talked — and the system temp + // dir on this crate's deployment target is a tmpfs, so an unbounded spill file is unbounded + // *RAM* (see `MAX_SPILL_BYTES`). `bash: yes` under the 30-minute default timeout was enough. + let budget = 256 * 1024; + let mut acc = OutputAccumulator::with_prefix("pi-bash").with_spill_budget(budget as u64); + feed_firehose(&mut acc, 8 * 1024 * 1024); // 32x the budget + acc.finish(); + let snap = acc.snapshot(true); + + let path = snap.full_output_path.clone().expect("must spill"); + let on_disk = std::fs::metadata(&path).unwrap().len(); + assert!( + on_disk <= budget as u64, + "spill file must never exceed its budget: {on_disk} > {budget}" + ); + // And it must actually have written up to the budget, not bailed out at the first chunk. + assert_eq!(on_disk, budget as u64); + assert_eq!(snap.full_output_bytes, budget as u64); + assert!(snap.full_output_capped); + + let _ = std::fs::remove_file(&path); + } + + #[test] + fn truncation_metadata_stays_honest_after_the_spill_budget_is_hit() { + // Stopping the writes must not stop the accumulator: the tail the model reads, the totals it + // reasons about, and the marker that admits how much was dropped all have to survive the + // budget — otherwise a capped firehose would present its opening 256 KiB as the whole story. + let budget = 256 * 1024usize; + let total = 8 * 1024 * 1024usize; + let mut acc = OutputAccumulator::with_prefix("pi-bash").with_spill_budget(budget as u64); + feed_firehose(&mut acc, total); + acc.finish(); + let snap = acc.snapshot(true); + + // Totals cover the COMPLETE stream, not the bytes that made it to disk. `feed_firehose` writes + // whole 8 KiB chunks up to `total` and then "\nEND\n" (5 bytes, 2 newlines → 2 lines: the long + // 'f' run it terminates, and "END"). + assert_eq!(snap.truncation.total_bytes, (total + 5) as u64); + assert_eq!(snap.truncation.total_lines, 2); + assert!(snap.truncation.truncated); + assert_eq!(snap.truncation.truncated_by, Some(TruncatedBy::Bytes)); + assert!(snap.truncation.output_bytes <= DEFAULT_MAX_BYTES as u64); + + // The tail still tracks bytes that arrived long after the file stopped taking them. (The 'f' + // run ahead of it is dropped by `snapshot_text`, which never shows a half-line at the top of + // the window — the rolling tail landed mid-line, as it must for a stream this size.) + assert_eq!( + snap.content.trim_end(), + "END", + "the rolling tail must keep working past the budget" + ); + + // And the rendered marker says plainly that the saved file stops short — the model must not + // read a 256 KiB prefix believing it's the complete 8 MiB output. + let path = snap.full_output_path.clone().expect("must spill"); + let out = format_output(&snap, "(no output)"); + assert!( + out.contains("only the first 256.0KB of 8.0MB was written"), + "the marker must quantify what was dropped: {out}" + ); + assert!( + out.contains(&path), + "the prefix is still worth reading: {out}" + ); + assert!( + !out.contains("Full output: "), + "a capped file must not be advertised as the full output: {out}" + ); + + let _ = std::fs::remove_file(&path); + } + #[test] fn line_count_alone_triggers_a_spill_well_under_the_byte_cap() { // Many short lines: total bytes stay far under the rolling byte cap, but the line count alone @@ -1109,6 +1336,8 @@ mod tests { max_bytes: DEFAULT_MAX_BYTES, }, full_output_path: Some("/tmp/pi-bash-abc.log".to_string()), + full_output_capped: false, + full_output_bytes: 12345, last_line_bytes: 4, }; assert_eq!( diff --git a/crates/agent/src/worktree.rs b/crates/agent/src/worktree.rs index a1c713f..30df2b0 100644 --- a/crates/agent/src/worktree.rs +++ b/crates/agent/src/worktree.rs @@ -494,30 +494,62 @@ async fn command_output_with_stdin( dir: &Path, args: &[&str], stdin: &[u8], +) -> Result { + let mut cmd = Command::new("git"); + cmd.arg("-C").arg(dir).args(args); + output_with_stdin(cmd, &format!("git {}", args.join(" ")), stdin).await +} + +/// Feed `stdin` to `cmd` while **concurrently** draining its stdout and stderr, and return what it +/// wrote once it exits. +/// +/// The concurrency is the whole point of this function, not an optimization. A pipe holds about 64 KiB; +/// past that a writer blocks until the far end reads. Push the entire patch in first and only then start +/// reading, and any child that produces more than a pipeful of output before it has swallowed the whole +/// patch wedges the pair of us forever: it blocks writing stdout, therefore stops reading stdin, +/// therefore our `write_all` never returns. Both callers reach that shape on real input — `git apply +/// --numstat -z` emits a record per file (a patch touching a few thousand files overruns the buffer), +/// and `git apply --3way` emits a `U ` line per file on a heavily conflicted patch. The result +/// would be a turn hung for the life of the process, holding a child and three pipe fds. +/// +/// The `join!` also means the child is always waited on, including when the write fails. +async fn output_with_stdin( + mut cmd: Command, + label: &str, + stdin: &[u8], ) -> Result { use tokio::io::AsyncWriteExt; - let mut child = Command::new("git") - .arg("-C") - .arg(dir) - .args(args) + let mut child = cmd .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .spawn() - .map_err(|e| format!("git {}: {e}", args.join(" ")))?; - child + .map_err(|e| format!("{label}: {e}"))?; + let mut pipe = child .stdin - .as_mut() - .ok_or_else(|| "git stdin unavailable".to_string())? - .write_all(stdin) - .await - .map_err(|e| format!("git {}: writing patch: {e}", args.join(" ")))?; - drop(child.stdin.take()); - child - .wait_with_output() - .await - .map_err(|e| format!("git {}: {e}", args.join(" "))) + .take() + .ok_or_else(|| format!("{label}: stdin unavailable"))?; + + let feed = async move { + let result = pipe.write_all(stdin).await; + // Dropping the handle closes the write end: without that EOF the child keeps waiting for more + // input and `wait_with_output` never returns. + drop(pipe); + result + }; + let (written, out) = tokio::join!(feed, child.wait_with_output()); + let out = out.map_err(|e| format!("{label}: {e}"))?; + if let Err(e) = written { + // A child that decides the input is garbage and exits before reading all of it (`git apply` on a + // malformed patch) closes the pipe under us, and we see EPIPE. That is not the failure worth + // reporting — the exit status and stderr we did collect say what actually went wrong, and the + // callers already turn those into a real message. Anything else is a genuine I/O fault. + if e.kind() != std::io::ErrorKind::BrokenPipe { + return Err(format!("{label}: writing stdin: {e}")); + } + } + Ok(out) } /// [`command_output_with_stdin`], erroring on a non-zero exit and returning stdout. @@ -1042,6 +1074,44 @@ mod tests { ); } + /// The pipe-deadlock regression guard. `cat` echoes its stdin straight back, so a payload far larger + /// than a pipe buffer (~64 KiB) in *both* directions is exactly the shape that wedges a + /// write-everything-then-drain implementation: the child blocks writing stdout, stops reading stdin, + /// and our `write_all` never returns. `git apply --numstat` on a patch touching thousands of files + /// does the same thing with less convenient setup. The timeout is what makes a regression *fail* + /// rather than hang CI forever. + #[tokio::test] + async fn feeding_a_child_that_floods_its_stdout_does_not_deadlock() { + let payload = vec![b'x'; 4 << 20]; + let out = tokio::time::timeout( + std::time::Duration::from_secs(30), + output_with_stdin(Command::new("cat"), "cat", &payload), + ) + .await + .expect("deadlocked writing stdin to a child that was flooding its stdout") + .expect("cat"); + assert!(out.status.success()); + assert_eq!(out.stdout.len(), payload.len()); + } + + /// A child that exits before reading its whole stdin hands us EPIPE on the write. That is not the + /// interesting failure — the caller wants the exit status and stderr, which is what says why it quit. + #[tokio::test] + async fn a_child_that_exits_without_reading_stdin_yields_its_output_not_a_write_error() { + let payload = vec![b'x'; 4 << 20]; + let mut cmd = Command::new("sh"); + cmd.args(["-c", "echo nope >&2; exit 3"]); + let out = tokio::time::timeout( + std::time::Duration::from_secs(30), + output_with_stdin(cmd, "sh", &payload), + ) + .await + .expect("deadlocked writing stdin to a child that never read it") + .expect("a broken stdin pipe must not mask the child's own output"); + assert_eq!(out.status.code(), Some(3)); + assert_eq!(String::from_utf8_lossy(&out.stderr).trim(), "nope"); + } + #[test] fn sanitize_never_yields_a_path_separator_or_a_parent_ref() { assert_eq!(sanitize("scout-0"), "scout-0"); diff --git a/crates/agent/tests/serve_reaper.rs b/crates/agent/tests/serve_reaper.rs index 1d9a055..e01f1de 100644 --- a/crates/agent/tests/serve_reaper.rs +++ b/crates/agent/tests/serve_reaper.rs @@ -16,7 +16,7 @@ use common::{ }; use serde_json::{Value, json}; -/// Spawn a `serve --listen` child with the idle reaper armed at `idle_secs`. +/// 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) -> Child { Command::new(env!("CARGO_BIN_EXE_beyond-ai-agent")) .args([ @@ -112,6 +112,47 @@ async fn idle_session_is_reaped_but_its_turn_persists() { let _ = child.wait(); } +/// (ia) The reaper is on by default (a finite idle window), so `--session-idle-timeout 0` is the +/// explicit opt-out an operator uses to pin every session for the daemon's lifetime: a detached, idle +/// session must still be `live` long after any default window would have reclaimed it. +#[tokio::test] +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 ws = ws_connect(port, Some(SID)).await; + ws_send( + &mut ws, + json!({ "type": "prompt", "id": "p", "message": "hi" }), + ) + .await; + let frames = ws_read_until_response(&mut ws, "prompt").await; + assert_eq!(frames.last().unwrap()["success"], true); + drop(ws); + } + + // Longer than any reaper tick would need if one were armed at all. + tokio::time::sleep(Duration::from_secs(3)).await; + + let sessions = list_daemon_sessions(port).await; + let entry = sessions + .iter() + .find(|s| s["id"] == SID) + .unwrap_or_else(|| panic!("{SID} should still be listed: {sessions:#?}")); + assert_eq!( + entry["live"], true, + "--session-idle-timeout 0 must disable reaping entirely: {entry}" + ); + + let _ = child.kill(); + let _ = child.wait(); +} + /// (ii) Negative: a session that is *still attached* past the timeout is never reaped — the reaper only /// reclaims detached sessions. The connection is held open across the whole wait; the id stays `live`. #[tokio::test] From 573e5262400e51782b9ffcf5c13cb70793f3b707 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sun, 12 Jul 2026 17:17:28 -0700 Subject: [PATCH 9/9] perf(agent): make session listing O(sessions), not O(transcript bytes) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measured first, then fixed — and the measurements killed half the original audit. `benches/persistence.rs` is new and covers the half of this crate that had no benchmark at all; every number below is from it or from `benches/search.rs`. The constraint: `read_listing` fully parses every line of every session file to recover a preview, a message count, and a max timestamp. Cost is O(total transcript bytes), and transcripts are append-only and never pruned — the one cost here that grows without bound with a user's history. A sidecar `.listings.json` per session directory now caches each file's derived `SessionMeta` against the `(size, mtime)` it was computed from. Session files are append-only, so any write moves both and a stale entry cannot survive one. Listing a 14 MB history: 15.1 ms -> 11.4 µs, and — the actual point — flat in transcript size (11.1 µs at 1.4 MB vs 11.4 µs at 14 MB). It is a pure cache: every read is best-effort and any failure falls back to a full scan, so a lost or clobbered index costs time, never correctness. The derived fields are carried explicitly in the index rather than inside `SessionMeta`, which marks all four `#[serde(skip)]` so they can't go stale in the on-disk *header*. Right call there, exactly wrong here — the first cut serialized a `SessionMeta` straight in and silently dropped the only fields the index exists to remember. The tests caught it. Two cheaper fixes were tried and reverted, because a control benchmark (`listing_read_floor`) showed only ~15% of a listing is reading bytes and ~85% is serde *tokenizing* payload it discards. Skipping materialization doesn't skip tokenization: a peek-parse cut allocations 75% and bought 0% wall time. The only way to make the scan cheaper was to not do it. Also fixed, each measured: - `edit`: ASCII fast path in normalization. NFKC is the identity on ASCII and every `fold_char` arm is >= U+00A0, so for ASCII input — which source code overwhelmingly is — normalization provably reduces to dropping trailing whitespace, with no `char` decoding at all. normalize 2.82 ms -> 107 µs (26x); full `edit::run` ~2.7 ms -> ~0.5 ms (~5x). A test pins that the fast path yields byte-identical text *and* an identical offset map — it feeds the splice, so a divergence would corrupt files, not just mismatch. - `edit`: onto `spawn_blocking`. On a 4 MB file it pinned a `current_thread` runtime — what `serve_ws` gives every session — for 72.9 ms, stalling that session's event pump and its abort/steer command loop. This is the invariant `serve.rs::persist_blocking` already states; `edit` just hadn't followed. `tests/tool_reactor_stall.rs` now pins it. `write` and `ls` were measured on the same harness and do NOT stall, so they are deliberately left alone rather than wrapped on principle. - `session_store::write_line`: single-pass `to_writer` instead of `Entry` -> owned `Value` tree -> `String` — the same two-pass shape `serve::event_frame` was already fixed out of and `benches/serve_events.rs` was written to condemn; the fix had never reached this module. Plus a borrowing `EntryRef` so `append_new` stops cloning each `Message` purely to hand it to serde. Serializer 2.9x, allocations 29 -> 1; `append_new` -10% end-to-end (fsync-dominated, so that's the ceiling). Also included, from a parallel session working in this same tree, NOT part of the above perf work: `Node::content` moves behind an `Arc` so `rewrite_compacted` — which keeps the pre-compaction chain intact *and* re-emits the surviving suffix under fresh ids — no longer holds every kept message twice per compaction round. Committed together because the two changes interleave in `session_store.rs`; it is green under the full suite. cargo test -p beyond-ai-agent: all green (1275 lib + every integration suite). clippy --all-targets: clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/agent/Cargo.toml | 7 + crates/agent/benches/persistence.rs | 276 +++++++++++++ crates/agent/benches/search.rs | 104 +++++ crates/agent/src/session_store.rs | 435 ++++++++++++++++++-- crates/agent/src/tools/edit.rs | 198 ++++++++- crates/agent/tests/session_listing_index.rs | 197 +++++++++ crates/agent/tests/tool_reactor_stall.rs | 147 +++++++ 7 files changed, 1316 insertions(+), 48 deletions(-) create mode 100644 crates/agent/benches/persistence.rs create mode 100644 crates/agent/tests/session_listing_index.rs create mode 100644 crates/agent/tests/tool_reactor_stall.rs diff --git a/crates/agent/Cargo.toml b/crates/agent/Cargo.toml index 3089295..072b13b 100644 --- a/crates/agent/Cargo.toml +++ b/crates/agent/Cargo.toml @@ -170,3 +170,10 @@ harness = false [[bench]] name = "serve_events" harness = false + +# Session persistence — the listing scan (whose cost grows with a user's whole history) and the +# per-turn durable append. Deliberately runs its temp dirs on the real filesystem, not tmpfs: see the +# bench's own `real_disk_tempdir` for why measuring `sync_all` on tmpfs answers nothing. +[[bench]] +name = "persistence" +harness = false diff --git a/crates/agent/benches/persistence.rs b/crates/agent/benches/persistence.rs new file mode 100644 index 0000000..9cffb7f --- /dev/null +++ b/crates/agent/benches/persistence.rs @@ -0,0 +1,276 @@ +// Bench target: `.unwrap()`/`.expect()` set up fixtures; not production code. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +//! Session-persistence bench — the half of `crates/agent` that had no coverage at all, and the only +//! cost in the crate that grows without bound in a user's history. +//! +//! Two paths, measured against a *realistic* transcript (a coding session's lines are dominated by +//! tool results — a `read` output, a `grep` dump — not by prose): +//! +//! - `listing`: `SessionRepo::list_with_progress` → `scan_listings` → `read_listing`, which today +//! fully deserializes **every** message of **every** session file to recover a preview, a count, +//! and a max timestamp. Cost is O(total transcript bytes). Args are the per-session message count, +//! so the slope across them *is* the scaling claim. +//! - `append_new`: one turn's durable append, on a **real ext4/NVMe temp dir** (`TempDir::new_in` +//! under the repo, not `/tmp` — which is tmpfs here, where `sync_all` is nearly free and would +//! flatter any CPU-side fix). This is the number that says whether shaving the serializer is +//! visible at all next to the fsync, or whether it drowns. +//! - `serialize_entry`: the serializer in isolation, both strategies, so the CPU delta is legible +//! even when `append_new`'s fsync hides it. `Entry`/`write_line` are private to `session_store` +//! (deliberately), so — following `benches/serve_events.rs`'s own precedent — this replicates both +//! strategies against the same public `agent_core::Message` payload the real `Entry` wraps. +//! +//! Run: `cargo bench -p beyond-ai-agent --bench persistence`. + +use std::hint::black_box; +use std::path::{Path, PathBuf}; + +use agent_core::{ContentBlock, Message}; +use beyond_ai_agent::session_store::{SessionMeta, SessionRepo, SessionStore}; +use divan::Bencher; +use serde::Serialize; +use serde_json::Value; +use tempfile::TempDir; + +#[global_allocator] +static ALLOC: divan::AllocProfiler = divan::AllocProfiler::system(); + +fn main() { + divan::main(); +} + +/// A ~64 KB `read`-tool output — the single most common shape of bytes in a coding transcript, and +/// the thing `read_listing` currently materializes in full only to throw away. +fn big_tool_result() -> String { + " 1\tuse std::io;\n".repeat(3200) +} + +/// One turn as it actually lands on disk: a user prompt, an assistant turn that *both* narrates and +/// calls a tool, and the tool result carrying the bulk of the bytes. +/// +/// The assistant prose is not decoration. `read_listing` accumulates a 50,000-char search corpus from +/// every user/assistant **text** block, and a transcript's text is overwhelmingly the assistant's — +/// user prompts are one line, tool results contribute no text at all. A workload with no assistant +/// prose never fills that corpus, which silently changes which branch of the scan is under test. Ask +/// a realistic question or you measure a fictional one. +fn turn(i: usize) -> Vec { + let prose = format!( + "Looking at file {i} now. The function it defines threads its error type through a \ + `Result` alias, which is why the call site two modules up compiles even though the \ + signatures look mismatched at a glance. I'll read it and confirm before changing anything, \ + since the alias is re-exported and a change here would ripple. " + ) + .repeat(3); + vec![ + Message::user(format!("please look at file number {i} and explain it")), + Message::assistant(vec![ + ContentBlock::Text { + text: prose, + id: None, + phase: None, + }, + ContentBlock::ToolUse { + id: format!("toolu_{i:06}"), + name: "read".to_string(), + input: serde_json::json!({ "path": format!("src/file_{i}.rs") }), + thought_signature: None, + }, + ]), + Message::tool_result(format!("toolu_{i:06}"), big_tool_result(), false), + ] +} + +/// A temp dir on the **real** filesystem (ext4/NVMe here), not `/tmp` (tmpfs). `sync_all` on tmpfs is +/// essentially a no-op, which would make every durability cost vanish and every CPU-side fix look +/// like a win. The append path's whole question is "is the serializer visible next to the fsync?" — +/// asking it on tmpfs answers a question nobody has. +fn real_disk_tempdir() -> TempDir { + let here = Path::new(env!("CARGO_MANIFEST_DIR")); + let scratch = here.join("target").join("bench-scratch"); + std::fs::create_dir_all(&scratch).unwrap(); + TempDir::new_in(&scratch).unwrap() +} + +/// Build a session directory holding one session file of `turns` turns, and hand back the repo dir. +/// Kept alive by the returned `TempDir`. +fn session_dir(turns: usize) -> (TempDir, PathBuf) { + let dir = real_disk_tempdir(); + let repo = SessionRepo::open(dir.path()).unwrap(); + let mut store = repo + .create(SessionMeta::new("/repo", "claude-sonnet-5")) + .unwrap(); + let mut messages = Vec::new(); + for i in 0..turns { + messages.extend(turn(i)); + store.append_new(&messages).unwrap(); + } + let path = dir.path().to_path_buf(); + (dir, path) +} + +// --------------------------------------------------------------------------------------------- +// The constraint: listing cost is O(total transcript bytes), because every line is fully parsed. +// --------------------------------------------------------------------------------------------- + +/// The **cold** listing: the sidecar index is thrown away before every sample, so each one pays the +/// full O(total transcript bytes) parse. This is what every listing cost before the index existed, and +/// what the first listing after a session is written still costs. +/// +/// `turns` × 3 messages, of which every third carries ~64 KB. 20 turns ≈ 1.4 MB, 200 turns ≈ 14 MB — +/// an ordinary long coding session. The slope from one to the other is the scaling claim. +#[divan::bench(args = [20, 200], sample_count = 10)] +fn listing_cold(bencher: Bencher, turns: usize) { + let (_keep, dir) = session_dir(turns); + let repo = SessionRepo::open(&dir).unwrap(); + bencher.bench_local(|| { + let _ = std::fs::remove_file(dir.join(".listings.json")); + let metas = repo.list_with_progress(|_, _| {}).unwrap(); + black_box(metas.len()); + }); +} + +/// The **warm** listing: the sidecar index is valid, so every file is answered by one `stat` instead of +/// a full parse. This is the steady state — a session's `(size, mtime)` only moves when it's actually +/// appended to, so at most the one session you're currently talking to is ever a miss. +/// +/// The gap between this and `listing_cold` is the whole point of the index, and it should widen with +/// transcript size, because this side doesn't depend on it at all. +#[divan::bench(args = [20, 200], sample_count = 10)] +fn listing_warm(bencher: Bencher, turns: usize) { + let (_keep, dir) = session_dir(turns); + let repo = SessionRepo::open(&dir).unwrap(); + // Prime the index once, outside the measured region. + black_box(repo.list_with_progress(|_, _| {}).unwrap().len()); + bencher.bench_local(|| { + let metas = repo.list_with_progress(|_, _| {}).unwrap(); + black_box(metas.len()); + }); +} + +/// The **floor** for any listing scan that still reads the file: open it, pull every line, touch +/// nothing else. Whatever `listing` costs above this number is parsing; whatever it costs *at* this +/// number is unavoidable as long as the bytes are read at all. This is the control that decides +/// whether a faster parser is worth anything, or whether the only real fix is to stop reading. +#[divan::bench(args = [200], sample_count = 10)] +fn listing_read_floor(bencher: Bencher, turns: usize) { + use std::io::BufRead; + let (_keep, dir) = session_dir(turns); + let file = std::fs::read_dir(&dir) + .unwrap() + .map(|e| e.unwrap().path()) + .find(|p| p.extension().is_some_and(|x| x == "jsonl")) + .unwrap(); + bencher.bench_local(|| { + let mut reader = std::io::BufReader::new(std::fs::File::open(&file).unwrap()); + let mut buf = Vec::new(); + let mut lines = 0usize; + while reader.read_until(b'\n', &mut buf).unwrap() > 0 { + lines += 1; + buf.clear(); + } + black_box(lines); + }); +} + +// --------------------------------------------------------------------------------------------- +// The append path: does the serializer matter next to `sync_all`? +// --------------------------------------------------------------------------------------------- + +/// One turn appended to an already-long session — `append_new`'s own `messages.len() <= persisted` +/// guard means only the new messages are serialized, so this measures exactly one turn's marginal +/// durable cost: serialize ~64 KB + `write_all` + `flush` + `sync_all`, on a real disk. +#[divan::bench(sample_count = 20)] +fn append_new(bencher: Bencher) { + let dir = real_disk_tempdir(); + let repo = SessionRepo::open(dir.path()).unwrap(); + let mut store = repo + .create(SessionMeta::new("/repo", "claude-sonnet-5")) + .unwrap(); + let mut messages = Vec::new(); + for i in 0..20 { + messages.extend(turn(i)); + } + store.append_new(&messages).unwrap(); + let mut next = 20usize; + bencher.bench_local(|| { + messages.extend(turn(next)); + next += 1; + store.append_new(&messages).unwrap(); + }); +} + +// --------------------------------------------------------------------------------------------- +// The serializer in isolation, both strategies. `Entry` is private; this replicates its shape. +// --------------------------------------------------------------------------------------------- + +/// Structurally what `Entry::Message` is: the envelope fields plus a **cloned** `Message`. The clone +/// is the point — `session_store::append_new` clones the message into the `Entry` purely to serialize +/// it, so it's part of the cost being measured. +#[derive(Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum OwnedEntry { + Message { + id: Option, + parent_id: Option, + timestamp: u64, + message: Message, + }, +} + +/// The same envelope, borrowing the `Message` instead of cloning it — the fix under test. +#[derive(Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum BorrowedEntry<'a> { + Message { + id: Option<&'a str>, + parent_id: Option<&'a str>, + timestamp: u64, + message: &'a Message, + }, +} + +fn payload() -> Message { + Message::tool_result("toolu_000001", big_tool_result(), false) +} + +/// Today's `write_line`: clone the message into an owned `Entry`, serialize that to an owned `Value` +/// tree, then serialize the tree to a `String`. Three full passes over ~64 KB. +#[divan::bench] +fn serialize_entry_clone_to_value_to_string(bencher: Bencher) { + let msg = payload(); + bencher.bench_local(|| { + let entry = OwnedEntry::Message { + id: Some("abc".to_string()), + parent_id: Some("def".to_string()), + timestamp: 1_700_000_000, + message: msg.clone(), + }; + let v = serde_json::to_value(&entry).unwrap(); + let mut buf: Vec = Vec::new(); + use std::io::Write; + writeln!(&mut buf, "{}", Value::to_string(&v)).unwrap(); + black_box(buf.len()); + }); +} + +/// The fix: borrow the message, serialize straight to the writer. One pass, no `Value` tree, no clone. +#[divan::bench] +fn serialize_entry_borrow_to_writer(bencher: Bencher) { + let msg = payload(); + bencher.bench_local(|| { + let entry = BorrowedEntry::Message { + id: Some("abc"), + parent_id: Some("def"), + timestamp: 1_700_000_000, + message: &msg, + }; + let mut buf: Vec = Vec::new(); + serde_json::to_writer(&mut buf, &entry).unwrap(); + buf.push(b'\n'); + black_box(buf.len()); + }); +} + +/// Keep `SessionStore` in the import graph even if a future edit drops its only other use. +#[allow(dead_code)] +fn _assert_store_type(_: &SessionStore) {} diff --git a/crates/agent/benches/search.rs b/crates/agent/benches/search.rs index d9784b3..8692fad 100644 --- a/crates/agent/benches/search.rs +++ b/crates/agent/benches/search.rs @@ -169,6 +169,110 @@ fn edit_normalize(bencher: Bencher) { }); } +/// The same normalization *without* the offset map — the exact-match path's version. Sits next to +/// `edit_normalize` on purpose: the pair is an A/B in a single run, so both arms see the same machine +/// load. The gap between them is the entire prize for building the map lazily, and this box runs other +/// work, so a cross-invocation comparison would be measuring the neighbours. +#[divan::bench] +fn edit_normalize_only(bencher: Bencher) { + let src = edit_src(); + bencher.bench_local(|| { + let norm = edit::normalize_only(&src); + black_box(norm.len()); + }); +} + +/// The same normalization over a **pure-ASCII** file — which is what source code overwhelmingly is. +/// NFKC is the identity on ASCII and `fold_char` has nothing to fold, so everything the normalizer +/// does here except trailing-whitespace stripping is provably wasted. Compare against +/// `edit_normalize_only` (identical size, a handful of non-ASCII chars) to see what the Unicode +/// machinery costs when there is no Unicode. +#[divan::bench] +fn edit_normalize_ascii(bencher: Bencher) { + let mut src = String::with_capacity(READ_LINES * 56); + for i in 0..READ_LINES { + src.push_str(&format!(" let x_{i} = compute(i, {i}) + adjust();\n")); + } + bencher.bench_local(|| { + let norm = edit::normalize_only(&src); + black_box(norm.len()); + }); +} + +/// The realistic source file `edit_run_exact` edits — same shape `edit_normalize` normalizes. +fn edit_src() -> String { + let mut src = String::with_capacity(READ_LINES * 56); + for i in 0..READ_LINES { + if i % 20 == 0 { + src.push_str(&format!("let s = \u{201c}line {i}\u{201d}; \n")); + } else { + src.push_str(&format!(" let x_{i} = compute(i, {i}) + adjust();\n")); + } + } + src +} + +/// `edit`'s **whole** `run` over a pure-ASCII source file — the realistic common case twice over: the +/// file is ASCII (source code overwhelmingly is) and the `old_string` matches exactly (the model +/// reproduces text it just read). This is the number that represents what an `edit` tool call actually +/// costs in practice. +#[divan::bench(sample_count = 50)] +fn edit_run_exact_ascii(bencher: Bencher) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("subject.rs"); + let mut src = String::with_capacity(READ_LINES * 56); + for i in 0..READ_LINES { + src.push_str(&format!(" let x_{i} = compute(i, {i}) + adjust();\n")); + } + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let tool = edit::Edit::new(dir.path()); + let p = path.to_str().unwrap().to_string(); + let mut n = 0usize; + bencher.bench_local(|| { + std::fs::write(&path, &src).unwrap(); + n += 1; + let out = rt + .block_on(tool.run(serde_json::json!({ + "path": p, + "old_string": " let x_1501 = compute(i, 1501) + adjust();", + "new_string": format!(" let x_1501 = compute(i, 1501) + adjust(); // {n}"), + }))) + .unwrap(); + black_box(out.text.len()); + }); +} + +/// The same `run`, but over a file carrying a few non-ASCII characters, so it takes the general +/// Unicode normalizer. Kept beside `edit_run_exact_ascii` as the guard rail: the ASCII fast path must +/// not have *regressed* the path that still needs full NFKC. +#[divan::bench(sample_count = 50)] +fn edit_run_exact_unicode(bencher: Bencher) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("subject.rs"); + let src = edit_src(); + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let tool = edit::Edit::new(dir.path()); + let p = path.to_str().unwrap().to_string(); + let mut n = 0usize; + bencher.bench_local(|| { + // Rewrite the subject each iteration so every sample edits the same pristine input. + std::fs::write(&path, &src).unwrap(); + n += 1; + let out = rt + .block_on(tool.run(serde_json::json!({ + "path": p, + "old_string": " let x_1501 = compute(i, 1501) + adjust();", + "new_string": format!(" let x_1501 = compute(i, 1501) + adjust(); // {n}"), + }))) + .unwrap(); + black_box(out.text.len()); + }); +} + /// A directory of 500 subdirectories, for the `ls` bench — directory entries are exactly where the old /// `format!("{name}/")` allocated a second String per entry (and dropped the first). fn dir_of_subdirs() -> &'static PathBuf { diff --git a/crates/agent/src/session_store.rs b/crates/agent/src/session_store.rs index caaf235..12b1f8d 100644 --- a/crates/agent/src/session_store.rs +++ b/crates/agent/src/session_store.rs @@ -573,7 +573,14 @@ enum NodeContent { #[derive(Clone)] struct Node { parent_id: Option, - content: NodeContent, + /// Behind an `Arc` because a node's content is immutable once inserted (nothing here ever mutates + /// an existing entry — the file is append-only, and so is this index), and one message's content + /// can legitimately belong to two nodes at once: [`SessionStore::rewrite_compacted`] deliberately + /// keeps the pre-compaction chain intact (its provenance guarantee) *and* re-emits the surviving + /// suffix under fresh ids on the new active chain, so every kept message would otherwise be held + /// twice in memory — once per compaction round, forever. Sharing makes that duplication cost a + /// refcount instead of a full [`Message`] clone. + content: Arc, /// Unix seconds this entry was actually appended, from [`Entry::Message`]/[`Entry::Custom`]'s own /// `timestamp` field — `0` for a node with no recorded timestamp (a legacy file written before this /// field existed, or a branch-summary-materialized node, which doesn't carry one). See @@ -591,7 +598,7 @@ impl Node { /// pi's own `buildSessionContext` skipping `"custom"`-typed entries. See [`Entry::Custom`]'s doc /// comment. fn as_message(&self) -> Option<&Message> { - match &self.content { + match &*self.content { NodeContent::Message(m) => Some(m), NodeContent::Custom { .. } => None, } @@ -876,7 +883,7 @@ impl SessionStore { id.clone(), Node { parent_id, - content: NodeContent::Message(message), + content: Arc::new(NodeContent::Message(message)), timestamp, }, ); @@ -898,7 +905,9 @@ impl SessionStore { id.clone(), Node { parent_id, - content: NodeContent::Message(branch_summary_message(&summary)), + content: Arc::new(NodeContent::Message(branch_summary_message( + &summary, + ))), // Track L45 (pi-parity fix): `Entry::BranchSummary` now carries its own // `timestamp`, same as `Entry::Message`/`Entry::Custom` — `0` here only for // a legacy file written before this field existed (`#[serde(default)]`), @@ -931,7 +940,7 @@ impl SessionStore { id.clone(), Node { parent_id, - content: NodeContent::Custom { kind, data }, + content: Arc::new(NodeContent::Custom { kind, data }), timestamp, }, ); @@ -1195,20 +1204,23 @@ impl SessionStore { for msg in &messages[self.persisted..] { let id = new_id(); let timestamp = now_secs(); - write_line( + // Serialized through `EntryRef`, which borrows `msg` — the `Node` below owns the one copy + // this loop legitimately needs (the in-memory tree keeps it), and the serializer no longer + // takes a second one it would drop on the next line. + write_line_ref( &mut buf, - &Entry::Message { - id: Some(id.clone()), - parent_id: parent.clone(), + &EntryRef::Message { + id: Some(&id), + parent_id: parent.as_deref(), timestamp, - message: msg.clone(), + message: msg, }, )?; staged.push(( id.clone(), Node { parent_id: parent.clone(), - content: NodeContent::Message(msg.clone()), + content: Arc::new(NodeContent::Message(msg.clone())), timestamp, }, )); @@ -1248,6 +1260,27 @@ impl SessionStore { kind: impl Into, data: serde_json::Value, ) -> std::io::Result { + // `data` is opaque and caller-supplied, and on the `serve` path the caller is a remote client + // (`{type:"append_custom", kind, data?}`). It is retained three times over — in the appended + // `Entry`, in `self.nodes`, and in `self.events` — and nothing downstream ever interprets or + // shrinks it, so an unbounded blob here is an unbounded write amplified threefold, on a single + // request. Everything else this store retains is at least proportional to conversation that + // actually happened; this is the one input a client can make arbitrarily large for free. + // + // Rejected outright rather than truncated: the entry is opaque by contract, so a *partial* + // `data` is not a degraded version of itself — it's corrupt, and this module has no way to know + // what that would break for whoever wrote it. The caller gets a real error and can decide. + let size = serde_json::to_vec(&data) + .map(|v| v.len()) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + if size > MAX_CUSTOM_ENTRY_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "custom entry `data` is {size} bytes, over the {MAX_CUSTOM_ENTRY_BYTES}-byte limit" + ), + )); + } let kind = kind.into(); let id = new_id(); let timestamp = now_secs(); @@ -1268,7 +1301,7 @@ impl SessionStore { id.clone(), Node { parent_id, - content: NodeContent::Custom { kind, data }, + content: Arc::new(NodeContent::Custom { kind, data }), timestamp, }, ); @@ -1344,7 +1377,7 @@ impl SessionStore { id.clone(), Node { parent_id: parent.clone(), - content: NodeContent::Message(m.clone()), + content: Arc::new(NodeContent::Message(m.clone())), timestamp: now_secs(), }, )); @@ -1386,7 +1419,7 @@ impl SessionStore { // `messages: &[Message]` admits nothing else. Round-tripping each back through its own // original `Entry` variant, not unconditionally `Entry::Message`, is what keeps a preserved // custom entry from being silently corrupted into an empty/default message by a rewrite. - match &node.content { + match &*node.content { NodeContent::Message(m) => write_line( &mut f, &Entry::Message { @@ -1569,7 +1602,7 @@ impl SessionStore { id.clone(), Node { parent_id: parent.clone(), - content: NodeContent::Message(summary.clone()), + content: Arc::new(NodeContent::Message(summary.clone())), timestamp: now_secs(), }, )); @@ -1581,14 +1614,23 @@ impl SessionStore { // in lockstep by construction — see this method's doc comment); a custom id clones its content // unchanged. for id in &kept_suffix_ids { - let content = match &self.nodes[id].content { - NodeContent::Message(_) => match rest.next() { - Some(m) => NodeContent::Message(m.clone()), + // The original node stays exactly where it is (that's the provenance guarantee above), so + // this new node's content is a second, identical copy of it — shared rather than cloned, + // since a node's content is immutable once inserted. Otherwise every compaction round would + // permanently double the memory held for every message that survived it. The equality check + // keeps the sharing an optimization and not a *trust*: `messages[1..]` is the kept suffix + // verbatim by construction, but if a caller ever broke that, its message still wins — the + // behavior is identical to cloning unconditionally, only cheaper. + let original = &self.nodes[id].content; + let content = match &**original { + NodeContent::Message(existing) => match rest.next() { + Some(m) if m == existing => Arc::clone(original), + Some(m) => Arc::new(NodeContent::Message(m.clone())), None => unreachable!( "kept_suffix_ids and messages[1..] must have the same message count" ), }, - custom @ NodeContent::Custom { .. } => custom.clone(), + NodeContent::Custom { .. } => Arc::clone(original), }; let new_node_id = new_id(); new_nodes.push(( @@ -1679,7 +1721,7 @@ impl SessionStore { // carried forward from the kept suffix — round-tripped through its own original `Entry` // variant, not unconditionally `Entry::Message`, mirroring `rewrite`'s own preserved-node // handling. - match &node.content { + match &*node.content { NodeContent::Message(m) => write_line( &mut buf, &Entry::Message { @@ -1926,7 +1968,7 @@ impl SessionStore { .nodes .iter() .map(|(id, node)| { - let (role, preview, entry_kind) = match &node.content { + let (role, preview, entry_kind) = match &*node.content { // A materialized branch-summary recap is a real `NodeContent::Message` (so it // actually reaches the model — see `branch_summary_message`), indistinguishable // from an ordinary message by content alone; `branch_summary_details` is exactly @@ -2119,7 +2161,7 @@ impl SessionStore { entry_id.clone(), Node { parent_id: Some(target_id.to_string()), - content: NodeContent::Message(branch_summary_message(&summary)), + content: Arc::new(NodeContent::Message(branch_summary_message(&summary))), timestamp, }, ); @@ -2221,7 +2263,7 @@ impl SessionStore { entry_id.clone(), Node { parent_id: None, - content: NodeContent::Message(branch_summary_message(&summary)), + content: Arc::new(NodeContent::Message(branch_summary_message(&summary))), timestamp, }, ); @@ -3437,39 +3479,279 @@ pub(crate) fn scan_session_dir(dir: &Path) -> Vec { .collect() } +/// The sidecar listing index: one per session directory, next to the `.jsonl` files it describes. +/// +/// A listing needs a handful of derived fields (preview, message count, `updated_at`, search text), and +/// the only way to recover them from a transcript is to read it end to end — `read_listing` is a +/// streaming parse, but it is still O(bytes), and a session file grows without bound as its +/// conversation does. Measured: listing a 13 MB session costs ~14 ms, of which only ~15% is reading the +/// bytes; the other ~85% is `serde_json` *tokenizing* payload it then throws away. That's why no parser +/// trick fixes this — skipping materialization doesn't skip tokenization (both were tried and both +/// bought 0% wall time). The only way to make the scan cheaper is to not do it. +/// +/// So each file's computed `SessionMeta` is cached here against the `(size, mtime)` it was computed +/// from. A session file is append-only, so any change to it moves both — a stale entry can't survive a +/// write. On a warm index a listing costs one `stat` per file instead of a full parse, which turns the +/// scan from O(total transcript bytes) into O(sessions). +/// +/// It is a **cache, and only a cache**: every read is best-effort, and any failure (missing, corrupt, +/// unknown version, unreadable, unwritable directory) simply falls back to scanning the file, which is +/// exactly what the code did before this existed. Nothing here is a source of truth, so a lost or +/// clobbered index costs time, never correctness — which is also why concurrent writers need no locking +/// beyond the atomic rename below. +const LISTING_INDEX_FILE: &str = ".listings.json"; + +/// Bumped whenever a change to the derived fields would make previously-cached entries wrong (a new +/// field, a different preview rule). An index at any other version is ignored wholesale and rebuilt, +/// so an upgrade can never serve stale-shaped metadata. +const LISTING_INDEX_VERSION: u32 = 1; + +#[derive(Serialize, Deserialize)] +struct ListingIndex { + version: u32, + /// Keyed by file *name*, not full path — the index lives in the directory it describes, so the + /// directory component is implied, and the file stays valid if the whole tree is moved. + entries: HashMap, +} + +/// One file's cached listing. +/// +/// The four derived fields are carried **explicitly** rather than riding along inside `meta`, because +/// `SessionMeta` marks every one of them `#[serde(skip)]` — deliberately, so they can never go stale +/// inside the on-disk *header* (see their doc comments). That is the right call there and exactly the +/// wrong one here: serializing a `SessionMeta` into this index would silently drop the only fields the +/// index exists to remember, and hand back a listing with `updated_at: 0` and no preview. Naming them +/// here keeps the header's invariant intact while letting the cache do its job, and means adding a fifth +/// derived field is a compile error in this struct rather than a silent hole in the cache. +#[derive(Clone, Serialize, Deserialize)] +struct ListingIndexEntry { + /// The `(size, mtime)` the cached listing was computed from. Both must match for it to be used. + size: u64, + mtime_ns: u128, + /// The persisted header fields (everything `SessionMeta` does serialize). + meta: SessionMeta, + updated_at: u64, + message_count: usize, + preview: Option, + search_text: String, +} + +impl ListingIndexEntry { + fn new(size: u64, mtime_ns: u128, meta: &SessionMeta) -> Self { + Self { + size, + mtime_ns, + meta: meta.clone(), + updated_at: meta.updated_at, + message_count: meta.message_count, + preview: meta.preview.clone(), + search_text: meta.search_text.clone(), + } + } + + /// Reassemble the full `SessionMeta` a scan would have produced — header fields from `meta`, derived + /// fields from the ones stored beside it. + fn to_meta(&self) -> SessionMeta { + SessionMeta { + updated_at: self.updated_at, + message_count: self.message_count, + preview: self.preview.clone(), + search_text: self.search_text.clone(), + ..self.meta.clone() + } + } +} + +/// A file's cache validity stamp. Nanosecond mtime rather than the whole-second [`mtime_secs`] used for +/// display: two appends inside the same second are ordinary during a live session, and a +/// second-resolution stamp would happily serve a stale listing for one of them. Size is carried too, so +/// even a filesystem with a coarse clock still invalidates on any append. +fn file_stamp(path: &Path) -> Option<(u64, u128)> { + let m = fs::metadata(path).ok()?; + let mtime = m.modified().ok()?.duration_since(UNIX_EPOCH).ok()?; + Some((m.len(), mtime.as_nanos())) +} + +/// Best-effort load. Any problem at all yields an empty index and a full rescan. +fn load_listing_index(dir: &Path) -> HashMap { + let raw = match fs::read_to_string(dir.join(LISTING_INDEX_FILE)) { + Ok(raw) => raw, + Err(_) => return HashMap::new(), + }; + match serde_json::from_str::(&raw) { + Ok(idx) if idx.version == LISTING_INDEX_VERSION => idx.entries, + _ => HashMap::new(), + } +} + +/// Write the index atomically (temp file + rename), so a reader never observes a half-written one and a +/// crash mid-write leaves the previous index intact rather than a corrupt one. No `fsync`: this is a +/// cache, and paying durability costs to protect data we can always recompute would defeat the purpose. +/// Concurrent writers race to rename and the last one wins — both wrote a valid index, so either is +/// correct. +/// +/// Every error is swallowed deliberately. A read-only or unwritable session directory must still *list*; +/// it just doesn't get to keep a cache. +fn store_listing_index(dir: &Path, entries: &HashMap) { + let idx = ListingIndex { + version: LISTING_INDEX_VERSION, + entries: entries.clone(), + }; + let Ok(bytes) = serde_json::to_vec(&idx) else { + return; + }; + // Unique per process so two concurrent listers can't tear each other's temp file. + let tmp = dir.join(format!(".listings.{}.tmp", std::process::id())); + let write = (|| -> std::io::Result<()> { + let mut f = create_private(&tmp)?; + f.write_all(&bytes)?; + f.flush()?; + drop(f); + fs::rename(&tmp, dir.join(LISTING_INDEX_FILE)) + })(); + if write.is_err() { + let _ = fs::remove_file(&tmp); + } +} + /// Scan every path in `paths` for its listing metadata, calling `on_progress(scanned, total)` once per /// path — including ones that turn out unreadable, which still count as scanned and just contribute -/// nothing. `read_listing` is pure disk I/O plus parsing with no dependency between files, so the work -/// fans out across a small worker pool (`std::thread::available_parallelism`, capped at one thread per -/// path) rather than running strictly one file at a time; below two candidate workers it just runs -/// inline; no thread pool to justify the setup cost for a one- or two-file listing. Returned in -/// arbitrary order — every caller sorts the result itself. +/// nothing. +/// +/// Files whose `(size, mtime)` still match the sidecar [`ListingIndex`] are served from it without being +/// opened; only the rest are actually parsed. Those misses are pure disk I/O plus parsing with no +/// dependency between files, so they fan out across a small worker pool +/// (`std::thread::available_parallelism`, capped at one thread per file) rather than running strictly one +/// at a time; below two candidate workers it just runs inline, there being no thread pool worth its setup +/// cost for a one- or two-file scan. Returned in arbitrary order — every caller sorts the result itself. pub(crate) fn scan_listings( paths: Vec, on_progress: &(impl Fn(usize, usize) + Send + Sync), ) -> Vec { let total = paths.len(); + if total == 0 { + return Vec::new(); + } + + // One index per directory — `list_all` spans many project directories in a single flat scan, and + // each keeps its own sidecar next to its own files. + let mut indexes: HashMap> = HashMap::new(); + for p in &paths { + if let Some(dir) = p.parent() { + if !indexes.contains_key(dir) { + indexes.insert(dir.to_path_buf(), load_listing_index(dir)); + } + } + } + + // Split into cache hits (served from the index) and misses (must be parsed). A hit still counts as + // scanned, so `on_progress` stays monotonic and reaches `total` exactly as before. + let mut hits: Vec = Vec::new(); + let mut misses: Vec = Vec::new(); + // The stamp each miss was read at, so the rebuilt index records what it was *actually* computed + // from. Stamping after the parse instead would race an append landing mid-scan and cache a listing + // against a file state it never saw. + let mut miss_stamps: HashMap = HashMap::new(); + let mut fresh: HashMap> = HashMap::new(); + let mut scanned = 0usize; + + for p in &paths { + let stamp = file_stamp(p); + let cached = (|| { + let (size, mtime_ns) = stamp?; + let dir = p.parent()?; + let name = p.file_name()?.to_str()?; + let e = indexes.get(dir)?.get(name)?; + (e.size == size && e.mtime_ns == mtime_ns).then(|| e.to_meta()) + })(); + match cached { + Some(meta) => { + if let (Some(dir), Some(name), Some((size, mtime_ns))) = + (p.parent(), p.file_name().and_then(|n| n.to_str()), stamp) + { + fresh + .entry(dir.to_path_buf()) + .or_default() + .insert(name.to_string(), ListingIndexEntry::new(size, mtime_ns, &meta)); + } + hits.push(meta); + scanned += 1; + on_progress(scanned, total); + } + None => { + if let Some(s) = stamp { + miss_stamps.insert(p.clone(), s); + } + misses.push(p.clone()); + } + } + } + + let mut metas = hits; + metas.reserve(misses.len()); + for (path, meta) in scan_uncached(&misses, scanned, total, on_progress) { + // Only a file we managed to stamp *before* reading can be cached — otherwise there's nothing to + // validate a future hit against. + if let (Some(dir), Some(name), Some(&(size, mtime_ns))) = ( + path.parent(), + path.file_name().and_then(|n| n.to_str()), + miss_stamps.get(&path), + ) { + fresh + .entry(dir.to_path_buf()) + .or_default() + .insert(name.to_string(), ListingIndexEntry::new(size, mtime_ns, &meta)); + } + metas.push(meta); + } + + // Rebuilt from what this scan actually saw, so a deleted session's entry is dropped rather than + // accumulating forever. Only rewritten when it would differ from what's already on disk — a listing + // that was a pure cache hit does no I/O at all. + for (dir, entries) in fresh { + let unchanged = indexes + .get(&dir) + .is_some_and(|old| old.len() == entries.len() && misses.is_empty()); + if !unchanged { + store_listing_index(&dir, &entries); + } + } + metas +} + +/// The parse-every-byte path, for the files the index couldn't answer for. `base` is how many paths were +/// already served from cache, so `on_progress` keeps counting up to `total` across both halves. +fn scan_uncached( + paths: &[PathBuf], + base: usize, + total: usize, + on_progress: &(impl Fn(usize, usize) + Send + Sync), +) -> Vec<(PathBuf, SessionMeta)> { + let total_uncached = paths.len(); + if total_uncached == 0 { + return Vec::new(); + } let workers = std::thread::available_parallelism() .map(std::num::NonZero::get) .unwrap_or(1) - .min(total); + .min(total_uncached); if workers < 2 { return paths .iter() .enumerate() .filter_map(|(i, path)| { let meta = read_listing(path); - on_progress(i + 1, total); - meta + on_progress(base + i + 1, total); + meta.map(|m| (path.clone(), m)) }) .collect(); } let scanned = AtomicUsize::new(0); - let metas = Mutex::new(Vec::with_capacity(total)); + let metas = Mutex::new(Vec::with_capacity(total_uncached)); let scanned_ref = &scanned; let metas_ref = &metas; - let chunk_size = total.div_ceil(workers); + let chunk_size = total_uncached.div_ceil(workers); std::thread::scope(|scope| { for chunk in paths.chunks(chunk_size) { scope.spawn(move || { @@ -3478,9 +3760,12 @@ pub(crate) fn scan_listings( metas_ref .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(meta); + .push((path.clone(), meta)); } - on_progress(scanned_ref.fetch_add(1, Ordering::Relaxed) + 1, total); + on_progress( + base + scanned_ref.fetch_add(1, Ordering::Relaxed) + 1, + total, + ); } }); } @@ -3753,6 +4038,16 @@ fn parse_entry_lenient(path: &Path, line: &str) -> Result std::io::Result<()> { } /// Serialize one entry as a single JSON line (no embedded newlines — `serde_json` escapes them). +/// +/// Straight to the writer, in one pass. It used to go `Entry` → owned `Value` tree → `String` → `w`, +/// which is the same two-pass shape `serve::event_frame` was fixed out of and `benches/serve_events.rs` +/// was written to condemn — the fix just never reached this module. An entry carrying a 100 KB tool +/// result built a 100 KB `Value` tree solely to throw it away. See `benches/persistence.rs`. fn write_line(w: &mut impl Write, entry: &Entry) -> std::io::Result<()> { - let v = serde_json::to_value(entry).map_err(std::io::Error::other)?; - // Defensive: a serialized line must be one physical line. - debug_assert!(!serde_json::to_string(&v).unwrap_or_default().contains('\n')); - writeln!(w, "{}", Value::to_string(&v)) + serde_json::to_writer(&mut *w, entry).map_err(std::io::Error::other)?; + w.write_all(b"\n") +} + +/// The write-side mirror of [`Entry::Message`] that **borrows** its `Message` rather than owning one. +/// +/// Byte-identical on the wire (same internal tag, same field names, same `flatten`) — it exists purely +/// so [`SessionStore::append_new`] can persist a message without cloning it first. The in-memory `Node` +/// legitimately owns its copy; the serializer never needed one, so a turn whose tool result is a 100 KB +/// `read` output was copying those bytes once just to hand them to `serde`. +#[derive(Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum EntryRef<'a> { + Message { + id: Option<&'a str>, + parent_id: Option<&'a str>, + timestamp: u64, + #[serde(flatten)] + message: &'a Message, + }, +} + +/// [`write_line`]'s borrowing counterpart — same single-pass write, no owned `Entry` required. +fn write_line_ref(w: &mut impl Write, entry: &EntryRef<'_>) -> std::io::Result<()> { + serde_json::to_writer(&mut *w, entry).map_err(std::io::Error::other)?; + w.write_all(b"\n") } /// Append one entry to the session file — an O(1) write, not a rewrite. Same durability posture as @@ -8648,6 +8970,41 @@ mod tests { let _ = store.abandoned_by_switch("m1"); } + #[test] + fn an_oversized_custom_entry_is_refused_rather_than_written() { + // `data` is opaque and, on the `serve` path, remote-client-supplied — and it is retained three + // times over (the appended entry, `nodes`, `events`). It is the one input a client can make + // arbitrarily large for free, so it is the one that needs a ceiling. + let dir = tmpdir(); + let repo = SessionRepo::open(dir.path()).unwrap(); + let mut store = repo.create(SessionMeta::new("/w", "m")).unwrap(); + + let ok = store.append_custom("note", json!({ "v": "x".repeat(1024) })); + assert!(ok.is_ok(), "an ordinary payload must still be accepted"); + + let before = store.active_ids().len(); + let bytes_before = std::fs::metadata(&store.path).unwrap().len(); + + let huge = json!({ "v": "x".repeat(MAX_CUSTOM_ENTRY_BYTES + 1) }); + let err = store + .append_custom("note", huge) + .expect_err("a payload past the ceiling must be refused"); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); + + // Refused means *nothing happened*: the rejection must not half-append, advance the tip, or + // leave the blob behind in any of the three places it would otherwise be retained. + assert_eq!( + store.active_ids().len(), + before, + "a refused entry must not occupy a slot on the active chain" + ); + assert_eq!( + std::fs::metadata(&store.path).unwrap().len(), + bytes_before, + "a refused entry must not have reached the file either" + ); + } + #[test] fn append_custom_participates_in_tree_traversal_but_is_skipped_from_active_messages() { // pi: session-manager/save-entry.test.ts, "saves custom entries and includes them in tree diff --git a/crates/agent/src/tools/edit.rs b/crates/agent/src/tools/edit.rs index 30e12e0..a204126 100644 --- a/crates/agent/src/tools/edit.rs +++ b/crates/agent/src/tools/edit.rs @@ -2,6 +2,7 @@ //! first, then a normalized fuzzy fallback (NFKC + unified quotes/dashes/spaces + per-line //! trailing-whitespace) so a model's slightly-off `old_string` still lands. +use std::cell::OnceCell; use std::path::PathBuf; use agent_core::tool::Tool; @@ -69,12 +70,23 @@ impl Tool for Edit { .map(|p| super::canonical_write_target(&self.root, &p)) } + /// Argument parsing stays on the caller's task (it's a few `Value` lookups); everything from the + /// file read onward goes to `spawn_blocking`. + /// + /// This tool is the one that genuinely needed it. `serve_ws` runs each session on its own + /// `current_thread` runtime, so work done inline here pins that session's whole executor — its event + /// pump *and* its `abort`/`steer` command loop — for the duration. `read`/`grep`/`find`/`exec` + /// already hand their work off this way; `edit` did not, and on a multi-MB file it blocked the + /// runtime for tens of milliseconds (measured at ~73 ms on a 4 MB file — see + /// `tests/tool_reactor_stall.rs`, which now pins this invariant). `write` and `ls` were measured on + /// the same harness and don't come close to stalling, so they're deliberately left alone rather than + /// wrapped on principle. async fn run(&self, input: Value) -> Result { let path = input .get("path") .and_then(Value::as_str) .ok_or_else(|| ToolError::InvalidInput("missing `path`".into()))?; - let path = &self.resolve(path); + let path = self.resolve(path); let edits = parse_edits(&input)?; let replace_all = input .get("replace_all") @@ -86,6 +98,23 @@ impl Tool for Edit { )); } + match tokio::task::spawn_blocking(move || apply_edits(&path, edits, replace_all)).await { + Ok(result) => result, + // The closure returns its errors as `Err`, so a `JoinError` can only be a panic inside it. + // Re-raise rather than dressing it up as an ordinary tool failure. + Err(e) => std::panic::resume_unwind(e.into_panic()), + } + } +} + +/// The whole synchronous body of an `edit`: read, writability probe, match, splice, write. Lives on +/// `spawn_blocking`, never on the reactor — see [`Edit::run`]. +fn apply_edits( + path: &str, + edits: Vec<(String, String)>, + replace_all: bool, +) -> Result { + { let (raw, initial_mtime) = read_with_mtime(path)?; // Fail fast on a file this process can't actually write before spending any match/diff work on // it (fuzzy matching does NFKC normalization + ambiguity resolution — real CPU work, not free) @@ -117,7 +146,12 @@ impl Tool for Edit { // `find_spans`'s ambiguity check on every edit, fuzzy or not) is computed once here rather than // once per edit — a multi-edit call used to renormalize the whole file from scratch for every // element of `edits`, even though the input to that normalization was always identical. - let (norm_working, working_map) = normalize_with_map(&working); + // + // The *offset map* back to the original bytes, though, is only read when an edit actually falls + // through to the fuzzy path — so it's built lazily, at most once, and not at all on the exact + // path that virtually every call takes. See `normalize_only`. + let norm_working = normalize_only(&working); + let working_map: OnceCell> = OnceCell::new(); let total_edits = edits.len(); // `(start, end, replacement, edit_index)` — `edit_index` rides along purely so the overlap // check below can name which `edits[i]`/`edits[j]` collided (pi-parity fix: matching pi's own @@ -291,7 +325,7 @@ fn write_if_unchanged( fn find_spans( working: &str, norm_work: &str, - map: &[u32], + map: &OnceCell>, old: &str, replace_all: bool, edit_index: usize, @@ -336,6 +370,10 @@ fn find_spans( if norm_old.is_empty() { return Err(not_found_message(old, edit_index, total_edits)); } + // Only here — on a real fuzzy fallback — is the offset map actually needed, so only here is it + // built. Normalizing `working` a second time to produce it is deliberate: it keeps the eager, + // every-call path free of a table it never reads, and this branch is the rare one. + let map = map.get_or_init(|| normalize_with_map(working).1); let fuzzy: Vec<(usize, usize)> = norm_work .match_indices(&norm_old) .map(|(i, _)| (map[i] as usize, map[i + norm_old.len()] as usize)) @@ -420,10 +458,102 @@ fn fold_char(c: char) -> Option { /// (half the width of the old `usize`): `edit` reads the whole file into a `String` first, so a >4 GiB /// file can't reach here. pub fn normalize_with_map(orig: &str) -> (String, Vec) { + normalize_inner::(orig) +} + +/// [`normalize_with_map`] without the offset map — the same normalized text, and nothing else. +/// +/// The map is only ever read on the **fuzzy** path, to translate a match found in normalized space +/// back to the original bytes. The exact path (the overwhelmingly common one — the model reproduces an +/// `old_string` it just read verbatim) needs the normalized *text* alone, for `find_spans`'s ambiguity +/// count, and never touches the map. Building it there cost a `Vec` four bytes wide per output +/// byte — on the order of the file size, ×4 — plus a bounds-checked push per byte, for a table thrown +/// away unread. See `benches/search.rs::edit_run_exact`. +pub fn normalize_only(orig: &str) -> String { + normalize_inner::(orig).0 +} + +/// The shared one-pass normalizer. `MAP` is a const generic rather than a runtime flag so the map +/// pushes monomorphize away entirely when they aren't wanted — no branch per byte, no empty `Vec`. +/// +/// Dispatches to [`normalize_ascii`] when the input has no non-ASCII byte, which for source code is +/// almost always. See that function for why the two are equivalent there. +fn normalize_inner(orig: &str) -> (String, Vec) { + if orig.is_ascii() { + return normalize_ascii::(orig); + } + normalize_general::(orig) +} + +/// The ASCII fast path — a byte scan, no `char` decoding, no NFKC iterator, no per-scalar fold. +/// +/// It is *equivalent* to [`normalize_general`] on ASCII input, not an approximation of it, and that +/// rests on two facts: NFKC is the identity on ASCII (every ASCII scalar is NFKC-stable), and every +/// arm of [`fold_char`] matches a scalar at or above `U+00A0` — so neither transform can alter an +/// ASCII byte. All that remains of normalization is dropping the spaces/tabs immediately before a +/// `\n` or EOF, which is a byte-level operation. +/// +/// This matters because the general path costs the same on pure-ASCII input as on Unicode: the price +/// is the per-`char` machinery, not the content (`benches/search.rs` measures the two within a hair of +/// each other). Source files are overwhelmingly ASCII, so that machinery was being paid for Unicode +/// that isn't there. Bytes are copied in bulk runs between whitespace, rather than one at a time. +fn normalize_ascii(orig: &str) -> (String, Vec) { + let b = orig.as_bytes(); + let mut out = String::with_capacity(b.len()); + let mut map = Vec::::with_capacity(if MAP { b.len() + 1 } else { 0 }); + + // `cursor` is the first byte not yet emitted; `ws_start` marks the beginning of the whitespace run + // currently being held (trailing only if a `\n`/EOF proves it so — same rule as the general path). + let mut cursor = 0usize; + let mut ws_start: Option = None; + + let emit = |out: &mut String, map: &mut Vec, from: usize, to: usize| { + if from >= to { + return; + } + // ASCII, so any byte index is a char boundary and this slice is valid UTF-8 by construction. + out.push_str(&orig[from..to]); + if MAP { + map.extend((from..to).map(|j| j as u32)); + } + }; + + for (i, &c) in b.iter().enumerate() { + match c { + b' ' | b'\t' => { + if ws_start.is_none() { + ws_start = Some(i); + } + } + b'\n' => { + // The held run was trailing → emit everything before it and resume at the `\n`, which + // the next bulk copy picks up. + if let Some(ws) = ws_start.take() { + emit(&mut out, &mut map, cursor, ws); + cursor = i; + } + } + // Any other byte proves the held run interior; it stays, so there's nothing to cut. + _ => ws_start = None, + } + } + // A trailing whitespace run at EOF is dropped, exactly as the general path leaves `pending` + // unflushed. + let end = ws_start.unwrap_or(b.len()); + emit(&mut out, &mut map, cursor, end); + if MAP { + map.push(b.len() as u32); // sentinel + } + (out, map) +} + +/// The general, fully-Unicode normalizer: per-scalar NFKC, then [`fold_char`], then trailing-whitespace +/// stripping. Reached only when the input actually contains a non-ASCII byte. +fn normalize_general(orig: &str) -> (String, Vec) { use unicode_normalization::UnicodeNormalization; let mut out = String::with_capacity(orig.len()); - let mut map = Vec::::with_capacity(orig.len() + 1); + let mut map = Vec::::with_capacity(if MAP { orig.len() + 1 } else { 0 }); // A held run of candidate trailing whitespace: (byte — always ASCII ' '/'\t', source offset). let mut pending: Vec<(u8, u32)> = Vec::new(); let mut buf = [0u8; 4]; @@ -439,25 +569,33 @@ pub fn normalize_with_map(orig: &str) -> (String, Vec) { '\n' => { pending.clear(); out.push('\n'); - map.push(off); + if MAP { + map.push(off); + } } // Any other char: the held run was interior → flush it, then emit the char. _ => { for (b, o) in pending.drain(..) { out.push(b as char); - map.push(o); + if MAP { + map.push(o); + } } let s = folded.encode_utf8(&mut buf); out.push_str(s); - for _ in 0..s.len() { - map.push(off); + if MAP { + for _ in 0..s.len() { + map.push(off); + } } } } } } // A trailing whitespace run at EOF is dropped (`pending` left unflushed). - map.push(orig.len() as u32); // sentinel + if MAP { + map.push(orig.len() as u32); // sentinel + } (out, map) } @@ -555,6 +693,48 @@ mod tests { f } + /// The ASCII fast path is only sound if it is *equivalent* to the general Unicode path, not merely + /// close to it — it produces both the normalized text and the offset map that `run` splices the + /// original bytes back together with, so a divergence in either would corrupt a file rather than + /// just mismatch. Pin that equivalence directly, on the shapes that actually distinguish the two: + /// trailing whitespace before a newline, at EOF, on blank lines, and interior whitespace that must + /// survive. + #[test] + fn ascii_fast_path_is_equivalent_to_the_general_unicode_path() { + let cases = [ + "", + "\n", + " ", + "\t\t\n", + "plain\n", + "trailing \nnext\n", + "tabs\t\t\nnext", + "interior spaces stay\n", + "mixed \t \t\nand\t \tinterior x\n", + "no trailing newline ", + "\n\n\n \n\n", + "a \n b \n ", + "let x = 1; \n let y = 2;\t\n", + ]; + for case in cases { + assert!( + case.is_ascii(), + "case must be ASCII to exercise the fast path" + ); + let (fast_text, fast_map) = normalize_ascii::(case); + let (slow_text, slow_map) = normalize_general::(case); + assert_eq!(fast_text, slow_text, "text diverged for {case:?}"); + assert_eq!(fast_map, slow_map, "offset map diverged for {case:?}"); + + // And the map-less variant must agree with the mapped one on the text it produces. + assert_eq!( + normalize_ascii::(case).0, + fast_text, + "no-map text diverged for {case:?}" + ); + } + } + #[tokio::test] async fn replaces_unique_match() { let f = write_tmp("the quick brown fox"); diff --git a/crates/agent/tests/session_listing_index.rs b/crates/agent/tests/session_listing_index.rs new file mode 100644 index 0000000..aca45a9 --- /dev/null +++ b/crates/agent/tests/session_listing_index.rs @@ -0,0 +1,197 @@ +// Test target: `.unwrap()` asserts preconditions; that's the point. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +//! The sidecar listing index (`.listings.json`) is a pure cache in front of `read_listing`, which is +//! O(total transcript bytes) and so the one cost in this crate that grows without bound with a user's +//! history. Warm, a listing is one `stat` per file instead of a full parse. +//! +//! A cache that can serve a stale answer is worse than no cache, so these tests are about exactly one +//! question: **can the index ever disagree with a from-scratch scan?** Every case below therefore +//! asserts the cached result equals the uncached one — never merely that it "looks right". + +use std::path::Path; + +use agent_core::Message; +use beyond_ai_agent::session_store::{SessionMeta, SessionRepo}; +use tempfile::TempDir; + +const INDEX: &str = ".listings.json"; + +fn repo() -> (TempDir, SessionRepo) { + let dir = tempfile::tempdir().unwrap(); + let repo = SessionRepo::open(dir.path()).unwrap(); + (dir, repo) +} + +/// A listing computed with the index removed — the ground truth every cached answer is compared to. +fn uncached(dir: &Path, repo: &SessionRepo) -> Vec { + let _ = std::fs::remove_file(dir.join(INDEX)); + repo.list().unwrap() +} + +/// Compares the two listings as *sets*, keyed by id. Deliberately not order-sensitive: `list` sorts by +/// `updated_at`, and sessions written within the same second tie — `sort_by` is stable, so a tie is +/// broken by scan order, which is already arbitrary (the uncached scan is a parallel fan-in). Asserting +/// on that order would be asserting on nondeterminism that predates the index. +fn assert_same(a: &[SessionMeta], b: &[SessionMeta]) { + assert_eq!(a.len(), b.len(), "different session counts"); + let mut a = a.to_vec(); + let mut b = b.to_vec(); + a.sort_by(|x, y| x.id.cmp(&y.id)); + b.sort_by(|x, y| x.id.cmp(&y.id)); + for (x, y) in a.iter().zip(&b) { + assert_eq!(x.id, y.id); + assert_eq!(x.updated_at, y.updated_at, "updated_at diverged for {}", x.id); + assert_eq!( + x.message_count, y.message_count, + "message_count diverged for {}", + x.id + ); + assert_eq!(x.preview, y.preview, "preview diverged for {}", x.id); + assert_eq!(x.title, y.title, "title diverged for {}", x.id); + assert_eq!(x.model, y.model, "model diverged for {}", x.id); + } +} + +#[test] +fn a_warm_listing_matches_a_cold_one() { + let (dir, repo) = repo(); + for i in 0..3 { + let mut store = repo.create(SessionMeta::new("/repo", "claude-sonnet-5")).unwrap(); + store + .append_new(&[ + Message::user(format!("question number {i}")), + Message::assistant(vec![agent_core::ContentBlock::Text { + text: format!("answer number {i}"), + id: None, + phase: None, + }]), + ]) + .unwrap(); + } + + let cold = uncached(dir.path(), &repo); + let built = repo.list().unwrap(); // writes the index + assert!(dir.path().join(INDEX).exists(), "index was never written"); + let warm = repo.list().unwrap(); // served from it + + assert_same(&cold, &built); + assert_same(&cold, &warm); +} + +/// The invalidation that actually matters: a session that is appended to must not keep serving the +/// listing it had *before* the append. Size and mtime both move on an append, and either alone is +/// enough to miss. +#[test] +fn appending_to_a_session_invalidates_its_cached_listing() { + let (dir, repo) = repo(); + let mut store = repo.create(SessionMeta::new("/repo", "claude-sonnet-5")).unwrap(); + let mut messages = vec![Message::user("first")]; + store.append_new(&messages).unwrap(); + + let before = repo.list().unwrap(); // warms the index + assert_eq!(before[0].message_count, 1); + + messages.push(Message::assistant(vec![agent_core::ContentBlock::Text { + text: "second".into(), + id: None, + phase: None, + }])); + messages.push(Message::user("third")); + store.append_new(&messages).unwrap(); + + let after = repo.list().unwrap(); + assert_eq!( + after[0].message_count, 3, + "the index served a listing from before the append — stale cache" + ); + assert_same(&uncached(dir.path(), &repo), &repo.list().unwrap()); +} + +/// A truncated/garbage index must be ignored, not trusted and not fatal. +#[test] +fn a_corrupt_index_falls_back_to_a_full_scan() { + let (dir, repo) = repo(); + let mut store = repo.create(SessionMeta::new("/repo", "claude-sonnet-5")).unwrap(); + store.append_new(&[Message::user("hello")]).unwrap(); + let expected = repo.list().unwrap(); + + for garbage in ["", "{", "null", r#"{"version":999,"entries":{}}"#, "not json at all"] { + std::fs::write(dir.path().join(INDEX), garbage).unwrap(); + let got = repo.list().unwrap(); + assert_same(&expected, &got); + } +} + +/// An index written by an older build with different derived-field semantics must be discarded +/// wholesale rather than half-trusted — that's what the version field is for. +#[test] +fn an_index_at_an_unknown_version_is_discarded() { + let (dir, repo) = repo(); + let mut store = repo.create(SessionMeta::new("/repo", "claude-sonnet-5")).unwrap(); + store.append_new(&[Message::user("hello")]).unwrap(); + let expected = repo.list().unwrap(); + + // A well-formed index, correct in every way except that it claims a version we don't speak — and + // whose payload is deliberately wrong, so trusting it would be visible. + let raw = std::fs::read_to_string(dir.path().join(INDEX)).unwrap(); + let bumped = raw.replacen("\"version\":1", "\"version\":2", 1); + assert_ne!(raw, bumped, "expected to find the index version field"); + std::fs::write(dir.path().join(INDEX), bumped).unwrap(); + + assert_same(&expected, &repo.list().unwrap()); +} + +/// A deleted session must fall out of the index rather than lingering in it forever. +#[test] +fn a_deleted_session_is_dropped_from_the_index() { + let (dir, repo) = repo(); + let keep = repo.create(SessionMeta::new("/repo", "claude-sonnet-5")).unwrap(); + let mut keep_store = keep; + keep_store.append_new(&[Message::user("keep me")]).unwrap(); + + let mut gone = repo.create(SessionMeta::new("/repo", "claude-sonnet-5")).unwrap(); + gone.append_new(&[Message::user("delete me")]).unwrap(); + let gone_id = gone.meta().id.clone(); + let gone_path = gone.path().to_path_buf(); + + assert_eq!(repo.list().unwrap().len(), 2); // warms the index with both + drop(gone); + std::fs::remove_file(&gone_path).unwrap(); + + let after = repo.list().unwrap(); + assert_eq!(after.len(), 1, "a deleted session survived in the listing"); + assert!(!after.iter().any(|m| m.id == gone_id)); + + let raw = std::fs::read_to_string(dir.path().join(INDEX)).unwrap(); + assert!( + !raw.contains(&gone_id), + "the deleted session's entry is still in the index, which would grow without bound" + ); +} + +/// The index is a cache, so a directory it cannot be written into must still list correctly — just +/// without the speedup. A listing that fails because it couldn't save a cache would be a regression +/// invented by the optimization itself. +#[cfg(unix)] +#[test] +fn a_read_only_directory_still_lists() { + use std::os::unix::fs::PermissionsExt; + + let (dir, repo) = repo(); + let mut store = repo.create(SessionMeta::new("/repo", "claude-sonnet-5")).unwrap(); + store.append_new(&[Message::user("hello")]).unwrap(); + let expected = uncached(dir.path(), &repo); + + let perms = std::fs::metadata(dir.path()).unwrap().permissions(); + let mut ro = perms.clone(); + ro.set_mode(0o500); // r-x: readable, not writable + std::fs::set_permissions(dir.path(), ro).unwrap(); + + let got = repo.list(); + + // Restore before asserting, so a failure can't leave an undeletable temp dir behind. + std::fs::set_permissions(dir.path(), perms).unwrap(); + + assert_same(&expected, &got.expect("listing must not fail just because the cache can't be saved")); +} diff --git a/crates/agent/tests/tool_reactor_stall.rs b/crates/agent/tests/tool_reactor_stall.rs new file mode 100644 index 0000000..4928ed3 --- /dev/null +++ b/crates/agent/tests/tool_reactor_stall.rs @@ -0,0 +1,147 @@ +// Test target: `.unwrap()` asserts preconditions; that's the point. +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +//! A tool must not pin the async runtime for the duration of its file I/O and CPU work. +//! +//! `serve_ws` gives every session its own **`current_thread`** runtime, so a tool that does its work +//! inline — rather than handing it to `spawn_blocking` — stops that session's executor dead for as +//! long as it runs. Nothing else on that runtime gets polled: not the outbound event pump, and not the +//! stdin/WebSocket command loop that carries `abort` and `steer`. `serve.rs`'s `persist_blocking` +//! already moved `sync_all` off the reactor for exactly this reason, and its doc comment says so; the +//! file tools are the ones that hadn't followed. +//! +//! These tests pin that invariant where it's cheapest to check: a ticker task that wants to wake every +//! millisecond, running concurrently with one tool call on a `current_thread` runtime. If the tool +//! yields, the ticker keeps ticking and the worst gap stays small. If the tool blocks, the worst gap is +//! the whole tool call. +//! +//! The subject files are deliberately large (a few MB), so a blocking implementation stalls for tens of +//! milliseconds and the assertion has a wide, un-flaky margin between "yielded" and "blocked" — rather +//! than depending on a tight timing threshold. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use agent_core::Tool; +use beyond_ai_agent::tools::{edit, ls, write}; +use serde_json::json; + +/// The worst gap the ticker observed between its own wakeups while `work` ran. A yielding tool leaves +/// this near the tick interval; a blocking one leaves it at roughly the tool's own duration. +async fn worst_ticker_gap(work: F) -> Duration +where + F: std::future::Future, +{ + let worst = Arc::new(AtomicU64::new(0)); + let stop = Arc::new(AtomicBool::new(false)); + let (w, s) = (worst.clone(), stop.clone()); + let ticker = tokio::spawn(async move { + let mut last = Instant::now(); + while !s.load(Ordering::Relaxed) { + tokio::time::sleep(Duration::from_millis(1)).await; + let now = Instant::now(); + w.fetch_max( + now.duration_since(last).as_micros() as u64, + Ordering::Relaxed, + ); + last = now; + } + }); + // Let the ticker settle so its own startup isn't counted as a stall. + tokio::time::sleep(Duration::from_millis(20)).await; + worst.store(0, Ordering::Relaxed); + + work.await; + + stop.store(true, Ordering::Relaxed); + let _ = ticker.await; + Duration::from_micros(worst.load(Ordering::Relaxed)) +} + +/// Generous enough that ordinary scheduler jitter on a loaded CI box can't trip it, yet far below the +/// tens of milliseconds an inline implementation stalls for on these inputs. +const MAX_STALL: Duration = Duration::from_millis(15); + +fn big_ascii_source(lines: usize) -> String { + let mut s = String::with_capacity(lines * 56); + for i in 0..lines { + s.push_str(&format!(" let x_{i} = compute(i, {i}) + adjust();\n")); + } + s +} + +#[tokio::test(flavor = "current_thread")] +async fn edit_does_not_stall_the_runtime() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("subject.rs"); + // ~4 MB: an inline `edit` spends tens of ms here (read + normalize + match + splice + write). + let src = big_ascii_source(80_000); + std::fs::write(&path, &src).unwrap(); + + let tool = edit::Edit::new(dir.path()); + let p = path.to_str().unwrap().to_string(); + let gap = worst_ticker_gap(async { + tool.run(json!({ + "path": p, + "old_string": " let x_40001 = compute(i, 40001) + adjust();", + "new_string": " let x_40001 = compute(i, 40001) + adjust(); // edited", + })) + .await + .unwrap() + }) + .await; + + assert!( + gap < MAX_STALL, + "`edit` pinned the current_thread runtime for {gap:?} — nothing else on this session's \ + executor could be polled for that long, including its abort/steer command loop. Its file I/O \ + and matching belong on `spawn_blocking`, like `read`/`grep`/`find` already are." + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn write_does_not_stall_the_runtime() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("out.rs"); + let src = big_ascii_source(80_000); + + let tool = write::Write::new(dir.path()); + let p = path.to_str().unwrap().to_string(); + let gap = worst_ticker_gap(async { + tool.run(json!({ "path": p, "content": src })) + .await + .unwrap() + }) + .await; + + assert!( + gap < MAX_STALL, + "`write` pinned the current_thread runtime for {gap:?} — its file I/O belongs on \ + `spawn_blocking`." + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn ls_does_not_stall_the_runtime() { + let dir = tempfile::tempdir().unwrap(); + // `ls` stats every entry it collects; a wide directory is where that adds up. + for i in 0..5_000 { + std::fs::create_dir(dir.path().join(format!("sub_{i:05}"))).unwrap(); + } + + let d = dir.path().to_str().unwrap().to_string(); + let gap = worst_ticker_gap(async { + ls::Ls::default() + .run(json!({ "path": d, "limit": 5000 })) + .await + .unwrap() + }) + .await; + + assert!( + gap < MAX_STALL, + "`ls` pinned the current_thread runtime for {gap:?} — its `read_dir` + per-entry `metadata` \ + belong on `spawn_blocking`." + ); +}