Bound AI conversation restoration memory growth (APP-5257)

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;
@@ -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,50 @@ 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;
+
+/// 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
+/// 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<CommandBlockInfo>,
+    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<CommandBlockInfo>, usize) {
+        (self.blocks.into(), self.total_seen)
+    }
+}
+
 pub(crate) fn artifact_from_fork_proto(
     proto_artifact: &api::message::artifact_event::ConversationArtifact,
 ) -> Option<Artifact> {
@@ -3914,30 +3959,78 @@ 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<u8> {
-        let s = s.replace("\r\n", "\n");
-        s.replace('\n', "\r\n").into_bytes()
+    /// which is consistent with how we serialize real terminal blocks. When `max_lines` is
+    /// 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<usize>) -> Vec<u8> {
+        let tail = match max_lines {
+            Some(max_lines) => Self::tail_lines(s, max_lines),
+            None => s,
+        };
+        let normalized = tail.replace("\r\n", "\n");
+        normalized.replace('\n', "\r\n").into_bytes()
     }
 
-    /// Extracts all shell command blocks, in order, from the conversation's API task
-    /// messages.
+    /// 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
+    }
+
+    /// 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.
     ///
     /// This includes:
     /// - RunShellCommand tool calls that completed
     /// - Attachments from UserQuery/SystemQuery messages
     /// - Context blocks from UserQuery/SystemQuery/ToolCallResult messages
     ///
     /// Returns CommandBlockInfo with command, output, exit_code, and optional ai_metadata.
-    fn extract_command_blocks(&self) -> Vec<CommandBlockInfo> {
-        let mut command_blocks = Vec::new();
+    fn extract_command_blocks(&self) -> (Vec<CommandBlockInfo>, 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.
@@ -3962,7 +4055,7 @@ impl AIConversation {
             &mut seen_command_ids,
         );
 
-        command_blocks
+        command_blocks.into_vec_with_total_seen()
     }
 
     /// Extracts command blocks from a list of messages.
@@ -3973,7 +4066,7 @@ impl AIConversation {
         &self,
         messages: &[api::Message],
         message_id_to_exchange: &HashMap<&str, &AIAgentExchange>,
-        command_blocks: &mut Vec<CommandBlockInfo>,
+        command_blocks: &mut RecentCommandBlocks,
         seen_command_ids: &mut HashSet<String>,
     ) {
         // Build a map from tool_call_id to (RunShellCommandResult, result_message_id, result_proto_timestamp)
@@ -4094,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::<SerializedAIMetadata>::into(
@@ -4157,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(),
@@ -4198,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(),
@@ -4220,11 +4313,12 @@ impl AIConversation {
     pub fn to_serialized_blocklist_items(&self) -> Vec<SerializedBlockListItem> {
         let mut serialized_blocks = Vec::new();
 
-        // Extract all command blocks from the task messages
-        let command_blocks = self.extract_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 {} command blocks for conversation {}",
+            "Extracted {} of {} command blocks for conversation {}",
             command_blocks.len(),
+            total_command_blocks,
             self.id()
         );
 
@@ -4257,8 +4351,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,