Rig setup
This commit is contained in:
@@ -1733,7 +1733,7 @@ impl BlocklistAIActionModel {
|
||||
}
|
||||
self.provider_tool_executions.extend(refs);
|
||||
self.executor.update(ctx, |executor, ctx| {
|
||||
executor.mark_restored_actions(&recovery_action_ids, ctx);
|
||||
executor.mark_restored_actions(conversation_id, &recovery_action_ids, ctx);
|
||||
});
|
||||
self.queue_actions(actions, conversation_id, ctx);
|
||||
Ok(())
|
||||
@@ -1981,7 +1981,7 @@ impl BlocklistAIActionModel {
|
||||
return;
|
||||
};
|
||||
for action in actions_to_cancel.drain(..).collect_vec() {
|
||||
log::info!(
|
||||
crate::ai::tool_diagnostics::tool_debug!(
|
||||
"Canceling pending action of type {:?} conversation_id={conversation_id:?} action_id={:?}, reason={:?}",
|
||||
AIAgentActionTypeDiscriminants::from(&action.action),
|
||||
action.id,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -356,16 +356,17 @@ impl CLISubagentController {
|
||||
drop(terminal_model);
|
||||
|
||||
let provider_accepted_completion = completion.as_ref().is_some_and(|completion| {
|
||||
let provider_completion = PendingProviderCommandCompletion::new(
|
||||
completion.completed_command.block_id.clone(),
|
||||
completion.initial_requested_command_action_id.clone(),
|
||||
completion.completed_command.command.clone(),
|
||||
completion.completed_command.grid_contents.clone(),
|
||||
completion.exit_code,
|
||||
);
|
||||
me.controller.update(ctx, |controller, ctx| {
|
||||
controller.accept_provider_command_completion(
|
||||
controller.offer_provider_command_completion(
|
||||
completion.conversation_id,
|
||||
PendingProviderCommandCompletion::new(
|
||||
completion.completed_command.block_id.clone(),
|
||||
completion.initial_requested_command_action_id.clone(),
|
||||
completion.completed_command.command.clone(),
|
||||
completion.completed_command.grid_contents.clone(),
|
||||
completion.exit_code,
|
||||
),
|
||||
provider_completion,
|
||||
ctx,
|
||||
)
|
||||
})
|
||||
|
||||
@@ -544,7 +544,7 @@ impl ActiveProviderRun {
|
||||
|
||||
const ACTIVE_PROVIDER_RUN_SNAPSHOT_VERSION: u32 = 1;
|
||||
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
struct ProviderProjectionTarget {
|
||||
task_id: TaskId,
|
||||
exchange_id: AIAgentExchangeId,
|
||||
@@ -684,6 +684,40 @@ struct QueuedProviderRun {
|
||||
request_params: api::RequestParams,
|
||||
}
|
||||
|
||||
struct PreparedQueuedProviderRunRestoration {
|
||||
snapshot: QueuedProviderRunSnapshot,
|
||||
root_task_id: TaskId,
|
||||
base_provider_config: crate::ai::provider::ProviderConfig,
|
||||
cli_provider_config: crate::ai::provider::ProviderConfig,
|
||||
request_params: api::RequestParams,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
struct QueuedProviderRunSnapshot {
|
||||
run_id: ProviderRunId,
|
||||
projection_target: ProviderProjectionTarget,
|
||||
did_input_contain_user_query: bool,
|
||||
supported_tools_override: Option<Vec<i32>>,
|
||||
}
|
||||
|
||||
const QUEUED_PROVIDER_RUN_SNAPSHOT_VERSION: u32 = 2;
|
||||
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
struct AbandonedProviderGenerationSnapshot {
|
||||
run_id: ProviderRunId,
|
||||
projection_target: ProviderProjectionTarget,
|
||||
response_stream_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
struct QueuedProviderRunsOnlySnapshot {
|
||||
version: u32,
|
||||
active_run_id: ProviderRunId,
|
||||
#[serde(default)]
|
||||
abandoned_generation: Option<AbandonedProviderGenerationSnapshot>,
|
||||
queued_follow_ups: Vec<QueuedProviderRunSnapshot>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ActiveProviderRunCheckpoint {
|
||||
run: ProviderRun,
|
||||
@@ -751,6 +785,8 @@ struct ActiveProviderRunSnapshot {
|
||||
pending_monitor_observation: Option<PendingProviderMonitorObservation>,
|
||||
pending_command_completion: Option<PendingProviderCommandCompletion>,
|
||||
monitor_prose_continuations: usize,
|
||||
#[serde(default)]
|
||||
queued_follow_ups: Vec<QueuedProviderRunSnapshot>,
|
||||
}
|
||||
|
||||
impl ActiveProviderRunSnapshot {
|
||||
@@ -791,6 +827,7 @@ impl ActiveProviderRunSnapshot {
|
||||
pending_monitor_observation: slot.pending_monitor_observation.clone(),
|
||||
pending_command_completion: slot.pending_command_completion.clone(),
|
||||
monitor_prose_continuations: slot.monitor_prose_continuations,
|
||||
queued_follow_ups: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -803,6 +840,10 @@ impl ActiveProviderRunSnapshot {
|
||||
snapshot.version
|
||||
));
|
||||
}
|
||||
snapshot
|
||||
.run
|
||||
.validate_restored_state()
|
||||
.map_err(|error| error.to_string())?;
|
||||
Ok(snapshot)
|
||||
}
|
||||
|
||||
@@ -997,6 +1038,10 @@ fn recoverable_run_agents_call_ids(
|
||||
fn normalize_restored_provider_snapshot(
|
||||
snapshot: &mut ActiveProviderRunSnapshot,
|
||||
) -> Result<(), String> {
|
||||
snapshot
|
||||
.run
|
||||
.validate_restored_state()
|
||||
.map_err(|error| error.to_string())?;
|
||||
let recoverable_call_ids = recoverable_run_agents_call_ids(snapshot)?;
|
||||
let normalization = snapshot
|
||||
.run
|
||||
@@ -1090,6 +1135,19 @@ fn apply_restored_provider_command_evidence(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn merge_completion_offered_during_restore(
|
||||
prepared: &mut ActiveProviderRunSnapshot,
|
||||
latest: ActiveProviderRunSnapshot,
|
||||
) {
|
||||
if latest.run.id() != prepared.run.id() {
|
||||
return;
|
||||
}
|
||||
prepared.pending_command_completion = latest.pending_command_completion;
|
||||
if prepared.pending_command_completion.is_some() {
|
||||
prepared.pending_monitor_observation = None;
|
||||
}
|
||||
}
|
||||
|
||||
fn restored_projection_was_initialized(
|
||||
has_output: bool,
|
||||
has_server_output_id: bool,
|
||||
@@ -1104,6 +1162,20 @@ fn restored_projection_was_initialized(
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_queued_provider_history(
|
||||
request_params: &mut api::RequestParams,
|
||||
conversation: &AIConversation,
|
||||
) {
|
||||
request_params.tasks = conversation.compute_active_tasks();
|
||||
request_params.root_task_id = Some(conversation.get_root_task_id().to_string());
|
||||
if conversation.is_child_agent_conversation() {
|
||||
request_params.orchestration_enabled = false;
|
||||
}
|
||||
request_params.message_history = conversation.bedrock_message_history().to_vec();
|
||||
request_params.tool_result_archive = conversation.tool_result_archive().to_vec();
|
||||
request_params.progressive_summary = conversation.progressive_summary().map(str::to_owned);
|
||||
}
|
||||
|
||||
fn provider_execution_matches_active_work(
|
||||
run_id: &ProviderRunId,
|
||||
active_work_id: Option<&ExternalWorkId>,
|
||||
@@ -4493,11 +4565,15 @@ impl BlocklistAIController {
|
||||
.all_inputs()
|
||||
.any(|input| input.is_passive_request());
|
||||
|
||||
// Make sure there's no existing response stream for the conversation. If
|
||||
// there is, something has gone wrong.
|
||||
if self
|
||||
// A same-conversation direct-provider follow-up is allowed to create its exchange while
|
||||
// the cancelled generation is still terminalizing. Its provider run is queued below and
|
||||
// cannot take the active slot until cleanup removes the old generation. Other overlapping
|
||||
// streams remain invalid.
|
||||
let has_in_flight_response = self
|
||||
.in_flight_response_streams
|
||||
.has_active_stream_for_conversation(conversation_id, ctx)
|
||||
.has_active_stream_for_conversation(conversation_id, ctx);
|
||||
if has_in_flight_response
|
||||
&& !self.provider_generation_is_terminalizing_for_follow_up(conversation_id)
|
||||
{
|
||||
send_telemetry_from_ctx!(
|
||||
TelemetryEvent::AIInputNotSent {
|
||||
@@ -4764,6 +4840,9 @@ impl BlocklistAIController {
|
||||
cli_provider_config,
|
||||
request_params: request_params.clone(),
|
||||
});
|
||||
if let Err(error) = self.persist_active_provider_run(conversation_data.id, ctx) {
|
||||
log::error!("Failed to persist queued provider follow-up: {error}");
|
||||
}
|
||||
} else {
|
||||
self.active_provider_runs.insert(conversation_data.id, slot);
|
||||
self.prepare_active_provider_run(
|
||||
@@ -4822,6 +4901,22 @@ impl BlocklistAIController {
|
||||
Ok((conversation_data.id, response_stream_id))
|
||||
}
|
||||
|
||||
fn provider_generation_is_terminalizing_for_follow_up(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
) -> bool {
|
||||
self.active_provider_runs
|
||||
.get(&conversation_id)
|
||||
.is_some_and(|slot| {
|
||||
matches!(
|
||||
slot.cancellation_reason,
|
||||
Some(CancellationReason::FollowUpSubmitted {
|
||||
is_for_same_conversation: true,
|
||||
})
|
||||
) && self.in_flight_response_streams.has_stream(&slot.stream_id)
|
||||
})
|
||||
}
|
||||
|
||||
fn schedule_restored_provider_runs(
|
||||
&mut self,
|
||||
conversation_ids: &[AIConversationId],
|
||||
@@ -4875,6 +4970,79 @@ impl BlocklistAIController {
|
||||
self.restoring_provider_runs.remove(&conversation_id);
|
||||
return;
|
||||
};
|
||||
if let Ok(snapshot) = serde_json::from_str::<QueuedProviderRunsOnlySnapshot>(&snapshot_json)
|
||||
{
|
||||
if !matches!(snapshot.version, 1 | QUEUED_PROVIDER_RUN_SNAPSHOT_VERSION) {
|
||||
self.fail_restored_provider_run(
|
||||
conversation_id,
|
||||
format!(
|
||||
"unsupported queued provider run snapshot version {}",
|
||||
snapshot.version
|
||||
),
|
||||
ctx,
|
||||
);
|
||||
return;
|
||||
}
|
||||
let validation = history_model
|
||||
.as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.ok_or_else(|| "queued provider conversation is missing".to_string())
|
||||
.and_then(|conversation| {
|
||||
if conversation.agent_backend() != &AgentBackend::Provider {
|
||||
return Err("queued provider run belongs to a non-provider conversation"
|
||||
.to_string());
|
||||
}
|
||||
Ok(())
|
||||
});
|
||||
if let Err(error) = validation {
|
||||
self.fail_restored_provider_run(conversation_id, error, ctx);
|
||||
return;
|
||||
}
|
||||
let Some(abandoned_generation) = snapshot.abandoned_generation else {
|
||||
self.fail_restored_provider_run(
|
||||
conversation_id,
|
||||
"queued-only provider snapshot is missing abandoned generation identity"
|
||||
.to_string(),
|
||||
ctx,
|
||||
);
|
||||
return;
|
||||
};
|
||||
if abandoned_generation.run_id != snapshot.active_run_id {
|
||||
self.fail_restored_provider_run(
|
||||
conversation_id,
|
||||
"queued provider abandoned generation identity mismatch".to_string(),
|
||||
ctx,
|
||||
);
|
||||
return;
|
||||
}
|
||||
if let Err(error) = validate_queued_provider_run_snapshots(
|
||||
Some(&snapshot.active_run_id),
|
||||
Some(&abandoned_generation.projection_target),
|
||||
&snapshot.queued_follow_ups,
|
||||
) {
|
||||
self.fail_restored_provider_run(conversation_id, error, ctx);
|
||||
return;
|
||||
}
|
||||
if let Err(error) = self.reconcile_abandoned_provider_generation(
|
||||
conversation_id,
|
||||
abandoned_generation,
|
||||
ctx,
|
||||
) {
|
||||
self.fail_restored_provider_run(conversation_id, error, ctx);
|
||||
return;
|
||||
}
|
||||
if let Err(error) = self.restore_queued_provider_follow_ups(
|
||||
conversation_id,
|
||||
snapshot.queued_follow_ups,
|
||||
ctx,
|
||||
) {
|
||||
self.fail_restored_provider_run(conversation_id, error, ctx);
|
||||
return;
|
||||
}
|
||||
self.restoring_provider_runs.remove(&conversation_id);
|
||||
self.start_next_queued_provider_run(conversation_id, ctx);
|
||||
return;
|
||||
}
|
||||
let mut snapshot = match ActiveProviderRunSnapshot::parse(&snapshot_json) {
|
||||
Ok(snapshot) => snapshot,
|
||||
Err(error) => {
|
||||
@@ -4986,6 +5154,63 @@ impl BlocklistAIController {
|
||||
);
|
||||
}
|
||||
|
||||
fn reconcile_abandoned_provider_generation(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
abandoned: AbandonedProviderGenerationSnapshot,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> Result<(), String> {
|
||||
let stream_id = ResponseStreamId::from_persisted(abandoned.response_stream_id);
|
||||
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
|
||||
let existing_target = history_model
|
||||
.conversation(&conversation_id)
|
||||
.ok_or_else(|| "queued provider conversation is missing".to_string())?
|
||||
.provider_projection_target(&stream_id);
|
||||
if existing_target.is_none() {
|
||||
history_model
|
||||
.rebind_provider_projection(
|
||||
conversation_id,
|
||||
&abandoned.projection_target.task_id,
|
||||
abandoned.projection_target.exchange_id,
|
||||
stream_id.clone(),
|
||||
self.terminal_surface_id,
|
||||
ctx,
|
||||
)
|
||||
.map_err(|error| {
|
||||
format!("failed to rebind abandoned provider projection: {error:?}")
|
||||
})?;
|
||||
}
|
||||
let target = history_model
|
||||
.conversation(&conversation_id)
|
||||
.and_then(|conversation| conversation.provider_projection_target(&stream_id))
|
||||
.expect("abandoned provider projection was rebound");
|
||||
if target
|
||||
!= (
|
||||
abandoned.projection_target.task_id.clone(),
|
||||
abandoned.projection_target.exchange_id,
|
||||
)
|
||||
{
|
||||
return Err(
|
||||
"abandoned provider generation projection identity mismatch".to_string()
|
||||
);
|
||||
}
|
||||
history_model.mark_response_stream_cancelled(
|
||||
&stream_id,
|
||||
conversation_id,
|
||||
self.terminal_surface_id,
|
||||
CancellationReason::FollowUpSubmitted {
|
||||
is_for_same_conversation: true,
|
||||
},
|
||||
ctx,
|
||||
);
|
||||
history_model
|
||||
.conversation_mut(&conversation_id)
|
||||
.expect("queued provider conversation was validated")
|
||||
.cleanup_completed_response_stream(&stream_id);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn reconcile_restored_provider_command(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
@@ -5037,10 +5262,12 @@ impl BlocklistAIController {
|
||||
|| self.active_provider_runs.contains_key(&conversation_id)
|
||||
{
|
||||
self.restoring_provider_runs.remove(&conversation_id);
|
||||
self.restoring_provider_command_completions
|
||||
.remove(&conversation_id);
|
||||
return;
|
||||
}
|
||||
let PreparedRestoredProviderRun {
|
||||
snapshot,
|
||||
mut snapshot,
|
||||
profiles,
|
||||
projection_was_initialized,
|
||||
} = match result {
|
||||
@@ -5050,6 +5277,22 @@ impl BlocklistAIController {
|
||||
return;
|
||||
}
|
||||
};
|
||||
// Completion can arrive while provider runtimes are being rebuilt. Reload only that
|
||||
// mailbox from the durable snapshot so the prepared run cannot overwrite the offer.
|
||||
if let Some(latest_snapshot) = BlocklistAIHistoryModel::as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.and_then(AIConversation::active_provider_run_json)
|
||||
.and_then(|json| ActiveProviderRunSnapshot::parse(json).ok())
|
||||
{
|
||||
merge_completion_offered_during_restore(&mut snapshot, latest_snapshot);
|
||||
}
|
||||
if let Some(completion) = self
|
||||
.restoring_provider_command_completions
|
||||
.remove(&conversation_id)
|
||||
{
|
||||
snapshot.pending_command_completion = Some(completion);
|
||||
snapshot.pending_monitor_observation = None;
|
||||
}
|
||||
let ActiveProviderRunSnapshot {
|
||||
version: _,
|
||||
run: provider_run,
|
||||
@@ -5069,6 +5312,7 @@ impl BlocklistAIController {
|
||||
pending_monitor_observation,
|
||||
pending_command_completion,
|
||||
monitor_prose_continuations,
|
||||
queued_follow_ups,
|
||||
} = snapshot;
|
||||
let run_id = provider_run.id().clone();
|
||||
let transcript = provider_run.transcript();
|
||||
@@ -5175,6 +5419,12 @@ impl BlocklistAIController {
|
||||
monitor_prose_continuations,
|
||||
},
|
||||
);
|
||||
if let Err(error) =
|
||||
self.restore_queued_provider_follow_ups(conversation_id, queued_follow_ups, ctx)
|
||||
{
|
||||
self.fail_active_provider_run(conversation_id, error, ctx);
|
||||
return;
|
||||
}
|
||||
self.restoring_provider_runs.remove(&conversation_id);
|
||||
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
|
||||
history_model.update_conversation_status(
|
||||
@@ -5191,6 +5441,194 @@ impl BlocklistAIController {
|
||||
self.resume_restored_provider_run(conversation_id, ctx);
|
||||
}
|
||||
|
||||
fn restore_queued_provider_follow_ups(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
snapshots: Vec<QueuedProviderRunSnapshot>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> Result<(), String> {
|
||||
let (active_run_id, active_projection_target) = self
|
||||
.active_provider_runs
|
||||
.get(&conversation_id)
|
||||
.map_or((None, None), |slot| {
|
||||
(Some(&slot.run_id), Some(&slot.projection_target))
|
||||
});
|
||||
validate_queued_provider_run_snapshots(
|
||||
active_run_id,
|
||||
active_projection_target,
|
||||
&snapshots,
|
||||
)?;
|
||||
let history_model = BlocklistAIHistoryModel::handle(ctx);
|
||||
let mut prepared_runs = Vec::with_capacity(snapshots.len());
|
||||
for snapshot in snapshots {
|
||||
let (request_input, conversation_data, message_history, tool_result_archive, summary) = {
|
||||
let conversation = history_model
|
||||
.as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.ok_or_else(|| "queued provider conversation is missing".to_string())?;
|
||||
let exchange = conversation
|
||||
.get_task(&snapshot.projection_target.task_id)
|
||||
.and_then(|task| task.exchange(snapshot.projection_target.exchange_id))
|
||||
.ok_or_else(|| "queued provider projection exchange is missing".to_string())?;
|
||||
let request_input = RequestInput {
|
||||
conversation_id,
|
||||
input_messages: HashMap::from([(
|
||||
snapshot.projection_target.task_id.clone(),
|
||||
exchange.input.clone(),
|
||||
)]),
|
||||
working_directory: exchange.working_directory.clone(),
|
||||
model_id: exchange.model_id.clone(),
|
||||
coding_model_id: exchange.coding_model_id.clone(),
|
||||
cli_agent_model_id: exchange.cli_agent_model_id.clone(),
|
||||
computer_use_model_id: exchange.computer_use_model_id.clone(),
|
||||
shared_session_response_initiator: exchange.response_initiator.clone(),
|
||||
request_start_ts: exchange.start_time,
|
||||
supported_tools_override: snapshot
|
||||
.supported_tools_override
|
||||
.as_ref()
|
||||
.map(|tools| {
|
||||
tools
|
||||
.iter()
|
||||
.map(|tool| {
|
||||
ToolType::try_from(*tool).map_err(|_| {
|
||||
format!("queued provider tool type {tool} is invalid")
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
})
|
||||
.transpose()?,
|
||||
};
|
||||
let conversation_data = api::ConversationData {
|
||||
id: conversation_id,
|
||||
tasks: conversation.compute_active_tasks(),
|
||||
server_conversation_token: conversation.server_conversation_token().cloned(),
|
||||
forked_from_conversation_token: conversation
|
||||
.forked_from_server_conversation_token()
|
||||
.cloned(),
|
||||
ambient_agent_task_id: self.ambient_agent_task_id,
|
||||
existing_suggestions: history_model
|
||||
.as_ref(ctx)
|
||||
.existing_suggestions_for_conversation(conversation_id)
|
||||
.cloned(),
|
||||
};
|
||||
(
|
||||
request_input,
|
||||
conversation_data,
|
||||
conversation.bedrock_message_history().to_vec(),
|
||||
conversation.tool_result_archive().to_vec(),
|
||||
conversation.progressive_summary().map(str::to_owned),
|
||||
)
|
||||
};
|
||||
let mut request_params = api::RequestParams::new(
|
||||
Some(self.terminal_surface_id),
|
||||
SessionContext::from_session(self.active_session.as_ref(ctx), ctx),
|
||||
&request_input,
|
||||
conversation_data,
|
||||
None,
|
||||
ctx,
|
||||
);
|
||||
request_params.message_history = message_history;
|
||||
request_params.tool_result_archive = tool_result_archive;
|
||||
request_params.progressive_summary = summary;
|
||||
let conversation = BlocklistAIHistoryModel::as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.expect("queued provider conversation was validated");
|
||||
let root_task_id = conversation.get_root_task_id().clone();
|
||||
request_params.root_task_id = Some(root_task_id.to_string());
|
||||
if conversation.is_child_agent_conversation() {
|
||||
request_params.orchestration_enabled = false;
|
||||
}
|
||||
let base_provider_config =
|
||||
ResponseStream::resolve_provider_config(request_params.model.as_str(), ctx);
|
||||
let cli_provider_config = ResponseStream::resolve_provider_config(
|
||||
request_params.cli_agent_model.as_str(),
|
||||
ctx,
|
||||
);
|
||||
prepared_runs.push(PreparedQueuedProviderRunRestoration {
|
||||
snapshot,
|
||||
root_task_id,
|
||||
base_provider_config,
|
||||
cli_provider_config,
|
||||
request_params,
|
||||
});
|
||||
}
|
||||
|
||||
// Do not rebind exchanges or register streams until every queued entry validates and its
|
||||
// request can be rebuilt. A malformed later entry must not make an earlier one executable.
|
||||
for prepared in prepared_runs {
|
||||
let PreparedQueuedProviderRunRestoration {
|
||||
snapshot,
|
||||
root_task_id,
|
||||
base_provider_config,
|
||||
cli_provider_config,
|
||||
request_params,
|
||||
} = prepared;
|
||||
let ai_identifiers = AIIdentifiers {
|
||||
client_conversation_id: Some(conversation_id),
|
||||
model_id: Some(request_params.model.clone()),
|
||||
..AIIdentifiers::default()
|
||||
};
|
||||
let response_stream = ctx.add_model(|ctx| {
|
||||
ResponseStream::new_provider_projection(request_params.clone(), ai_identifiers, ctx)
|
||||
});
|
||||
let stream_id = response_stream.as_ref(ctx).id().clone();
|
||||
let response_stream_clone = response_stream.clone();
|
||||
let did_input_contain_user_query = snapshot.did_input_contain_user_query;
|
||||
ctx.subscribe_to_model(&response_stream, move |me, _, event, ctx| {
|
||||
let _ = me.handle_response_stream_event(
|
||||
did_input_contain_user_query,
|
||||
event,
|
||||
&response_stream_clone,
|
||||
ctx,
|
||||
);
|
||||
});
|
||||
history_model
|
||||
.update(ctx, |history_model, ctx| {
|
||||
history_model.rebind_provider_projection(
|
||||
conversation_id,
|
||||
&snapshot.projection_target.task_id,
|
||||
snapshot.projection_target.exchange_id,
|
||||
stream_id.clone(),
|
||||
self.terminal_surface_id,
|
||||
ctx,
|
||||
)
|
||||
})
|
||||
.map_err(|error| {
|
||||
format!("failed to rebind queued provider projection: {error:?}")
|
||||
})?;
|
||||
self.in_flight_response_streams
|
||||
.register_additional_stream(stream_id.clone(), response_stream.clone());
|
||||
self.queued_provider_runs
|
||||
.entry(conversation_id)
|
||||
.or_default()
|
||||
.push_back(QueuedProviderRun {
|
||||
slot: ActiveProviderRunSlot {
|
||||
stream_id,
|
||||
response_stream,
|
||||
did_input_contain_user_query,
|
||||
run_id: snapshot.run_id,
|
||||
root_task_id,
|
||||
projection_target: snapshot.projection_target,
|
||||
run: None,
|
||||
checkpoint: None,
|
||||
turn_control: None,
|
||||
cancellation_reason: None,
|
||||
committed_provider_batch: None,
|
||||
finished_provider_batch: None,
|
||||
command_action_refs: HashMap::new(),
|
||||
command_monitor: None,
|
||||
pending_monitor_observation: None,
|
||||
pending_command_completion: None,
|
||||
monitor_prose_continuations: 0,
|
||||
},
|
||||
base_provider_config,
|
||||
cli_provider_config,
|
||||
request_params,
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn resume_restored_provider_run(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
@@ -5253,6 +5691,8 @@ impl BlocklistAIController {
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
self.restoring_provider_runs.remove(&conversation_id);
|
||||
self.restoring_provider_command_completions
|
||||
.remove(&conversation_id);
|
||||
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
|
||||
history_model.update_conversation_status_with_error(
|
||||
self.terminal_surface_id,
|
||||
@@ -5382,7 +5822,42 @@ impl BlocklistAIController {
|
||||
.active_provider_runs
|
||||
.get(&conversation_id)
|
||||
.ok_or_else(|| "active provider run disappeared before persistence".to_string())?;
|
||||
let snapshot = ActiveProviderRunSnapshot::from_slot(slot)?;
|
||||
let queued_follow_ups = self.queued_provider_run_snapshots(conversation_id);
|
||||
let mut snapshot = match ActiveProviderRunSnapshot::from_slot(slot) {
|
||||
Ok(snapshot) => snapshot,
|
||||
Err(error) if slot.run.is_none() && slot.checkpoint.is_none() => {
|
||||
let persisted = BlocklistAIHistoryModel::as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.and_then(AIConversation::active_provider_run_json);
|
||||
if let Some(persisted) = persisted {
|
||||
if let Ok(snapshot) = ActiveProviderRunSnapshot::parse(persisted) {
|
||||
if snapshot.run.id() != &slot.run_id {
|
||||
return Err(
|
||||
"persisted provider run identity does not match active slot".into(),
|
||||
);
|
||||
}
|
||||
snapshot
|
||||
} else {
|
||||
return self.persist_queued_provider_runs_only(
|
||||
conversation_id,
|
||||
slot,
|
||||
queued_follow_ups,
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
let _ = error;
|
||||
return self.persist_queued_provider_runs_only(
|
||||
conversation_id,
|
||||
slot,
|
||||
queued_follow_ups,
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
snapshot.queued_follow_ups = queued_follow_ups;
|
||||
let json = serde_json::to_string(&snapshot)
|
||||
.map_err(|error| format!("failed to serialize active provider run: {error}"))?;
|
||||
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
|
||||
@@ -5392,6 +5867,55 @@ impl BlocklistAIController {
|
||||
})
|
||||
}
|
||||
|
||||
fn queued_provider_run_snapshots(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
) -> Vec<QueuedProviderRunSnapshot> {
|
||||
self.queued_provider_runs
|
||||
.get(&conversation_id)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.map(|queued| QueuedProviderRunSnapshot {
|
||||
run_id: queued.slot.run_id.clone(),
|
||||
projection_target: queued.slot.projection_target.clone(),
|
||||
did_input_contain_user_query: queued.slot.did_input_contain_user_query,
|
||||
supported_tools_override: queued
|
||||
.request_params
|
||||
.supported_tools_override
|
||||
.as_ref()
|
||||
.map(|tools| tools.iter().map(|tool| *tool as i32).collect()),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn persist_queued_provider_runs_only(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
active_slot: &ActiveProviderRunSlot,
|
||||
queued_follow_ups: Vec<QueuedProviderRunSnapshot>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> Result<(), String> {
|
||||
if queued_follow_ups.is_empty() {
|
||||
return Err("provider run is not prepared".to_string());
|
||||
}
|
||||
let json = serde_json::to_string(&QueuedProviderRunsOnlySnapshot {
|
||||
version: QUEUED_PROVIDER_RUN_SNAPSHOT_VERSION,
|
||||
active_run_id: active_slot.run_id.clone(),
|
||||
abandoned_generation: Some(AbandonedProviderGenerationSnapshot {
|
||||
run_id: active_slot.run_id.clone(),
|
||||
projection_target: active_slot.projection_target.clone(),
|
||||
response_stream_id: active_slot.stream_id.as_str().to_owned(),
|
||||
}),
|
||||
queued_follow_ups,
|
||||
})
|
||||
.map_err(|error| format!("failed to serialize queued provider runs: {error}"))?;
|
||||
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
|
||||
history_model
|
||||
.persist_active_provider_run_json(conversation_id, Some(json), ctx)
|
||||
.map_err(|error| format!("failed to persist queued provider runs: {error:?}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn clear_persisted_active_provider_run(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
@@ -6538,8 +7062,10 @@ impl BlocklistAIController {
|
||||
{
|
||||
return;
|
||||
}
|
||||
if let Err(error) = self.clear_persisted_active_provider_run(conversation_id, ctx) {
|
||||
log::error!("Failed to clear persisted provider run during cleanup: {error}");
|
||||
if !self.queued_provider_runs.contains_key(&conversation_id) {
|
||||
if let Err(error) = self.clear_persisted_active_provider_run(conversation_id, ctx) {
|
||||
log::error!("Failed to clear persisted provider run during cleanup: {error}");
|
||||
}
|
||||
}
|
||||
self.active_provider_runs.remove(&conversation_id);
|
||||
self.restoring_provider_runs.remove(&conversation_id);
|
||||
@@ -6566,6 +7092,11 @@ impl BlocklistAIController {
|
||||
conversation_id: AIConversationId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
// The old generation must be gone before its successor can own the conversation slot.
|
||||
// This also makes delayed cleanup callbacks harmless: cleanup checks the stream identity.
|
||||
if self.active_provider_runs.contains_key(&conversation_id) {
|
||||
return;
|
||||
}
|
||||
let next = self
|
||||
.queued_provider_runs
|
||||
.get_mut(&conversation_id)
|
||||
@@ -6580,6 +7111,12 @@ impl BlocklistAIController {
|
||||
let Some(next) = next else {
|
||||
return;
|
||||
};
|
||||
let mut request_params = next.request_params;
|
||||
if let Some(conversation) =
|
||||
BlocklistAIHistoryModel::as_ref(ctx).conversation(&conversation_id)
|
||||
{
|
||||
refresh_queued_provider_history(&mut request_params, conversation);
|
||||
}
|
||||
let stream_id = next.slot.stream_id.clone();
|
||||
self.active_provider_runs.insert(conversation_id, next.slot);
|
||||
self.prepare_active_provider_run(
|
||||
@@ -6587,7 +7124,7 @@ impl BlocklistAIController {
|
||||
stream_id,
|
||||
next.base_provider_config,
|
||||
next.cli_provider_config,
|
||||
next.request_params,
|
||||
request_params,
|
||||
ctx,
|
||||
);
|
||||
}
|
||||
@@ -6689,6 +7226,19 @@ impl BlocklistAIController {
|
||||
self.active_provider_runs.contains_key(&conversation_id)
|
||||
}
|
||||
|
||||
pub(super) fn offer_provider_command_completion(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
completion: PendingProviderCommandCompletion,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> bool {
|
||||
if self.active_provider_runs.contains_key(&conversation_id) {
|
||||
self.accept_provider_command_completion(conversation_id, completion, ctx)
|
||||
} else {
|
||||
self.persist_restoring_provider_command_completion(conversation_id, completion, ctx)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn accept_provider_command_completion(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
@@ -6778,6 +7328,64 @@ impl BlocklistAIController {
|
||||
true
|
||||
}
|
||||
|
||||
pub(super) fn persist_restoring_provider_command_completion(
|
||||
&mut self,
|
||||
conversation_id: AIConversationId,
|
||||
mut completion: PendingProviderCommandCompletion,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) -> bool {
|
||||
if self.active_provider_runs.contains_key(&conversation_id) {
|
||||
return self.accept_provider_command_completion(conversation_id, completion, ctx);
|
||||
}
|
||||
let history_model = BlocklistAIHistoryModel::handle(ctx);
|
||||
let Some(snapshot_json) = history_model
|
||||
.as_ref(ctx)
|
||||
.conversation(&conversation_id)
|
||||
.and_then(AIConversation::active_provider_run_json)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let Ok(mut snapshot) = ActiveProviderRunSnapshot::parse(snapshot_json) else {
|
||||
return false;
|
||||
};
|
||||
if !provider_command_completion_matches(
|
||||
snapshot.run.id(),
|
||||
&snapshot.command_action_refs,
|
||||
snapshot.command_monitor.as_ref(),
|
||||
&completion.block_id,
|
||||
completion.initial_requested_command_action_id.as_ref(),
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
if completion.command.is_empty() {
|
||||
completion.command = snapshot
|
||||
.command_monitor
|
||||
.as_ref()
|
||||
.map(|monitor| monitor.command.clone())
|
||||
.unwrap_or_default();
|
||||
}
|
||||
match self
|
||||
.restoring_provider_command_completions
|
||||
.get(&conversation_id)
|
||||
{
|
||||
Some(existing) => return existing == &completion,
|
||||
None => {
|
||||
self.restoring_provider_command_completions
|
||||
.insert(conversation_id, completion.clone());
|
||||
}
|
||||
}
|
||||
snapshot.pending_monitor_observation = None;
|
||||
snapshot.pending_command_completion = Some(completion);
|
||||
match self.persist_provider_run_snapshot(conversation_id, &snapshot, ctx) {
|
||||
Ok(()) => true,
|
||||
Err(error) => {
|
||||
log::error!("Failed to persist completion for restoring provider run: {error}");
|
||||
// The in-memory mailbox remains the exactly-once owner until restore installs it.
|
||||
true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn has_active_stream_for_conversation(
|
||||
&self,
|
||||
conversation_id: AIConversationId,
|
||||
|
||||
@@ -148,7 +148,7 @@ impl PendingResponseStreams {
|
||||
false
|
||||
} else {
|
||||
for response_stream in streams_to_cancel.into_iter() {
|
||||
log::info!(
|
||||
crate::ai::tool_diagnostics::tool_debug!(
|
||||
"Canceling active stream for conversation_id={conversation_id:?}, \
|
||||
reason={reason}"
|
||||
);
|
||||
|
||||
@@ -56,6 +56,10 @@ impl ResponseStreamId {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub(crate) fn from_persisted(value: String) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
|
||||
pub fn for_shared_session(init_event: &response_event::StreamInit) -> Self {
|
||||
// Make the stream ID unique per viewing by appending a local UUID
|
||||
// This prevents collisions when replaying the same conversation multiple times
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -296,6 +296,31 @@ fn mark_run_agents_child_removed(
|
||||
true
|
||||
}
|
||||
|
||||
fn run_agents_event_matches_card(
|
||||
event: &RunAgentsExecutorEvent,
|
||||
conversation_id: Option<AIConversationId>,
|
||||
action_id: &AIAgentActionId,
|
||||
) -> bool {
|
||||
let (event_conversation_id, event_action_id) = match event {
|
||||
RunAgentsExecutorEvent::SpawningStarted {
|
||||
conversation_id,
|
||||
action_id,
|
||||
..
|
||||
}
|
||||
| RunAgentsExecutorEvent::SpawningFinished {
|
||||
conversation_id,
|
||||
action_id,
|
||||
} => (*conversation_id, action_id),
|
||||
RunAgentsExecutorEvent::ChildConversationCreated {
|
||||
action_id,
|
||||
parent_conversation_id,
|
||||
..
|
||||
} => (*parent_conversation_id, action_id),
|
||||
};
|
||||
|
||||
Some(event_conversation_id) == conversation_id && event_action_id == action_id
|
||||
}
|
||||
|
||||
pub struct RunAgentsCardView {
|
||||
action_id: AIAgentActionId,
|
||||
state: RunAgentsEditState,
|
||||
@@ -441,32 +466,34 @@ impl RunAgentsCardView {
|
||||
});
|
||||
|
||||
let action_id_for_subscription = action_id.clone();
|
||||
ctx.subscribe_to_model(&run_agents_executor, move |me, _, event, ctx| match event {
|
||||
RunAgentsExecutorEvent::SpawningStarted {
|
||||
action_id,
|
||||
snapshot,
|
||||
} if action_id == &action_id_for_subscription => {
|
||||
me.spawning = Some(*snapshot);
|
||||
ctx.notify();
|
||||
let conversation_id_for_subscription = block_model.conversation_id(ctx);
|
||||
ctx.subscribe_to_model(&run_agents_executor, move |me, _, event, ctx| {
|
||||
if !run_agents_event_matches_card(
|
||||
event,
|
||||
conversation_id_for_subscription,
|
||||
&action_id_for_subscription,
|
||||
) {
|
||||
return;
|
||||
}
|
||||
RunAgentsExecutorEvent::SpawningFinished { action_id }
|
||||
if action_id == &action_id_for_subscription =>
|
||||
{
|
||||
me.spawning = None;
|
||||
ctx.notify();
|
||||
|
||||
match event {
|
||||
RunAgentsExecutorEvent::SpawningStarted { snapshot, .. } => {
|
||||
me.spawning = Some(*snapshot);
|
||||
ctx.notify();
|
||||
}
|
||||
RunAgentsExecutorEvent::SpawningFinished { .. } => {
|
||||
me.spawning = None;
|
||||
ctx.notify();
|
||||
}
|
||||
RunAgentsExecutorEvent::ChildConversationCreated {
|
||||
agent_name,
|
||||
child_conversation_id,
|
||||
..
|
||||
} => {
|
||||
me.link_child_conversation(agent_name, *child_conversation_id);
|
||||
ctx.notify();
|
||||
}
|
||||
}
|
||||
RunAgentsExecutorEvent::ChildConversationCreated {
|
||||
action_id,
|
||||
agent_name,
|
||||
child_conversation_id,
|
||||
..
|
||||
} if action_id == &action_id_for_subscription => {
|
||||
me.link_child_conversation(agent_name, *child_conversation_id);
|
||||
ctx.notify();
|
||||
}
|
||||
RunAgentsExecutorEvent::SpawningStarted { .. }
|
||||
| RunAgentsExecutorEvent::SpawningFinished { .. }
|
||||
| RunAgentsExecutorEvent::ChildConversationCreated { .. } => {}
|
||||
});
|
||||
|
||||
let history_model = BlocklistAIHistoryModel::handle(ctx);
|
||||
|
||||
@@ -10,9 +10,11 @@ use warp_util::local_or_remote_path::LocalOrRemotePath;
|
||||
|
||||
use super::{
|
||||
has_run_agents_child, link_run_agents_child, mark_run_agents_child_removed,
|
||||
sync_run_agents_children, RunAgentsChildState, RunAgentsEditState,
|
||||
run_agents_event_matches_card, sync_run_agents_children, RunAgentsChildState,
|
||||
RunAgentsEditState, RunAgentsExecutorEvent,
|
||||
};
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
use crate::ai::agent::AIAgentActionId;
|
||||
use crate::ai::blocklist::inline_action::orchestration_controls::OrchestrationEditState;
|
||||
|
||||
fn make_request(harness: &str, mode: RunAgentsExecutionMode) -> RunAgentsRequest {
|
||||
@@ -305,6 +307,31 @@ fn live_child_links_and_removal_survive_streaming_config_sync() {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn child_created_with_duplicate_action_id_only_matches_parent_conversation() {
|
||||
let card_conversation_id = AIConversationId::new();
|
||||
let other_conversation_id = AIConversationId::new();
|
||||
let child_conversation_id = AIConversationId::new();
|
||||
let duplicate_action_id = AIAgentActionId::from("duplicate-action".to_string());
|
||||
let event = RunAgentsExecutorEvent::ChildConversationCreated {
|
||||
action_id: duplicate_action_id.clone(),
|
||||
agent_name: "child".to_string(),
|
||||
parent_conversation_id: other_conversation_id,
|
||||
child_conversation_id,
|
||||
};
|
||||
|
||||
assert!(!run_agents_event_matches_card(
|
||||
&event,
|
||||
Some(card_conversation_id),
|
||||
&duplicate_action_id,
|
||||
));
|
||||
assert!(run_agents_event_matches_card(
|
||||
&event,
|
||||
Some(other_conversation_id),
|
||||
&duplicate_action_id,
|
||||
));
|
||||
}
|
||||
|
||||
mod format_terminal_state_tests {
|
||||
use super::super::{format_terminal_state, StatusKind};
|
||||
use super::*;
|
||||
|
||||
Reference in New Issue
Block a user