diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 29736744c..5c2078a63 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/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!(*backend.restore_count.lock().unwrap(), 1); + assert_eq!(retry.state, WorkerRestoreState::Accepted); + 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!(*restoring_backend.restore_count.lock().unwrap(), 1); - let err = restored + assert_eq!(retry.state, WorkerRestoreState::Accepted); + assert_eq!(*restoring_backend.restore_count.lock().unwrap(), 2); + restored .send_input( &worker.worker_ref, - WorkerInput::user("after failed restore"), + WorkerInput::user("after reconciled restore"), ) - .unwrap_err(); - assert!(matches!( - err, - RuntimeError::WorkerExecutionUnavailable { .. } - )); + .unwrap(); let _ = std::fs::remove_dir_all(root); } diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 8447fc726..9b8f4c9ed 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/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; + 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 { working_directory_materializer: Option>, runtime: Mutex>, workers: Mutex>, + 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() { - Ok(workers) => workers, - Err(_) => { - let cleanup = self.cleanup_unconnected_controller( - &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) - }; - } - }; + let mut workers = self + .workers + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); 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, - worker_state: None, + handle: Some(existing_handle), + 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 { diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 27112fe5f..721977812 100644 --- a/crates/workspace-server/src/server.rs +++ b/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:?}" ); }