public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rig setup

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/action_model/execute.rs
+17 -8
View File
@@ -316,7 +316,7 @@ pub struct BlocklistAIActionExecutor {
wait_for_events_executor: ModelHandle<WaitForEventsExecutor>,
/// The actions currently executing asynchronously, scoped by conversation and action ID.
async_executing_actions: AsyncExecutingActions,
restored_action_ids: HashSet<AIAgentActionId>,
restored_action_ids: HashSet<AsyncExecutingActionKey>,
/// Reference to the terminal model for checking session sharing state.
terminal_model: Arc<FairMutex<TerminalModel>>,
@@ -428,12 +428,18 @@ impl BlocklistAIActionExecutor {
pub fn mark_restored_actions(
&mut self,
conversation_id: AIConversationId,
action_ids: &HashSet<AIAgentActionId>,
ctx: &mut ModelContext<Self>,
) {
self.restored_action_ids.extend(action_ids.iter().cloned());
self.restored_action_ids.extend(
action_ids
.iter()
.cloned()
.map(|action_id| (conversation_id, action_id)),
);
self.run_agents_executor.update(ctx, |executor, _| {
executor.mark_recovery_actions(action_ids);
executor.mark_recovery_actions(conversation_id, action_ids);
});
}
@@ -753,7 +759,8 @@ impl BlocklistAIActionExecutor {
action.id,
std::mem::discriminant(&action.action)
);
let is_restored = self.restored_action_ids.remove(&action.id);
let action_key = (conversation_id, action.id.clone());
let is_restored = self.restored_action_ids.remove(&action_key);
let action_clone = action.clone();
let execution = match &action.action {
AIAgentActionType::RequestCommandOutput { .. }
@@ -1054,7 +1061,7 @@ impl BlocklistAIActionExecutor {
.remove(conversation_id, action_id)
{
let action_kind = AIAgentActionTypeDiscriminants::from(&running.action.action);
log::info!(
crate::ai::tool_diagnostics::tool_debug!(
"Canceling running async action of type {action_kind:?} action_id={action_id:?}, reason={reason:?}"
);
if let Some(backtrace) = crate::ai::tool_diagnostics::capture_backtrace() {
@@ -1075,11 +1082,11 @@ impl BlocklistAIActionExecutor {
});
} else if matches!(running.action.action, AIAgentActionType::RunAgents(..)) {
self.run_agents_executor.update(ctx, |executor, ctx| {
executor.cancel_execution(&running.action.id, ctx);
executor.cancel_execution(conversation_id, &running.action.id, ctx);
});
} else if matches!(running.action.action, AIAgentActionType::StartAgent { .. }) {
self.start_agent_executor.update(ctx, |executor, _| {
executor.cancel_execution(&running.action.id);
executor.cancel_execution(conversation_id, &running.action.id);
});
} else if let AIAgentActionType::WaitForEvents { tool_call_id, .. } =
&running.action.action
@@ -1123,7 +1130,9 @@ impl BlocklistAIActionExecutor {
}
fn should_autoexecute(&self, input: ExecuteActionInput, ctx: &mut ModelContext<Self>) -> bool {
if self.restored_action_ids.contains(&input.action.id)
if self
.restored_action_ids
.contains(&(input.conversation_id, input.action.id.clone()))
|| cfg!(feature = "bedrock_smoke_test")
{
return true;
app/src/ai/blocklist/action_model/execute/run_agents.rs
+49 -19
View File
@@ -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);
app/src/ai/blocklist/action_model/execute/run_agents_tests.rs
+8 -6
View File
@@ -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.
app/src/ai/blocklist/action_model/execute/start_agent.rs
+14 -4
View File
@@ -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(
app/src/ai/blocklist/action_model/execute/start_agent_tests.rs
+60 -1
View File
@@ -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 {