fix: reconcile worker restores by operation

crates/worker-runtime/src/execution.rs

@@ -388,6 +388,13 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
         Ok(())
     }
 
+    fn reconcile_restore(
+        &self,
+        request: WorkerExecutionRestoreRequest,
+    ) -> WorkerExecutionSpawnResult {
+        self.restore_worker(request)
+    }
+
     fn restore_worker(
         &self,
         _request: WorkerExecutionRestoreRequest,
@@ -608,6 +615,13 @@ impl WorkerExecutionBackendRef {
         self.backend.preflight_restore(request)
     }
 
+    pub(crate) fn reconcile_restore(
+        &self,
+        request: WorkerExecutionRestoreRequest,
+    ) -> WorkerExecutionSpawnResult {
+        self.backend.reconcile_restore(request)
+    }
+
     pub(crate) fn restore_worker(
         &self,
         request: WorkerExecutionRestoreRequest,

crates/worker-runtime/src/runtime.rs

@@ -1668,19 +1668,25 @@ impl Runtime {
         worker_ref: &WorkerRef,
         mode: WorkerRestoreMode,
     ) -> Result<RuntimeWorkerRestoreResult, RuntimeError> {
-        let operation_id = WorkerLifecycleOperationId::new();
+        let fresh_operation_id = WorkerLifecycleOperationId::new();
-        let (backend, request) = {
+        let (backend, request, operation_id, operation_mode, reconcile) = {
             let state = self.lock()?;
             state.ensure_running()?;
             let worker = state.worker(worker_ref)?;
-            if worker.has_pending_lifecycle_operation() {
+            if worker.pending_stop.is_some() {
                 return Ok(RuntimeWorkerRestoreResult::failed(
                     WorkerRestoreState::ReconciliationRequired,
-                    "worker_restore_reconciliation_pending",
+                    "worker_stop_reconciliation_pending",
-                    "Worker restore reconciliation is still pending",
+                    "Worker stop reconciliation is still pending",
                 ));
             }
-            if worker.execution_handle.is_some() {
+            let pending_restore = worker.pending_restore;
+            let reconcile = pending_restore.is_some();
+            let operation_id = pending_restore
+                .map(|pending| pending.operation_id)
+                .unwrap_or(fresh_operation_id);
+            let operation_mode = pending_restore.map(|pending| pending.mode).unwrap_or(mode);
+            if !reconcile && worker.execution_handle.is_some() {
                 if worker.status.is_active() {
                     return Ok(RuntimeWorkerRestoreResult::accepted(worker.detail()));
                 }
@@ -1695,7 +1701,8 @@ impl Runtime {
             // before Workspace Backend renewed operation-scoped Repository
             // access). Preserve its automatic intent so an explicit restore can
             // serve as the authorized retry for that state.
-            if mode == WorkerRestoreMode::Explicit
+            if !reconcile
+                && mode == WorkerRestoreMode::Explicit
                 && worker.execution_metadata_available
                 && worker.status != WorkerStatus::Stopped
                 && worker.restore_intent != WorkerRestoreIntent::Automatic
@@ -1706,7 +1713,8 @@ impl Runtime {
                     "Worker restore requires a stopped Worker",
                 ));
             }
-            if mode == WorkerRestoreMode::Automatic
+            if !reconcile
+                && mode == WorkerRestoreMode::Automatic
                 && (!worker.status.is_active()
                     || worker.restore_intent != WorkerRestoreIntent::Automatic)
             {
@@ -1743,10 +1751,10 @@ impl Runtime {
                 workdir_attachments: BTreeMap::new(),
                 config_bundle: None,
             };
-            (backend, request)
+            (backend, request, operation_id, operation_mode, reconcile)
         };
 
-        if let Err(_result) = backend.preflight_restore(&request) {
+        if !reconcile && let Err(_result) = backend.preflight_restore(&request) {
             return Ok(RuntimeWorkerRestoreResult::failed(
                 WorkerRestoreState::Rejected,
                 "worker_restore_preflight_rejected",
@@ -1754,7 +1762,10 @@ impl Runtime {
             ));
         }
 
-        if let Err(error) = self.begin_restore_operation(worker_ref, operation_id, mode) {
+        if !reconcile
+            && let Err(error) =
+                self.begin_restore_operation(worker_ref, operation_id, operation_mode)
+        {
             if matches!(error, RuntimeError::StoreCommitOutcomeUnknown { .. }) {
                 return Ok(RuntimeWorkerRestoreResult::failed(
                     WorkerRestoreState::ReconciliationRequired,
@@ -1765,7 +1776,12 @@ impl Runtime {
             return Err(error);
         }
 
-        match backend.restore_worker(request) {
+        let restore_result = if reconcile {
+            backend.reconcile_restore(request)
+        } else {
+            backend.restore_worker(request)
+        };
+        match restore_result {
             WorkerExecutionSpawnResult::Connected {
                 handle,
                 worker_state,
@@ -1860,6 +1876,13 @@ impl Runtime {
                 }
             }
             WorkerExecutionSpawnResult::Rejected(_result) => {
+                if reconcile {
+                    return Ok(RuntimeWorkerRestoreResult::failed(
+                        WorkerRestoreState::ReconciliationRequired,
+                        "worker_restore_reconciliation_rejected",
+                        "Worker restore reconciliation was not completed",
+                    ));
+                }
                 if let Err(error) = self.settle_restore_rollback(worker_ref) {
                     tracing::error!(
                         worker_id = %worker_ref.worker_id,
@@ -2922,6 +2945,32 @@ impl Runtime {
             }
         }
 
+        let pending_restores = {
+            let state = self.lock()?;
+            state
+                .workers
+                .values()
+                .filter_map(|worker| {
+                    worker
+                        .pending_restore
+                        .map(|pending| (worker.worker_ref.clone(), pending.mode))
+                })
+                .collect::<Vec<_>>()
+        };
+        for (worker_ref, mode) in pending_restores {
+            let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
+            let _operation_guard = operation_lock
+                .lock()
+                .map_err(|_| RuntimeError::StatePoisoned)?;
+            if let Err(error) = self.restore_worker_under_lock(&worker_ref, mode) {
+                tracing::warn!(
+                    worker_id = %worker_ref.worker_id,
+                    error = %error,
+                    "Worker restore reconciliation remains pending after Runtime startup"
+                );
+            }
+        }
+
         let candidates = {
             let state = self.lock()?;
             if state.execution_backend.is_none() {
@@ -8806,18 +8855,20 @@ mod tests {
         let retry = runtime
             .restore_worker_operation(&created.worker_ref)
             .unwrap();
-        assert_eq!(retry.state, WorkerRestoreState::ReconciliationRequired);
+        assert_eq!(retry.state, WorkerRestoreState::Accepted);
-        assert_eq!(*backend.restore_count.lock().unwrap(), 1);
+        assert_eq!(*backend.restore_count.lock().unwrap(), 2);
         assert_eq!(
+            backend.restore_operation_ids.lock().unwrap().as_slice(),
+            &[operation_id, operation_id]
+        );
+        assert!(
             runtime
                 .lock()
                 .unwrap()
                 .worker(&created.worker_ref)
                 .unwrap()
                 .pending_restore
-                .unwrap()
+                .is_none()
-                .operation_id,
-            operation_id
         );
     }
 
@@ -9226,24 +9277,28 @@ mod tests {
         let retry = restarted
             .restore_worker_operation(&worker.worker_ref)
             .unwrap();
-        assert_eq!(retry.state, WorkerRestoreState::ReconciliationRequired);
+        assert_eq!(retry.state, WorkerRestoreState::Accepted);
-        assert_eq!(*restarted_backend.restore_count.lock().unwrap(), 0);
+        assert_eq!(*restarted_backend.restore_count.lock().unwrap(), 1);
         assert_eq!(
+            restarted_backend
+                .restore_operation_ids
+                .lock()
+                .unwrap()
+                .as_slice(),
+            &[operation_id]
+        );
+        assert!(
             restarted
                 .lock()
                 .unwrap()
                 .worker(&worker.worker_ref)
                 .unwrap()
                 .pending_restore
-                .unwrap()
+                .is_none()
-                .operation_id,
-            operation_id
         );
-        let delete_error = restarted.delete_worker(&worker.worker_ref).unwrap_err();
+        assert_eq!(
-        assert!(
+            restarted.worker_detail(&worker.worker_ref).unwrap().status,
-            delete_error
+            WorkerStatus::Idle
-                .to_string()
-                .contains("pending restore reconciliation")
         );
 
         let _ = std::fs::remove_dir_all(root);
@@ -9298,37 +9353,32 @@ mod tests {
         let retry = restored
             .restore_worker_operation(&worker.worker_ref)
             .unwrap();
-        assert_eq!(retry.state, WorkerRestoreState::ReconciliationRequired);
+        assert_eq!(retry.state, WorkerRestoreState::Accepted);
-        assert_eq!(*restoring_backend.restore_count.lock().unwrap(), 1);
+        assert_eq!(*restoring_backend.restore_count.lock().unwrap(), 2);
-
-        drop(restored);
-        let restarted_backend = Arc::new(TestExecutionBackend::default());
-        let restarted = Runtime::with_fs_store_and_execution_backend(
-            crate::fs_store::FsRuntimeStoreOptions {
-                root: root.clone(),
-                runtime_id: "test-runtime".to_string(),
-                display_name: None,
-            },
-            restarted_backend.clone(),
-        )
-        .unwrap();
-        assert_eq!(*restarted_backend.restore_count.lock().unwrap(), 0);
         assert_eq!(
-            restarted
+            restoring_backend
+                .restore_operation_ids
+                .lock()
+                .unwrap()
+                .as_slice(),
+            &[operation_id, operation_id]
+        );
+        assert!(
+            restored
                 .lock()
                 .unwrap()
                 .worker(&worker.worker_ref)
                 .unwrap()
                 .pending_restore
-                .unwrap()
+                .is_none()
-                .operation_id,
-            operation_id
         );
-        let retry = restarted
+        restored
-            .restore_worker_operation(&worker.worker_ref)
+            .send_input(
+                &worker.worker_ref,
+                WorkerInput::user("after restore reconciliation"),
+            )
             .unwrap();
-        assert_eq!(retry.state, WorkerRestoreState::ReconciliationRequired);
+        restored.stop_worker(&worker.worker_ref, None).unwrap();
-        assert_eq!(*restarted_backend.restore_count.lock().unwrap(), 0);
 
         let _ = std::fs::remove_dir_all(root);
     }

crates/worker-runtime/src/worker_backend.rs

@@ -472,6 +472,17 @@ impl ProfileRuntimeWorkerFactory {
             .join(uuid::Uuid::now_v7().to_string()))
     }
 
+    fn worker_restore_run_dir(
+        &self,
+        worker_ref: &WorkerRef,
+        operation_id: crate::execution::WorkerLifecycleOperationId,
+    ) -> Result<PathBuf, String> {
+        Ok(self
+            .worker_aggregate_dir(worker_ref)?
+            .join("runs")
+            .join(format!("restore-{operation_id}")))
+    }
+
     fn runtime_worker_name_for_ref(worker_ref: &crate::identity::WorkerRef) -> String {
         format!("worker-runtime-{}", worker_ref.worker_id)
     }
@@ -1423,7 +1434,8 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
         }
 
         let workspace_client = worker.workspace_client_handle();
-        let run_dir = self.worker_run_dir(&request.worker_ref)?;
+        let run_dir =
+            self.worker_restore_run_dir(&request.worker_ref, request.operation_id)?;
         let bash_output_dir = bash_output_dir_for_worker_id(&request.worker_ref.worker_id);
         let started = PreparedWorker::new(
             worker,
@@ -1567,6 +1579,8 @@ struct RuntimeWorkerExecution {
     tasks: RuntimeExecutionTaskScope,
     worker_state: Arc<RwLock<protocol::WorkerStateSnapshot>>,
     workspace_client: Option<Arc<dyn WorkspaceClient>>,
+    restore_operation_id: Option<crate::execution::WorkerLifecycleOperationId>,
+    workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
 }
 
 /// `worker-runtime` execution backend backed by real `worker` crate Workers.
@@ -1836,6 +1850,8 @@ where
         tasks: RuntimeExecutionTaskScope,
         worker_state: Arc<RwLock<protocol::WorkerStateSnapshot>>,
         workspace_client: Option<Arc<dyn WorkspaceClient>>,
+        restore_operation_id: Option<crate::execution::WorkerLifecycleOperationId>,
+        workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
     ) -> Result<WorkerExecutionHandle, String> {
         let mut workers = self
             .workers
@@ -1853,6 +1869,8 @@ where
                 tasks,
                 worker_state,
                 workspace_client,
+                restore_operation_id,
+                workdir_attachments,
             },
         );
         Ok(WorkerExecutionHandle::new(
@@ -1907,6 +1925,7 @@ where
         handle: WorkerHandle,
         shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
         controller_task: tokio::task::JoinHandle<()>,
+        restore_operation_id: Option<crate::execution::WorkerLifecycleOperationId>,
         workdir_attachments: BTreeMap<WorkdirAttachmentAlias, WorkingDirectoryBinding>,
         workspace_client: Option<Arc<dyn WorkspaceClient>>,
     ) -> WorkerExecutionSpawnResult {
@@ -1921,6 +1940,15 @@ where
         let worker_state_snapshot = handle.shared_state.snapshot();
         let worker_state = Arc::new(RwLock::new(worker_state_snapshot));
         let tasks = RuntimeExecutionTaskScope::new(controller_task);
+        let workdir_attachment_statuses = workdir_attachments
+            .iter()
+            .map(
+                |(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
+                    alias: alias.clone(),
+                    working_directory: binding.status(),
+                },
+            )
+            .collect::<Vec<_>>();
         #[cfg(feature = "ws-server")]
         {
             let mut events = streams.events;
@@ -1987,13 +2015,8 @@ where
                             Err(_) => {
                                 let retained_state =
                                     worker_state.read().ok().map(|state| state.clone());
-                                let retained_workdir_attachments = workdir_attachments
+                                let retained_workdir_attachments =
-                                    .iter()
+                                    workdir_attachment_statuses.clone();
-                                    .map(|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
-                                        alias: alias.clone(),
-                                        working_directory: binding.status(),
-                                    })
-                                    .collect();
                                 let retained_handle = self
                                     .retain_uncertain_unconnected_controller(
                                         &worker_ref,
@@ -2002,6 +2025,8 @@ where
                                         tasks,
                                         worker_state,
                                         workspace_client,
+                                        restore_operation_id,
+                                        retained_workdir_attachments.clone(),
                                     )
                                     .ok();
                                 WorkerExecutionSpawnResult::ReconciliationRequired {
@@ -2072,21 +2097,15 @@ where
                 tasks,
                 worker_state,
                 workspace_client,
+                restore_operation_id,
+                workdir_attachments: workdir_attachment_statuses.clone(),
             },
         );
 
         WorkerExecutionSpawnResult::Connected {
             handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()),
             worker_state: connected_worker_state,
-            workdir_attachments: workdir_attachments
+            workdir_attachments: workdir_attachment_statuses,
-                .into_iter()
-                .map(
-                    |(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
-                        alias,
-                        working_directory: binding.status(),
-                    },
-                )
-                .collect(),
         }
     }
 }
@@ -2338,6 +2357,7 @@ where
                     controller.handle,
                     controller.shutdown,
                     controller.controller_task,
+                    None,
                     BTreeMap::new(),
                     Some(controller.workspace_client),
                 );
@@ -2465,6 +2485,7 @@ where
             controller.handle,
             controller.shutdown,
             controller.controller_task,
+            None,
             workdir_attachments,
             Some(controller.workspace_client),
         )
@@ -2577,10 +2598,25 @@ where
             .workers
             .lock()
             .unwrap_or_else(|poisoned| poisoned.into_inner());
-        if workers.contains_key(&request.worker_ref) {
+        if let Some(existing) = workers.get(&request.worker_ref) {
+            if existing.restore_operation_id == Some(request.operation_id) {
+                let worker_state = existing
+                    .worker_state
+                    .read()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner)
+                    .clone();
+                return WorkerExecutionSpawnResult::Connected {
+                    handle: WorkerExecutionHandle::new(
+                        request.worker_ref.clone(),
+                        self.backend_id(),
+                    ),
+                    worker_state,
+                    workdir_attachments: existing.workdir_attachments.clone(),
+                };
+            }
             return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::busy(
                 WorkerExecutionOperation::Restore,
-                "Worker is already connected to execution backend",
+                "Worker is already connected to execution backend by another lifecycle operation",
             ));
         }
         drop(workers);
@@ -2671,6 +2707,7 @@ where
         let factory = self.factory.clone();
         let bridge_context = request.context.clone();
         let worker_ref = request.worker_ref.clone();
+        let operation_id = request.operation_id;
         let restore_result = self
             .run_cancellable_on_adapter_runtime(self.spawn_restore_timeout, async move {
                 factory.restore_controller(request).await
@@ -2698,6 +2735,7 @@ where
             controller.handle,
             controller.shutdown,
             controller.controller_task,
+            Some(operation_id),
             workdir_attachments,
             Some(controller.workspace_client),
         )
@@ -3555,7 +3593,13 @@ mod tests {
     fn test_execution_context(worker_ref: WorkerRef) -> WorkerExecutionContext {
         WorkerExecutionContext::new(
             worker_ref,
-            Arc::new(|_, _| panic!("unused test event sink")),
+            Arc::new(|worker_ref, payload| {
+                Ok(crate::observation::WorkerObservationEvent::new(
+                    1,
+                    worker_ref,
+                    payload,
+                ))
+            }),
         )
     }
 
@@ -4562,6 +4606,102 @@ mod tests {
         drop(restored_backend);
     }
 
+    #[test]
+    #[serial_test::serial(worker_allocation)]
+    fn restore_reconciliation_reuses_the_operation_controller() {
+        let client = MockClient::sequential(Vec::new());
+        let runtime_base = tempfile::tempdir().unwrap();
+        let cwd = tempfile::tempdir().unwrap();
+        let store = tempfile::tempdir().unwrap();
+        let observed_workspace_clients = Arc::new(Mutex::new(Vec::new()));
+        let factory = MockFactory {
+            client,
+            runtime_base: runtime_base.path().to_path_buf(),
+            cwd: cwd.path().to_path_buf(),
+            store_dir: store.path().join("sessions"),
+            worker_metadata_dir: store.path().join("workers"),
+            observed_cwds: Arc::new(Mutex::new(Vec::new())),
+            observed_workspace_clients: observed_workspace_clients.clone(),
+        };
+        let backend = Arc::new(WorkerRuntimeExecutionBackend::new(factory).unwrap());
+        let runtime = EmbeddedRuntime::with_execution_backend(
+            RuntimeOptions::default(),
+            backend.clone(),
+        )
+        .unwrap();
+        runtime.store_config_bundle(test_bundle()).unwrap();
+        let request = create_request("idempotent restore reconciliation");
+        let worker = runtime.create_worker(request.clone()).unwrap();
+        runtime.stop_worker(&worker.worker_ref, None).unwrap();
+        let operation_id = crate::execution::WorkerLifecycleOperationId::new();
+        let restore_request = WorkerExecutionRestoreRequest {
+            operation_id,
+            worker_ref: worker.worker_ref.clone(),
+            request,
+            workspace_scope: None,
+            context: test_execution_context(worker.worker_ref.clone()),
+            previous_workdir_attachments: Vec::new(),
+            logical_workdir_attachments: Vec::new(),
+            workdir_attachments: BTreeMap::new(),
+            config_bundle: None,
+        };
+
+        let first = backend.restore_worker(restore_request.clone());
+        let first_handle = match first {
+            WorkerExecutionSpawnResult::Connected { handle, .. } => handle,
+            other => panic!("initial restore was not connected: {other:?}"),
+        };
+        let factory_calls_after_restore = observed_workspace_clients.lock().unwrap().len();
+
+        let reconciled = backend.reconcile_restore(restore_request);
+        assert!(matches!(
+            reconciled,
+            WorkerExecutionSpawnResult::Connected { .. }
+        ));
+        assert_eq!(
+            observed_workspace_clients.lock().unwrap().len(),
+            factory_calls_after_restore,
+            "reconciliation must not construct a second controller"
+        );
+        let workers = backend.workers.lock().unwrap();
+        let execution = workers.get(&worker.worker_ref).unwrap();
+        assert_eq!(execution.restore_operation_id, Some(operation_id));
+        assert_eq!(workers.len(), 1);
+        drop(workers);
+
+        assert!(backend.stop_worker(&first_handle).is_accepted());
+    }
+
+    #[test]
+    fn profile_restore_run_directory_is_operation_scoped() {
+        let root = tempfile::tempdir().unwrap();
+        let runtime_store_dir = root.path().join("runtime");
+        let factory = ProfileRuntimeWorkerFactory::new(root.path())
+            .with_runtime_store_dir(&runtime_store_dir);
+        let worker_ref = WorkerRef::new(WorkerId::now_v7());
+        let operation_id = crate::execution::WorkerLifecycleOperationId::new();
+
+        let first = factory
+            .worker_restore_run_dir(&worker_ref, operation_id)
+            .unwrap();
+        let retry = factory
+            .worker_restore_run_dir(&worker_ref, operation_id)
+            .unwrap();
+        let different = factory
+            .worker_restore_run_dir(
+                &worker_ref,
+                crate::execution::WorkerLifecycleOperationId::new(),
+            )
+            .unwrap();
+
+        assert_eq!(first, retry);
+        assert_ne!(first, different);
+        assert_eq!(
+            first.file_name().unwrap().to_string_lossy(),
+            format!("restore-{operation_id}")
+        );
+    }
+
     #[test]
     fn builtin_profile_selector_is_not_double_prefixed() {
         assert_eq!(