From b7a22146a0c2ea4e0a2c4bcb0324ecefa42ce077 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 20 Oct 2026 06:00:58 +0900 Subject: [PATCH] fix: recover confirmed worker restore failures --- crates/worker-runtime/src/execution.rs | 19 + crates/worker-runtime/src/runtime.rs | 16 +- crates/worker-runtime/src/worker_backend.rs | 1510 ++++++++++++----- .../worker_backend_restore_failure_tests.rs | 945 +++++++++++ crates/worker/src/bootstrap.rs | 6 +- crates/worker/src/controller.rs | 20 +- crates/worker/src/lib.rs | 11 +- .../worker/src/runtime/worker_allocation.rs | 5 +- .../src/runtime/worker_allocation/error.rs | 20 + .../runtime/worker_allocation/lifecycle.rs | 557 +++++- .../src/runtime/worker_allocation/mutate.rs | 25 +- .../src/runtime/worker_allocation/table.rs | 9 + crates/worker/src/worker.rs | 788 ++++++++- crates/worker/tests/controller_test.rs | 205 +++ .../tests/restored_parent_control_tests.rs | 5 +- .../src/server_assignment_attention_tests.rs | 11 +- .../src/server_assignment_delivery_tests.rs | 11 +- .../src/server_workspace_config_wip_tests.rs | 5 +- 18 files changed, 3621 insertions(+), 547 deletions(-) create mode 100644 crates/worker-runtime/src/worker_backend_restore_failure_tests.rs diff --git a/crates/worker-runtime/src/execution.rs b/crates/worker-runtime/src/execution.rs index f15d5ca4e..56892006b 100644 --- a/crates/worker-runtime/src/execution.rs +++ b/crates/worker-runtime/src/execution.rs @@ -630,6 +630,16 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static { Ok(()) } + /// Completed resource-owner evidence retained while Runtime settlement may + /// still need retry. This is not a new request or execution authorization. + fn completed_restore_failure( + &self, + _operation_id: WorkerLifecycleOperationId, + _worker_ref: &WorkerRef, + ) -> Option { + None + } + fn reconcile_restore( &self, request: WorkerExecutionRestoreRequest, @@ -898,6 +908,15 @@ impl WorkerExecutionBackendRef { self.backend.preflight_restore(request) } + pub(crate) fn completed_restore_failure( + &self, + operation_id: WorkerLifecycleOperationId, + worker_ref: &WorkerRef, + ) -> Option { + self.backend + .completed_restore_failure(operation_id, worker_ref) + } + pub(crate) fn reconcile_restore( &self, request: WorkerExecutionRestoreRequest, diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 2b671c6e5..9f41c41f4 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -2489,6 +2489,14 @@ impl Runtime { .map(|result| (Some(result), None)) } + #[cfg(all(test, feature = "fs-store"))] + pub(crate) fn test_runtime_store(&self) -> FsRuntimeStore { + match &self.lock().unwrap().persistence { + RuntimePersistence::Fs(store) => store.clone(), + RuntimePersistence::Memory => panic!("expected filesystem Runtime"), + } + } + #[cfg(test)] pub(crate) fn test_restore_request( &self, @@ -2944,7 +2952,11 @@ impl Runtime { return Err(error); } - let restore_result = if reconcile { + let completed_failure = backend.completed_restore_failure(operation_id, worker_ref); + let has_completed_failure = completed_failure.is_some(); + let restore_result = if let Some(result) = completed_failure { + result + } else if reconcile { backend.reconcile_restore(request) } else { backend.restore_worker(request) @@ -3069,7 +3081,7 @@ impl Runtime { } } WorkerExecutionSpawnResult::Rejected(result) => { - if reconcile { + if reconcile && !has_completed_failure { return Ok(RuntimeWorkerRestoreResult::failed( WorkerRestoreState::ReconciliationRequired, "worker_restore_reconciliation_rejected", diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 74044cde9..55617d312 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -175,6 +175,163 @@ pub struct RuntimeWorkerController { pub workspace_client: Arc, } +/// Failure facts supplied by the resource owner, never inferred from error text. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RuntimeRestoreStage { + Preflight, + Metadata, + Reconstruction, + Binding, + Feature, + Controller, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RuntimeRestoreResources { + None, + WorkerStartup, + WorkdirBinding, + Unproven, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum RuntimeRestoreCleanupFailure { + #[error("restore resource ownership unavailable")] + Ownership, + #[error("restore Workdir close failed")] + Workdir, + #[error("restore run artifact cleanup failed")] + Artifacts, + #[error("restore allocation release unconfirmed")] + Allocation, + #[error("restore Feature task settlement unconfirmed")] + Feature, + #[error("factory cannot confirm restore cleanup")] + Unsupported, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RuntimeRestoreCleanup { + NotNeeded, + Required, + Confirmed, + Failed(RuntimeRestoreCleanupFailure), +} + +/// Diagnostics are fixed owner-selected causes, not arbitrary provider messages, +/// paths, configuration, or credentials. Cleanup cannot overwrite the first cause. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RuntimeRestoreFailure { + pub stage: RuntimeRestoreStage, + pub resources: RuntimeRestoreResources, + cause: &'static str, + cleanup: RuntimeRestoreCleanup, + io_kind: Option, + cleanup_failure: Option, +} + +impl RuntimeRestoreFailure { + pub fn refused(stage: RuntimeRestoreStage, cause: &'static str) -> Self { + Self { + stage, + resources: RuntimeRestoreResources::None, + cause, + cleanup: RuntimeRestoreCleanup::NotNeeded, + io_kind: None, + cleanup_failure: None, + } + } + + pub fn unproven(stage: RuntimeRestoreStage, cause: &'static str) -> Self { + Self { + stage, + resources: RuntimeRestoreResources::Unproven, + cause, + cleanup: RuntimeRestoreCleanup::Required, + io_kind: None, + cleanup_failure: None, + } + } + + fn started(stage: RuntimeRestoreStage, cause: &'static str) -> Self { + Self { + stage, + resources: RuntimeRestoreResources::WorkerStartup, + cause, + cleanup: RuntimeRestoreCleanup::Required, + io_kind: None, + cleanup_failure: None, + } + } + + fn needs_cleanup(&self) -> bool { + matches!( + self.cleanup, + RuntimeRestoreCleanup::Required | RuntimeRestoreCleanup::Failed(_) + ) + } +} + +impl std::fmt::Display for RuntimeRestoreFailure { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let cause: String = self.cause.chars().take(160).collect(); + write!( + f, + "restore {:?}: {}; io={:?}; resources={:?}; cleanup={:?}; prior_cleanup={:?}", + self.stage, cause, self.io_kind, self.resources, self.cleanup, self.cleanup_failure + ) + } +} + +fn restore_worker_cause(error: &WorkerError) -> &'static str { + match error { + WorkerError::ConstructionCleanup { source, .. } => restore_worker_cause(source), + WorkerError::SegmentEmpty { .. } => "saved segment empty", + WorkerError::SegmentSystemPromptMissing { .. } => "saved segment system prompt missing", + WorkerError::WorkerMetadataPending { .. } => "saved Worker metadata pending", + WorkerError::WorkerMetadataMissing { .. } => "saved Worker metadata missing", + WorkerError::WorkerMetadataInactive { .. } => "saved Worker metadata inactive", + WorkerError::WorkerMetadataManifestSnapshotMissing { .. } => { + "saved manifest snapshot missing" + } + WorkerError::ScopeLock(_) => "Worker allocation acquisition failed", + WorkerError::Store(_) => "saved segment read failed", + WorkerError::WorkerStore(_) => "Worker metadata read failed", + WorkerError::Provider(_) => "saved model provider initialization failed", + WorkerError::SubjektivSessionAttribution { .. } => "restored attribution failed", + WorkerError::FeatureInstall(_) | WorkerError::FeatureLifecycle(_) => { + "restored Feature installation failed" + } + _ => "saved Worker reconstruction failed", + } +} + +async fn finish_restore_failure( + factory: &F, + request: &WorkerExecutionRestoreRequest, + mut failure: RuntimeRestoreFailure, +) -> RuntimeRestoreFailure { + if failure.needs_cleanup() { + let cleanup = if failure.resources == RuntimeRestoreResources::WorkdirBinding { + // Mutable materialization is outside the factory's cleanup authority. + // Preserve the exact accepted operation for its resource owner; never + // infer rollback from a factory that did not acquire these bindings. + Err(RuntimeRestoreCleanupFailure::Unsupported) + } else { + factory.cleanup_failed_restore(request).await + }; + failure.cleanup = match cleanup { + Ok(()) => RuntimeRestoreCleanup::Confirmed, + Err(error) => { + failure.cleanup_failure.get_or_insert(error); + RuntimeRestoreCleanup::Failed(error) + } + }; + } + tracing::warn!(failure = %failure, "Worker restore failed; operation-owned cleanup evidence retained"); + failure +} + /// Factory seam used by [`WorkerRuntimeExecutionBackend`] to construct a real /// controller-backed Worker for a Runtime catalog entry. #[async_trait] @@ -235,7 +392,7 @@ pub trait RuntimeWorkerFactory: Send + Sync + 'static { async fn restore_controller( &self, request: WorkerExecutionRestoreRequest, - ) -> Result; + ) -> Result; /// Prove cleanup of a restore future that completed without returning a /// controller. An error is not absence evidence; unsupported factories keep @@ -243,8 +400,8 @@ pub trait RuntimeWorkerFactory: Send + Sync + 'static { async fn cleanup_failed_restore( &self, _request: &WorkerExecutionRestoreRequest, - ) -> Result<(), String> { - Err("Runtime Worker factory cannot prove failed restore cleanup".to_string()) + ) -> Result<(), RuntimeRestoreCleanupFailure> { + Err(RuntimeRestoreCleanupFailure::Unsupported) } /// Verify/release factory-owned resources when the backend registry has no @@ -508,20 +665,73 @@ pub struct ProfileRuntimeWorkerFactory { workspace_request_clients: Arc>, embedded_worker_mutation_dispatcher: Option>, controller_transport: WorkerControllerTransport, - failed_restore_sessions: Arc< - Mutex< - HashMap< - WorkerRef, - ( - crate::execution::WorkerLifecycleOperationId, - Arc, - ), - >, - >, - >, + failed_restore_resources: Arc>>, + #[cfg(test)] + restore_faults: Arc, +} + +#[derive(Clone)] +struct FailedRestoreResources { + operation_id: crate::execution::WorkerLifecycleOperationId, + sessions: Vec>, + allocation: Option, + feature_tasks: Option, +} + +#[cfg(test)] +#[derive(Default)] +struct RestoreFaults { + fail_at: Mutex>, + executions: std::sync::atomic::AtomicUsize, + cleanup_fail: AtomicBool, + socket_failure: AtomicBool, + factory_gate: Mutex>>, + factory_entered: Mutex>>, + cleanup_gate: Mutex>>, + cleanup_entered: Mutex>>, } impl ProfileRuntimeWorkerFactory { + fn retain_construction_failure( + &self, + request: &WorkerExecutionRestoreRequest, + error: &WorkerError, + ) { + let mut resources = self + .failed_restore_resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let Some(owner) = resources.get_mut(&request.worker_ref) else { + return; + }; + if owner.operation_id != request.operation_id { + return; + } + match error { + WorkerError::ConstructionCleanup { + allocation, + sessions, + .. + } => { + owner.allocation = allocation.clone(); + owner.sessions.push(sessions.clone()); + } + WorkerError::ScopeLock(error) => owner.allocation = error.allocation_cleanup_handle(), + _ => {} + } + } + + #[cfg(test)] + fn restore_checkpoint(&self, stage: RuntimeRestoreStage) -> Result<(), String> { + let mut failure = self.restore_faults.fail_at.lock().unwrap(); + if *failure == Some(stage) { + failure.take(); + Err("injected restore boundary failure".into()) + } else { + Ok(()) + } + } + pub fn new(profile_base_dir: impl Into) -> Self { let profile_base_dir = profile_base_dir.into(); Self { @@ -536,7 +746,9 @@ impl ProfileRuntimeWorkerFactory { workspace_request_clients: Arc::new(HashMap::new()), embedded_worker_mutation_dispatcher: None, controller_transport: WorkerControllerTransport::InProcess, - failed_restore_sessions: Arc::new(Mutex::new(HashMap::new())), + failed_restore_resources: Arc::new(Mutex::new(HashMap::new())), + #[cfg(test)] + restore_faults: Arc::new(RestoreFaults::default()), } } @@ -1708,265 +1920,403 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { async fn restore_controller( &self, request: WorkerExecutionRestoreRequest, - ) -> Result { - self.preflight_restore(&request).await?; - let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref); - let only_binding = (request.workdir_attachments.len() == 1) - .then(|| request.workdir_attachments.values().next()) - .flatten(); - let filesystem_authority = only_binding - .map(|binding| { - WorkerFilesystemAuthority::local( - binding.root().to_path_buf(), - binding.cwd().to_path_buf(), - ) - }) - .unwrap_or(WorkerFilesystemAuthority::None); - let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request( - &request.request, - self.runtime_id.as_deref(), - ); - let observation_runtime_id = self.runtime_id.clone(); - let observation_workspace_id = request - .request - .workspace_api - .as_ref() - .map(|api| api.workspace_id.clone()); - let observation_grants = request.request.worker_observation_grants.clone(); - let observation_enabled = request.request.worker_observation_enabled; - let workspace_request_client = request - .request - .workspace_api - .as_ref() - .and_then(|api| self.workspace_request_clients.get(&api.workspace_id)); - let workspace_context = workspace_backend_ref.worker_context( - &request.worker_ref, - request.workspace_scope.as_ref(), - self.worker_mutation_identity.as_ref(), - workspace_request_client, - self.embedded_worker_mutation_dispatcher.as_ref(), - Some(self.prompt_projection_cache.clone()), - ); - - let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?; - let session_dir = worker_aggregate_dir.join("session"); - let session_store = - WorkerSessionStore::new_with_retention(&session_dir, self.session_retention()?) - .map_err(|err| { - format!( - "failed to initialize canonical Worker Session store at {}: {err}", - session_dir.display() - ) - })?; - let worker_metadata_store = - WorkerAggregateStore::new(&worker_aggregate_dir, worker_name.clone()).map_err( - |err| { - format!( - "failed to initialize canonical Worker metadata store at {}: {err}", - worker_aggregate_dir.display() - ) - }, - )?; - let metadata = worker_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())?; - let (manifest, loader) = Self::manifest_for_restore(&metadata, &request.request)?; - let store = CombinedStore::new(session_store, worker_metadata_store); - - let mut worker = match Worker::restore_from_worker_metadata_with_context( - &worker_name, - manifest.clone(), - store.clone(), - loader.clone(), - workspace_context.clone(), - filesystem_authority.clone(), - ) - .await - { - Ok(worker) => worker, - Err(WorkerError::WorkerMetadataPending { .. }) - if request.request.initial_input.is_none() => - { - let pending_loader = if workspace_context.workspace_id().is_some() { - let bundle = request.config_bundle.as_ref().ok_or_else(|| { - "pending Workspace Worker restore requires operation-owned launch material; generic restore must not reconstruct it from current Workspace config" - .to_string() - })?; - let resolution = self - .observe_bundle_prompt_projection( - bundle, - observation_workspace_id.as_deref(), - )? - .ok_or_else(|| { - "pending Workspace Worker restore requires a saved Workspace Prompt projection" - .to_string() - })?; - loader - .clone() - .with_effective_catalog(resolution.projection.catalog.clone()) - } else { - loader.clone() - }; - Worker::restore_pending_from_worker_metadata_with_context( - &worker_name, - manifest.clone(), - store, - pending_loader, - workspace_context, - filesystem_authority, - ) - .await - .map_err(|err| format!("failed to recreate pending Worker from metadata: {err}"))? - } - Err(err) => return Err(format!("failed to restore Worker from metadata: {err}")), - }; - if let Some(binding) = request.request.backend_job.as_ref() { - worker - .bind_backend_job(binding.clone()) - .map_err(str::to_string)?; - } - validate_worker_memory_settings(worker.manifest(), &request.request)?; - worker - .finalize_subjektiv_session_attribution( - request.request.subjektiv_attached, - SubjektivSessionAttributionLifecycle::RestoredSession, + ) -> Result { + self.preflight_restore(&request).await.map_err(|_| { + RuntimeRestoreFailure::refused( + RuntimeRestoreStage::Preflight, + "restore preflight refused", ) - .map_err(|error| format!("finalize restored Worker Session attribution: {error}"))?; - let flow_transition_enabled = worker.manifest().feature.flow.enabled; - let workdir_sessions = restored_workdir_router( - &request.workdir_attachments, - &request.logical_workdir_attachments, - &request.request.workdir_attachments, - worker.scope().clone(), - worker.workspace_client_handle(), - )?; - worker.bind_workdir_sessions(workdir_sessions); - if let (Some(runtime_id), Some(workspace_id)) = - (observation_runtime_id, observation_workspace_id.clone()) - && observation_enabled - { - let mut providers: Vec> = vec![Arc::new( - WorkspaceClientWorkerObservationProvider::new(worker.workspace_client_handle()), - )]; - if !observation_grants.is_empty() { - providers.push(Arc::new(RuntimeGrantedWorkerObservationProvider { - runtime_id, - workspace_id, - grants: observation_grants.into_iter().take(100).collect(), - hub: self.observation_hub.clone(), - })); - } - worker.bind_worker_observation_provider(Some(Arc::new( - CompositeWorkerObservationProvider::new(providers), - ))); - } - if flow_transition_enabled { - let report = worker - .install_runtime_flow_transition_feature() - .map_err(|error| format!("install Flow transition feature: {error}"))?; - if report.reports.iter().any(|report| !report.installed) { - return Err(format!( - "install Flow transition feature failed: {:?}", - report.reports - )); - } - } - - let workspace_client = worker.workspace_client_handle(); - 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 cleanup_sessions = worker.workdir_sessions(); - let started = PreparedWorker::new( - worker, - WorkerBootstrapLayout::RuntimeManagedRun { - run_dir: run_dir.clone(), - bash_output_dir, - }, - self.controller_transport, - ) - .start() - .await - .map_err(|error| match error { - WorkerBootstrapError::Worker(source) => { - format!("failed to prepare restored Worker: {source}") - } - WorkerBootstrapError::Controller { - source, - cleanup_failed, - } => { - if cleanup_failed { - self.failed_restore_sessions - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert( - request.worker_ref.clone(), - (request.operation_id, cleanup_sessions), - ); - } - format!( - "failed to spawn restored Worker controller in {}: {source}", - run_dir.display() - ) - } })?; - let (handle, shutdown_rx, controller_task) = - (started.handle, started.shutdown, started.controller_task); - if flow_transition_enabled { - handle.shared_state.enable_flow_transition(); + #[cfg(test)] + self.restore_faults + .executions + .fetch_add(1, Ordering::SeqCst); + let mut stage = RuntimeRestoreStage::Metadata; + let mut io_kind = None; + // Reconstruction can write before failing; only the resource owner can + // confirm cleanup, even when no Worker value was returned. + let result: Result = async { + let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref); + let only_binding = (request.workdir_attachments.len() == 1) + .then(|| request.workdir_attachments.values().next()) + .flatten(); + let filesystem_authority = only_binding + .map(|binding| { + WorkerFilesystemAuthority::local( + binding.root().to_path_buf(), + binding.cwd().to_path_buf(), + ) + }) + .unwrap_or(WorkerFilesystemAuthority::None); + let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request( + &request.request, + self.runtime_id.as_deref(), + ); + let observation_runtime_id = self.runtime_id.clone(); + let observation_workspace_id = request + .request + .workspace_api + .as_ref() + .map(|api| api.workspace_id.clone()); + let observation_grants = request.request.worker_observation_grants.clone(); + let observation_enabled = request.request.worker_observation_enabled; + let workspace_request_client = request + .request + .workspace_api + .as_ref() + .and_then(|api| self.workspace_request_clients.get(&api.workspace_id)); + let workspace_context = workspace_backend_ref.worker_context( + &request.worker_ref, + request.workspace_scope.as_ref(), + self.worker_mutation_identity.as_ref(), + workspace_request_client, + self.embedded_worker_mutation_dispatcher.as_ref(), + Some(self.prompt_projection_cache.clone()), + ); + + let worker_aggregate_dir = self + .worker_aggregate_dir(&request.worker_ref) + .map_err(|_| "Worker aggregate root unavailable")?; + let session_dir = worker_aggregate_dir.join("session"); + let session_store = WorkerSessionStore::new_with_retention( + &session_dir, + self.session_retention() + .map_err(|_| "Session retention unavailable")?, + ) + .map_err(|_| "canonical Session store initialization failed")?; + let worker_metadata_store = + WorkerAggregateStore::new(&worker_aggregate_dir, worker_name.clone()) + .map_err(|_| "canonical metadata store initialization failed")?; + let metadata = worker_metadata_store + .read_by_name(&worker_name) + .map_err(|_| "Worker metadata read failed")? + .ok_or("saved Worker metadata unavailable")?; + let (manifest, loader) = Self::manifest_for_restore(&metadata, &request.request) + .map_err(|_| "saved manifest resolution failed")?; + let store = CombinedStore::new(session_store, worker_metadata_store); + + stage = RuntimeRestoreStage::Reconstruction; + self.failed_restore_resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert( + request.worker_ref.clone(), + FailedRestoreResources { + operation_id: request.operation_id, + sessions: Vec::new(), + allocation: None, + feature_tasks: None, + }, + ); + #[cfg(test)] + self.restore_checkpoint(stage) + .map_err(|_| "injected restore boundary failure")?; + let mut worker = match Worker::restore_from_worker_metadata_with_context( + &worker_name, + manifest.clone(), + store.clone(), + loader.clone(), + workspace_context.clone(), + filesystem_authority.clone(), + ) + .await + { + Ok(worker) => worker, + Err(WorkerError::WorkerMetadataPending { .. }) + if request.request.initial_input.is_none() => + { + let pending_loader = if workspace_context.workspace_id().is_some() { + let bundle = request + .config_bundle + .as_ref() + .ok_or("pending restore launch material missing")?; + let resolution = self + .observe_bundle_prompt_projection( + bundle, + observation_workspace_id.as_deref(), + ) + .map_err(|_| "pending restore Prompt projection invalid")? + .ok_or("pending restore Prompt projection missing")?; + loader + .clone() + .with_effective_catalog(resolution.projection.catalog.clone()) + } else { + loader.clone() + }; + Worker::restore_pending_from_worker_metadata_with_context( + &worker_name, + manifest.clone(), + store, + pending_loader, + workspace_context, + filesystem_authority, + ) + .await + .map_err(|error| { + self.retain_construction_failure(&request, &error); + restore_worker_cause(&error) + })? + } + Err(error) => { + self.retain_construction_failure(&request, &error); + return Err(restore_worker_cause(&error)); + } + }; + let feature_tasks = worker.retain_startup_feature_cleanup(); + self.failed_restore_resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert( + request.worker_ref.clone(), + FailedRestoreResources { + operation_id: request.operation_id, + sessions: vec![worker.workdir_sessions()], + allocation: worker.restore_allocation_cleanup_handle(), + feature_tasks: Some(feature_tasks), + }, + ); + #[cfg(test)] + { + let gate = self.restore_faults.factory_gate.lock().unwrap().take(); + if let Some(gate) = gate { + if let Some(entered) = + self.restore_faults.factory_entered.lock().unwrap().take() + { + let _ = entered.send(()); + } + gate.notified().await; + } + } + stage = RuntimeRestoreStage::Binding; + #[cfg(test)] + self.restore_checkpoint(stage) + .map_err(|_| "injected restore boundary failure")?; + if let Some(binding) = request.request.backend_job.as_ref() { + worker + .bind_backend_job(binding.clone()) + .map_err(|_| "restored Backend Job binding failed")?; + } + validate_worker_memory_settings(worker.manifest(), &request.request) + .map_err(|_| "saved memory settings mismatch")?; + worker + .finalize_subjektiv_session_attribution( + request.request.subjektiv_attached, + SubjektivSessionAttributionLifecycle::RestoredSession, + ) + .map_err(|error| restore_worker_cause(&error))?; + let flow_transition_enabled = worker.manifest().feature.flow.enabled; + let workdir_sessions = restored_workdir_router( + &request.workdir_attachments, + &request.logical_workdir_attachments, + &request.request.workdir_attachments, + worker.scope().clone(), + worker.workspace_client_handle(), + ) + .map_err(|_| "restored Workdir binding failed")?; + self.failed_restore_resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get_mut(&request.worker_ref) + .expect("retained restore owner") + .sessions + .push(workdir_sessions.clone()); + worker + .workdir_sessions() + .close_all() + .await + .map_err(|_| "initial Workdir close failed")?; + worker.bind_workdir_sessions(workdir_sessions); + if let (Some(runtime_id), Some(workspace_id)) = + (observation_runtime_id, observation_workspace_id.clone()) + && observation_enabled + { + let mut providers: Vec> = vec![Arc::new( + WorkspaceClientWorkerObservationProvider::new(worker.workspace_client_handle()), + )]; + if !observation_grants.is_empty() { + providers.push(Arc::new(RuntimeGrantedWorkerObservationProvider { + runtime_id, + workspace_id, + grants: observation_grants.into_iter().take(100).collect(), + hub: self.observation_hub.clone(), + })); + } + worker.bind_worker_observation_provider(Some(Arc::new( + CompositeWorkerObservationProvider::new(providers), + ))); + } + stage = RuntimeRestoreStage::Feature; + #[cfg(test)] + self.restore_checkpoint(stage) + .map_err(|_| "injected restore boundary failure")?; + if flow_transition_enabled { + let report = worker + .install_runtime_flow_transition_feature() + .map_err(|_| "Flow transition installation failed")?; + if report.reports.iter().any(|report| !report.installed) { + return Err("Flow transition installation refused"); + } + } + + stage = RuntimeRestoreStage::Controller; + #[cfg(test)] + self.restore_checkpoint(stage) + .map_err(|_| "injected restore boundary failure")?; + let workspace_client = worker.workspace_client_handle(); + let run_dir = self + .worker_restore_run_dir(&request.worker_ref, request.operation_id) + .map_err(|_| "restore run directory unavailable")?; + let bash_output_dir = bash_output_dir_for_worker_id(&request.worker_ref.worker_id); + let transport = self.controller_transport; + #[cfg(test)] + let transport = if self.restore_faults.socket_failure.load(Ordering::SeqCst) { + WorkerControllerTransport::UnixSocket + } else { + transport + }; + let started = PreparedWorker::new( + worker, + WorkerBootstrapLayout::RuntimeManagedRun { + run_dir: run_dir.clone(), + bash_output_dir, + }, + transport, + ) + .start() + .await + .map_err(|error| match error { + WorkerBootstrapError::Worker(error) => restore_worker_cause(&error), + WorkerBootstrapError::Controller { source, .. } => { + io_kind = Some(source.kind()); + "restored Worker controller startup failed" + } + })?; + self.failed_restore_resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&request.worker_ref); + let (handle, shutdown_rx, controller_task) = + (started.handle, started.shutdown, started.controller_task); + if flow_transition_enabled { + handle.shared_state.enable_flow_transition(); + } + Ok(RuntimeWorkerController { + handle, + shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))), + controller_task, + workspace_client, + }) } - Ok(RuntimeWorkerController { - handle, - shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))), - controller_task, - workspace_client, + .await; + result.map_err(|cause| { + let mut failure = RuntimeRestoreFailure::started(stage, cause); + failure.io_kind = io_kind; + failure }) } async fn cleanup_failed_restore( &self, request: &WorkerExecutionRestoreRequest, - ) -> Result<(), String> { + ) -> Result<(), RuntimeRestoreCleanupFailure> { + #[cfg(test)] + { + let gate = self.restore_faults.cleanup_gate.lock().unwrap().take(); + if let Some(gate) = gate { + if let Some(entered) = self.restore_faults.cleanup_entered.lock().unwrap().take() { + let _ = entered.send(()); + } + gate.notified().await; + } + if self.restore_faults.cleanup_fail.load(Ordering::SeqCst) { + return Err(RuntimeRestoreCleanupFailure::Workdir); + } + } let retained = self - .failed_restore_sessions + .failed_restore_resources .lock() - .map_err(|_| "failed restore sessions lock is poisoned".to_string())? + .map_err(|_| RuntimeRestoreCleanupFailure::Ownership)? .get(&request.worker_ref) .cloned(); - if let Some((operation_id, sessions)) = retained { - if operation_id != request.operation_id { - return Err("failed restore sessions belong to another operation".to_string()); + if let Some(resources) = retained { + if resources.operation_id != request.operation_id { + return Err(RuntimeRestoreCleanupFailure::Ownership); + } + if let Some(feature_tasks) = resources.feature_tasks { + feature_tasks + .shutdown() + .await + .map_err(|_| RuntimeRestoreCleanupFailure::Feature)?; + } + for sessions in resources.sessions { + sessions + .close_all() + .await + .map_err(|_| RuntimeRestoreCleanupFailure::Workdir)?; + } + if let Some(allocation) = resources.allocation { + allocation + .release() + .map_err(|_| RuntimeRestoreCleanupFailure::Allocation)?; + } else { + // Reconstruction may fail before returning its allocation guard. + // Absence is checked at the allocation authority, not inferred from Err. + let path = worker::runtime::worker_allocation::default_allocation_path() + .map_err(|_| RuntimeRestoreCleanupFailure::Allocation)?; + let table = worker::runtime::worker_allocation::LockFileGuard::open(&path) + .map_err(|_| RuntimeRestoreCleanupFailure::Allocation)?; + if table + .data() + .find(&Self::runtime_worker_name_for_ref(&request.worker_ref)) + .is_some() + { + return Err(RuntimeRestoreCleanupFailure::Allocation); + } } - sessions - .close_all() - .await - .map_err(|error| format!("failed restore session cleanup: {error}"))?; } // restore_controller starts the Controller only as its final fallible // step. A completed error leaves no live Controller; remove only this // operation's disposable run artifacts, never Session/metadata authority. - self.stopped_restore_run_dirs(&request.worker_ref)?; - let run_dir = self.worker_restore_run_dir(&request.worker_ref, request.operation_id)?; - Self::remove_restore_run_artifacts(&run_dir)?; - self.failed_restore_sessions + self.stopped_restore_run_dirs(&request.worker_ref) + .map_err(|_| RuntimeRestoreCleanupFailure::Artifacts)?; + let run_dir = self + .worker_restore_run_dir(&request.worker_ref, request.operation_id) + .map_err(|_| RuntimeRestoreCleanupFailure::Artifacts)?; + Self::remove_restore_run_artifacts(&run_dir) + .map_err(|_| RuntimeRestoreCleanupFailure::Artifacts)?; + self.failed_restore_resources .lock() - .map_err(|_| "failed restore sessions lock is poisoned".to_string())? + .map_err(|_| RuntimeRestoreCleanupFailure::Ownership)? .remove(&request.worker_ref); + self.observation_hub.get(&request.worker_ref); Ok(()) } async fn reconcile_stopped_worker(&self, worker_ref: &WorkerRef) -> Result<(), String> { let retained = self - .failed_restore_sessions + .failed_restore_resources .lock() .map_err(|_| "failed restore sessions lock is poisoned".to_string())? .get(worker_ref) .cloned(); - if let Some((_, sessions)) = retained.as_ref() { - sessions - .close_all() - .await - .map_err(|error| format!("failed restore session cleanup: {error}"))?; + if let Some(resources) = retained.as_ref() { + if let Some(feature_tasks) = resources.feature_tasks.as_ref() { + feature_tasks + .shutdown() + .await + .map_err(|_| "failed restore Feature cleanup".to_string())?; + } + for sessions in &resources.sessions { + sessions + .close_all() + .await + .map_err(|_| "failed restore Workdir cleanup".to_string())?; + } + if let Some(allocation) = resources.allocation.as_ref() { + allocation + .release() + .map_err(|_| "failed restore allocation cleanup".to_string())?; + } } // A restart loses the map, not the operation-owned filesystem artifacts. // Validate every candidate before deleting any, and retain map/journal @@ -1975,7 +2325,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { Self::remove_restore_run_artifacts(&run_dir)?; } if retained.is_some() { - self.failed_restore_sessions + self.failed_restore_resources .lock() .map_err(|_| "failed restore sessions lock is poisoned".to_string())? .remove(worker_ref); @@ -2310,7 +2660,7 @@ struct RuntimeWorkerExecution { struct PendingRuntimeRestore { request: WorkerExecutionRestoreRequest, task: tokio::sync::Mutex>>, - result: Mutex>>, + result: Mutex>>, failure: Mutex>, } @@ -2332,8 +2682,14 @@ impl PendingRuntimeRestore { } Ok(Err(error)) => { task.take(); - let message = - format!("restore factory task failed: {error}; cleanup is unproven"); + let message = format!( + "restore factory task failed ({}); cleanup is unproven", + if error.is_panic() { + "panic" + } else { + "cancelled" + } + ); *self .failure .lock() @@ -2363,6 +2719,18 @@ pub struct WorkerRuntimeExecutionBackend { runtime: Mutex>, workers: Mutex>, pending_restores: Mutex>>, + // One completed execution receipt per Worker, separate from live resources. + // Runtime owns durable request replay; retain this until the next admitted + // operation so a failed canonical settlement cannot re-execute the factory. + completed_restore_failures: Mutex< + HashMap< + WorkerRef, + ( + crate::execution::WorkerLifecycleOperationId, + RuntimeRestoreFailure, + ), + >, + >, worker_locks: Mutex>>>, spawn_restore_timeout: Duration, subject_reservations: Mutex>, @@ -2386,6 +2754,7 @@ where runtime: Mutex::new(Some(runtime)), workers: Mutex::new(HashMap::new()), pending_restores: Mutex::new(HashMap::new()), + completed_restore_failures: Mutex::new(HashMap::new()), worker_locks: Mutex::new(HashMap::new()), spawn_restore_timeout: SPAWN_RESTORE_TASK_TIMEOUT, subject_reservations: Mutex::new(HashMap::new()), @@ -2519,6 +2888,82 @@ where .map_err(|err| format!("worker adapter task did not complete: {err}"))? } + fn preflight_restore_inputs( + &self, + request: &WorkerExecutionRestoreRequest, + ) -> Result<(), WorkerExecutionResult> { + if !request.previous_workdir_attachments.is_empty() { + let materializer = self + .working_directory_materializer + .as_ref() + .ok_or_else(|| { + WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + "Persisted Worker Workdir attachments cannot be restored by this Runtime", + ) + })?; + for previous in &request.previous_workdir_attachments { + let relative_cwd = request + .request + .workdir_attachments + .iter() + .find(|claim| claim.alias == previous.alias) + .and_then(|claim| claim.relative_cwd.as_deref()); + // Recovery attachments remain CleanupPending across restart. + // Validate current bindability/identity, not the old state name. + // Preflight is read-only; actual restore binds again below. + materializer.preflight_bind_working_directory( + &previous.working_directory.summary.working_directory_id, + relative_cwd, + ).map_err(|message| { + WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + format!("Persisted Worker Workdir attachment `{}` is unavailable: {message}", previous.alias), + ) + })?; + } + } + if request.previous_workdir_attachments.is_empty() + && !request.request.workdir_attachment_requests.is_empty() + { + return Err(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + "Persisted Worker Workdir allocations are unavailable", + )); + } + if request.previous_workdir_attachments.is_empty() + && !request.request.workdir_attachments.is_empty() + && self.working_directory_materializer.is_none() + { + return Err(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + "Persisted Worker Workdir claims 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(()) + } + fn pending_restore( &self, worker_ref: &WorkerRef, @@ -2534,10 +2979,94 @@ where pending: &Arc, timeout: Duration, ) -> Result<(), String> { - let pending = Arc::clone(pending); + let owner = Arc::clone(pending); self.run_cancellable_on_adapter_runtime(timeout + Duration::from_secs(1), async move { - pending.wait(timeout).await + owner.wait(timeout).await }) + .map_err(|message| { + match pending.result.lock().ok().and_then(|result| { + result + .as_ref() + .and_then(|result| result.as_ref().err()) + .cloned() + }) { + Some(failure) => format!("{failure}; {message}"), + None => message, + } + }) + } + + // Binding failures cannot borrow the factory's cleanup proof. Retain the + // exact operation and any already acquired bindings for their actual owner. + fn retain_restore_binding_failure( + &self, + mut request: WorkerExecutionRestoreRequest, + bindings: BTreeMap, + ) -> WorkerExecutionSpawnResult { + request.workdir_attachments = bindings; + let mut failure = RuntimeRestoreFailure::unproven( + RuntimeRestoreStage::Binding, + "restore Workdir binding failed; materializer reconciliation required", + ); + failure.resources = RuntimeRestoreResources::WorkdirBinding; + failure.cleanup = RuntimeRestoreCleanup::Failed(RuntimeRestoreCleanupFailure::Unsupported); + failure.cleanup_failure = Some(RuntimeRestoreCleanupFailure::Unsupported); + self.pending_restores + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert( + request.worker_ref.clone(), + Arc::new(PendingRuntimeRestore { + request, + task: tokio::sync::Mutex::new(None), + result: Mutex::new(Some(Err(failure.clone()))), + failure: Mutex::new(None), + }), + ); + WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + failure.to_string(), + ), + worker_state: None, + workdir_attachments: Vec::new(), + } + } + + // Called only after joining the previous owned task, under the Worker lock. + // A caller timeout never cancels this close/join future. + fn retry_pending_restore_cleanup( + &self, + pending: &Arc, + ) -> Result<(), String> { + let failure = pending + .result + .lock() + .map_err(|_| "restore result lock is poisoned")? + .as_ref() + .and_then(|result| result.as_ref().err()) + .filter(|failure| failure.needs_cleanup()) + .cloned(); + if let Some(failure) = failure { + let owner = Arc::clone(pending); + let factory = Arc::clone(&self.factory); + let mut task = pending + .task + .try_lock() + .map_err(|_| "restore task is still owned")?; + if task.is_some() { + return Err("restore task has not been joined".into()); + } + *task = Some(self.spawn_on_adapter_runtime(async move { + let failure = + finish_restore_failure(factory.as_ref(), &owner.request, failure).await; + *owner + .result + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Err(failure)); + })?); + } + Ok(()) } fn remove_pending_restore(&self, worker_ref: &WorkerRef) -> Result<(), String> { @@ -3230,6 +3759,12 @@ where if let Err(message) = self.wait_pending_restore(&pending, Duration::from_secs(5)) { return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); } + if let Err(message) = self + .retry_pending_restore_cleanup(&pending) + .and_then(|()| self.wait_pending_restore(&pending, Duration::from_secs(5))) + { + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); + } if self.workers.lock().is_err() { return WorkerExecutionResult::errored( WorkerExecutionOperation::Stop, @@ -3264,17 +3799,14 @@ where Vec::new(), ); } - Some(Err(_)) => { - let factory = Arc::clone(&self.factory); - let request = pending.request.clone(); - if let Err(message) = self - .run_cancellable_on_adapter_runtime(RUNTIME_TASK_TIMEOUT, async move { - factory.cleanup_failed_restore(&request).await - }) - { + Some(Err(failure)) => { + if !matches!( + failure.cleanup, + RuntimeRestoreCleanup::NotNeeded | RuntimeRestoreCleanup::Confirmed + ) { return WorkerExecutionResult::errored( WorkerExecutionOperation::Stop, - message, + failure.to_string(), ); } } @@ -4067,76 +4599,29 @@ where } drop(workers); - if !request.previous_workdir_attachments.is_empty() { - let materializer = self - .working_directory_materializer - .as_ref() - .ok_or_else(|| { - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - "Persisted Worker Workdir attachments cannot be restored by this Runtime", - ) - })?; - for previous in &request.previous_workdir_attachments { - let relative_cwd = request - .request - .workdir_attachments - .iter() - .find(|claim| claim.alias == previous.alias) - .and_then(|claim| claim.relative_cwd.as_deref()); - // Recovery attachments remain CleanupPending across restart. - // Validate current bindability/identity, not the old state name. - // Preflight is read-only; actual restore binds again below. - materializer.preflight_bind_working_directory( - &previous.working_directory.summary.working_directory_id, - relative_cwd, - ).map_err(|message| { - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - format!("Persisted Worker Workdir attachment `{}` is unavailable: {message}", previous.alias), - ) - })?; - } - } - if request.previous_workdir_attachments.is_empty() - && !request.request.workdir_attachment_requests.is_empty() - { - return Err(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - "Persisted Worker Workdir allocations are unavailable", - )); - } - if request.previous_workdir_attachments.is_empty() - && !request.request.workdir_attachments.is_empty() - && self.working_directory_materializer.is_none() - { - return Err(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - "Persisted Worker Workdir claims 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", - )); - } - } + self.preflight_restore_inputs(request) + } - 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(()) + fn completed_restore_failure( + &self, + operation_id: crate::execution::WorkerLifecycleOperationId, + worker_ref: &WorkerRef, + ) -> Option { + let completed = self + .completed_restore_failures + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let (completed_id, failure) = completed.get(worker_ref)?; + if *completed_id != operation_id { + return None; + } + let result = + WorkerExecutionResult::errored(WorkerExecutionOperation::Restore, failure.to_string()); + Some(if failure.cleanup == RuntimeRestoreCleanup::NotNeeded { + WorkerExecutionSpawnResult::Rejected(result) + } else { + WorkerExecutionSpawnResult::RolledBack(result) + }) } fn restore_worker( @@ -4162,6 +4647,11 @@ where )); } }; + if let Some(result) = + self.completed_restore_failure(request.operation_id, &request.worker_ref) + { + return result; + } let workers = self.workers.lock().map_err(|_| { WorkerExecutionResult::errored( WorkerExecutionOperation::Restore, @@ -4206,93 +4696,138 @@ where } drop(workers); - let mut workdir_attachments = BTreeMap::new(); - if !request.previous_workdir_attachments.is_empty() { - let Some(materializer) = self.working_directory_materializer.as_ref() else { - return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - "persisted Worker has Workdir attachments, but no materializer is configured for this Runtime backend", - )); - }; - for status in request.previous_workdir_attachments.clone() { - let relative_cwd = request + // An already accepted operation owns its bindings. Re-observing mutable + // Workdir materialization must not turn an unknown operation into a new + // read-only refusal or allocate another candidate. + let admitted_pending = match self.pending_restore(&request.worker_ref) { + Ok(Some(pending)) if pending.request.operation_id == request.operation_id => { + Some(pending) + } + Ok(Some(_)) => { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::busy( + WorkerExecutionOperation::Restore, + "another restore operation still owns pending resources", + ), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + Ok(None) => None, + Err(message) => { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message, + ), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + }; + if let Some(pending) = admitted_pending.as_ref() { + request = pending.request.clone(); + } else { + if let Err(result) = self.preflight_restore_inputs(&request) { + return WorkerExecutionSpawnResult::Rejected(result); + } + let aliases = if request.previous_workdir_attachments.is_empty() { + request .request .workdir_attachments .iter() - .find(|claim| claim.alias == status.alias) - .and_then(|claim| claim.relative_cwd.as_deref()); - match materializer.bind_working_directory( - &status.working_directory.summary.working_directory_id, - relative_cwd, - ) { - Ok(binding) => { - if workdir_attachments - .insert(status.alias.clone(), binding) - .is_some() - { - return WorkerExecutionSpawnResult::Rejected( - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - format!( - "duplicate persisted Workdir attachment alias `{}`", - status.alias - ), - ), - ); - } - } - Err(error) => { - return WorkerExecutionSpawnResult::Rejected( - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - error.to_string(), - ), - ); - } - } - } - } else if !request.request.workdir_attachments.is_empty() { - let Some(materializer) = self.working_directory_materializer.as_ref() else { + .map(|claim| &claim.alias) + .collect::>() + } else { + request + .previous_workdir_attachments + .iter() + .map(|status| &status.alias) + .collect::>() + }; + if aliases + .iter() + .collect::>() + .len() + != aliases.len() + { return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( WorkerExecutionOperation::Restore, - "persisted Worker has Workdir claims, but no materializer is configured for this Runtime backend", + "duplicate restore Workdir attachment alias", )); - }; - for claim in request.request.workdir_attachments.clone() { - match materializer.bind_working_directory( - &claim.working_directory_id, - claim.relative_cwd.as_deref(), - ) { - Ok(binding) => { - if workdir_attachments - .insert(claim.alias.clone(), binding) - .is_some() - { - return WorkerExecutionSpawnResult::Rejected( - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - format!("duplicate Workdir attachment alias `{}`", claim.alias), - ), - ); + } + let mut workdir_attachments = BTreeMap::new(); + if !request.previous_workdir_attachments.is_empty() { + let Some(materializer) = self.working_directory_materializer.as_ref() else { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + "persisted Worker has Workdir attachments, but no materializer is configured for this Runtime backend", + )); + }; + for status in request.previous_workdir_attachments.clone() { + let relative_cwd = request + .request + .workdir_attachments + .iter() + .find(|claim| claim.alias == status.alias) + .and_then(|claim| claim.relative_cwd.as_deref()); + match materializer.bind_working_directory( + &status.working_directory.summary.working_directory_id, + relative_cwd, + ) { + Ok(binding) => { + if workdir_attachments + .insert(status.alias.clone(), binding) + .is_some() + { + return self + .retain_restore_binding_failure(request, workdir_attachments); + } + } + Err(_) => { + return self + .retain_restore_binding_failure(request, workdir_attachments); } } - Err(error) => { - return WorkerExecutionSpawnResult::Rejected( - WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - error.to_string(), - ), - ); + } + } else if !request.request.workdir_attachments.is_empty() { + let Some(materializer) = self.working_directory_materializer.as_ref() else { + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( + WorkerExecutionOperation::Restore, + "persisted Worker has Workdir claims, but no materializer is configured for this Runtime backend", + )); + }; + for claim in request.request.workdir_attachments.clone() { + match materializer.bind_working_directory( + &claim.working_directory_id, + claim.relative_cwd.as_deref(), + ) { + Ok(binding) => { + if workdir_attachments + .insert(claim.alias.clone(), binding) + .is_some() + { + return self + .retain_restore_binding_failure(request, workdir_attachments); + } + } + Err(_) => { + return self + .retain_restore_binding_failure(request, workdir_attachments); + } } } } + request.workdir_attachments = workdir_attachments.clone(); } - request.workdir_attachments = workdir_attachments.clone(); - let worker_ref = request.worker_ref.clone(); let operation_id = request.operation_id; + let mut existing_pending = false; let pending = match self.pending_restore(&worker_ref) { - Ok(Some(pending)) if pending.request.operation_id == operation_id => pending, + Ok(Some(pending)) if pending.request.operation_id == operation_id => { + existing_pending = true; + pending + } Ok(Some(_)) => { return WorkerExecutionSpawnResult::ReconciliationRequired { result: WorkerExecutionResult::busy( @@ -4304,16 +4839,10 @@ where }; } Ok(None) => { - let factory = Arc::clone(&self.factory); - let preflight_request = request.clone(); - if let Err(message) = self.run_on_adapter_runtime(async move { - factory.preflight_restore(&preflight_request).await - }) { - return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected( - WorkerExecutionOperation::Restore, - message, - )); - } + self.completed_restore_failures + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&worker_ref); let pending = Arc::new(PendingRuntimeRestore { request: request.clone(), task: tokio::sync::Mutex::new(None), @@ -4335,7 +4864,23 @@ where let factory = Arc::clone(&self.factory); let result_owner = Arc::clone(&pending); let task = self.spawn_on_adapter_runtime(async move { - let result = factory.restore_controller(request).await; + let result = match factory.restore_controller(request.clone()).await { + Ok(controller) => Ok(controller), + Err(mut failure) => { + if failure.resources == RuntimeRestoreResources::None + && !request.workdir_attachments.is_empty() + { + failure.resources = RuntimeRestoreResources::WorkdirBinding; + failure.cleanup = RuntimeRestoreCleanup::Required; + } + *result_owner + .result + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Some(Err(failure.clone())); + Err(finish_restore_failure(factory.as_ref(), &request, failure).await) + } + }; let mut retained = result_owner .result .lock() @@ -4377,6 +4922,21 @@ where workdir_attachments: Vec::new(), }; } + if existing_pending { + if let Err(message) = self + .retry_pending_restore_cleanup(&pending) + .and_then(|()| self.wait_pending_restore(&pending, self.spawn_restore_timeout)) + { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message, + ), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + } let controller = { let mut retained = match pending.result.lock() { Ok(retained) => retained, @@ -4396,12 +4956,39 @@ where Ok(controller) => controller, Err(_) => unreachable!(), }, - Some(Err(message)) => { + Some(Err(failure)) => { + let result = WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + failure.to_string(), + ); + let cleanup = failure.cleanup; + if matches!( + cleanup, + RuntimeRestoreCleanup::NotNeeded | RuntimeRestoreCleanup::Confirmed + ) { + self.completed_restore_failures + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(worker_ref.clone(), (operation_id, failure.clone())); + drop(retained); + if let Err(message) = self.remove_pending_restore(&worker_ref) { + return WorkerExecutionSpawnResult::ReconciliationRequired { + result: WorkerExecutionResult::errored( + WorkerExecutionOperation::Restore, + message, + ), + worker_state: None, + workdir_attachments: Vec::new(), + }; + } + return if cleanup == RuntimeRestoreCleanup::NotNeeded { + WorkerExecutionSpawnResult::Rejected(result) + } else { + WorkerExecutionSpawnResult::RolledBack(result) + }; + } return WorkerExecutionSpawnResult::ReconciliationRequired { - result: WorkerExecutionResult::errored( - WorkerExecutionOperation::Restore, - message.clone(), - ), + result, worker_state: None, workdir_attachments: Vec::new(), }; @@ -6508,7 +7095,7 @@ mod tests { } Err(error) => error, }; - assert!(error.contains(expected), "{error}"); + assert_eq!(error.stage, RuntimeRestoreStage::Preflight); assert_eq!(persisted_files(&runtime_store), before); assert!( factory @@ -6647,10 +7234,7 @@ mod tests { } Err(error) => error, }; - assert!( - error.contains("subjektiv settings snapshot mismatch"), - "{error}" - ); + assert_eq!(error.stage, RuntimeRestoreStage::Preflight); assert_eq!(persisted_files(&runtime_store), before); assert!(factory.observation_hub.workers.lock().unwrap().is_empty()); let run_dir = factory @@ -7421,7 +8005,7 @@ mod tests { async fn restore_controller( &self, _request: WorkerExecutionRestoreRequest, - ) -> Result { + ) -> Result { panic!("create timeout must not restore a Worker") } } @@ -7584,7 +8168,7 @@ mod tests { async fn restore_controller( &self, request: WorkerExecutionRestoreRequest, - ) -> Result { + ) -> Result { let request = WorkerExecutionSpawnRequest { worker_ref: request.worker_ref, request: request.request, @@ -7593,7 +8177,12 @@ mod tests { workdir_attachments: request.workdir_attachments, config_bundle: request.config_bundle, }; - self.spawn_controller(request).await + self.spawn_controller(request).await.map_err(|_| { + RuntimeRestoreFailure::unproven( + RuntimeRestoreStage::Controller, + "mock controller startup failed", + ) + }) } } @@ -7901,8 +8490,11 @@ mod tests { async fn restore_controller( &self, _request: WorkerExecutionRestoreRequest, - ) -> Result { - Err("restore failed".to_string()) + ) -> Result { + Err(RuntimeRestoreFailure::unproven( + RuntimeRestoreStage::Reconstruction, + "restore failed", + )) } } @@ -7926,7 +8518,7 @@ mod tests { async fn restore_controller( &self, request: WorkerExecutionRestoreRequest, - ) -> Result { + ) -> Result { *self.restore_id.lock().unwrap() = Some(request.operation_id); let runs = self .runtime_store @@ -7940,14 +8532,17 @@ mod tests { b"incomplete restore artifact", ) .unwrap(); - Err("injected restore failure with persisted artifacts".to_string()) + Err(RuntimeRestoreFailure::unproven( + RuntimeRestoreStage::Controller, + "injected restore failure with persisted artifacts", + )) } async fn cleanup_failed_restore( &self, _: &WorkerExecutionRestoreRequest, - ) -> Result<(), String> { - Err("injected restore artifact cleanup failure".to_string()) + ) -> Result<(), RuntimeRestoreCleanupFailure> { + Err(RuntimeRestoreCleanupFailure::Artifacts) } } @@ -8054,7 +8649,7 @@ mod tests { assert!( backend .factory - .failed_restore_sessions + .failed_restore_resources .lock() .unwrap() .is_empty() @@ -8107,16 +8702,18 @@ mod tests { fs::create_dir_all(run.parent().unwrap()).unwrap(); fs::write(&run, "invalid restore artifact").unwrap(); let sessions = Arc::new(WorkdirSessionRouter::new()); - factory.failed_restore_sessions.lock().unwrap().insert( + factory.failed_restore_resources.lock().unwrap().insert( worker_ref.clone(), - (request.operation_id, Arc::clone(&sessions)), + FailedRestoreResources { + operation_id: request.operation_id, + sessions: vec![Arc::clone(&sessions)], + allocation: None, + feature_tasks: None, + }, ); assert!( - factory - .cleanup_failed_restore(&request) - .await - .unwrap_err() - .contains("unsafe restore artifact") + factory.cleanup_failed_restore(&request).await.unwrap_err() + == RuntimeRestoreCleanupFailure::Artifacts ); assert!( factory @@ -8126,20 +8723,20 @@ mod tests { .contains("unsafe restore artifact") ); let retained = factory - .failed_restore_sessions + .failed_restore_resources .lock() .unwrap() .get(&worker_ref) .cloned() .unwrap(); - assert_eq!(retained.0, request.operation_id); - assert!(Arc::ptr_eq(&retained.1, &sessions)); + assert_eq!(retained.operation_id, request.operation_id); + assert!(Arc::ptr_eq(&retained.sessions[0], &sessions)); fs::remove_file(&run).unwrap(); fs::create_dir(&run).unwrap(); fs::write(run.join("worker.out.log"), "stale restore artifact").unwrap(); factory.reconcile_stopped_worker(&worker_ref).await.unwrap(); assert!(!run.exists()); - assert!(factory.failed_restore_sessions.lock().unwrap().is_empty()); + assert!(factory.failed_restore_resources.lock().unwrap().is_empty()); } #[cfg(unix)] @@ -8206,6 +8803,9 @@ mod tests { ); } + #[cfg(feature = "fs-store")] + include!("worker_backend_restore_failure_tests.rs"); + struct PendingFailureFactory { restore_calls: Arc, cleanup_calls: Arc, @@ -8223,20 +8823,23 @@ mod tests { async fn restore_controller( &self, _request: WorkerExecutionRestoreRequest, - ) -> Result { + ) -> Result { self.restore_calls.fetch_add(1, Ordering::SeqCst); tokio::time::sleep(Duration::from_millis(50)).await; - Err("injected uncertain restore side effect".to_string()) + Err(RuntimeRestoreFailure::unproven( + RuntimeRestoreStage::Reconstruction, + "injected uncertain restore side effect", + )) } async fn cleanup_failed_restore( &self, _request: &WorkerExecutionRestoreRequest, - ) -> Result<(), String> { + ) -> Result<(), RuntimeRestoreCleanupFailure> { self.cleanup_calls.fetch_add(1, Ordering::SeqCst); if self.cleanup_allowed.load(Ordering::SeqCst) { Ok(()) } else { - Err("injected pending restore cleanup failure".to_string()) + Err(RuntimeRestoreCleanupFailure::Workdir) } } } @@ -8307,12 +8910,7 @@ mod tests { failed.outcome, crate::execution::WorkerExecutionOutcome::Errored ); - assert!( - failed - .message - .unwrap() - .contains("pending restore cleanup failure") - ); + assert!(failed.message.unwrap().contains("Workdir")); assert_eq!( backend .pending_restore(&worker_ref) @@ -8330,7 +8928,7 @@ mod tests { "cleanup then catalog-commit retry must converge" ); assert_eq!(restore_calls.load(Ordering::SeqCst), 1); - assert_eq!(cleanup_calls.load(Ordering::SeqCst), 2); + assert_eq!(cleanup_calls.load(Ordering::SeqCst), 3); } struct LateControllerFactory { @@ -8351,7 +8949,7 @@ mod tests { async fn restore_controller( &self, request: WorkerExecutionRestoreRequest, - ) -> Result { + ) -> Result { let first = self.restores.fetch_add(1, Ordering::SeqCst) == 0; if first { self.release.notified().await; @@ -8831,7 +9429,7 @@ mod tests { Ok(_) => panic!("legacy Workspace Worker restore unexpectedly succeeded"), Err(error) => error, }; - assert!(error.contains("replacement Worker is required"), "{error}"); + assert_eq!(error.stage, RuntimeRestoreStage::Preflight); } #[test] diff --git a/crates/worker-runtime/src/worker_backend_restore_failure_tests.rs b/crates/worker-runtime/src/worker_backend_restore_failure_tests.rs new file mode 100644 index 000000000..af39d3619 --- /dev/null +++ b/crates/worker-runtime/src/worker_backend_restore_failure_tests.rs @@ -0,0 +1,945 @@ +// Integration coverage through EmbeddedRuntime -> real Profile factory -> real +// Worker/bootstrap, local Git Workdir provider, and allocation table. Model I/O +// is not invoked; this is a deterministic ordinary suite, not product E2E. + +struct ProfileFailureFixture { + runtime: EmbeddedRuntime, + backend: Arc, + worker: crate::catalog::WorkerDetail, + _repo: tempfile::TempDir, + root: tempfile::TempDir, +} + +impl ProfileFailureFixture { + fn new(timeout: Duration) -> Self { + let root = tempfile::tempdir().unwrap(); + let repo = create_clean_repo(); + let store = root.path().join("private-restore-test/runtime"); + let backend = Arc::new( + WorkerRuntimeExecutionBackend::new( + ProfileRuntimeWorkerFactory::new(root.path()).with_runtime_store_dir(&store), + ) + .unwrap() + .with_spawn_restore_timeout(timeout) + .with_working_directory_materializer(RuntimeGitMaterializer::new( + root.path().join("workdirs"), + )), + ); + let runtime = EmbeddedRuntime::with_fs_store_and_execution_backend( + crate::fs_store::FsRuntimeStoreOptions { + root: store, + runtime_id: "restore-failure-test".into(), + display_name: None, + }, + backend.clone(), + ) + .unwrap(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let mut request = create_request("restore failure recovery"); + request.profile = ProfileSelector::Builtin("default".into()); + request.workdir_attachment_requests = vec![WorkingDirectoryAttachmentRequest { + alias: WorkdirAttachmentAlias::new("checkout").unwrap(), + working_directory: working_directory_request(repo.path()), + }]; + let worker = runtime.create_worker(request).unwrap(); + runtime.stop_worker(&worker.worker_ref, None).unwrap(); + Self { + root, + _repo: repo, + runtime, + backend, + worker, + } + } + + fn fail(&self, stage: RuntimeRestoreStage) { + *self.backend.factory.restore_faults.fail_at.lock().unwrap() = Some(stage); + } + + fn assert_released(&self) { + assert!(self.backend.pending_restores.lock().unwrap().is_empty()); + assert!(self.backend.workers.lock().unwrap().is_empty()); + assert!( + self.backend + .factory + .failed_restore_resources + .lock() + .unwrap() + .is_empty() + ); + assert!( + self.backend + .factory + .observation_hub + .workers + .lock() + .unwrap() + .is_empty() + ); + let table = worker::runtime::worker_allocation::LockFileGuard::open( + &worker::runtime::worker_allocation::default_allocation_path().unwrap(), + ) + .unwrap(); + assert!( + table + .data() + .find(&ProfileRuntimeWorkerFactory::runtime_worker_name_for_ref( + &self.worker.worker_ref + )) + .is_none() + ); + let aggregate = self + .backend + .factory + .worker_aggregate_dir(&self.worker.worker_ref) + .unwrap(); + assert!(fs::read_dir(aggregate.join("runs")).unwrap().all(|entry| { + !entry + .unwrap() + .file_name() + .to_string_lossy() + .starts_with("restore-") + })); + assert_eq!( + self.runtime + .worker_detail(&self.worker.worker_ref) + .unwrap() + .status, + crate::catalog::WorkerStatus::Stopped + ); + } + + fn fresh_restore_and_old_replay( + &self, + request: runtime_api::WorkerRestoreRequest, + failed: crate::runtime::RuntimeWorkerRestoreResult, + ) { + let calls = self + .backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst); + let replay = self + .runtime + .restore_worker_operation(&self.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + serde_json::to_value(&replay).unwrap(), + serde_json::to_value(&failed).unwrap() + ); + assert_eq!( + self.backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + calls + ); + let fresh = self.runtime.test_restore_request(&self.worker.worker_ref); + assert_ne!(fresh.request_id, request.request_id); + assert_ne!( + fresh.expected_observation_token, + request.expected_observation_token + ); + let restored = self + .runtime + .restore_worker_operation(&self.worker.worker_ref, fresh) + .unwrap(); + assert_eq!( + restored.state, + server_api::WorkerRestoreState::Accepted, + "{restored:?}" + ); + assert_eq!( + self.backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + calls + 1 + ); + assert_eq!(self.backend.workers.lock().unwrap().len(), 1); + let replay = self + .runtime + .restore_worker_operation(&self.worker.worker_ref, request) + .unwrap(); + assert_eq!( + serde_json::to_value(&replay).unwrap(), + serde_json::to_value(&failed).unwrap() + ); + assert_eq!( + self.backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + calls + 1 + ); + // Only tear down after the fresh explicit Restore has already succeeded. + self.runtime + .stop_worker(&self.worker.worker_ref, None) + .unwrap(); + } +} + +#[test] +#[serial_test::serial(worker_allocation)] +fn profile_restore_multistage_failure_cleans_then_fresh_restore_without_stop() { + for stage in [ + RuntimeRestoreStage::Reconstruction, + RuntimeRestoreStage::Binding, + RuntimeRestoreStage::Feature, + RuntimeRestoreStage::Controller, + ] { + let fixture = ProfileFailureFixture::new(Duration::from_secs(5)); + fixture.fail(stage); + let request = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::RolledBack, + "{stage:?}: {result:?}" + ); + let diagnostic = serde_json::to_string(&result).unwrap(); + assert!( + diagnostic.contains("injected restore boundary failure"), + "{diagnostic}" + ); + assert!(diagnostic.contains(&format!("{stage:?}")), "{diagnostic}"); + assert!(!diagnostic.contains(&fixture.root.path().to_string_lossy().to_string())); + assert!(diagnostic.len() < 1024); + fixture.assert_released(); + eprintln!( + "restore-fault stage={stage:?} outcome=RolledBack pending=0 controllers=0 allocations=0" + ); + fixture.fresh_restore_and_old_replay(request, result); + } +} + +#[test] +#[serial_test::serial(worker_allocation)] +fn profile_restore_cleanup_failure_retains_owner_and_original_cause_until_same_operation_recovers() +{ + let fixture = ProfileFailureFixture::new(Duration::from_secs(5)); + fixture.fail(RuntimeRestoreStage::Feature); + fixture + .backend + .factory + .restore_faults + .cleanup_fail + .store(true, Ordering::SeqCst); + let request = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let failed = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + failed.state, + server_api::WorkerRestoreState::ReconciliationRequired + ); + let pending = fixture + .backend + .pending_restore(&fixture.worker.worker_ref) + .unwrap() + .unwrap(); + let message = pending + .result + .lock() + .unwrap() + .as_ref() + .unwrap() + .as_ref() + .err() + .unwrap() + .to_string(); + assert!( + message.contains("injected restore boundary failure"), + "{message}" + ); + assert!(message.contains("Workdir"), "{message}"); + let owned = fixture + .backend + .factory + .failed_restore_resources + .lock() + .unwrap() + .get(&fixture.worker.worker_ref) + .cloned() + .unwrap(); + assert_eq!(owned.operation_id, pending.request.operation_id); + assert!(owned.allocation.is_some()); + assert!( + owned + .sessions + .iter() + .any(|sessions| !sessions.aliases().is_empty()) + ); + let repeat = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + repeat.state, + server_api::WorkerRestoreState::ReconciliationRequired + ); + assert!(Arc::ptr_eq( + &pending, + &fixture + .backend + .pending_restore(&fixture.worker.worker_ref) + .unwrap() + .unwrap() + )); + assert_eq!( + fixture + .backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + 1 + ); + fixture + .backend + .factory + .restore_faults + .cleanup_fail + .store(false, Ordering::SeqCst); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::RolledBack, + "{result:?}" + ); + assert!( + serde_json::to_string(&result) + .unwrap() + .contains("injected restore boundary failure") + ); + assert!( + owned + .sessions + .iter() + .all(|sessions| sessions.aliases().is_empty()) + ); + fixture.assert_released(); + eprintln!( + "restore-fault cleanup=Workdir owner-retained=true retry=same-operation outcome=RolledBack" + ); + fixture.fresh_restore_and_old_replay(request, result); +} + +#[test] +#[serial_test::serial(worker_allocation)] +fn profile_restore_cleanup_timeout_keeps_running_join_and_late_failure_reconciles_without_execution() + { + let fixture = ProfileFailureFixture::new(Duration::from_secs(2)); + fixture.fail(RuntimeRestoreStage::Feature); + let release = Arc::new(tokio::sync::Notify::new()); + let (entered_tx, entered) = mpsc::channel(); + *fixture + .backend + .factory + .restore_faults + .cleanup_gate + .lock() + .unwrap() = Some(release.clone()); + *fixture + .backend + .factory + .restore_faults + .cleanup_entered + .lock() + .unwrap() = Some(entered_tx); + let request = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + entered.recv_timeout(Duration::from_secs(5)).unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::ReconciliationRequired + ); + let pending = fixture + .backend + .pending_restore(&fixture.worker.worker_ref) + .unwrap() + .unwrap(); + assert!( + pending + .result + .lock() + .unwrap() + .as_ref() + .unwrap() + .as_ref() + .err() + .unwrap() + .to_string() + .contains("injected restore boundary failure") + ); + assert!(pending.task.try_lock().unwrap().is_some()); + for _ in 0..2 { + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::ReconciliationRequired + ); + assert!(Arc::ptr_eq( + &pending, + &fixture + .backend + .pending_restore(&fixture.worker.worker_ref) + .unwrap() + .unwrap() + )); + } + assert_eq!( + fixture + .backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + 1 + ); + release.notify_one(); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::RolledBack, + "{result:?}" + ); + fixture.assert_released(); + eprintln!( + "restore-fault cleanup=timeout retained-task=true late-completion=joined executions=1" + ); + fixture.fresh_restore_and_old_replay(request, result); +} + +#[cfg(unix)] +#[test] +#[serial_test::serial(worker_allocation)] +fn profile_restore_real_socket_start_failure_removes_partial_run_and_allocation() { + let fixture = ProfileFailureFixture::new(Duration::from_secs(5)); + fixture + .backend + .factory + .restore_faults + .socket_failure + .store(true, Ordering::SeqCst); + let request = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::RolledBack, + "{result:?}" + ); + let diagnostic = serde_json::to_string(&result).unwrap(); + assert!( + diagnostic.contains("controller startup failed"), + "{diagnostic}" + ); + assert!(diagnostic.contains("InvalidInput"), "{diagnostic}"); + assert!(!diagnostic.contains("private-restore-test")); + fixture.assert_released(); + fixture + .backend + .factory + .restore_faults + .socket_failure + .store(false, Ordering::SeqCst); + eprintln!( + "restore-fault stage=socket-start real-bootstrap=true partial-run=removed allocation=released" + ); + fixture.fresh_restore_and_old_replay(request, result); +} + +#[test] +#[serial_test::serial(worker_allocation)] +fn profile_restore_read_only_refusal_replays_and_fresh_observation_can_restore() { + let fixture = ProfileFailureFixture::new(Duration::from_secs(5)); + let aggregate = fixture + .backend + .factory + .worker_aggregate_dir(&fixture.worker.worker_ref) + .unwrap(); + fs::rename(aggregate.join("session"), aggregate.join("saved-session")).unwrap(); + let before = persisted_files(&aggregate); + let request = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::Rejected, + "{result:?}" + ); + let mut after = persisted_files(&aggregate); + let mut before = before; + // The Runtime rejection receipt is expected to commit; source authority is not. + before.remove(Path::new("worker.json")); + after.remove(Path::new("worker.json")); + assert_eq!(after, before); + assert_eq!( + fixture + .backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + 0 + ); + fs::rename(aggregate.join("saved-session"), aggregate.join("session")).unwrap(); + // Creation's stopped observation entry is weak; querying it evicts it. + assert!( + fixture + .backend + .factory + .observation_hub + .get(&fixture.worker.worker_ref) + .is_none() + ); + fixture.assert_released(); + fixture.fresh_restore_and_old_replay(request, result); +} + +#[test] +#[serial_test::serial(worker_allocation)] +fn profile_restore_real_artifact_cleanup_failure_retains_same_owner_and_cause() { + let fixture = ProfileFailureFixture::new(Duration::from_secs(2)); + fixture.fail(RuntimeRestoreStage::Binding); + let release = Arc::new(tokio::sync::Notify::new()); + let (entered_tx, entered) = mpsc::channel(); + *fixture + .backend + .factory + .restore_faults + .factory_gate + .lock() + .unwrap() = Some(release.clone()); + *fixture + .backend + .factory + .restore_faults + .factory_entered + .lock() + .unwrap() = Some(entered_tx); + let request = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + entered.recv_timeout(Duration::from_secs(5)).unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::ReconciliationRequired + ); + let pending = fixture + .backend + .pending_restore(&fixture.worker.worker_ref) + .unwrap() + .unwrap(); + let run = fixture + .backend + .factory + .worker_restore_run_dir(&fixture.worker.worker_ref, pending.request.operation_id) + .unwrap(); + fs::write(&run, "invalid disposable run artifact").unwrap(); + release.notify_one(); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::ReconciliationRequired, + "{result:?}" + ); + let cause = pending + .result + .lock() + .unwrap() + .as_ref() + .unwrap() + .as_ref() + .err() + .unwrap() + .clone(); + assert_eq!(cause.stage, RuntimeRestoreStage::Binding); + assert_eq!( + cause.cleanup_failure, + Some(RuntimeRestoreCleanupFailure::Artifacts) + ); + assert!( + fixture + .backend + .factory + .failed_restore_resources + .lock() + .unwrap() + .contains_key(&fixture.worker.worker_ref) + ); + assert!(fixture.backend.workers.lock().unwrap().is_empty()); + assert_eq!( + fixture + .backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + 1 + ); + fs::remove_file(&run).unwrap(); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::RolledBack, + "{result:?}" + ); + assert!(result.message.as_ref().unwrap().contains("Artifacts")); + fixture.assert_released(); + eprintln!( + "restore-fault cleanup=real-artifact-failure late-factory=true owner=retained original-cause=preserved" + ); + fixture.fresh_restore_and_old_replay(request, result); +} + +#[test] +#[serial_test::serial(worker_allocation)] +fn profile_restore_stop_race_joins_late_factory_failure_without_second_controller() { + let fixture = ProfileFailureFixture::new(Duration::from_secs(2)); + fixture.fail(RuntimeRestoreStage::Binding); + let release = Arc::new(tokio::sync::Notify::new()); + let (entered_tx, entered) = mpsc::channel(); + *fixture + .backend + .factory + .restore_faults + .factory_gate + .lock() + .unwrap() = Some(release.clone()); + *fixture + .backend + .factory + .restore_faults + .factory_entered + .lock() + .unwrap() = Some(entered_tx); + let request = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request) + .unwrap(); + entered.recv_timeout(Duration::from_secs(5)).unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::ReconciliationRequired + ); + let pending = fixture + .backend + .pending_restore(&fixture.worker.worker_ref) + .unwrap() + .unwrap(); + assert!(pending.result.lock().unwrap().is_none()); + let runtime = fixture.runtime.clone(); + let worker_ref = fixture.worker.worker_ref.clone(); + let (started_tx, started) = mpsc::channel(); + let stopping = std::thread::spawn(move || { + started_tx.send(()).unwrap(); + runtime.stop_worker(&worker_ref, None) + }); + started.recv_timeout(Duration::from_secs(5)).unwrap(); + // The factory still owns the allocation and cannot attest stopped yet. + assert!( + !pending + .task + .try_lock() + .map(|task| task.as_ref().is_some_and(|task| task.is_finished())) + .unwrap_or(false) + ); + release.notify_one(); + stopping.join().unwrap().unwrap(); + fixture.assert_released(); + assert_eq!( + fixture + .backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + 1 + ); + let fresh = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, fresh) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::Accepted, + "{result:?}" + ); + assert_eq!(fixture.backend.workers.lock().unwrap().len(), 1); + fixture + .runtime + .stop_worker(&fixture.worker.worker_ref, None) + .unwrap(); + eprintln!("restore-fault race=Stop late-factory=joined controllers-at-most=1"); +} + +#[test] +#[serial_test::serial(worker_allocation)] +fn profile_restore_failure_settlement_write_retry_preserves_completed_cause_without_execution() { + for fault in [ + crate::fs_store::AtomicWriteFault::BeforeRename, + crate::fs_store::AtomicWriteFault::AfterRename, + ] { + let fixture = ProfileFailureFixture::new(Duration::from_secs(2)); + fixture.fail(RuntimeRestoreStage::Binding); + let release = Arc::new(tokio::sync::Notify::new()); + let (entered_tx, entered) = mpsc::channel(); + *fixture + .backend + .factory + .restore_faults + .factory_gate + .lock() + .unwrap() = Some(release.clone()); + *fixture + .backend + .factory + .restore_faults + .factory_entered + .lock() + .unwrap() = Some(entered_tx); + let request = fixture + .runtime + .test_restore_request(&fixture.worker.worker_ref); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + entered.recv_timeout(Duration::from_secs(5)).unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::ReconciliationRequired + ); + let operation_id = fixture + .backend + .pending_restore(&fixture.worker.worker_ref) + .unwrap() + .unwrap() + .request + .operation_id; + fixture + .runtime + .test_runtime_store() + .fail_next_worker_write(fault); + release.notify_one(); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::ReconciliationRequired, + "{result:?}" + ); + assert_eq!( + result.reason_code.as_deref(), + Some("worker_restore_rollback_commit_failed") + ); + assert!( + fixture + .backend + .pending_restore(&fixture.worker.worker_ref) + .unwrap() + .is_none() + ); + assert!( + fixture + .backend + .factory + .failed_restore_resources + .lock() + .unwrap() + .is_empty() + ); + let saved = fixture + .backend + .completed_restore_failure(operation_id, &fixture.worker.worker_ref) + .unwrap(); + assert!( + matches!(saved, WorkerExecutionSpawnResult::RolledBack(ref failure) + if failure.message.as_ref().unwrap().contains("injected restore boundary failure")) + ); + let result = fixture + .runtime + .restore_worker_operation(&fixture.worker.worker_ref, request.clone()) + .unwrap(); + assert_eq!( + result.state, + server_api::WorkerRestoreState::RolledBack, + "{result:?}" + ); + assert!( + result + .message + .as_ref() + .unwrap() + .contains("injected restore boundary failure") + ); + assert_eq!( + fixture + .backend + .factory + .restore_faults + .executions + .load(Ordering::SeqCst), + 1 + ); + fixture.assert_released(); + eprintln!( + "restore-fault settlement={fault:?} completed-cause=retained replay-executions=1" + ); + fixture.fresh_restore_and_old_replay(request, result); + } +} + +struct RacedBindingFactory { + changed_checkout: PathBuf, + preflights: AtomicUsize, +} + +#[async_trait] +impl RuntimeWorkerFactory for RacedBindingFactory { + async fn spawn_controller( + &self, + _: WorkerExecutionSpawnRequest, + ) -> Result { + panic!("binding-only failure must not create a controller") + } + async fn preflight_restore(&self, _: &WorkerExecutionRestoreRequest) -> Result<(), String> { + // Inject another owner's checkout change AFTER the real materializer's + // read-only checks, but before its mutable bind. Recovery stays real. + if self.preflights.fetch_add(1, Ordering::SeqCst) == 0 { + fs::remove_dir_all(&self.changed_checkout).unwrap(); + } + Ok(()) + } + async fn restore_controller( + &self, + _: WorkerExecutionRestoreRequest, + ) -> Result { + panic!("factory execution must not run after partial binding failed") + } +} + +#[test] +fn partial_restore_binding_failure_keeps_exact_operation_without_factory_cleanup_claim() { + let root = tempfile::tempdir().unwrap(); + let repo = create_clean_repo(); + let materializer = RuntimeGitMaterializer::new(root.path()); + let first = materializer + .create(&working_directory_request(repo.path())) + .unwrap(); + let second = materializer + .create(&working_directory_request(repo.path())) + .unwrap(); + let mut request = direct_restore_request(WorkerRef::new(WorkerId::now_v7())); + request.previous_workdir_attachments = vec![ + WorkingDirectoryAttachmentStatus { + alias: WorkdirAttachmentAlias::new("first").unwrap(), + working_directory: first.status(), + }, + WorkingDirectoryAttachmentStatus { + alias: WorkdirAttachmentAlias::new("second").unwrap(), + working_directory: second.status(), + }, + ]; + let backend = WorkerRuntimeExecutionBackend::new(RacedBindingFactory { + changed_checkout: second.root().to_path_buf(), + preflights: AtomicUsize::new(0), + }) + .unwrap() + .with_working_directory_materializer(materializer); + assert!(matches!( + backend.restore_worker(request.clone()), + WorkerExecutionSpawnResult::ReconciliationRequired { .. } + )); + let pending = backend + .pending_restore(&request.worker_ref) + .unwrap() + .unwrap(); + assert_eq!(pending.request.operation_id, request.operation_id); + assert_eq!(pending.request.workdir_attachments.len(), 1); + assert_eq!( + pending + .request + .workdir_attachments + .values() + .next() + .unwrap() + .root(), + first.root() + ); + for _ in 0..2 { + assert!(matches!( + backend.reconcile_restore(request.clone()), + WorkerExecutionSpawnResult::ReconciliationRequired { .. } + )); + } + assert_eq!(backend.factory.preflights.load(Ordering::SeqCst), 1); + assert!(!backend.stop_worker(&request.worker_ref).is_accepted()); + assert!(Arc::ptr_eq( + &pending, + &backend + .pending_restore(&request.worker_ref) + .unwrap() + .unwrap() + )); + assert!(backend.workers.lock().unwrap().is_empty()); + eprintln!( + "restore-fault stage=mutable-binding outcome=ReconciliationRequired partial-bindings=retained factory-executions=0" + ); + // Fixture teardown drops the retained passive local binding. This is not + // evidence that credential/policy materialization has been rolled back. +} diff --git a/crates/worker/src/bootstrap.rs b/crates/worker/src/bootstrap.rs index d027aa8c9..ba8afef39 100644 --- a/crates/worker/src/bootstrap.rs +++ b/crates/worker/src/bootstrap.rs @@ -271,6 +271,7 @@ where St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, { let cleanup_sessions = worker.workdir_sessions(); + let factory_owned_cleanup = worker.startup_feature_cleanup_handle().is_some(); let controller = match layout { WorkerBootstrapLayout::Direct { runtime_base, @@ -305,7 +306,10 @@ where controller_task, }), Err(source) => { - let cleanup_failed = cleanup_sessions.close_all().await.is_err(); + // Factory-owned task scopes must settle before sessions close. + // Cleanup remains unconfirmed here; preserve the startup source. + let cleanup_failed = + factory_owned_cleanup || cleanup_sessions.close_all().await.is_err(); Err(WorkerBootstrapError::Controller { source, cleanup_failed, diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index d8be6da1a..8c4640bf9 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -854,6 +854,7 @@ impl WorkerController { St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, { let sessions = worker.workdir_sessions(); + let factory_owned_cleanup = worker.startup_feature_cleanup_handle().is_some(); let result = Self::spawn_initialized( worker, runtime_base, @@ -863,7 +864,10 @@ impl WorkerController { transport, ) .await; + // An opted-in factory must quiesce its captured Feature tasks before + // closing their Workdir sessions; this wrapper cannot confirm that order. if result.is_err() + && !factory_owned_cleanup && let Err(error) = sessions.close_all().await { tracing::warn!(%error, "Workdir sessions close after controller startup failure failed"); @@ -1092,6 +1096,7 @@ impl WorkerController { let pause_tx = worker.engine_mut().pause_sender(); let notify_buffer = worker.notify_buffer_handle(); + worker.finish_startup_feature_cleanup_capture(); let controller_task = tokio::spawn(controller_loop( worker, method_rx, @@ -4424,21 +4429,6 @@ mod tests { assert!(feature_install < materialize_gate); assert!(materialize_gate < materialize); assert!(materialize < exposure); - let observer = startup.find("let command_observer = fs_for_view").unwrap(); - let recover = startup.find(".recover_unfinished_compaction()").unwrap(); - let files = startup - .find("runtime_dir.write_status(&shared_state).await?") - .unwrap(); - let loop_start = startup.find("tokio::spawn(controller_loop(").unwrap(); - assert!(recover < observer && files < observer && observer < loop_start); - let socket_startup = &startup[observer..loop_start]; - let abort = socket_startup.find("observer.abort()").unwrap(); - let join = socket_startup.find("observer.await").unwrap(); - let failure = socket_startup.find("return Err(error)").unwrap(); - assert!( - abort < join && join < failure, - "startup cannot detach its observer on error" - ); assert!( startup[materialize_gate..materialize] .contains("needs_initial_session_head_materialization") diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 10cf5a6c5..54145dbf9 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -69,9 +69,10 @@ pub use session_history::{ pub use shared_state::WorkerSharedState; pub use worker::{ LocalWorkingDirectory, SubjektivSessionAttributionLifecycle, Worker, WorkerError, - WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceBinaryRequest, - WorkspaceBinaryResponse, WorkspaceClient, WorkspaceClientError, WorkspaceId, WorkspaceIdError, - WorkspacePromptCatalogResolution, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, - WorkspaceServerOperation, WorkspaceWorkerDiscoveryRequest, apply_worker_manifest, - marker_workspace_client, unavailable_workspace_client, + WorkerFilesystemAuthority, WorkerRunResult, WorkerStartupFeatureCleanup, + WorkerWorkspaceContext, WorkspaceBinaryRequest, WorkspaceBinaryResponse, WorkspaceClient, + WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspacePromptCatalogResolution, + WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, WorkspaceServerOperation, + WorkspaceWorkerDiscoveryRequest, apply_worker_manifest, marker_workspace_client, + unavailable_workspace_client, }; diff --git a/crates/worker/src/runtime/worker_allocation.rs b/crates/worker/src/runtime/worker_allocation.rs index c7445cdaa..ab6bc3455 100644 --- a/crates/worker/src/runtime/worker_allocation.rs +++ b/crates/worker/src/runtime/worker_allocation.rs @@ -21,8 +21,9 @@ pub use error::ScopeLockError; #[cfg(test)] pub(crate) use lifecycle::install_top_level_at_for_test; pub use lifecycle::{ - ScopeAllocationGuard, SegmentActivationCommitError, SegmentActivationGuard, SegmentLockInfo, - adopt_allocation, install_top_level, install_top_level_with_deny, lookup_segment, + ScopeAllocationCleanup, ScopeAllocationGuard, SegmentActivationCommitError, + SegmentActivationGuard, SegmentLockInfo, adopt_allocation, install_top_level, + install_top_level_with_deny, lookup_segment, }; pub use mutate::{ delegate_scope, reclaim_delegated_scope, reclaim_stale, reclaim_stale_with, register_worker, diff --git a/crates/worker/src/runtime/worker_allocation/error.rs b/crates/worker/src/runtime/worker_allocation/error.rs index 601202693..334e0125f 100644 --- a/crates/worker/src/runtime/worker_allocation/error.rs +++ b/crates/worker/src/runtime/worker_allocation/error.rs @@ -11,6 +11,14 @@ use session_store::SegmentId; pub enum ScopeLockError { #[error("I/O error on workers.json: {0}")] Io(#[from] io::Error), + #[error("worker `{worker_name}` allocation still has a live guard")] + AllocationGuardAttached { worker_name: String }, + #[error("allocation acquisition was published but its save failed: {source}")] + AcquisitionSaveFailed { + #[source] + source: io::Error, + cleanup: super::lifecycle::ScopeAllocationCleanup, + }, #[error("worker `{0}` is already allocated")] DuplicateWorkerName(String), #[error("requested scope `{}` conflicts with worker allocation `{competitor}` rule `{}`", .rule.target.display(), .competitor_rule.target.display())] @@ -46,3 +54,15 @@ pub enum ScopeLockError { actual: Option, }, } + +impl ScopeLockError { + /// Retain the release authority for an owned acquisition whose table rename + /// succeeded before save returned an error. Refusals and unpublished saves + /// do not confer ownership of any preexisting allocation. + pub fn allocation_cleanup_handle(&self) -> Option { + match self { + Self::AcquisitionSaveFailed { cleanup, .. } => Some(cleanup.clone()), + _ => None, + } + } +} diff --git a/crates/worker/src/runtime/worker_allocation/lifecycle.rs b/crates/worker/src/runtime/worker_allocation/lifecycle.rs index cab09ec18..ba6505008 100644 --- a/crates/worker/src/runtime/worker_allocation/lifecycle.rs +++ b/crates/worker/src/runtime/worker_allocation/lifecycle.rs @@ -3,6 +3,7 @@ //! up on drop. use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; use manifest::ScopeRule; use session_store::SegmentId; @@ -11,13 +12,82 @@ use super::error::ScopeLockError; use super::mutate::release_worker; use super::table::{LockFileGuard, default_allocation_path}; -/// Owned allocation: on drop, opens the lock file and releases this -/// Worker's entry. The guard keeps only the name + lock-file path; it -/// does not hold the `flock` for the Worker's lifetime. +/// Owned allocation shared with any retained cleanup handles. Only the final +/// owner's drop attempts best-effort release; no `flock` is held for its lifetime. #[derive(Debug)] pub struct ScopeAllocationGuard { + owner: Arc, +} + +/// Retains an allocation's release authority even after its Worker is consumed +/// or dropped. A successful controller transfer should drop this handle without +/// explicitly releasing the live Worker's allocation. +#[derive(Debug, Clone)] +pub struct ScopeAllocationCleanup { + owner: Arc, +} + +#[derive(Debug)] +struct ScopeAllocationOwner { worker_name: String, lock_path: PathBuf, + state: Mutex, +} + +#[derive(Debug)] +struct ScopeAllocationOwnerState { + guard_attached: bool, + released: bool, +} + +impl ScopeAllocationOwner { + fn release(&self) -> Result<(), ScopeLockError> { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.guard_attached { + return Err(ScopeLockError::AllocationGuardAttached { + worker_name: self.worker_name.clone(), + }); + } + if state.released { + return Ok(()); + } + let mut guard = LockFileGuard::open(&self.lock_path)?; + if guard.data().find(&self.worker_name).is_some() { + release_worker(&mut guard, &self.worker_name)?; + } else { + // A previous save may have removed the entry but failed after rename. + // Complete the durable save before reporting release as confirmed. + guard.save()?; + } + state.released = true; + Ok(()) + } +} + +impl Drop for ScopeAllocationOwner { + fn drop(&mut self) { + let _ = self.release(); + } +} + +impl Drop for ScopeAllocationGuard { + fn drop(&mut self) { + self.owner + .state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .guard_attached = false; + } +} + +impl ScopeAllocationCleanup { + /// Release only after the factory has returned an error and no live + /// Controller can still use this allocation. An attached guard rejects + /// release. Success is idempotent across all clones; failure retains the + /// authority for an explicit retry. + pub fn release(&self) -> Result<(), ScopeLockError> { + self.owner.release() + } } /// Exclusive authority to activate a replacement Segment for one allocated Worker. @@ -56,12 +126,58 @@ impl SegmentActivationCommitError { } impl ScopeAllocationGuard { + fn new(worker_name: String, lock_path: PathBuf) -> Self { + Self { + owner: Arc::new(ScopeAllocationOwner { + worker_name, + lock_path, + state: Mutex::new(ScopeAllocationOwnerState { + guard_attached: true, + released: false, + }), + }), + } + } + + fn commit_acquisition(self, table: &mut LockFileGuard) -> Result { + let mut published = false; + match table.save_with_publication(&mut published) { + Ok(()) => Ok(self), + Err(source) => { + let cleanup = if published { + Some(self.cleanup_handle()) + } else { + // A staged entry is not ours until rename. Disarm before + // detaching so Drop cannot remove a preexisting allocation + // (or reopen the table lock still held by this caller). + self.owner + .state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .released = true; + None + }; + drop(self); + match cleanup { + Some(cleanup) => Err(ScopeLockError::AcquisitionSaveFailed { source, cleanup }), + None => Err(ScopeLockError::Io(source)), + } + } + } + } + pub fn worker_name(&self) -> &str { - &self.worker_name + &self.owner.worker_name } pub fn lock_path(&self) -> &Path { - &self.lock_path + &self.owner.lock_path + } + + pub fn cleanup_handle(&self) -> ScopeAllocationCleanup { + ScopeAllocationCleanup { + owner: Arc::clone(&self.owner), + } } /// Lock the machine-wide table and validate a Segment replacement before @@ -71,24 +187,25 @@ impl ScopeAllocationGuard { expected: SegmentId, replacement: SegmentId, ) -> Result { - let guard = LockFileGuard::open(&self.lock_path)?; + let guard = LockFileGuard::open(self.lock_path())?; let actual = guard .data() .allocations .iter() - .find(|allocation| allocation.worker_name == self.worker_name) + .find(|allocation| allocation.worker_name == self.worker_name()) .and_then(|allocation| allocation.segment_id); if actual != Some(expected) { return Err(ScopeLockError::SegmentChanged { - worker_name: self.worker_name.clone(), + worker_name: self.worker_name().to_owned(), expected, actual, }); } if let Some(existing) = guard.data().allocations.iter().find(|allocation| { - allocation.segment_id == Some(replacement) && allocation.worker_name != self.worker_name + allocation.segment_id == Some(replacement) + && allocation.worker_name != self.worker_name() }) { return Err(ScopeLockError::SegmentConflict { segment_id: replacement, @@ -99,7 +216,7 @@ impl ScopeAllocationGuard { Ok(SegmentActivationGuard { guard, - worker_name: self.worker_name.clone(), + worker_name: self.worker_name().to_owned(), expected, replacement, }) @@ -145,14 +262,6 @@ impl SegmentActivationGuard { } } -impl Drop for ScopeAllocationGuard { - fn drop(&mut self) { - if let Ok(mut guard) = LockFileGuard::open(&self.lock_path) { - let _ = release_worker(&mut guard, &self.worker_name); - } - } -} - /// Open the default lock file, register a top-level Worker, and return a /// guard that will release the allocation on drop. pub fn install_top_level( @@ -185,7 +294,7 @@ pub fn install_top_level_with_deny( ) -> Result { let lock_path = default_allocation_path()?; let mut guard = LockFileGuard::open(&lock_path)?; - super::mutate::register_worker_with_deny( + super::mutate::prepare_worker_with_deny( &mut guard, worker_name.clone(), pid, @@ -194,10 +303,7 @@ pub fn install_top_level_with_deny( scope_deny, segment_id, )?; - Ok(ScopeAllocationGuard { - worker_name, - lock_path, - }) + ScopeAllocationGuard::new(worker_name, lock_path).commit_acquisition(&mut guard) } #[cfg(test)] @@ -210,7 +316,7 @@ pub(crate) fn install_top_level_at_for_test( segment_id: SegmentId, ) -> Result { let mut guard = LockFileGuard::open(&lock_path)?; - super::mutate::register_worker_with_deny( + super::mutate::prepare_worker_with_deny( &mut guard, worker_name.clone(), pid, @@ -219,10 +325,7 @@ pub(crate) fn install_top_level_at_for_test( Vec::new(), segment_id, )?; - Ok(ScopeAllocationGuard { - worker_name, - lock_path, - }) + ScopeAllocationGuard::new(worker_name, lock_path).commit_acquisition(&mut guard) } /// Take ownership of an existing allocation that was pre-registered by @@ -240,6 +343,15 @@ pub fn adopt_allocation( segment_id: SegmentId, ) -> Result { let lock_path = default_allocation_path()?; + adopt_allocation_at(lock_path, worker_name, new_pid, segment_id) +} + +fn adopt_allocation_at( + lock_path: PathBuf, + worker_name: String, + new_pid: u32, + segment_id: SegmentId, +) -> Result { let mut guard = LockFileGuard::open(&lock_path)?; let alloc = guard .data_mut() @@ -247,11 +359,7 @@ pub fn adopt_allocation( .ok_or_else(|| ScopeLockError::UnknownWorker(worker_name.clone()))?; alloc.pid = new_pid; alloc.segment_id = Some(segment_id); - guard.save()?; - Ok(ScopeAllocationGuard { - worker_name, - lock_path, - }) + ScopeAllocationGuard::new(worker_name, lock_path).commit_acquisition(&mut guard) } /// Information about a Worker that currently holds an allocation for a @@ -329,6 +437,388 @@ mod tests { } } + #[test] + fn retained_cleanup_handle_defers_release_until_final_owner_drop() { + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + let guard = install_top_level_at_for_test( + lock_path.clone(), + "retained".into(), + std::process::id(), + sock("retained"), + Vec::new(), + sid(), + ) + .unwrap(); + let cleanup = guard.cleanup_handle(); + let retained = cleanup.clone(); + + drop(guard); + drop(cleanup); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("retained").is_some()); + } + drop(retained); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("retained").is_none()); + } + + #[test] + fn dropping_cleanup_handle_after_transfer_preserves_live_guard() { + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + let segment = sid(); + let guard = install_top_level_at_for_test( + lock_path.clone(), + "transferred".into(), + std::process::id(), + sock("transferred"), + Vec::new(), + segment, + ) + .unwrap(); + + drop(guard.cleanup_handle()); + // A successful controller transfer leaves the guard usable, including + // the existing activation protocol, without retaining a cleanup handle. + let replacement = sid(); + guard + .begin_segment_activation(segment, replacement) + .unwrap() + .commit() + .unwrap(); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!( + table.data().find("transferred").unwrap().segment_id, + Some(replacement) + ); + } + drop(guard); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("transferred").is_none()); + } + + #[test] + fn cleanup_rejects_live_guard_then_releases_only_detached_allocation() { + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + let guard = install_top_level_at_for_test( + lock_path.clone(), + "live".into(), + std::process::id(), + sock("live"), + Vec::new(), + sid(), + ) + .unwrap(); + let _competitor = install_top_level_at_for_test( + lock_path.clone(), + "competitor".into(), + std::process::id(), + sock("competitor"), + Vec::new(), + sid(), + ) + .unwrap(); + let cleanup = guard.cleanup_handle(); + assert!(matches!( + cleanup.release(), + Err(ScopeLockError::AllocationGuardAttached { worker_name }) if worker_name == "live" + )); + drop(cleanup.clone()); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("live").is_some()); + assert!(table.data().find("competitor").is_some()); + } + + drop(guard); + cleanup.release().unwrap(); + cleanup.clone().release().unwrap(); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("live").is_none()); + assert!(table.data().find("competitor").is_some()); + } + + #[test] + fn published_install_failure_retains_cleanup_across_error_drop_and_retry() { + use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test}; + + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + let _competitor = install_top_level_at_for_test( + lock_path.clone(), + "competitor".into(), + std::process::id(), + sock("competitor"), + Vec::new(), + sid(), + ) + .unwrap(); + let segment = sid(); + fail_next_save_at_for_test(&lock_path, SaveFaultPoint::AfterRename); + let error = install_top_level_at_for_test( + lock_path.clone(), + "failed".into(), + std::process::id(), + sock("failed"), + Vec::new(), + segment, + ) + .unwrap_err(); + assert!(matches!( + &error, + ScopeLockError::AcquisitionSaveFailed { .. } + )); + let cleanup = error.allocation_cleanup_handle().unwrap(); + let retry = error.allocation_cleanup_handle().unwrap(); + drop(error); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!( + table.data().find("failed").unwrap().segment_id, + Some(segment) + ); + assert!(table.data().find("competitor").is_some()); + } + + fail_next_save_at_for_test(&lock_path, SaveFaultPoint::BeforeRename); + assert!(matches!(cleanup.release(), Err(ScopeLockError::Io(_)))); + drop(cleanup); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("failed").is_some()); + } + retry.release().unwrap(); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("failed").is_none()); + assert!(table.data().find("competitor").is_some()); + } + + #[test] + fn refused_install_has_no_cleanup_authority_over_existing_allocation() { + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + let segment = sid(); + let _competitor = install_top_level_at_for_test( + lock_path.clone(), + "competitor".into(), + std::process::id(), + sock("competitor"), + vec![write_rule("/shared", true)], + segment, + ) + .unwrap(); + let cases = [ + ("competitor", sid(), Vec::new()), + ("other", segment, Vec::new()), + ("other", sid(), vec![write_rule("/shared", true)]), + ]; + for (name, requested_segment, scope) in cases { + let error = install_top_level_at_for_test( + lock_path.clone(), + name.into(), + std::process::id(), + sock("refused"), + scope, + requested_segment, + ) + .unwrap_err(); + assert!(matches!( + &error, + ScopeLockError::DuplicateWorkerName(_) + | ScopeLockError::SegmentConflict { .. } + | ScopeLockError::WriteConflict { .. } + )); + assert!(error.allocation_cleanup_handle().is_none()); + drop(error); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!(table.data().allocations.len(), 1); + assert_eq!( + table.data().find("competitor").unwrap().socket, + sock("competitor") + ); + assert_eq!( + table.data().find("competitor").unwrap().segment_id, + Some(segment) + ); + } + } + + #[test] + fn unpublished_install_failure_has_no_cleanup_authority() { + use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test}; + + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + let _competitor = install_top_level_at_for_test( + lock_path.clone(), + "competitor".into(), + std::process::id(), + sock("competitor"), + Vec::new(), + sid(), + ) + .unwrap(); + fail_next_save_at_for_test(&lock_path, SaveFaultPoint::BeforeRename); + let error = install_top_level_at_for_test( + lock_path.clone(), + "failed".into(), + std::process::id(), + sock("failed"), + Vec::new(), + sid(), + ) + .unwrap_err(); + assert!(matches!(&error, ScopeLockError::Io(_))); + assert!(error.allocation_cleanup_handle().is_none()); + drop(error); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("failed").is_none()); + assert!(table.data().find("competitor").is_some()); + } + + #[test] + fn published_adoption_failure_retains_cleanup_for_retry() { + use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test}; + + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + { + let mut table = LockFileGuard::open(&lock_path).unwrap(); + delegate_placeholder(&mut table, "child", std::process::id()); + delegate_placeholder(&mut table, "competitor", std::process::id()); + } + let segment = sid(); + fail_next_save_at_for_test(&lock_path, SaveFaultPoint::AfterRename); + let error = adopt_allocation_at( + lock_path.clone(), + "child".into(), + std::process::id(), + segment, + ) + .unwrap_err(); + let cleanup = error.allocation_cleanup_handle().unwrap(); + drop(error); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!( + table.data().find("child").unwrap().segment_id, + Some(segment) + ); + } + fail_next_save_at_for_test(&lock_path, SaveFaultPoint::AfterRename); + assert!(matches!(cleanup.release(), Err(ScopeLockError::Io(_)))); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("child").is_none()); + assert!(table.data().find("competitor").is_some()); + } + cleanup.release().unwrap(); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("child").is_none()); + assert!(table.data().find("competitor").is_some()); + } + + #[test] + fn unpublished_adoption_failure_preserves_placeholder_without_cleanup_authority() { + use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test}; + + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + let placeholder_pid = std::process::id(); + { + let mut table = LockFileGuard::open(&lock_path).unwrap(); + delegate_placeholder(&mut table, "child", placeholder_pid); + } + fail_next_save_at_for_test(&lock_path, SaveFaultPoint::BeforeRename); + let error = adopt_allocation_at( + lock_path.clone(), + "child".into(), + placeholder_pid.wrapping_add(1), + sid(), + ) + .unwrap_err(); + assert!(matches!(&error, ScopeLockError::Io(_))); + assert!(error.allocation_cleanup_handle().is_none()); + drop(error); + let table = LockFileGuard::open(&lock_path).unwrap(); + let allocation = table.data().find("child").unwrap(); + assert_eq!(allocation.pid, placeholder_pid); + assert_eq!(allocation.segment_id, None); + } + + #[test] + fn failed_cleanup_release_retains_authority_for_retry() { + use super::super::table::{SaveFaultPoint, fail_next_save_at_for_test}; + + for point in [SaveFaultPoint::BeforeRename, SaveFaultPoint::AfterRename] { + let dir = TempDir::new().unwrap(); + let lock_path = dir.path().join("workers.json"); + let guard = install_top_level_at_for_test( + lock_path.clone(), + "failed".into(), + std::process::id(), + sock("failed"), + vec![read_rule("/shared", true)], + sid(), + ) + .unwrap(); + let _other = install_top_level_at_for_test( + lock_path.clone(), + "other".into(), + std::process::id(), + sock("other"), + vec![read_rule("/shared", true)], + sid(), + ) + .unwrap(); + let cleanup = guard.cleanup_handle(); + let retry = cleanup.clone(); + drop(guard); + fail_next_save_at_for_test(&lock_path, point); + + assert!(matches!(cleanup.release(), Err(ScopeLockError::Io(_)))); + drop(cleanup); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!( + table.data().find("failed").is_some(), + point == SaveFaultPoint::BeforeRename + ); + assert!(table.data().find("other").is_some()); + } + retry.release().unwrap(); + { + let table = LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("failed").is_none()); + assert!(table.data().find("other").is_some()); + } + + // Success disarms every clone and final Drop, even if the name is + // subsequently allocated to a different Worker. + let _replacement = install_top_level_at_for_test( + lock_path.clone(), + "failed".into(), + std::process::id(), + sock("replacement"), + vec![read_rule("/shared", true)], + sid(), + ) + .unwrap(); + retry.clone().release().unwrap(); + drop(retry); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!( + table.data().find("failed").unwrap().socket, + sock("replacement") + ); + assert!(table.data().find("other").is_some()); + } + } + #[test] fn adopt_allocation_rewrites_pid_and_releases_on_drop() { let dir = TempDir::new().unwrap(); @@ -358,6 +848,7 @@ mod tests { let dir = TempDir::new().unwrap(); let _sandbox = RuntimeDirSandbox::new(dir.path()); let err = adopt_allocation("ghost".into(), 42, sid()).unwrap_err(); + assert!(err.allocation_cleanup_handle().is_none()); assert!(matches!(err, ScopeLockError::UnknownWorker(ref n) if n == "ghost")); } diff --git a/crates/worker/src/runtime/worker_allocation/mutate.rs b/crates/worker/src/runtime/worker_allocation/mutate.rs index 846adca98..9ba22f057 100644 --- a/crates/worker/src/runtime/worker_allocation/mutate.rs +++ b/crates/worker/src/runtime/worker_allocation/mutate.rs @@ -57,6 +57,30 @@ pub fn register_worker_with_deny( scope_allow: Vec, scope_deny: Vec, segment_id: SegmentId, +) -> Result<(), ScopeLockError> { + prepare_worker_with_deny( + guard, + worker_name, + pid, + socket, + scope_allow, + scope_deny, + segment_id, + )?; + guard.save()?; + Ok(()) +} + +/// Admit and stage a registration without publishing it. Owned installation +/// constructs its cleanup authority after admission and before committing. +pub(super) fn prepare_worker_with_deny( + guard: &mut LockFileGuard, + worker_name: String, + pid: u32, + socket: PathBuf, + scope_allow: Vec, + scope_deny: Vec, + segment_id: SegmentId, ) -> Result<(), ScopeLockError> { reclaim_stale(guard); if guard.data().find(&worker_name).is_some() { @@ -101,7 +125,6 @@ pub fn register_worker_with_deny( delegated_from: None, segment_id: Some(segment_id), }); - guard.save()?; Ok(()) } diff --git a/crates/worker/src/runtime/worker_allocation/table.rs b/crates/worker/src/runtime/worker_allocation/table.rs index b9d0a16ae..a8420551d 100644 --- a/crates/worker/src/runtime/worker_allocation/table.rs +++ b/crates/worker/src/runtime/worker_allocation/table.rs @@ -245,6 +245,14 @@ impl LockFileGuard { /// Persist with atomic replacement while the separate allocation lock stays /// held. A crash exposes either the old complete table or the new one. pub fn save(&mut self) -> io::Result<()> { + let mut published = false; + self.save_with_publication(&mut published) + } + + /// Report this save's rename while the caller still holds the table lock. + /// In-memory mutation alone is not evidence that an acquisition took place. + pub(super) fn save_with_publication(&mut self, published: &mut bool) -> io::Result<()> { + *published = false; let json = serde_json::to_vec_pretty(&self.data).map_err(io::Error::other)?; let mut temp_os = self.data_path.as_os_str().to_owned(); temp_os.push(".tmp"); @@ -264,6 +272,7 @@ impl LockFileGuard { )); } fs::rename(&temp_path, &self.data_path)?; + *published = true; #[cfg(test)] if take_save_fault(&self.data_path, SaveFaultPoint::AfterRename) { return Err(io::Error::other( diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 6af8d337b..92a7e7047 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -533,7 +533,7 @@ use crate::prompt::source::PromptCatalogSource; use crate::prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate}; use crate::runtime::dir; use crate::runtime::worker_allocation::{ - self, ScopeAllocationGuard, ScopeLockError, SegmentActivationGuard, + self, ScopeAllocationCleanup, ScopeAllocationGuard, ScopeLockError, SegmentActivationGuard, }; use crate::skill::{SkillActivationResponse, SkillClientError}; #[cfg(test)] @@ -2951,6 +2951,69 @@ struct PendingCompactionCleanup { compaction_id: String, } +/// Retains startup task scopes after a consuming controller startup fails. +/// +/// Call `shutdown` only after startup has failed with no live Controller. Keep +/// this handle on errors or timeout and retry; it owns any in-flight settlement +/// future, including join handles drained from the registries. On successful +/// controller transfer, drop the factory's handle without shutting it down. +#[derive(Clone)] +pub struct WorkerStartupFeatureCleanup { + inner: Arc, +} + +struct StartupFeatureCleanupInner { + registries: Mutex>, + settlement: tokio::sync::Mutex< + Option>>, + >, +} + +impl WorkerStartupFeatureCleanup { + fn new(registry: FeatureBackgroundTaskRegistry) -> Self { + Self { + inner: Arc::new(StartupFeatureCleanupInner { + registries: Mutex::new(vec![registry]), + settlement: tokio::sync::Mutex::new(None), + }), + } + } + + fn retain(&self, registry: FeatureBackgroundTaskRegistry) { + self.inner + .registries + .lock() + .expect("startup Feature cleanup poisoned") + .push(registry); + } + + /// Stops and joins every captured registry. Cancelling this call preserves + /// its settlement future for the next call, rather than dropping task joins. + pub async fn shutdown(&self) -> Result<(), crate::hook::HookError> { + let mut settlement = self.inner.settlement.lock().await; + let pending = settlement.get_or_insert_with(|| { + let registries = self + .inner + .registries + .lock() + .expect("startup Feature cleanup poisoned") + .clone(); + Box::pin(async move { + let mut failure = None; + for registry in registries { + if let Err(error) = registry.shutdown().await { + failure.get_or_insert(error); + } + } + failure.map_or(Ok(()), Err) + }) + }); + let result = pending.as_mut().await; + *settlement = None; + result + } +} + /// An independent agent execution unit. /// /// Holds a [`Engine`] directly and persists session state via @@ -2997,6 +3060,8 @@ pub struct Worker { hook_registry: Option>, /// Executable background tasks registered by successfully installed features. feature_background_tasks: FeatureBackgroundTaskRegistry, + /// Present only while an opted-in factory is capturing startup task scopes. + startup_feature_cleanup: Option, /// Enabled chat invocation metadata and executable handlers. feature_invocations: FeatureInvocationRegistry, /// Internal Workers install an explicit Feature composition and disable @@ -3087,9 +3152,8 @@ pub struct Worker { /// Workers built via `from_manifest` / `from_manifest_spawned` / /// `restore_from_manifest` (production paths); `None` for the /// low-level `Worker::new` constructor used in tests, which bypasses - /// the registry. Kept purely for its `Drop` impl, which releases - /// the allocation when the Worker is dropped. - #[allow(dead_code)] + /// the registry. Cleanup handles can retain the allocation beyond Worker + /// drop; only the final shared owner attempts best-effort release. scope_allocation: Option, /// Socket path of the spawning Worker. `Some` only for Workers built via /// `from_manifest_spawned`. Consumed by the controller to fire @@ -3155,7 +3219,7 @@ pub struct Worker { impl Drop for Worker { fn drop(&mut self) { // A failed activation authority nests the same machine-wide lock used by - // `ScopeAllocationGuard::drop`. Release the nested guard first so Worker + // allocation owner's final Drop. Release the nested guard first so Worker // teardown cannot self-deadlock while removing its outer allocation. let authority = match self.failed_activation_authority.get_mut() { Ok(authority) => authority.take(), @@ -3165,6 +3229,47 @@ impl Drop for Worker { } } +impl Worker { + /// Opt in before a consuming startup call. The handle captures the current + /// task registry and each subsequently installed registry until successful + /// controller transfer ends capture. Repeated calls share the same capture. + pub fn retain_startup_feature_cleanup(&mut self) -> WorkerStartupFeatureCleanup { + self.startup_feature_cleanup + .get_or_insert_with(|| { + WorkerStartupFeatureCleanup::new(self.feature_background_tasks.clone()) + }) + .clone() + } + + /// Observes explicit factory cleanup ownership without enabling capture. + pub fn startup_feature_cleanup_handle(&self) -> Option { + self.startup_feature_cleanup.clone() + } + + /// End factory capture immediately before transferring to a live Controller. + /// The factory must also drop its handle after startup returns success. + pub fn finish_startup_feature_cleanup_capture(&mut self) { + self.startup_feature_cleanup = None; + } + + /// Retain cleanup authority before passing the reconstructed Worker into a + /// consuming factory. Explicit release is only valid after factory failure + /// with no live Controller; drop the handle on successful controller transfer. + pub fn restore_allocation_cleanup_handle(&self) -> Option { + self.scope_allocation + .as_ref() + .map(ScopeAllocationGuard::cleanup_handle) + } + + fn construction_failure(&self, error: WorkerError) -> WorkerError { + WorkerError::ConstructionCleanup { + source: Box::new(error), + allocation: self.restore_allocation_cleanup_handle(), + sessions: self.workdir_sessions.clone(), + } + } +} + impl Worker { pub async fn stop_feature_runtime(&mut self, reason: impl Into) { if let Some(hooks) = self.hook_registry.clone() { @@ -3352,6 +3457,7 @@ impl Worker { hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + startup_feature_cleanup: None, feature_invocations: FeatureInvocationRegistry::default(), manifest_lifecycle_features_enabled: true, interceptor_installed: false, @@ -3850,6 +3956,9 @@ impl Worker { return report; } self.feature_background_tasks = report.background_tasks.clone(); + if let Some(cleanup) = &self.startup_feature_cleanup { + cleanup.retain(self.feature_background_tasks.clone()); + } self.feature_invocations = report.chat_invocations.clone(); for instruction in report.installed_instruction_contributions() { self.register_feature_instruction(instruction); @@ -7906,7 +8015,7 @@ where let workdir_sessions = workdir_sessions_from_authority(&common.filesystem_authority, &scope); - let mut worker = Self { + let worker = Self { backend_job: std::sync::OnceLock::new(), manifest, engine: Some(worker), @@ -7925,6 +8034,7 @@ where hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + startup_feature_cleanup: None, feature_invocations: FeatureInvocationRegistry::default(), manifest_lifecycle_features_enabled: true, interceptor_installed: false, @@ -7966,10 +8076,15 @@ where log_writer: None, notification_committer: None, }; - worker.apply_permissions_from_manifest(); - worker.apply_prune_from_manifest(); - worker.write_worker_metadata_pending()?; - Ok(worker) + worker.finish_pending_construction() + } + + fn finish_pending_construction(mut self) -> Result { + self.apply_permissions_from_manifest(); + self.apply_prune_from_manifest(); + self.write_worker_metadata_pending() + .map_err(|error| self.construction_failure(error))?; + Ok(self) } /// Complete the Host-owned Session-to-Subject attribution boundary before a @@ -8306,6 +8421,7 @@ where hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + startup_feature_cleanup: None, feature_invocations: FeatureInvocationRegistry::default(), manifest_lifecycle_features_enabled: false, interceptor_installed: false, @@ -8438,6 +8554,7 @@ where hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + startup_feature_cleanup: None, feature_invocations: FeatureInvocationRegistry::default(), manifest_lifecycle_features_enabled: true, interceptor_installed: false, @@ -8751,6 +8868,11 @@ where if state.system_prompt.is_none() { return Err(WorkerError::SegmentSystemPromptMissing { segment_id }); } + let restored_history_entries = + restore_history_entries(session_id, segment_id, &raw_entries).map_err(|error| { + WorkerError::InvalidState(format!("restore typed Worker session history: {error}")) + })?; + let flow_runtime_state = restored_flow_runtime_state(&state.extensions)?; let mirror_entries: Vec = raw_entries.clone(); let scope_config = effective_restore_scope_config(&store, &manifest)?; @@ -8797,10 +8919,6 @@ where // (the Worker's one and only write path that prepends a summary at // history[0]). Restoring the anchor lets Anthropic re-use a // stable cache prefix for long-lived restored sessions. - let restored_history_entries = - restore_history_entries(session_id, segment_id, &raw_entries).map_err(|error| { - WorkerError::InvalidState(format!("restore typed Worker session history: {error}")) - })?; let restored_history = restored_history_entries .iter() .map(|entry| entry.item.clone()) @@ -8839,14 +8957,13 @@ where filesystem_authority: common.filesystem_authority, workdir_sessions, workspace_context: common.workspace_context, - flow_runtime_state: Arc::new(Mutex::new(restored_flow_runtime_state( - &state.extensions, - )?)), + flow_runtime_state: Arc::new(Mutex::new(flow_runtime_state)), flow_feature_enabled: false, scope, hook_builder: HookRegistryBuilder::new(), hook_registry: None, feature_background_tasks: FeatureBackgroundTaskRegistry::default(), + startup_feature_cleanup: None, feature_invocations: FeatureInvocationRegistry::default(), manifest_lifecycle_features_enabled: true, interceptor_installed: false, @@ -8896,17 +9013,24 @@ where worker .session .restore_pending_activations(&state.extensions); - worker - .pending_submission_handle() - .reconcile_uploaded_file_pins()?; - worker.apply_permissions_from_manifest(); - worker.apply_prune_from_manifest(); - worker.write_worker_metadata_active(SegmentLocation { - session_id, - segment_id, - })?; - worker.reconcile_restored_delegations().await?; - Ok(worker) + worker.finish_restored_construction().await + } + + async fn finish_restored_construction(mut self) -> Result { + self.pending_submission_handle() + .reconcile_uploaded_file_pins() + .map_err(|error| self.construction_failure(error.into()))?; + self.apply_permissions_from_manifest(); + self.apply_prune_from_manifest(); + self.write_worker_metadata_active(SegmentLocation { + session_id: self.session_id(), + segment_id: self.segment_id(), + }) + .map_err(|error| self.construction_failure(error))?; + self.reconcile_restored_delegations() + .await + .map_err(|error| self.construction_failure(error))?; + Ok(self) } async fn reconcile_restored_delegations(&mut self) -> Result<(), WorkerError> { @@ -9745,6 +9869,16 @@ fn compact_failure_category(error: &WorkerError) -> CompactFailureCategory { #[derive(Debug, thiserror::Error)] pub enum WorkerError { + /// A constructor failed after acquiring resources. Retain these handles until + /// explicit cleanup succeeds; the Worker has dropped before this is returned. + #[error("{source}")] + ConstructionCleanup { + #[source] + source: Box, + allocation: Option, + sessions: Arc, + }, + #[error("invalid durable Worker state: {0}")] InvalidState(String), @@ -12224,6 +12358,604 @@ mod build_summary_prompt_tests { assert_eq!(restored.active_run_turn_count, None); } + #[tokio::test] + async fn restore_cleanup_handle_retains_allocation_after_worker_drop() { + let dir = tempfile::tempdir().unwrap(); + let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); + let mut worker = Worker::new( + minimal_manifest(), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), + store, + WorkerWorkspaceContext::no_workspace(), + WorkerFilesystemAuthority::None, + Scope::empty(), + ) + .await + .unwrap(); + assert!(worker.restore_allocation_cleanup_handle().is_none()); + let lock_path = dir.path().join("workers.json"); + worker.scope_allocation = Some( + worker_allocation::install_top_level_at_for_test( + lock_path.clone(), + "test-worker".into(), + std::process::id(), + dir.path().join("worker.sock"), + Vec::new(), + worker.segment_id(), + ) + .unwrap(), + ); + let cleanup = worker.restore_allocation_cleanup_handle().unwrap(); + + drop(worker); + { + let table = worker_allocation::LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("test-worker").is_some()); + } + cleanup.release().unwrap(); + let table = worker_allocation::LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find("test-worker").is_none()); + } + + // Constructor finalization uses real Worker, allocation, Workdir, and storage + // implementations. Only the model and metadata-write failure boundary are substituted. + mod construction_cleanup_tests { + use super::*; + + type ConstructionStore = + session_store::CombinedStore; + + async fn constructed_worker( + temp: &tempfile::TempDir, + ) -> Worker, ConstructionStore> { + let store = session_store::CombinedStore::new( + session_store::FsStore::new(temp.path().join("sessions")).unwrap(), + session_store::FsWorkerStore::new(temp.path().join("workers")).unwrap(), + ); + let mut worker = Worker::new( + minimal_manifest(), + Engine::, Mutable, SessionHistoryMetadata>::new_annotated( + Box::new(NoopClient), + ), + store, + WorkerWorkspaceContext::no_workspace(), + WorkerFilesystemAuthority::local( + temp.path().to_path_buf(), + temp.path().to_path_buf(), + ), + Scope::writable(temp.path()).unwrap(), + ) + .await + .unwrap(); + worker.scope_allocation = Some( + worker_allocation::install_top_level_at_for_test( + temp.path().join("workers.json"), + worker.manifest.worker.name.clone(), + std::process::id(), + temp.path().join("worker.sock"), + Vec::new(), + worker.segment_id(), + ) + .unwrap(), + ); + worker + } + + #[tokio::test] + async fn post_construction_failures_retain_resources_through_failed_release_retry() { + for stage in [ + "pending_metadata", + "restore_pins", + "restore_metadata", + "restore_delegations", + ] { + let temp = tempfile::tempdir().unwrap(); + let mut worker = constructed_worker(&temp).await; + if stage != "pending_metadata" { + worker.ensure_segment_head().await.unwrap(); + } + let name = worker.manifest.worker.name.clone(); + let lock_path = temp.path().join("workers.json"); + let sessions = Arc::downgrade(&worker.workdir_sessions()); + let session = worker.workdir_sessions().resolve(None).unwrap().session; + let writes = Arc::new(AtomicUsize::new(0)); + let writer_calls = writes.clone(); + let metadata_path = temp + .path() + .join("workers") + .join(&name) + .join("metadata.json"); + worker.worker_metadata_writer = Some(Arc::new(move |_| { + writer_calls.fetch_add(1, Ordering::SeqCst); + if stage == "restore_delegations" { + // A successful publication followed by an unreadable retained + // record forces the subsequent delegation reconciliation to fail. + std::fs::create_dir_all(&metadata_path)?; + Ok(()) + } else { + Err(WorkerStoreError::Io(std::io::Error::other( + "metadata write refused", + ))) + } + })); + if stage == "restore_pins" { + let pins_dir = temp + .path() + .join("sessions") + .join(worker.session_id().to_string()) + .join("artifacts") + .join("paste"); + std::fs::create_dir_all(pins_dir.join(".aggregate.lock")).unwrap(); + } + let result = if stage == "pending_metadata" { + worker.finish_pending_construction() + } else { + worker.finish_restored_construction().await + }; + let error = match result { + Err(error) => error, + Ok(_) => panic!("{stage} must fail after Worker construction"), + }; + let WorkerError::ConstructionCleanup { + source, + allocation, + sessions: retained, + } = &error + else { + panic!("{stage} lost construction cleanup authority: {error:?}"); + }; + if stage == "restore_pins" { + assert!(matches!( + source.as_ref(), + WorkerError::Store(StoreError::Io(_)) + )); + assert_eq!(writes.load(Ordering::SeqCst), 0); + } else { + assert!(matches!( + source.as_ref(), + WorkerError::WorkerStore(WorkerStoreError::Io(_)) + )); + assert_eq!(writes.load(Ordering::SeqCst), 1); + } + assert_eq!(error.to_string(), source.to_string()); + assert_eq!( + std::error::Error::source(&error).unwrap().to_string(), + source.to_string() + ); + assert!(Arc::ptr_eq(&sessions.upgrade().unwrap(), retained)); + let allocation = allocation.clone().expect("constructed Worker's allocation"); + let retained = retained.clone(); + drop(error); + { + let table = worker_allocation::LockFileGuard::open(&lock_path).unwrap(); + assert!( + table.data().find(&name).is_some(), + "{stage} released during Worker Drop" + ); + } + assert!( + session + .stat(workdir::StatRequest { + path: workdir::WorkdirPath::root() + }) + .await + .is_ok() + ); + retained.close_all().await.unwrap(); + assert!(retained.is_empty()); + assert!( + session + .stat(workdir::StatRequest { + path: workdir::WorkdirPath::root() + }) + .await + .is_err() + ); + + worker_allocation::fail_next_save_for_test(&lock_path); + assert!(matches!(allocation.release(), Err(ScopeLockError::Io(_)))); + { + let table = worker_allocation::LockFileGuard::open(&lock_path).unwrap(); + assert!( + table.data().find(&name).is_some(), + "{stage} lost failed release" + ); + } + allocation.release().unwrap(); + allocation.clone().release().unwrap(); + let table = worker_allocation::LockFileGuard::open(&lock_path).unwrap(); + assert!(table.data().find(&name).is_none()); + } + } + + #[tokio::test] + async fn malformed_restore_state_is_rejected_before_filesystem_preparation() { + for corrupt_flow in [false, true] { + let temp = tempfile::tempdir().unwrap(); + let store = session_store::CombinedStore::new( + session_store::FsStore::new(temp.path().join("sessions")).unwrap(), + session_store::FsWorkerStore::new(temp.path().join("workers")).unwrap(), + ); + let session_id = session_store::new_session_id(); + let segment_id = session_store::new_segment_id(); + let mut entries = vec![LogEntry::AnnotatedSegmentStart { + ts: 0, + session_id, + system_prompt: Some("saved prompt".into()), + config: RequestConfig::default(), + history: Vec::new(), + forked_from: None, + compacted_from: None, + }]; + if corrupt_flow { + entries.push(LogEntry::Extension { + ts: 1, + domain: FLOW_RUNTIME_EXTENSION_DOMAIN.into(), + payload: serde_json::Value::Null, + }); + } else { + entries.push(LogEntry::ToolResultCorrected { + ts: 1, + entry: to_logged_history_entry(&history_entry( + Item::user_message("correction without a target"), + WorkerHistoryProvenance::LegacyUnknown, + )), + }); + } + store + .create_segment(session_id, segment_id, &entries) + .unwrap(); + let result = Worker::, _>::restore_from_manifest_with_context_and_model_client( + session_id, + segment_id, + minimal_manifest(), + store, + PromptCatalogSource::builtins_only(), + WorkerWorkspaceContext::no_workspace(), + // If preparation runs first, it must reject this missing root, + // without touching the process-global allocation table. + WorkerFilesystemAuthority::local(temp.path().join("missing"), temp.path().join("missing")), + Some(Box::new(NoopClient)), + ).await; + let error = match result { + Err(error) => error, + Ok(_) => panic!("corrupt restore state must be rejected"), + }; + let WorkerError::InvalidState(message) = error else { + panic!("restore parsed state too late: {error:?}"); + }; + assert!(message.contains(if corrupt_flow { + "Flow runtime state" + } else { + "typed Worker session history" + })); + } + } + } + + // Real Feature installation, task registries, Workdir provider/router, and + // PreparedWorker startup; only the model and task completion barrier are test implementations. + mod startup_feature_cleanup_tests { + use super::*; + use crate::feature::background::{ + BackgroundTaskCancellation, BackgroundTaskContext, BackgroundTaskSpec, + FeatureBackgroundTask, + }; + use crate::feature::{ + BackgroundTaskDeclaration, FeatureDescriptor, FeatureId, FeatureInstallContext, + FeatureInstallError, FeatureModule, + }; + use crate::hook::{HookError, HookErrorCategory}; + use tokio::sync::Notify; + + type StartupStore = + session_store::CombinedStore; + + #[derive(Default)] + struct TaskSignals { + entered: Notify, + cancelled: AtomicBool, + release: Notify, + exited: AtomicBool, + session: Mutex>, + } + + #[derive(Clone)] + struct BarrierTask { + signals: Arc, + _lifetime: Arc<()>, + panic_on_release: bool, + } + + #[async_trait] + impl FeatureBackgroundTask for BarrierTask { + async fn run( + &self, + _context: BackgroundTaskContext, + cancellation: BackgroundTaskCancellation, + ) -> Result<(), HookError> { + self.signals.entered.notify_one(); + cancellation.cancelled().await; + self.signals.cancelled.store(true, Ordering::Release); + self.signals.release.notified().await; + assert!(!self.panic_on_release, "injected startup task join failure"); + let session = self.signals.session.lock().unwrap().clone(); + if let Some(session) = session { + assert!( + session_open(&session).await, + "Feature task must retain Workdir access until it is quiescent" + ); + } + self.signals.exited.store(true, Ordering::Release); + Ok(()) + } + } + + struct BarrierFeature(BarrierTask); + + impl FeatureModule for BarrierFeature { + fn descriptor(&self) -> FeatureDescriptor { + FeatureDescriptor::builtin("startup-barrier", "Startup barrier") + .with_background_task(BackgroundTaskDeclaration::worker_managed( + "barrier", + "Startup task completion barrier", + )) + } + + fn install( + &self, + context: &mut FeatureInstallContext<'_>, + ) -> Result<(), FeatureInstallError> { + context.background_tasks().register( + BackgroundTaskSpec::single_flight( + BackgroundTaskDeclaration::worker_managed( + "barrier", + "Startup task completion barrier", + ), + Duration::from_secs(60), + ), + self.0.clone(), + ) + } + } + + async fn worker(temp: &tempfile::TempDir) -> Worker { + let store = session_store::CombinedStore::new( + session_store::FsStore::new(temp.path().join("sessions")).unwrap(), + session_store::FsWorkerStore::new(temp.path().join("workers")).unwrap(), + ); + Worker::new( + minimal_manifest(), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), + store, + WorkerWorkspaceContext::no_workspace(), + WorkerFilesystemAuthority::None, + Scope::empty(), + ) + .await + .unwrap() + } + + fn install( + worker: &mut Worker, + panic_on_release: bool, + ) -> (Arc, std::sync::Weak<()>) { + let signals = Arc::new(TaskSignals::default()); + let lifetime = Arc::new(()); + let weak = Arc::downgrade(&lifetime); + let report = worker.install_features(FeatureRegistryBuilder::new().with_module( + BarrierFeature(BarrierTask { + signals: signals.clone(), + _lifetime: lifetime, + panic_on_release, + }), + )); + assert!(!report.has_errors()); + (signals, weak) + } + + async fn start(worker: &Worker, signals: &TaskSignals) { + worker + .feature_background_tasks + .start( + &FeatureId::builtin("startup-barrier"), + "barrier", + HookInvocationContext { + workspace_id: None, + worker_id: "startup-owner".into(), + session_id: worker.session_id().to_string(), + run_id: None, + turn_index: None, + call_id: None, + }, + ) + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), signals.entered.notified()) + .await + .unwrap(); + } + + fn bind_session( + worker: &mut Worker, + temp: &tempfile::TempDir, + ) -> workdir::WorkdirSessionHandle { + let session: workdir::WorkdirSessionHandle = + Arc::new(workdir::LocalWorkdirSession::materialized_bound( + workdir::Workdir::new("startup-task-cleanup"), + temp.path().to_path_buf(), + temp.path().to_path_buf(), + SharedScope::new(Scope::writable(temp.path()).unwrap()), + workdir::WorkdirSessionCapabilities::ALL, + )); + worker.bind_single_workdir_session(Some(session.clone())); + session + } + + async fn session_open(session: &workdir::WorkdirSessionHandle) -> bool { + session + .stat(workdir::StatRequest { + path: workdir::WorkdirPath::root(), + }) + .await + .is_ok() + } + + async fn fail_consuming_start( + worker: Worker, + temp: &tempfile::TempDir, + ) { + let factory_owned_cleanup = worker.startup_feature_cleanup_handle().is_some(); + let blocked_parent = temp.path().join("blocked-runtime"); + std::fs::write(&blocked_parent, "not a directory").unwrap(); + let result = crate::PreparedWorker::new( + worker, + crate::WorkerBootstrapLayout::RuntimeManagedRun { + run_dir: blocked_parent.join("1"), + bash_output_dir: temp.path().join("spill"), + }, + crate::WorkerControllerTransport::InProcess, + ) + .start() + .await; + let Err(crate::WorkerBootstrapError::Controller { + source, + cleanup_failed, + }) = result + else { + panic!("runtime parent file must reject controller startup"); + }; + assert!(matches!( + source.kind(), + std::io::ErrorKind::NotADirectory | std::io::ErrorKind::AlreadyExists + )); + assert_eq!(cleanup_failed, factory_owned_cleanup); + } + + #[tokio::test] + async fn consumed_worker_failure_retains_all_task_joins_across_cleanup_timeout() { + let temp = tempfile::tempdir().unwrap(); + let mut worker = worker(&temp).await; + let session = bind_session(&mut worker, &temp); + let sessions = worker.workdir_sessions(); + let (first, first_lifetime) = install(&mut worker, false); + *first.session.lock().unwrap() = Some(session.clone()); + start(&worker, &first).await; + let cleanup = worker.retain_startup_feature_cleanup(); + let retry_owner = cleanup.clone(); + let (second, second_lifetime) = install(&mut worker, false); + *second.session.lock().unwrap() = Some(session.clone()); + start(&worker, &second).await; + fail_consuming_start(worker, &temp).await; + + assert!(!sessions.aliases().is_empty()); + assert!(session_open(&session).await); + assert!(first_lifetime.upgrade().is_some()); + assert!(second_lifetime.upgrade().is_some()); + assert!( + tokio::time::timeout(Duration::from_millis(50), cleanup.shutdown()) + .await + .is_err() + ); + assert!(first.cancelled.load(Ordering::Acquire)); + drop(cleanup); + + // A fresh shutdown of the drained registry would falsely succeed + // once the second task finishes, abandoning the first task's join. + second.release.notify_one(); + assert!( + tokio::time::timeout(Duration::from_millis(50), retry_owner.shutdown()) + .await + .is_err() + ); + assert!(!first.exited.load(Ordering::Acquire)); + assert!(session_open(&session).await); + first.release.notify_one(); + tokio::time::timeout(Duration::from_secs(5), retry_owner.shutdown()) + .await + .unwrap() + .unwrap(); + assert!(first.exited.load(Ordering::Acquire)); + assert!(second.exited.load(Ordering::Acquire)); + assert!(session_open(&session).await); + sessions.close_all().await.unwrap(); + assert!(sessions.aliases().is_empty()); + assert!(!session_open(&session).await); + drop(retry_owner); + assert!(first_lifetime.upgrade().is_none()); + assert!(second_lifetime.upgrade().is_none()); + } + + #[tokio::test] + async fn startup_without_factory_cleanup_owner_still_closes_sessions_on_error() { + let temp = tempfile::tempdir().unwrap(); + let mut worker = worker(&temp).await; + let session = bind_session(&mut worker, &temp); + let sessions = worker.workdir_sessions(); + assert!(session_open(&session).await); + fail_consuming_start(worker, &temp).await; + assert!(sessions.aliases().is_empty()); + assert!(!session_open(&session).await); + } + + #[tokio::test] + async fn cleanup_error_settles_other_registries_and_keeps_owner_for_retry() { + let temp = tempfile::tempdir().unwrap(); + let mut worker = worker(&temp).await; + let (first, first_lifetime) = install(&mut worker, true); + start(&worker, &first).await; + let cleanup = worker.retain_startup_feature_cleanup(); + let (second, second_lifetime) = install(&mut worker, false); + start(&worker, &second).await; + fail_consuming_start(worker, &temp).await; + first.release.notify_one(); + second.release.notify_one(); + + let error = tokio::time::timeout(Duration::from_secs(5), cleanup.shutdown()) + .await + .unwrap() + .unwrap_err(); + assert_eq!(error.category, HookErrorCategory::Internal); + assert!(second.exited.load(Ordering::Acquire)); + assert!(first_lifetime.upgrade().is_some()); + assert!(second_lifetime.upgrade().is_some()); + tokio::time::timeout(Duration::from_secs(5), cleanup.shutdown()) + .await + .unwrap() + .unwrap(); + drop(cleanup); + assert!(first_lifetime.upgrade().is_none()); + assert!(second_lifetime.upgrade().is_none()); + } + + #[tokio::test] + async fn successful_controller_transfer_releases_old_startup_task_scopes() { + let temp = tempfile::tempdir().unwrap(); + let mut worker = worker(&temp).await; + let (_signals, old_scope) = install(&mut worker, false); + let cleanup = worker.retain_startup_feature_cleanup(); + let started = crate::PreparedWorker::new( + worker, + crate::WorkerBootstrapLayout::RuntimeManagedRun { + run_dir: temp.path().join("runs/1"), + bash_output_dir: temp.path().join("spill"), + }, + crate::WorkerControllerTransport::InProcess, + ) + .start() + .await + .unwrap(); + + assert!(old_scope.upgrade().is_some()); + drop(cleanup); + let released = old_scope.upgrade().is_none(); + stop_body_publication_controller(started).await; + assert!( + released, + "live Controller must not retain factory startup capture" + ); + } + } + #[tokio::test] async fn post_metadata_allocation_failure_fences_appends_and_restart_converges() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 1938f0b18..e4b8d7041 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -759,6 +759,211 @@ async fn controller_startup_failure_closes_bound_workdir_session() { )); } +mod socket_startup_failure { + //! Connect the real controller startup, command observer, socket bind and session rollback. + //! Only the Workdir telemetry boundary and LLM are substituted. This ordinary controller + //! integration target uses an isolated directory at the socket path to reject the real bind. + use super::*; + use std::path::PathBuf; + use tokio::sync::broadcast; + use workdir::WorkdirSession; + + #[derive(Debug)] + struct ObservedSession { + local: LocalWorkdirSession, + socket_path: PathBuf, + events: broadcast::Sender, + subscriptions: AtomicUsize, + receivers_at_close: Mutex>, + } + + #[async_trait] + impl WorkdirSession for ObservedSession { + fn workdir(&self) -> &Workdir { + self.local.workdir() + } + + fn capabilities(&self) -> WorkdirSessionCapabilities { + self.local.capabilities() + } + + async fn stat( + &self, + request: workdir::StatRequest, + ) -> Result { + self.local.stat(request).await + } + + async fn read( + &self, + request: workdir::ReadRequest, + ) -> Result { + self.local.read(request).await + } + + async fn write( + &self, + request: workdir::WriteRequest, + ) -> Result { + self.local.write(request).await + } + + async fn edit( + &self, + request: workdir::EditRequest, + ) -> Result { + self.local.edit(request).await + } + + async fn list( + &self, + request: workdir::ListRequest, + ) -> Result { + self.local.list(request).await + } + + async fn glob( + &self, + request: workdir::GlobRequest, + ) -> Result { + self.local.glob(request).await + } + + async fn grep( + &self, + request: workdir::GrepRequest, + ) -> Result { + self.local.grep(request).await + } + + async fn start_command( + &self, + request: CommandRequest, + ) -> Result { + self.local.start_command(request).await + } + + async fn command_status( + &self, + handle: workdir::CommandHandle, + ) -> Result { + self.local.command_status(handle).await + } + + async fn command_output( + &self, + request: CommandOutputRequest, + ) -> Result { + self.local.command_output(request).await + } + + async fn cancel_command(&self, handle: workdir::CommandHandle) -> Result<(), WorkdirError> { + self.local.cancel_command(handle).await + } + + fn subscribe_command_events(&self) -> Option> { + // Runtime creation and initial file writes must already have succeeded. A directory + // survives SocketServer's stale-file removal and makes the subsequent Unix bind fail. + std::fs::create_dir(&self.socket_path).unwrap(); + self.subscriptions.fetch_add(1, Ordering::SeqCst); + Some(self.events.subscribe()) + } + + async fn close(&self) -> Result<(), WorkdirError> { + // Record before any await: aborting without joining must not pass by getting a + // scheduler turn during the provider's asynchronous close. The sender stays alive + // even after close, so channel closure cannot accidentally dispose of the observer. + self.receivers_at_close + .lock() + .unwrap() + .push(self.events.receiver_count()); + self.local.close().await + } + } + + #[tokio::test(flavor = "current_thread")] + async fn socket_bind_failure_joins_command_observer_before_closing_session() { + let root = tempfile::tempdir().unwrap(); + let run_dir = root.path().join("runs/1"); + let socket_path = run_dir.join("worker.sock"); + let client = MockClient::new(simple_text_events()); + let manifest = WorkerManifest::from_toml(MANIFEST_TOML).unwrap(); + let store = CombinedStore::new( + FsStore::new(root.path().join("sessions")).unwrap(), + FsWorkerStore::new(root.path().join("workers")).unwrap(), + ); + let engine = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated( + client.clone(), + ); + let mut worker = Worker::new( + manifest, + engine, + store, + WorkerWorkspaceContext::local_filesystem(None), + WorkerFilesystemAuthority::local(root.path().into(), root.path().into()), + manifest::Scope::writable(root.path()).unwrap(), + ) + .await + .unwrap(); + let (events, unused_receiver) = broadcast::channel(8); + drop(unused_receiver); + let session = Arc::new(ObservedSession { + local: LocalWorkdirSession::materialized_bound( + Workdir::new("socket-startup-failure"), + root.path().into(), + root.path().into(), + worker.scope().clone(), + WorkdirSessionCapabilities::ALL, + ), + socket_path: socket_path.clone(), + events, + subscriptions: AtomicUsize::new(0), + receivers_at_close: Mutex::new(Vec::new()), + }); + worker.bind_single_workdir_session(Some(session.clone())); + let sessions = worker.workdir_sessions(); + + let result = tokio::time::timeout( + std::time::Duration::from_secs(5), + WorkerController::spawn_runtime_managed_run( + worker, + &run_dir, + &root.path().join("bash-output"), + ), + ) + .await + .expect("socket startup and rollback must settle"); + let error = match result { + Err(error) => error, + Ok(_) => panic!("a directory at the Unix socket path must reject controller startup"), + }; + assert_eq!(error.kind(), std::io::ErrorKind::AddrInUse); + assert_eq!(session.subscriptions.load(Ordering::SeqCst), 1); + assert_eq!(*session.receivers_at_close.lock().unwrap(), vec![0]); + assert_eq!(session.events.receiver_count(), 0); + assert!(sessions.aliases().is_empty()); + // No controller retains the Worker/Engine after the failed startup. + assert_eq!(Arc::strong_count(&client.call_count), 1); + assert_eq!(client.call_count.load(Ordering::SeqCst), 0); + assert!(socket_path.is_dir()); + assert!(tokio::net::UnixStream::connect(&socket_path).await.is_err()); + assert!(matches!( + session + .start_command(CommandRequest { + command: "printf unreachable".to_owned(), + timeout_secs: 5, + output_limit: 1024, + cwd: workdir::WorkdirPath::root(), + spill_dir: None, + tool_call_id: None, + }) + .await, + Err(WorkdirError::Unavailable(_)) + )); + } +} + async fn wait_for_status(handle: &WorkerHandle, status: WorkerStatus) { let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2); loop { diff --git a/crates/workspace-server/src/server/tests/restored_parent_control_tests.rs b/crates/workspace-server/src/server/tests/restored_parent_control_tests.rs index 0c5ca25dc..8e5c4d779 100644 --- a/crates/workspace-server/src/server/tests/restored_parent_control_tests.rs +++ b/crates/workspace-server/src/server/tests/restored_parent_control_tests.rs @@ -85,7 +85,10 @@ impl RuntimeWorkerFactory for CapturingFactory { async fn restore_controller( &self, request: worker_runtime::execution::WorkerExecutionRestoreRequest, - ) -> std::result::Result { + ) -> std::result::Result< + RuntimeWorkerController, + worker_runtime::worker_backend::RuntimeRestoreFailure, + > { let worker_ref = request.worker_ref.clone(); let controller = self.inner.restore_controller(request).await?; self.captures diff --git a/crates/workspace-server/src/server_assignment_attention_tests.rs b/crates/workspace-server/src/server_assignment_attention_tests.rs index 2f5c36ca8..da46ba67d 100644 --- a/crates/workspace-server/src/server_assignment_attention_tests.rs +++ b/crates/workspace-server/src/server_assignment_attention_tests.rs @@ -494,9 +494,18 @@ impl worker_runtime::worker_backend::RuntimeWorkerFactory for AttentionWorkerFac async fn restore_controller( &self, request: worker_runtime::execution::WorkerExecutionRestoreRequest, - ) -> std::result::Result { + ) -> std::result::Result< + worker_runtime::worker_backend::RuntimeWorkerController, + worker_runtime::worker_backend::RuntimeRestoreFailure, + > { self.controller(request.worker_ref.worker_id.to_string(), true) .await + .map_err(|_| { + worker_runtime::worker_backend::RuntimeRestoreFailure::unproven( + worker_runtime::worker_backend::RuntimeRestoreStage::Reconstruction, + "attention fixture reconstruction failed", + ) + }) } } diff --git a/crates/workspace-server/src/server_assignment_delivery_tests.rs b/crates/workspace-server/src/server_assignment_delivery_tests.rs index 5d397814a..ff5adfca2 100644 --- a/crates/workspace-server/src/server_assignment_delivery_tests.rs +++ b/crates/workspace-server/src/server_assignment_delivery_tests.rs @@ -832,9 +832,18 @@ impl worker_runtime::worker_backend::RuntimeWorkerFactory for DeliveryWorkerFact async fn restore_controller( &self, request: delivery_execution::WorkerExecutionRestoreRequest, - ) -> std::result::Result { + ) -> std::result::Result< + worker_runtime::worker_backend::RuntimeWorkerController, + worker_runtime::worker_backend::RuntimeRestoreFailure, + > { self.controller(request.worker_ref.worker_id.to_string(), true) .await + .map_err(|_| { + worker_runtime::worker_backend::RuntimeRestoreFailure::unproven( + worker_runtime::worker_backend::RuntimeRestoreStage::Reconstruction, + "delivery fixture reconstruction failed", + ) + }) } } diff --git a/crates/workspace-server/src/server_workspace_config_wip_tests.rs b/crates/workspace-server/src/server_workspace_config_wip_tests.rs index 1e989edd8..0714b0f22 100644 --- a/crates/workspace-server/src/server_workspace_config_wip_tests.rs +++ b/crates/workspace-server/src/server_workspace_config_wip_tests.rs @@ -584,7 +584,10 @@ impl worker_runtime::worker_backend::RuntimeWorkerFactory for ConfigCaptureFacto async fn restore_controller( &self, request: worker_runtime::execution::WorkerExecutionRestoreRequest, - ) -> std::result::Result { + ) -> std::result::Result< + worker_runtime::worker_backend::RuntimeWorkerController, + worker_runtime::worker_backend::RuntimeRestoreFailure, + > { self.inner.restore_controller(request).await } }