From c914d74ab5febfef20c32f6acada44c460edc6d1 Mon Sep 17 00:00:00 2001 From: Adam Dalloul <47503782+Adam-Dalloul@users.noreply.github.com> Date: Thu, 10 Sep 2026 08:12:02 -0700 Subject: [PATCH 1/2] fix(acp): bound the tool-call payload a session snapshot carries `active_tool_calls` is cleared in exactly one place, `TurnComplete`, so a turn that keeps working accumulates an entry for every tool call it has ever made, and `to_snapshot` copied all of them whole. Issue #380 sampled a turn that had not reached a `TurnComplete`: 1814 entries / 23.6 MB, then 1983 entries / 26.7 MB. Each entry holds up to the 64 KiB per-event output cap plus the agent's rendered content and any base64 image data, so the payload the attach path serves had no bound, and the desktop client that parses it on every attach stopped responding. Nothing could simply drop the finished entries: `denormalizeSnapshot` resolves each `ToolCallRef` in the live message through this list and skips a block whose id is missing, so a removed entry is a tool card missing from the middle of the in-flight turn. Every call still ships, in the same order, with its id, kind, label, status and meta. What is now bounded is the result payload (input / output / content / locations / images), and only on calls that already reached a terminal status: the newest keep everything until a 2 MiB budget is spent, older ones ship without it, and their results stay durable in the agent's transcript. Pending and in-progress calls are never trimmed at any size, since their partial output exists nowhere else and their count tracks live concurrency rather than turn length. Sizing reuses the escape-aware, allocation-free estimator the per-event cap already uses, so a payload means the same number of bytes on both paths. Measured on the reported shape (one turn, prose interleaved, ~12 KB of output and rendered content per call): 1983 calls went from 24.06 MiB to 2.62 MiB, 4000 calls from 48.53 MiB to 3.24 MiB, with the newest 170 calls keeping their full payload. --- src-tauri/src/acp/event_stream.rs | 46 ++-- src-tauri/src/acp/session_state.rs | 335 ++++++++++++++++++++++++++++- 2 files changed, 360 insertions(+), 21 deletions(-) diff --git a/src-tauri/src/acp/event_stream.rs b/src-tauri/src/acp/event_stream.rs index 3f67fde64a..a2fa14b619 100644 --- a/src-tauri/src/acp/event_stream.rs +++ b/src-tauri/src/acp/event_stream.rs @@ -199,7 +199,10 @@ impl RecentEventsBuffer { /// because this feeds the per-event size cap: an escape-heavy payload (tool /// output full of quotes/newlines, say) serializes much larger than its raw byte /// length and must still be recognized as oversized. -fn json_str_len(s: &str) -> usize { +/// +/// Shared with `SessionState::snapshot_tool_calls`, which spends a byte budget +/// over the same payloads and has to size them the same way. +pub(super) fn json_str_len(s: &str) -> usize { let mut extra = 0usize; for b in s.bytes() { match b { @@ -243,7 +246,7 @@ fn number_size(n: &serde_json::Number) -> usize { /// adds: brackets/braces, the `:` after each key, and the `,` BETWEEN elements. /// Computed without serializing and never undercounting, so it stays a safe /// proxy for the per-event size cap even for dense arrays/objects. -fn json_value_size(value: &serde_json::Value) -> usize { +pub(super) fn json_value_size(value: &serde_json::Value) -> usize { match value { serde_json::Value::Null => 4, serde_json::Value::Bool(b) => { @@ -271,11 +274,11 @@ fn json_value_size(value: &serde_json::Value) -> usize { } } -fn opt_str_size(s: &Option) -> usize { +pub(super) fn opt_str_size(s: &Option) -> usize { s.as_ref().map_or(0, |v| json_str_len(v)) } -fn opt_json_size(v: &Option) -> usize { +pub(super) fn opt_json_size(v: &Option) -> usize { v.as_ref().map_or(0, json_value_size) } @@ -290,20 +293,27 @@ fn opt_json_size(v: &Option) -> usize { /// allocation — far cheaper than the full-envelope `serde_json::to_vec` this /// replaced — and image events are infrequent (not on the per-token path). fn images_size(images: &Option>) -> usize { - images.as_ref().map_or(0, |imgs| { - // `PER_IMAGE_STRUCT` conservatively bounds each object's keys/braces and - // the trailing comma; `+ 2` is the array brackets. - const PER_IMAGE_STRUCT: usize = 48; - 2 + imgs - .iter() - .map(|img| { - PER_IMAGE_STRUCT - + json_str_len(&img.data) - + json_str_len(&img.mime_type) - + opt_str_size(&img.uri) - }) - .sum::() - }) + images + .as_ref() + .map_or(0, |imgs| images_slice_size(imgs.as_slice())) +} + +/// [`images_size`] over a plain slice. Split out so +/// `SessionState::snapshot_tool_calls`, whose `ToolCallState.images` is a bare +/// `Vec`, sizes an image list exactly the way the per-event cap does. +pub(super) fn images_slice_size(images: &[ToolCallImageInfo]) -> usize { + // `PER_IMAGE_STRUCT` conservatively bounds each object's keys/braces and + // the trailing comma; `+ 2` is the array brackets. + const PER_IMAGE_STRUCT: usize = 48; + 2 + images + .iter() + .map(|img| { + PER_IMAGE_STRUCT + + json_str_len(&img.data) + + json_str_len(&img.mime_type) + + opt_str_size(&img.uri) + }) + .sum::() } /// Byte size of a single user-message block, including its `{"type":..,..}` diff --git a/src-tauri/src/acp/session_state.rs b/src-tauri/src/acp/session_state.rs index 529d5b51aa..f4492d77d2 100644 --- a/src-tauri/src/acp/session_state.rs +++ b/src-tauri/src/acp/session_state.rs @@ -1,7 +1,7 @@ //! 会话级状态结构。后端权威：流式累积、in-flight tool calls、待处理 permission 等 //! 全部住在这里。Phase 2 的 snapshot 端点直接从此处读取 live 部分。 -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::path::PathBuf; use std::sync::Arc; @@ -9,7 +9,10 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use crate::acp::delegation::types::{BlockedKind, BlockedOn}; -use crate::acp::event_stream::{ConnectionEventStream, RecentEventsBuffer}; +use crate::acp::event_stream::{ + images_slice_size, json_str_len, json_value_size, opt_json_size, opt_str_size, + ConnectionEventStream, RecentEventsBuffer, +}; use crate::acp::feedback::{FeedbackItem, FeedbackStatus}; use crate::acp::plan_approval::PendingPlanApprovalState; use crate::acp::question::PendingQuestionState; @@ -1842,6 +1845,128 @@ impl SessionState { } } + /// Wire copy of `active_tool_calls`, with the bulky RESULT payload of + /// already-finished calls bounded by `MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES`. + /// + /// Nothing removes a completed call from `active_tool_calls`: the map is + /// cleared in one place, `TurnComplete`. So a single long agentic turn + /// accumulates one entry per tool call it has ever made, and the snapshot + /// used to carry all of them whole. Issue #380 sampled a turn that had not + /// reached `TurnComplete`: 1814 entries / 23.6 MB, then 1983 entries (1964 + /// completed, 18 failed, 1 running) / 26.7 MB. Each entry holds up to + /// `MAX_SINGLE_EMIT_BYTES` (64 KiB) of tool output plus the agent's + /// rendered `content` and, for image tools, base64 image data, so the + /// payload had no bound at all — and the desktop client that parses it on + /// every attach stopped responding. + /// + /// Every call still ships, in the same (id-sorted) order, with its id, + /// kind, label, status and meta. That is the part nothing else can supply: + /// `denormalizeSnapshot` resolves each `LiveContentBlock::ToolCallRef` in + /// `live_message.content` through this list and DROPS a block whose id is + /// missing, so an entry left out is a tool card missing from the middle of + /// the in-flight turn. + /// + /// What the budget bounds is `input` / `output` / `content` / `locations` / + /// `images`, and only on calls that already reached a terminal status: + /// + /// * Pending / InProgress calls are never trimmed, at any size. They are + /// the ones the attaching client has to keep rendering and revising from + /// live events, their partial output exists nowhere else, and their count + /// tracks live concurrency rather than turn length. + /// * Terminal calls keep everything, newest first, until the budget is + /// spent; the older ones then ship without those fields. A finished + /// call's result is durable in the agent's own transcript, which is what + /// the conversation reloads from. + /// + /// The budget is spent, not enforced per entry, so the last entry admitted + /// may carry the total past it by its own size (one oversized image tool + /// call). In-flight entries are counted against the budget but never + /// trimmed by it. + fn snapshot_tool_calls(&self) -> Vec { + // Arrival order comes from the live message: `push_tool_call_ref_if_absent` + // anchors exactly one `ToolCallRef` per call, in the order the agent + // opened them. `active_tool_calls` itself is keyed by id, which says + // nothing about age. + let mut arrival: BTreeMap<&str, usize> = BTreeMap::new(); + if let Some(live) = self.live_message.as_ref() { + for (i, block) in live.content.iter().enumerate() { + if let LiveContentBlock::ToolCallRef { tool_call_id } = block { + arrival.entry(tool_call_id.as_str()).or_insert(i); + } + } + } + + // (arrival rank, id, payload bytes). An id with no anchoring ref sorts + // newest, so the fail-safe direction is "keep everything" — today that + // cannot happen, because both `ToolCall` and `ToolCallUpdate` anchor a + // ref for the id they upsert. + let mut ordered: Vec<(usize, &str, usize)> = self + .active_tool_calls + .iter() + .map(|(id, tc)| { + ( + arrival.get(id.as_str()).copied().unwrap_or(usize::MAX), + id.as_str(), + tool_call_payload_bytes(tc), + ) + }) + .collect(); + + let total = ordered + .iter() + .fold(0usize, |acc, (_, _, bytes)| acc.saturating_add(*bytes)); + if total <= MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES { + // The ordinary turn: nothing to trim, wire shape byte-identical. + return self.active_tool_calls.values().cloned().collect(); + } + + ordered.sort_unstable(); + let mut trimmed: BTreeSet<&str> = BTreeSet::new(); + let mut spent = 0usize; + for (_, id, bytes) in ordered.iter().rev() { + let terminal = matches!( + self.active_tool_calls[*id].status, + ToolCallStatus::Completed | ToolCallStatus::Failed + ); + if terminal && spent >= MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES { + trimmed.insert(*id); + continue; + } + spent = spent.saturating_add(*bytes); + } + + self.active_tool_calls + .values() + .map(|tc| { + if !trimmed.contains(tc.id.as_str()) { + return tc.clone(); + } + // Listed field by field (no `..tc.clone()`) so a field added to + // `ToolCallState` later has to be classified here as identity + // or as payload, instead of silently riding an unbounded value + // back onto the wire. + ToolCallState { + id: tc.id.clone(), + kind: tc.kind.clone(), + label: tc.label.clone(), + status: tc.status.clone(), + input: None, + output: None, + content: None, + locations: None, + // Kept: the delegation broker writes the parent↔child + // binding here (`meta["codeg.delegation"]`), which is what + // re-anchors an inline sub-thread on a mid-turn attach. + // Bounded by contract — a small status object, not output. + meta: tc.meta.clone(), + images: Vec::new(), + // `#[serde(skip)]` — never on the wire either way. + raw_input_chunks: Vec::new(), + } + }) + .collect() + } + /// 拷贝出对外可见的 wire-friendly snapshot。Phase 2 snapshot 端点直接调用此方法。 pub fn to_snapshot(&self) -> LiveSessionSnapshot { LiveSessionSnapshot { @@ -1851,7 +1976,7 @@ impl SessionState { status: self.status.clone(), external_id: self.external_id.clone(), live_message: self.live_message.clone(), - active_tool_calls: self.active_tool_calls.values().cloned().collect(), + active_tool_calls: self.snapshot_tool_calls(), pending_permission: self.pending_permission.clone(), pending_question: self.pending_question.clone(), pending_plan_approval: self.pending_plan_approval.clone(), @@ -2018,6 +2143,37 @@ fn u32_is_zero(v: &u32) -> bool { *v == 0 } +/// Byte budget for the bulky tool-call payload one snapshot may carry. See +/// [`SessionState::snapshot_tool_calls`] for what it does and does not bound. +/// +/// 2 MiB leaves the newest ~150 finished calls of a typical read/edit/grep turn +/// intact (and, at the 64 KiB per-call ceiling `MAX_SINGLE_EMIT_BYTES` imposes, +/// at least the newest 32 in the worst case) — far more than a client attaching +/// mid-turn has on screen — while holding the #380 session's snapshot at ~2.6 MB +/// instead of 24 MB. +const MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES: usize = 2 * 1024 * 1024; + +/// The part of a `ToolCallState` that grows with what the tool actually did — +/// what [`MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES`] bounds. Excludes the identity fields +/// (id / kind / label / status / meta), which every entry keeps. +/// +/// Sized with the same escape-aware, allocation-free accounting the per-event +/// cap uses (`event_stream`), so "this call's payload" means the same number of +/// bytes on both paths. +fn tool_call_payload_bytes(tc: &ToolCallState) -> usize { + let output = match tc.output.as_ref() { + Some(ToolCallOutput::Text { content }) => json_str_len(content), + Some(ToolCallOutput::Error { message }) => json_str_len(message), + Some(ToolCallOutput::Json { value }) => json_value_size(value), + None => 0, + }; + opt_json_size(&tc.input) + .saturating_add(output) + .saturating_add(opt_str_size(&tc.content)) + .saturating_add(opt_json_size(&tc.locations)) + .saturating_add(images_slice_size(&tc.images)) +} + /// Last non-empty line of `s`, trimmed. `None` if every line is blank. fn last_nonempty_line(s: &str) -> Option<&str> { s.lines().map(str::trim).rev().find(|l| !l.is_empty()) @@ -3405,6 +3561,179 @@ mod tests { assert_eq!(ids, vec!["tc-a", "tc-m", "tc-z"]); } + /// Open a tool call and finish it, with `output_bytes` of tool output. + /// `settled` false leaves it in progress (still streaming its output). + fn run_tool_call(s: &mut SessionState, id: &str, output_bytes: usize, settled: bool) { + s.apply_event(&AcpEvent::ToolCall { + tool_call_id: id.into(), + title: format!("Read src/{id}.rs"), + kind: "read".into(), + status: "in_progress".into(), + content: None, + raw_input: Some(format!("{{\"file_path\":\"src/{id}.rs\"}}")), + raw_output: None, + locations: Some(serde_json::json!([{ "path": format!("src/{id}.rs") }])), + meta: Some(serde_json::json!({ "codeg.delegation": { "status": "completed" } })), + images: None, + }); + let status = if settled { "completed" } else { "in_progress" }; + s.apply_event(&AcpEvent::ToolCallUpdate { + tool_call_id: id.into(), + title: None, + status: Some(status.to_string()), + content: None, + raw_input: None, + raw_output: Some("o".repeat(output_bytes)), + raw_output_append: None, + locations: None, + meta: None, + images: None, + }); + } + + /// The ordinary turn stays byte-identical: nothing is trimmed while the + /// table fits the budget, so the wire shape is exactly what it always was. + #[test] + fn snapshot_carries_every_tool_call_whole_while_it_fits_the_budget() { + let mut s = fresh_state(); + for i in 0..20 { + run_tool_call(&mut s, &format!("tc-{i:03}"), 4 * 1024, true); + } + + let snap = s.to_snapshot(); + assert_eq!(snap.active_tool_calls.len(), 20); + for tc in &snap.active_tool_calls { + let live = &s.active_tool_calls[&tc.id]; + assert_eq!( + serde_json::to_value(tc).unwrap(), + serde_json::to_value(live).unwrap(), + "{} must ship exactly as held", + tc.id + ); + } + } + + /// #380: `active_tool_calls` is cleared only at `TurnComplete`, so a turn + /// that keeps working accumulates every tool call it ever made — the + /// reporter sampled 1983 entries and a 26.7 MB snapshot on a turn that had + /// not reached one, and the client that parses that on attach stopped + /// responding. + /// + /// The snapshot now spends a byte budget over the finished calls' payload, + /// newest first, WITHOUT dropping an entry: every `ToolCallRef` in the live + /// message must still resolve, or the reattaching client renders the + /// in-flight turn with tool cards missing from the middle of it. + #[test] + fn snapshot_bounds_finished_tool_call_payload_on_a_long_turn() { + const CALLS: usize = 300; + const OUTPUT_BYTES: usize = 16 * 1024; + let mut s = fresh_state(); + for i in 0..CALLS { + run_tool_call(&mut s, &format!("tc-{i:04}"), OUTPUT_BYTES, true); + } + // …and one still running, the call the attaching client has to keep + // rendering from live events. + run_tool_call(&mut s, "tc-running", OUTPUT_BYTES, false); + + // What the state holds, and what shipping it whole would have cost. + assert_eq!(s.active_tool_calls.len(), CALLS + 1); + let held: usize = s + .active_tool_calls + .values() + .map(tool_call_payload_bytes) + .sum(); + assert!( + held > 2 * MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES, + "the turn must hold well past the budget for this to test anything (held {held})" + ); + + let snap = s.to_snapshot(); + let wire = serde_json::to_string(&snap).expect("serialize snapshot"); + assert!( + wire.len() < MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES * 3 / 2, + "snapshot must stay near the budget, got {} bytes for {held} bytes held", + wire.len() + ); + + // Nothing is dropped: same count, same (id-sorted) order, and every + // `ToolCallRef` block in the live message still resolves. + assert_eq!(snap.active_tool_calls.len(), CALLS + 1); + let wire_ids: Vec<&str> = snap + .active_tool_calls + .iter() + .map(|tc| tc.id.as_str()) + .collect(); + let held_ids: Vec<&str> = s.active_tool_calls.keys().map(String::as_str).collect(); + assert_eq!(wire_ids, held_ids); + let by_id: std::collections::BTreeMap<&str, &ToolCallState> = snap + .active_tool_calls + .iter() + .map(|tc| (tc.id.as_str(), tc)) + .collect(); + for block in &s.live_message.as_ref().expect("live message").content { + if let LiveContentBlock::ToolCallRef { tool_call_id } = block { + assert!( + by_id.contains_key(tool_call_id.as_str()), + "{tool_call_id} is anchored in the live message but missing from the snapshot" + ); + } + } + + // The running call keeps its partial output at any size — nothing else + // has it. So does the newest finished one. + assert!(by_id["tc-running"].output.is_some()); + assert!(by_id[format!("tc-{:04}", CALLS - 1).as_str()].output.is_some()); + + // The oldest finished calls ship without the payload, but keep every + // field that identifies the card. + let oldest = by_id["tc-0000"]; + assert!(oldest.output.is_none(), "oldest call must shed its output"); + assert!(oldest.content.is_none()); + assert!(oldest.input.is_none()); + assert!(oldest.locations.is_none()); + assert!(oldest.images.is_empty()); + assert_eq!(oldest.label, "Read src/tc-0000.rs"); + assert_eq!(oldest.kind, ToolKind::Read); + assert_eq!(oldest.status, ToolCallStatus::Completed); + assert!( + oldest.meta.is_some(), + "delegation meta re-anchors an inline sub-thread on attach" + ); + + // And the trimming is a tail, not a purge: the budget is actually spent + // on the recent calls rather than thrown away. + let kept = snap + .active_tool_calls + .iter() + .filter(|tc| tc.output.is_some()) + .count(); + assert!( + kept > 32 && kept < CALLS, + "expected a bounded recent window to keep its output, got {kept}" + ); + } + + /// The budget bounds finished work only. A single running call bigger than + /// the whole budget still ships whole: its output exists nowhere else yet, + /// and its count tracks live concurrency, not how long the turn has run. + #[test] + fn snapshot_never_trims_a_running_tool_call() { + let mut s = fresh_state(); + for i in 0..200 { + run_tool_call(&mut s, &format!("tc-{i:04}"), 16 * 1024, true); + } + run_tool_call(&mut s, "tc-huge", MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES + 1024, false); + + let snap = s.to_snapshot(); + let huge = snap + .active_tool_calls + .iter() + .find(|tc| tc.id == "tc-huge") + .expect("running call present"); + assert_eq!(huge.status, ToolCallStatus::InProgress); + assert!(huge.output.is_some(), "a running call is never trimmed"); + } + #[test] fn tool_call_content_field_is_preserved_on_state() { let mut s = fresh_state(); From 8d983aac50ca6db458c9ae182deb80650b277361 Mon Sep 17 00:00:00 2001 From: xintaofei Date: Tue, 22 Sep 2026 00:02:33 +0800 Subject: [PATCH 2/2] fix(acp): keep a generated image out of the snapshot's trim MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Shedding a finished call's `images` does not show the card with less in it, it inverts what the card says. `isImageGenerationToolCall` classifies the call from the `label` the trim keeps, and `generated-images-block` renders an image-generation block whose `image` is null under a terminal status as "image generation failed" — so a trimmed success reported a failure, for the one payload the user would then go looking for. Image bytes are carried whole again, which is what the snapshot does today, so an image-heavy turn costs what it always did. They are still counted against the budget, so they push older result payload out first, and a trimmed entry's retained bytes now stay in the tally rather than leaving it — otherwise the budget measured something the wire does not carry. Also states the two things the bound rests on: that the agent reports a terminal status (pinned by a new test, since the carve-out for in-flight calls is what makes an agent that never did unbounded), and that a trimmed entry loses the input-shape-only identity signals — codex collab capsules and Kimi TodoList writes fall back to a generic card until the conversation reloads. --- src-tauri/src/acp/session_state.rs | 176 ++++++++++++++++++++++++++--- 1 file changed, 158 insertions(+), 18 deletions(-) diff --git a/src-tauri/src/acp/session_state.rs b/src-tauri/src/acp/session_state.rs index f4492d77d2..61f734692a 100644 --- a/src-tauri/src/acp/session_state.rs +++ b/src-tauri/src/acp/session_state.rs @@ -1866,8 +1866,8 @@ impl SessionState { /// missing, so an entry left out is a tool card missing from the middle of /// the in-flight turn. /// - /// What the budget bounds is `input` / `output` / `content` / `locations` / - /// `images`, and only on calls that already reached a terminal status: + /// What the budget bounds is `input` / `output` / `content` / `locations`, + /// and only on calls that already reached a terminal status: /// /// * Pending / InProgress calls are never trimmed, at any size. They are /// the ones the attaching client has to keep rendering and revising from @@ -1878,10 +1878,44 @@ impl SessionState { /// call's result is durable in the agent's own transcript, which is what /// the conversation reloads from. /// + /// `images` is NOT bounded, and deliberately: dropping the bytes does not + /// degrade the card, it inverts it. The frontend reads an image-generation + /// block with `image: null` and a terminal status as a FAILED generation + /// (`generated-images-block.tsx`, and `isImageGenerationToolCall` still + /// classifies the call from the `label` this keeps), so a trimmed success + /// would render "image generation failed". Carrying them whole is exactly + /// what the snapshot does today, so nothing here is a regression; what an + /// image-heavy turn costs is the same as before this function existed. The + /// bytes are still counted below, so they push older RESULT payload out + /// first. + /// /// The budget is spent, not enforced per entry, so the last entry admitted - /// may carry the total past it by its own size (one oversized image tool - /// call). In-flight entries are counted against the budget but never - /// trimmed by it. + /// may carry the total past it by its own size. In-flight entries and the + /// image data every entry keeps are counted against the budget but never + /// trimmed by it — which is why the floor can exceed it, and why the + /// accounting counts a trimmed entry's retained bytes rather than dropping + /// it from the tally. + /// + /// ## Assumption this rests on + /// + /// Bounding only terminal calls assumes the agent reports one. Every agent + /// codeg ships does (`upsert_tool_call` inserts at `Pending` and the + /// adapter's completion update moves it), but one that never did would + /// leave the table untrimmable at any length. Pinned by + /// `snapshot_ships_an_all_unsettled_table_whole` so a future change has to + /// confront the assumption rather than inherit it. + /// + /// ## Known degradation + /// + /// A trimmed entry loses `input`, which is one of the signals the client + /// infers a tool's identity from. `label`, `kind` and `meta` survive and + /// carry that identity for everything with an authoritative marker + /// (delegation companions, claudeCode/qoder/grok meta, the OpenCode name), + /// but two input-shape-only classifications fall back to a generic tool + /// card on a mid-turn attach of an over-budget turn: codex collab capsules + /// (`isCodexCollabInput`) and Kimi `TodoList` writes + /// (`kimiTodoWriteEntries`). Cosmetic and self-healing — the conversation + /// renders from the transcript on reload. fn snapshot_tool_calls(&self) -> Vec { // Arrival order comes from the live message: `push_tool_call_ref_if_absent` // anchors exactly one `ToolCallRef` per call, in the order the agent @@ -1896,7 +1930,7 @@ impl SessionState { } } - // (arrival rank, id, payload bytes). An id with no anchoring ref sorts + // (arrival rank, id, trimmable bytes). An id with no anchoring ref sorts // newest, so the fail-safe direction is "keep everything" — today that // cannot happen, because both `ToolCall` and `ToolCallUpdate` anchor a // ref for the id they upsert. @@ -1907,14 +1941,20 @@ impl SessionState { ( arrival.get(id.as_str()).copied().unwrap_or(usize::MAX), id.as_str(), - tool_call_payload_bytes(tc), + tool_call_trimmable_bytes(tc), ) }) .collect(); + // Both halves, because the question the early return answers is + // "would shipping this whole be over budget", and the image bytes are + // part of what ships either way. let total = ordered .iter() - .fold(0usize, |acc, (_, _, bytes)| acc.saturating_add(*bytes)); + .fold(0usize, |acc, (_, id, bytes)| { + acc.saturating_add(*bytes) + .saturating_add(images_slice_size(&self.active_tool_calls[*id].images)) + }); if total <= MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES { // The ordinary turn: nothing to trim, wire shape byte-identical. return self.active_tool_calls.values().cloned().collect(); @@ -1924,15 +1964,21 @@ impl SessionState { let mut trimmed: BTreeSet<&str> = BTreeSet::new(); let mut spent = 0usize; for (_, id, bytes) in ordered.iter().rev() { + let tc = &self.active_tool_calls[*id]; + // Counted on every entry, trimmed or not: images ship regardless + // (see the doc comment), so leaving them out of the tally would + // make the budget measure something the wire does not carry. + let kept = images_slice_size(&tc.images); let terminal = matches!( - self.active_tool_calls[*id].status, + tc.status, ToolCallStatus::Completed | ToolCallStatus::Failed ); if terminal && spent >= MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES { trimmed.insert(*id); + spent = spent.saturating_add(kept); continue; } - spent = spent.saturating_add(*bytes); + spent = spent.saturating_add(*bytes).saturating_add(kept); } self.active_tool_calls @@ -1959,7 +2005,10 @@ impl SessionState { // re-anchors an inline sub-thread on a mid-turn attach. // Bounded by contract — a small status object, not output. meta: tc.meta.clone(), - images: Vec::new(), + // Kept: an image-generation block whose `image` is null and + // whose status is terminal renders as a FAILED generation, + // so dropping these would report a success as a failure. + images: tc.images.clone(), // `#[serde(skip)]` — never on the wire either way. raw_input_chunks: Vec::new(), } @@ -2151,16 +2200,22 @@ fn u32_is_zero(v: &u32) -> bool { /// at least the newest 32 in the worst case) — far more than a client attaching /// mid-turn has on screen — while holding the #380 session's snapshot at ~2.6 MB /// instead of 24 MB. +/// +/// A ceiling on what is DROPPABLE, not a hard cap on the message: in-flight +/// calls and every call's image data are counted against it but never trimmed +/// by it, so a turn holding more of those than the budget ships more than the +/// budget. That is the same size it ships today. const MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES: usize = 2 * 1024 * 1024; -/// The part of a `ToolCallState` that grows with what the tool actually did — -/// what [`MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES`] bounds. Excludes the identity fields -/// (id / kind / label / status / meta), which every entry keeps. +/// The part of a `ToolCallState` a trim actually removes: the RESULT payload, +/// which grows with what the tool did. Excludes the identity fields (id / kind / +/// label / status / meta) and `images`, which every entry keeps whatever the +/// budget says — see [`SessionState::snapshot_tool_calls`] for why. /// /// Sized with the same escape-aware, allocation-free accounting the per-event /// cap uses (`event_stream`), so "this call's payload" means the same number of /// bytes on both paths. -fn tool_call_payload_bytes(tc: &ToolCallState) -> usize { +fn tool_call_trimmable_bytes(tc: &ToolCallState) -> usize { let output = match tc.output.as_ref() { Some(ToolCallOutput::Text { content }) => json_str_len(content), Some(ToolCallOutput::Error { message }) => json_str_len(message), @@ -2171,7 +2226,6 @@ fn tool_call_payload_bytes(tc: &ToolCallState) -> usize { .saturating_add(output) .saturating_add(opt_str_size(&tc.content)) .saturating_add(opt_json_size(&tc.locations)) - .saturating_add(images_slice_size(&tc.images)) } /// Last non-empty line of `s`, trimmed. `None` if every line is blank. @@ -3640,7 +3694,7 @@ mod tests { let held: usize = s .active_tool_calls .values() - .map(tool_call_payload_bytes) + .map(tool_call_trimmable_bytes) .sum(); assert!( held > 2 * MAX_SNAPSHOT_TOOL_PAYLOAD_BYTES, @@ -3691,7 +3745,6 @@ mod tests { assert!(oldest.content.is_none()); assert!(oldest.input.is_none()); assert!(oldest.locations.is_none()); - assert!(oldest.images.is_empty()); assert_eq!(oldest.label, "Read src/tc-0000.rs"); assert_eq!(oldest.kind, ToolKind::Read); assert_eq!(oldest.status, ToolCallStatus::Completed); @@ -3734,6 +3787,93 @@ mod tests { assert!(huge.output.is_some(), "a running call is never trimmed"); } + /// A generated image survives the trim, however old the call is. + /// + /// `isImageGenerationToolCall` classifies the call from the `label` the + /// trim keeps, and `generated-images-block.tsx` renders an + /// image-generation block whose `image` is null under a terminal status as + /// "image generation failed". So shedding the bytes would not show less, + /// it would report a success as a failure — and an image is exactly the + /// payload a user would then go looking for. + #[test] + fn snapshot_keeps_a_generated_image_on_the_oldest_trimmed_call() { + let mut s = fresh_state(); + // The oldest call, and the one carrying the image. + s.apply_event(&AcpEvent::ToolCall { + tool_call_id: "tc-image".into(), + title: "Image generation".into(), + kind: "other".into(), + status: "in_progress".into(), + content: None, + raw_input: Some("{\"prompt\":\"a cat\"}".into()), + raw_output: None, + locations: None, + meta: None, + images: None, + }); + s.apply_event(&AcpEvent::ToolCallUpdate { + tool_call_id: "tc-image".into(), + title: None, + status: Some("completed".into()), + content: None, + raw_input: None, + raw_output: Some("o".repeat(16 * 1024)), + raw_output_append: None, + locations: None, + meta: None, + images: Some(vec![ToolCallImageInfo { + data: "R0lGODlhAQABAAAAACw=".repeat(64), + mime_type: "image/png".into(), + uri: None, + }]), + }); + // …then enough finished work after it to push it out of the budget. + for i in 0..300 { + run_tool_call(&mut s, &format!("tc-{i:04}"), 16 * 1024, true); + } + + let snap = s.to_snapshot(); + let image_call = snap + .active_tool_calls + .iter() + .find(|tc| tc.id == "tc-image") + .expect("the image call still ships"); + assert!( + image_call.output.is_none(), + "it is old enough to be trimmed, which is what makes this a test" + ); + assert_eq!( + image_call.images.len(), + 1, + "a trimmed success must not come back as a failed generation" + ); + assert_eq!(image_call.label, "Image generation"); + } + + /// The bound covers terminal calls only, so a table that never reaches one + /// ships whole however long it gets. + /// + /// Every agent codeg ships reports a terminal status, which is what makes + /// the carve-out for in-flight calls safe. This pins the assumption rather + /// than leaving it implicit: an agent that stopped reporting one would + /// turn this test red instead of silently restoring #380. + #[test] + fn snapshot_ships_an_all_unsettled_table_whole() { + let mut s = fresh_state(); + for i in 0..300 { + run_tool_call(&mut s, &format!("tc-{i:04}"), 16 * 1024, false); + } + + let snap = s.to_snapshot(); + assert_eq!(snap.active_tool_calls.len(), 300); + assert!( + snap.active_tool_calls + .iter() + .all(|tc| tc.output.is_some() && tc.status == ToolCallStatus::InProgress), + "nothing unsettled is trimmed at any size" + ); + } + #[test] fn tool_call_content_field_is_preserved_on_state() { let mut s = fresh_state();