Merge branch 'dev/fix-lag'

This commit is contained in:
Ryan Ward
2026-08-20 08:53:48 -05:00
5 changed files with 100 additions and 0 deletions
+11
View File
@@ -793,6 +793,11 @@ impl AIConversation {
) { ) {
use crate::ai::bedrock::convert::{ContentPart, MessageContent}; use crate::ai::bedrock::convert::{ContentPart, MessageContent};
// Cap the archive to prevent unbounded growth. The archive is only used by
// `recall_tool_history` which already truncates individual results to 50K chars,
// so retaining the most recent entries is sufficient for lookup.
const MAX_TOOL_RESULT_ARCHIVE_ENTRIES: usize = 400;
let mut pending_tool_uses: Vec<crate::ai::bedrock::convert::ConversationMessage> = let mut pending_tool_uses: Vec<crate::ai::bedrock::convert::ConversationMessage> =
Vec::new(); Vec::new();
@@ -831,6 +836,12 @@ impl AIConversation {
for tool_use in pending_tool_uses { for tool_use in pending_tool_uses {
self.tool_result_archive.push(tool_use); self.tool_result_archive.push(tool_use);
} }
// Evict oldest entries if the archive exceeds the cap.
if self.tool_result_archive.len() > MAX_TOOL_RESULT_ARCHIVE_ENTRIES {
let excess = self.tool_result_archive.len() - MAX_TOOL_RESULT_ARCHIVE_ENTRIES;
self.tool_result_archive.drain(0..excess);
}
} }
pub fn append_to_bedrock_history( pub fn append_to_bedrock_history(
+53
View File
@@ -49,6 +49,7 @@ use self::execute::{
BlocklistAIActionExecutor, BlocklistAIActionExecutorEvent, NotExecutedReason, BlocklistAIActionExecutor, BlocklistAIActionExecutorEvent, NotExecutedReason,
RunningActionPhase, TryExecuteResult, RunningActionPhase, TryExecuteResult,
}; };
use super::history_model::BlocklistAIHistoryEvent;
use super::BlocklistAIHistoryModel; use super::BlocklistAIHistoryModel;
use crate::ai::agent::conversation::{AIConversationId, ConversationStatus}; use crate::ai::agent::conversation::{AIConversationId, ConversationStatus};
use crate::ai::agent::{ use crate::ai::agent::{
@@ -855,6 +856,29 @@ impl BlocklistAIActionModel {
} }
}); });
let history_model = BlocklistAIHistoryModel::handle(ctx);
ctx.subscribe_to_model(&history_model, |me, _, event, ctx| {
match event {
BlocklistAIHistoryEvent::RemoveConversation {
conversation_id, ..
}
| BlocklistAIHistoryEvent::DeletedConversation {
conversation_id, ..
} => {
me.cleanup_conversation_state(*conversation_id, ctx);
}
BlocklistAIHistoryEvent::ClearedConversationsForTerminalSurface {
cleared_conversation_ids,
..
} => {
for conversation_id in cleared_conversation_ids {
me.cleanup_conversation_state(*conversation_id, ctx);
}
}
_ => {}
}
});
Self { Self {
pending_actions: Default::default(), pending_actions: Default::default(),
finished_action_results: Default::default(), finished_action_results: Default::default(),
@@ -2158,6 +2182,35 @@ impl BlocklistAIActionModel {
}); });
} }
/// Removes all per-conversation state for a conversation that has been removed or cleared.
/// This prevents unbounded growth of `past_action_results` and other maps over long sessions.
fn cleanup_conversation_state(
&mut self,
conversation_id: AIConversationId,
ctx: &mut ModelContext<Self>,
) {
self.past_action_results
.retain(|(conv_id, _), _| *conv_id != conversation_id);
self.provider_tool_executions
.retain(|(conv_id, _), _| *conv_id != conversation_id);
self.denied_permissions
.retain(|(conv_id, _)| *conv_id != conversation_id);
self.pending_actions.remove(&conversation_id);
self.running_actions.remove(&conversation_id);
self.finished_action_results.remove(&conversation_id);
self.finished_tool_results.remove(&conversation_id);
self.provider_finished_action_results
.retain(|(conv_id, _), _| *conv_id != conversation_id);
self.action_order.remove(&conversation_id);
self.pending_preprocessed_actions.remove(&conversation_id);
self.not_ready_actions
.actions
.retain(|(conv_id, _)| *conv_id != conversation_id);
self.executor.update(ctx, |executor, ctx| {
executor.cleanup_conversation(conversation_id, ctx);
});
}
#[cfg(test)] #[cfg(test)]
pub(super) fn push_pending_action_for_test( pub(super) fn push_pending_action_for_test(
&mut self, &mut self,
@@ -1129,6 +1129,23 @@ impl BlocklistAIActionExecutor {
} }
} }
/// Removes all per-conversation retained state from sub-executors.
/// Called when a conversation is removed or cleared from memory.
pub fn cleanup_conversation(
&mut self,
conversation_id: AIConversationId,
ctx: &mut ModelContext<Self>,
) {
self.restored_action_ids
.retain(|(conv_id, _)| *conv_id != conversation_id);
self.async_executing_actions
.0
.retain(|(conv_id, _), _| *conv_id != conversation_id);
self.run_agents_executor.update(ctx, |executor, _| {
executor.cleanup_conversation(conversation_id);
});
}
fn should_autoexecute(&self, input: ExecuteActionInput, ctx: &mut ModelContext<Self>) -> bool { fn should_autoexecute(&self, input: ExecuteActionInput, ctx: &mut ModelContext<Self>) -> bool {
if self if self
.restored_action_ids .restored_action_ids
@@ -155,6 +155,15 @@ impl RunAgentsExecutor {
self.terminal_view_id self.terminal_view_id
} }
/// Removes all per-conversation state for a conversation that has been removed.
pub(super) fn cleanup_conversation(&mut self, conversation_id: AIConversationId) {
self.launched_agents.remove(&conversation_id);
self.pending
.retain(|(conv_id, _), _| *conv_id != conversation_id);
self.recovery_action_ids
.retain(|(conv_id, _)| *conv_id != conversation_id);
}
/// Cancels the parent tool wait without cancelling independently-running children. /// Cancels the parent tool wait without cancelling independently-running children.
pub(super) fn cancel_execution( pub(super) fn cancel_execution(
&mut self, &mut self,
@@ -628,6 +628,16 @@ impl ResponseStream {
let projection_model = params.model.clone(); let projection_model = params.model.clone();
let projection_messages_sent = params.messages_sent.clone(); let projection_messages_sent = params.messages_sent.clone();
// For direct-provider projections, drop the heavy history fields from params
// after logging. Only `tool_results` is needed post-init (for `has_error_tool_results`
// and `tool_result_count`). This prevents retaining the full conversation history
// for the lifetime of the response stream.
let mut params = params;
params.message_history = Vec::new();
params.tool_result_archive = Vec::new();
params.input = Vec::new();
params.global_rules = Vec::new();
params.mcp_context = None;
Self { Self {
id: response_stream_id, id: response_stream_id,
runtime_capabilities: RuntimeCapabilities::provider(), runtime_capabilities: RuntimeCapabilities::provider(),