Harden restore preflight and retry reconciliation

crates/worker-runtime/src/runtime.rs

@@ -1483,13 +1483,6 @@ impl Runtime {
             let state = self.lock()?;
             state.ensure_running()?;
             let worker = state.worker(worker_ref)?;
-            if worker.execution_bound && !worker.execution_metadata_available {
-                return Ok(RuntimeWorkerRestoreResult::failed(
-                    WorkerRestoreState::ReconciliationRequired,
-                    "worker_restore_reconciliation_pending",
-                    "Worker restore cleanup or reconciliation is still pending",
-                ));
-            }
             if worker.execution_handle.is_some() {
                 if worker.status.is_active() {
                     return Ok(RuntimeWorkerRestoreResult::accepted(worker.detail()));
@@ -7520,11 +7513,12 @@ mod tests {
             runtime.worker_detail(&created.worker_ref).unwrap().status,
             WorkerStatus::Idle
         );
+        backend.restore_result.lock().unwrap().take();
         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);
     }
 
     #[cfg(feature = "fs-store")]
@@ -7898,21 +7892,18 @@ mod tests {
         let restored_worker = restored.worker_detail(&worker.worker_ref).unwrap();
         assert_eq!(restored_worker.status, WorkerStatus::Idle);
         assert!(!restored_worker.execution_metadata_available);
+        restoring_backend.restore_result.lock().unwrap().take();
         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);
-        let err = restored
+        restored
             .send_input(
                 &worker.worker_ref,
-                WorkerInput::user("after failed restore"),
+                WorkerInput::user("after reconciled restore"),
             )
-            .unwrap_err();
+            .unwrap();
-        assert!(matches!(
-            err,
-            RuntimeError::WorkerExecutionUnavailable { .. }
-        ));
 
         let _ = std::fs::remove_dir_all(root);
     }

crates/worker-runtime/src/worker_backend.rs

@@ -59,7 +59,9 @@ fn next_internal_command(
         .max(floor);
     Ok(WorkerCommandEnvelope::new(command_id))
 }
-use session_store::{CombinedStore, WorkerAggregateStore, WorkerSessionStore};
+use session_store::{
+    CombinedStore, WorkerAggregateStore, WorkerMetadataStore, WorkerSessionStore,
+};
 #[cfg(test)]
 use session_store::{FsStore, FsWorkerStore};
 use tokio::runtime::Runtime;
@@ -122,6 +124,13 @@ pub trait RuntimeWorkerFactory: Send + Sync + 'static {
         request: WorkerExecutionSpawnRequest,
     ) -> Result<RuntimeWorkerController, String>;
 
+    async fn preflight_restore(
+        &self,
+        _request: &WorkerExecutionRestoreRequest,
+    ) -> Result<(), String> {
+        Ok(())
+    }
+
     async fn restore_controller(
         &self,
         request: WorkerExecutionRestoreRequest,
@@ -1002,6 +1011,41 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
         })
     }
 
+    async fn preflight_restore(
+        &self,
+        request: &WorkerExecutionRestoreRequest,
+    ) -> Result<(), String> {
+        let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref);
+        let (mut manifest, _) = Self::restore_fallback_manifest(&worker_name)?;
+        bind_workspace_memory_settings(&mut manifest, &request.request)?;
+        let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?;
+        if !worker_aggregate_dir.is_dir() {
+            return Err("Persisted Worker aggregate metadata is unavailable".to_string());
+        }
+        if !worker_aggregate_dir.join("session").is_dir() {
+            return Err("Persisted Worker Session metadata is unavailable".to_string());
+        }
+        let metadata_store = WorkerAggregateStore::new(worker_aggregate_dir, &worker_name)
+            .map_err(|error| format!("failed to read Worker aggregate metadata: {error}"))?;
+        let metadata = metadata_store
+            .read_by_name(&worker_name)
+            .map_err(|error| format!("failed to read Worker metadata: {error}"))?
+            .ok_or_else(|| "Persisted Worker metadata is unavailable".to_string())?;
+        if request.request.workspace_api.is_some()
+            && metadata
+                .active
+                .as_ref()
+                .is_some_and(|active| active.segment_id.is_none())
+            && request.config_bundle.is_none()
+        {
+            return Err(
+                "Pending Workspace Worker restore requires current profile launch authority"
+                    .to_string(),
+            );
+        }
+        Ok(())
+    }
+
     async fn restore_controller(
         &self,
         request: WorkerExecutionRestoreRequest,
@@ -1310,6 +1354,7 @@ pub struct WorkerRuntimeExecutionBackend<F = ProfileRuntimeWorkerFactory> {
     working_directory_materializer: Option<Arc<dyn WorkingDirectoryMaterializer>>,
     runtime: Mutex<Option<Runtime>>,
     workers: Mutex<HashMap<crate::identity::WorkerRef, RuntimeWorkerExecution>>,
+    restore_lock: Mutex<()>,
     spawn_restore_timeout: Duration,
 }
 
@@ -1338,6 +1383,7 @@ where
             working_directory_materializer: None,
             runtime: Mutex::new(Some(runtime)),
             workers: Mutex::new(HashMap::new()),
+            restore_lock: Mutex::new(()),
             spawn_restore_timeout: SPAWN_RESTORE_TASK_TIMEOUT,
         })
     }
@@ -1760,40 +1806,19 @@ where
             let _ = bridge_context;
         }
 
