diff --git a/doc/DATABASE.md b/doc/DATABASE.md index 46a104d61a6..e87e2067315 100644 --- a/doc/DATABASE.md +++ b/doc/DATABASE.md @@ -795,6 +795,59 @@ blocks concurrent task or run updates. It allows audit inserts to retain their foreign-key `KEY SHARE` locks without waiting on identity acquisition. The audit foreign keys and their deletion behavior remain enforced. +## Remote workspace row locks + +The runtime workspace service takes company `KEY SHARE`, current task `UPDATE`, +run `UPDATE`, then workspace `UPDATE` locks. It rechecks the run's task binding +under the run lock. A shared workspace can retain a different original source +task; new provisioning and snapshot operations belong to the current run's task. +Existing operation receipts keep their recorded attribution. + +| Owning transaction | Locks after company, current task and run | +| --- | --- | +| Workspace reservation | Shared-selection advisory lock when needed, workspace | +| Preparation publication | Workspace, source lease, existing restore receipt | +| Snapshot reservation and publication | Workspace, operation and snapshot rows | +| Provisioning reservation | Workspace, existing operation or new receipt | + +`workspace_operations` has ordinary company, task, run and workspace foreign +keys. An insert takes `KEY SHARE` locks on those parents even when the service +does not explicitly select them. Locking the parents first prevents their +foreign-key checks, or a later run-reference write, from reversing the order. +File preservation already takes task, run, workspace and lease locks in that +order. Confirmed lease destruction locks matching workspaces before the lease; +its workspace-location update takes no task or run lock afterward. + +Provisioning settlement and save-invocation settlement lock the workspace and +update existing receipt metadata. They neither add a parent foreign-key +reference nor lock a task or run afterward. Snapshot upload and retention own +snapshot rows; the snapshot ledger has no parent foreign keys. These paths do +not need an otherwise unused parent lock. + +Administrative import reservation takes company `KEY SHARE`, original source +task `UPDATE`, then workspace `UPDATE`, and rechecks that source binding. This +also applies to an already-imported no-op, which can overlap runtime saves. +Import publication takes the same company and recorded source-task locks before +its legacy coordinator and run locks, then the workspace and saved allocation +leases. Native claim, status and cancellation owners do not share one row-lock +order. Import therefore requests those legacy coordinator, run and lease locks +with `NOWAIT`. A real `55P03` aborts the publication transaction and reports +`remote_workspace_migration_blocked:execution_rows_busy`. The journal, uploaded +archive and local files remain available for a later administrative retry. Other +database errors propagate. No automatic retry or lock-strength reduction occurs. + +The workspace fingerprint pins the recorded source. An arriving reservation can +reach a legacy workspace before refusing it, so publication must follow this +order too. The lifecycle advisory lock alone does not fence runtime operations. +Publication still rechecks current writers and unresolved work. Offline imports +also require the deployment-stopped maintenance fence. Ordinary save requires an +existing prepared remote location and refuses the unpublished legacy authority. + +The workspace lock still serializes generations and publication. Foreign keys, +ownership and epoch checks, and transaction-aware work notifications remain in +place. Remote preparation, command execution, archive capture and object upload +run outside these transactions. No workspace or parent row lock spans that I/O. + ## Execution admission accounting The execution-admission module (`server/src/modules/execution-admission/`) owns diff --git a/server/src/services/remote-workspace-migration-concurrency.test.ts b/server/src/services/remote-workspace-migration-concurrency.test.ts new file mode 100644 index 00000000000..d2fa8bbf097 --- /dev/null +++ b/server/src/services/remote-workspace-migration-concurrency.test.ts @@ -0,0 +1,188 @@ +import { randomUUID } from "node:crypto"; +import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { eq, sql } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import { agents, closeRegisteredClients, companies, createDb, environmentLeases, executionWorkspaceSnapshots, + executionWorkspaces, heartbeatRuns, issues, nativeRunFinalizations, workspaceOperations } from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; +import { createStorageService } from "../storage/service.js"; +import { createLocalDiskStorageProvider } from "../storage/local-disk-provider.js"; +import { migrateRemoteWorkspace } from "./remote-workspace-migration.js"; +import { commitNativeStatusDecision } from "./native-runtime/status-decision-committer.js"; +import { NATIVE_STATUS_ARBITER_POLICY_VERSION } from "./native-runtime/status-arbiter.js"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise(ready => { resolve = ready; }); + return { promise, resolve }; +} + +const support = await getEmbeddedPostgresTestSupport(); +describe.skipIf(!support.supported)("legacy workspace publication ownership contention", () => { + let database: Awaited>; + let db: ReturnType, root: string; + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("workspace-import-order-"); + db = createDb(database.connectionString); + root = await mkdtemp(path.join(os.tmpdir(), "workspace-import-order-")); + vi.stubEnv("PAPERCLIP_SECRETS_MASTER_KEY", "62".repeat(32)); + vi.stubEnv("PAPERCLIP_EPHEMERAL_MIGRATION", "true"); + }, 60_000); + afterAll(async () => { + if (database) await closeRegisteredClients(database.connectionString); + await database?.cleanup(); + if (root) await rm(root, { recursive: true, force: true }); + vi.unstubAllEnvs(); + }); + + async function fixture() { + const companyId = randomUUID(), runId = randomUUID(), workspaceId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Legacy publication", issuePrefix: `L${companyId.slice(0, 6)}` }); + const [agent] = await db.insert(agents).values({ companyId, name: "Historical worker", adapterType: "paperclip_runner" }).returning(); + const [issue] = await db.insert(issues).values({ companyId, title: "Retain local work", status: "blocked" }).returning(); + const cwd = path.join(root, workspaceId); await mkdir(cwd); + await writeFile(path.join(cwd, "saved.txt"), "original local work\n"); + await db.insert(executionWorkspaces).values({ id: workspaceId, companyId, sourceIssueId: issue.id, + mode: "isolated_workspace", strategyType: "project_primary", providerType: "local_fs", name: "Legacy files", cwd, providerRef: cwd }); + await db.update(issues).set({ executionWorkspaceId: workspaceId }).where(eq(issues.id, issue.id)); + const reference = { schema: "paperclip.native-workspace-sync/v3", state: "prepared", workspaceId, + leaseId: randomUUID(), providerLeaseId: randomUUID(), remoteCwd: "/workspaces/task", + descriptorSha256: "12".repeat(32), baselineSha256: "34".repeat(32), finalHostSha256: null, resourceDisposition: "stop_and_retain" }; + await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId: agent.id, nativeIssueId: issue.id, + runtimeMode: "native", status: "failed", finishedAt: new Date(), runnerProfileJson: { nativeWorkspaceSync: reference } }); + await db.insert(nativeRunFinalizations).values({ companyId, issueId: issue.id, runId, phase: "terminal_failure" }); + const storage = createStorageService(createLocalDiskStorageProvider(path.join(root, `${workspaceId}-objects`))); + return { companyId, workspaceId, runId, issue, cwd, reference, storage, + expectedProfile: { nativeWorkspaceSync: reference } as Record, expectedPhase: "terminal_failure" }; + } + + async function retained(f: Awaited>) { + const [workspace] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, f.workspaceId)); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + const [coordinator] = await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + expect(workspace.metadata?.remoteWorkspace).toBeUndefined(); + expect(workspace.metadata?.remoteWorkspaceMigration).toMatchObject({ state: "running" }); + expect(run.runnerProfileJson).toEqual(f.expectedProfile); + expect(coordinator).toMatchObject({ phase: f.expectedPhase, leaseOwner: null, leaseExpiresAt: null }); + expect(await readFile(path.join(f.cwd, "saved.txt"), "utf8")).toBe("original local work\n"); + const snapshots = await db.select().from(executionWorkspaceSnapshots).where(eq(executionWorkspaceSnapshots.executionWorkspaceId, f.workspaceId)); + expect(snapshots).toEqual([expect.objectContaining({ status: "uploaded" })]); + expect((await f.storage.headObject(f.companyId, snapshots[0].objectKey!)).exists).toBe(true); + expect(await db.select().from(workspaceOperations).where(eq(workspaceOperations.executionWorkspaceId, f.workspaceId))) + .toEqual([expect.objectContaining({ phase: "remote_workspace_migration", status: "running" })]); + } + + it("refuses a busy publication beside an actual stale native status commit", async () => { + const f = await fixture(); + const migrationDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "import-native-migration" }); + const nativeDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "import-native-status" }); + await migrationDb.execute(sql`select set_config('deadlock_timeout', '50ms', false)`); + const uploadReady = deferred(), uploadRelease = deferred(), lockReady = deferred(), lockRelease = deferred(); + const storage = { ...f.storage, putFile: async (input: Parameters[0]) => { + const saved = await f.storage.putFile(input); uploadReady.resolve(); await uploadRelease.promise; return saved; + } }; + let migrationSettled = false; + const migrating = migrateRemoteWorkspace({ db: migrationDb, storage, companyId: f.companyId, + workspaceId: f.workspaceId, apply: true, writersStopped: true }) + .then(value => ({ value }), error => ({ error })).finally(() => { migrationSettled = true; }); + let barrier: Promise | undefined, committing: Promise | undefined; + try { + await uploadReady.promise; + barrier = db.transaction(async tx => { + await tx.select({ id: nativeRunFinalizations.runId }).from(nativeRunFinalizations) + .where(eq(nativeRunFinalizations.runId, f.runId)).for("update"); + lockReady.resolve(); await lockRelease.promise; + }); + await lockReady.promise; + committing = commitNativeStatusDecision({ db: nativeDb, companyId: f.companyId, issueId: f.issue.id, + runId: f.runId, assessmentId: randomUUID(), priorStatus: f.issue.status, + priorStatusVersion: f.issue.statusVersion + 1, priorDecisionId: null, + decision: { policyVersion: NATIVE_STATUS_ARBITER_POLICY_VERSION, statusAction: "preserve", toStatus: "blocked", + reasonCode: "stale_status_control", unblockDescriptor: null, effects: [] } }) + .then(value => ({ value }), error => ({ error })); + const waiting = async (application: string) => (await db.execute<{ query: string }>(sql`select query from pg_stat_activity + where datname = current_database() and application_name = ${application} and wait_event_type = 'Lock'`))[0]; + await expect.poll(async () => (await waiting("import-native-status"))?.query, { timeout: 5000 }).toContain('"native_run_finalizations"'); + uploadRelease.resolve(); + await expect.poll(async () => migrationSettled || Boolean(await waiting("import-native-migration")), { timeout: 5000 }).toBe(true); + const migrationWait = await waiting("import-native-migration"); + lockRelease.resolve(); await barrier; + const results = await Promise.all([migrating, committing]); + const codes = results.map(result => (result as { error?: { cause?: { code?: string } } }).error?.cause?.code); + if (codes.includes("40P01")) console.info("import-native-postgres-evidence", JSON.stringify({ migrationQuery: migrationWait?.query, + errors: results.map(result => (result as { error?: { cause?: unknown } }).error?.cause) })); + expect(codes).not.toContain("40P01"); + expect(results[0]).toMatchObject({ error: expect.objectContaining({ message: "remote_workspace_migration_blocked:execution_rows_busy" }) }); + expect(results[1]).toMatchObject({ error: expect.objectContaining({ code: "native_status_race" }) }); + await retained(f); + await expect(migrateRemoteWorkspace({ db, storage: f.storage, companyId: f.companyId, + workspaceId: f.workspaceId, apply: true, writersStopped: true })).resolves.toMatchObject({ applied: true }); + expect(await db.select().from(executionWorkspaceSnapshots).where(eq(executionWorkspaceSnapshots.executionWorkspaceId, f.workspaceId))) + .toEqual([expect.objectContaining({ status: "committed" })]); + } finally { + uploadRelease.resolve(); lockRelease.resolve(); await barrier; await Promise.all([migrating, committing]); + await migrationDb.$client.end(); await nativeDb.$client.end(); + } + }, 30_000); + + it.each(["run", "lease"] as const)("holds publication while the saved %s is busy", async kind => { + const f = await fixture(); + if (kind === "lease") { + await db.insert(environmentLeases).values({ id: f.reference.leaseId, companyId: f.companyId, heartbeatRunId: f.runId, + issueId: f.issue.id, provider: "daytona", providerLeaseId: f.reference.providerLeaseId, + status: "released", metadata: { remoteCwd: f.reference.remoteCwd } }); + f.expectedProfile = {}; f.expectedPhase = "observed"; + await db.update(heartbeatRuns).set({ runnerProfileJson: {}, resultJson: { workspaceRestoreRecovery: { + schema: "paperclip.workspace-restore-recovery.v1", leaseIds: [f.reference.leaseId] } } }).where(eq(heartbeatRuns.id, f.runId)); + await db.update(nativeRunFinalizations).set({ phase: "observed" }).where(eq(nativeRunFinalizations.runId, f.runId)); + } + const ready = deferred(), release = deferred(); + const barrier = db.transaction(async tx => { + if (kind === "run") await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)).for("update"); + else await tx.select({ id: environmentLeases.id }).from(environmentLeases).where(eq(environmentLeases.id, f.reference.leaseId)).for("update"); + ready.resolve(); await release.promise; + }); + await ready.promise; + let settled = false; + const migrating = migrateRemoteWorkspace({ db, storage: f.storage, companyId: f.companyId, + workspaceId: f.workspaceId, apply: true, writersStopped: true }) + .then(value => ({ value }), error => ({ error })).finally(() => { settled = true; }); + try { + await expect.poll(() => settled, { timeout: 5000 }).toBe(true); + expect(await migrating).toMatchObject({ error: expect.objectContaining({ message: "remote_workspace_migration_blocked:execution_rows_busy" }) }); + await retained(f); + release.resolve(); await barrier; + await expect(migrateRemoteWorkspace({ db, storage: f.storage, companyId: f.companyId, + workspaceId: f.workspaceId, apply: true, writersStopped: true })).resolves.toMatchObject({ applied: true }); + } finally { release.resolve(); await barrier; await migrating; } + }, 30_000); + + it("rolls back publication writes and preserves an unexpected database error", async () => { + const f = await fixture(); + await db.update(nativeRunFinalizations).set({ phase: "observed", leaseOwner: "expired-controller", + leaseExpiresAt: new Date(Date.now() - 60_000) }).where(eq(nativeRunFinalizations.runId, f.runId)); + const [before] = await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + await db.execute(sql.raw(`create function import_publication_failure() returns trigger language plpgsql as $$ + begin raise exception 'publication constraint failed' using errcode = '23514'; end $$`)); + await db.execute(sql.raw(`create trigger import_publication_failure before update on execution_workspaces for each row + when (new.id = '${f.workspaceId}'::uuid and new.metadata ? 'remoteWorkspace') execute function import_publication_failure()`)); + try { + await expect(migrateRemoteWorkspace({ db, storage: f.storage, companyId: f.companyId, + workspaceId: f.workspaceId, apply: true, writersStopped: true })).rejects.toMatchObject({ cause: { code: "23514" } }); + const [workspace] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, f.workspaceId)); + const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + const [coordinator] = await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, f.runId)); + expect(workspace.metadata?.remoteWorkspace).toBeUndefined(); + expect(run.runnerProfileJson).toEqual(f.expectedProfile); + expect(coordinator).toEqual(before); + expect(await db.select().from(executionWorkspaceSnapshots).where(eq(executionWorkspaceSnapshots.executionWorkspaceId, f.workspaceId))) + .toEqual([expect.objectContaining({ status: "uploaded" })]); + expect(await readFile(path.join(f.cwd, "saved.txt"), "utf8")).toBe("original local work\n"); + } finally { + await db.execute(sql`drop trigger import_publication_failure on execution_workspaces`); + await db.execute(sql`drop function import_publication_failure()`); + } + }, 30_000); +}); diff --git a/server/src/services/remote-workspace-migration.ts b/server/src/services/remote-workspace-migration.ts index 9ffc19581b4..3b0e4448001 100644 --- a/server/src/services/remote-workspace-migration.ts +++ b/server/src/services/remote-workspace-migration.ts @@ -4,7 +4,7 @@ import os from "node:os"; import path from "node:path"; import { and, eq, inArray, ne, or, sql } from "drizzle-orm"; import { - environmentLeases, executionWorkspaceRepositories, executionWorkspaceRuntimeLeases, executionWorkspaceSnapshots, executionWorkspaces, + companies, environmentLeases, executionWorkspaceRepositories, executionWorkspaceRuntimeLeases, executionWorkspaceSnapshots, executionWorkspaces, heartbeatRuns, issueRecoveryActions, issues, nativeRunFinalizations, projects, workspaceOperations, workspaceRuntimeServices, type Db, } from "@paperclipai/db"; import { readRemoteWorkspaceAuthority, remoteWorkspaceSourceSchema, type RemoteWorkspaceAuthority, type RemoteWorkspaceSource } from "@paperclipai/shared"; @@ -27,6 +27,12 @@ import { inspectNativeWorkspaceMigration, nativeWorkspaceMigrationMarker } from const SCHEMA = "paperclip.remote-workspace-migration.v1"; type Workspace = typeof executionWorkspaces.$inferSelect; type ReadDb = Pick; +async function lockMigrationParents(db: ReadDb, companyId: string, issueId: string | null) { + await db.select({ id: companies.id }).from(companies).where(eq(companies.id, companyId)).for("key share"); + if (issueId) await db.select({ id: issues.id }).from(issues).where(and( + eq(issues.companyId, companyId), eq(issues.id, issueId))).for("update"); +} + const object = (value: unknown): Record => value && typeof value === "object" && !Array.isArray(value) ? value as Record : {}; function canonical(value: unknown): string { if (Array.isArray(value)) return `[${value.map(canonical).join(",")}]`; @@ -229,9 +235,13 @@ export async function migrateRemoteWorkspace(input: { if (!input.apply) return { ...await inventoryRemoteWorkspaceMigration(input.db, input), applied: false }; return withExecutionWorkspaceLifecycleLock(input.db, input.workspaceId, async lock => { const reserved = await input.db.transaction(async tx => { + const [selected] = await tx.select({ sourceIssueId: executionWorkspaces.sourceIssueId }).from(executionWorkspaces) + .where(and(eq(executionWorkspaces.companyId, input.companyId), eq(executionWorkspaces.id, input.workspaceId))); + if (!selected) throw new Error("remote_workspace_migration_workspace_missing"); + await lockMigrationParents(tx, input.companyId, selected.sourceIssueId); const [workspace] = await tx.select().from(executionWorkspaces).where(and(eq(executionWorkspaces.companyId, input.companyId), eq(executionWorkspaces.id, input.workspaceId))).for("update"); if (!workspace) throw new Error("remote_workspace_migration_workspace_missing"); - if (workspace.sourceIssueId) await tx.select({ id: issues.id }).from(issues).where(and(eq(issues.companyId, workspace.companyId), eq(issues.id, workspace.sourceIssueId))).for("update"); + if (workspace.sourceIssueId !== selected.sourceIssueId) throw new Error("remote_workspace_migration_source_changed"); const current = await inspect(tx as unknown as Db, input); if (current.authority) return { workspace, authority: current.authority, operation: null, source: current.source, sourceFingerprint: "", issueBinding: current.issueBinding, retainedDirectory: null, archivedGitSource: null, legacyRecoverySources: [] as LegacyWorkspaceRecoverySource[] }; if (current.reasons.length) throw new Error(`remote_workspace_migration_blocked:${current.reasons.join(",")}`); @@ -319,14 +329,12 @@ export async function migrateRemoteWorkspace(input: { input.signal?.throwIfAborted(); await lock.execute(sql`select 1`); const publication = await input.db.transaction(async tx => { + await lockMigrationParents(tx, input.companyId, reserved.workspace.sourceIssueId); for (const source of reserved.legacyRecoverySources) { await tx.select({ id: nativeRunFinalizations.runId }).from(nativeRunFinalizations) - .where(and(eq(nativeRunFinalizations.companyId, input.companyId), eq(nativeRunFinalizations.runId, source.runId))).for("update"); + .where(and(eq(nativeRunFinalizations.companyId, input.companyId), eq(nativeRunFinalizations.runId, source.runId))).for("update", { noWait: true }); await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns) - .where(and(eq(heartbeatRuns.companyId, input.companyId), eq(heartbeatRuns.id, source.runId))).for("update"); - for (const allocation of source.recoveryAllocations ?? []) await tx.select({ id: environmentLeases.id }) - .from(environmentLeases).where(and(eq(environmentLeases.companyId, input.companyId), - eq(environmentLeases.id, allocation.leaseId))).for("update"); + .where(and(eq(heartbeatRuns.companyId, input.companyId), eq(heartbeatRuns.id, source.runId))).for("update", { noWait: true }); } const [workspace] = await tx.select().from(executionWorkspaces).where(and(eq(executionWorkspaces.companyId, input.companyId), eq(executionWorkspaces.id, input.workspaceId))).for("update"); if (!workspace) throw new Error("remote_workspace_migration_workspace_missing"); @@ -334,7 +342,11 @@ export async function migrateRemoteWorkspace(input: { if (existing?.snapshotId === ref.id) return null; if (existing || fingerprint(workspace) !== reserved.sourceFingerprint || object(workspace.metadata?.remoteWorkspaceMigration).operationId !== operation.id) throw new Error("remote_workspace_migration_publication_changed"); - if (workspace.sourceIssueId) await tx.select({ id: issues.id }).from(issues).where(eq(issues.id, workspace.sourceIssueId)).for("update"); + for (const source of reserved.legacyRecoverySources) { + for (const allocation of source.recoveryAllocations ?? []) await tx.select({ id: environmentLeases.id }) + .from(environmentLeases).where(and(eq(environmentLeases.companyId, input.companyId), + eq(environmentLeases.id, allocation.leaseId))).for("update", { noWait: true }); + } const validation = await inspect(tx as unknown as Db, { ...input, operationId: operation.id }); if (validation.issueBinding !== reserved.issueBinding) throw new Error("remote_workspace_migration_source_changed"); if (Boolean(validation.retainedDirectory) !== Boolean(reserved.retainedDirectory)) throw new Error("remote_workspace_migration_source_changed"); @@ -395,6 +407,10 @@ export async function migrateRemoteWorkspace(input: { return (await persistActivity(tx as unknown as Db, { companyId: workspace.companyId, actorType: "system", actorId: "workspace-migration", action: "workspace.remote_storage_migrated", entityType: "execution_workspace", entityId: workspace.id, details: { snapshotId: ref.id, operationId: operation.id, manifestDigest: ref.manifestDigest } })).publication; + }).catch(error => { + const failure = object(error), cause = object(failure.cause); + if ((cause.code ?? failure.code) !== "55P03") throw error; + throw new Error("remote_workspace_migration_blocked:execution_rows_busy", { cause: error }); }); if (publication) publishActivity(publication); return { ...await inventoryRemoteWorkspaceMigration(input.db, input), applied: true, replayed: false }; diff --git a/server/src/services/remote-workspace-save-concurrency.test.ts b/server/src/services/remote-workspace-save-concurrency.test.ts new file mode 100644 index 00000000000..9ecf9f04d3b --- /dev/null +++ b/server/src/services/remote-workspace-save-concurrency.test.ts @@ -0,0 +1,489 @@ +import { execFile } from "node:child_process"; +import { randomUUID } from "node:crypto"; +import { cp, mkdir, mkdtemp, readFile, realpath, rm, writeFile } from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { promisify } from "node:util"; +import { and, eq, sql } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import { + agents, closeRegisteredClients, companies, createDb, environmentLeases, environments, + executionWorkspaces, executionWorkspaceSnapshots, heartbeatRuns, issues, workspaceOperations, +} from "@paperclipai/db"; +import type { AdapterSandboxExecutionTarget } from "@paperclipai/adapter-utils/execution-target"; +import type { EnvironmentLease } from "@paperclipai/shared"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; +import { createStorageService } from "../storage/service.js"; +import { createLocalDiskStorageProvider } from "../storage/local-disk-provider.js"; +import { preserveRemoteWorkspacePreparationFailure } from "./native-runtime/native-workspace-finalization-state.js"; +import { + createSavedWorkspaceLoader, prepareRemoteWorkspace, readRemoteWorkspaceAuthority, + provisionRemoteWorkspace, reserveRemoteWorkspace, saveRemoteWorkspace, +} from "./remote-workspaces.js"; +import { remoteTerminationReceipt } from "./remote-execution-termination.js"; +import { migrateRemoteWorkspace } from "./remote-workspace-migration.js"; + +const exec = promisify(execFile); +function deferred() { + let resolve!: () => void; + const promise = new Promise(ready => { resolve = ready; }); + return { promise, resolve }; +} + +const support = await getEmbeddedPostgresTestSupport(); +describe.skipIf(!support.supported)("remote workspace save parent lock order", () => { + let database: Awaited>; + let db: ReturnType; + let root: string; + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("workspace-save-order-"); + db = createDb(database.connectionString); + root = await realpath(await mkdtemp(path.join(os.tmpdir(), "workspace-save-order-"))); + vi.stubEnv("PAPERCLIP_SECRETS_MASTER_KEY", "62".repeat(32)); + vi.stubEnv("DATABASE_URL", database.connectionString); + }, 60_000); + afterAll(async () => { + if (database) await closeRegisteredClients(database.connectionString); + await database?.cleanup(); + if (root) await rm(root, { recursive: true, force: true }); + vi.unstubAllEnvs(); + }); + + async function preparedFailure() { + const companyId = randomUUID(), agentId = randomUUID(), issueId = randomUUID(), runId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Save lock order", issuePrefix: `S${companyId.slice(0, 6)}` }); + await db.insert(agents).values({ id: agentId, companyId, name: "Worker", adapterType: "paperclip_runner" }); + await db.insert(issues).values({ id: issueId, companyId, title: "Preserve saved work", status: "blocked", assigneeAgentId: agentId }); + await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, issueId, runtimeMode: "legacy", + status: "running", executionStage: "preparing", contextSnapshot: { issueId } }); + const [environment] = await db.insert(environments).values({ name: runId, driver: "sandbox" }).returning(); + const [lease] = await db.insert(environmentLeases).values({ companyId, issueId, heartbeatRunId: runId, + environmentId: environment.id, provider: "daytona", providerLeaseId: randomUUID(), status: "active" }).returning(); + const storage = { ...createStorageService(createLocalDiskStorageProvider(path.join(root, runId, "objects"))), provider: "s3" as const }; + const remoteCwd = path.join(root, runId, "remote"); + await mkdir(remoteCwd, { recursive: true }); + const transfer: NonNullable["syncIn"]> = async operations => { + for (const operation of operations) for (const file of operation.files) { + await mkdir(path.dirname(file.targetPath), { recursive: true }); + await cp(file.sourcePath, file.targetPath, { recursive: file.kind === "directory" }); + } + return { operations: operations.map(operation => ({ operationId: operation.operationId, + filesTransferred: operation.files.length, bytesTransferred: 0 })) }; + }; + const target: AdapterSandboxExecutionTarget = { + kind: "remote", transport: "sandbox", leaseId: lease.id, remoteCwd, timeoutMs: 30_000, + runner: { syncIn: transfer, syncOut: transfer, execute: async input => { + const result = await exec(input.command, input.args ?? [], { cwd: input.cwd ?? remoteCwd, + env: { ...process.env, ...input.env }, timeout: input.timeoutMs, maxBuffer: 8 * 1024 * 1024 }); + return { exitCode: 0, stdout: result.stdout, stderr: result.stderr, signal: null, timedOut: false, + pid: process.pid, startedAt: new Date().toISOString() }; + } }, + }; + const reserved = await reserveRemoteWorkspace({ db, storage, companyId, issueId, runId, agentId, + environmentId: environment.id, projectId: null, projectWorkspaceId: null, existingWorkspaceId: null, source: { kind: "empty" } }); + const prepared = await prepareRemoteWorkspace({ db, storage, workspace: reserved, runId, + lease: lease as EnvironmentLease, target }); + await writeFile(path.join(prepared.remoteRoot, "saved.txt"), "one preserved result\n"); + await db.update(heartbeatRuns).set({ status: "failed", finishedAt: new Date(), errorCode: "setup_failed", + resultJson: { workspaceRestoreFailure: "restore_failed", workspaceRestoreRecovery: { + schema: "paperclip.workspace-restore-recovery.v1", leaseIds: [lease.id], + } } }).where(eq(heartbeatRuns.id, runId)); + const [stoppedLease] = await db.update(environmentLeases).set({ status: "released", releasedAt: new Date(), cleanupStatus: "success", + metadata: { ...prepared.lease.metadata, remoteExecutionTermination: remoteTerminationReceipt(lease, + { providerLeaseId: lease.providerLeaseId, state: "stopped" }) } }).where(eq(environmentLeases.id, lease.id)).returning(); + const [workspace] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, reserved.id)); + const snapshots = await db.select().from(executionWorkspaceSnapshots).where(eq(executionWorkspaceSnapshots.executionWorkspaceId, workspace.id)); + return { companyId, agentId, issueId, runId, workspace, environmentId: environment.id, + lease: stoppedLease as EnvironmentLease, target, storage, snapshotCount: snapshots.length, + generation: Math.max(0, ...snapshots.map(snapshot => snapshot.generation)) }; + } + + it.each([ + { mode: "preserve", stage: "reservation" }, + { mode: "reserve", stage: "reservation" }, + { mode: "preserve", stage: "publication" }, + ] as const)("saves once beside actual $mode ($stage)", async ({ mode, stage }) => { + const f = await preparedFailure(); + const nextRunId = randomUUID(); + if (mode === "reserve") await db.insert(heartbeatRuns).values({ id: nextRunId, companyId: f.companyId, + agentId: f.agentId, issueId: f.issueId, runtimeMode: "legacy", status: "running", contextSnapshot: { issueId: f.issueId } }); + const saveDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "save-order-save" }); + const peerDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "save-order-preserve" }); + await saveDb.execute(sql`select set_config('deadlock_timeout', '50ms', false)`); + const barrierKey = 7540291; + await db.execute(sql.raw(`create function save_order_gate() returns trigger language plpgsql as $$ + begin perform pg_advisory_xact_lock(${barrierKey}); return new; end $$`)); + const triggerTable = stage === "reservation" ? "workspace_operations" : "execution_workspace_snapshots"; + await db.execute(sql.raw(stage === "reservation" + ? `create trigger save_order_gate before insert on workspace_operations for each row + when (new.heartbeat_run_id = '${f.runId}'::uuid and new.phase = 'remote_workspace_snapshot') execute function save_order_gate()` + : `create trigger save_order_gate before update on execution_workspace_snapshots for each row + when (new.source_run_id = '${f.runId}'::uuid and new.status = 'committed' and old.status = 'uploaded') execute function save_order_gate()`)); + const ready = deferred(), release = deferred(); + const barrier = db.transaction(async tx => { + await tx.execute(sql`select pg_advisory_xact_lock(${barrierKey})`); + ready.resolve(); await release.promise; + }); + await ready.promise; + let saving: Promise>>> | undefined; + let preserving: Promise> | undefined; + const reserveNext = () => reserveRemoteWorkspace({ db: peerDb, storage: f.storage, + companyId: f.companyId, issueId: f.issueId, runId: nextRunId, agentId: f.agentId, + environmentId: f.environmentId, projectId: null, projectWorkspaceId: null, + existingWorkspaceId: f.workspace.id, source: { kind: "empty" } }); + const settle = (work: Promise): Promise> => work.then( + value => ({ status: "fulfilled", value }), reason => ({ status: "rejected", reason })); + try { + saving = settle(saveRemoteWorkspace({ db: saveDb, storage: f.storage, workspace: f.workspace, + runId: f.runId, lease: f.lease, target: f.target, assertOwned: async () => undefined })); + const activity = async (application: string) => (await db.execute<{ pid: number; query: string; wait_event: string; blockers: number[] }>(sql` + select pid, query, wait_event, pg_blocking_pids(pid) as blockers from pg_stat_activity + where datname = current_database() and application_name = ${application} and wait_event_type = 'Lock'`))[0]; + await expect.poll(async () => (await activity("save-order-save"))?.wait_event, { timeout: 5000 }).toBe("advisory"); + const savePid = (await activity("save-order-save"))!.pid; + preserving = settle(mode === "preserve" ? preserveRemoteWorkspacePreparationFailure(peerDb, f.lease) : reserveNext()); + await expect.poll(async () => (await activity("save-order-preserve"))?.blockers.includes(savePid), { timeout: 5000 }).toBe(true); + const peer = await activity("save-order-preserve"); + release.resolve(); await barrier; + const results = await Promise.all([saving, preserving]); + const errors = results.filter((result, index) => result.status === "rejected" && (index === 0 || mode === "preserve")) + .map(result => (result as PromiseRejectedResult).reason as Error & { cause?: { code?: string; detail?: string; where?: string } }); + if (errors.length) console.info("save-order-postgres-evidence", JSON.stringify({ mode, stage, peerQuery: peer?.query, + errors: errors.map(error => ({ code: error.cause?.code, detail: error.cause?.detail, where: error.cause?.where })) })); + expect(errors).toEqual([]); + const saved = (results[0] as PromiseFulfilledResult>>).value; + const restarted = createDb(database.connectionString); + const [workspace] = await restarted.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, f.workspace.id)); + expect(readRemoteWorkspaceAuthority(workspace.metadata)).toMatchObject({ dirty: false, snapshotId: saved.snapshotId }); + const snapshots = await restarted.select().from(executionWorkspaceSnapshots).where(and(eq(executionWorkspaceSnapshots.companyId, f.companyId), + eq(executionWorkspaceSnapshots.sourceRunId, f.runId))); + expect(snapshots).toHaveLength(f.snapshotCount + 1); + expect(snapshots.filter(snapshot => snapshot.status === "committed")) + .toEqual([expect.objectContaining({ id: saved.snapshotId, generation: f.generation + 1 })]); + expect(await restarted.select().from(workspaceOperations).where(and(eq(workspaceOperations.companyId, f.companyId), + eq(workspaceOperations.heartbeatRunId, f.runId), eq(workspaceOperations.phase, "remote_workspace_snapshot"), + sql`${workspaceOperations.metadata}->>'boundary' = 'final'`))) + .toEqual([expect.objectContaining({ status: "succeeded" })]); + const loaded = await createSavedWorkspaceLoader(restarted, f.storage)({ companyId: f.companyId, + workspaceId: f.workspace.id, snapshotId: saved.snapshotId }); + try { expect(await readFile(path.join(loaded.root, "saved.txt"), "utf8")).toBe("one preserved result\n"); } + finally { await loaded.cleanup(); } + if (mode === "reserve") { + if (results[1].status === "rejected") expect(results[1].reason).toMatchObject({ details: { code: "remote_workspace_writer_pending" } }); + const admitted = await reserveNext(), repeated = await reserveNext(); + expect(readRemoteWorkspaceAuthority(admitted.metadata)).toMatchObject({ writerRunId: nextRunId, + dirty: false, snapshotId: saved.snapshotId }); + expect(readRemoteWorkspaceAuthority(repeated.metadata)).toEqual(readRemoteWorkspaceAuthority(admitted.metadata)); + } + } finally { + release.resolve(); await barrier; + await Promise.all([saving, preserving]); + await db.execute(sql`drop trigger save_order_gate on ${sql.identifier(triggerTable)}`); + await db.execute(sql`drop function save_order_gate()`); + await saveDb.$client.end(); await peerDb.$client.end(); + } + }, 30_000); + + it("releases database locks while a provisioning command is in flight", async () => { + const f = await preparedFailure(); + await db.update(heartbeatRuns).set({ status: "running", finishedAt: null }).where(eq(heartbeatRuns.id, f.runId)); + const authority = readRemoteWorkspaceAuthority(f.workspace.metadata)!; + const ready = deferred(), release = deferred(); + const execute = f.target.runner!.execute.bind(f.target.runner!); + const command = "printf 'one effect\\n' >> in-flight.txt"; + let calls = 0; + f.target.runner!.execute = async input => { + if (input.command === "bash" && input.args?.[1] === command) { + calls++; ready.resolve(); await release.promise; + } + return execute(input); + }; + const provision = () => provisionRemoteWorkspace({ db, companyId: f.companyId, workspaceId: f.workspace.id, + runId: f.runId, epoch: authority.epoch, remoteCwd: authority.location!.remoteCwd, target: f.target, commands: [command] }); + const provisioning = provision(); + try { + await ready.promise; + await saveRemoteWorkspace({ db, storage: f.storage, workspace: f.workspace, runId: f.runId, + lease: f.lease, target: f.target, boundary: "checkpoint", assertOwned: async () => undefined }); + expect(calls).toBe(1); + release.resolve(); await provisioning; await provision(); + expect(calls).toBe(1); + expect(await readFile(path.join(authority.location!.remoteCwd, "in-flight.txt"), "utf8")).toBe("one effect\n"); + } finally { release.resolve(); await provisioning; } + }, 30_000); + + it("holds an unknown provisioning effect after a lost command response", async () => { + const f = await preparedFailure(); + await db.update(heartbeatRuns).set({ status: "running", finishedAt: null }).where(eq(heartbeatRuns.id, f.runId)); + const authority = readRemoteWorkspaceAuthority(f.workspace.metadata)!; + const execute = f.target.runner!.execute.bind(f.target.runner!); + const command = "printf 'one effect\\n' >> unknown-effect.txt"; + let calls = 0; + f.target.runner!.execute = async input => { + const result = await execute(input); + if (input.command === "bash" && input.args?.[1] === command) { calls++; throw new Error("command response lost"); } + return result; + }; + const provision = () => provisionRemoteWorkspace({ db, companyId: f.companyId, workspaceId: f.workspace.id, + runId: f.runId, epoch: authority.epoch, remoteCwd: authority.location!.remoteCwd, target: f.target, commands: [command] }); + await expect(provision()).rejects.toThrow("command response lost"); + await expect(provision()).rejects.toThrow("remote_workspace_provision_unverified"); + expect(calls).toBe(1); + expect(await readFile(path.join(authority.location!.remoteCwd, "unknown-effect.txt"), "utf8")).toBe("one effect\n"); + const [workspace] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, f.workspace.id)); + expect(readRemoteWorkspaceAuthority(workspace.metadata)).toEqual(authority); + expect(await db.select().from(workspaceOperations).where(and(eq(workspaceOperations.companyId, f.companyId), + eq(workspaceOperations.heartbeatRunId, f.runId), eq(workspaceOperations.phase, "remote_workspace_provision")))) + .toEqual([expect.objectContaining({ status: "running", metadata: expect.objectContaining({ uncertain: true }) })]); + }, 30_000); + + it("does not admit Save before an administrative import publishes remote authority", async () => { + const f = await preparedFailure(), workspaceId = randomUUID(); + const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Import legacy files" }).returning(); + const cwd = path.join(root, workspaceId); await mkdir(cwd); + await writeFile(path.join(cwd, "local.txt"), "retained local work\n"); + const [workspace] = await db.insert(executionWorkspaces).values({ id: workspaceId, companyId: f.companyId, + sourceIssueId: issue.id, mode: "isolated_workspace", strategyType: "project_primary", providerType: "local_fs", + name: "Legacy files", cwd, providerRef: cwd }).returning(); + await db.update(issues).set({ executionWorkspaceId: workspaceId }).where(eq(issues.id, issue.id)); + const ready = deferred(), release = deferred(); + const storage = { ...f.storage, putFile: async (input: Parameters[0]) => { + const saved = await f.storage.putFile(input); ready.resolve(); await release.promise; return saved; + } }; + const migrating = migrateRemoteWorkspace({ db, storage, companyId: f.companyId, workspaceId, apply: true }) + .then(value => ({ value }), error => ({ error })); + try { + await ready.promise; + await expect(saveRemoteWorkspace({ db, storage, workspace, runId: f.runId, + lease: f.lease, target: f.target, assertOwned: async () => undefined })).rejects.toThrow("remote_workspace_save_binding_mismatch"); + await db.insert(heartbeatRuns).values({ companyId: f.companyId, agentId: f.agentId, issueId: issue.id, + runtimeMode: "legacy", status: "running", contextSnapshot: { issueId: issue.id } }); + release.resolve(); + expect(await migrating).toMatchObject({ error: expect.objectContaining({ message: expect.stringContaining("active_execution") }) }); + const [current] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, workspaceId)); + expect(readRemoteWorkspaceAuthority(current.metadata)).toBeNull(); + expect(await readFile(path.join(cwd, "local.txt"), "utf8")).toBe("retained local work\n"); + } finally { release.resolve(); await migrating; } + }, 30_000); + + it("rechecks the migration source task after waiting for the workspace", async () => { + const f = await preparedFailure(); + const [other] = await db.insert(issues).values({ companyId: f.companyId, title: "Changed source" }).returning(); + const peerDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "save-order-migration-binding" }); + const ready = deferred(), release = deferred(); + const barrier = db.transaction(async tx => { + await tx.select({ id: executionWorkspaces.id }).from(executionWorkspaces) + .where(eq(executionWorkspaces.id, f.workspace.id)).for("update"); + ready.resolve(); await release.promise; + await tx.update(executionWorkspaces).set({ sourceIssueId: other.id }).where(eq(executionWorkspaces.id, f.workspace.id)); + }); + await ready.promise; + const migrating = migrateRemoteWorkspace({ db: peerDb, storage: f.storage, companyId: f.companyId, + workspaceId: f.workspace.id, apply: true }).then(value => ({ value }), error => ({ error })); + try { + await expect.poll(async () => (await db.execute<{ waiting: boolean }>(sql`select exists(select 1 from pg_stat_activity + where datname = current_database() and application_name = 'save-order-migration-binding' + and wait_event_type = 'Lock' and query like '%"execution_workspaces"%') as waiting`))[0]?.waiting, + { timeout: 5000 }).toBe(true); + release.resolve(); await barrier; + expect(await migrating).toMatchObject({ error: expect.objectContaining({ message: "remote_workspace_migration_source_changed" }) }); + const [workspace] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, f.workspace.id)); + expect(readRemoteWorkspaceAuthority(workspace.metadata)).toEqual(readRemoteWorkspaceAuthority(f.workspace.metadata)); + expect(workspace.metadata?.remoteWorkspaceMigration).toBeUndefined(); + expect(await readFile(path.join(readRemoteWorkspaceAuthority(workspace.metadata)!.location!.remoteCwd, "saved.txt"), "utf8")) + .toBe("one preserved result\n"); + } finally { release.resolve(); await barrier; await migrating; await peerDb.$client.end(); } + }, 30_000); + + it("keeps import publication and an arriving legacy reservation in parent order", async () => { + const f = await preparedFailure(), workspaceId = randomUUID(); + const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Import before admission" }).returning(); + const cwd = path.join(root, workspaceId); await mkdir(cwd); + await writeFile(path.join(cwd, "local.txt"), "retained local work\n"); + await db.insert(executionWorkspaces).values({ id: workspaceId, companyId: f.companyId, sourceIssueId: issue.id, + mode: "isolated_workspace", strategyType: "project_primary", providerType: "local_fs", name: "Legacy files", cwd, providerRef: cwd }); + await db.update(issues).set({ executionWorkspaceId: workspaceId }).where(eq(issues.id, issue.id)); + const migrationDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "save-order-import-peer" }); + const reserveDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "save-order-import-reserve" }); + await reserveDb.execute(sql`select set_config('deadlock_timeout', '50ms', false)`); + const uploadReady = deferred(), uploadRelease = deferred(), lockReady = deferred(), lockRelease = deferred(); + const storage = { ...f.storage, putFile: async (input: Parameters[0]) => { + const saved = await f.storage.putFile(input); uploadReady.resolve(); await uploadRelease.promise; return saved; + } }; + const migrating = migrateRemoteWorkspace({ db: migrationDb, storage, companyId: f.companyId, workspaceId, apply: true }) + .then(value => ({ value }), error => ({ error })); + let barrier: Promise | undefined, reserving: Promise<{ error: unknown } | { value: unknown }> | undefined; + try { + await uploadReady.promise; + barrier = db.transaction(async tx => { + await tx.select({ id: executionWorkspaces.id }).from(executionWorkspaces) + .where(eq(executionWorkspaces.id, workspaceId)).for("update"); + lockReady.resolve(); await lockRelease.promise; + }); + await lockReady.promise; + const [run] = await db.insert(heartbeatRuns).values({ companyId: f.companyId, agentId: f.agentId, + issueId: issue.id, runtimeMode: "legacy", status: "running", contextSnapshot: { issueId: issue.id } }).returning(); + uploadRelease.resolve(); + const waiting = async (application: string) => (await db.execute<{ query: string; blockers: number[] }>(sql` + select query, pg_blocking_pids(pid) as blockers from pg_stat_activity + where datname = current_database() and application_name = ${application} and wait_event_type = 'Lock'`))[0]; + await expect.poll(async () => (await waiting("save-order-import-peer"))?.query, { timeout: 5000 }).toContain('"execution_workspaces"'); + reserving = reserveRemoteWorkspace({ db: reserveDb, storage: f.storage, companyId: f.companyId, issueId: issue.id, + runId: run.id, agentId: f.agentId, environmentId: f.environmentId, projectId: null, + projectWorkspaceId: null, existingWorkspaceId: workspaceId, source: { kind: "empty" } }) + .then(value => ({ value }), error => ({ error })); + await expect.poll(async () => Boolean(await waiting("save-order-import-reserve")), { timeout: 5000 }).toBe(true); + const peerWait = await waiting("save-order-import-peer"), reserveWait = await waiting("save-order-import-reserve"); + lockRelease.resolve(); await barrier; + const results = await Promise.all([migrating, reserving]); + const databaseErrors = results.flatMap(result => "error" in result && (result.error as { cause?: { code?: string } })?.cause?.code + ? [result.error as { cause: { code: string; detail?: string; where?: string } }] : []); + if (databaseErrors.length) console.info("save-order-import-postgres-evidence", JSON.stringify({ + peerQuery: peerWait?.query, reserveQuery: reserveWait?.query, errors: databaseErrors.map(error => error.cause) })); + expect(databaseErrors).toEqual([]); + expect(results[0]).toMatchObject({ error: expect.objectContaining({ message: expect.stringContaining("active_execution") }) }); + expect(results[1]).toMatchObject({ error: expect.objectContaining({ details: { code: "remote_workspace_migration_required" } }) }); + const [workspace] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, workspaceId)); + expect(readRemoteWorkspaceAuthority(workspace.metadata)).toBeNull(); + expect(await readFile(path.join(cwd, "local.txt"), "utf8")).toBe("retained local work\n"); + const snapshots = await db.select().from(executionWorkspaceSnapshots).where(eq(executionWorkspaceSnapshots.executionWorkspaceId, workspaceId)); + expect(snapshots).toEqual([expect.objectContaining({ status: "uploaded" })]); + expect((await f.storage.headObject(f.companyId, snapshots[0].objectKey!)).exists).toBe(true); + } finally { + uploadRelease.resolve(); lockRelease.resolve(); await barrier; await Promise.all([migrating, reserving]); + await migrationDb.$client.end(); await reserveDb.$client.end(); + } + }, 30_000); + + it.each(["prepare", "provision", "migrate"] as const)("saves beside actual %s started first", async mode => { + const f = await preparedFailure(); + await db.update(heartbeatRuns).set({ status: "running", finishedAt: null, resultJson: null, errorCode: null }) + .where(eq(heartbeatRuns.id, f.runId)); + const metadata = { ...f.lease.metadata }; delete metadata.remoteExecutionTermination; + const [lease] = await db.update(environmentLeases).set({ status: "active", releasedAt: null, cleanupStatus: null, metadata }) + .where(eq(environmentLeases.id, f.lease.id)).returning(); + const authority = readRemoteWorkspaceAuthority(f.workspace.metadata)!; + const saveDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "save-order-reverse-save" }); + const peerDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "save-order-reverse-peer" }); + await saveDb.execute(sql`select set_config('deadlock_timeout', '50ms', false)`); + const ready = deferred(), release = deferred(); + const barrier = db.transaction(async tx => { + await tx.select({ id: executionWorkspaces.id }).from(executionWorkspaces) + .where(eq(executionWorkspaces.id, f.workspace.id)).for("update"); + ready.resolve(); await release.promise; + }); + await ready.promise; + const command = "printf 'effect\\n' >> provision-count.txt"; + const provision = () => provisionRemoteWorkspace({ db: peerDb, companyId: f.companyId, workspaceId: f.workspace.id, + runId: f.runId, epoch: authority.epoch, remoteCwd: authority.location!.remoteCwd, target: f.target, commands: [command] }); + const settle = (work: Promise): Promise> => work.then( + value => ({ status: "fulfilled", value }), reason => ({ status: "rejected", reason })); + let saving: Promise>>> | undefined; + let peer: Promise> | undefined; + try { + const waiting = async (application: string) => (await db.execute<{ pid: number; query: string; blockers: number[] }>(sql` + select pid, query, pg_blocking_pids(pid) as blockers from pg_stat_activity + where datname = current_database() and application_name = ${application} and wait_event_type = 'Lock'`))[0]; + peer = settle(mode === "prepare" ? prepareRemoteWorkspace({ db: peerDb, storage: f.storage, + workspace: f.workspace, runId: f.runId, lease: lease as EnvironmentLease, target: f.target }) + : mode === "provision" ? provision() : migrateRemoteWorkspace({ db: peerDb, storage: f.storage, + companyId: f.companyId, workspaceId: f.workspace.id, apply: true })); + await expect.poll(async () => (await waiting("save-order-reverse-peer"))?.query, + { timeout: 5000 }).toContain('"execution_workspaces"'); + const peerWait = await waiting("save-order-reverse-peer"); + saving = settle(saveRemoteWorkspace({ db: saveDb, storage: f.storage, workspace: f.workspace, + runId: f.runId, lease: lease as EnvironmentLease, target: f.target, boundary: "checkpoint", assertOwned: async () => undefined })); + await expect.poll(async () => Boolean(await waiting("save-order-reverse-save")), { timeout: 5000 }).toBe(true); + const saveWait = await waiting("save-order-reverse-save"); + release.resolve(); await barrier; + const results = await Promise.all([saving, peer]); + const errors = results.filter(result => result.status === "rejected").map(result => + (result as PromiseRejectedResult).reason as Error & { cause?: { code?: string; detail?: string; where?: string } }); + if (errors.length) console.info("save-order-reverse-postgres-evidence", JSON.stringify({ mode, + peerQuery: peerWait?.query, saveQuery: saveWait?.query, + errors: errors.map(error => ({ code: error.cause?.code, detail: error.cause?.detail, where: error.cause?.where })) })); + expect(errors).toEqual([]); + if (mode === "migrate") expect((results[1] as PromiseFulfilledResult).value).toMatchObject({ applied: true, replayed: true }); + if (mode === "provision") { + await provision(); + expect(await readFile(path.join(authority.location!.remoteCwd, "provision-count.txt"), "utf8")).toBe("effect\n"); + expect(await db.select().from(workspaceOperations).where(and(eq(workspaceOperations.companyId, f.companyId), + eq(workspaceOperations.heartbeatRunId, f.runId), eq(workspaceOperations.phase, "remote_workspace_provision")))) + .toEqual([expect.objectContaining({ issueId: f.issueId, status: "succeeded", metadata: expect.objectContaining({ uncertain: false }) })]); + } + const [current] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, f.workspace.id)); + const saved = await saveRemoteWorkspace({ db: saveDb, storage: f.storage, workspace: current, + runId: f.runId, lease: lease as EnvironmentLease, target: f.target, assertOwned: async () => undefined }); + const loaded = await createSavedWorkspaceLoader(db, f.storage)({ companyId: f.companyId, + workspaceId: f.workspace.id, snapshotId: saved.snapshotId }); + try { + expect(await readFile(path.join(loaded.root, "saved.txt"), "utf8")).toBe("one preserved result\n"); + if (mode === "provision") expect(await readFile(path.join(loaded.root, "provision-count.txt"), "utf8")).toBe("effect\n"); + } finally { await loaded.cleanup(); } + } finally { + release.resolve(); await barrier; await Promise.all([saving, peer]); + await saveDb.$client.end(); await peerDb.$client.end(); + } + }, 30_000); + + it.each([false, true])("preserves immutable task identity while waiting for its lock (foreign company=%s)", async foreignCompany => { + const f = await preparedFailure(); + const companyId = foreignCompany ? randomUUID() : f.companyId; + if (foreignCompany) await db.insert(companies).values({ id: companyId, name: "Other owner", issuePrefix: `S${companyId.slice(0, 6)}` }); + const [otherIssue] = await db.insert(issues).values({ companyId, title: "Other task", status: "blocked" }).returning(); + const saveDb = createDb(database.connectionString, { maxConnections: 1, applicationName: "save-order-binding" }); + const ready = deferred(), release = deferred(); + const barrier = db.transaction(async tx => { + await tx.select({ id: issues.id }).from(issues).where(eq(issues.id, f.issueId)).for("update"); + ready.resolve(); await release.promise; + }); + await ready.promise; + const saving = saveRemoteWorkspace({ db: saveDb, storage: f.storage, workspace: f.workspace, + runId: f.runId, lease: f.lease, target: f.target, assertOwned: async () => undefined }) + .then(value => ({ value }), error => ({ error })); + try { + await expect.poll(async () => (await db.execute<{ waiting: boolean }>(sql`select exists(select 1 from pg_stat_activity + where datname = current_database() and application_name = 'save-order-binding' + and wait_event_type = 'Lock' and query like '%"issues"%') as waiting`))[0]?.waiting, + { timeout: 5000 }).toBe(true); + await expect(db.update(heartbeatRuns).set({ issueId: otherIssue.id }).where(and( + eq(heartbeatRuns.companyId, f.companyId), eq(heartbeatRuns.id, f.runId)))) + .rejects.toMatchObject({ cause: { code: "P0001", message: "run task binding is immutable" } }); + release.resolve(); await barrier; + const result = await saving; + expect(result).toHaveProperty("value.snapshotId"); + const [source] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId)); + expect(source.issueId).toBe(f.issueId); + const [workspace] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, f.workspace.id)); + expect(readRemoteWorkspaceAuthority(workspace.metadata)).toMatchObject({ dirty: false, writerRunId: null }); + expect(await db.select().from(executionWorkspaceSnapshots).where(eq(executionWorkspaceSnapshots.executionWorkspaceId, workspace.id))) + .toHaveLength(f.snapshotCount + 1); + expect(await db.select().from(workspaceOperations).where(and(eq(workspaceOperations.companyId, f.companyId), + eq(workspaceOperations.heartbeatRunId, f.runId), sql`${workspaceOperations.metadata}->>'boundary' = 'final'`))) + .toEqual([expect.objectContaining({ issueId: f.issueId, status: "succeeded" })]); + } finally { + release.resolve(); await barrier; await saving; await saveDb.$client.end(); + } + }, 30_000); + + it("retains the uploaded snapshot when the workspace writer changes before publication", async () => { + const f = await preparedFailure(); + const authority = readRemoteWorkspaceAuthority(f.workspace.metadata)!; + const next = { ...authority, epoch: authority.epoch + 1, writerRunId: randomUUID() }; + let checks = 0; + await expect(saveRemoteWorkspace({ db, storage: f.storage, workspace: f.workspace, + runId: f.runId, lease: f.lease, target: f.target, assertOwned: async () => { + if (++checks === 3) await db.update(executionWorkspaces).set({ metadata: { ...f.workspace.metadata, remoteWorkspace: next } }) + .where(eq(executionWorkspaces.id, f.workspace.id)); + } })).rejects.toThrow("remote_workspace_publication_fenced"); + const [workspace] = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.id, f.workspace.id)); + expect(readRemoteWorkspaceAuthority(workspace.metadata)).toEqual(next); + const snapshots = await db.select().from(executionWorkspaceSnapshots).where(eq(executionWorkspaceSnapshots.executionWorkspaceId, workspace.id)); + expect(snapshots).toHaveLength(f.snapshotCount + 1); + const uploaded = snapshots.filter(snapshot => snapshot.status === "uploaded"); + expect(uploaded).toHaveLength(1); + expect((await f.storage.headObject(f.companyId, uploaded[0].objectKey!)).exists).toBe(true); + expect(snapshots.filter(snapshot => snapshot.status === "committed")).toHaveLength(1); + expect(await readFile(path.join(authority.location!.remoteCwd, "saved.txt"), "utf8")).toBe("one preserved result\n"); + }, 30_000); +}); diff --git a/server/src/services/remote-workspaces.integration.test.ts b/server/src/services/remote-workspaces.integration.test.ts index c5eeef78d43..ba19ea07e70 100644 --- a/server/src/services/remote-workspaces.integration.test.ts +++ b/server/src/services/remote-workspaces.integration.test.ts @@ -825,7 +825,12 @@ suite("remote workspace lifecycle with real Postgres and independent filesystem const newLease = (await db.select().from(environmentLeases).where(eq(environmentLeases.id, secondLease)))[0]! as EnvironmentLease; const nextTarget = inherited ? { ...await target("child-sandbox"), leaseId: secondLease, sandboxId: newLease.providerLeaseId! } : first; - const continued = await prepareRemoteWorkspace({ db, storage, workspace: reused, runId: secondRun, lease: newLease, target: nextTarget }); + const continued = await prepareRemoteWorkspace({ db, storage, workspace: reused, runId: secondRun, lease: newLease, target: nextTarget, + runtimeProvisionCommand: "printf 'continued effect\\n' >> continued-provision.txt" }); + expect(await fs.readFile(path.join(continued.remoteRoot, "continued-provision.txt"), "utf8")).toBe("continued effect\n"); + const [provisioned] = await db.select().from(workspaceOperations).where(and(eq(workspaceOperations.heartbeatRunId, secondRun), + eq(workspaceOperations.phase, "remote_workspace_provision"))); + expect(provisioned).toMatchObject({ issueId: secondIssue, status: "succeeded" }); expect(await fs.readFile(path.join(continued.remoteRoot, "shared-result.txt"), "utf8")).toBe("from the first task\n"); await fs.appendFile(path.join(continued.remoteRoot, "shared-result.txt"), "from the child task\n"); const childSaved = await saveRemoteWorkspace({ db, storage, workspace: continued.workspace, runId: secondRun, diff --git a/server/src/services/remote-workspaces.ts b/server/src/services/remote-workspaces.ts index 6efda0aa5af..c7c4a00f535 100644 --- a/server/src/services/remote-workspaces.ts +++ b/server/src/services/remote-workspaces.ts @@ -7,7 +7,7 @@ import os from "node:os"; import path from "node:path"; import { and, eq, inArray, sql } from "drizzle-orm"; import { - environmentLeases, executionWorkspaces, executionWorkspaceSnapshots, heartbeatRuns, issues, workspaceOperations, type Db, + companies, environmentLeases, executionWorkspaces, executionWorkspaceSnapshots, heartbeatRuns, issues, workspaceOperations, type Db, } from "@paperclipai/db"; import { readRemoteWorkspaceAuthority, remoteWorkspaceSourceSchema, @@ -127,6 +127,28 @@ async function currentWorkspace(tx: Transaction, companyId: string, workspaceId: return workspace; } +async function lockWorkspaceRunParents(tx: Transaction, companyId: string, runId: string, expectedIssueId?: string | null) { + const fields = { nativeIssueId: heartbeatRuns.nativeIssueId, issueId: heartbeatRuns.issueId, + contextIssueId: sql`${heartbeatRuns.contextSnapshot}->'issueId'` }; + const [selected] = await tx.select(fields).from(heartbeatRuns).where(and( + eq(heartbeatRuns.companyId, companyId), eq(heartbeatRuns.id, runId))); + if (!selected) throw new Error("remote_workspace_run_missing"); + const issueId = selected.nativeIssueId ?? selected.issueId ?? (typeof selected.contextIssueId === "string" ? selected.contextIssueId : null); + if (expectedIssueId !== undefined && issueId !== expectedIssueId) throw new Error("remote_workspace_run_binding_changed"); + const [company] = await tx.select({ id: companies.id }).from(companies).where(eq(companies.id, companyId)).for("key share"); + if (!company) throw new Error("remote_workspace_run_missing"); + if (issueId) { + const [issue] = await tx.select({ id: issues.id }).from(issues).where(and( + eq(issues.companyId, companyId), eq(issues.id, issueId))).for("update"); + if (!issue) throw new Error("remote_workspace_run_binding_changed"); + } + const [current] = await tx.select(fields).from(heartbeatRuns).where(and( + eq(heartbeatRuns.companyId, companyId), eq(heartbeatRuns.id, runId))).for("update"); + const currentIssueId = current?.nativeIssueId ?? current?.issueId ?? (typeof current?.contextIssueId === "string" ? current.contextIssueId : null); + if (!current || currentIssueId !== issueId) throw new Error("remote_workspace_run_binding_changed"); + return issueId; +} + export class RemoteWorkspaceWriterActiveError extends HttpError { constructor( readonly workspaceId: string, @@ -149,8 +171,9 @@ export async function reserveRemoteWorkspace(input: { assertRemoteWorkspacePersistence(input.storage ?? getStorageService()); const source = remoteWorkspaceSourceSchema.parse(input.source); return input.db.transaction(async tx => { + await tx.select({ id: companies.id }).from(companies).where(eq(companies.id, input.companyId)).for("key share"); const [issue] = await tx.select().from(issues).where(and(eq(issues.companyId, input.companyId), eq(issues.id, input.issueId))).for("update"); - const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.companyId, input.companyId), eq(heartbeatRuns.id, input.runId))); + const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.companyId, input.companyId), eq(heartbeatRuns.id, input.runId))).for("update"); if (!issue || !run || run.agentId !== input.agentId || run.status !== "running" || ![run.nativeIssueId, run.issueId, object(run.contextSnapshot).issueId].includes(input.issueId)) throw forbidden("Workspace preparation requires its current task execution."); if (input.fresh && (issue.executionWorkspaceId !== input.fresh.previousWorkspaceId @@ -433,6 +456,7 @@ export async function prepareRemoteWorkspace(input: { } await input.assertOwned?.(); const prepared = await input.db.transaction(async tx => { + await lockWorkspaceRunParents(tx, input.workspace.companyId, input.runId); const workspace = await currentWorkspace(tx, input.workspace.companyId, input.workspace.id); const current = readRemoteWorkspaceAuthority(workspace.metadata); if (!current || current.epoch !== initial.epoch || current.writerRunId !== input.runId) throw new Error("remote_workspace_authority_changed"); @@ -565,6 +589,7 @@ export async function saveRemoteWorkspace(input: { await input.assertOwned(); const invocationId = randomUUID(); const operation = await input.db.transaction(async tx => { + const issueId = await lockWorkspaceRunParents(tx, input.workspace.companyId, input.runId); const workspace = await currentWorkspace(tx, input.workspace.companyId, input.workspace.id); const current = readRemoteWorkspaceAuthority(workspace.metadata); if (!current || current.epoch !== expected.epoch || (current.writerRunId !== input.runId && current.dirty)) throw new Error("remote_workspace_authority_changed"); @@ -573,6 +598,7 @@ export async function saveRemoteWorkspace(input: { eq(workspaceOperations.phase, "remote_workspace_snapshot"), sql`${workspaceOperations.metadata}->>'boundary' = ${boundary}`, sql`${workspaceOperations.metadata}->>'epoch' = ${String(current.epoch)}`)).limit(1).for("update"); if (previous) { + if (previous.issueId !== issueId) throw new Error("remote_workspace_run_binding_changed"); if (object(previous.metadata?.retiredSave).schema === RETIRED_SAVE_SCHEMA) throw new Error("remote_workspace_snapshot_retired"); const work = readSaveWork(previous.metadata); const [registered] = await tx.update(workspaceOperations).set({ metadata: { ...previous.metadata, @@ -580,11 +606,6 @@ export async function saveRemoteWorkspace(input: { unknown: work?.unknown ?? true, quiescedAt: null } }, updatedAt: new Date() }).where(eq(workspaceOperations.id, previous.id)).returning(); return registered!; } - const [run] = await tx.select({ nativeIssueId: heartbeatRuns.nativeIssueId, issueId: heartbeatRuns.issueId, - contextSnapshot: heartbeatRuns.contextSnapshot }).from(heartbeatRuns).where(and( - eq(heartbeatRuns.companyId, workspace.companyId), eq(heartbeatRuns.id, input.runId))); - if (!run) throw new Error("remote_workspace_run_missing"); - const issueId = run.nativeIssueId ?? run.issueId ?? (typeof run.contextSnapshot?.issueId === "string" ? run.contextSnapshot.issueId : null); // The workspace lock serializes generation reservations. Failed/deleted // snapshots still own their generation, and an operation can be committed // before beginWorkspaceSnapshot creates its ledger row. Count both journals @@ -627,6 +648,7 @@ export async function saveRemoteWorkspace(input: { await input.assertOwned(); const at = new Date(); await input.db.transaction(async tx => { + await lockWorkspaceRunParents(tx, input.workspace.companyId, input.runId, operation.issueId); const workspace = await currentWorkspace(tx, input.workspace.companyId, input.workspace.id); const current = readRemoteWorkspaceAuthority(workspace.metadata); if (!current || current.epoch !== expected.epoch || current.writerRunId !== input.runId || current.snapshotId !== snapshot.parentSnapshotId) { @@ -787,6 +809,7 @@ export async function provisionRemoteWorkspace(input: { await input.assertOwned?.(); const digest = createHash("sha256").update(command).digest("hex"); const operation = await input.db.transaction(async tx => { + const issueId = await lockWorkspaceRunParents(tx, input.companyId, input.runId); const [workspace] = await tx.select().from(executionWorkspaces).where(and(eq(executionWorkspaces.companyId, input.companyId), eq(executionWorkspaces.id, input.workspaceId))).for("update"); const authority = readRemoteWorkspaceAuthority(workspace?.metadata); const [run] = await tx.select({ status: heartbeatRuns.status }).from(heartbeatRuns).where(and(eq(heartbeatRuns.companyId, input.companyId), eq(heartbeatRuns.id, input.runId))); @@ -800,7 +823,7 @@ export async function provisionRemoteWorkspace(input: { return null; } const [created] = await tx.insert(workspaceOperations).values({ companyId: input.companyId, executionWorkspaceId: input.workspaceId, - heartbeatRunId: input.runId, issueId: workspace!.sourceIssueId, phase: "remote_workspace_provision", status: "running", + heartbeatRunId: input.runId, issueId, phase: "remote_workspace_provision", status: "running", metadata: { epoch: input.epoch, index, commandSha256: digest, uncertain: true } }).returning(); return created!; });