public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 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/agent/conversation.rs
+16 -12
View File
@@ -2008,19 +2008,23 @@ impl AIConversation {
.sum()
}
pub fn action(&self, action_id: &AIAgentActionId) -> Option<AIAgentAction> {
self.task_store.tasks().find_map(|task| {
task.exchanges().find_map(|exchange| {
let output = exchange.output_status.output()?;
output.get().messages.iter().find_map(|step| match step {
AIAgentOutputMessage {
message: AIAgentOutputMessageType::Action(action),
..
} if &action.id == action_id => Some(action.clone()),
AIAgentOutputMessage { .. } => None,
})
})
})
}
pub fn contains_action(&self, action_id: &AIAgentActionId) -> bool {
self.task_store.tasks().any(|task| {
task.exchanges()
.any(|exchange| {
let Some(output) = exchange.output_status.output()
else {
return false;
};
output.get().messages.iter().any(|step| {
matches!(step, AIAgentOutputMessage{ message: AIAgentOutputMessageType::Action(AIAgentAction { id, .. }), .. } if id == action_id)
})
})
})
self.action(action_id).is_some()
}
/// Returns the exchange ID that contains the given action ID, if any.
app/src/ai/blocklist/action_model.rs
+4
View File
@@ -1546,6 +1546,7 @@ impl BlocklistAIActionModel {
pub(super) fn queue_provider_actions(
&mut self,
actions: Vec<AIAgentAction>,
recovery_action_ids: HashSet<AIAgentActionId>,
conversation_id: AIConversationId,
batch: &PendingToolBatch,
ctx: &mut ModelContext<Self>,
@@ -1562,6 +1563,9 @@ impl BlocklistAIActionModel {
}
}
self.provider_tool_executions.extend(refs);
self.executor.update(ctx, |executor, ctx| {
executor.mark_restored_actions(&recovery_action_ids, ctx);
});
self.queue_actions(actions, conversation_id, ctx);
Ok(())
}
app/src/ai/blocklist/action_model/execute.rs
+22 -1
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,
},
);
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) => {
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
+132 -8
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,6 +584,19 @@ 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 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,
@@ -543,9 +620,20 @@ impl RunAgentsExecutor {
RunAgentsResult::Denied { reason },
));
}
HashMap::new()
};
let receiver =
self.dispatch_prepared_run_agents(action_id, request, parent_conversation_id, ctx);
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 {
app/src/ai/blocklist/controller.rs
+83 -8
View File
@@ -21,9 +21,9 @@ use anyhow::anyhow;
use chrono::{DateTime, Local};
use futures::channel::oneshot;
use galaxy_agent_core::{
turn_control, ExternalWorkId, PendingToolBatch, ProviderRun, ProviderRunFailureKind,
ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState, ToolLoopGuard,
TurnCommand, TurnCommandSender, TurnRequest,
turn_control, ExternalWorkId, PendingToolBatch, PendingToolCallState, ProviderRun,
ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState,
ToolLoopGuard, TurnCommand, TurnCommandSender, TurnRequest,
};
use galaxy_core::assertions::safe_assert;
use input_context::{input_context_for_request, parse_context_attachments};
@@ -930,12 +930,39 @@ fn record_provider_batch_signal(
committed_work_id.as_ref() == Some(work_id) && finished_work_id.as_ref() == Some(work_id)
}
fn recoverable_run_agents_call_ids(
snapshot: &ActiveProviderRunSnapshot,
) -> Result<HashSet<String>, String> {
let ProviderRunState::AwaitingTools { batch } = snapshot.run.state() else {
return Ok(HashSet::new());
};
batch
.calls
.iter()
.filter(|pending| {
matches!(
pending.state,
PendingToolCallState::Executing | PendingToolCallState::RecoveryPending
)
})
.try_fold(HashSet::new(), |mut call_ids, pending| {
let action = snapshot
.action_context
.action_from_tool_call(&pending.call)?;
if matches!(action.action, AIAgentActionType::RunAgents(_)) {
call_ids.insert(pending.call.id.clone());
}
Ok(call_ids)
})
}
fn normalize_restored_provider_snapshot(
snapshot: &mut ActiveProviderRunSnapshot,
) -> Result<(), String> {
let recoverable_call_ids = recoverable_run_agents_call_ids(snapshot)?;
let normalization = snapshot
.run
.normalize_after_restore()
.normalize_after_restore_with_recoverable_calls(&recoverable_call_ids)
.map_err(|error| error.to_string())?;
let interrupted_call_ids = normalization
.interrupted_call_ids
@@ -5404,10 +5431,19 @@ impl BlocklistAIController {
.calls
.iter()
.filter(|pending| pending.state.result().is_none())
.map(|pending| run.action_context.action_from_tool_call(&pending.call))
.map(|pending| {
run.action_context
.action_from_tool_call(&pending.call)
.map(|action| {
(
action,
matches!(pending.state, PendingToolCallState::RecoveryPending),
)
})
})
.collect::<Result<Vec<_>, _>>()
});
let actions = match conversion {
let converted_actions = match conversion {
Some(Ok(actions)) => actions,
Some(Err(message)) => {
self.fail_active_provider_run(conversation_id, message, ctx);
@@ -5418,7 +5454,38 @@ impl BlocklistAIController {
let stream_id = self.active_provider_runs[&conversation_id]
.stream_id
.clone();
for action in &actions {
let mut recovery_action_ids = HashSet::new();
let mut actions = Vec::with_capacity(converted_actions.len());
for (mut action, is_recovery) in converted_actions {
if is_recovery {
let restored_action = BlocklistAIHistoryModel::as_ref(ctx)
.conversation(&conversation_id)
.and_then(|conversation| conversation.action(&action.id));
let Some(restored_action) = restored_action else {
self.fail_active_provider_run(
conversation_id,
format!(
"restored RunAgents action {} is missing from conversation history",
action.id
),
ctx,
);
return;
};
if !matches!(restored_action.action, AIAgentActionType::RunAgents(_)) {
self.fail_active_provider_run(
conversation_id,
format!(
"restored provider action {} no longer matches RunAgents history",
action.id
),
ctx,
);
return;
}
recovery_action_ids.insert(action.id.clone());
action = restored_action;
} else {
let apply_result =
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
history_model.apply_domain_tool_proposal(
@@ -5438,6 +5505,8 @@ impl BlocklistAIController {
return;
}
}
actions.push(action);
}
if let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) {
slot.command_action_refs.extend(
actions
@@ -5464,7 +5533,13 @@ impl BlocklistAIController {
return;
}
let queue_result = self.action_model.update(ctx, |action_model, ctx| {
action_model.queue_provider_actions(actions, conversation_id, &batch, ctx)
action_model.queue_provider_actions(
actions,
recovery_action_ids,
conversation_id,
&batch,
ctx,
)
});
if let Err(error) = queue_result {
self.fail_active_provider_run(conversation_id, error.to_string(), ctx);
app/src/ai/blocklist/controller_tests.rs
+54
View File
@@ -312,6 +312,60 @@ fn restore_normalization_removes_interrupted_command_correlation() {
));
}
#[test]
fn restore_normalization_preserves_executing_run_agents_for_recovery() {
let conversation_id = AIConversationId::new();
let mut snapshot = provider_snapshot(conversation_id);
let Some(ProviderRunStep::CallModel(call)) = snapshot.run.next_step().unwrap() else {
panic!("expected provider model call");
};
snapshot
.run
.accept_model_turn(
&call.work_id,
CompletedModelTurn {
assistant_content: vec![ContentPart::Text("I will run child agents.".to_owned())],
tool_calls: vec![ToolCall {
id: "run-agents-call".to_owned(),
name: "run_agents".to_owned(),
arguments: serde_json::json!({
"summary": "Run child agents",
"base_prompt": "Shared instructions",
"agent_run_configs": [{
"name": "child",
"prompt": "Do work",
"title": "Child",
}],
}),
}],
usage: Usage::default(),
stop_reason: StopReason::Completed,
advertised_tools: BTreeSet::from(["run_agents".to_owned()]),
},
)
.unwrap();
let Some(ProviderRunStep::DispatchTools(batch)) = snapshot.run.next_step().unwrap() else {
panic!("expected provider tool batch");
};
snapshot
.run
.start_tool(&batch.work_id, "run-agents-call")
.unwrap();
let recoverable = super::recoverable_run_agents_call_ids(&snapshot).unwrap();
assert_eq!(recoverable.len(), 1);
assert!(recoverable.contains("run-agents-call"));
super::normalize_restored_provider_snapshot(&mut snapshot).unwrap();
let ProviderRunState::AwaitingTools { batch } = snapshot.run.state() else {
panic!("recovered RunAgents call should keep the provider batch pending");
};
assert!(matches!(
batch.calls[0].state,
galaxy_agent_core::PendingToolCallState::RecoveryPending
));
}
#[test]
fn restore_normalization_reproposes_permission_without_losing_correlation() {
let conversation_id = AIConversationId::new();
crates/galaxy_agent_core/src/provider_run.rs
+37 -6
View File
@@ -139,6 +139,8 @@ pub enum PendingToolCallState {
decision: PermissionDecision,
},
Executing,
/// External work survived a restart and must be reattached before it can complete.
RecoveryPending,
Resolved {
result: ToolResult,
},
@@ -151,7 +153,8 @@ impl PendingToolCallState {
Self::Proposed
| Self::PermissionPending { .. }
| Self::Approved { .. }
| Self::Executing => None,
| Self::Executing
| Self::RecoveryPending => None,
}
}
@@ -161,6 +164,7 @@ impl PendingToolCallState {
Self::PermissionPending { .. } => "permission_pending",
Self::Approved { .. } => "approved",
Self::Executing => "executing",
Self::RecoveryPending => "recovery_pending",
Self::Resolved { .. } => "resolved",
}
}
@@ -300,6 +304,7 @@ pub enum ModelFailureDisposition {
pub struct ProviderRunRestoreNormalization {
pub permission_call_ids_reset: Vec<String>,
pub interrupted_call_ids: Vec<String>,
pub recovery_call_ids: Vec<String>,
pub committed_tool_batch: bool,
}
@@ -499,6 +504,13 @@ impl ProviderRun {
pub fn normalize_after_restore(
&mut self,
) -> Result<ProviderRunRestoreNormalization, ProviderRunProtocolError> {
self.normalize_after_restore_with_recoverable_calls(&HashSet::new())
}
pub fn normalize_after_restore_with_recoverable_calls(
&mut self,
recoverable_call_ids: &HashSet<String>,
) -> Result<ProviderRunRestoreNormalization, ProviderRunProtocolError> {
let ProviderRunState::AwaitingTools { batch } = &mut self.state else {
return Ok(ProviderRunRestoreNormalization::default());
@@ -513,6 +525,19 @@ impl ProviderRun {
.push(pending.call.id.clone());
pending.state = PendingToolCallState::Proposed;
}
PendingToolCallState::Executing
if recoverable_call_ids.contains(&pending.call.id) =>
{
normalization
.recovery_call_ids
.push(pending.call.id.clone());
pending.state = PendingToolCallState::RecoveryPending;
}
PendingToolCallState::RecoveryPending => {
normalization
.recovery_call_ids
.push(pending.call.id.clone());
}
PendingToolCallState::Approved { .. } | PendingToolCallState::Executing => {
normalization
.interrupted_call_ids
@@ -719,7 +744,8 @@ impl ProviderRun {
}
PendingToolCallState::PermissionPending { .. }
| PendingToolCallState::Approved { .. }
| PendingToolCallState::Executing => {
| PendingToolCallState::Executing
| PendingToolCallState::RecoveryPending => {
Err(invalid_tool_transition(call, "permission request"))
}
}
@@ -742,7 +768,8 @@ impl ProviderRun {
}
PendingToolCallState::Proposed
| PendingToolCallState::Approved { .. }
| PendingToolCallState::Executing => {
| PendingToolCallState::Executing
| PendingToolCallState::RecoveryPending => {
return Err(invalid_tool_transition(call, "permission resolution"));
}
};
@@ -782,7 +809,9 @@ impl ProviderRun {
) -> Result<(), ProviderRunProtocolError> {
let call = self.pending_tool_call_mut(work_id, call_id)?;
match &call.state {
PendingToolCallState::Proposed | PendingToolCallState::Approved { .. } => {
PendingToolCallState::Proposed
| PendingToolCallState::Approved { .. }
| PendingToolCallState::RecoveryPending => {
call.state = PendingToolCallState::Executing;
Ok(())
}
@@ -806,7 +835,8 @@ impl ProviderRun {
match &call.state {
PendingToolCallState::Proposed
| PendingToolCallState::Approved { .. }
| PendingToolCallState::Executing => {
| PendingToolCallState::Executing
| PendingToolCallState::RecoveryPending => {
call.state = PendingToolCallState::Resolved { result };
Ok(())
}
@@ -837,7 +867,8 @@ impl ProviderRun {
PendingToolCallState::Proposed
| PendingToolCallState::PermissionPending { .. }
| PendingToolCallState::Approved { .. }
| PendingToolCallState::Executing => {
| PendingToolCallState::Executing
| PendingToolCallState::RecoveryPending => {
call.state = PendingToolCallState::Resolved {
result: ToolResult {
call_id: call_id.to_string(),
crates/galaxy_agent_core/src/provider_run_tests.rs
+67
View File
@@ -758,6 +758,73 @@ fn restore_normalization_reproposes_permissions_and_interrupts_unsafe_tools() {
);
}
#[test]
fn restore_normalization_preserves_selected_executing_tools_for_recovery() {
let mut run = run();
let batch = accept_tool_turn(
&mut run,
tool_turn(
vec![
tool_call("run-agents", "run_agents"),
tool_call("read", "read_files"),
],
&["run_agents", "read_files"],
),
);
run.start_tool(&batch.work_id, "run-agents").unwrap();
run.start_tool(&batch.work_id, "read").unwrap();
let normalization = run
.normalize_after_restore_with_recoverable_calls(&HashSet::from(["run-agents".to_string()]))
.unwrap();
assert_eq!(normalization.recovery_call_ids, vec!["run-agents"]);
assert_eq!(normalization.interrupted_call_ids, vec!["read"]);
let ProviderRunState::AwaitingTools { batch } = run.state() else {
panic!("recovered tool batch should remain pending");
};
assert!(matches!(
batch.calls[0].state,
PendingToolCallState::RecoveryPending
));
assert_eq!(
batch.calls[1].state.result().unwrap().status,
ToolResultStatus::Error
);
}
#[test]
fn recovery_pending_tool_survives_another_restore_and_completes_once() {
let mut run = run();
let batch = accept_tool_turn(
&mut run,
tool_turn(vec![tool_call("run-agents", "run_agents")], &["run_agents"]),
);
run.start_tool(&batch.work_id, "run-agents").unwrap();
run.normalize_after_restore_with_recoverable_calls(&HashSet::from(["run-agents".to_string()]))
.unwrap();
let serialized = serde_json::to_string(&run).unwrap();
let mut restored: ProviderRun = serde_json::from_str(&serialized).unwrap();
let normalization = restored.normalize_after_restore().unwrap();
assert_eq!(normalization.recovery_call_ids, vec!["run-agents"]);
restored
.complete_tool(
&batch.work_id,
successful_result("run-agents", "children completed"),
)
.unwrap();
assert!(matches!(
restored.complete_tool(
&batch.work_id,
successful_result("run-agents", "duplicate completion"),
),
Err(ProviderRunProtocolError::DuplicateToolUpdate { .. })
));
restored.commit_tool_batch(&batch.work_id).unwrap();
assert_eq!(restored.state().phase(), ProviderRunPhase::ReadyToCallModel);
}
#[test]
fn restore_normalization_commits_a_fully_resolved_batch() {
let mut run = run();