fix(subagent): bound worker output so Fleet fanout cannot exhaust TUI memory

Fixes the #3882 report of codewhale-tui at ~15 GB during Fleet use. The
sub-agent loop bypassed the root engine's tool-output spillover, so one
multi-MB build log entered messages raw and was then multiplied: every
per-step checkpoint cloned the whole history, snapshots/projections
cloned each checkpoint, the full-transcript handle serialized messages
plus a second copy inside the embedded checkpoint, and subagents.v1.json
persisted it all.

Three bounds, all with the full output still recoverable from disk:

- Tool results pass through the shared spillover writer before entering
  messages (head + footer naming the on-disk path; a worker-events line
  records where output spilled). Unlike the root loop, oversized error
  output spills too - sub-agent errors are routinely full build logs.
- SubAgentCheckpoint keeps a byte-bounded message tail (256 KiB budget,
  always at least the newest message so continuability is preserved)
  and records omitted_messages; message_count stays the true total.
  Old persisted records load unchanged via serde defaults.
- The subagent_full_transcript handle keeps a 1 MiB bounded tail plus
  true counts, and embeds the checkpoint without duplicating its
  messages.

Regression tests simulate the fanout shape: four workers emitting
~2.4 MB outputs keep total resident state under 2 MB while every full
output remains readable from the spillover files.

