fix: retain exact restore ownership in test execution backends

crates/tui/src/backend_worker_picker.rs

@@ -785,6 +785,12 @@ mod tests {
     #[derive(Default)]
     struct PendingRuntimeBackend {
         operations: std::sync::Mutex<Vec<worker_runtime::execution::WorkerLifecycleOperationId>>,
+        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<worker_runtime::execution::RestoreRecoveryDisposition> {
+            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(

crates/workspace-server/src/server.rs

@@ -48392,6 +48392,12 @@ mod tests {
             std::sync::Mutex<Vec<worker_runtime::execution::WorkerLifecycleOperationId>>,
         reconcile_operations:
             std::sync::Mutex<Vec<worker_runtime::execution::WorkerLifecycleOperationId>>,
+        restore_owners: std::sync::Mutex<
+            std::collections::HashMap<
+                worker_runtime::execution::WorkerLifecycleOperationId,
+                worker_runtime::identity::WorkerRef,
+            >,
+        >,
         inputs: std::sync::Mutex<Vec<(worker_runtime::identity::WorkerRef, String)>>,
         protocol_methods:
             Arc<std::sync::Mutex<Vec<(worker_runtime::identity::WorkerRef, protocol::Method)>>>,
@@ -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<worker_runtime::execution::RestoreRecoveryDisposition> {
+            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,