fix: drain background forks on shutdown and fork in the background from the CLI

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

backend/Cargo.lock

@@ -15551,6 +15551,7 @@ dependencies = [
  "strum",
  "tokio",
  "tokio-postgres",
+ "tokio-util",
  "tracing",
  "url",
  "uuid",

backend/Cargo.toml

@@ -517,7 +517,7 @@ serde_urlencoded = "^0"
 astral-tokio-tar = "^0.6.2"
 tempfile = "^3"
 x509-parser = "^0.16"
-tokio-util = { version = "=0.7.17", features = ["io"] }
+tokio-util = { version = "=0.7.17", features = ["io", "rt"] }
 json-pointer = "^0"
 itertools = "^0.14.0"
 regex = "^1"

backend/windmill-api-workspaces/Cargo.toml

@@ -50,6 +50,7 @@ sha2.workspace = true
 sqlx.workspace = true
 tokio.workspace = true
 tokio-postgres.workspace = true
+tokio-util.workspace = true
 tracing.workspace = true
 url.workspace = true
 uuid.workspace = true

backend/windmill-api-workspaces/src/workspaces.rs

@@ -9079,7 +9079,7 @@ async fn create_workspace_fork(
         )));
     };
 
-    tokio::spawn(async move {
+    BACKGROUND_FORKS.spawn(async move {
         let heartbeat = {
             let (db, fork_id) = (db.clone(), fork_id.clone());
             tokio::spawn(async move {
@@ -9139,6 +9139,12 @@ async fn create_workspace_fork(
 const FORK_HEARTBEAT_SECS: u64 = 10;
 const FORK_HEARTBEAT_STALE_SECS: u64 = 60;
 
+/// Forks being created in the background. The server waits for them on shutdown: a copy cut short
+/// never drops the data table databases it already made, and a retry under the same id is refused
+/// while they exist.
+pub static BACKGROUND_FORKS: std::sync::LazyLock<tokio_util::task::TaskTracker> =
+    std::sync::LazyLock::new(tokio_util::task::TaskTracker::new);
+
 #[derive(Deserialize)]
 struct CreateWorkspaceForkQuery {
     /// Return once the request is validated and create the fork in the background; its progress

backend/windmill-api/src/lib.rs

@@ -1279,6 +1279,15 @@ pub async fn run_server(
     });
 
     server.await?;
+    let forks = &windmill_api_workspaces::workspaces::BACKGROUND_FORKS;
+    forks.close();
+    if !forks.is_empty() {
+        tracing::info!(
+            "server off. waiting for {} fork(s) being created",
+            forks.len()
+        );
+        forks.wait().await;
+    }
     #[cfg(feature = "agent_worker_server")]
     for (i, bg_processor) in agent_workers_bg_processor.into_iter().enumerate() {
         tracing::info!("server off. shutting down agent worker bg processor {i}");

cli/src/commands/workspace/fork.ts

@@ -22,6 +22,34 @@ import {
   readConfigFile,
 } from "../../core/conf.ts";
 
+/**
+ * Wait for a fork requested with `background: true`. A server without background forks ignores
+ * the flag and has already created the fork when the request returns; it has no status route.
+ */
+async function waitForForkCreation(parentWorkspace: string, forkId: string) {
+  let failedPolls = 0;
+  while (true) {
+    await new Promise((resolve) => setTimeout(resolve, 1500));
+    let result: Awaited<ReturnType<typeof wmill.getForkCreationStatus>>;
+    try {
+      result = await wmill.getForkCreationStatus({
+        workspace: parentWorkspace,
+        forkWorkspaceId: forkId,
+      });
+      failedPolls = 0;
+    } catch (e) {
+      if ((e as { status?: number }).status === 404) return;
+      // A poll lost to a proxy says nothing about the fork, which is still being created.
+      if (++failedPolls >= 5) throw e;
+      continue;
+    }
+    if (result.status === "completed") return;
+    if (result.status === "failed") {
+      throw new Error(result.error ?? "Unknown error");
+    }
+  }
+}
+
 async function createWorkspaceFork(
   opts: GlobalOptions & {
     createWorkspaceName: string | undefined;
@@ -376,8 +404,9 @@ async function createWorkspaceFork(
 
   // --- Create the fork workspace ---
   try {
-    const result = await wmill.createWorkspaceFork({
+    await wmill.createWorkspaceFork({
       workspace: workspace.workspaceId,
+      background: true,
       requestBody: {
         id: trueWorkspaceId,
         name: opts.createWorkspaceName ?? workspaceName ?? trueWorkspaceId,
@@ -385,8 +414,9 @@ async function createWorkspaceFork(
         forked_datatables: forkedDatatables,
       },
     });
+    await waitForForkCreation(workspace.workspaceId, trueWorkspaceId);
 
-    log.info(colors.green(`✅ ${result}`));
+    log.info(colors.green(`✅ Created forked workspace ${trueWorkspaceId}`));
   } catch (error) {
     log.error(
       colors.red(`Failed to create forked workspace: ${(error as Error).message}`),