From a7852eb4327749e30fb13206b2914675fdc07583 Mon Sep 17 00:00:00 2001 From: Anthony Ronning <101225832+AnthonyRonning@users.noreply.github.com> Date: Tue, 10 Nov 2026 23:34:43 -0800 Subject: [PATCH] Agent: keep a streamed chunk's cost flat as the transcript grows Every streamed event bounded the transcript window by adding up the weight of every loaded row, cloning each row's id for the lookup. The window now keeps a running byte total (RowWeights): a row that arrives or changes adjusts it, and only a bulk change weighs every row. The limits and what gets evicted are unchanged. With application Vim on, following the stream also collected every row's id to read the last one; it now searches from the end. Co-Authored-By: Claude --- apps/maple-agent/app/src/ui/chat/mod.rs | 4 +- .../maple-agent/app/src/ui/chat/navigation.rs | 11 +- apps/maple-agent/app/src/ui/chat/paging.rs | 221 ++++++++++++++++-- apps/maple-agent/app/src/ui/chat/tests.rs | 22 ++ 4 files changed, 238 insertions(+), 20 deletions(-) diff --git a/apps/maple-agent/app/src/ui/chat/mod.rs b/apps/maple-agent/app/src/ui/chat/mod.rs index fe2ac4888..1149b51b5 100644 --- a/apps/maple-agent/app/src/ui/chat/mod.rs +++ b/apps/maple-agent/app/src/ui/chat/mod.rs @@ -6132,6 +6132,7 @@ impl ChatScreen { } let old = self.timeline[index].id.clone(); self.local_error_text.remove(&old); + self.pages.weights.remove(&old); self.timeline_index.remove(&old); self.timeline[index].id = item.id.clone(); self.timeline_index.insert(item.id.clone(), (index, 0)); @@ -6143,7 +6144,6 @@ impl ChatScreen { { return None; } - self.pages.weights.remove(&item.id); if item.revision > 0 { if let Some((index, _)) = self.timeline_index.get(&item.id) && self.timeline[*index].revision >= item.revision @@ -6237,7 +6237,7 @@ impl ChatScreen { let id = self.timeline[index].id.clone(); let anchor = self.scroll_anchor(); let count = self.timeline.len(); - self.bound_transcript_window(self.list_state.is_following_tail()); + self.bound_after_row_change(index, self.list_state.is_following_tail()); if self.timeline.len() != count { self.rebuild_window(anchor); } diff --git a/apps/maple-agent/app/src/ui/chat/navigation.rs b/apps/maple-agent/app/src/ui/chat/navigation.rs index c11a242d9..9e1a842f9 100644 --- a/apps/maple-agent/app/src/ui/chat/navigation.rs +++ b/apps/maple-agent/app/src/ui/chat/navigation.rs @@ -930,10 +930,17 @@ impl ChatScreen { if !self.application_vim.follows(Some(&task_id)) { return; } - if let Some(item_id) = self.navigable_timeline_ids().last().cloned() { + // Runs for every streamed chunk: look for the newest row from the + // end instead of collecting every row's id. + let newest = self + .timeline + .iter() + .rev() + .find(|item| timeline_item_is_navigable(item)); + if let Some(item) = newest { self.application_vim .transcript_by_task - .insert(task_id, item_id); + .insert(task_id, item.id.clone()); } } diff --git a/apps/maple-agent/app/src/ui/chat/paging.rs b/apps/maple-agent/app/src/ui/chat/paging.rs index 8d1c47f55..9be8e675e 100644 --- a/apps/maple-agent/app/src/ui/chat/paging.rs +++ b/apps/maple-agent/app/src/ui/chat/paging.rs @@ -8,6 +8,7 @@ use crate::ui::task::Failure as UiFailure; use base64::Engine; use gpui::{Context, ListOffset}; use maple_agent::agent::AgentTimelineItem; +use std::collections::hash_map::Entry; use std::collections::{HashMap, HashSet}; const MAX_ROWS: usize = 2000; @@ -22,7 +23,79 @@ pub(super) struct TranscriptPages { pub live_has_more: bool, pub generation: u64, pub expanding: HashSet, - pub weights: HashMap, + pub weights: RowWeights, +} + +/// What each loaded row counts against the byte budget, and the sum. A row +/// that changes adjusts the sum in place, so bounding the window after a +/// streamed chunk never rescans the transcript. +#[derive(Default)] +pub(super) struct RowWeights { + by_id: HashMap, + /// Always the sum of `by_id`. + total: usize, +} + +impl RowWeights { + pub(super) fn total(&self) -> usize { + self.total + } + + /// Record what `item` weighs now, in place of what it weighed before. + pub(super) fn set(&mut self, item: &AgentTimelineItem) { + let bytes = item_bytes(item); + match self.by_id.get_mut(&item.id) { + Some(previous) => { + self.total = self.total - *previous + bytes; + *previous = bytes; + } + None => { + self.by_id.insert(item.id.clone(), bytes); + self.total += bytes; + } + } + } + + /// Forget a row that left the window or took another key. + pub(super) fn remove(&mut self, id: &str) { + if let Some(bytes) = self.by_id.remove(id) { + self.total -= bytes; + } + } + + pub(super) fn clear(&mut self) { + self.by_id.clear(); + self.total = 0; + } + + fn retain(&mut self, mut keep: impl FnMut(&str) -> bool) { + let total = &mut self.total; + self.by_id.retain(|id, bytes| { + let kept = keep(id); + if !kept { + *total -= *bytes; + } + kept + }); + } + + /// Weigh exactly the rows of `timeline`, for when rows arrived or left + /// in bulk. A row weighed before keeps that weight: `set` refreshes it + /// whenever the row changes. + fn weigh(&mut self, timeline: &[AgentTimelineItem]) { + let mut known = std::mem::take(&mut self.by_id); + self.by_id.reserve(timeline.len()); + self.total = 0; + for item in timeline { + let (id, bytes) = known + .remove_entry(&item.id) + .unwrap_or_else(|| (item.id.clone(), item_bytes(item))); + if let Entry::Vacant(slot) = self.by_id.entry(id) { + slot.insert(bytes); + self.total += bytes; + } + } + } } impl TranscriptPages { @@ -199,7 +272,7 @@ impl ChatScreen { .retain(|id| self.timeline_index.contains_key(id)); self.pages .weights - .retain(|id, _| self.timeline_index.contains_key(id)); + .retain(|id| self.timeline_index.contains_key(id)); self.list_state .reset_with_uniform_height(self.timeline.len(), super::TRANSCRIPT_ROW_ESTIMATE); if let Some((id, mut offset)) = anchor @@ -210,20 +283,25 @@ impl ChatScreen { } } + /// Bound the window after rows arrived or left in bulk. pub(super) fn bound_transcript_window(&mut self, keep_tail: bool) { - let mut bytes = self - .timeline - .iter() - .map(|item| { - *self - .pages - .weights - .entry(item.id.clone()) - .or_insert_with(|| item_bytes(item)) - }) - .sum::(); + self.pages.weights.weigh(&self.timeline); + self.evict_past_budget(keep_tail); + } + + /// Bound the window after the row at `index` arrived or changed. Every + /// other row keeps its weight, so this costs the same at any transcript + /// length; streaming calls it once per chunk. + pub(super) fn bound_after_row_change(&mut self, index: usize, keep_tail: bool) { + self.pages.weights.set(&self.timeline[index]); + self.evict_past_budget(keep_tail); + } + + fn evict_past_budget(&mut self, keep_tail: bool) { let mut changed = false; - while self.timeline.len() > MAX_ROWS || (bytes > MAX_BYTES && self.timeline.len() > 1) { + while self.timeline.len() > MAX_ROWS + || (self.pages.weights.total() > MAX_BYTES && self.timeline.len() > 1) + { let item = if keep_tail { self.pages.before = true; self.timeline.remove(0) @@ -232,7 +310,7 @@ impl ChatScreen { self.timeline.pop().unwrap() }; self.local_error_text.remove(&item.id); - bytes = bytes.saturating_sub(item_bytes(&item)); + self.pages.weights.remove(&item.id); changed = true; } if changed { @@ -347,6 +425,7 @@ impl ChatScreen { this.pages.weights.remove(&key); this.local_error_text.remove(&key); this.timeline[index] = item; + this.pages.weights.set(&this.timeline[index]); this.markdown_cache.clear(); this.derived.clear(); this.messages = Default::default(); @@ -373,9 +452,119 @@ mod tests { use super::*; use crate::backend::AgentBackend; use gpui::AppContext; - use maple_agent::agent::{AgentSessionSummary, AgentTaskState}; + use maple_agent::agent::{ + AgentLiveRowIdentity, AgentPersistedRowIdentity, AgentSessionSummary, AgentTaskState, + }; use std::sync::Arc; + fn chat_with_task(cx: &mut gpui::TestAppContext) -> gpui::Entity { + let _settings = crate::settings::SETTINGS_IO_LOCK.lock(); + cx.executor().allow_parking(); + let backend = + Arc::new(AgentBackend::for_test("http://127.0.0.1:9".into(), String::new()).unwrap()); + cx.new(|cx| { + let mut chat = ChatScreen::new_without_start(backend, "user".into(), cx); + chat.selected_session = Some("task".into()); + chat + }) + } + + fn message(text: &str) -> AgentTimelineItem { + serde_json::from_value(serde_json::json!({ + "id": "message", "itemType": "message", "text": text, "createdMs": 1, + "merge": "replace" + })) + .unwrap() + } + + fn stored(row_seq: i64, text: &str) -> AgentTimelineItem { + AgentTimelineItem { + persisted: Some(AgentPersistedRowIdentity { + task_id: "task".into(), + rows_epoch: 0, + row_seq, + update_seq: row_seq, + }), + ..message(text) + } + } + + fn live(occurrence: u64, text: &str) -> AgentTimelineItem { + AgentTimelineItem { + live: Some(AgentLiveRowIdentity { + run_id: "run".into(), + occurrence, + }), + ..message(text) + } + } + + /// The running total must match a fresh sum over the loaded rows. + #[track_caller] + fn assert_weights_exact(chat: &ChatScreen) { + assert_eq!( + chat.pages.weights.total(), + chat.timeline.iter().map(item_bytes).sum::() + ); + assert_eq!(chat.pages.weights.by_id.len(), chat.timeline.len()); + } + + #[gpui::test] + fn row_weights_follow_streamed_rekeyed_and_evicted_rows(cx: &mut gpui::TestAppContext) { + let chat = chat_with_task(cx); + chat.update(cx, |chat, _| { + chat.replace_timeline((1..=3).map(|seq| stored(seq, "history")).collect()); + assert_weights_exact(chat); + + // A live row arrives, grows with each chunk, and can shrink. + for text in ["An", "An answer", "An answer that grows", "Edited"] { + chat.apply_timeline_item("task", live(1, text)); + assert_eq!(chat.timeline.len(), 4); + assert_eq!(chat.timeline[3].text.as_deref(), Some(text)); + assert_weights_exact(chat); + } + + // Its commit moves the same row to its stored key. + let mut committed = stored(4, "Committed"); + committed.live = live(1, "").live; + chat.apply_timeline_item("task", committed); + assert_eq!(chat.timeline.len(), 4); + assert_eq!(chat.timeline[3].id, "stored:task:0:4"); + assert_weights_exact(chat); + + // Rows past the cap push the oldest out of the window. + for occurrence in 2..=(MAX_ROWS as u64 + 2) { + chat.apply_timeline_item("task", live(occurrence, "step")); + } + assert_eq!(chat.timeline.len(), MAX_ROWS); + assert!(chat.pages.before && !chat.pages.after); + assert_eq!(chat.timeline[0].id, "live:run:3"); + assert_weights_exact(chat); + }); + } + + #[gpui::test] + fn a_streamed_row_past_the_byte_budget_evicts_the_oldest_row(cx: &mut gpui::TestAppContext) { + let chat = chat_with_task(cx); + chat.update(cx, |chat, _| { + let large = "x".repeat(MAX_BYTES / 4); + chat.replace_timeline((1..=3).map(|seq| stored(seq, &large)).collect()); + assert_eq!(chat.timeline.len(), 3); + + // The fourth quarter fits only until the row outgrows it. + chat.apply_timeline_item("task", live(1, &large[..large.len() / 2])); + assert_eq!(chat.timeline.len(), 4); + assert!(!chat.pages.before); + + chat.apply_timeline_item("task", live(1, &large)); + assert_eq!(chat.timeline.len(), 3); + assert_eq!(chat.timeline[0].id, "stored:task:0:2"); + assert_eq!(chat.timeline[2].text.as_deref(), Some(large.as_str())); + assert!(chat.pages.before && !chat.pages.after); + assert_weights_exact(chat); + }); + } + #[test] fn a_late_history_reply_cannot_finish_the_next_tasks_request() { let mut pages = TranscriptPages { diff --git a/apps/maple-agent/app/src/ui/chat/tests.rs b/apps/maple-agent/app/src/ui/chat/tests.rs index 7ba8eee9d..70fad5e66 100644 --- a/apps/maple-agent/app/src/ui/chat/tests.rs +++ b/apps/maple-agent/app/src/ui/chat/tests.rs @@ -12338,6 +12338,28 @@ mod state_tests { }); } + #[gpui::test] + fn test_application_vim_follows_a_stream_to_its_newest_selectable_row(cx: &mut TestAppContext) { + cx.executor().allow_parking(); + let chat = screen(cx); + chat.update(cx, |this, cx| { + this.application_vim_enabled = true; + this.replace_timeline(vec![user_item("u1", "question")]); + assert_eq!(this.selected_transcript_id(), Some("u1")); + + this.apply_incoming_item("s1", item("a1", "message", Some("answer")), cx); + assert_eq!(this.selected_transcript_id(), Some("a1")); + + // Rows the transcript does not show never take the selection. + this.apply_incoming_item("s1", item("internal", "internal", Some("state")), cx); + this.apply_incoming_item("s1", todo_item("todo", serde_json::json!([])), cx); + assert_eq!(this.selected_transcript_id(), Some("a1")); + + this.apply_incoming_item("s1", item("tool", "tool", None), cx); + assert_eq!(this.selected_transcript_id(), Some("tool")); + }); + } + #[gpui::test] fn test_returning_from_settings_reclaims_chat_focus_without_a_preference_change( cx: &mut TestAppContext,