diff --git a/server/src/services/heartbeat/workspace-preparation.test.ts b/server/src/services/heartbeat/workspace-preparation.test.ts index 49602d95422..b51d8713be5 100644 --- a/server/src/services/heartbeat/workspace-preparation.test.ts +++ b/server/src/services/heartbeat/workspace-preparation.test.ts @@ -177,9 +177,17 @@ describe.skipIf(!support.supported)("heartbeat workspace preparation boundary", expect(call).not.toHaveBeenCalled(); }); - it.each([false, true])("keeps a fresh task reservation stable through repeated preparation (project: %s)", async project => { + it.each([ + { project: false, savedSource: false }, + { project: true, savedSource: true }, + { project: true, savedSource: false }, + ])("keeps a fresh task reservation stable through repeated preparation (project: $project, saved source: $savedSource)", async ({ project, savedSource }) => { const { input, issue, envOrchestrator } = await fixture(); if (!project) await db.update(issues).set({ projectId: null, projectWorkspaceId: null }).where(eq(issues.id, issue.id)); + else if (!savedSource) { + await db.update(issues).set({ projectWorkspaceId: null }).where(eq(issues.id, issue.id)); + input.task.issueRef = { ...input.task.issueRef!, projectWorkspaceId: null }; + } const pluginKey = "test.fresh-task"; await db.insert(plugins).values({ pluginKey, packageName: "test-tasks", version: "1.0.0", apiVersion: 1, status: "ready", categories: ["workspace"], manifestJson: { id: pluginKey, apiVersion: 1, version: "1.0.0", @@ -208,6 +216,8 @@ describe.skipIf(!support.supported)("heartbeat workspace preparation boundary", metadata: { runnerTaskFreshWorkspace: { attemptId: input.run.id } } }); const rows = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.companyId, input.agent.companyId)); expect(rows).toHaveLength(1); + const [task] = await db.select().from(issues).where(eq(issues.id, issue.id)); + expect(task.projectWorkspaceId).toBe(project ? issue.projectWorkspaceId : null); }); const readIssue = (id: string) => db.select().from(issues).where(eq(issues.id, id)).then(rows => rows[0]!); const readRun = (id: string) => db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, id)).then(rows => rows[0]!); diff --git a/server/src/services/provider-fresh-workspace.ts b/server/src/services/provider-fresh-workspace.ts index 31a70b8cd57..e3e0ecfca0e 100644 --- a/server/src/services/provider-fresh-workspace.ts +++ b/server/src/services/provider-fresh-workspace.ts @@ -6,8 +6,8 @@ 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 { and, asc, eq } from "drizzle-orm"; +import { executionWorkspaces, heartbeatRuns, issues, projectWorkspaces, workspaceOperations, type Db } from "@paperclipai/db"; import { freshWorkspaceSchema, type PluginFreshWorkspace } from "@paperclipai/plugin-sdk"; import { resolvePaperclipInstanceRoot } from "../home-paths.js"; import { getStorageService } from "../storage/index.js"; @@ -102,10 +102,19 @@ export async function reserveFreshPluginTaskWorkspace(input: { eq(heartbeatRuns.companyId, input.companyId))); if (!task || !run || run.agentId !== input.agentId || run.status !== "running" || ![run.issueId, run.nativeIssueId].includes(input.issueId)) throw forbidden("Workspace preparation no longer has an active run"); - await assertTaskWorkspaceExecutionAccess(tx, { type: "agent", source: "agent_jwt", + const actor = { type: "agent" as const, source: "agent_jwt" as const, companyId: input.companyId, agentId: input.agentId, runId: input.runId, - onBehalfOfUserId: input.responsibleUserId === "local-board" ? null : input.responsibleUserId, - }, input.companyId, { issueId: input.issueId, projectWorkspaceId: source.projectWorkspaceId }); + onBehalfOfUserId: input.responsibleUserId === "local-board" ? null : input.responsibleUserId }; + if (!task.projectWorkspaceId && task.projectId && task.projectId === source.projectId && source.projectWorkspaceId) { + await assertTaskWorkspaceExecutionAccess(tx, actor, input.companyId, { issueId: input.issueId }); + const [defaultSource] = await tx.select({ id: projectWorkspaces.id }).from(projectWorkspaces).where(and( + eq(projectWorkspaces.companyId, input.companyId), eq(projectWorkspaces.projectId, task.projectId))) + .orderBy(asc(projectWorkspaces.createdAt), asc(projectWorkspaces.id)).limit(1); + if (defaultSource?.id === source.projectWorkspaceId) await tx.update(issues).set({ projectWorkspaceId: defaultSource.id }) + .where(eq(issues.id, task.id)); + } + await assertTaskWorkspaceExecutionAccess(tx, actor, input.companyId, { + issueId: input.issueId, projectWorkspaceId: source.projectWorkspaceId }); const [reserved] = await tx.select().from(executionWorkspaces).where(and( eq(executionWorkspaces.companyId, input.companyId), eq(executionWorkspaces.sourceIssueId, input.issueId), eq(executionWorkspaces.id, input.fresh.attemptId)));