Tests: cargo test -p codewhale-tui --bin codewhale-tui --locked subagent
Tests: cargo test -p codewhale-tui --bin codewhale-tui --locked fleet
This commit is contained in:
Hunter B
2026-07-01 20:53:43 -07:00
parent 8bd23dcc94
commit fa7c4b0553
2 changed files with 333 additions and 3 deletions
+144 -3
View File
@@ -42,6 +42,7 @@ use crate::tools::spec::{
use crate::tools::todo::SharedTodoList;
#[cfg(test)]
use crate::tools::todo::TodoList;
use crate::tools::truncate::{SPILLOVER_HEAD_BYTES, SPILLOVER_THRESHOLD_BYTES, maybe_spillover};
use crate::tui::app::ReasoningEffort;
use crate::utils::spawn_supervised;
use crate::worker_profile::{ModelRoute, ShellPolicy, ToolScope, WorkerRuntimeProfile};
@@ -116,6 +117,20 @@ const DEFAULT_STEP_API_TIMEOUT: Duration =
const COMPLETED_AGENT_RETENTION: Duration = Duration::from_secs(60 * 60);
const MAX_AGENT_WORKER_RECORDS: usize = 256;
const MAX_AGENT_WORKER_EVENTS_PER_RECORD: usize = 128;
/// Byte budget for the message tail retained in a [`SubAgentCheckpoint`]
/// (#3882). Checkpoints fire on every step of every worker and are cloned
/// into snapshots, projections, and `subagents.v1.json`; an unbounded
/// `messages` clone turns one large tool output into many resident copies
/// under Fleet fanout. The checkpoint keeps the most recent messages within
/// this budget (always at least the last one, so continuability is
/// preserved) and records how many older messages were omitted. Full tool
/// outputs remain recoverable from the spillover files on disk.
const SUBAGENT_CHECKPOINT_MESSAGE_BUDGET_BYTES: usize = 256 * 1024;
/// Byte budget for the message tail embedded in a `subagent_full_transcript`
/// handle (#3882). One handle is retained in memory per agent; the payload
/// keeps a bounded tail plus the true `message_count` so inspection stays
/// useful without pinning a whole unbounded transcript in RAM.
const SUBAGENT_TRANSCRIPT_MESSAGE_BUDGET_BYTES: usize = 1024 * 1024;
const SUBAGENT_STATE_SCHEMA_VERSION: u32 = 1;
const SUBAGENT_STATE_FILE: &str = "subagents.v1.json";
const SUBAGENT_WORKTREE_ROOT_DIR: &str = ".codewhale-worktrees";
@@ -1278,6 +1293,13 @@ struct AgentUsageBudgetScope {
}
/// Durable recovery point for an interrupted sub-agent session.
///
/// `messages` is a byte-bounded tail (#3882), not the full history:
/// checkpoints fire per step and are cloned into snapshots/persistence, so an
/// unbounded clone multiplies large tool outputs under Fleet fanout.
/// `message_count` records the true total and `omitted_messages` how many of
/// the oldest were dropped from this snapshot; spilled tool outputs remain on
/// disk under the spillover directory.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SubAgentCheckpoint {
pub checkpoint_id: String,
@@ -1290,6 +1312,15 @@ pub struct SubAgentCheckpoint {
pub created_at_ms: u64,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub messages: Vec<Message>,
/// Oldest messages omitted from `messages` to honor the checkpoint byte
/// budget. `0` for records written before v0.8.67 (serde default).
#[serde(default, skip_serializing_if = "is_zero")]
pub omitted_messages: usize,
}
#[allow(clippy::trivially_copy_pass_by_ref)]
fn is_zero(n: &usize) -> bool {
*n == 0
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -4332,6 +4363,18 @@ async fn insert_subagent_full_transcript_handle(
duration_ms: u64,
fork_context: bool,
) -> VarHandle {
// Byte-bound the retained transcript (#3882): the handle store keeps this
// payload resident per agent, and the checkpoint already carries its own
// bounded message tail — embedding it verbatim would duplicate that tail
// inside one payload. Keep checkpoint metadata, drop its messages, and
// record how much of the true history the bounded tail omits.
let (bounded_messages, omitted_messages) =
bounded_tail_messages(messages, SUBAGENT_TRANSCRIPT_MESSAGE_BUDGET_BYTES);
let checkpoint_meta = checkpoint.map(|checkpoint| SubAgentCheckpoint {
omitted_messages: checkpoint.message_count,
messages: Vec::new(),
..checkpoint.clone()
});
let payload = json!({
"kind": "subagent_full_transcript",
"agent_id": agent_id,
@@ -4343,13 +4386,96 @@ async fn insert_subagent_full_transcript_handle(
"steps_taken": steps_taken,
"duration_ms": duration_ms,
"assignment": assignment,
"checkpoint": checkpoint,
"messages": messages,
"checkpoint": checkpoint_meta,
"message_count": messages.len(),
"omitted_messages": omitted_messages,
"messages": bounded_messages,
});
let mut store = runtime.context.runtime.handle_store.lock().await;
store.insert_json(format!("agent:{agent_id}"), "full_transcript", payload)
}
/// Bound a sub-agent tool result before it enters `messages` (#3882).
///
/// The root engine applies spillover in `turn_loop.rs`; the sub-agent loop
/// bypassed it, so one multi-MB build log became many resident copies across
/// child messages, checkpoints, transcript handles, and persistence — the
/// Fleet fanout memory blow-up. Over-threshold content (successes AND
/// errors: sub-agent error output is routinely a full build log, so the root
/// loop's pass-errors-through rationale does not hold here) is written to the
/// shared spillover directory and replaced inline by a bounded head plus a
/// footer naming the on-disk path.
///
/// Returns the (possibly bounded) content and the spillover path when one was
/// written. Spillover write failures degrade to passing the original content
/// through, mirroring `apply_spillover`.
fn bound_subagent_tool_result(
agent_id: &str,
tool_id: &str,
content: String,
) -> (String, Option<PathBuf>) {
if content.len() <= SPILLOVER_THRESHOLD_BYTES {
return (content, None);
}
let spill_id = format!("sa_{agent_id}_{tool_id}");
match maybe_spillover(
&spill_id,
&content,
SPILLOVER_THRESHOLD_BYTES,
SPILLOVER_HEAD_BYTES,
) {
Ok(Some((head, path))) => {
let footer = format!(
"\n\n[Sub-agent tool output truncated: {head_kib} KiB of {total_kib} KiB shown. \
Full output saved to {path}. Use `read_file` on that path if you need the \
elided output.]",
head_kib = head.len() / 1024,
total_kib = content.len() / 1024,
path = path.display(),
);
(format!("{head}{footer}"), Some(path))
}
Ok(None) => (content, None),
Err(err) => {
tracing::warn!(
target: "subagent",
?err,
agent_id,
tool_id,
"sub-agent spillover write failed; passing original content through"
);
(content, None)
}
}
}
/// Rough serialized size of one message, used for checkpoint/transcript byte
/// budgets. Exact JSON size via serde; unserializable messages (should not
/// happen) count as 1 KiB so they still consume budget.
fn approximate_message_bytes(message: &Message) -> usize {
serde_json::to_string(message).map_or(1024, |s| s.len())
}
/// Keep the most recent messages whose combined approximate size fits
/// `budget_bytes`. Always keeps at least the final message (even if it alone
/// exceeds the budget) so a non-empty history stays continuable. Returns the
/// retained tail and how many older messages were omitted.
fn bounded_tail_messages(messages: &[Message], budget_bytes: usize) -> (Vec<Message>, usize) {
let mut kept_rev: Vec<Message> = Vec::new();
let mut used = 0usize;
for message in messages.iter().rev() {
let size = approximate_message_bytes(message);
if !kept_rev.is_empty() && used.saturating_add(size) > budget_bytes {
break;
}
used = used.saturating_add(size);
kept_rev.push(message.clone());
}
kept_rev.reverse();
let omitted = messages.len().saturating_sub(kept_rev.len());
(kept_rev, omitted)
}
fn build_subagent_checkpoint(
agent_id: &str,
reason: impl Into<String>,
@@ -4359,6 +4485,8 @@ fn build_subagent_checkpoint(
) -> SubAgentCheckpoint {
let created_at_ms = epoch_millis_now();
let checkpoint_id = format!("{agent_id}:step:{steps_taken}:ts:{created_at_ms}");
let (bounded_messages, omitted_messages) =
bounded_tail_messages(messages, SUBAGENT_CHECKPOINT_MESSAGE_BUDGET_BYTES);
SubAgentCheckpoint {
checkpoint_id: checkpoint_id.clone(),
agent_id: agent_id.to_string(),
@@ -4368,7 +4496,8 @@ fn build_subagent_checkpoint(
steps_taken,
message_count: messages.len(),
created_at_ms,
messages: messages.to_vec(),
messages: bounded_messages,
omitted_messages,
}
}
@@ -5143,6 +5272,18 @@ async fn run_subagent(
Err(_) => format!("Error: Tool {tool_name} timed out"),
};
let tool_ok = !result.starts_with("Error:");
let (result, spilled_to) = bound_subagent_tool_result(&agent_id, &tool_id, result);
if let Some(path) = spilled_to.as_ref() {
record_agent_progress(
runtime,
&agent_id,
format!(
"{}: tool '{tool_display_name}' output spilled to {}",
format_step_counter(steps, max_steps),
path.display()
),
);
}
record_agent_progress(
runtime,
&agent_id,
+189
View File
@@ -5215,3 +5215,192 @@ fn cleanup_due_gates_write_locked_cleanup_to_a_bounded_cadence() {
"a zero interval is always due"
);
}
// ── #3882: bounded sub-agent output under Fleet fanout ─────────────────────
/// Serialize-and-restore guard for the shared spillover test root, mirroring
/// the pattern in `tools::truncate::tests`.
fn with_spillover_root<F: FnOnce()>(root: &std::path::Path, f: F) {
let _guard = crate::tools::truncate::TEST_SPILLOVER_GUARD
.lock()
.unwrap_or_else(|err| err.into_inner());
let prior = crate::tools::truncate::set_test_spillover_root(Some(root.to_path_buf()));
struct Restore(Option<std::path::PathBuf>);
impl Drop for Restore {
fn drop(&mut self) {
crate::tools::truncate::set_test_spillover_root(self.0.take());
}
}
let _restore = Restore(prior);
f();
}
#[test]
fn bounded_tail_messages_keeps_recent_within_budget_and_counts_omitted() {
let messages: Vec<Message> = (0..10)
.map(|i| text_message("user", &format!("{i}:{}", "x".repeat(10_000))))
.collect();
let (kept, omitted) = bounded_tail_messages(&messages, 35_000);
assert!(!kept.is_empty());
assert_eq!(kept.len() + omitted, messages.len());
assert!(omitted > 0, "a 100 KB history must not fit a 35 KB budget");
// The tail is the most recent slice, in order.
let last_kept = message_text(kept.last().expect("tail non-empty"));
assert!(
last_kept.starts_with("9:"),
"kept tail must end at the newest message"
);
let total: usize = kept.iter().map(approximate_message_bytes).sum();
assert!(
total <= 35_000 + 11_000,
"kept tail exceeds budget by more than one message: {total}"
);
}
#[test]
fn bounded_tail_messages_always_keeps_the_final_message() {
let messages = vec![
text_message("user", &"a".repeat(50_000)),
text_message("assistant", &"b".repeat(50_000)),
];
let (kept, omitted) = bounded_tail_messages(&messages, 10);
assert_eq!(
kept.len(),
1,
"the newest message survives even over budget"
);
assert_eq!(omitted, 1);
assert!(message_text(&kept[0]).starts_with('b'));
}
#[test]
fn checkpoints_are_byte_bounded_under_fanout_scale_output() {
// Simulates the #3882 report shape: a worker whose tool results are
// multi-MB build logs. Without bounding, every per-step checkpoint clone
// carried the whole history; the persisted fleet file and every snapshot
// multiplied it further.
let huge = "error: expected `;`\n".repeat(120_000); // ~2.3 MB per message
let messages: Vec<Message> = (0..6).map(|_| text_message("user", &huge)).collect();
let checkpoint = make_checkpoint("fleet-worker-1", 6, messages.clone());
assert_eq!(checkpoint.message_count, messages.len());
assert!(checkpoint.omitted_messages > 0);
assert!(
!checkpoint.messages.is_empty(),
"checkpoint must stay continuable"
);
let serialized = serde_json::to_string(&checkpoint).expect("serialize checkpoint");
assert!(
serialized.len() <= SUBAGENT_CHECKPOINT_MESSAGE_BUDGET_BYTES + huge.len() + 64 * 1024,
"checkpoint JSON must be bounded, got {} bytes",
serialized.len()
);
// The raw history is ~14 MB; the checkpoint must not carry it.
assert!(
serialized.len() < 4 * 1024 * 1024,
"checkpoint JSON should be far below the raw transcript size, got {} bytes",
serialized.len()
);
}
#[test]
fn checkpoint_without_omitted_field_still_deserializes() {
// Records persisted before v0.8.67 carry no omitted_messages key.
let legacy = r#"{
"checkpoint_id": "a:step:1:ts:1",
"agent_id": "a",
"continuation_handle": "agent:a:checkpoint:a:step:1:ts:1",
"reason": "interrupted",
"continuable": true,
"steps_taken": 1,
"message_count": 1,
"created_at_ms": 1
}"#;
let checkpoint: SubAgentCheckpoint =
serde_json::from_str(legacy).expect("legacy checkpoint should load");
assert_eq!(checkpoint.omitted_messages, 0);
}
#[test]
fn subagent_tool_results_spill_to_disk_and_stay_bounded_inline() {
let tmp = tempdir().expect("tempdir");
with_spillover_root(tmp.path(), || {
let raw = "cargo build noise line\n".repeat(220_000); // ~5 MB
let raw_len = raw.len();
let (inline, spilled) =
bound_subagent_tool_result("fleet-worker-1", "call-42", raw.clone());
let path = spilled.expect("multi-MB output must spill");
// Model-visible content is bounded to head + footer.
assert!(inline.len() <= crate::tools::truncate::SPILLOVER_HEAD_BYTES + 1024);
assert!(inline.contains("Sub-agent tool output truncated"));
assert!(inline.contains(&path.display().to_string()));
assert!(inline.contains("read_file"));
// Full output remains recoverable from disk.
let on_disk = std::fs::read_to_string(&path).expect("spill file readable");
assert_eq!(on_disk.len(), raw_len);
// Small outputs pass through untouched, no spill file.
let (small, spilled) =
bound_subagent_tool_result("fleet-worker-1", "call-43", "ok".to_string());
assert_eq!(small, "ok");
assert!(spilled.is_none());
// Oversized error output is bounded too: sub-agent errors are
// routinely full build logs, unlike the root loop's short errors.
let (bounded_err, spilled) =
bound_subagent_tool_result("fleet-worker-1", "call-44", format!("Error: {raw}"));
assert!(spilled.is_some());
assert!(bounded_err.len() <= crate::tools::truncate::SPILLOVER_HEAD_BYTES + 1024);
assert!(bounded_err.starts_with("Error:"));
});
}
#[test]
fn fanout_of_workers_with_huge_outputs_keeps_resident_state_bounded() {
// Acceptance shape for #3882: multiple workers, each emitting multi-MB
// tool output. Model-visible content and per-worker checkpoints stay
// bounded while every full output is recoverable from disk.
let tmp = tempdir().expect("tempdir");
with_spillover_root(tmp.path(), || {
let huge = "warning: unused import `std::mem`\n".repeat(70_000); // ~2.4 MB
let mut resident_bytes = 0usize;
for worker in 0..4 {
let agent_id = format!("fleet-worker-{worker}");
let mut messages = Vec::new();
for call in 0..3 {
let (inline, spilled) =
bound_subagent_tool_result(&agent_id, &format!("call-{call}"), huge.clone());
let path = spilled.expect("should spill");
assert_eq!(
std::fs::read_to_string(&path).expect("readable").len(),
huge.len()
);
resident_bytes += inline.len();
messages.push(text_message("user", &inline));
}
let checkpoint = make_checkpoint(&agent_id, 3, messages);
let serialized = serde_json::to_string(&checkpoint).expect("serialize");
assert!(
serialized.len() <= SUBAGENT_CHECKPOINT_MESSAGE_BUDGET_BYTES + 128 * 1024,
"worker {worker} checkpoint too large: {} bytes",
serialized.len()
);
resident_bytes += serialized.len();
}
// 4 workers × 3 calls × ~2.4 MB ≈ 29 MB raw. Bounded resident state
// must stay under 2 MB total.
assert!(
resident_bytes < 2 * 1024 * 1024,
"resident bytes not bounded: {resident_bytes}"
);
});
}