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

fix: recover run agents after restart

Browse Source
This commit is contained in:
rkw6086
2026-08-15 22:37:14 -05:00
parent 93d6172072
commit a5a3361e7f
11 changed files with 731 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/action_model/execute.rs
+30 -9
View File
@@ -24,6 +24,7 @@ pub(super) mod use_computer;
pub(super) mod wait_for_events;
use std::any::Any;
use std::collections::HashSet;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
@@ -295,6 +296,7 @@ pub struct BlocklistAIActionExecutor {
/// We track them per action rather than as a single slot so multiple actions from the same
/// parallel phase can complete independently.
async_executing_actions: std::collections::HashMap<AIAgentActionId, AsyncExecutingAction>,
restored_action_ids: HashSet<AIAgentActionId>,
/// Reference to the terminal model for checking session sharing state.
terminal_model: Arc<FairMutex<TerminalModel>>,
@@ -382,6 +384,7 @@ impl BlocklistAIActionExecutor {
use_computer_executor,
request_computer_use_executor,
async_executing_actions: Default::default(),
restored_action_ids: Default::default(),
terminal_model,
read_skill_executor,
fetch_conversation_executor,
@@ -399,6 +402,17 @@ impl BlocklistAIActionExecutor {
.map(|running| &running.action)
}
pub fn mark_restored_actions(
&mut self,
action_ids: &HashSet<AIAgentActionId>,
ctx: &mut ModelContext<Self>,
) {
self.restored_action_ids.extend(action_ids.iter().cloned());
self.run_agents_executor.update(ctx, |executor, _| {
executor.mark_recovery_actions(action_ids);
});
}
pub(super) fn has_running_ask_user_question(&self, conversation_id: AIConversationId) -> bool {
self.async_executing_actions.values().any(|running| {
running.conversation_id == conversation_id
@@ -710,6 +724,7 @@ impl BlocklistAIActionExecutor {
action.id,
std::mem::discriminant(&action.action)
);
let is_restored = self.restored_action_ids.remove(&action.id);
let action_clone = action.clone();
let execution = match &action.action {
AIAgentActionType::RequestCommandOutput { .. }
@@ -904,10 +919,12 @@ impl BlocklistAIActionExecutor {
conversation_id,
},
);
ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction {
action_id: action_id.clone(),
conversation_id,
});
if !is_restored {
ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction {
action_id: action_id.clone(),
conversation_id,
});
}
log::info!("[tool-debug] try_to_execute_action: spawning ASYNC execution for action_id={:?}", action_id);
ctx.spawn(execute_future, move |me, result, ctx| {
let Some(running) = me.async_executing_actions.remove(&action_id) else {
@@ -933,10 +950,12 @@ impl BlocklistAIActionExecutor {
TryExecuteResult::ExecutedAsync
}
AnyActionExecution::Sync(action_result) => {
ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction {
action_id: action_id.clone(),
conversation_id,
});
if !is_restored {
ctx.emit(BlocklistAIActionExecutorEvent::ExecutingAction {
action_id: action_id.clone(),
conversation_id,
});
}
ctx.emit(BlocklistAIActionExecutorEvent::FinishedAction {
result: Arc::new(AIAgentActionResult {
id: action_id,
@@ -1035,7 +1054,9 @@ impl BlocklistAIActionExecutor {
}
fn should_autoexecute(&self, input: ExecuteActionInput, ctx: &mut ModelContext<Self>) -> bool {
if cfg!(feature = "bedrock_smoke_test") {
if self.restored_action_ids.contains(&input.action.id)
|| cfg!(feature = "bedrock_smoke_test")
{
return true;
}
match input.action.action {
app/src/ai/blocklist/action_model/execute/run_agents.rs
+155 -31
View File
@@ -69,6 +69,7 @@ struct ExistingLaunchedAgent {
pub struct RunAgentsExecutor {
pending: HashMap<AIAgentActionId, PendingRunAgents>,
recovery_action_ids: HashSet<AIAgentActionId>,
launched_agents: HashMap<AIConversationId, HashMap<String, ExistingLaunchedAgent>>,
start_agent_executor: ModelHandle<StartAgentExecutor>,
terminal_view_id: EntityId,
@@ -119,6 +120,7 @@ impl RunAgentsExecutor {
});
Self {
pending: HashMap::new(),
recovery_action_ids: HashSet::new(),
launched_agents: HashMap::new(),
start_agent_executor,
terminal_view_id,
@@ -129,21 +131,25 @@ impl RunAgentsExecutor {
self.pending.contains_key(action_id)
}
pub fn mark_recovery_actions(&mut self, action_ids: &HashSet<AIAgentActionId>) {
self.recovery_action_ids.extend(action_ids.iter().cloned());
}
pub(crate) fn terminal_view_id(&self) -> EntityId {
self.terminal_view_id
}
/// Cancels a pending run so publication completion cannot fan out children.
/// Cancels the parent tool wait without cancelling independently-running children.
pub(super) fn cancel_execution(
&mut self,
action_id: &AIAgentActionId,
ctx: &mut ModelContext<Self>,
) {
if matches!(
self.pending.get(action_id),
Some(PendingRunAgents::Publishing)
) {
self.pending.remove(action_id);
self.recovery_action_ids.remove(action_id);
self.start_agent_executor.update(ctx, |executor, _| {
executor.cancel_dispatches_for_action(action_id);
});
if self.pending.remove(action_id).is_some() {
ctx.emit(RunAgentsExecutorEvent::SpawningFinished {
action_id: action_id.clone(),
});
@@ -276,6 +282,7 @@ impl RunAgentsExecutor {
action_id_for_wait.clone(),
request,
parent_conversation_id,
HashMap::new(),
sender,
ctx,
)
@@ -285,11 +292,48 @@ impl RunAgentsExecutor {
receiver
}
fn dispatch_recovered_run_agents(
&mut self,
action_id: AIAgentActionId,
request: RunAgentsRequest,
parent_conversation_id: AIConversationId,
recovery_children: HashMap<String, AIConversationId>,
ctx: &mut ModelContext<Self>,
) -> async_channel::Receiver<RunAgentsResult> {
let (sender, receiver) = async_channel::bounded(1);
if self.pending.contains_key(&action_id) {
let _ = sender.try_send(RunAgentsResult::Cancelled);
return receiver;
}
if let Err(error) = validate_request(&request) {
let _ = sender.try_send(RunAgentsResult::Failure { error });
return receiver;
}
let snapshot = RunAgentsSpawningSnapshot {
agent_count: request.agent_run_configs.len(),
};
ctx.emit(RunAgentsExecutorEvent::SpawningStarted {
action_id: action_id.clone(),
snapshot,
});
self.dispatch_children_for_prepared_request(
action_id,
request,
parent_conversation_id,
recovery_children,
sender,
ctx,
);
receiver
}
fn dispatch_children_for_prepared_request(
&mut self,
action_id: AIAgentActionId,
request: RunAgentsRequest,
parent_conversation_id: AIConversationId,
mut recovery_children: HashMap<String, AIConversationId>,
sender: async_channel::Sender<RunAgentsResult>,
ctx: &mut ModelContext<Self>,
) {
@@ -329,6 +373,23 @@ impl RunAgentsExecutor {
let mut slots: Vec<ChildSlot> = Vec::with_capacity(agent_run_configs.len());
for cfg in &agent_run_configs {
let normalized_name = normalize_agent_name(&cfg.name)
.expect("validated RunAgents requests have non-empty agent names");
if let Some(child_conversation_id) = recovery_children.remove(&normalized_name) {
let dispatch = self.start_agent_executor.update(ctx, |executor, exec_ctx| {
executor.reattach(
action_id.clone(),
cfg.name.clone(),
parent_conversation_id,
child_conversation_id,
parent_run_id.clone(),
exec_ctx,
)
});
slots.push(ChildSlot::Pending(dispatch));
continue;
}
let prompt = compose_run_agents_child_prompt(&base_prompt, &cfg.prompt);
let mode = match run_agents_to_start_agent_mode(
&run_execution_mode,
@@ -443,6 +504,9 @@ impl RunAgentsExecutor {
outcomes
},
move |me, outcomes, ctx| {
if !me.is_pending(&action_id_for_aggr) {
return;
}
let agents: Vec<RunAgentsAgentOutcome> = agent_run_configs_for_result
.iter()
.zip(outcomes)
@@ -520,32 +584,56 @@ impl RunAgentsExecutor {
let mut request = request.clone();
let action_id = id.clone();
let parent_conversation_id = input.conversation_id;
if let Some(reason) = prepare_request_for_execution(
&mut request,
parent_conversation_id,
self.terminal_view_id,
&self.launched_agents,
ctx,
) {
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Warn,
"RunAgents execution denied",
serde_json::json!({
"event": "run_agents_execution_denied",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"reason": remote_logging::sanitize_error(&reason),
}),
);
return ActionExecution::Sync(AIAgentActionResultType::RunAgents(
RunAgentsResult::Denied { reason },
));
}
let is_recovery = self.recovery_action_ids.remove(&action_id);
let receiver =
self.dispatch_prepared_run_agents(action_id, request, parent_conversation_id, ctx);
let recovery_children = if is_recovery {
prepare_recovery_request_for_execution(&mut request, parent_conversation_id, ctx);
match recovery_children_by_name(parent_conversation_id, ctx) {
Ok(children) => children,
Err(error) => {
return ActionExecution::Sync(AIAgentActionResultType::RunAgents(
RunAgentsResult::Failure { error },
));
}
}
} else {
if let Some(reason) = prepare_request_for_execution(
&mut request,
parent_conversation_id,
self.terminal_view_id,
&self.launched_agents,
ctx,
) {
#[cfg(not(target_family = "wasm"))]
log_run_agents_event(
ctx,
RemoteLogLevel::Warn,
"RunAgents execution denied",
serde_json::json!({
"event": "run_agents_execution_denied",
"action_id": action_id.to_string(),
"parent_conversation_id": parent_conversation_id.to_string(),
"reason": remote_logging::sanitize_error(&reason),
}),
);
return ActionExecution::Sync(AIAgentActionResultType::RunAgents(
RunAgentsResult::Denied { reason },
));
}
HashMap::new()
};
let receiver = if is_recovery {
self.dispatch_recovered_run_agents(
action_id,
request,
parent_conversation_id,
recovery_children,
ctx,
)
} else {
self.dispatch_prepared_run_agents(action_id, request, parent_conversation_id, ctx)
};
ActionExecution::new_async(
async move { receiver.recv().await },
@@ -771,6 +859,42 @@ fn prepare_request_for_execution(
None
}
fn prepare_recovery_request_for_execution(
request: &mut RunAgentsRequest,
parent_conversation_id: AIConversationId,
ctx: &ModelContext<RunAgentsExecutor>,
) {
normalize_request_for_local_execution(request);
resolve_request_from_approved_config(request, parent_conversation_id, ctx);
populate_default_auth_secret_for_execution(request, ctx);
}
fn recovery_children_by_name(
parent_conversation_id: AIConversationId,
ctx: &ModelContext<RunAgentsExecutor>,
) -> Result<HashMap<String, AIConversationId>, String> {
let mut children_by_name = HashMap::new();
for conversation in
BlocklistAIHistoryModel::as_ref(ctx).child_conversations_of(parent_conversation_id)
{
let Some(name) = conversation.agent_name() else {
continue;
};
let Some(normalized_name) = normalize_agent_name(name) else {
continue;
};
if children_by_name
.insert(normalized_name.clone(), conversation.id())
.is_some()
{
return Err(format!(
"Cannot recover child agent '{name}': multiple persisted child conversations have the same name."
));
}
}
Ok(children_by_name)
}
fn duplicate_launched_agents_reason(
request: &RunAgentsRequest,
parent_conversation_id: AIConversationId,
app/src/ai/blocklist/action_model/execute/run_agents_tests.rs
+160
View File
@@ -358,6 +358,166 @@ fn validate_request_rejects_remote_dispatch() {
);
}
#[test]
fn recovered_run_agents_reattaches_existing_child_and_launches_only_missing_child() {
App::test((), |mut app| async move {
let state = initialize_run_agents_test(&mut app, ExecutionMode::Sdk);
let terminal_view_id = EntityId::new();
let history = BlocklistAIHistoryModel::handle(&app);
let existing_child_id = history.update(&mut app, |history, ctx| {
history.start_new_child_conversation(
terminal_view_id,
"child".to_string(),
state.conversation_id,
None,
ctx,
)
});
let captured = subscribe_to_start_agent_requests(&mut app, &state.start_agent_executor);
let mut action = remote_run_agents_action("oz");
let AIAgentActionType::RunAgents(request) = &mut action.action else {
panic!("expected run_agents action");
};
request.agent_run_configs.push(RunAgentsAgentRunConfig {
name: "missing-child".to_string(),
prompt: "Do separate work".to_string(),
title: String::new(),
});
state.executor.update(&mut app, |executor, _| {
executor.mark_recovery_actions(&HashSet::from([action.id.clone()]));
});
let execution = state.executor.update(&mut app, |executor, ctx| {
executor
.execute(
ExecuteActionInput {
action: &action,
conversation_id: state.conversation_id,
},
ctx,
)
.into()
});
let AnyActionExecution::Async {
execute_future,
on_complete,
} = execution
else {
panic!("expected async recovery execution");
};
let missing_request = captured.read(&app, |captured, _| {
assert_eq!(captured.0.len(), 1);
assert_eq!(captured.0[0].name, "missing-child");
captured.0[0].clone()
});
history.update(&mut app, |history, ctx| {
history.update_conversation_status(
terminal_view_id,
existing_child_id,
crate::ai::agent::conversation::ConversationStatus::Success,
ctx,
);
});
let missing_child_id = history.update(&mut app, |history, ctx| {
history.start_new_child_conversation(
terminal_view_id,
"missing-child".to_string(),
state.conversation_id,
None,
ctx,
)
});
history.update(&mut app, |history, ctx| {
history.record_new_conversation_request_complete(
missing_request.id,
missing_child_id,
ctx,
);
history.update_conversation_status(
terminal_view_id,
missing_child_id,
crate::ai::agent::conversation::ConversationStatus::Success,
ctx,
);
});
let async_result = execute_future.await;
let result = app.update(|ctx| on_complete(async_result, ctx));
let AIAgentActionResultType::RunAgents(RunAgentsResult::Launched { agents, .. }) = result
else {
panic!("expected recovered RunAgents result");
};
assert_eq!(agents.len(), 2);
assert!(matches!(
&agents[0].kind,
RunAgentsAgentOutcomeKind::Launched { agent_id }
if agent_id == &existing_child_id.to_string()
));
assert!(matches!(
&agents[1].kind,
RunAgentsAgentOutcomeKind::Launched { agent_id }
if agent_id == &missing_child_id.to_string()
));
});
}
#[test]
fn cancelling_recovered_run_agents_keeps_persisted_child_running() {
App::test((), |mut app| async move {
let state = initialize_run_agents_test(&mut app, ExecutionMode::Sdk);
let terminal_view_id = EntityId::new();
let history = BlocklistAIHistoryModel::handle(&app);
let child_id = history.update(&mut app, |history, ctx| {
history.start_new_child_conversation(
terminal_view_id,
"child".to_string(),
state.conversation_id,
None,
ctx,
)
});
let action = remote_run_agents_action("oz");
state.executor.update(&mut app, |executor, _| {
executor.mark_recovery_actions(&HashSet::from([action.id.clone()]));
});
let execution = state.executor.update(&mut app, |executor, ctx| {
executor
.execute(
ExecuteActionInput {
action: &action,
conversation_id: state.conversation_id,
},
ctx,
)
.into()
});
let AnyActionExecution::Async {
execute_future,
on_complete,
} = execution
else {
panic!("expected async recovery execution");
};
state.executor.update(&mut app, |executor, ctx| {
executor.cancel_execution(&action.id, ctx);
});
let async_result = execute_future.await;
let result = app.update(|ctx| on_complete(async_result, ctx));
assert!(matches!(
result,
AIAgentActionResultType::RunAgents(RunAgentsResult::Cancelled)
));
history.read(&app, |history, _| {
assert!(matches!(
history.conversation(&child_id).map(|child| child.status()),
Some(crate::ai::agent::conversation::ConversationStatus::InProgress)
));
});
});
}
#[test]
fn completion_slots_are_polled_concurrently_and_preserve_request_order() {
App::test((), |_app| async move {
app/src/ai/blocklist/action_model/execute/start_agent.rs
+39
View File
@@ -713,6 +713,45 @@ impl StartAgentExecutor {
}
}
pub fn reattach(
&mut self,
action_id: AIAgentActionId,
name: String,
parent_conversation_id: AIConversationId,
child_conversation_id: AIConversationId,
parent_run_id: Option<String>,
ctx: &mut ModelContext<Self>,
) -> StartAgentDispatch {
let wait_policy = if parent_run_id.is_none() {
StartAgentWaitPolicy::Completion
} else {
StartAgentWaitPolicy::Startup
};
let (sender, receiver) = async_channel::bounded(1);
let request_id = self.next_request_id();
self.pending.insert(
request_id,
PendingStartAgent {
action_id,
run_agents_child_name: Some(name),
parent_conversation_id,
child_conversation_id: Some(child_conversation_id),
sender,
wait_policy,
},
);
self.record_child_conversation(request_id, child_conversation_id, ctx);
StartAgentDispatch {
receiver,
wait_policy,
}
}
pub fn cancel_dispatches_for_action(&mut self, action_id: &AIAgentActionId) {
self.pending
.retain(|_, pending| &pending.action_id != action_id);
}
pub(super) fn preprocess_action(
&mut self,
_action: PreprocessActionInput,
app/src/ai/blocklist/action_model/execute/start_agent_tests.rs
+73
View File
@@ -946,6 +946,79 @@ fn run_agents_dispatch_publishes_only_run_agents_child_link() {
});
}
#[test]
fn reattach_reuses_persisted_child_without_launching_another_agent() {
App::test((), |mut app| async move {
initialize_history_persistence_for_tests(&mut app);
let terminal_view_id = EntityId::new();
let history_model = app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test());
let executor = app.add_model(StartAgentExecutor::new);
let captured_prompts = capture_start_agent_prompts(&mut app, &executor);
let captured_links = capture_run_agents_child_links(&mut app, &executor);
let parent_conversation_id = history_model.update(&mut app, |history_model, ctx| {
history_model.start_new_conversation(terminal_view_id, false, false, false, ctx)
});
let child_conversation_id = history_model.update(&mut app, |history_model, ctx| {
history_model.start_new_child_conversation(
terminal_view_id,
"child".to_string(),
parent_conversation_id,
None,
ctx,
)
});
let action_id = AIAgentActionId::from("run-agents-action".to_string());
let dispatch = executor.update(&mut app, |executor, ctx| {
executor.reattach(
action_id.clone(),
"child".to_string(),
parent_conversation_id,
child_conversation_id,
None,
ctx,
)
});
assert_eq!(dispatch.wait_policy, StartAgentWaitPolicy::Completion);
assert!(matches!(
dispatch.receiver.try_recv(),
Err(async_channel::TryRecvError::Empty)
));
captured_prompts.read(&app, |captured, _| {
assert!(captured.0.is_empty());
});
captured_links.read(&app, |captured, _| {
assert_eq!(
captured.0,
vec![(
action_id,
"child".to_string(),
parent_conversation_id,
child_conversation_id,
)]
);
});
history_model.update(&mut app, |history_model, ctx| {
history_model.update_conversation_status(
terminal_view_id,
child_conversation_id,
ConversationStatus::Success,
ctx,
);
});
assert!(matches!(
dispatch.receiver.try_recv(),
Ok(StartAgentOutcome::Completed { agent_id, .. })
if agent_id == child_conversation_id.to_string()
));
executor.read(&app, |executor, _| {
assert!(executor.pending.is_empty());
});
});
}
#[test]
fn execute_waits_for_direct_provider_child_and_returns_its_output() {
App::test((), |mut app| async move {