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

Fix provider batch completion race

Browse Source
This commit is contained in:
rkw6086
2026-08-15 07:13:07 -05:00
parent b079f036fa
commit 642cb6adc1
2 changed files with 142 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/controller.rs
+80 -5
View File
@@ -642,6 +642,7 @@ struct ActiveProviderRunSlot {
turn_control: Option<TurnCommandSender>,
cancellation_reason: Option<CancellationReason>,
committed_provider_batch: Option<ExternalWorkId>,
finished_provider_batch: Option<ExternalWorkId>,
command_action_refs: HashMap<AIAgentActionId, ProviderToolExecutionRef>,
command_monitor: Option<ProviderCommandMonitorState>,
pending_monitor_observation: Option<PendingProviderMonitorObservation>,
@@ -706,6 +707,8 @@ struct ActiveProviderRunSnapshot {
did_input_contain_user_query: bool,
persistence_offset: usize,
committed_provider_batch: Option<ExternalWorkId>,
#[serde(default)]
finished_provider_batch: Option<ExternalWorkId>,
command_action_refs: HashMap<AIAgentActionId, ProviderToolExecutionRef>,
command_monitor: Option<ProviderCommandMonitorState>,
pending_monitor_observation: Option<PendingProviderMonitorObservation>,
@@ -744,6 +747,7 @@ impl ActiveProviderRunSnapshot {
did_input_contain_user_query: slot.did_input_contain_user_query,
persistence_offset: checkpoint.persistence_offset,
committed_provider_batch: slot.committed_provider_batch.clone(),
finished_provider_batch: slot.finished_provider_batch.clone(),
command_action_refs: slot.command_action_refs.clone(),
command_monitor: slot.command_monitor.clone(),
pending_monitor_observation: slot.pending_monitor_observation.clone(),
@@ -813,6 +817,13 @@ impl ActiveProviderRunSnapshot {
{
return Err("committed provider batch belongs to a different run".to_string());
}
if self
.finished_provider_batch
.as_ref()
.is_some_and(|work_id| &work_id.run_id != run_id)
{
return Err("finished provider batch belongs to a different run".to_string());
}
for (action_id, execution_ref) in &self.command_action_refs {
if execution_ref.conversation_id != conversation_id
|| &execution_ref.run_id != run_id
@@ -891,6 +902,34 @@ enum ProviderFinishedActionDisposition {
Resume,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ProviderBatchSignal {
BatchCommitted,
ActionsFinished,
}
fn record_provider_batch_signal(
committed_work_id: &mut Option<ExternalWorkId>,
finished_work_id: &mut Option<ExternalWorkId>,
work_id: &ExternalWorkId,
signal: ProviderBatchSignal,
) -> bool {
if committed_work_id
.as_ref()
.is_some_and(|recorded_work_id| recorded_work_id != work_id)
|| finished_work_id
.as_ref()
.is_some_and(|recorded_work_id| recorded_work_id != work_id)
{
return false;
}
match signal {
ProviderBatchSignal::BatchCommitted => *committed_work_id = Some(work_id.clone()),
ProviderBatchSignal::ActionsFinished => *finished_work_id = Some(work_id.clone()),
}
committed_work_id.as_ref() == Some(work_id) && finished_work_id.as_ref() == Some(work_id)
}
fn normalize_restored_provider_snapshot(
snapshot: &mut ActiveProviderRunSnapshot,
) -> Result<(), String> {
@@ -923,6 +962,7 @@ fn normalize_restored_provider_snapshot(
}
snapshot.committed_provider_batch = None;
}
snapshot.finished_provider_batch = None;
Ok(())
}
@@ -4342,6 +4382,7 @@ impl BlocklistAIController {
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,
@@ -4629,6 +4670,7 @@ impl BlocklistAIController {
did_input_contain_user_query,
persistence_offset,
committed_provider_batch,
finished_provider_batch,
command_action_refs,
command_monitor,
pending_monitor_observation,
@@ -4721,6 +4763,7 @@ impl BlocklistAIController {
turn_control: None,
cancellation_reason: None,
committed_provider_batch,
finished_provider_batch,
command_action_refs,
command_monitor,
pending_monitor_observation,
@@ -5281,6 +5324,8 @@ impl BlocklistAIController {
};
match block {
ProviderRunBlock::Tools(batch) => {
slot.committed_provider_batch = None;
slot.finished_provider_batch = None;
slot.run = Some(run);
if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) {
self.fail_active_provider_run(
@@ -5445,10 +5490,17 @@ impl BlocklistAIController {
) {
return;
}
let mut should_resume = false;
let should_drive = match run.coordinator.apply_tool_lifecycle(execution_ref, event) {
Ok(ProviderToolLifecycleOutcome::Pending) => false,
Ok(ProviderToolLifecycleOutcome::BatchCommitted) => {
slot.committed_provider_batch = Some(execution_ref.work_id());
let work_id = execution_ref.work_id();
should_resume = record_provider_batch_signal(
&mut slot.committed_provider_batch,
&mut slot.finished_provider_batch,
&work_id,
ProviderBatchSignal::BatchCommitted,
);
false
}
Err(error) => {
@@ -5471,7 +5523,9 @@ impl BlocklistAIController {
);
return;
}
if should_drive {
if should_resume {
self.handle_provider_actions_finished(conversation_id, execution_ref, ctx);
} else if should_drive {
self.drive_active_provider_run(conversation_id, ctx);
}
}
@@ -5655,11 +5709,31 @@ impl BlocklistAIController {
))
})
.unwrap_or(ProviderFinishedActionDisposition::Ignore);
let work_id = execution_ref.work_id();
let should_resume = if disposition == ProviderFinishedActionDisposition::Ignore {
false
} else {
let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else {
return;
};
record_provider_batch_signal(
&mut slot.committed_provider_batch,
&mut slot.finished_provider_batch,
&work_id,
ProviderBatchSignal::ActionsFinished,
)
};
if disposition == ProviderFinishedActionDisposition::AwaitBatchCommit {
if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) {
self.fail_active_provider_run(
conversation_id,
format!("failed to persist finished provider action phase: {error}"),
ctx,
);
}
return;
}
let work_id = execution_ref.work_id();
let results = (disposition == ProviderFinishedActionDisposition::Resume).then(|| {
let results = should_resume.then(|| {
self.action_model
.as_ref(ctx)
.provider_finished_action_results(conversation_id, &work_id)
@@ -5674,7 +5748,7 @@ impl BlocklistAIController {
self.fail_active_provider_run(conversation_id, error, ctx);
return;
}
if disposition != ProviderFinishedActionDisposition::Resume {
if !should_resume {
return;
}
@@ -5682,6 +5756,7 @@ impl BlocklistAIController {
return;
};
slot.committed_provider_batch = None;
slot.finished_provider_batch = None;
let Some(mut run) = slot.run.take() else {
return;
};
app/src/ai/blocklist/controller_tests.rs
+62
View File
@@ -141,6 +141,7 @@ fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProvider
did_input_contain_user_query: true,
persistence_offset: 0,
committed_provider_batch: None,
finished_provider_batch: None,
command_action_refs: HashMap::new(),
command_monitor: None,
pending_monitor_observation: None,
@@ -599,6 +600,67 @@ fn provider_finished_action_only_resumes_its_committed_batch() {
);
}
#[test]
fn provider_batch_resumes_after_both_signals_in_either_order() {
let work_id = ExternalWorkId {
run_id: ProviderRunId::new("current"),
epoch: RunEpoch::new(3),
};
for signals in [
[
super::ProviderBatchSignal::ActionsFinished,
super::ProviderBatchSignal::BatchCommitted,
],
[
super::ProviderBatchSignal::BatchCommitted,
super::ProviderBatchSignal::ActionsFinished,
],
] {
let mut committed = None;
let mut finished = None;
assert!(!super::record_provider_batch_signal(
&mut committed,
&mut finished,
&work_id,
signals[0],
));
assert!(super::record_provider_batch_signal(
&mut committed,
&mut finished,
&work_id,
signals[1],
));
}
}
#[test]
fn provider_batch_signals_do_not_cross_work_ids() {
let current = ExternalWorkId {
run_id: ProviderRunId::new("current"),
epoch: RunEpoch::new(3),
};
let stale = ExternalWorkId {
run_id: current.run_id.clone(),
epoch: RunEpoch::new(2),
};
let mut committed = None;
let mut finished = None;
assert!(!super::record_provider_batch_signal(
&mut committed,
&mut finished,
&current,
super::ProviderBatchSignal::BatchCommitted,
));
assert!(!super::record_provider_batch_signal(
&mut committed,
&mut finished,
&stale,
super::ProviderBatchSignal::ActionsFinished,
));
assert_eq!(committed, Some(current));
assert_eq!(finished, None);
}
#[test]
fn provider_boundary_prioritizes_completion_and_waits_for_committed_results() {
assert_eq!(