diff --git a/server/src/services/environment-task-runtime.ts b/server/src/services/environment-task-runtime.ts index 1a99a51741c..9da7e111669 100644 --- a/server/src/services/environment-task-runtime.ts +++ b/server/src/services/environment-task-runtime.ts @@ -1,3 +1,4 @@ +import { stageProviderFreshWorkspace } from "./provider-fresh-workspace.js"; import { readWorkspaceProviderBinding } from "@paperclipai/shared"; import { settlePluginTaskCredentialHandoff } from "./plugin-task-credential-writeback.js"; import { and, eq, inArray, or, sql } from "drizzle-orm"; @@ -105,9 +106,14 @@ export async function executeEnvironmentTask(db: Db, workers: PluginWorkerManage throw new Error("Managed workspace execution binding changed"); } } - const freshWorkspace = metadata.freshWorkspace === undefined ? undefined : freshWorkspaceSchema.parse(metadata.freshWorkspace); + let freshWorkspace = metadata.freshWorkspace === undefined ? undefined : freshWorkspaceSchema.parse(metadata.freshWorkspace); if (freshWorkspace && driver.freshWorkspaceVersion !== 1 && (launches || operation.kind === "describe" || operation.kind === "check_project_auth")) throw new Error("Plugin-provided fresh workspace preparation unavailable"); + if (freshWorkspace && operation.kind === "prepare") { + if (!lease.issueId) throw new Error("Fresh workspace task ownership is missing"); + freshWorkspace = await stageProviderFreshWorkspace({ db, companyId: input.companyId, + issueId: lease.issueId, ...freshWorkspace, existing: freshWorkspace }); + } const projectIds = scopesProjects ? operation.projectIds : []; if (activeProbe && projectIds.length > 0) throw new Error("Plugin Runner probe cannot use project mounts"); if (projectIds.length > 0) { diff --git a/server/src/services/plugin-workspaces.ts b/server/src/services/plugin-workspaces.ts index 09eb306a3bb..9c1a3f09ccb 100644 --- a/server/src/services/plugin-workspaces.ts +++ b/server/src/services/plugin-workspaces.ts @@ -1,4 +1,5 @@ import { effectiveProjectWorkspacePolicy, isWorkspaceProvider } from "./project-workspace-provider.js"; +import { stageProviderFreshWorkspace } from "./provider-fresh-workspace.js"; import { randomUUID, createHash } from "node:crypto"; import { and, eq, sql } from "drizzle-orm"; import { @@ -86,13 +87,20 @@ export async function callWorkspaceProvider( const pluginId = binding?.pluginId ?? input.pluginId; if (!pluginId || !workspace.projectId) throw conflict("Workspace storage ownership is missing"); - await provider(db, workers, workspace.companyId, pluginId); + const selectedProvider = await provider(db, workers, workspace.companyId, pluginId); + if (input.operation.kind === "prepare" && input.operation.fresh && selectedProvider.manifestJson.workspaceProvider?.freshWorkspaceVersion !== 1) + throw conflict("Fresh workspace preparation is unavailable"); const generation = typeof workspace.metadata?.lifecycleGeneration === "number" ? workspace.metadata.lifecycleGeneration : 0; const operationGeneration = generation + (input.advanceGeneration ? 1 : 0); const operation = workspaceProviderOperationSchema.parse(input.operation); + if (operation.kind === "prepare" && operation.fresh) { + if (!workspace.sourceIssueId) throw conflict("Fresh workspace task ownership is missing"); + operation.fresh = await stageProviderFreshWorkspace({ db, companyId: workspace.companyId, + issueId: workspace.sourceIssueId, ...operation.fresh, existing: operation.fresh }); + } const params: PluginWorkspaceParams = { pluginId, companyId: workspace.companyId, @@ -871,10 +879,11 @@ export async function restorePluginWorkspace( source, mode: workspace.mode, strategy: workspace.strategyType, - branchName: workspace.branchName, - branchMode: "existing", + branchName: workspace.metadata?.workspaceProviderFresh ? null : workspace.branchName, + branchMode: workspace.metadata?.workspaceProviderFresh ? "create" : "existing", baseRef: workspace.baseRef, existing: binding, + ...(workspace.metadata?.workspaceProviderFresh ? { fresh: workspace.metadata.workspaceProviderFresh as PluginFreshWorkspace } : {}), }); const [pending] = await db.transaction(async (tx) => { await notifyTerminalWorkspaceWork(tx, `workspace:${workspace.id}`); diff --git a/server/src/services/provider-fresh-workspace.test.ts b/server/src/services/provider-fresh-workspace.test.ts index 4fab83f3032..74f7d35f41b 100644 --- a/server/src/services/provider-fresh-workspace.test.ts +++ b/server/src/services/provider-fresh-workspace.test.ts @@ -25,11 +25,13 @@ describe("fresh provider recovery transfer staging", () => { const first = await stageProviderFreshWorkspace(input); expect(first.recoverySnapshots).toHaveLength(1); expect(await readFile(first.recoverySnapshots[0].archive.path)).toEqual(archive); - expect(await stageProviderFreshWorkspace(input)).toEqual(first); - expect(state.getObject).toHaveBeenCalledTimes(2); + expect(await stageProviderFreshWorkspace({ ...input, existing: first })).toEqual(first); + await rm(first.recoverySnapshots[0].archive.path); + expect(await stageProviderFreshWorkspace({ ...input, existing: first })).toEqual(first); + expect(state.getObject).toHaveBeenCalledTimes(3); await chmod(first.recoverySnapshots[0].archive.path, 0o600); await writeFile(first.recoverySnapshots[0].archive.path, "corrupt"); await expect(stageProviderFreshWorkspace(input)).rejects.toThrow(); - expect(state.getObject).toHaveBeenCalledTimes(2); + expect(state.getObject).toHaveBeenCalledTimes(3); }); }); diff --git a/server/src/services/provider-fresh-workspace.ts b/server/src/services/provider-fresh-workspace.ts index 65a32384e59..0b69cba8f3e 100644 --- a/server/src/services/provider-fresh-workspace.ts +++ b/server/src/services/provider-fresh-workspace.ts @@ -4,6 +4,7 @@ import { createReadStream, createWriteStream } from "node:fs"; import { lstat, mkdir, realpath, rename, rm } from "node:fs/promises"; import path from "node:path"; import { pipeline } from "node:stream/promises"; +import { Transform } from "node:stream"; import { and, eq } from "drizzle-orm"; import { executionWorkspaces, heartbeatRuns, issues, workspaceOperations, type Db } from "@paperclipai/db"; import { freshWorkspaceSchema, type PluginFreshWorkspace } from "@paperclipai/plugin-sdk"; @@ -17,10 +18,11 @@ import { SystemRecoveryFailure } from "./system-recovery-state.js"; export async function stageProviderFreshWorkspace(input: { db: Db; companyId: string; issueId: string; attemptId: string; sourceRunId: string; - previousWorkspaceId: string | null; + previousWorkspaceId: string | null; existing?: PluginFreshWorkspace; }): Promise { const fresh: PluginFreshWorkspace = freshWorkspaceSchema.parse({ attemptId: input.attemptId, sourceRunId: input.sourceRunId, previousWorkspaceId: input.previousWorkspaceId, recoverySnapshots: [] }); + if (input.existing?.recoverySnapshots.length === 0) return freshWorkspaceSchema.parse(input.existing); const storage = getStorageService(); const sources = await recoveryRunAncestry({ ...input, sourceRunIds: [input.sourceRunId] }); const stagingRoot = path.join(resolvePaperclipInstanceRoot(), "runtime", "provider-recovery", input.companyId, fresh.attemptId); @@ -32,6 +34,7 @@ export async function stageProviderFreshWorkspace(input: { eq(workspaceOperations.heartbeatRunId, sourceRunId), eq(workspaceOperations.phase, "workspace_recovery_snapshot"), eq(workspaceOperations.status, "succeeded"))); for (const operation of operations) { + if (input.existing && !input.existing.recoverySnapshots.some(snapshot => snapshot.operationId === operation.id)) continue; const saved = readSnapshot(operation.metadata, input.companyId, sourceRunId); const directory = path.join(root, operation.id); await mkdir(directory, { recursive: true, mode: 0o700 }); @@ -50,7 +53,12 @@ export async function stageProviderFreshWorkspace(input: { const temporary = `${destination}.${process.pid}.${randomUUID()}.tmp`; const object = await storage.getObject(input.companyId, expected.objectKey); try { - await pipeline(object.stream, createWriteStream(temporary, { flags: "wx", mode: 0o400 })); + let received = 0; + const bounded = new Transform({ transform(bytes, _encoding, callback) { + received += bytes.length; + callback(received > expected.byteSize ? new SystemRecoveryFailure("snapshot_verification_failed", "service_required") : null, bytes); + } }); + await pipeline(object.stream, bounded, createWriteStream(temporary, { flags: "wx", mode: 0o400 })); await verify(temporary, expected); await rename(temporary, destination); } finally { object.stream.destroy(); await rm(temporary, { force: true }); } @@ -62,6 +70,8 @@ export async function stageProviderFreshWorkspace(input: { } } fresh.recoverySnapshots.sort((a, b) => a.operationId.localeCompare(b.operationId)); + if (input.existing && !isDeepStrictEqual(fresh, input.existing)) + throw new SystemRecoveryFailure("snapshot_verification_failed", "service_required"); return fresh; }