From 10c8feec70f868c4ebcd26a0ded96195cba9ed76 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 12:46:11 +0000 Subject: [PATCH 1/3] Bound AI conversation restoration memory growth (APP-5257) Cap the per-block stylized output and the number of command blocks materialized when restoring an AI conversation, mirroring the bound already applied to real terminal block serialization. CHANGELOG-BUG-FIX: Fixed unbounded memory growth when restoring AI conversations with very large command output or long history. --- app/src/ai/agent/conversation.rs | 39 ++++- app/src/ai/agent/conversation_tests.rs | 234 +++++++++++++++++++++++-- app/src/terminal/model/block.rs | 6 +- 3 files changed, 257 insertions(+), 22 deletions(-) diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index ab1a28f5d4a..58164c7b0d3 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -69,7 +69,8 @@ use crate::persistence::model::{ use crate::server::ids::ServerId; use crate::terminal::general_settings::GeneralSettings; use crate::terminal::model::block::{ - AgentInteractionMetadata, AgentViewVisibility, BlockId, SerializedAIMetadata, SerializedBlock, + AgentInteractionMetadata, AgentViewVisibility, BlockId, MAX_SERIALIZED_STYLIZED_OUTPUT_LINES, + SerializedAIMetadata, SerializedBlock, }; use crate::ui_components::icons::Icon; use crate::workspaces::user_profiles::UserProfileWithUID; @@ -448,6 +449,13 @@ pub struct AIConversation { task_sync_mode: TaskSyncMode, } +/// Maximum number of command blocks materialized per conversation on restore, keeping the +/// most recent. Restoring reads full command output straight from the persisted task +/// messages, which still hold everything summarization has since dropped from the AI's own +/// context, so a long-running or repeatedly summarized conversation would otherwise +/// materialize its entire history at once. +const MAX_RESTORED_COMMAND_BLOCKS: usize = 100; + pub(crate) fn artifact_from_fork_proto( proto_artifact: &api::message::artifact_event::ConversationArtifact, ) -> Option { @@ -3914,9 +3922,17 @@ impl AIConversation { } /// Normalize all newlines to CRLF so restored blocks render lines starting at column 0, - /// which is consistent with how we serialize real terminal blocks. - fn to_stylized_bytes(s: &str) -> Vec { + /// which is consistent with how we serialize real terminal blocks. When `max_lines` is + /// set, keeps only the most recent lines. + fn to_stylized_bytes(s: &str, max_lines: Option) -> Vec { let s = s.replace("\r\n", "\n"); + let s = match max_lines { + Some(max_lines) => { + let lines: Vec<&str> = s.split('\n').collect(); + lines[lines.len().saturating_sub(max_lines)..].join("\n") + } + None => s, + }; s.replace('\n', "\r\n").into_bytes() } @@ -4220,11 +4236,15 @@ impl AIConversation { pub fn to_serialized_blocklist_items(&self) -> Vec { let mut serialized_blocks = Vec::new(); - // Extract all command blocks from the task messages - let command_blocks = self.extract_command_blocks(); + // Extract all command blocks from the task messages, keeping only the most recent + // MAX_RESTORED_COMMAND_BLOCKS. + let mut command_blocks = self.extract_command_blocks(); + let total_command_blocks = command_blocks.len(); + command_blocks.drain(0..total_command_blocks.saturating_sub(MAX_RESTORED_COMMAND_BLOCKS)); log::info!( - "Extracted {} command blocks for conversation {}", + "Extracted {} of {} command blocks for conversation {}", command_blocks.len(), + total_command_blocks, self.id() ); @@ -4257,8 +4277,11 @@ impl AIConversation { let serialized_block = SerializedBlock { id: BlockId::new(), - stylized_command: Self::to_stylized_bytes(&command_block.command), - stylized_output: Self::to_stylized_bytes(&command_block.output), + stylized_command: Self::to_stylized_bytes(&command_block.command, None), + stylized_output: Self::to_stylized_bytes( + &command_block.output, + Some(MAX_SERIALIZED_STYLIZED_OUTPUT_LINES), + ), pwd, git_head: None, git_branch_name: None, diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index a9988c2bad1..7bcc3266668 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/app/src/ai/agent/conversation_tests.rs @@ -7,9 +7,11 @@ use warpui::{App, SingletonEntity}; use super::{ AIConversation, AIConversationAutoexecuteMode, AIConversationId, ConversationStatus, - ConversationUsageTotals, RecordingSpanStatus, RestoreConversationError, + ConversationUsageTotals, MAX_RESTORED_COMMAND_BLOCKS, MAX_SERIALIZED_STYLIZED_OUTPUT_LINES, + RecordingSpanStatus, RestoreConversationError, SerializedBlockListItem, artifact_from_fork_proto, footer_model_token_usage, }; +use crate::ai::agent::task::helper::MessageExt; use crate::ai::artifacts::Artifact; use crate::ai::llms::LLMPreferences; use crate::auth::AuthStateProvider; @@ -223,19 +225,90 @@ fn stop_recording_error_result(message: &str) -> api::message::tool_call_result: }) } fn restored_conversation_with_messages(messages: Vec) -> AIConversation { - AIConversation::new_restored( - AIConversationId::new(), - vec![api::Task { - id: "root-task".to_string(), - messages, - dependencies: None, - description: String::new(), - summary: String::new(), - server_data: String::new(), - }], - None, + restored_conversation_with_tasks(vec![api::Task { + id: "root-task".to_string(), + messages, + dependencies: None, + description: String::new(), + summary: String::new(), + server_data: String::new(), + }]) +} + +fn restored_conversation_with_tasks(tasks: Vec) -> AIConversation { + AIConversation::new_restored(AIConversationId::new(), tasks, None).unwrap() +} + +fn run_shell_command_tool_call(command: &str) -> api::message::tool_call::Tool { + api::message::tool_call::Tool::RunShellCommand(api::message::tool_call::RunShellCommand { + command: command.to_string(), + is_read_only: false, + uses_pager: false, + citations: vec![], + is_risky: false, + wait_until_complete_value: None, + risk_category: 0, + }) +} + +#[allow(deprecated)] +fn run_shell_command_finished_result( + command_id: &str, + output: &str, +) -> api::message::tool_call_result::Result { + api::message::tool_call_result::Result::RunShellCommand(api::RunShellCommandResult { + command: String::new(), + output: String::new(), + exit_code: 0, + result: Some(api::run_shell_command_result::Result::CommandFinished( + api::ShellCommandFinished { + command_id: command_id.to_string(), + output: output.to_string(), + exit_code: 0, + start_ts: None, + finish_ts: None, + }, + )), + }) +} + +/// A pair of messages recording a completed shell command: the tool call and its result. +fn run_shell_command_messages(index: usize, output: &str) -> Vec { + let tool_call_id = format!("call-{index}"); + vec![ + tool_call_message( + &format!("call-msg-{index}"), + "req", + &tool_call_id, + run_shell_command_tool_call(&format!("cmd-{index}")), + ), + tool_call_result_message( + &format!("result-msg-{index}"), + "req", + &tool_call_id, + run_shell_command_finished_result(&format!("command-{index}"), output), + ), + ] +} + +fn summarization_subagent_tool_call(task_id: &str) -> api::message::tool_call::Tool { + api::message::tool_call::Tool::Subagent(api::message::tool_call::Subagent { + task_id: task_id.to_string(), + payload: String::new(), + metadata: Some(api::message::tool_call::subagent::Metadata::Summarization( + (), + )), + }) +} + +/// Returns the command block's stylized command and output, decoded as UTF-8, from a +/// [`SerializedBlockListItem`] produced by [`AIConversation::to_serialized_blocklist_items`]. +fn command_and_output(item: &SerializedBlockListItem) -> (String, String) { + let SerializedBlockListItem::Command { block } = item; + ( + String::from_utf8(block.stylized_command.clone()).unwrap(), + String::from_utf8(block.stylized_output.clone()).unwrap(), ) - .unwrap() } fn agent_output_message(id: &str, request_id: &str) -> api::Message { @@ -1540,3 +1613,138 @@ fn fetched_memories_dedupes_keeping_first_position_and_latest_data() { ] ); } + +#[test] +fn to_serialized_blocklist_items_truncates_long_output_to_most_recent_lines() { + let line_count = MAX_SERIALIZED_STYLIZED_OUTPUT_LINES + 10; + let output = (0..line_count) + .map(|i| format!("line-{i}")) + .collect::>() + .join("\n"); + let conversation = restored_conversation_with_messages(run_shell_command_messages(0, &output)); + + let items = conversation.to_serialized_blocklist_items(); + + assert_eq!(items.len(), 1); + let (_, restored_output) = command_and_output(&items[0]); + let restored_lines: Vec<&str> = restored_output.split("\r\n").collect(); + assert_eq!(restored_lines.len(), MAX_SERIALIZED_STYLIZED_OUTPUT_LINES); + assert_eq!(restored_lines[0], "line-10"); + assert_eq!( + restored_lines[restored_lines.len() - 1], + format!("line-{}", line_count - 1) + ); +} + +#[test] +fn to_serialized_blocklist_items_does_not_mutate_persisted_task_messages() { + let long_output = (0..MAX_SERIALIZED_STYLIZED_OUTPUT_LINES + 10) + .map(|i| format!("line-{i}")) + .collect::>() + .join("\n"); + let conversation = + restored_conversation_with_messages(run_shell_command_messages(0, &long_output)); + + // This only produces a derived, display-time view; it must not touch the persisted task + // messages that `write_updated_conversation_state` reads from. + let _ = conversation.to_serialized_blocklist_items(); + + let persisted_output = conversation + .get_root_task() + .and_then(|task| task.source()) + .and_then(|source| { + source.messages.iter().find_map(|message| { + let result = message.tool_call_result()?; + match &result.result { + Some(api::message::tool_call_result::Result::RunShellCommand(cmd_result)) => { + match &cmd_result.result { + Some(api::run_shell_command_result::Result::CommandFinished( + finished, + )) => Some(finished.output.clone()), + _ => None, + } + } + _ => None, + } + }) + }) + .unwrap(); + assert_eq!(persisted_output, long_output); +} + +#[test] +fn to_serialized_blocklist_items_caps_total_blocks_to_most_recent() { + let total_commands = MAX_RESTORED_COMMAND_BLOCKS + 5; + let messages = (0..total_commands) + .flat_map(|i| run_shell_command_messages(i, "output")) + .collect(); + let conversation = restored_conversation_with_messages(messages); + + let items = conversation.to_serialized_blocklist_items(); + + assert_eq!(items.len(), MAX_RESTORED_COMMAND_BLOCKS); + let (first_command, _) = command_and_output(&items[0]); + let (last_command, _) = command_and_output(&items[items.len() - 1]); + assert_eq!( + first_command, + format!("cmd-{}", total_commands - MAX_RESTORED_COMMAND_BLOCKS) + ); + assert_eq!(last_command, format!("cmd-{}", total_commands - 1)); +} + +#[test] +fn to_serialized_blocklist_items_bounds_output_resurrected_from_summarized_history() { + let old_output = (0..MAX_SERIALIZED_STYLIZED_OUTPUT_LINES + 10) + .map(|i| format!("old-line-{i}")) + .collect::>() + .join("\n"); + let mut old_command_messages = run_shell_command_messages(0, &old_output); + for message in &mut old_command_messages { + message.task_id = "summary-task".to_string(); + } + let subtask = api::Task { + id: "summary-task".to_string(), + messages: old_command_messages, + dependencies: Some(api::task::Dependencies { + parent_task_id: "root-task".to_string(), + }), + description: String::new(), + summary: String::new(), + server_data: String::new(), + }; + let mut root_messages = vec![tool_call_message( + "summarize-call", + "req", + "summarize", + summarization_subagent_tool_call("summary-task"), + )]; + root_messages.extend(run_shell_command_messages(1, "new-output")); + let root_task = api::Task { + id: "root-task".to_string(), + messages: root_messages, + dependencies: None, + description: String::new(), + summary: String::new(), + server_data: String::new(), + }; + let conversation = restored_conversation_with_tasks(vec![root_task, subtask]); + + let items = conversation.to_serialized_blocklist_items(); + + // Both the resurrected pre-summarization command and the new one are still restored... + assert_eq!(items.len(), 2); + let (old_command, old_restored_output) = command_and_output(&items[0]); + let (new_command, new_restored_output) = command_and_output(&items[1]); + assert_eq!(old_command, "cmd-0"); + assert_eq!(new_command, "cmd-1"); + assert_eq!(new_restored_output, "new-output"); + + // ...but the resurrected command's output is bounded the same as any other block, since + // restore reads it straight from the task messages summarization left untouched. + let old_restored_lines: Vec<&str> = old_restored_output.split("\r\n").collect(); + assert_eq!( + old_restored_lines.len(), + MAX_SERIALIZED_STYLIZED_OUTPUT_LINES + ); + assert_eq!(old_restored_lines[0], "old-line-10"); +} diff --git a/app/src/terminal/model/block.rs b/app/src/terminal/model/block.rs index 8f5781ebc2f..4afd1acf673 100644 --- a/app/src/terminal/model/block.rs +++ b/app/src/terminal/model/block.rs @@ -108,7 +108,11 @@ impl TranscriptScope { } } -pub(super) const MAX_SERIALIZED_STYLIZED_OUTPUT_LINES: usize = 5000; +/// Cap on stored `stylized_output` lines, applied both when serializing a live `Block` (see +/// `impl From<&Block> for SerializedBlock`) and when materializing a restored AI conversation's +/// command blocks (`AIConversation::to_serialized_blocklist_items`), so both restoration paths +/// bound memory the same way. +pub(crate) const MAX_SERIALIZED_STYLIZED_OUTPUT_LINES: usize = 5000; /// Number of max lines to store that aren't stylized. We only store 50 lines as we only need /// non-stylized lines for command corrections and notifications whereas we need more lines for the From 6d24d5e2de32a033d5eb47e46da0f2e57f321f3b Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:22:38 +0000 Subject: [PATCH 2/3] Bound peak memory during conversation-block extraction/truncation - Evict oldest command blocks during traversal (RecentCommandBlocks) instead of collecting the full history and truncating afterward, so peak memory is bounded to the retained window rather than the whole conversation. - Locate the retained output tail by scanning backward for the line boundary instead of materializing the untruncated string first. - Restate the shared line-cap constant's doc comment without enumerating its call sites. - Add regression tests for the extraction accumulator, the tail-finding helper, and the restoration-plan placement boundary. --- app/src/ai/agent/conversation.rs | 91 ++++++++++++++----- app/src/ai/agent/conversation_tests.rs | 80 +++++++++++++++- .../conversation_restoration_tests.rs | 20 ++++ app/src/terminal/model/block.rs | 6 +- 4 files changed, 167 insertions(+), 30 deletions(-) diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 58164c7b0d3..c1c02852d84 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{HashMap, HashSet, VecDeque}; use ai::agent::orchestration_config::{OrchestrationConfig, OrchestrationConfigStatus}; use ai::document::AIDocumentId; @@ -456,6 +456,38 @@ pub struct AIConversation { /// materialize its entire history at once. const MAX_RESTORED_COMMAND_BLOCKS: usize = 100; +/// Retains only the most recent [`MAX_RESTORED_COMMAND_BLOCKS`] pushed to it, evicting the +/// oldest as soon as that cap is exceeded. Collecting every command first and truncating +/// afterward would still allocate a clone of every historical command's output before any of +/// it could be freed; evicting during traversal keeps that peak bounded to the retained window +/// instead of the conversation's full history. +struct RecentCommandBlocks { + blocks: VecDeque, + total_seen: usize, +} + +impl RecentCommandBlocks { + fn new() -> Self { + Self { + blocks: VecDeque::with_capacity(MAX_RESTORED_COMMAND_BLOCKS), + total_seen: 0, + } + } + + fn push(&mut self, block: CommandBlockInfo) { + self.total_seen += 1; + self.blocks.push_back(block); + if self.blocks.len() > MAX_RESTORED_COMMAND_BLOCKS { + self.blocks.pop_front(); + } + } + + /// Consumes this accumulator, returning the retained blocks and the total number pushed. + fn into_vec_with_total_seen(self) -> (Vec, usize) { + (self.blocks.into(), self.total_seen) + } +} + pub(crate) fn artifact_from_fork_proto( proto_artifact: &api::message::artifact_event::ConversationArtifact, ) -> Option { @@ -3923,21 +3955,39 @@ impl AIConversation { /// Normalize all newlines to CRLF so restored blocks render lines starting at column 0, /// which is consistent with how we serialize real terminal blocks. When `max_lines` is - /// set, keeps only the most recent lines. + /// set, keeps only the most recent lines; the retained tail is located without allocating + /// a copy of the untruncated string, so a single oversized output can't force a transient + /// full-size copy. fn to_stylized_bytes(s: &str, max_lines: Option) -> Vec { - let s = s.replace("\r\n", "\n"); - let s = match max_lines { - Some(max_lines) => { - let lines: Vec<&str> = s.split('\n').collect(); - lines[lines.len().saturating_sub(max_lines)..].join("\n") - } + let tail = match max_lines { + Some(max_lines) => Self::tail_lines(s, max_lines), None => s, }; - s.replace('\n', "\r\n").into_bytes() + let normalized = tail.replace("\r\n", "\n"); + normalized.replace('\n', "\r\n").into_bytes() + } + + /// Returns the suffix of `s` made up of at most its last `max_lines` `\n`-delimited lines, + /// found by scanning backward for the boundary instead of splitting the whole string. + fn tail_lines(s: &str, max_lines: usize) -> &str { + if max_lines == 0 { + return ""; + } + let mut newlines_seen = 0; + for (i, byte) in s.bytes().enumerate().rev() { + if byte == b'\n' { + newlines_seen += 1; + if newlines_seen == max_lines { + return &s[i + 1..]; + } + } + } + s } - /// Extracts all shell command blocks, in order, from the conversation's API task - /// messages. + /// Extracts the most recent [`MAX_RESTORED_COMMAND_BLOCKS`] shell command blocks, in order, + /// from the conversation's API task messages, plus the total number found before that cap + /// was applied. /// /// This includes: /// - RunShellCommand tool calls that completed @@ -3945,15 +3995,15 @@ impl AIConversation { /// - Context blocks from UserQuery/SystemQuery/ToolCallResult messages /// /// Returns CommandBlockInfo with command, output, exit_code, and optional ai_metadata. - fn extract_command_blocks(&self) -> Vec { - let mut command_blocks = Vec::new(); + fn extract_command_blocks(&self) -> (Vec, usize) { + let mut command_blocks = RecentCommandBlocks::new(); // Get the root task's API messages. let Some(root_task) = self.get_root_task() else { - return command_blocks; + return command_blocks.into_vec_with_total_seen(); }; let Some(api_task) = root_task.source() else { - return command_blocks; + return command_blocks.into_vec_with_total_seen(); }; // Build a map from message ID to exchange for timestamp lookups. @@ -3978,7 +4028,7 @@ impl AIConversation { &mut seen_command_ids, ); - command_blocks + command_blocks.into_vec_with_total_seen() } /// Extracts command blocks from a list of messages. @@ -3989,7 +4039,7 @@ impl AIConversation { &self, messages: &[api::Message], message_id_to_exchange: &HashMap<&str, &AIAgentExchange>, - command_blocks: &mut Vec, + command_blocks: &mut RecentCommandBlocks, seen_command_ids: &mut HashSet, ) { // Build a map from tool_call_id to (RunShellCommandResult, result_message_id, result_proto_timestamp) @@ -4236,11 +4286,8 @@ impl AIConversation { pub fn to_serialized_blocklist_items(&self) -> Vec { let mut serialized_blocks = Vec::new(); - // Extract all command blocks from the task messages, keeping only the most recent - // MAX_RESTORED_COMMAND_BLOCKS. - let mut command_blocks = self.extract_command_blocks(); - let total_command_blocks = command_blocks.len(); - command_blocks.drain(0..total_command_blocks.saturating_sub(MAX_RESTORED_COMMAND_BLOCKS)); + // Extraction itself retains only the most recent MAX_RESTORED_COMMAND_BLOCKS. + let (command_blocks, total_command_blocks) = self.extract_command_blocks(); log::info!( "Extracted {} of {} command blocks for conversation {}", command_blocks.len(), diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index 7bcc3266668..69349c289ff 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/app/src/ai/agent/conversation_tests.rs @@ -1,15 +1,17 @@ use std::collections::HashMap; use ai::api_keys::{ApiKeyManager, CustomEndpointParams, CustomEndpointSchema}; +use warp_core::command::ExitCode; use warp_core::features::FeatureFlag; use warp_multi_agent_api as api; use warpui::{App, SingletonEntity}; use super::{ - AIConversation, AIConversationAutoexecuteMode, AIConversationId, ConversationStatus, - ConversationUsageTotals, MAX_RESTORED_COMMAND_BLOCKS, MAX_SERIALIZED_STYLIZED_OUTPUT_LINES, - RecordingSpanStatus, RestoreConversationError, SerializedBlockListItem, - artifact_from_fork_proto, footer_model_token_usage, + AIConversation, AIConversationAutoexecuteMode, AIConversationId, CommandBlockInfo, + ConversationStatus, ConversationUsageTotals, MAX_RESTORED_COMMAND_BLOCKS, + MAX_SERIALIZED_STYLIZED_OUTPUT_LINES, RecentCommandBlocks, RecordingSpanStatus, + RestoreConversationError, SerializedBlockListItem, artifact_from_fork_proto, + footer_model_token_usage, }; use crate::ai::agent::task::helper::MessageExt; use crate::ai::artifacts::Artifact; @@ -1748,3 +1750,73 @@ fn to_serialized_blocklist_items_bounds_output_resurrected_from_summarized_histo ); assert_eq!(old_restored_lines[0], "old-line-10"); } + +fn command_block_info(command: &str) -> CommandBlockInfo { + CommandBlockInfo { + command: command.to_string(), + output: String::new(), + exit_code: ExitCode::from(0), + ai_metadata: None, + message_id: format!("{command}-message"), + start_ts: None, + completed_ts: None, + } +} + +#[test] +fn recent_command_blocks_evicts_oldest_and_counts_total_pushed() { + let total = MAX_RESTORED_COMMAND_BLOCKS + 5; + let mut recent = RecentCommandBlocks::new(); + for i in 0..total { + recent.push(command_block_info(&format!("cmd-{i}"))); + } + + let (blocks, total_seen) = recent.into_vec_with_total_seen(); + + assert_eq!(total_seen, total); + assert_eq!(blocks.len(), MAX_RESTORED_COMMAND_BLOCKS); + assert_eq!( + blocks.first().unwrap().command, + format!("cmd-{}", total - MAX_RESTORED_COMMAND_BLOCKS) + ); + assert_eq!(blocks.last().unwrap().command, format!("cmd-{}", total - 1)); +} + +#[test] +fn recent_command_blocks_keeps_everything_under_the_cap() { + let mut recent = RecentCommandBlocks::new(); + recent.push(command_block_info("cmd-0")); + recent.push(command_block_info("cmd-1")); + + let (blocks, total_seen) = recent.into_vec_with_total_seen(); + + assert_eq!(total_seen, 2); + assert_eq!(blocks.len(), 2); + assert_eq!(blocks[0].command, "cmd-0"); + assert_eq!(blocks[1].command, "cmd-1"); +} + +#[test] +fn tail_lines_keeps_full_string_when_under_the_cap() { + assert_eq!(AIConversation::tail_lines("a\nb", 2), "a\nb"); + assert_eq!(AIConversation::tail_lines("a\nb", 5), "a\nb"); +} + +#[test] +fn tail_lines_finds_the_exact_boundary() { + let lines: Vec = (0..10).map(|i| format!("line-{i}")).collect(); + let s = lines.join("\n"); + + assert_eq!(AIConversation::tail_lines(&s, 3), "line-7\nline-8\nline-9"); +} + +#[test] +fn tail_lines_handles_a_trailing_newline() { + // "a\nb\n" splits into three `\n`-delimited pieces: "a", "b", and a trailing empty piece. + assert_eq!(AIConversation::tail_lines("a\nb\n", 2), "b\n"); +} + +#[test] +fn tail_lines_with_zero_max_lines_returns_empty() { + assert_eq!(AIConversation::tail_lines("a\nb\nc", 0), ""); +} diff --git a/app/src/terminal/conversation_restoration_tests.rs b/app/src/terminal/conversation_restoration_tests.rs index bcae222e099..04052d3b964 100644 --- a/app/src/terminal/conversation_restoration_tests.rs +++ b/app/src/terminal/conversation_restoration_tests.rs @@ -194,6 +194,26 @@ fn sorted_tail_multiple_exchanges() { // ── Edge cases ──────────────────────────────────────────────────────── +// ── Retained-window boundary (APP-5257) ──────────────────────────────── +// `AIConversation::to_serialized_blocklist_items` keeps only the most recent +// `MAX_RESTORED_COMMAND_BLOCKS` command blocks. An exchange whose command fell +// outside that window no longer has a matching block, but it must still be +// placed immediately before the oldest *retained* command block rather than +// vanishing or being reordered. + +#[test] +fn exchange_older_than_the_retained_window_attaches_to_the_oldest_retained_block() { + // The retained window starts at t=100; an exchange from t=1 (long before any + // retained block, because its own command was evicted by the cap) must attach + // to the oldest retained block rather than being dropped. + let blocks = vec![(bi(0), ts(100)), (bi(1), ts(200)), (bi(2), ts(300))]; + let exchanges = vec![ts(1)]; + + let result = find_block_indices_for_exchange_timestamps(&blocks, &exchanges); + + assert_eq!(result, vec![Some(bi(0))]); +} + #[test] fn empty_blocks_returns_none_for_all_exchanges() { let blocks: Vec<(BlockIndex, chrono::DateTime)> = vec![]; diff --git a/app/src/terminal/model/block.rs b/app/src/terminal/model/block.rs index 4afd1acf673..c9b8fba2a70 100644 --- a/app/src/terminal/model/block.rs +++ b/app/src/terminal/model/block.rs @@ -108,10 +108,8 @@ impl TranscriptScope { } } -/// Cap on stored `stylized_output` lines, applied both when serializing a live `Block` (see -/// `impl From<&Block> for SerializedBlock`) and when materializing a restored AI conversation's -/// command blocks (`AIConversation::to_serialized_blocklist_items`), so both restoration paths -/// bound memory the same way. +/// Maximum number of lines retained in a serialized block's `stylized_output`, keeping the +/// most recent. pub(crate) const MAX_SERIALIZED_STYLIZED_OUTPUT_LINES: usize = 5000; /// Number of max lines to store that aren't stylized. We only store 50 lines as we only need From b18baf6c27ce2a3d1e7d22fd10a391e62ba142d4 Mon Sep 17 00:00:00 2001 From: "warp-agent-staging[bot]" <240773466+warp-agent-staging[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:43:11 +0000 Subject: [PATCH 3/3] Truncate command output before accumulating, add byte-ceiling backstop - Fix: the count cap bounded RecentCommandBlocks entries, but each CommandBlockInfo still held a full-size clone of its output because truncation only happened later in to_serialized_blocklist_items. Apply truncated_output() at all three extraction sites (RunShellCommand result, UserQuery attachment, deprecated context executed_shell_commands) before constructing CommandBlockInfo. - Add MAX_RESTORED_COMMAND_OUTPUT_BYTES (1 MB, mirroring the existing MAX_FILE_READ_BYTES convention) as a backstop for a single pathologically long line that a line cap alone cannot bound. - Add regression tests asserting truncation happens during extraction (not just in the final serialized output) for all three sources, plus byte-ceiling and UTF-8 boundary tests. - Fix a brace-nesting slip from an earlier revision that had silently moved the context-block extraction outside the message loop (caught by the compiler before landing). --- app/src/ai/agent/conversation.rs | 33 +++++- app/src/ai/agent/conversation_tests.rs | 151 ++++++++++++++++++++++++- 2 files changed, 178 insertions(+), 6 deletions(-) diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index c1c02852d84..6606cafd247 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -456,6 +456,11 @@ pub struct AIConversation { /// materialize its entire history at once. const MAX_RESTORED_COMMAND_BLOCKS: usize = 100; +/// Backstop against a single pathologically long line (e.g. no `\n` at all), which +/// [`MAX_SERIALIZED_STYLIZED_OUTPUT_LINES`] alone cannot bound. Mirrors the per-file byte +/// ceiling already used for local file-read context (`MAX_FILE_READ_BYTES`). +const MAX_RESTORED_COMMAND_OUTPUT_BYTES: usize = 1_000_000; + /// Retains only the most recent [`MAX_RESTORED_COMMAND_BLOCKS`] pushed to it, evicting the /// oldest as soon as that cap is exceeded. Collecting every command first and truncating /// afterward would still allocate a clone of every historical command's output before any of @@ -3985,6 +3990,28 @@ impl AIConversation { s } + /// Clones only the retained tail of a command's output, so a [`CommandBlockInfo`] + /// accumulated during extraction never holds a full-size copy of an oversized command's + /// output. Applies the byte ceiling before the line cap so a single pathologically long + /// line (no `\n` at all) is still bounded. + fn truncated_output(output: &str) -> String { + let byte_capped = Self::tail_bytes(output, MAX_RESTORED_COMMAND_OUTPUT_BYTES); + Self::tail_lines(byte_capped, MAX_SERIALIZED_STYLIZED_OUTPUT_LINES).to_string() + } + + /// Returns the suffix of `s` containing at most `max_bytes` bytes, snapped forward to the + /// nearest UTF-8 character boundary so the result is always a valid `&str`. + fn tail_bytes(s: &str, max_bytes: usize) -> &str { + if s.len() <= max_bytes { + return s; + } + let mut start = s.len() - max_bytes; + while !s.is_char_boundary(start) { + start += 1; + } + &s[start..] + } + /// Extracts the most recent [`MAX_RESTORED_COMMAND_BLOCKS`] shell command blocks, in order, /// from the conversation's API task messages, plus the total number found before that cap /// was applied. @@ -4160,7 +4187,7 @@ impl AIConversation { command_blocks.push(CommandBlockInfo { command: command.clone(), - output: command_output.clone(), + output: Self::truncated_output(command_output), exit_code: ExitCode::from(*exit_code), ai_metadata: Some( serde_json::to_string(&Some(Into::::into( @@ -4223,7 +4250,7 @@ impl AIConversation { .or(msg_ts); command_blocks.push(CommandBlockInfo { command: cmd.command.clone(), - output: cmd.output.clone(), + output: Self::truncated_output(&cmd.output), exit_code: ExitCode::from(cmd.exit_code), ai_metadata: None, message_id: message_id.clone(), @@ -4264,7 +4291,7 @@ impl AIConversation { .or(msg_ts); command_blocks.push(CommandBlockInfo { command: executed_shell_command.command.clone(), - output: executed_shell_command.output.clone(), + output: Self::truncated_output(&executed_shell_command.output), exit_code: ExitCode::from(executed_shell_command.exit_code), ai_metadata: None, message_id: message_id.clone(), diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index 69349c289ff..69a615b8d61 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/app/src/ai/agent/conversation_tests.rs @@ -9,9 +9,9 @@ use warpui::{App, SingletonEntity}; use super::{ AIConversation, AIConversationAutoexecuteMode, AIConversationId, CommandBlockInfo, ConversationStatus, ConversationUsageTotals, MAX_RESTORED_COMMAND_BLOCKS, - MAX_SERIALIZED_STYLIZED_OUTPUT_LINES, RecentCommandBlocks, RecordingSpanStatus, - RestoreConversationError, SerializedBlockListItem, artifact_from_fork_proto, - footer_model_token_usage, + MAX_RESTORED_COMMAND_OUTPUT_BYTES, MAX_SERIALIZED_STYLIZED_OUTPUT_LINES, RecentCommandBlocks, + RecordingSpanStatus, RestoreConversationError, SerializedBlockListItem, + artifact_from_fork_proto, footer_model_token_usage, }; use crate::ai::agent::task::helper::MessageExt; use crate::ai::artifacts::Artifact; @@ -1820,3 +1820,148 @@ fn tail_lines_handles_a_trailing_newline() { fn tail_lines_with_zero_max_lines_returns_empty() { assert_eq!(AIConversation::tail_lines("a\nb\nc", 0), ""); } + +fn executed_shell_command( + command_id: &str, + command: &str, + output: &str, +) -> api::ExecutedShellCommand { + api::ExecutedShellCommand { + command: command.to_string(), + output: output.to_string(), + exit_code: 0, + command_id: command_id.to_string(), + started_ts: None, + finished_ts: None, + is_auto_attached: false, + } +} + +fn user_query_with_attachment( + id: &str, + request_id: &str, + attachment_key: &str, + cmd: api::ExecutedShellCommand, +) -> api::Message { + api::Message { + fetched_memories: vec![], + id: id.to_string(), + task_id: "root-task".to_string(), + server_message_data: String::new(), + citations: vec![], + message: Some(api::message::Message::UserQuery(api::message::UserQuery { + query: String::new(), + context: None, + referenced_attachments: HashMap::from([( + attachment_key.to_string(), + api::Attachment { + value: Some(api::attachment::Value::ExecutedShellCommand(cmd)), + }, + )]), + mode: None, + intended_agent: Default::default(), + })), + request_id: request_id.to_string(), + timestamp: None, + } +} + +#[allow(deprecated)] +fn user_query_with_context_executed_shell_command( + id: &str, + request_id: &str, + cmd: api::ExecutedShellCommand, +) -> api::Message { + api::Message { + fetched_memories: vec![], + id: id.to_string(), + task_id: "root-task".to_string(), + server_message_data: String::new(), + citations: vec![], + message: Some(api::message::Message::UserQuery(api::message::UserQuery { + query: String::new(), + context: Some(api::InputContext { + executed_shell_commands: vec![cmd], + ..Default::default() + }), + referenced_attachments: HashMap::new(), + mode: None, + intended_agent: Default::default(), + })), + request_id: request_id.to_string(), + timestamp: None, + } +} + +fn long_output(line_count: usize, prefix: &str) -> String { + (0..line_count) + .map(|i| format!("{prefix}{i}")) + .collect::>() + .join("\n") +} + +/// Asserts that a single extracted command block's output was truncated to the most recent +/// `MAX_SERIALIZED_STYLIZED_OUTPUT_LINES` lines *during extraction* (i.e. before it's ever +/// serialized), not merely truncated later when serialized. +fn assert_extracted_output_bounded(conversation: &AIConversation, line_prefix: &str) { + let (blocks, _total_seen) = conversation.extract_command_blocks(); + assert_eq!(blocks.len(), 1); + let stored_lines: Vec<&str> = blocks[0].output.split('\n').collect(); + assert_eq!(stored_lines.len(), MAX_SERIALIZED_STYLIZED_OUTPUT_LINES); + assert_eq!(stored_lines[0], format!("{line_prefix}10")); +} + +#[test] +fn extract_command_blocks_truncates_run_shell_command_output_before_accumulating() { + let output = long_output(MAX_SERIALIZED_STYLIZED_OUTPUT_LINES + 10, "line-"); + let conversation = restored_conversation_with_messages(run_shell_command_messages(0, &output)); + + assert_extracted_output_bounded(&conversation, "line-"); +} + +#[test] +fn extract_command_blocks_truncates_attachment_output_before_accumulating() { + let output = long_output(MAX_SERIALIZED_STYLIZED_OUTPUT_LINES + 10, "line-"); + let cmd = executed_shell_command("cmd-1", "cat big.log", &output); + let conversation = restored_conversation_with_messages(vec![user_query_with_attachment( + "user-0", + "req", + "attachment-1", + cmd, + )]); + + assert_extracted_output_bounded(&conversation, "line-"); +} + +#[test] +fn extract_command_blocks_truncates_context_executed_shell_command_output_before_accumulating() { + let output = long_output(MAX_SERIALIZED_STYLIZED_OUTPUT_LINES + 10, "line-"); + let cmd = executed_shell_command("cmd-1", "cat big.log", &output); + let conversation = + restored_conversation_with_messages(vec![user_query_with_context_executed_shell_command( + "user-0", "req", cmd, + )]); + + assert_extracted_output_bounded(&conversation, "line-"); +} + +#[test] +fn truncated_output_applies_byte_ceiling_to_a_single_line_with_no_newlines() { + // A pathologically long single line has no `\n` at all, so `tail_lines` alone can't bound + // it; the byte ceiling must still cap it. + let huge_single_line = "x".repeat(MAX_RESTORED_COMMAND_OUTPUT_BYTES + 100); + + let truncated = AIConversation::truncated_output(&huge_single_line); + + assert_eq!(truncated.len(), MAX_RESTORED_COMMAND_OUTPUT_BYTES); +} + +#[test] +fn tail_bytes_snaps_forward_to_a_utf8_character_boundary() { + // "é" is 2 bytes, starting right after "a" (1 byte). Cutting at max_bytes=2 lands exactly + // on that boundary; max_bytes=1 would land mid-character and must snap forward past it. + let s = "aé"; + assert_eq!(AIConversation::tail_bytes(s, 2), "é"); + assert_eq!(AIConversation::tail_bytes(s, 1), ""); + assert_eq!(AIConversation::tail_bytes(s, 100), s); +}