Rig setup

This commit is contained in:
Ryan Ward
2026-08-18 01:40:11 -05:00
parent 56e3b51d48
commit a203d359bf
16 changed files with 2495 additions and 89 deletions
@@ -68,8 +68,8 @@ struct ExistingLaunchedAgent {
}
pub struct RunAgentsExecutor {
pending: HashMap<AIAgentActionId, PendingRunAgents>,
recovery_action_ids: HashSet<AIAgentActionId>,
pending: HashMap<(AIConversationId, AIAgentActionId), PendingRunAgents>,
recovery_action_ids: HashSet<(AIConversationId, AIAgentActionId)>,
launched_agents: HashMap<AIConversationId, HashMap<String, ExistingLaunchedAgent>>,
start_agent_executor: ModelHandle<StartAgentExecutor>,
terminal_view_id: EntityId,
@@ -78,10 +78,12 @@ pub struct RunAgentsExecutor {
/// Lifecycle events for in-flight dispatches.
pub enum RunAgentsExecutorEvent {
SpawningStarted {
conversation_id: AIConversationId,
action_id: AIAgentActionId,
snapshot: RunAgentsSpawningSnapshot,
},
SpawningFinished {
conversation_id: AIConversationId,
action_id: AIAgentActionId,
},
ChildConversationCreated {
@@ -127,12 +129,26 @@ impl RunAgentsExecutor {
}
}
pub fn is_pending(&self, action_id: &AIAgentActionId) -> bool {
self.pending.contains_key(action_id)
pub fn is_pending(
&self,
conversation_id: AIConversationId,
action_id: &AIAgentActionId,
) -> bool {
self.pending
.contains_key(&(conversation_id, action_id.clone()))
}
pub fn mark_recovery_actions(&mut self, action_ids: &HashSet<AIAgentActionId>) {
self.recovery_action_ids.extend(action_ids.iter().cloned());
pub fn mark_recovery_actions(
&mut self,
conversation_id: AIConversationId,
action_ids: &HashSet<AIAgentActionId>,
) {
self.recovery_action_ids.extend(
action_ids
.iter()
.cloned()
.map(|action_id| (conversation_id, action_id)),
);
}
pub(crate) fn terminal_view_id(&self) -> EntityId {
@@ -142,18 +158,21 @@ impl RunAgentsExecutor {
/// Cancels the parent tool wait without cancelling independently-running children.
pub(super) fn cancel_execution(
&mut self,
conversation_id: AIConversationId,
action_id: &AIAgentActionId,
ctx: &mut ModelContext<Self>,
) {
self.recovery_action_ids.remove(action_id);
let action_key = (conversation_id, action_id.clone());
self.recovery_action_ids.remove(&action_key);
let detached_dispatches = self.start_agent_executor.update(ctx, |executor, _| {
executor.cancel_dispatches_for_action(action_id)
executor.cancel_dispatches_for_action(conversation_id, action_id)
});
log::info!(
crate::ai::tool_diagnostics::tool_debug!(
"RunAgents cancellation detached {detached_dispatches} pending child dispatch(es) for action {action_id}"
);
if self.pending.remove(action_id).is_some() {
if self.pending.remove(&action_key).is_some() {
ctx.emit(RunAgentsExecutorEvent::SpawningFinished {
conversation_id,
action_id: action_id.clone(),
});
}
@@ -224,7 +243,8 @@ impl RunAgentsExecutor {
) -> async_channel::Receiver<RunAgentsResult> {
let (sender, receiver) = async_channel::bounded(1);
if self.pending.contains_key(&action_id) {
let action_key = (parent_conversation_id, action_id.clone());
if self.pending.contains_key(&action_key) {
log::warn!("RunAgentsExecutor: dispatch reentered for {action_id:?}; rejecting");
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
@@ -265,7 +285,7 @@ impl RunAgentsExecutor {
agent_count: request.agent_run_configs.len(),
};
self.pending
.insert(action_id.clone(), PendingRunAgents::Publishing);
.insert(action_key, PendingRunAgents::Publishing);
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
@@ -280,6 +300,7 @@ impl RunAgentsExecutor {
}),
);
ctx.emit(RunAgentsExecutorEvent::SpawningStarted {
conversation_id: parent_conversation_id,
action_id: action_id.clone(),
snapshot,
});
@@ -294,7 +315,7 @@ impl RunAgentsExecutor {
request
},
move |me, request, ctx| {
if !me.is_pending(&action_id_for_wait) {
if !me.is_pending(parent_conversation_id, &action_id_for_wait) {
return;
}
me.dispatch_children_for_prepared_request(
@@ -320,7 +341,7 @@ impl RunAgentsExecutor {
ctx: &mut ModelContext<Self>,
) -> async_channel::Receiver<RunAgentsResult> {
let (sender, receiver) = async_channel::bounded(1);
if self.pending.contains_key(&action_id) {
if self.is_pending(parent_conversation_id, &action_id) {
let _ = sender.try_send(RunAgentsResult::Cancelled);
return receiver;
}
@@ -333,6 +354,7 @@ impl RunAgentsExecutor {
agent_count: request.agent_run_configs.len(),
};
ctx.emit(RunAgentsExecutorEvent::SpawningStarted {
conversation_id: parent_conversation_id,
action_id: action_id.clone(),
snapshot,
});
@@ -356,8 +378,10 @@ impl RunAgentsExecutor {
sender: async_channel::Sender<RunAgentsResult>,
ctx: &mut ModelContext<Self>,
) {
self.pending
.insert(action_id.clone(), PendingRunAgents::Spawning);
self.pending.insert(
(parent_conversation_id, action_id.clone()),
PendingRunAgents::Spawning,
);
let parent_run_id = BlocklistAIHistoryModel::as_ref(ctx)
.conversation(&parent_conversation_id)
.and_then(|c| c.run_id());
@@ -527,7 +551,7 @@ impl RunAgentsExecutor {
resolved_slots
},
move |me, resolved_slots, ctx| {
if !me.is_pending(&action_id_for_aggr) {
if !me.is_pending(parent_conversation_id_for_result, &action_id_for_aggr) {
return;
}
let timed_out_request_ids = resolved_slots
@@ -603,8 +627,12 @@ impl RunAgentsExecutor {
execution_mode: launched_mode,
agents,
};
me.pending.remove(&action_id_for_aggr);
me.pending.remove(&(
parent_conversation_id_for_result,
action_id_for_aggr.clone(),
));
ctx.emit(RunAgentsExecutorEvent::SpawningFinished {
conversation_id: parent_conversation_id_for_result,
action_id: action_id_for_aggr,
});
let _ = sender.try_send(result);
@@ -624,7 +652,9 @@ impl RunAgentsExecutor {
let mut request = request.clone();
let action_id = id.clone();
let parent_conversation_id = input.conversation_id;
let is_recovery = self.recovery_action_ids.remove(&action_id);
let is_recovery = self
.recovery_action_ids
.remove(&(parent_conversation_id, action_id.clone()));
let recovery_children = if is_recovery {
prepare_recovery_request_for_execution(&mut request, parent_conversation_id, ctx);
@@ -386,7 +386,8 @@ fn recovered_run_agents_reattaches_existing_child_and_launches_only_missing_chil
title: String::new(),
});
state.executor.update(&mut app, |executor, _| {
executor.mark_recovery_actions(&HashSet::from([action.id.clone()]));
executor
.mark_recovery_actions(state.conversation_id, &HashSet::from([action.id.clone()]));
});
let execution = state.executor.update(&mut app, |executor, ctx| {
@@ -481,7 +482,8 @@ fn cancelling_recovered_run_agents_keeps_persisted_child_running() {
});
let action = remote_run_agents_action("oz");
state.executor.update(&mut app, |executor, _| {
executor.mark_recovery_actions(&HashSet::from([action.id.clone()]));
executor
.mark_recovery_actions(state.conversation_id, &HashSet::from([action.id.clone()]));
});
let execution = state.executor.update(&mut app, |executor, ctx| {
executor
@@ -503,7 +505,7 @@ fn cancelling_recovered_run_agents_keeps_persisted_child_running() {
};
state.executor.update(&mut app, |executor, ctx| {
executor.cancel_execution(&action.id, ctx);
executor.cancel_execution(state.conversation_id, &action.id, ctx);
});
let async_result = execute_future.await;
let result = app.update(|ctx| on_complete(async_result, ctx));
@@ -1090,9 +1092,9 @@ fn cancel_during_plan_publication_does_not_dispatch_children() {
// The action is awaiting plan publication, so it's pending but no children dispatched yet.
assert!(matches!(execution, AnyActionExecution::Async { .. }));
state.executor.update(&mut app, |executor, ctx| {
assert!(executor.is_pending(&action_id));
executor.cancel_execution(&action_id, ctx);
assert!(!executor.is_pending(&action_id));
assert!(executor.is_pending(state.conversation_id, &action_id));
executor.cancel_execution(state.conversation_id, &action_id, ctx);
assert!(!executor.is_pending(state.conversation_id, &action_id));
});
// Finish publishing the plan, which resolves the wait the dispatch was blocked on.
@@ -788,12 +788,18 @@ impl StartAgentExecutor {
self.pending.contains_key(&request_id)
}
pub fn cancel_dispatches_for_action(&mut self, action_id: &AIAgentActionId) -> usize {
pub fn cancel_dispatches_for_action(
&mut self,
conversation_id: AIConversationId,
action_id: &AIAgentActionId,
) -> usize {
let request_ids = self
.pending
.iter()
.filter_map(|(request_id, pending)| {
(&pending.action_id == action_id).then_some(*request_id)
(pending.parent_conversation_id == conversation_id
&& &pending.action_id == action_id)
.then_some(*request_id)
})
.collect::<Vec<_>>();
let detached_count = request_ids.len();
@@ -805,8 +811,12 @@ impl StartAgentExecutor {
/// Cancels only the caller's pending tool wait. A child that was already created keeps
/// running independently and remains available in conversation history.
pub(super) fn cancel_execution(&mut self, action_id: &AIAgentActionId) {
self.cancel_dispatches_for_action(action_id);
pub(super) fn cancel_execution(
&mut self,
conversation_id: AIConversationId,
action_id: &AIAgentActionId,
) {
self.cancel_dispatches_for_action(conversation_id, action_id);
}
pub(super) fn preprocess_action(
@@ -986,7 +986,7 @@ fn cancelling_standalone_start_agent_drops_dispatch_but_keeps_child_running() {
});
executor.update(&mut app, |executor, _| {
executor.cancel_execution(&action.id);
executor.cancel_execution(parent_conversation_id, &action.id);
});
executor.read(&app, |executor, _| assert!(executor.pending.is_empty()));
@@ -1006,6 +1006,65 @@ fn cancelling_standalone_start_agent_drops_dispatch_but_keeps_child_running() {
});
}
#[test]
fn cancelling_duplicate_action_id_detaches_only_the_matching_conversation() {
App::test((), |mut app| async move {
initialize_history_persistence_for_tests(&mut app);
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test());
let executor = app.add_model(StartAgentExecutor::new);
let terminal_view_id = EntityId::new();
let first_conversation = history_model.update(&mut app, |history, ctx| {
history.start_new_conversation(terminal_view_id, false, false, false, ctx)
});
let second_conversation = history_model.update(&mut app, |history, ctx| {
history.start_new_conversation(terminal_view_id, false, false, false, ctx)
});
let action = build_start_agent_action(
StartAgentVersion::V1,
StartAgentExecutionMode::local_with_defaults(),
);
let first = executor.update(&mut app, |executor, ctx| {
executor
.execute(
ExecuteActionInput {
action: &action,
conversation_id: first_conversation,
},
ctx,
)
.into()
});
let second = executor.update(&mut app, |executor, ctx| {
executor
.execute(
ExecuteActionInput {
action: &action,
conversation_id: second_conversation,
},
ctx,
)
.into()
});
assert!(matches!(first, AnyActionExecution::Async { .. }));
assert!(matches!(second, AnyActionExecution::Async { .. }));
executor.update(&mut app, |executor, _| {
executor.cancel_execution(first_conversation, &action.id);
assert_eq!(executor.pending.len(), 1);
assert_eq!(
executor
.pending
.values()
.next()
.unwrap()
.parent_conversation_id,
second_conversation
);
});
});
}
#[test]
fn removing_direct_provider_child_resolves_pending_wait() {
App::test((), |mut app| async move {