diff --git a/server/src/services/heartbeat/workspace-preparation.test.ts b/server/src/services/heartbeat/workspace-preparation.test.ts index aa115dfc429..49602d95422 100644 --- a/server/src/services/heartbeat/workspace-preparation.test.ts +++ b/server/src/services/heartbeat/workspace-preparation.test.ts @@ -177,18 +177,20 @@ describe.skipIf(!support.supported)("heartbeat workspace preparation boundary", expect(call).not.toHaveBeenCalled(); }); - it("keeps a fresh projectless reservation stable through repeated preparation", async () => { + it.each([false, true])("keeps a fresh task reservation stable through repeated preparation (project: %s)", async project => { const { input, issue, envOrchestrator } = await fixture(); - await db.update(issues).set({ projectId: null, projectWorkspaceId: null }).where(eq(issues.id, issue.id)); - const pluginKey = "test.fresh-projectless"; + if (!project) await db.update(issues).set({ projectId: null, projectWorkspaceId: null }).where(eq(issues.id, issue.id)); + 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", displayName: "Tasks", description: "Tasks", author: "Test", categories: ["workspace"], capabilities: ["environment.drivers.register"], entrypoints: { worker: "./worker.js" }, environmentDrivers: [{ driverKey: "tasks", displayName: "Task VM", configSchema: { type: "object" }, supportsTasks: true, freshWorkspaceVersion: 1 }] } }); input.agent.adapterType = "paperclip_runner"; - input.task.issueRef = { ...input.task.issueRef!, projectId: null, projectWorkspaceId: null }; - input.task.context = { issueId: issue.id }; + if (!project) { + input.task.issueRef = { ...input.task.issueRef!, projectId: null, projectWorkspaceId: null }; + input.task.context = { issueId: issue.id }; + } input.task.freshWorkspace = { mode: "fresh_isolated_workspace", retryIntentId: input.run.id, sourceRunId: input.run.id, retainedSourceIds: [] }; input.policy.requestedExecutionWorkspaceMode = input.policy.effectiveExecutionWorkspaceMode = "isolated_workspace"; input.config.mergedConfig = { workspaceStrategy: { type: "project_primary" } }; @@ -201,7 +203,9 @@ describe.skipIf(!support.supported)("heartbeat workspace preparation boundary", expect(read).toHaveBeenCalledWith(input.agent.companyId, input.run.id); expect(acquire).toHaveBeenCalledTimes(2); for (const [args] of acquire.mock.calls) expect(args.persistedExecutionWorkspace).toMatchObject({ id: input.run.id, - providerType: "adapter_managed", metadata: { runnerTaskFreshWorkspace: { attemptId: input.run.id } } }); + providerType: "adapter_managed", projectId: project ? issue.projectId : null, + projectWorkspaceId: project ? issue.projectWorkspaceId : null, + metadata: { runnerTaskFreshWorkspace: { attemptId: input.run.id } } }); const rows = await db.select().from(executionWorkspaces).where(eq(executionWorkspaces.companyId, input.agent.companyId)); expect(rows).toHaveLength(1); }); diff --git a/server/src/services/heartbeat/workspace-preparation.ts b/server/src/services/heartbeat/workspace-preparation.ts index 03ef7b10999..68ce0e2db5b 100644 --- a/server/src/services/heartbeat/workspace-preparation.ts +++ b/server/src/services/heartbeat/workspace-preparation.ts @@ -417,9 +417,11 @@ async function prepareHeartbeatWorkspaceLocked(db: Db, input: HeartbeatWorkspace baseRef: typeof parseObject(mergedConfig.workspaceStrategy).baseRef === "string" ? String(parseObject(mergedConfig.workspaceStrategy).baseRef) : null, }); - const freshTaskReservationId = providerFresh && pluginTaskWorkspace && !pluginPrepared && !resolvedWorkspace.projectId && issueId + const freshTaskReservationId = providerFresh && pluginTaskWorkspace && !pluginPrepared && issueId ? await reserveFreshPluginTaskWorkspace({ db, companyId: agent.companyId, issueId, runId: run.id, agentId: agent.id, - fresh: providerFresh, responsibleUserId, expectedBindingRevision: previousProviderBindingRevision, owner: taskWorkspaceOwner! }) : null; + fresh: providerFresh, responsibleUserId, expectedBindingRevision: previousProviderBindingRevision, owner: taskWorkspaceOwner!, + source: { projectId: resolvedWorkspace.projectId, projectWorkspaceId: resolvedWorkspace.workspaceId, + repoUrl: resolvedWorkspace.repoUrl, baseRef: resolvedWorkspace.repoRef } }) : null; const freshTaskReservation = freshTaskReservationId ? await executionWorkspacesSvc.getStoredById(agent.companyId, freshTaskReservationId) : null; const copyOnlyProvisioning = pluginPrepared ? true : await timeWorkspace("workspace_provision_policy", () => workspaceRequiresCopyOnlyProvisioning(db, { companyId: agent.companyId, diff --git a/server/src/services/provider-fresh-workspace.ts b/server/src/services/provider-fresh-workspace.ts index 04f418f7878..31a70b8cd57 100644 --- a/server/src/services/provider-fresh-workspace.ts +++ b/server/src/services/provider-fresh-workspace.ts @@ -92,7 +92,9 @@ export async function reserveFreshPluginTaskWorkspace(input: { db: Db; companyId: string; issueId: string; runId: string; agentId: string; responsibleUserId: string | null; fresh: PluginFreshWorkspace; expectedBindingRevision: number; owner: Record; + source?: { projectId: string | null; projectWorkspaceId: string | null; repoUrl: string | null; baseRef: string | null }; }) { + const source = input.source ?? { projectId: null, projectWorkspaceId: null, repoUrl: null, baseRef: null }; return input.db.transaction(async tx => { const [task] = await tx.select().from(issues).where(and(eq(issues.id, input.issueId), eq(issues.companyId, input.companyId))).for("update"); @@ -103,14 +105,15 @@ export async function reserveFreshPluginTaskWorkspace(input: { await assertTaskWorkspaceExecutionAccess(tx, { type: "agent", source: "agent_jwt", companyId: input.companyId, agentId: input.agentId, runId: input.runId, onBehalfOfUserId: input.responsibleUserId === "local-board" ? null : input.responsibleUserId, - }, input.companyId, { issueId: input.issueId }); + }, 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))); if (reserved) { if (task.executionWorkspaceId !== reserved.id || reserved.providerType !== "adapter_managed" || !isDeepStrictEqual(reserved.metadata?.runnerTaskFreshWorkspace, input.fresh) || - !isDeepStrictEqual(reserved.metadata?.runnerTaskWorkspaceOwner, input.owner)) + !isDeepStrictEqual(reserved.metadata?.runnerTaskWorkspaceOwner, input.owner) || + Object.entries(source).some(([key, value]) => reserved[key as keyof typeof source] !== value)) throw conflict("Task workspace binding changed; retry preparation"); return reserved.id; } @@ -118,7 +121,7 @@ export async function reserveFreshPluginTaskWorkspace(input: { throw conflict("Task workspace binding changed; retry preparation"); await notifyTerminalWorkspaceWork(tx, `workspace:${input.fresh.attemptId}`); await tx.insert(executionWorkspaces).values({ id: input.fresh.attemptId, - companyId: input.companyId, sourceIssueId: input.issueId, projectId: null, projectWorkspaceId: null, + companyId: input.companyId, sourceIssueId: input.issueId, ...source, mode: "isolated_workspace", strategyType: "project_primary", name: "Fresh task workspace", status: "active", providerType: "adapter_managed", cwd: `/workspaces/${input.fresh.attemptId}`, metadata: { runnerTaskFreshWorkspace: input.fresh, runnerTaskWorkspaceOwner: input.owner },