From 68c913646fa6a489a10343e0f6d2e5cd14156c07 Mon Sep 17 00:00:00 2001 From: Andy Carlson <2yinyang2@gmail.com> Date: Thu, 25 Jun 2026 15:24:20 -0700 Subject: [PATCH 1/3] perf(ai): O(1) exchange_with_id via exchange_id_index --- app/src/ai/agent/conversation.rs | 7 +------ app/src/ai/agent/task_store.rs | 33 ++++++++++++++++++++------------ 2 files changed, 22 insertions(+), 18 deletions(-) diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index e53ba5f15d..a0fa74985e 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -1440,12 +1440,7 @@ impl AIConversation { #[cfg_attr(target_family = "wasm", allow(unused))] pub fn exchange_with_id(&self, exchange_id: AIAgentExchangeId) -> Option<&AIAgentExchange> { - for task in self.task_store.tasks() { - if let Some(exchange) = task.exchanges().find(|exchange| exchange.id == exchange_id) { - return Some(exchange); - } - } - None + self.task_store.exchange_by_id(exchange_id) } /// Returns the exchange that preceded the exchange with the given id, if there is one. diff --git a/app/src/ai/agent/task_store.rs b/app/src/ai/agent/task_store.rs index 70ef8c96f9..09b8052fc0 100644 --- a/app/src/ai/agent/task_store.rs +++ b/app/src/ai/agent/task_store.rs @@ -20,6 +20,8 @@ pub struct TaskStore { root_task_id: TaskId, tasks: HashMap, linearized_refs: Vec, + /// Map from exchange ID to its position in `linearized_refs` for faster lookup. + exchange_id_index: HashMap, /// If the root task was upgraded from an optimistic (client-generated) ID /// to a server-assigned ID, stores the original optimistic ID so that /// deferred event handlers referencing the stale ID can still resolve @@ -33,6 +35,7 @@ impl TaskStore { let mut store = Self { tasks: HashMap::new(), linearized_refs: Vec::new(), + exchange_id_index: HashMap::new(), root_task_id: root_task_id.clone(), optimistic_root_task_id: None, }; @@ -47,6 +50,7 @@ impl TaskStore { let mut store = Self { tasks, linearized_refs: Vec::new(), + exchange_id_index: HashMap::new(), root_task_id, optimistic_root_task_id: None, }; @@ -107,24 +111,15 @@ impl TaskStore { None } - /// Modifies a task via the provided closure and rebuilds the exchange index - /// if exchanges changed. + /// Modifies a task via the provided closure and rebuilds the exchange index. pub fn modify_task( &mut self, task_id: &TaskId, f: impl FnOnce(&mut Task) -> R, ) -> Option { - let exchange_count_before = self.tasks.get(task_id)?.exchanges_len(); let task = self.tasks.get_mut(task_id)?; let result = f(task); - let exchange_count_after = self - .tasks - .get(task_id) - .map(|t| t.exchanges_len()) - .unwrap_or(0); - if exchange_count_before != exchange_count_after { - self.rebuild_linearized_refs_index(); - } + self.rebuild_linearized_refs_index(); Some(result) } @@ -152,6 +147,11 @@ impl TaskStore { self.insert(root_task); } + pub fn exchange_by_id(&self, exchange_id: AIAgentExchangeId) -> Option<&AIAgentExchange> { + let idx = self.exchange_id_index.get(&exchange_id)?; + self.lookup_exchange(self.linearized_refs.get(*idx)?) + } + pub fn first_exchange(&self) -> Option<&AIAgentExchange> { self.linearized_refs .first() @@ -272,7 +272,7 @@ impl TaskStore { pub fn remove(&mut self, task_id: &TaskId) -> Option { let task = self.tasks.remove(task_id)?; - self.linearized_refs.retain(|r| &r.task_id != task_id); + self.rebuild_linearized_refs_index(); Some(task) } @@ -286,6 +286,15 @@ impl TaskStore { /// Rebuilds the linearized index from scratch using DFS traversal. fn rebuild_linearized_refs_index(&mut self) { self.linearized_refs = Self::build_linearized_refs(&self.tasks, &self.root_task_id); + self.exchange_id_index = self + .linearized_refs + .iter() + .enumerate() + .filter_map(|(idx, r)| { + let exchange = self.lookup_exchange(r)?; + Some((exchange.id, idx)) + }) + .collect(); } /// Builds linearized exchange refs via DFS traversal without mutating self. From bba29d60989fde9c34055cbf071d4a0e3b7066fd Mon Sep 17 00:00:00 2001 From: Andy Carlson <2yinyang2@gmail.com> Date: Fri, 26 Jun 2026 16:49:11 -0700 Subject: [PATCH 2/3] keep modify_task performant --- app/src/ai/agent/task_store.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/app/src/ai/agent/task_store.rs b/app/src/ai/agent/task_store.rs index 09b8052fc0..baca1cccd8 100644 --- a/app/src/ai/agent/task_store.rs +++ b/app/src/ai/agent/task_store.rs @@ -111,15 +111,24 @@ impl TaskStore { None } - /// Modifies a task via the provided closure and rebuilds the exchange index. + /// Modifies a task via the provided closure and rebuilds the exchange index if the exchange + /// count changes. pub fn modify_task( &mut self, task_id: &TaskId, f: impl FnOnce(&mut Task) -> R, ) -> Option { let task = self.tasks.get_mut(task_id)?; + let exchange_count_before = task.exchanges_len(); let result = f(task); - self.rebuild_linearized_refs_index(); + let exchange_count_after = self + .tasks + .get(task_id) + .map(|t| t.exchanges_len()) + .unwrap_or(0); + if exchange_count_before != exchange_count_after { + self.rebuild_linearized_refs_index(); + } Some(result) } From 75061d8460e7040a966f753cd5c5c68d6fdce616 Mon Sep 17 00:00:00 2001 From: Andy Carlson <2yinyang2@gmail.com> Date: Fri, 26 Jun 2026 17:14:49 -0700 Subject: [PATCH 3/3] fix --- app/src/ai/agent/conversation.rs | 1 + app/src/ai/agent/task_store.rs | 38 +++++++++++++++++++++----------- 2 files changed, 26 insertions(+), 13 deletions(-) diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index a0fa74985e..1e3424caa9 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -631,6 +631,7 @@ impl AIConversation { task.reassign_exchange_ids(); }); } + self.task_store.rebuild_exchange_id_index(); } pub fn is_viewing_shared_session(&self) -> bool { diff --git a/app/src/ai/agent/task_store.rs b/app/src/ai/agent/task_store.rs index baca1cccd8..428e8227f2 100644 --- a/app/src/ai/agent/task_store.rs +++ b/app/src/ai/agent/task_store.rs @@ -20,8 +20,7 @@ pub struct TaskStore { root_task_id: TaskId, tasks: HashMap, linearized_refs: Vec, - /// Map from exchange ID to its position in `linearized_refs` for faster lookup. - exchange_id_index: HashMap, + exchange_id_index: HashMap, /// If the root task was upgraded from an optimistic (client-generated) ID /// to a server-assigned ID, stores the original optimistic ID so that /// deferred event handlers referencing the stale ID can still resolve @@ -157,8 +156,29 @@ impl TaskStore { } pub fn exchange_by_id(&self, exchange_id: AIAgentExchangeId) -> Option<&AIAgentExchange> { - let idx = self.exchange_id_index.get(&exchange_id)?; - self.lookup_exchange(self.linearized_refs.get(*idx)?) + let exchange_ref = self.exchange_id_index.get(&exchange_id)?; + self.lookup_exchange(exchange_ref) + } + + pub(super) fn rebuild_exchange_id_index(&mut self) { + self.exchange_id_index = self + .tasks + .values() + .flat_map(|task| { + let task_id = task.id().clone(); + task.exchanges() + .enumerate() + .map(move |(exchange_index, exchange)| { + ( + exchange.id, + ExchangeRef { + task_id: task_id.clone(), + exchange_index, + }, + ) + }) + }) + .collect(); } pub fn first_exchange(&self) -> Option<&AIAgentExchange> { @@ -295,15 +315,7 @@ impl TaskStore { /// Rebuilds the linearized index from scratch using DFS traversal. fn rebuild_linearized_refs_index(&mut self) { self.linearized_refs = Self::build_linearized_refs(&self.tasks, &self.root_task_id); - self.exchange_id_index = self - .linearized_refs - .iter() - .enumerate() - .filter_map(|(idx, r)| { - let exchange = self.lookup_exchange(r)?; - Some((exchange.id, idx)) - }) - .collect(); + self.rebuild_exchange_id_index(); } /// Builds linearized exchange refs via DFS traversal without mutating self.