Retain shared machines while new workspaces prepare

apps/server/src/services/machines/provider-orchestration.ts

@@ -21,6 +21,7 @@ import {
   machineHasStartingThreadLaunch,
   machineHasProvisioningEnvironment,
   machineHasLiveThreads,
+  machineHasPendingThreads,
   updateHost,
 } from "@bb/db";
 import { jsonValueSchema, type Host, type JsonValue } from "@bb/domain";
@@ -1158,6 +1159,7 @@ export function requestAutomaticMachineRemoval(
   }
   if (row.type !== "ephemeral") return false;
   if (
+    machineHasPendingThreads(deps.db, hostId) ||
     machineHasLiveThreadLaunch(deps.db, hostId) ||
     machineHasLiveThreads(deps.db, hostId)
   ) {

apps/server/src/services/threads/thread-create.ts

@@ -46,7 +46,10 @@ import {
   resolveStableThreadRequestEnvironment,
   type ResolvedStableThreadRequestEnvironment,
 } from "./thread-request-eligibility.js";
-import { resolveThreadEnvironmentPlacement } from "./thread-environment-placement.js";
+import {
+  requireEnvironmentPlacementHost,
+  resolveThreadEnvironmentPlacement,
+} from "./thread-environment-placement.js";
 import {
   buildProviderThreadExecutionDefaults,
   resolveCreateThreadEnvironment,
@@ -392,6 +395,12 @@ async function createPendingThreadAndAttemptFirstDispatch(
       startedOnBehalfOf: args.request.startedOnBehalfOf,
       titleProvided: Boolean(args.request.title),
     };
+    const placementHostId = hostIdForEnvironmentIntent(
+      deps,
+      args.environmentIntent,
+    );
+    if (placementHostId !== null)
+      requireEnvironmentPlacementHost(deps, placementHostId);
     setThreadStartupContext(deps.db, {
       threadId: thread.id,
       startupContext: JSON.stringify({ kind: "pending", ...startContext }),

apps/server/src/services/threads/thread-environment-placement.ts

@@ -175,6 +175,21 @@ export async function parseProviderInputs(
   return value.data;
 }
 
+export function requireEnvironmentPlacementHost(
+  deps: PlacementDeps,
+  hostId: string,
+) {
+  const host = requireNonDestroyedHostWithStatus(deps, hostId);
+  if (host.lifecycle.phase === "removing") {
+    throw new ApiError(
+      409,
+      "machine_removing",
+      "Machine is being removed and cannot accept new environments",
+    );
+  }
+  return host;
+}
+
 export async function completeProviderSelection(
   deps: PlacementDeps,
   record: PluginEnvironmentProviderRecord,
@@ -185,17 +200,7 @@ export async function completeProviderSelection(
   const requires = record.provider.requires;
   let machine: EnvironmentMachineSelection;
   if (selection.machine.type === "existing") {
-    const host = requireNonDestroyedHostWithStatus(
-      deps,
-      selection.machine.hostId,
-    );
-    if (host.lifecycle.phase === "removing") {
-      throw new ApiError(
-        409,
-        "machine_removing",
-        "Machine is being removed and cannot accept new environments",
-      );
-    }
+    requireEnvironmentPlacementHost(deps, selection.machine.hostId);
     if (requires.projectCheckout) {
       requireSourceForHost(deps, projectId, selection.machine.hostId);
     }
@@ -1070,6 +1075,7 @@ export function prepareProviderEnvironment(
     statusMessage?: string;
   } = {},
 ): ProviderEnvironmentCreationDecision {
+  requireEnvironmentPlacementHost(deps, context.host.id);
   const now = Date.now();
   const policy = record.provider.policy;
   const previous =

apps/server/test/services/machines/provider-orchestration.test.ts

@@ -7,6 +7,7 @@ import {
   getEnvironment,
   hosts,
   archiveThread,
+  setThreadStartupContext,
   updateHost,
 } from "@bb/db";
 import { createDeferredPromise } from "@bb/test-helpers";
@@ -465,6 +466,40 @@ describe("machine retirement", () => {
       }),
   );
 
+  it("preserves explicit removal of a machine with an unattached pending start", async () =>
+    withTestHarness(async (harness) => {
+      const remove = vi.fn(async () => ({ status: "removed" as const }));
+      installMachineProvider({ ephemeral: true, remove });
+      const { host } = seedHostSession(harness.deps);
+      const { project } = seedProjectWithSource(harness.deps, {
+        hostId: host.id,
+      });
+      const thread = seedThread(harness.deps, {
+        projectId: project.id,
+        status: "pending",
+      });
+      setThreadStartupContext(harness.db, {
+        threadId: thread.id,
+        startupContext: JSON.stringify({
+          kind: "pending",
+          environmentIntent: {
+            type: "provider",
+            machine: { type: "existing", hostId: host.id },
+          },
+        }),
+      });
+      updateHost(harness.db, harness.hub, host.id, {
+        type: "ephemeral",
+        machineProviderId: "test-machine",
+        resource: {},
+      });
+      expect(requestAutomaticMachineRemoval(harness.deps, host.id)).toBe(false);
+      expect(requestMachineRemoval(harness.deps, host.id)).toBe(true);
+      await sweepProviderMachine(harness.deps, host.id);
+      expect(remove).toHaveBeenCalledOnce();
+      expect(getHost(harness.db, host.id)?.phase).toBe("destroyed");
+    }));
+
   it("keeps a persistent machine with no threads", async () =>
     withTestHarness(async (harness) => {
       installMachineProvider();

apps/server/test/threads/environment-providers.test.ts

@@ -1,3 +1,9 @@
+import {
+  prepareProviderEnvironment,
+  resolveProviderOperationContext,
+} from "../../src/services/threads/thread-environment-placement.js";
+import { cancelProviderEnvironmentCreation } from "../../src/services/environments/environment-engine.js";
+import { sweepProviderMachine } from "../../src/services/machines/provider-orchestration.js";
 import { stopThreadForCurrentState } from "../../src/services/threads/thread-lifecycle.js";
 import { createDeferredPromise } from "@bb/test-helpers";
 import { resolveGitCheckoutAvailability } from "../../src/services/environments/provider-availability.js";
@@ -12,6 +18,7 @@ import {
   createProjectSource,
   ensurePersonalProject,
   getEnvironment,
+  getHost,
   getNonDestroyedHostByLaunchKey,
   getPreparingEnvironment,
   getDefaultProjectSource,
@@ -643,6 +650,206 @@ function readyAt(host: { id: string }): TestProviderDecision {
   };
 }
 
+describe("shared machine preparation retention", () => {
+  it.each([
+    ["archive", false],
+    ["delete-project", false],
+    ["archive", true],
+    ["delete-project", true],
+  ] as const)(
+    "retains unfinished work during %s (scheduled: %s)",
+    async (action, scheduled) => {
+      await withTestHarness(async (harness) => {
+        const { host, project, environment } = seedTargetFixture(
+          harness,
+          "shared-machine",
+        );
+        const owner = seedThread(harness.deps, {
+          projectId: project.id,
+          environmentId: environment.id,
+        });
+        const otherProject = seedProjectWithSource(harness.deps, {
+          hostId: host.id,
+        }).project;
+        const remove = vi.fn(async () => ({ status: "removed" as const }));
+        const machine = {
+          pluginId: "cloud",
+          provider: validatePluginMachineProviderDeclaration({
+            id: "test-machine",
+            displayName: "Test machine",
+            description: "Test machine",
+            icon: "Terminal",
+            ephemeral: true,
+            create: async () => {
+              throw new Error("Unexpected machine allocation");
+            },
+            reconcileCleanup: async () => ({ status: "removed" }),
+            remove,
+          }),
+        };
+        setPluginMachineProviderBridge({
+          listMachineProviders: () => [machine],
+          getMachineProvider: (id) =>
+            id === machine.provider.id ? machine : undefined,
+          invokeProvider: async (_pluginId, _label, run) => ({
+            ok: true,
+            value: await run(),
+          }),
+          decisionTimeoutMs: 10000,
+        });
+        updateHost(harness.db, harness.hub, host.id, {
+          type: "ephemeral",
+          machineProviderId: machine.provider.id,
+          launchKey: owner.id,
+          resource: {},
+        });
+        const entered = createDeferredPromise<void>();
+        const release = createDeferredPromise<void>();
+        installTarget({
+          provision: async () => {
+            entered.resolve();
+            await release.promise;
+            return { action: "reject", message: "Setup failed" };
+          },
+        });
+        let nextThreadId: string | null = null;
+        try {
+          const next = await createThreadFromRequest(harness.deps, {
+            projectId: otherProject.id,
+            environment: {
+              type: "provider",
+              environmentProviderId: PROVIDER_ID,
+              machine: { type: "existing", hostId: host.id },
+              inputs: null,
+            },
+            input: textInput("Prepare shared workspace"),
+            providerId: "codex",
+            model: "requested-model",
+            origin: "app",
+            startedOnBehalfOf: null,
+            ...(scheduled ? { sendAt: Date.now() + 60000 } : {}),
+          });
+          nextThreadId = next.id;
+          if (!scheduled) await entered.promise;
+          expect(getThread(harness.db, next.id)).toMatchObject({
+            environmentId: null,
+            status: scheduled ? "pending" : "starting",
+          });
+          const response = await harness.app.request(
+            action === "archive"
+              ? `/api/v1/threads/${owner.id}/archive-all`
+              : `/api/v1/projects/${project.id}`,
+            { method: action === "archive" ? "POST" : "DELETE" },
+          );
+          expect(response.status).toBe(200);
+          await sweepProviderMachine(harness.deps, host.id);
+          expect(remove).not.toHaveBeenCalled();
+          expect(getHost(harness.db, host.id)?.phase).toBe("active");
+          if (scheduled) {
+            const cancelled = await harness.app.request(
+              `/api/v1/threads/${next.id}/archive-all`,
+              { method: "POST" },
+            );
+            expect(cancelled.status).toBe(200);
+          } else {
+            expect(getPreparingEnvironment(harness.db, next.id)?.status).toBe(
+              "creating",
+            );
+            release.resolve();
+            await expect
+              .poll(() => getPreparingEnvironment(harness.db, next.id)?.status)
+              .toBe("error");
+            await advanceThreadProvisioning(harness.deps, {
+              threadId: next.id,
+            });
+            expect(getThread(harness.db, next.id)?.status).toBe("error");
+          }
+          await sweepProviderMachine(harness.deps, host.id);
+          expect(remove).toHaveBeenCalledTimes(1);
+          expect(getHost(harness.db, host.id)?.phase).toBe("destroyed");
+        } finally {
+          release.resolve();
+          if (nextThreadId !== null)
+            await cancelProviderEnvironmentCreation(harness.deps, nextThreadId);
+        }
+      });
+    },
+  );
+
+  it("rechecks removal after resolving a preparation context", async () => {
+    await withTestHarness(async (harness) => {
+      const { host, project } = seedTargetFixture(
+        harness,
+        "removing-before-reservation",
+      );
+      const provision = vi.fn(() => readyAt(host));
+      installTarget({ provision });
+      const record = listEnvironmentProviders()[0];
+      const thread = seedThread(harness.deps, {
+        projectId: project.id,
+        status: "starting",
+      });
+      const context = await resolveProviderOperationContext(
+        harness.deps,
+        thread,
+        {
+          type: "provider",
+          environmentProviderId: PROVIDER_ID,
+          machine: { type: "existing", hostId: host.id },
+          inputs: null,
+          selectionResolved: true,
+        },
+        record,
+      );
+      if (context === null) throw new Error("Missing preparation context");
+      updateHost(harness.db, harness.hub, host.id, { phase: "removing" });
+      expect(() =>
+        prepareProviderEnvironment(harness.deps, record, context),
+      ).toThrow(/remov/i);
+      expect(getPreparingEnvironment(harness.db, thread.id)).toBeNull();
+      expect(provision).not.toHaveBeenCalled();
+    });
+  });
+
+  it("rejects a stale selection when removal wins during validation", async () => {
+    await withTestHarness(async (harness) => {
+      const { host, project } = seedTargetFixture(
+        harness,
+        "removing-after-validation",
+      );
+      const entered = createDeferredPromise<void>();
+      const release = createDeferredPromise<void>();
+      const provision = vi.fn(() => readyAt(host));
+      installTarget({
+        provision,
+        validate: async () => {
+          entered.resolve();
+          await release.promise;
+          return { action: "accept" };
+        },
+      });
+      const creating = createTargetThread(harness, {
+        hostId: host.id,
+        projectId: project.id,
+      });
+      const rejected = expect(creating).rejects.toThrow(/remov/i);
+      try {
+        await entered.promise;
+        updateHost(harness.db, harness.hub, host.id, { phase: "removing" });
+      } finally {
+        release.resolve();
+      }
+      await rejected;
+      expect(provision).not.toHaveBeenCalled();
+      expect(
+        listEnvironments(harness.db, { hostId: host.id }).filter(
+          (row) => row.ownerThreadId !== null,
+        ),
+      ).toEqual([]);
+    });
+  });
+});
+
 describe("environment providers are asked inside provisioning", () => {
   it("refuses placement on a machine being removed", async () => {
     await withTestHarness(async (harness) => {