diff --git a/code-rs/core/src/cgroup.rs b/code-rs/core/src/cgroup.rs index 689810ff0be..84d08ecd4a8 100644 --- a/code-rs/core/src/cgroup.rs +++ b/code-rs/core/src/cgroup.rs @@ -1,3 +1,4 @@ +#[cfg(target_os = "linux")] use std::path::{Path, PathBuf}; #[cfg(target_os = "linux")] @@ -172,24 +173,3 @@ pub(crate) fn best_effort_cleanup_exec_cgroup(pid: u32) { // Only remove the per-pid directory. The parent container stays. let _ = std::fs::remove_dir(&dir); } - -#[cfg(not(target_os = "linux"))] -pub(crate) fn default_exec_memory_max_bytes() -> Option { - None -} - -#[cfg(not(target_os = "linux"))] -pub(crate) fn best_effort_attach_self_to_exec_cgroup(_pid: u32, _memory_max_bytes: u64) {} - -#[cfg(not(target_os = "linux"))] -pub(crate) fn exec_cgroup_oom_killed(_pid: u32) -> Option { - None -} - -#[cfg(not(target_os = "linux"))] -pub(crate) fn exec_cgroup_memory_max_bytes(_pid: u32) -> Option { - None -} - -#[cfg(not(target_os = "linux"))] -pub(crate) fn best_effort_cleanup_exec_cgroup(_pid: u32) {} diff --git a/code-rs/core/src/exec.rs b/code-rs/core/src/exec.rs index 2d294d458b8..a5e038e34d1 100644 --- a/code-rs/core/src/exec.rs +++ b/code-rs/core/src/exec.rs @@ -564,24 +564,31 @@ async fn consume_truncated_output( combined_handle.await.map_err(CodexErr::from)? }; - let mut oom_killed = false; - let mut cgroup_memory_max_bytes: Option = None; - #[cfg(target_os = "linux")] - { - if !timed_out { - if let Some(pid) = pid { - if matches!(exit_status.signal(), Some(SIGKILL_CODE)) - && crate::cgroup::exec_cgroup_oom_killed(pid).unwrap_or(false) - { - oom_killed = true; - cgroup_memory_max_bytes = crate::cgroup::exec_cgroup_memory_max_bytes(pid); + let (oom_killed, cgroup_memory_max_bytes) = { + #[cfg(target_os = "linux")] + { + let mut oom_killed = false; + let mut cgroup_memory_max_bytes: Option = None; + if !timed_out { + if let Some(pid) = pid { + if matches!(exit_status.signal(), Some(SIGKILL_CODE)) + && crate::cgroup::exec_cgroup_oom_killed(pid).unwrap_or(false) + { + oom_killed = true; + cgroup_memory_max_bytes = crate::cgroup::exec_cgroup_memory_max_bytes(pid); + } } } + if let Some(pid) = pid { + crate::cgroup::best_effort_cleanup_exec_cgroup(pid); + } + (oom_killed, cgroup_memory_max_bytes) } - if let Some(pid) = pid { - crate::cgroup::best_effort_cleanup_exec_cgroup(pid); + #[cfg(not(target_os = "linux"))] + { + (false, None) } - } + }; Ok(RawExecToolCallOutput { exit_status, diff --git a/code-rs/core/src/spawn.rs b/code-rs/core/src/spawn.rs index 7d903356a57..f2a744cd92c 100644 --- a/code-rs/core/src/spawn.rs +++ b/code-rs/core/src/spawn.rs @@ -135,8 +135,6 @@ pub(crate) async fn spawn_child_async( StdioPolicy::RedirectForShellTool => crate::cgroup::default_exec_memory_max_bytes(), StdioPolicy::Inherit => None, }; - #[cfg(not(target_os = "linux"))] - let exec_memory_max_bytes: Option = None; cmd.pre_exec(move || { // Start a new process group let _ = libc::setpgid(0, 0); diff --git a/code-rs/tui/src/chatwidget.rs b/code-rs/tui/src/chatwidget.rs index a6aeb8a5433..0ad119ad84f 100644 --- a/code-rs/tui/src/chatwidget.rs +++ b/code-rs/tui/src/chatwidget.rs @@ -1597,6 +1597,9 @@ pub(crate) struct ChatWidget<'a> { active_exec_cell: Option, history_cells: Vec>, // Store all history in memory history_cell_ids: Vec>, + history_id_to_cell_index: HashMap, + history_id_index_dirty: bool, + last_history_state_mismatch_warn: Option, history_live_window: Option<(usize, usize)>, history_frozen_width: u16, history_frozen_count: usize, @@ -4847,10 +4850,12 @@ impl ChatWidget<'_> { self.history_cells = kept_cells; self.history_cell_ids = kept_ids; if removed_any { + self.mark_history_id_index_dirty(); self.invalidate_height_cache(); } } else if !self.history_cell_ids.is_empty() { self.history_cell_ids.clear(); + self.mark_history_id_index_dirty(); } // Send a redraw event to trigger UI update @@ -5143,7 +5148,9 @@ impl ChatWidget<'_> { Self::assign_history_id_inner(&mut cell, history_id); self.history_cells[idx] = cell; if idx < self.history_cell_ids.len() { + let old_id = self.history_cell_ids.get(idx).and_then(|slot| *slot); self.history_cell_ids[idx] = Some(history_id); + self.update_history_id_index_after_assignment(idx, old_id, Some(history_id)); } self.history_frozen_count = self.history_frozen_count.saturating_sub(1); self.update_render_request_seed(idx); @@ -5612,7 +5619,13 @@ impl ChatWidget<'_> { }); if let Some(id) = self.apply_mutation_to_cell_index(idx, mutation) { if idx < self.history_cell_ids.len() { + let old_id = self.history_cell_ids.get(idx).and_then(|slot| *slot); self.history_cell_ids[idx] = Some(id); + self.update_history_id_index_after_assignment( + idx, + old_id, + Some(id), + ); } self.maybe_hide_spinner(); return; @@ -5817,9 +5830,23 @@ impl ChatWidget<'_> { id, ); if idx < self.history_cell_ids.len() { + let old_id = self.history_cell_ids.get(idx).and_then(|slot| *slot); self.history_cell_ids[idx] = Some(id); + self.update_history_id_index_after_assignment(idx, old_id, Some(id)); } else { self.history_cell_ids.push(Some(id)); + if id != HistoryId::ZERO && !self.history_id_index_dirty { + let needs_dirty = match self.history_id_to_cell_index.entry(id) { + Entry::Vacant(entry) => { + entry.insert(idx); + false + } + Entry::Occupied(_) => true, + }; + if needs_dirty { + self.history_id_index_dirty = true; + } + } } self.context_cell_id = Some(id); self.mark_history_dirty(); @@ -6307,9 +6334,14 @@ impl ChatWidget<'_> { } let mut has_wait_running = false; - for (call_id, entry) in self.tools_state.running_custom_tools.iter() { - if let Some(idx) = running_tools::resolve_entry_index(self, entry, &call_id.0) - { + let running_entries: Vec<(ToolCallId, RunningToolEntry)> = self + .tools_state + .running_custom_tools + .iter() + .map(|(call_id, entry)| (call_id.clone(), *entry)) + .collect(); + for (call_id, entry) in running_entries { + if let Some(idx) = running_tools::resolve_entry_index(self, &entry, &call_id.0) { if let Some(cell) = self.history_cells.get(idx).and_then(|c| c .as_any() .downcast_ref::()) @@ -6595,6 +6627,9 @@ impl ChatWidget<'_> { context_last_sequence: None, context_browser_sequence: None, history_cell_ids: Vec::new(), + history_id_to_cell_index: HashMap::new(), + history_id_index_dirty: true, + last_history_state_mismatch_warn: None, history_live_window: None, history_frozen_width: 0, history_frozen_count: 0, @@ -6721,8 +6756,8 @@ impl ChatWidget<'_> { let notice_key = w.next_req_key_top(); let _ = w.history_insert_plain_state_with_key(notice_state, notice_key, "prelude"); if connecting_mcp && !w.test_mode { - // Render connecting status as a separate cell with standard gutter and spacing - w.history_push_top_next_req(history_cell::new_connecting_mcp_status()); + // Render connecting status as a background cell using the background helper. + w.push_background_before_next_output("\nConnecting MCP servers…"); } // Mark welcome as shown to avoid duplicating the Popular commands section // when SessionConfigured arrives shortly after. @@ -6964,6 +6999,9 @@ impl ChatWidget<'_> { context_last_sequence: None, context_browser_sequence: None, history_cell_ids: Vec::new(), + history_id_to_cell_index: HashMap::new(), + history_id_index_dirty: true, + last_history_state_mismatch_warn: None, history_live_window: None, history_frozen_width: 0, history_frozen_count: 0, @@ -9598,6 +9636,24 @@ impl ChatWidget<'_> { self.history_cells.insert(pos, cell); self.history_cell_ids.insert(pos, maybe_id); } + if append { + if let Some(id) = maybe_id.filter(|id| *id != HistoryId::ZERO) { + if !self.history_id_index_dirty { + let needs_dirty = match self.history_id_to_cell_index.entry(id) { + Entry::Vacant(entry) => { + entry.insert(pos); + false + } + Entry::Occupied(_) => true, + }; + if needs_dirty { + self.history_id_index_dirty = true; + } + } + } + } else { + self.mark_history_id_index_dirty(); + } // In terminal mode, App mirrors history lines into the native buffer. // Ensure order vector is also long enough for position after cell insert if self.cell_order_seq.len() < pos { @@ -9810,6 +9866,22 @@ impl ChatWidget<'_> { self.history_cells.insert(pos, cell); self.history_cell_ids.insert(pos, Some(id)); } + if append { + if id != HistoryId::ZERO && !self.history_id_index_dirty { + let needs_dirty = match self.history_id_to_cell_index.entry(id) { + Entry::Vacant(entry) => { + entry.insert(pos); + false + } + Entry::Occupied(_) => true, + }; + if needs_dirty { + self.history_id_index_dirty = true; + } + } + } else { + self.mark_history_id_index_dirty(); + } if self.cell_order_seq.len() < pos { self.cell_order_seq.resize( pos, @@ -10186,7 +10258,81 @@ impl ChatWidget<'_> { } } - fn cell_index_for_history_id(&self, id: HistoryId) -> Option { + fn mark_history_id_index_dirty(&mut self) { + self.history_id_index_dirty = true; + } + + fn rebuild_history_id_index_if_dirty(&mut self) { + if !self.history_id_index_dirty { + return; + } + + self.history_id_to_cell_index.clear(); + self.history_id_to_cell_index + .reserve(self.history_cell_ids.len()); + let mut duplicates = 0usize; + for (idx, maybe_id) in self.history_cell_ids.iter().enumerate() { + let Some(id) = *maybe_id else { + continue; + }; + if id == HistoryId::ZERO { + continue; + } + match self.history_id_to_cell_index.entry(id) { + Entry::Vacant(entry) => { + entry.insert(idx); + } + Entry::Occupied(_) => { + duplicates = duplicates.saturating_add(1); + } + } + } + self.history_id_index_dirty = false; + + if cfg!(debug_assertions) && duplicates > 0 { + tracing::warn!("history id index rebuild saw {duplicates} duplicate ids"); + } + } + + fn update_history_id_index_after_assignment( + &mut self, + idx: usize, + old_id: Option, + new_id: Option, + ) { + if old_id == new_id { + return; + } + if self.history_id_index_dirty { + return; + } + + if let Some(old_id) = old_id.filter(|id| *id != HistoryId::ZERO) { + if let Some(stored_idx) = self.history_id_to_cell_index.get(&old_id).copied() { + if stored_idx == idx { + self.history_id_to_cell_index.remove(&old_id); + } else { + self.mark_history_id_index_dirty(); + return; + } + } + } + + if let Some(new_id) = new_id.filter(|id| *id != HistoryId::ZERO) { + let needs_dirty = match self.history_id_to_cell_index.entry(new_id) { + Entry::Vacant(entry) => { + entry.insert(idx); + false + } + Entry::Occupied(entry) => *entry.get() != idx, + }; + if needs_dirty { + self.history_id_index_dirty = true; + } + } + } + + fn cell_index_for_history_id_slow(&self, id: HistoryId) -> Option { if let Some(idx) = self .history_cell_ids .iter() @@ -10196,13 +10342,74 @@ impl ChatWidget<'_> { } self.history_cells.iter().enumerate().find_map(|(idx, cell)| { - history_cell::record_from_cell(cell.as_ref()) + history_cell::record_from_cell(cell.as_ref()) .map(|record| record.id() == id) .filter(|matched| *matched) .map(|_| idx) }) } + fn cell_index_for_history_id(&mut self, id: HistoryId) -> Option { + if id == HistoryId::ZERO { + return None; + } + + self.rebuild_history_id_index_if_dirty(); + if let Some(idx) = self.history_id_to_cell_index.get(&id).copied() { + if self.history_cell_ids.get(idx) == Some(&Some(id)) { + return Some(idx); + } + + self.mark_history_id_index_dirty(); + self.rebuild_history_id_index_if_dirty(); + if let Some(idx) = self.history_id_to_cell_index.get(&id).copied() { + if self.history_cell_ids.get(idx) == Some(&Some(id)) { + return Some(idx); + } + } + } + + let idx = self.cell_index_for_history_id_slow(id)?; + if !self.history_id_index_dirty && self.history_cell_ids.get(idx) == Some(&Some(id)) { + self.history_id_to_cell_index.insert(id, idx); + } + Some(idx) + } + + fn should_warn_history_state_mismatch(&mut self) -> bool { + let now = Instant::now(); + let should_warn = self + .last_history_state_mismatch_warn + .map(|last| now.duration_since(last) >= Duration::from_secs(10)) + .unwrap_or(true); + if should_warn { + self.last_history_state_mismatch_warn = Some(now); + } + should_warn + } + + #[cfg(test)] + fn debug_validate_history_id_index(&mut self) { + self.rebuild_history_id_index_if_dirty(); + for (idx, maybe_id) in self.history_cell_ids.iter().enumerate() { + let Some(id) = *maybe_id else { + continue; + }; + if id == HistoryId::ZERO { + continue; + } + let mapped = self.history_id_to_cell_index.get(&id).copied(); + assert_eq!(mapped, Some(idx), "history id map missing {id:?}"); + } + for (id, idx) in &self.history_id_to_cell_index { + let stored = self + .history_cell_ids + .get(*idx) + .and_then(|slot| *slot); + assert_eq!(stored, Some(*id), "history id map mismatch for {id:?}"); + } + } + fn update_cell_from_record(&mut self, id: HistoryId, record: HistoryRecord) { if id == HistoryId::ZERO { tracing::debug!("skip update_cell_from_record: zero id"); @@ -10222,15 +10429,19 @@ impl ChatWidget<'_> { } if idx < self.history_cell_ids.len() { + let old_id = self.history_cell_ids.get(idx).and_then(|slot| *slot); self.history_cell_ids[idx] = Some(id); + self.update_history_id_index_after_assignment(idx, old_id, Some(id)); } self.invalidate_height_cache(); self.request_redraw(); } else { - tracing::warn!( - "history-state mismatch: unable to locate cell for id {:?}", - id - ); + if self.should_warn_history_state_mismatch() { + tracing::warn!( + "history-state mismatch: unable to locate cell for id {:?}", + id + ); + } } } @@ -10512,6 +10723,7 @@ impl ChatWidget<'_> { return; } + let old_id = self.history_cell_ids.get(idx).and_then(|slot| *slot); let record_idx = self .record_index_for_cell(idx) .unwrap_or_else(|| self.record_index_for_position(idx)); @@ -10528,6 +10740,10 @@ impl ChatWidget<'_> { } self.history_cells[idx] = cell; + let new_id = self.history_cell_ids.get(idx).and_then(|slot| *slot); + if old_id != new_id { + self.mark_history_id_index_dirty(); + } self.invalidate_height_cache(); self.request_redraw(); self.refresh_explore_trailing_flags(); @@ -10579,6 +10795,9 @@ impl ChatWidget<'_> { if idx < self.history_cell_ids.len() { self.history_cell_ids[idx] = maybe_id; } + if old_id != maybe_id { + self.mark_history_id_index_dirty(); + } if let Some(id) = old_id { self.history_render.invalidate_history_id(id); } else { @@ -10602,6 +10821,7 @@ impl ChatWidget<'_> { return; } + let len_before = self.history_cells.len(); let removed_id = self.history_cell_ids.get(idx).and_then(|id| *id); if let Some(record_idx) = self.record_index_for_cell(idx) { let _ = self @@ -10613,6 +10833,21 @@ impl ChatWidget<'_> { if idx < self.history_cell_ids.len() { self.history_cell_ids.remove(idx); } + if idx + 1 == len_before { + if !self.history_id_index_dirty { + if let Some(id) = removed_id.filter(|id| *id != HistoryId::ZERO) { + if let Some(stored_idx) = self.history_id_to_cell_index.get(&id).copied() { + if stored_idx == idx { + self.history_id_to_cell_index.remove(&id); + } else { + self.history_id_index_dirty = true; + } + } + } + } + } else { + self.mark_history_id_index_dirty(); + } if idx < self.cell_order_seq.len() { self.cell_order_seq.remove(idx); } @@ -11948,7 +12183,7 @@ impl ChatWidget<'_> { } if !force { if let Some(last) = self.history_snapshot_last_flush { - if last.elapsed() < Duration::from_millis(400) { + if last.elapsed() < Duration::from_millis(5000) { return; } } @@ -12515,6 +12750,8 @@ impl ChatWidget<'_> { self.history_cells.clear(); self.history_cell_ids.clear(); + self.history_id_to_cell_index.clear(); + self.history_id_index_dirty = true; self.history_live_window = None; self.history_frozen_width = 0; self.history_frozen_count = 0; @@ -30142,6 +30379,8 @@ use code_core::protocol::OrderMeta; ); chat.history_cells.clear(); chat.history_cell_ids.clear(); + chat.history_id_to_cell_index.clear(); + chat.history_id_index_dirty = true; chat.history_live_window = None; chat.history_frozen_width = 0; chat.history_frozen_count = 0; @@ -31924,6 +32163,69 @@ use code_core::protocol::OrderMeta; ); } + #[test] + fn history_id_index_stays_consistent_under_churn() { + let mut harness = ChatWidgetHarness::new(); + let chat = harness.chat(); + reset_history(chat); + + insert_plain_cell(chat, &["one"]); + insert_plain_cell(chat, &["two"]); + insert_plain_cell(chat, &["three"]); + chat.debug_validate_history_id_index(); + + let target_id = chat + .history_cell_ids + .get(1) + .and_then(|slot| *slot) + .expect("expected a history id at index 1"); + let record = chat + .history_state + .record(target_id) + .cloned() + .expect("expected record for history id"); + chat.update_cell_from_record(target_id, record); + chat.debug_validate_history_id_index(); + + let background = history_cell::new_background_event("bg".to_string()); + let insert_key = OrderKey { + req: 0, + out: -1, + seq: 0, + }; + let idx = chat.history_insert_with_key_global_tagged( + Box::new(background), + insert_key, + "background", + None, + ); + assert_eq!(idx, 0, "expected background insert at head"); + chat.debug_validate_history_id_index(); + + let replace_idx = 1.min(chat.history_cells.len().saturating_sub(1)); + let replacement = history_cell::plain_message_state_from_paragraphs( + PlainMessageKind::Plain, + PlainMessageRole::System, + ["replacement"], + ); + chat.history_replace_at( + replace_idx, + Box::new(history_cell::PlainHistoryCell::from_state(replacement)), + ); + chat.debug_validate_history_id_index(); + + if chat.history_cells.len() > 2 { + chat.history_remove_at(1); + } + chat.debug_validate_history_id_index(); + + if !chat.history_cells.is_empty() { + let last_idx = chat.history_cells.len().saturating_sub(1); + chat.history_remove_at(last_idx); + } + chat.debug_validate_history_id_index(); + } + #[test] fn ordering_keeps_new_answers_after_prior_backgrounds() { let mut harness = ChatWidgetHarness::new(); diff --git a/code-rs/tui/src/chatwidget/running_tools.rs b/code-rs/tui/src/chatwidget/running_tools.rs index b8a7c424dc0..cdfe513914c 100644 --- a/code-rs/tui/src/chatwidget/running_tools.rs +++ b/code-rs/tui/src/chatwidget/running_tools.rs @@ -97,7 +97,7 @@ pub(super) fn rehydrate(chat: &mut ChatWidget<'_>) { } pub(super) fn resolve_entry_index( - chat: &ChatWidget<'_>, + chat: &mut ChatWidget<'_>, entry: &RunningToolEntry, call_id: &str, ) -> Option { diff --git a/code-rs/tui/src/chatwidget/tools.rs b/code-rs/tui/src/chatwidget/tools.rs index c020237aef4..a519a41e947 100644 --- a/code-rs/tui/src/chatwidget/tools.rs +++ b/code-rs/tui/src/chatwidget/tools.rs @@ -64,8 +64,9 @@ pub(super) fn mcp_end(chat: &mut ChatWidget<'_>, ev: McpToolCallEndEvent, key: O .remove(&map_key); let resolved_idx = entry_removed .as_ref() - .and_then(|entry| running_tools::resolve_entry_index(chat, entry, &call_id)) - .or_else(|| running_tools::find_by_call_id(chat, &call_id)); + .and_then(|entry| running_tools::resolve_entry_index(chat, entry, &call_id)); + let resolved_idx = + resolved_idx.or_else(|| running_tools::find_by_call_id(chat, &call_id)); if let Some(idx) = resolved_idx { chat.history_replace_at(idx, Box::new(completed)); diff --git a/code-rs/tui/src/history_cell/mod.rs b/code-rs/tui/src/history_cell/mod.rs index be006fcfd63..1420f49e6a9 100644 --- a/code-rs/tui/src/history_cell/mod.rs +++ b/code-rs/tui/src/history_cell/mod.rs @@ -2242,19 +2242,6 @@ pub(crate) fn new_popular_commands_notice( plain_message_state_from_lines(lines, HistoryCellType::Notice) } -/// Background status cell shown during startup while external MCP servers -/// are being connected. Uses the standard background-event gutter (») -/// and inserts a blank line above the message for visual separation from -/// the Popular commands block. -pub(crate) fn new_connecting_mcp_status() -> BackgroundEventCell { - let record = BackgroundEventRecord { - id: HistoryId::ZERO, - title: String::new(), - description: "\nConnecting MCP servers…".to_string(), - }; - BackgroundEventCell::new(record) -} - pub(crate) fn new_user_prompt(message: String) -> PlainMessageState { let mut lines: Vec> = Vec::new(); lines.push(Line::from("user"));