fix: checkpoint workspace cleanup children

crates/workspace-server/src/server.rs

@@ -164,6 +164,7 @@ use crate::workdir_removal::{
     WorkdirRemovalOperationState, workdir_removal_intent,
 };
 use crate::workspace_catalog::{WorkspaceCatalogService, WorkspaceCreateRequest};
+use crate::workspace_deletion::WorkspaceDeletionStore;
 use crate::{Error, Result};
 use worker_runtime::catalog::{
     ConfigBundleRef, ProfileSelector, RepositoryMaterializationContext, RepositoryRefObservation,
@@ -746,7 +747,8 @@ impl WorkspaceWorkerRemoveExecutor {
             ));
         }
 
-        self.execute_target_removal(&runtime, &target, reason).await
+        self.execute_target_removal(&runtime, &target, reason, None)
+            .await
     }
 
     async fn execute_target_removal(
@@ -754,6 +756,7 @@ impl WorkspaceWorkerRemoveExecutor {
         runtime: &RuntimeRegistry,
         target: &RuntimeWorkerRef,
         reason: &str,
+        parent_workspace_deletion_operation_id: Option<&str>,
     ) -> std::result::Result<worker::WorkspaceResponse, String> {
         let remove_lock = {
             let mut locks = self
@@ -802,6 +805,14 @@ impl WorkspaceWorkerRemoveExecutor {
             } else {
                 prepared
             };
+            if let Some(parent_operation_id) = parent_workspace_deletion_operation_id {
+                self.store
+                    .append_workspace_deletion_child_operation(
+                        parent_operation_id,
+                        &prepared.plan.operation_id,
+                    )
+                    .map_err(|error| error.to_string())?;
+            }
             if close_worker_workdir_sessions(&self.workdir_sessions, &target)
                 .await
                 .is_err()
@@ -898,6 +909,14 @@ impl WorkspaceWorkerRemoveExecutor {
             Err(error) => return Ok(worker_retention_error_response(error)),
         };
 
+        if let Some(parent_operation_id) = parent_workspace_deletion_operation_id {
+            self.store
+                .append_workspace_deletion_child_operation(
+                    parent_operation_id,
+                    &prepared.plan.operation_id,
+                )
+                .map_err(|error| error.to_string())?;
+        }
         if close_worker_workdir_sessions(&self.workdir_sessions, &target)
             .await
             .is_err()
@@ -1177,19 +1196,21 @@ impl WorkspaceServerApi {
             let worker_key = worker.display_name.clone();
             let target = worker.worker;
             let response = WorkspaceWorkerRemoveExecutor::new(&api)
-                .execute_target_removal(api.runtime.as_ref(), &target, "Workspace deletion")
+                .execute_target_removal(
+                    api.runtime.as_ref(),
+                    &target,
+                    "Workspace deletion",
+                    Some(operation_id),
+                )
                 .await
                 .map_err(Error::Store)?;
-            if let Some(child_operation_id) = self.store.latest_worker_removal_operation_id(
-                &operation.workspace_id,
-                &target.runtime_id,
-                &target.worker_id,
-            )? {
             child_operation_ids = self
                 .store
-                    .append_workspace_deletion_child_operation(operation_id, &child_operation_id)?
+                .workspace_deletion_operation_for_recovery(operation_id)?
+                .ok_or_else(|| {
+                    Error::Store("Workspace deletion operation disappeared".to_string())
+                })?
                 .child_operation_ids;
-            }
             if response.status != 200 {
                 blockers.push(WorkspaceDeletionBlocker {
                     kind: WorkspaceDeletionBlockerKind::WorkerRemovalBlocked,
@@ -10873,6 +10894,8 @@ fn execute_workdir_removal_for_workspace_deletion(
         api.config_store
             .reserve_workdir_removal_operation(&intent)?
     };
+    api.config_store
+        .append_workspace_deletion_child_operation(parent_operation_id, &operation.operation_id)?;
     execute_reserved_workdir_removal(api, operation, false)
 }
 
@@ -17132,6 +17155,34 @@ mod tests {
         assert!(source.contains("api.recover_workspace_deletions().await?"));
     }
 
+    #[test]
+    fn workspace_deletion_checkpoints_child_operations_before_external_cleanup() {
+        let source = include_str!("server.rs");
+        let worker_start = source
+            .find("async fn execute_target_removal")
+            .expect("Worker removal helper");
+        let worker_source = &source[worker_start..];
+        let checkpoint = worker_source
+            .find("append_workspace_deletion_child_operation")
+            .expect("Worker child checkpoint");
+        let cleanup = worker_source
+            .find("close_worker_workdir_sessions")
+            .expect("Worker cleanup side effect");
+        assert!(checkpoint < cleanup);
+
+        let workdir_start = source
+            .find("fn execute_workdir_removal_for_workspace_deletion")
+            .expect("Workdir removal helper");
+        let workdir_source = &source[workdir_start..];
+        let checkpoint = workdir_source
+            .find("append_workspace_deletion_child_operation")
+            .expect("Workdir child checkpoint");
+        let cleanup = workdir_source
+            .find("execute_reserved_workdir_removal")
+            .expect("Workdir cleanup side effect");
+        assert!(checkpoint < cleanup);
+    }
+
     fn handler_source<'a>(source: &'a str, name: &str) -> &'a str {
         let start = source
             .find(&format!("async fn {name}"))

crates/workspace-server/src/workspace_deletion.rs

@@ -110,13 +110,6 @@ pub trait WorkspaceDeletionStore: Send + Sync {
         operation_id: &str,
     ) -> Result<Option<WorkspaceDeletionOperationResponse>>;
 
-    fn latest_worker_removal_operation_id(
-        &self,
-        workspace_id: &str,
-        runtime_id: &str,
-        worker_id: &str,
-    ) -> Result<Option<String>>;
-
     fn workspace_deletion_operation_for_recovery(
         &self,
         operation_id: &str,
@@ -313,25 +306,6 @@ impl WorkspaceDeletionStore for SqliteWorkspaceStore {
         })
     }
 
-    fn latest_worker_removal_operation_id(
-        &self,
-        workspace_id: &str,
-        runtime_id: &str,
-        worker_id: &str,
-    ) -> Result<Option<String>> {
-        self.with_conn(|conn| {
-            conn.query_row(
-                "SELECT operation_id FROM worker_removal_operations
-                 WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3
-                 ORDER BY created_at DESC LIMIT 1",
-                params![workspace_id, runtime_id, worker_id],
-                |row| row.get(0),
-            )
-            .optional()
-            .map_err(Into::into)
-        })
-    }
-
     fn workspace_deletion_operation_for_recovery(
         &self,
         operation_id: &str,
@@ -970,6 +944,10 @@ mod tests {
             .finalize_workspace_deletion(&request.operation_id)
             .expect("finalize");
         assert_eq!(completed.state, WorkspaceDeletionState::Succeeded);
+        assert_eq!(
+            completed.child_operation_ids,
+            vec!["child-operation-1".to_string()]
+        );
         let replayed = store
             .finalize_workspace_deletion(&request.operation_id)
             .expect("finalize replay");