diff --git a/crates/tui/src/backend_worker_picker.rs b/crates/tui/src/backend_worker_picker.rs index 7812f85e3..f8fa5336c 100644 --- a/crates/tui/src/backend_worker_picker.rs +++ b/crates/tui/src/backend_worker_picker.rs @@ -785,6 +785,12 @@ mod tests { #[derive(Default)] struct PendingRuntimeBackend { operations: std::sync::Mutex>, + owners: std::sync::Mutex< + std::collections::HashMap< + worker_runtime::execution::WorkerLifecycleOperationId, + worker_runtime::identity::WorkerRef, + >, + >, } impl worker_runtime::execution::WorkerExecutionBackend for PendingRuntimeBackend { @@ -806,6 +812,21 @@ mod tests { ) } + fn restore_recovery_disposition( + &self, + operation_id: worker_runtime::execution::WorkerLifecycleOperationId, + worker_ref: &worker_runtime::identity::WorkerRef, + ) -> Option { + use worker_runtime::execution::RestoreRecoveryDisposition as D; + Some( + if self.owners.lock().unwrap().get(&operation_id) == Some(worker_ref) { + D::Continue + } else { + D::Absent + }, + ) + } + fn restore_worker( &self, request: worker_runtime::execution::WorkerExecutionRestoreRequest, @@ -813,6 +834,10 @@ mod tests { use worker_runtime::execution::*; let mut operations = self.operations.lock().unwrap(); operations.push(request.operation_id); + self.owners + .lock() + .unwrap() + .insert(request.operation_id, request.worker_ref.clone()); if operations.len() == 1 { WorkerExecutionSpawnResult::ReconciliationRequired { result: WorkerExecutionResult::errored( diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index a8ddde657..175feb3e8 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -48392,6 +48392,12 @@ mod tests { std::sync::Mutex>, reconcile_operations: std::sync::Mutex>, + restore_owners: std::sync::Mutex< + std::collections::HashMap< + worker_runtime::execution::WorkerLifecycleOperationId, + worker_runtime::identity::WorkerRef, + >, + >, inputs: std::sync::Mutex>, protocol_methods: Arc>>, @@ -48425,6 +48431,7 @@ mod tests { restore_failure: std::sync::Mutex::new(None), restore_operations: std::sync::Mutex::new(Vec::new()), reconcile_operations: std::sync::Mutex::new(Vec::new()), + restore_owners: Default::default(), inputs: std::sync::Mutex::new(Vec::new()), protocol_methods: Arc::new(std::sync::Mutex::new(Vec::new())), stops: std::sync::Mutex::new(Vec::new()), @@ -48629,6 +48636,10 @@ mod tests { .lock() .unwrap() .push(request.operation_id); + self.restore_owners + .lock() + .unwrap() + .insert(request.operation_id, request.worker_ref.clone()); if let Some(result) = self.restore_failure.lock().unwrap().take() { return result; } @@ -48664,6 +48675,21 @@ mod tests { } } + fn restore_recovery_disposition( + &self, + operation_id: worker_runtime::execution::WorkerLifecycleOperationId, + worker_ref: &worker_runtime::identity::WorkerRef, + ) -> Option { + use worker_runtime::execution::RestoreRecoveryDisposition as D; + Some( + if self.restore_owners.lock().unwrap().get(&operation_id) == Some(worker_ref) { + D::Continue + } else { + D::Absent + }, + ) + } + fn reconcile_restore( &self, request: worker_runtime::execution::WorkerExecutionRestoreRequest,