fix(core/retention): avoid cloning large history during cleanup

code-rs/core/src/codex.rs

@@ -2381,7 +2381,7 @@ impl Session {
     /// This is called when a new user message arrives to keep history manageable
     async fn cleanup_old_status_items(&self) {
         let mut state = self.state.lock().unwrap();
-        let current_items = state.history.contents();
+        let current_items = state.history.take_contents();
 
         let (items_to_keep, stats) = if self.env_ctx_v2 {
             let policy = crate::retention::RetentionPolicy {
@@ -2391,7 +2391,8 @@ impl Session {
                 keep_latest_baseline: self.retention_config.keep_latest_baseline,
             };
 
-            let (kept, retention_stats) = crate::retention::apply_retention_policy(&current_items, &policy);
+            let (kept, retention_stats) =
+                crate::retention::apply_retention_policy_owned(current_items, &policy);
 
             crate::telemetry::global_telemetry().record_retention(&retention_stats);
 
@@ -2408,11 +2409,10 @@ impl Session {
 
             (kept, legacy_stats)
         } else {
-            prune_history_items(&current_items)
+            prune_history_items_owned(current_items)
         };
 
-        state.history = ConversationHistory::new();
-        state.history.record_items(&items_to_keep);
+        state.history.replace_filtered(items_to_keep);
         drop(state);
 
         if stats.any_removed() {
@@ -3761,6 +3761,7 @@ impl CleanupStats {
     }
 }
 
+#[cfg(test)]
 fn prune_history_items(current_items: &[ResponseItem]) -> (Vec<ResponseItem>, CleanupStats) {
     let mut real_user_messages = Vec::new();
     let mut status_messages = Vec::new();
@@ -3926,6 +3927,171 @@ fn prune_history_items(current_items: &[ResponseItem]) -> (Vec<ResponseItem>, Cl
     (items_to_keep, stats)
 }
 
+fn prune_history_items_owned(current_items: Vec<ResponseItem>) -> (Vec<ResponseItem>, CleanupStats) {
+    let mut real_user_messages = Vec::new();
+    let mut status_messages = Vec::new();
+    let mut env_baselines = Vec::new();
+    let mut env_deltas = Vec::new();
+    let mut browser_snapshot_messages = Vec::new();
+
+    const MAX_ENV_DELTAS: usize = 3;
+    const MAX_BROWSER_SNAPSHOTS: usize = 2;
+
+    for (idx, item) in current_items.iter().enumerate() {
+        if let ResponseItem::Message { role, content, .. } = item {
+            if role != "user" {
+                continue;
+            }
+
+            let has_status = content.iter().any(|c| {
+                if let ContentItem::InputText { text } = c {
+                    text.contains("== System Status ==")
+                        || text.contains("Current working directory:")
+                        || text.contains("Git branch:")
+                        || text.contains(ENVIRONMENT_CONTEXT_OPEN_TAG)
+                        || text.contains(ENVIRONMENT_CONTEXT_DELTA_OPEN_TAG)
+                        || text.contains(BROWSER_SNAPSHOT_OPEN_TAG)
+                } else {
+                    false
+                }
+            });
+
+            let has_screenshot = content
+                .iter()
+                .any(|c| matches!(c, ContentItem::InputImage { .. }));
+
+            let has_real_text = content.iter().any(|c| {
+                if let ContentItem::InputText { text } = c {
+                    !text.contains("== System Status ==")
+                        && !text.contains("Current working directory:")
+                        && !text.contains("Git branch:")
+                        && !text.trim().is_empty()
+                        && !text.contains(ENVIRONMENT_CONTEXT_OPEN_TAG)
+                        && !text.contains(ENVIRONMENT_CONTEXT_DELTA_OPEN_TAG)
+                        && !text.contains(BROWSER_SNAPSHOT_OPEN_TAG)
+                } else {
+                    false
+                }
+            });
+
+            let has_env_baseline = content.iter().any(|c| {
+                if let ContentItem::InputText { text } = c {
+                    text.contains(ENVIRONMENT_CONTEXT_OPEN_TAG)
+                        && !text.contains(ENVIRONMENT_CONTEXT_DELTA_OPEN_TAG)
+                } else {
+                    false
+                }
+            });
+
+            let has_env_delta = content.iter().any(|c| {
+                if let ContentItem::InputText { text } = c {
+                    text.contains(ENVIRONMENT_CONTEXT_DELTA_OPEN_TAG)
+                } else {
+                    false
+                }
+            });
+
+            let has_browser_snapshot = content.iter().any(|c| {
+                if let ContentItem::InputText { text } = c {
+                    text.contains(BROWSER_SNAPSHOT_OPEN_TAG)
+                } else {
+                    false
+                }
+            });
+
+            if has_real_text && !has_status && !has_screenshot {
+                real_user_messages.push(idx);
+            } else if has_status || has_screenshot {
+                status_messages.push(idx);
+            }
+
+            if has_env_baseline {
+                env_baselines.push(idx);
+            }
+            if has_env_delta {
+                env_deltas.push(idx);
+            }
+            if has_browser_snapshot {
+                browser_snapshot_messages.push(idx);
+            }
+        }
+    }
+
+    let mut screenshots_to_keep = std::collections::HashSet::new();
+    for &user_idx in real_user_messages.iter().rev().take(2) {
+        for &status_idx in status_messages.iter() {
+            if status_idx > user_idx {
+                if let Some(ResponseItem::Message { content, .. }) = current_items.get(status_idx)
+                {
+                    if content.iter().any(|c| matches!(c, ContentItem::InputImage { .. })) {
+                        screenshots_to_keep.insert(status_idx);
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
+    let baseline_to_keep = env_baselines.last().copied();
+    let env_deltas_to_keep: std::collections::HashSet<usize> = env_deltas
+        .iter()
+        .rev()
+        .take(MAX_ENV_DELTAS)
+        .copied()
+        .collect();
+    let browser_snapshots_to_keep: std::collections::HashSet<usize> = browser_snapshot_messages
+        .iter()
+        .rev()
+        .take(MAX_BROWSER_SNAPSHOTS)
+        .copied()
+        .collect();
+
+    let mut items_to_keep = Vec::new();
+    let mut removed_screenshots = 0usize;
+    let mut removed_status = 0usize;
+
+    for (idx, item) in current_items.into_iter().enumerate() {
+        let keep = if status_messages.contains(&idx) {
+            screenshots_to_keep.contains(&idx)
+                || browser_snapshots_to_keep.contains(&idx)
+                || baseline_to_keep == Some(idx)
+                || env_deltas_to_keep.contains(&idx)
+        } else {
+            true
+        };
+
+        if keep {
+            items_to_keep.push(item);
+        } else if let ResponseItem::Message { content, .. } = &item {
+            if content
+                .iter()
+                .any(|c| matches!(c, ContentItem::InputImage { .. }))
+            {
+                removed_screenshots += 1;
+            } else {
+                removed_status += 1;
+            }
+        }
+    }
+
+    let stats = CleanupStats {
+        removed_screenshots,
+        removed_status,
+        removed_env_baselines: env_baselines
+            .len()
+            .saturating_sub(if baseline_to_keep.is_some() { 1 } else { 0 }),
+        removed_env_deltas: env_deltas.len().saturating_sub(env_deltas_to_keep.len()),
+        removed_browser_snapshots: browser_snapshot_messages
+            .len()
+            .saturating_sub(browser_snapshots_to_keep.len()),
+        kept_recent_screenshots: screenshots_to_keep.len(),
+        kept_env_deltas: env_deltas_to_keep.len(),
+        kept_browser_snapshots: browser_snapshots_to_keep.len(),
+    };
+
+    (items_to_keep, stats)
+}
+
 impl Drop for Session {
     fn drop(&mut self) {
         // Interrupt any running turn when the session is dropped.

code-rs/core/src/conversation_history.rs

@@ -17,6 +17,14 @@ impl ConversationHistory {
         self.items.clone()
     }
 
+    /// Drains and returns the contents of the transcript.
+    ///
+    /// This avoids cloning large history items (e.g., base64 screenshots) when
+    /// callers need to re-apply retention policies.
+    pub(crate) fn take_contents(&mut self) -> Vec<ResponseItem> {
+        std::mem::take(&mut self.items)
+    }
+
     /// `items` is ordered from oldest to newest.
     pub(crate) fn record_items<I>(&mut self, items: I)
     where
@@ -39,6 +47,19 @@ impl ConversationHistory {
         self.items = items;
     }
 
+    /// Replace the entire history, filtering out any items that should not be
+    /// re-sent to the model.
+    pub(crate) fn replace_filtered(&mut self, items: Vec<ResponseItem>) {
+        self.items.clear();
+        self.items.reserve(items.len());
+        for item in items {
+            if !is_api_message(&item) {
+                continue;
+            }
+            self.items.push(item);
+        }
+    }
+
 }
 
 /// Anything that is not a system message or "reasoning" message is considered

code-rs/core/src/retention.rs

@@ -379,6 +379,134 @@ pub fn apply_retention_policy(
     (result, stats)
 }
 
+/// Owned variant of `apply_retention_policy` that avoids cloning large history items.
+///
+/// This is used by core-side cleanup paths where the history is already owned and we
+/// want to prevent transient memory spikes from duplicating large payloads (e.g. base64
+/// screenshots, environment context snapshots).
+pub fn apply_retention_policy_owned(
+    items: Vec<ResponseItem>,
+    policy: &RetentionPolicy,
+) -> (Vec<ResponseItem>, RetentionStats) {
+    let mut stats = RetentionStats::default();
+
+    let categorized: Vec<CategorizedItem> = items
+        .into_iter()
+        .enumerate()
+        .map(|(idx, item)| CategorizedItem::new(item, idx))
+        .collect();
+
+    let mut env_baselines = Vec::new();
+    let mut env_deltas = Vec::new();
+    let mut browser_snapshots = Vec::new();
+    let mut screenshots = Vec::new();
+    let mut status_messages = Vec::new();
+    let mut other_items = Vec::new();
+
+    for cat_item in categorized {
+        match cat_item.category {
+            ItemCategory::EnvBaseline => env_baselines.push(cat_item),
+            ItemCategory::EnvDelta => env_deltas.push(cat_item),
+            ItemCategory::BrowserSnapshot => browser_snapshots.push(cat_item),
+            ItemCategory::Screenshot => screenshots.push(cat_item),
+            ItemCategory::StatusMessage => status_messages.push(cat_item),
+            _ => other_items.push(cat_item),
+        }
+    }
+
+    let mut kept: Vec<CategorizedItem> = Vec::new();
+
+    // 1. Keep all non-categorized items (real messages, assistant, tool calls)
+    for item in other_items {
+        stats.bytes_kept += item.size_bytes;
+        kept.push(item);
+    }
+
+    // 2. Keep only the latest baseline (immutable)
+    if policy.keep_latest_baseline {
+        if let Some(latest) = env_baselines.pop() {
+            stats.bytes_kept += latest.size_bytes;
+            kept.push(latest);
+        }
+    }
+
+    for item in env_baselines {
+        stats.removed_env_baselines += 1;
+        stats.bytes_removed += item.size_bytes;
+    }
+
+    // 3. Keep last N deltas
+    if env_deltas.len() > policy.max_env_deltas {
+        let to_drop = env_deltas.len().saturating_sub(policy.max_env_deltas);
+        stats.removed_env_deltas += to_drop;
+        for item in env_deltas.drain(0..to_drop) {
+            stats.bytes_removed += item.size_bytes;
+        }
+    }
+    stats.kept_env_deltas = env_deltas.len();
+    for item in env_deltas {
+        stats.bytes_kept += item.size_bytes;
+        kept.push(item);
+    }
+
+    // 4. Keep last K browser snapshots
+    if browser_snapshots.len() > policy.max_browser_snapshots {
+        let to_drop = browser_snapshots.len().saturating_sub(policy.max_browser_snapshots);
+        stats.removed_browser_snapshots += to_drop;
+        for item in browser_snapshots.drain(0..to_drop) {
+            stats.bytes_removed += item.size_bytes;
+        }
+    }
+    stats.kept_browser_snapshots = browser_snapshots.len();
+    for item in browser_snapshots {
+        stats.bytes_kept += item.size_bytes;
+        kept.push(item);
+    }
+
+    // 5. Remove all status messages and old screenshots (legacy behavior)
+    for item in status_messages {
+        stats.removed_status += 1;
+        stats.bytes_removed += item.size_bytes;
+    }
+
+    // Keep only the most recent screenshot
+    if let Some(latest) = screenshots.pop() {
+        stats.kept_recent_screenshots = 1;
+        stats.bytes_kept += latest.size_bytes;
+        kept.push(latest);
+        for item in screenshots {
+            stats.removed_screenshots += 1;
+            stats.bytes_removed += item.size_bytes;
+        }
+    }
+
+    // 6. Apply byte budget constraint
+    kept.sort_by_key(|item| item.index);
+    if stats.bytes_kept > policy.max_total_bytes {
+        let mut cumulative_bytes = 0usize;
+        let mut final_items: Vec<CategorizedItem> = Vec::new();
+        let mut bytes_kept = stats.bytes_kept;
+
+        for item in kept.into_iter().rev() {
+            if cumulative_bytes.saturating_add(item.size_bytes) <= policy.max_total_bytes {
+                cumulative_bytes = cumulative_bytes.saturating_add(item.size_bytes);
+                final_items.push(item);
+            } else {
+                stats.dropped_for_budget += 1;
+                stats.bytes_removed += item.size_bytes;
+                bytes_kept = bytes_kept.saturating_sub(item.size_bytes);
+            }
+        }
+
+        stats.bytes_kept = bytes_kept;
+        final_items.sort_by_key(|item| item.index);
+        kept = final_items;
+    }
+
+    let result = kept.into_iter().map(|item| item.item).collect();
+    (result, stats)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;