-        let mut workers = match self.workers.lock() {
+        let mut workers = self
-            Ok(workers) => workers,
+            .workers
-            Err(_) => {
+            .lock()
-                let cleanup = self.cleanup_unconnected_controller(
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
-                    &handle,
-                    &shutdown,
-                    &tasks,
-                    &worker_state,
-                );
-                let result = WorkerExecutionResult::errored(
-                    operation,
-                    match &cleanup {
-                        Ok(()) => "worker adapter registry lock is poisoned".to_string(),
-                        Err(cleanup) => format!(
-                            "worker adapter registry lock is poisoned; controller cleanup failed: {cleanup}"
-                        ),
-                    },
-                );
-                return if operation == WorkerExecutionOperation::Restore {
-                    match cleanup {
-                        Ok(()) => WorkerExecutionSpawnResult::RolledBack(result),
-                        Err(_) => WorkerExecutionSpawnResult::ReconciliationRequired {
-                            result,
-                            handle: None,
-                            worker_state: None,
-                            working_directory: None,
-                        },
-                    }
-                } else {
-                    WorkerExecutionSpawnResult::Errored(result)
-                };
-            }
-        };
         if workers.contains_key(&worker_ref) {
+            let existing_worker_state = workers.get(&worker_ref).and_then(|worker| {
+                worker
+                    .worker_state
+                    .read()
+                    .ok()
+                    .map(|state| state.clone())
+            });
+            let existing_handle = WorkerExecutionHandle::new(worker_ref.clone(), self.backend_id());
             drop(workers);
             let cleanup = self.cleanup_unconnected_controller(
                 &handle,
@@ -1815,8 +1840,8 @@ where
                     Ok(()) => WorkerExecutionSpawnResult::RolledBack(result),
                     Err(_) => WorkerExecutionSpawnResult::ReconciliationRequired {
                         result,
-                        handle: None,
+                        handle: Some(existing_handle),
-                        worker_state: None,
+                        worker_state: existing_worker_state,
                         working_directory: None,
                     },
                 }
@@ -2125,6 +2150,54 @@ where
                 "Worker is already connected to execution backend",
             ));
         }
+        drop(workers);
+
+        if request.previous_working_directory.is_some()
+            && self.working_directory_materializer.is_none()
+        {
+            return Err(WorkerExecutionResult::rejected(
+                WorkerExecutionOperation::Restore,
+                "Persisted Worker Workdir binding cannot be restored by this Runtime",
+            ));
+        }
+        if request.previous_working_directory.is_none()
+            && request.request.working_directory_request.is_some()
+        {
+            return Err(WorkerExecutionResult::rejected(
+                WorkerExecutionOperation::Restore,
+                "Persisted Worker Workdir allocation is unavailable",
+            ));
+        }
+        if request.previous_working_directory.is_none()
+            && request.request.working_directory.is_some()
+            && self.working_directory_materializer.is_none()
+        {
+            return Err(WorkerExecutionResult::rejected(
+                WorkerExecutionOperation::Restore,
+                "Persisted Worker Workdir claim cannot be restored by this Runtime",
+            ));
+        }
+        if let Some(workspace_api) = request.request.workspace_api.as_ref() {
+            let Some(scope) = request.workspace_scope.as_ref() else {
+                return Err(WorkerExecutionResult::rejected(
+                    WorkerExecutionOperation::Restore,
+                    "Persisted Workspace restore authorization is unavailable",
+                ));
+            };
+            if scope.workspace_id != workspace_api.workspace_id {
+                return Err(WorkerExecutionResult::rejected(
+                    WorkerExecutionOperation::Restore,
+                    "Persisted Workspace restore authorization does not match Worker authority",
+                ));
+            }
+        }
+
+        let factory = Arc::clone(&self.factory);
+        let request = request.clone();
+        self.run_on_adapter_runtime(async move { factory.preflight_restore(&request).await })
+            .map_err(|message| {
+                WorkerExecutionResult::rejected(WorkerExecutionOperation::Restore, message)
+            })?;
         Ok(())
     }
 
@@ -2132,6 +2205,25 @@ where
         &self,
         mut request: WorkerExecutionRestoreRequest,
     ) -> WorkerExecutionSpawnResult {
+        // Serialize backend restore starts so duplicate detection is completed
+        // before any controller or bridge side effect. Runtime also serializes
+        // per Worker; this guard protects direct/concurrent backend callers.
+        let _restore_guard = self
+            .restore_lock
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        let workers = self
+            .workers
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        if workers.contains_key(&request.worker_ref) {
+            return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::busy(
+                WorkerExecutionOperation::Restore,
+                "Worker is already connected to execution backend",
+            ));
+        }
+        drop(workers);
+
         let working_directory = match request.previous_working_directory.clone() {
             Some(status) => {
                 let Some(materializer) = self.working_directory_materializer.as_ref() else {

crates/workspace-server/src/server.rs

@@ -20722,7 +20722,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn production_profile_backend_reports_uncertain_pending_orchestrator_restore() {
+    async fn production_profile_backend_preflight_rejects_pending_orchestrator_restore() {
         let workspace = tempfile::tempdir().unwrap();
         init_clean_git_workspace(workspace.path());
         let config = test_server_config(workspace.path());
@@ -20772,9 +20772,11 @@ mod tests {
             }),
         )
         .await
-        .expect_err("pending Workspace Orchestrator restore without durable Prompt must require reconciliation");
+        .expect_err(
+            "pending Workspace Orchestrator restore without durable Prompt must be rejected",
+        );
         assert!(
-            format!("{error:?}").contains("worker_restore_reconciliation_required"),
+            format!("{error:?}").contains("worker_restore_preflight_rejected"),
             "unexpected restore error: {error:?}"
         );
     }