fix: address lifecycle review findings

crates/worker-runtime/src/fs_store.rs

@@ -1725,12 +1725,36 @@ mod tests {
         let Some(root) = std::env::var_os("YOI_TEST_RUNTIME_STORE_LOCK_ROOT") else {
             return;
         };
+        let root = PathBuf::from(root);
+        if std::env::var_os("YOI_TEST_RUNTIME_STORE_LOCK_EXIT_WITHOUT_DROP").is_some() {
+            let _store = FsRuntimeStore::open_or_create(root, "runtime-test").unwrap();
+            std::process::exit(0);
+        }
         assert!(matches!(
-            FsRuntimeStore::open_or_create(PathBuf::from(root), "runtime-test").unwrap_err(),
+            FsRuntimeStore::open_or_create(root, "runtime-test").unwrap_err(),
             RuntimeError::RuntimeStoreAlreadyOpen { .. }
         ));
     }
 
+    #[test]
+    fn runtime_store_owner_lock_is_released_when_process_exits_without_drop() {
+        let parent = tempfile::tempdir().unwrap();
+        let root = parent.path().join("crashed-runtime-store");
+        let child = std::process::Command::new(std::env::current_exe().unwrap())
+            .arg("fs_store::tests::runtime_store_owner_lock_child_probe")
+            .arg("--exact")
+            .env("YOI_TEST_RUNTIME_STORE_LOCK_ROOT", &root)
+            .env("YOI_TEST_RUNTIME_STORE_LOCK_EXIT_WITHOUT_DROP", "1")
+            .output()
+            .unwrap();
+        assert!(
+            child.status.success(),
+            "process-exit lock probe failed: {}",
+            String::from_utf8_lossy(&child.stderr)
+        );
+        acquire_runtime_store_owner_lock(&root).unwrap();
+    }
+
     #[test]
     fn second_runtime_store_open_conflicts_before_store_mutation() {
         let parent = tempfile::tempdir().unwrap();

crates/worker-runtime/src/runtime.rs

@@ -162,6 +162,36 @@ pub struct Runtime {
     worker_operations: Arc<Mutex<BTreeMap<WorkerId, Arc<Mutex<()>>>>>,
 }
 
+struct WorkerOperationLease {
+    worker_id: WorkerId,
+    operation_lock: Arc<Mutex<()>>,
+    registry: Arc<Mutex<BTreeMap<WorkerId, Arc<Mutex<()>>>>>,
+}
+
+impl std::ops::Deref for WorkerOperationLease {
+    type Target = Mutex<()>;
+
+    fn deref(&self) -> &Self::Target {
+        self.operation_lock.as_ref()
+    }
+}
+
+impl Drop for WorkerOperationLease {
+    fn drop(&mut self) {
+        let mut registry = match self.registry.lock() {
+            Ok(registry) => registry,
+            Err(poisoned) => poisoned.into_inner(),
+        };
+        if Arc::strong_count(&self.operation_lock) == 2
+            && registry
+                .get(&self.worker_id)
+                .is_some_and(|current| Arc::ptr_eq(current, &self.operation_lock))
+        {
+            registry.remove(&self.worker_id);
+        }
+    }
+}
+
 impl Runtime {
     /// Create a memory-backed Runtime with generated identity.
     pub fn new_memory() -> Self {
@@ -1295,8 +1325,7 @@ impl Runtime {
         scope: &RuntimeWorkspaceScope,
         worker_ref: &WorkerRef,
     ) -> Result<WorkerDetail, RuntimeError> {
-        self.ensure_worker_in_workspace(scope, worker_ref)?;
-        self.restore_worker(worker_ref)
+        self.restore_worker_with_scope(worker_ref, Some(scope))
     }
 
     /// Attach a live execution to a persisted Worker definition.
@@ -1306,10 +1335,21 @@ impl Runtime {
     /// converges on its already-installed execution rather than spawning another
     /// controller.
     pub fn restore_worker(&self, worker_ref: &WorkerRef) -> Result<WorkerDetail, RuntimeError> {
+        self.restore_worker_with_scope(worker_ref, None)
+    }
+
+    fn restore_worker_with_scope(
+        &self,
+        worker_ref: &WorkerRef,
+        scope: Option<&RuntimeWorkspaceScope>,
+    ) -> Result<WorkerDetail, RuntimeError> {
         let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
         let _operation_guard = operation_lock
             .lock()
             .map_err(|_| RuntimeError::StatePoisoned)?;
+        if let Some(scope) = scope {
+            self.ensure_worker_in_workspace(scope, worker_ref)?;
+        }
         self.restore_worker_under_lock(worker_ref, WorkerRestoreMode::Explicit)
     }
 
@@ -1934,8 +1974,7 @@ impl Runtime {
         worker_ref: &WorkerRef,
         reason: Option<String>,
     ) -> Result<WorkerLifecycleAck, RuntimeError> {
-        self.ensure_worker_in_workspace(scope, worker_ref)?;
-        self.stop_worker(worker_ref, reason)
+        self.stop_worker_with_scope(worker_ref, reason, Some(scope))
     }
 
     /// Stop a Worker. Repeated stops are idempotent. The per-Worker lifecycle
@@ -1945,11 +1984,23 @@ impl Runtime {
         &self,
         worker_ref: &WorkerRef,
         reason: Option<String>,
+    ) -> Result<WorkerLifecycleAck, RuntimeError> {
+        self.stop_worker_with_scope(worker_ref, reason, None)
+    }
+
+    fn stop_worker_with_scope(
+        &self,
+        worker_ref: &WorkerRef,
+        reason: Option<String>,
+        scope: Option<&RuntimeWorkspaceScope>,
     ) -> Result<WorkerLifecycleAck, RuntimeError> {
         let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
         let _operation_guard = operation_lock
             .lock()
             .map_err(|_| RuntimeError::StatePoisoned)?;
+        if let Some(scope) = scope {
+            self.ensure_worker_in_workspace(scope, worker_ref)?;
+        }
 
         let (backend, handle) = {
             let state = self.lock()?;
@@ -2043,19 +2094,29 @@ impl Runtime {
         scope: &RuntimeWorkspaceScope,
         worker_ref: &WorkerRef,
     ) -> Result<WorkerDeleteResult, RuntimeError> {
-        self.ensure_worker_in_workspace(scope, worker_ref)?;
-        self.delete_worker(worker_ref)
+        self.delete_worker_with_scope(worker_ref, Some(scope))
     }
 
     /// Delete a non-running Worker from Runtime state and persisted Worker storage.
     pub fn delete_worker(
         &self,
         worker_ref: &WorkerRef,
+    ) -> Result<WorkerDeleteResult, RuntimeError> {
+        self.delete_worker_with_scope(worker_ref, None)
+    }
+
+    fn delete_worker_with_scope(
+        &self,
+        worker_ref: &WorkerRef,
+        scope: Option<&RuntimeWorkspaceScope>,
     ) -> Result<WorkerDeleteResult, RuntimeError> {
         let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
         let _operation_guard = operation_lock
             .lock()
             .map_err(|_| RuntimeError::StatePoisoned)?;
+        if let Some(scope) = scope {
+            self.ensure_worker_in_workspace(scope, worker_ref)?;
+        }
         let (backend, execution_handle) = {
             let state = self.lock()?;
             state.ensure_running()?;
@@ -2540,15 +2601,25 @@ impl Runtime {
         Ok(result)
     }
 
-    fn worker_operation_lock(&self, worker_id: WorkerId) -> Result<Arc<Mutex<()>>, RuntimeError> {
-        let mut operations = self
-            .worker_operations
-            .lock()
-            .map_err(|_| RuntimeError::StatePoisoned)?;
-        Ok(operations
-            .entry(worker_id)
-            .or_insert_with(|| Arc::new(Mutex::new(())))
-            .clone())
+    fn worker_operation_lock(
+        &self,
+        worker_id: WorkerId,
+    ) -> Result<WorkerOperationLease, RuntimeError> {
+        let operation_lock = {
+            let mut operations = self
+                .worker_operations
+                .lock()
+                .map_err(|_| RuntimeError::StatePoisoned)?;
+            operations
+                .entry(worker_id)
+                .or_insert_with(|| Arc::new(Mutex::new(())))
+                .clone()
+        };
+        Ok(WorkerOperationLease {
+            worker_id,
+            operation_lock,
+            registry: self.worker_operations.clone(),
+        })
     }
 
     fn lock(&self) -> Result<MutexGuard<'_, RuntimeState>, RuntimeError> {
@@ -5789,6 +5860,43 @@ mod tests {
         );
     }
 
+    #[test]
+    fn remove_waits_for_in_flight_restore_and_cannot_delete_live_execution() {
+        let backend = Arc::new(TestExecutionBackend::default());
+        let runtime = Arc::new(
+            Runtime::with_execution_backend(RuntimeOptions::default(), backend.clone()).unwrap(),
+        );
+        runtime.store_config_bundle(test_bundle()).unwrap();
+        let created = runtime
+            .create_worker(task_request("restore remove race"))
+            .unwrap();
+        runtime.stop_worker(&created.worker_ref, None).unwrap();
+        let gate = Arc::new(RestoreGate::default());
+        *backend.restore_gate.lock().unwrap() = Some(gate.clone());
+
+        let restore_runtime = runtime.clone();
+        let restore_ref = created.worker_ref.clone();
+        let restoring = std::thread::spawn(move || restore_runtime.restore_worker(&restore_ref));
+        assert!(gate.wait_for_entered(1, std::time::Duration::from_secs(2)));
+        let remove_runtime = runtime.clone();
+        let remove_ref = created.worker_ref.clone();
+        let removing = std::thread::spawn(move || remove_runtime.delete_worker(&remove_ref));
+        std::thread::sleep(std::time::Duration::from_millis(100));
+        assert_eq!(runtime.list_workers().unwrap().len(), 1);
+
+        gate.release();
+        assert_eq!(
+            restoring.join().unwrap().unwrap().status,
+            WorkerStatus::Idle
+        );
+        assert!(matches!(
+            removing.join().unwrap().unwrap_err(),
+            RuntimeError::InvalidRequest(message) if message.contains("must be stopped")
+        ));
+        assert_eq!(runtime.list_workers().unwrap().len(), 1);
+        assert!(runtime.worker_operations.lock().unwrap().is_empty());
+    }
+
     #[test]
     fn failed_stop_cleanup_retains_execution_for_retry() {
         let (runtime, backend) = runtime_and_backend();
@@ -5912,6 +6020,58 @@ mod tests {
                 .as_ref()
                 .map(|state| state.execution_generation)
         );
+        assert!(runtime.worker_operations.lock().unwrap().is_empty());
+    }
+
+    #[test]
+    fn operation_lock_registry_reclaims_idle_worker_entries() {
+        let runtime = runtime_with_backend();
+        for _ in 0..256 {
+            let lease = runtime.worker_operation_lock(WorkerId::now_v7()).unwrap();
+            let guard = lease.lock().unwrap();
+            drop(guard);
+            drop(lease);
+        }
+        assert!(runtime.worker_operations.lock().unwrap().is_empty());
+    }
+
+    #[test]
+    fn scoped_stop_revalidates_workspace_after_waiting_for_worker_lock() {
+        let (runtime, backend) = runtime_and_backend();
+        let runtime = Arc::new(runtime);
+        runtime.store_config_bundle(test_bundle()).unwrap();
+        let workspace_scope = scope("workspace-a", "server-a");
+        let worker = runtime
+            .create_worker_scoped(
+                &workspace_scope,
+                scoped_task_request("scoped stop race", "workspace-a"),
+            )
+            .unwrap();
+        let lease = runtime
+            .worker_operation_lock(worker.worker_ref.worker_id)
+            .unwrap();
+        let guard = lease.lock().unwrap();
+        let (started_tx, started_rx) = std::sync::mpsc::channel();
+        let stopping_runtime = runtime.clone();
+        let stopping_ref = worker.worker_ref.clone();
+        let stopping_scope = workspace_scope.clone();
+        let stopping = std::thread::spawn(move || {
+            started_tx.send(()).unwrap();
+            stopping_runtime.stop_worker_scoped(&stopping_scope, &stopping_ref, None)
+        });
+        started_rx.recv().unwrap();
+        {
+            let mut state = runtime.lock().unwrap();
+            state.worker_mut(&worker.worker_ref).unwrap().workspace_id =
+                Some("workspace-b".to_string());
+        }
+        drop(guard);
+        drop(lease);
+
+        let error = stopping.join().unwrap().unwrap_err();
+        assert!(matches!(error, RuntimeError::WorkerNotFound { .. }));
+        assert_eq!(*backend.stop_count.lock().unwrap(), 0);
+        assert!(runtime.worker_operations.lock().unwrap().is_empty());
     }
 
     #[test]

crates/worker-runtime/src/worker_backend.rs

@@ -1197,7 +1197,6 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
 #[derive(Clone)]
 struct RuntimeExecutionTaskScope {
     tasks: Arc<Mutex<Vec<RuntimeExecutionTask>>>,
-    terminal_failure: Arc<Mutex<Option<String>>>,
 }
 
 struct RuntimeExecutionTask {
@@ -1214,7 +1213,6 @@ impl RuntimeExecutionTaskScope {
                 task: controller_task,
                 abort_before_join: false,
             }])),
-            terminal_failure: Arc::new(Mutex::new(None)),
         }
     }
 
@@ -1241,15 +1239,8 @@ impl RuntimeExecutionTaskScope {
     }
 
     async fn join(&self) -> Result<(), String> {
-        if let Some(message) = self
-            .terminal_failure
-            .lock()
-            .map_err(|_| "execution task failure lock is poisoned".to_string())?
-            .clone()
-        {
-            return Err(message);
-        }
-
+        let mut first_failure = None;
+        let mut pending = Vec::new();
         loop {
             let next = {
                 let mut tasks = self
@@ -1263,7 +1254,7 @@ impl RuntimeExecutionTaskScope {
                 }
             };
             let Some(mut task) = next else {
-                return Ok(());
+                break;
             };
             let name = task.name;
             if task.abort_before_join {
@@ -1273,23 +1264,29 @@ impl RuntimeExecutionTaskScope {
                 Ok(Ok(())) => {}
                 Ok(Err(error)) if task.abort_before_join && error.is_cancelled() => {}
                 Ok(Err(error)) => {
-                    let message = format!("{name} task failed while stopping Worker: {error}");
-                    if let Ok(mut failure) = self.terminal_failure.lock() {
-                        *failure = Some(message.clone());
-                    }
-                    return Err(message);
+                    first_failure.get_or_insert_with(|| {
+                        format!("{name} task failed while stopping Worker: {error}")
+                    });
                 }
                 Err(_) => {
-                    self.tasks
-                        .lock()
-                        .map_err(|_| "execution task registry lock is poisoned".to_string())?
-                        .insert(0, task);
-                    return Err(format!(
-                        "{name} task did not stop before timeout; stop remains retryable"
-                    ));
+                    first_failure.get_or_insert_with(|| {
+                        format!("{name} task did not stop before timeout; stop remains retryable")
+                    });
+                    pending.push(task);
                 }
             }
         }
+        if !pending.is_empty() {
+            let mut tasks = match self.tasks.lock() {
+                Ok(tasks) => tasks,
+                Err(poisoned) => poisoned.into_inner(),
+            };
+            tasks.extend(pending);
+        }
+        match first_failure {
+            Some(message) => Err(message),
+            None => Ok(()),
+        }
     }
 }
 
@@ -2547,6 +2544,22 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn execution_scope_drains_remaining_tasks_after_join_failure_and_allows_retry() {
+        let failed = tokio::spawn(async { panic!("injected controller failure") });
+        let scope = RuntimeExecutionTaskScope::new(failed);
+        scope.push(
+            "protocol bridge",
+            tokio::spawn(std::future::pending()),
+            true,
+        );
+
+        let error = scope.join().await.unwrap_err();
+        assert!(error.contains("controller task failed"));
+        assert!(scope.tasks.lock().unwrap().is_empty());
+        scope.join().await.unwrap();
+    }
+
     fn adapter_command(
         backend: &WorkerRuntimeExecutionBackend<MockFactory>,
         worker_ref: &WorkerRef,
@@ -4197,7 +4210,36 @@ mod tests {
         let detail = runtime
             .create_worker(create_request("restore-after-stop"))
             .unwrap();
+        let failed_task = backend
+            .spawn_on_adapter_runtime(async { panic!("injected owned task failure") })
+            .unwrap();
+        backend
+            .workers
+            .lock()
+            .unwrap()
+            .get(&detail.worker_ref)
+            .unwrap()
+            .tasks
+            .push("injected failure", failed_task, false);
 
+        let first_stop = runtime.stop_worker(&detail.worker_ref, None).unwrap_err();
+        assert!(
+            first_stop
+                .to_string()
+                .contains("injected failure task failed")
+        );
+        assert_eq!(
+            runtime.worker_detail(&detail.worker_ref).unwrap().status,
+            crate::catalog::WorkerStatus::Idle
+        );
+        assert!(
+            backend
+                .workers
+                .lock()
+                .unwrap()
+                .contains_key(&detail.worker_ref),
+            "failed cleanup must retain retry authority"
+        );
         runtime.stop_worker(&detail.worker_ref, None).unwrap();
         assert_eq!(
             runtime.worker_detail(&detail.worker_ref).unwrap().status,

crates/worker/src/controller.rs

@@ -2383,6 +2383,10 @@ async fn controller_loop<C, St>(
         }
     }
 
+    // Close method admission before any fallible child/session cleanup. Existing
+    // senders then fail instead of accepting work that this execution can no
+    // longer process.
+    drop(method_rx);
     let had_socket_server = _socket_server.is_some();
     if let Some(socket_server) = _socket_server {
         socket_server.shutdown().await;

crates/worker/tests/controller_test.rs

@@ -2219,6 +2219,43 @@ async fn status_json_reflects_worker_name() {
 // Socket transport tests
 // ---------------------------------------------------------------------------
 
+#[tokio::test]
+async fn shutdown_closes_method_admission_before_terminal_confirmation() {
+    let worker = make_worker(MockClient::new(simple_text_events())).await;
+    let runtime_base = tempfile::tempdir().unwrap();
+    let bash_output_dir = runtime_base.path().join("bash-output");
+    let (handle, mut shutdown_rx) =
+        WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir)
+            .await
+            .unwrap();
+    handle
+        .send(Method::Shutdown {
+            command: worker_command(&handle),
+        })
+        .await
+        .unwrap();
+
+    tokio::time::timeout(std::time::Duration::from_secs(5), async {
+        loop {
+            tokio::select! {
+                biased;
+                result = handle.send(Method::ListRewindTargets) => {
+                    if result.is_err() {
+                        break;
+                    }
+                }
+                result = &mut shutdown_rx => {
+                    result.expect("controller shutdown signal should remain open");
+                    panic!("method admission remained open until terminal confirmation");
+                }
+            }
+        }
+    })
+    .await
+    .expect("method admission did not close during shutdown");
+    shutdown_rx.await.unwrap();
+}
+
 #[tokio::test]
 async fn shutdown_joins_socket_server_with_active_connection() {
     use tokio::net::UnixStream;