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:
@@ -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,
|
||||
|
||||
@@ -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}"
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user