fix(workspace-server): fence cleanup removal races

crates/workspace-server/src/server.rs

@@ -684,9 +684,9 @@ impl WorkerRemovalService {
         }
     }
 
-    /// Terminal catalog deletion shared by cleanup and spawn compensation.
+    /// Terminal catalog deletion for rollback-style compensation. Normal cleanup
-    /// The Store atomically commits membership, reservation, link and projection
+    /// and WorkerRemove use the retention commit below; compensation preserves
-    /// state; this service owns the only production-side publication step.
+    /// its pre-create semantics while sharing the atomic catalog contract.
     fn commit_catalog_removal(&self, target: &RuntimeWorkerRef) -> crate::Result<()> {
         let commit = self
             .store
@@ -699,76 +699,31 @@ impl WorkerRemovalService {
             candidate.runtime_id.clone(),
             candidate.runtime_worker_id.clone(),
         );
-        let remove_lock = {
-            let mut locks = self
-                .worker_remove_locks
-                .lock()
-                .map_err(|_| Error::Store("Worker removal lock registry was poisoned".into()))?;
-            locks
-                .entry(target.clone())
-                .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
-                .clone()
-        };
-        let _remove_guard = remove_lock.lock().await;
-        let session_lock = {
-            let mut locks = self
-                .workdir_session_locks
-                .lock()
-                .map_err(|_| Error::Store("Workdir session lock registry was poisoned".into()))?;
-            locks
-                .entry(target.clone())
-                .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
-                .clone()
-        };
-        let _session_guard = session_lock.lock().await;
-        close_worker_workdir_sessions(&self.workdir_sessions, &target)
-            .await
-            .map_err(|message| Error::RuntimeOperationFailed {
-                runtime_id: target.runtime_id.clone(),
-                code: "workdir_session_close_failed".to_string(),
-                message,
-            })?;
         let runtime = self.runtime.upgrade().ok_or_else(|| {
             Error::Store("Workspace Runtime registry is unavailable during cleanup".to_string())
         })?;
-        match runtime.stop_worker(
+        let response = self
-            &target,
+            .execute_target_removal(
-            WorkerLifecycleRequest {
+                runtime.as_ref(),
-                reason: Some("cleanup worker before deletion".to_string()),
+                &target,
-                ticket_assignment: None,
+                &format!("Browser cleanup: {}", candidate.reason),
-            },
+                None,
-        ) {
+            )
-            Ok(result) if result.state == WorkerOperationState::Accepted => {}
+            .await
-            Ok(result) => {
+            .map_err(|message| {
-                return Err(ApiError::with_diagnostics(
+                cleanup_api_error(
-                    Error::RuntimeOperationFailed {
+                    candidate.runtime_id.as_str(),
-                        runtime_id: target.runtime_id.clone(),
+                    "workspace_cleanup_worker_removal_failed",
-                        code: "workspace_cleanup_worker_runtime_stop_rejected".to_string(),
+                    message.as_str(),
-                        message: "Runtime did not stop selected Worker before deletion".to_string(),
+                )
-                    },
+            })?;
-                    result.diagnostics,
+        if response.status != StatusCode::OK.as_u16() {
-                ));
+            return Err(cleanup_api_error(
-            }
+                candidate.runtime_id.as_str(),
-            Err(RuntimeRegistryError::UnknownWorker { .. }) => {}
+                "workspace_cleanup_worker_removal_rejected",
-            Err(error) => return Err(error.into_error().into()),
+                response.body.as_str(),
+            ));
         }
-        match runtime.delete_worker(&target) {
-            Ok(result) if result.deleted && result.state == WorkerOperationState::Accepted => {}
-            Ok(result) => {
-                return Err(ApiError::with_diagnostics(
-                    Error::RuntimeOperationFailed {
-                        runtime_id: target.runtime_id.clone(),
-                        code: "workspace_cleanup_worker_runtime_delete_rejected".to_string(),
-                        message: "Runtime did not delete selected Worker after stopping it"
-                            .to_string(),
-                    },
-                    result.diagnostics,
-                ));
-            }
-            Err(RuntimeRegistryError::UnknownWorker { .. }) => {}
-            Err(error) => return Err(error.into_error().into()),
-        }
-        self.commit_catalog_removal(&target)?;
         Ok(())
     }
 
@@ -12146,6 +12101,11 @@ fn build_runtime_cleanup_plan(
             &record.worker,
         )?;
         let is_running = live_running_worker_ids.contains(&record.worker);
+        let is_internal = api
+            .runtime
+            .worker(&record.worker)
+            .ok()
+            .is_some_and(|worker| worker.singleton_key.is_some());
         let pinned = record.retention_state == "pinned";
         let blocking_reason = if let Some(assignment) = current_assignment {
             Some(format!(
@@ -12155,6 +12115,8 @@ fn build_runtime_cleanup_plan(
             ))
         } else if pinned {
             Some("worker is pinned".to_string())
+        } else if is_internal {
+            Some("internal service Worker cannot be deleted".to_string())
         } else if is_running {
             Some("worker is running".to_string())
         } else {
@@ -27288,6 +27250,14 @@ mod tests {
         .await
         .expect("spawn compensation catalog removal broadcast");
         assert!(removal.revision > 0);
+        sync_worker_observation(&api, &summary).unwrap();
+        assert!(
+            api.store
+                .get_worker_registry(TEST_WORKSPACE_ID, &worker)
+                .unwrap()
+                .is_none(),
+            "a late pre-compensation observation must not clear the removal fence"
+        );
     }
 
     #[tokio::test]
@@ -28665,6 +28635,66 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn cleanup_service_rejects_internal_worker_before_runtime_removal() {
+        let workspace = tempfile::tempdir().unwrap();
+        init_clean_git_workspace(workspace.path());
+        let api = test_api(workspace.path()).await;
+        let Json(orchestrator) = scoped_start_workspace_orchestrator(
+            State(api.clone()),
+            AxumPath(ScopedWorkspacePath {
+                workspace_id: TEST_WORKSPACE_ID.to_string(),
+            }),
+        )
+        .await
+        .unwrap();
+        let worker = orchestrator.worker.unwrap();
+        let worker = RuntimeWorkerRef::new(worker.runtime_id, worker.worker_id);
+        api.runtime
+            .stop_worker(
+                &worker,
+                WorkerLifecycleRequest {
+                    reason: Some("prepare internal cleanup guard".to_string()),
+                    ticket_assignment: None,
+                },
+            )
+            .unwrap();
+        let plan = build_runtime_cleanup_plan(&api, worker.runtime_id.as_str())
+            .unwrap_or_else(|error| panic!("cleanup plan: {}", error.error));
+        let candidate = plan
+            .workers
+            .iter()
+            .find(|candidate| candidate.runtime_worker_id == worker.worker_id)
+            .expect("internal Worker cleanup candidate")
+            .clone();
+        assert_eq!(
+            candidate.blocking_reason.as_deref(),
+            Some("internal service Worker cannot be deleted")
+        );
+        let mut subscriber = api.worker_projection.subscribe();
+
+        let error = WorkerRemovalService::new(&api)
+            .execute_cleanup_removal(&candidate)
+            .await
+            .unwrap_err();
+        assert!(matches!(
+            error.error,
+            Error::RuntimeOperationFailed { ref code, .. }
+                if code == "workspace_cleanup_worker_removal_rejected"
+        ));
+        assert!(api.runtime.worker(&worker).is_ok());
+        assert!(
+            api.store
+                .get_worker_registry(TEST_WORKSPACE_ID, &worker)
+                .unwrap()
+                .is_some()
+        );
+        assert!(matches!(
+            subscriber.try_recv(),
+            Err(tokio::sync::broadcast::error::TryRecvError::Empty)
+        ));
+    }
+
     #[tokio::test]
     async fn cleanup_http_entry_publishes_atomic_catalog_removal_without_runtime_event() {
         let workspace = tempfile::tempdir().unwrap();
@@ -28708,9 +28738,54 @@ mod tests {
             )
             .unwrap();
         assert_eq!(stopped.state, WorkerOperationState::Accepted);
+        let worker_root = workspace
+            .path()
+            .join(".test-embedded-runtime-store/workers")
+            .join(&worker.worker_id);
+        fs::create_dir_all(worker_root.join("session/segments")).unwrap();
+        fs::write(
+            worker_root.join("session/session.json"),
+            serde_json::to_vec_pretty(&json!({
+                "schema_version": 1,
+                "session_id": "cleanup-removal-session"
+            }))
+            .unwrap(),
+        )
+        .unwrap();
+        fs::write(
+            worker_root.join("session/segments/segment-a.jsonl"),
+            b"cleanup retention evidence\n",
+        )
+        .unwrap();
         let summary = api.runtime.worker(&worker).unwrap();
         sync_worker_observation(&api, &summary).unwrap();
+        // Deliberately drop Runtime observation delivery: catalog removal itself
+        // must drive every connected Workspace subscription.
+        api.worker_projection.shutdown();
         let worker_id = worker.worker_id.clone();
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let app = build_inner_router(api.clone()).layer(Extension(test_browser_request_actor()));
+        let server = tokio::spawn(async move {
+            let _ = axum::serve(listener, app).await;
+        });
+        let subscription_url = format!("ws://{address}/api/w/{TEST_WORKSPACE_ID}/protocol/ws");
+        let (mut acting_socket, _) = connect_async(subscription_url.as_str()).await.unwrap();
+        let (mut observing_socket, _) = connect_async(subscription_url.as_str()).await.unwrap();
+        let acting_snapshot =
+            subscribe_workspace_worker_snapshot(&mut acting_socket, "cleanup-acting").await;
+        let observing_snapshot =
+            subscribe_workspace_worker_snapshot(&mut observing_socket, "cleanup-observing").await;
+        assert!(
+            acting_snapshot
+                .iter()
+                .any(|item| item.worker_id.as_str() == worker_id)
+        );
+        assert!(
+            observing_snapshot
+                .iter()
+                .any(|item| item.worker_id.as_str() == worker_id)
+        );
         let mut subscriber = api.worker_projection.subscribe();
         let plan = build_runtime_cleanup_plan(&api, EMBEDDED_WORKER_RUNTIME_ID)
             .unwrap_or_else(|error| panic!("cleanup plan: {}", error.error));
@@ -28739,20 +28814,51 @@ mod tests {
         .unwrap();
         assert_eq!(response.results[0].status, "deleted");
 
-        let event = tokio::time::timeout(std::time::Duration::from_secs(1), subscriber.recv())
+        let event = tokio::time::timeout(std::time::Duration::from_secs(1), async {
-            .await
+            loop {
-            .expect("catalog removal broadcast")
+                let event = subscriber.recv().await.expect("projection event");
-            .expect("projection event");
+                if event.changes.iter().any(|change| {
+                    matches!(change, crate::store::WorkerCatalogChange::Removed(removed) if removed == &worker)
+                }) {
+                    return event;
+                }
+            }
+        })
+        .await
+        .expect("catalog removal broadcast");
         assert!(matches!(
             event.changes.as_slice(),
             [crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker
         ));
+        let acting_revision = tokio::time::timeout(
+            std::time::Duration::from_secs(1),
+            next_workspace_worker_removal(&mut acting_socket, worker_id.as_str()),
+        )
+        .await
+        .expect("acting Workspace Worker subscription removal");
+        let observing_revision = tokio::time::timeout(
+            std::time::Duration::from_secs(1),
+            next_workspace_worker_removal(&mut observing_socket, worker_id.as_str()),
+        )
+        .await
+        .expect("independent Workspace Worker subscription removal");
+        assert_eq!(acting_revision, event.revision);
+        assert_eq!(observing_revision, event.revision);
         assert!(
             api.store
                 .get_worker_registry(&api.config.workspace_id, &worker)
                 .unwrap()
                 .is_none()
         );
+        let recovered = api
+            .store
+            .recover_worker_removal_execution(&api.config.workspace_id, &worker)
+            .unwrap()
+            .expect("Browser cleanup uses durable retention removal");
+        assert_eq!(
+            recovered.plan.state,
+            crate::retention::WorkerRemovalPlanState::Succeeded
+        );
         let (snapshot_revision, snapshot) = api
             .store
             .worker_registry_projection_snapshot(&api.config.workspace_id, 100)
@@ -28763,6 +28869,21 @@ mod tests {
                 .iter()
                 .all(|record| record.registry.worker != worker)
         );
+        let rest_workers = workers_response(api.clone()).unwrap();
+        assert!(
+            rest_workers
+                .items
+                .iter()
+                .all(|item| item.worker_id != worker_id)
+        );
+        let (mut fresh_socket, _) = connect_async(subscription_url.as_str()).await.unwrap();
+        let fresh_snapshot =
+            subscribe_workspace_worker_snapshot(&mut fresh_socket, "cleanup-fresh").await;
+        assert!(
+            fresh_snapshot
+                .iter()
+                .all(|item| item.worker_id.as_str() != worker_id)
+        );
 
         // A retry replays the durable fence at the same revision, healing an
         // in-process broadcast omitted after the first DB commit.
@@ -28778,6 +28899,7 @@ mod tests {
             replay.changes.as_slice(),
             [crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker
         ));
+        server.abort();
     }
 
     #[tokio::test]
@@ -32123,6 +32245,74 @@ mod tests {
         serde_json::from_str(&text).unwrap()
     }
 
+    type TestWebSocket = tokio_tungstenite::WebSocketStream<
+        tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
+    >;
+
+    async fn subscribe_workspace_worker_snapshot(
+        socket: &mut TestWebSocket,
+        request_id: &str,
+    ) -> Vec<protocol::subscription::SubscriptionWorker> {
+        let frame = protocol::subscription::SubscriptionFrame::new(
+            protocol::subscription::SubscriptionFramePayload::Request(
+                protocol::subscription::SubscriptionRequest::SubscribeEvents {
+                    request_id: protocol::subscription::SubscriptionRequestId::new(request_id)
+                        .unwrap(),
+                    selector: protocol::subscription::EventSubscriptionSelector::WorkspaceWorkers,
+                },
+            ),
+        );
+        socket
+            .send(Message::Text(serde_json::to_string(&frame).unwrap().into()))
+            .await
+            .unwrap();
+        loop {
+            let Message::Text(text) = socket.next().await.unwrap().unwrap() else {
+                continue;
+            };
+            let frame: protocol::subscription::SubscriptionFrame =
+                serde_json::from_str(text.as_str()).unwrap();
+            if let protocol::subscription::SubscriptionFramePayload::Response(
+                protocol::subscription::SubscriptionResponse::Subscribed {
+                    selector: protocol::subscription::EventSubscriptionSelector::WorkspaceWorkers,
+                    snapshot: protocol::subscription::SubscriptionSnapshot::Workers { workers },
+                    ..
+                },
+            ) = frame.payload
+            {
+                return workers;
+            }
+        }
+    }
+
+    async fn next_workspace_worker_removal(
+        socket: &mut TestWebSocket,
+        expected_worker_id: &str,
+    ) -> u64 {
+        loop {
+            let Message::Text(text) = socket.next().await.unwrap().unwrap() else {
+                continue;
+            };
+            let frame: protocol::subscription::SubscriptionFrame =
+                serde_json::from_str(text.as_str()).unwrap();
+            if let protocol::subscription::SubscriptionFramePayload::Event(
+                protocol::subscription::SubscriptionEvent::Event {
+                    subject_revision,
+                    payload:
+                        protocol::subscription::SubscriptionEventPayload::WorkerRemoved {
+                            worker_id,
+                            ..
+                        },
+                    ..
+                },
+            ) = frame.payload
+                && worker_id.as_str() == expected_worker_id
+            {
+                return subject_revision;
+            }
+        }
+    }
+
     async fn spawn_runtime_worker() -> (
         worker_runtime::Runtime,
         worker_runtime::identity::WorkerRef,

crates/workspace-server/src/store.rs

@@ -5576,14 +5576,17 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
         &self,
         record: &WorkerRegistryRecord,
     ) -> Result<WorkerRegistryProjectionCommit> {
-        self.with_conn(|conn| {
+        self.with_conn_mut(|conn| {
-            let tx = conn.unchecked_transaction()?;
+            let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
             let removal_blocks_upsert: bool = tx.query_row(
                 "SELECT EXISTS(
                     SELECT 1 FROM worker_removal_operations
                     WHERE workspace_id = ?1 AND runtime_id = ?2
                       AND worker_id = ?3
                       AND state IN ('executing', 'failed', 'succeeded')
+                    UNION ALL
+                    SELECT 1 FROM worker_registry_projection_removals
+                    WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3
                 )",
                 params![
                     record.workspace_id,
@@ -5647,10 +5650,6 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
                 &record.worker.worker_id,
                 &record.created_at,
             )?;
-            tx.execute(
-                "DELETE FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3",
-                params![record.workspace_id, record.worker.runtime_id, record.worker.worker_id],
-            )?;
             let changed_workers = if changed == 0 {
                 Vec::new()
             } else {
@@ -15171,6 +15170,29 @@ INSERT INTO worker_registry (
             [WorkerCatalogChange::Removed(worker)] if worker == &worker_ref
         ));
         assert_eq!(duplicate_removal.revision, first_removal.revision);
+        let stale_upsert = store
+            .upsert_worker_registry(&WorkerRegistryRecord {
+                workspace_id: "local-dev".to_string(),
+                worker: worker_ref.clone(),
+                display_name: "Stale Removed Worker".to_string(),
+                profile: None,
+                retention_state: "normal".to_string(),
+                transcript_ref: None,
+                session_ref: None,
+                summary_ref: None,
+                diagnostics_ref: None,
+                created_at: "1".to_string(),
+                updated_at: "stale-after-removal".to_string(),
+            })
+            .unwrap();
+        assert!(stale_upsert.changes.is_empty());
+        assert_eq!(stale_upsert.revision, first_removal.revision);
+        assert!(
+            store
+                .get_worker_registry("local-dev", &worker_ref)
+                .unwrap()
+                .is_none()
+        );
         let event = SubscriptionWorker {
             worker_id: SubscriptionWorkerId::new("removed").unwrap(),
             runtime_id: Some("embedded".to_string()),