From 38680969859ceef4620179a58c5718ece4614e73 Mon Sep 17 00:00:00 2001 From: Michael Yong Date: Fri, 11 Sep 2026 15:52:49 -0700 Subject: [PATCH] Retain shared machines for pending workspace preparation --- .../machines/provider-orchestration.ts | 2 + .../src/services/threads/thread-create.ts | 11 +- .../threads/thread-environment-placement.ts | 28 ++- .../machines/provider-orchestration.test.ts | 35 +++ .../threads/environment-providers.test.ts | 207 ++++++++++++++++++ packages/db/src/data/machines.ts | 56 ++++- packages/db/test/data/machines.test.ts | 179 ++++++++++++++- 7 files changed, 501 insertions(+), 17 deletions(-) diff --git a/apps/server/src/services/machines/provider-orchestration.ts b/apps/server/src/services/machines/provider-orchestration.ts index 45e8b793f7..ca79079a69 100644 --- a/apps/server/src/services/machines/provider-orchestration.ts +++ b/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) ) { diff --git a/apps/server/src/services/threads/thread-create.ts b/apps/server/src/services/threads/thread-create.ts index 5f0f2ed33e..eb0b7a2ea9 100644 --- a/apps/server/src/services/threads/thread-create.ts +++ b/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 }), diff --git a/apps/server/src/services/threads/thread-environment-placement.ts b/apps/server/src/services/threads/thread-environment-placement.ts index 07cacc2e6b..0685104ad2 100644 --- a/apps/server/src/services/threads/thread-environment-placement.ts +++ b/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 = diff --git a/apps/server/test/services/machines/provider-orchestration.test.ts b/apps/server/test/services/machines/provider-orchestration.test.ts index e58d4fb589..dc6046adfb 100644 --- a/apps/server/test/services/machines/provider-orchestration.test.ts +++ b/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(); diff --git a/apps/server/test/threads/environment-providers.test.ts b/apps/server/test/threads/environment-providers.test.ts index 06ef83be01..f25fd1646c 100644 --- a/apps/server/test/threads/environment-providers.test.ts +++ b/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(); + const release = createDeferredPromise(); + 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(); + const release = createDeferredPromise(); + 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) => { diff --git a/packages/db/src/data/machines.ts b/packages/db/src/data/machines.ts index 976f18eb83..42f75692d7 100644 --- a/packages/db/src/data/machines.ts +++ b/packages/db/src/data/machines.ts @@ -1,4 +1,4 @@ -import { and, eq, inArray, isNull, or } from "drizzle-orm"; +import { and, eq, inArray, isNull, or, sql } from "drizzle-orm"; import type { DbConnection, DbTransaction } from "../connection.js"; import { environments, hosts, threads } from "../schema.js"; @@ -51,6 +51,60 @@ export function machineHasLiveThreadLaunch( ); } +export function machineHasPendingThreads( + db: Connection, + hostId: string, +): boolean { + const intent = sql`case json_extract(${threads.startupContext}, '$.kind') + when 'pending' then json_extract(${threads.startupContext}, '$.environmentIntent') + when 'provisioning' then json_extract(${threads.startupContext}, '$.request.environmentIntent') + end`; + return ( + db + .select({ id: threads.id }) + .from(threads) + .leftJoin( + environments, + or( + and( + eq(environments.ownerThreadId, threads.id), + eq(threads.status, "starting"), + ), + and( + sql`json_extract(${intent}, '$.type') = 'reuse'`, + eq( + environments.id, + sql`json_extract(${intent}, '$.environmentId')`, + ), + ), + ), + ) + .where( + and( + inArray(threads.status, ["pending", "starting"]), + isNull(threads.archivedAt), + isNull(threads.deletedAt), + or( + and( + eq(environments.hostId, hostId), + inArray(environments.status, [ + "creating", + "provisioning", + "ready", + ]), + isNull(environments.teardownStatus), + ), + sql`json_extract(${intent}, '$.type') = 'provider' + and json_extract(${intent}, '$.machine.type') = 'existing' + and json_extract(${intent}, '$.machine.hostId') = ${hostId}`, + ), + ), + ) + .limit(1) + .get() !== undefined + ); +} + export function machineHasProvisioningEnvironment( db: Connection, hostId: string, diff --git a/packages/db/test/data/machines.test.ts b/packages/db/test/data/machines.test.ts index f5abc595f0..2b1564421c 100644 --- a/packages/db/test/data/machines.test.ts +++ b/packages/db/test/data/machines.test.ts @@ -1,9 +1,11 @@ import { describe, expect, it } from "vitest"; +import { createConnection } from "../../src/connection.js"; import { eq } from "drizzle-orm"; import { createEnvironment } from "../../src/data/environments.js"; import { updateHost, upsertHost } from "../../src/data/hosts.js"; import { machineHasLiveThreadLaunch, + machineHasPendingThreads, machineHasProvisioningEnvironment, machineHasStartingThreadLaunch, } from "../../src/data/machines.js"; @@ -36,10 +38,7 @@ describe("machine provisioning state", () => { expect(machineHasLiveThreadLaunch(db, host.id)).toBe(true); expect(machineHasStartingThreadLaunch(db, host.id)).toBe(true); for (const status of ["active", "idle", "error"] as const) { - db.update(threads) - .set({ status }) - .where(eq(threads.id, thread.id)) - .run(); + db.update(threads).set({ status }).where(eq(threads.id, thread.id)).run(); expect(machineHasStartingThreadLaunch(db, host.id)).toBe(false); expect(machineHasLiveThreadLaunch(db, host.id)).toBe(true); } @@ -67,3 +66,175 @@ describe("machine provisioning state", () => { expect(machineHasProvisioningEnvironment(db, host.id)).toBe(false); }); }); + +describe("pending machine ownership", () => { + it.each(["pending", "provisioning"] as const)( + "retains persisted %s placement after reloading the database", + (kind) => { + const { db, host, project } = setup(); + try { + const thread = createThread(db, noopNotifier, { + projectId: project.id, + providerId: "codex", + status: kind === "pending" ? "pending" : "starting", + }); + const environmentIntent = { + type: "provider", + machine: { type: "existing", hostId: host.id }, + }; + db.update(threads) + .set({ + startupContext: JSON.stringify( + kind === "pending" + ? { kind, environmentIntent } + : { kind, request: { environmentIntent } }, + ), + }) + .where(eq(threads.id, thread.id)) + .run(); + const restored = createConnection(db.$client.serialize()); + try { + expect(machineHasPendingThreads(restored, host.id)).toBe(true); + expect(machineHasPendingThreads(restored, "different-host")).toBe( + false, + ); + for (const patch of [ + { status: "error" }, + { status: "idle" }, + { archivedAt: 1 }, + { deletedAt: 1 }, + ] as const) { + restored + .update(threads) + .set({ + status: kind === "pending" ? "pending" : "starting", + archivedAt: null, + deletedAt: null, + ...patch, + }) + .where(eq(threads.id, thread.id)) + .run(); + expect(machineHasPendingThreads(restored, host.id)).toBe(false); + } + } finally { + restored.$client.close(); + } + } finally { + db.$client.close(); + } + }, + ); + + it("retains a preparing owner through ready-before-attachment but releases orphan or cancelled preparation", () => { + const { db, host, project } = setup(); + try { + const thread = createThread(db, noopNotifier, { + projectId: project.id, + providerId: "codex", + status: "starting", + }); + const environment = createEnvironment(db, noopNotifier, { + projectId: project.id, + hostId: host.id, + path: "/tmp/preparing", + status: "provisioning", + providerOwnsPath: false, + environmentProvider: null, + }); + db.update(environments) + .set({ ownerThreadId: thread.id }) + .where(eq(environments.id, environment.id)) + .run(); + for (const status of ["creating", "provisioning", "ready"] as const) { + db.update(environments) + .set({ status }) + .where(eq(environments.id, environment.id)) + .run(); + expect(machineHasPendingThreads(db, host.id)).toBe(true); + } + db.update(environments) + .set({ teardownStatus: "running" }) + .where(eq(environments.id, environment.id)) + .run(); + expect(machineHasPendingThreads(db, host.id)).toBe(false); + db.update(environments) + .set({ teardownStatus: null }) + .where(eq(environments.id, environment.id)) + .run(); + for (const patch of [ + { status: "pending" }, + { status: "error" }, + { status: "idle" }, + { archivedAt: 1 }, + { deletedAt: 1 }, + ] as const) { + db.update(threads) + .set({ + status: "starting", + archivedAt: null, + deletedAt: null, + ...patch, + }) + .where(eq(threads.id, thread.id)) + .run(); + expect(machineHasPendingThreads(db, host.id)).toBe(false); + } + db.delete(threads).where(eq(threads.id, thread.id)).run(); + expect(machineHasPendingThreads(db, host.id)).toBe(false); + } finally { + db.$client.close(); + } + }); + + it("finds an unattached pending reuse intent without retaining failed or destroyed workspaces", () => { + const { db, host, project } = setup(); + try { + const environment = createEnvironment(db, noopNotifier, { + projectId: project.id, + hostId: host.id, + path: "/tmp/reuse", + status: "ready", + providerOwnsPath: false, + environmentProvider: null, + }); + const thread = createThread(db, noopNotifier, { + projectId: project.id, + providerId: "codex", + status: "pending", + }); + db.update(threads) + .set({ + startupContext: JSON.stringify({ + kind: "pending", + environmentIntent: { type: "reuse", environmentId: environment.id }, + }), + }) + .where(eq(threads.id, thread.id)) + .run(); + expect(machineHasPendingThreads(db, host.id)).toBe(true); + const oldHost = upsertHost(db, noopNotifier, { name: "previous-host" }); + const oldEnvironment = createEnvironment(db, noopNotifier, { + projectId: project.id, + hostId: oldHost.id, + path: "/tmp/old", + status: "ready", + providerOwnsPath: false, + environmentProvider: null, + }); + db.update(environments) + .set({ ownerThreadId: thread.id }) + .where(eq(environments.id, oldEnvironment.id)) + .run(); + expect(machineHasPendingThreads(db, oldHost.id)).toBe(false); + for (const status of ["error", "destroyed"] as const) { + db.update(environments) + .set({ status }) + .where(eq(environments.id, environment.id)) + .run(); + expect(machineHasPendingThreads(db, host.id)).toBe(false); + } + } finally { + db.$client.close(); + } + }); +});