From 108792ac512fbd912b04ef12e73d5d77767b55e8 Mon Sep 17 00:00:00 2001 From: Dotta Date: Tue, 17 Mar 2026 10:12:44 -0500 Subject: [PATCH] Address remaining Greptile workspace review --- server/src/__tests__/work-products.test.ts | 95 ++++++++++++ server/src/services/heartbeat.ts | 162 ++++++++++++++------- server/src/services/work-products.ts | 88 ++++++----- 3 files changed, 252 insertions(+), 93 deletions(-) create mode 100644 server/src/__tests__/work-products.test.ts diff --git a/server/src/__tests__/work-products.test.ts b/server/src/__tests__/work-products.test.ts new file mode 100644 index 00000000..93e4a0fc --- /dev/null +++ b/server/src/__tests__/work-products.test.ts @@ -0,0 +1,95 @@ +import { describe, expect, it, vi } from "vitest"; +import { workProductService } from "../services/work-products.ts"; + +function createWorkProductRow(overrides: Partial> = {}) { + const now = new Date("2026-03-17T00:00:00.000Z"); + return { + id: "work-product-1", + companyId: "company-1", + projectId: "project-1", + issueId: "issue-1", + executionWorkspaceId: null, + runtimeServiceId: null, + type: "pull_request", + provider: "github", + externalId: null, + title: "PR 1", + url: "https://example.com/pr/1", + status: "open", + reviewState: "draft", + isPrimary: true, + healthStatus: "unknown", + summary: null, + metadata: null, + createdByRunId: null, + createdAt: now, + updatedAt: now, + ...overrides, + }; +} + +describe("workProductService", () => { + it("uses a transaction when creating a new primary work product", async () => { + const updatedWhere = vi.fn(async () => undefined); + const updateSet = vi.fn(() => ({ where: updatedWhere })); + const txUpdate = vi.fn(() => ({ set: updateSet })); + + const insertedRow = createWorkProductRow(); + const insertReturning = vi.fn(async () => [insertedRow]); + const insertValues = vi.fn(() => ({ returning: insertReturning })); + const txInsert = vi.fn(() => ({ values: insertValues })); + + const tx = { + update: txUpdate, + insert: txInsert, + }; + const transaction = vi.fn(async (callback: (input: typeof tx) => Promise) => await callback(tx)); + + const svc = workProductService({ transaction } as any); + const result = await svc.createForIssue("issue-1", "company-1", { + type: "pull_request", + provider: "github", + title: "PR 1", + status: "open", + reviewState: "draft", + isPrimary: true, + }); + + expect(transaction).toHaveBeenCalledTimes(1); + expect(txUpdate).toHaveBeenCalledTimes(1); + expect(txInsert).toHaveBeenCalledTimes(1); + expect(result?.id).toBe("work-product-1"); + }); + + it("uses a transaction when promoting an existing work product to primary", async () => { + const existingRow = createWorkProductRow({ isPrimary: false }); + + const selectWhere = vi.fn(async () => [existingRow]); + const selectFrom = vi.fn(() => ({ where: selectWhere })); + const txSelect = vi.fn(() => ({ from: selectFrom })); + + const updateReturning = vi + .fn() + .mockResolvedValue([createWorkProductRow({ reviewState: "ready_for_review" })]); + const updateWhere = vi.fn(() => ({ returning: updateReturning })); + const updateSet = vi.fn(() => ({ where: updateWhere })); + const txUpdate = vi.fn(() => ({ set: updateSet })); + + const tx = { + select: txSelect, + update: txUpdate, + }; + const transaction = vi.fn(async (callback: (input: typeof tx) => Promise) => await callback(tx)); + + const svc = workProductService({ transaction } as any); + const result = await svc.update("work-product-1", { + isPrimary: true, + reviewState: "ready_for_review", + }); + + expect(transaction).toHaveBeenCalledTimes(1); + expect(txSelect).toHaveBeenCalledTimes(1); + expect(txUpdate).toHaveBeenCalledTimes(2); + expect(result?.reviewState).toBe("ready_for_review"); + }); +}); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 76bac35c..43c5c05f 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -31,6 +31,7 @@ import { resolveDefaultAgentWorkspaceDir, resolveManagedProjectWorkspaceDir } fr import { summarizeHeartbeatRunResultJson } from "./heartbeat-run-summary.js"; import { buildWorkspaceReadyComment, + cleanupExecutionWorkspaceArtifacts, ensureRuntimeServicesForRun, persistAdapterManagedRuntimeServices, realizeExecutionWorkspace, @@ -1629,9 +1630,12 @@ export function heartbeatService(db: Db) { const taskKey = deriveTaskKey(context, null); const sessionCodec = getAdapterSessionCodec(agent.adapterType); const issueId = readNonEmptyString(context.issueId); - const issueAssigneeConfig = issueId + const issueContext = issueId ? await db .select({ + id: issues.id, + identifier: issues.identifier, + title: issues.title, projectId: issues.projectId, projectWorkspaceId: issues.projectWorkspaceId, executionWorkspaceId: issues.executionWorkspaceId, @@ -1645,17 +1649,17 @@ export function heartbeatService(db: Db) { .then((rows) => rows[0] ?? null) : null; const issueAssigneeOverrides = - issueAssigneeConfig && issueAssigneeConfig.assigneeAgentId === agent.id + issueContext && issueContext.assigneeAgentId === agent.id ? parseIssueAssigneeAdapterOverrides( - issueAssigneeConfig.assigneeAdapterOverrides, + issueContext.assigneeAdapterOverrides, ) : null; const isolatedWorkspacesEnabled = (await instanceSettings.getExperimental()).enableIsolatedWorkspaces; const issueExecutionWorkspaceSettings = isolatedWorkspacesEnabled - ? parseIssueExecutionWorkspaceSettings(issueAssigneeConfig?.executionWorkspaceSettings) + ? parseIssueExecutionWorkspaceSettings(issueContext?.executionWorkspaceSettings) : null; const contextProjectId = readNonEmptyString(context.projectId); - const executionProjectId = issueAssigneeConfig?.projectId ?? contextProjectId; + const executionProjectId = issueContext?.projectId ?? contextProjectId; const projectExecutionWorkspacePolicy = executionProjectId ? await db .select({ executionWorkspacePolicy: projects.executionWorkspacePolicy }) @@ -1702,20 +1706,16 @@ export function heartbeatService(db: Db) { agent.companyId, mergedConfig, ); - const issueRef = issueId - ? await db - .select({ - id: issues.id, - identifier: issues.identifier, - title: issues.title, - projectId: issues.projectId, - projectWorkspaceId: issues.projectWorkspaceId, - executionWorkspaceId: issues.executionWorkspaceId, - executionWorkspacePreference: issues.executionWorkspacePreference, - }) - .from(issues) - .where(and(eq(issues.id, issueId), eq(issues.companyId, agent.companyId))) - .then((rows) => rows[0] ?? null) + const issueRef = issueContext + ? { + id: issueContext.id, + identifier: issueContext.identifier, + title: issueContext.title, + projectId: issueContext.projectId, + projectWorkspaceId: issueContext.projectWorkspaceId, + executionWorkspaceId: issueContext.executionWorkspaceId, + executionWorkspacePreference: issueContext.executionWorkspacePreference, + } : null; const existingExecutionWorkspace = issueRef?.executionWorkspaceId ? await executionWorkspacesSvc.getById(issueRef.executionWorkspaceId) : null; @@ -1748,54 +1748,108 @@ export function heartbeatService(db: Db) { issueRef?.executionWorkspacePreference === "reuse_existing" && existingExecutionWorkspace && existingExecutionWorkspace.status !== "archived"; - const persistedExecutionWorkspace = shouldReuseExisting && existingExecutionWorkspace - ? await executionWorkspacesSvc.update(existingExecutionWorkspace.id, { - cwd: executionWorkspace.cwd, - repoUrl: executionWorkspace.repoUrl, - baseRef: executionWorkspace.repoRef, - branchName: executionWorkspace.branchName, - providerType: executionWorkspace.strategy === "git_worktree" ? "git_worktree" : "local_fs", - providerRef: executionWorkspace.worktreePath, - status: "active", - lastUsedAt: new Date(), - metadata: { - ...(existingExecutionWorkspace.metadata ?? {}), - source: executionWorkspace.source, - createdByRuntime: executionWorkspace.created, - }, - }) - : resolvedProjectId - ? await executionWorkspacesSvc.create({ - companyId: agent.companyId, - projectId: resolvedProjectId, - projectWorkspaceId: resolvedProjectWorkspaceId, - sourceIssueId: issueRef?.id ?? null, - mode: - executionWorkspaceMode === "isolated_workspace" - ? "isolated_workspace" - : executionWorkspaceMode === "operator_branch" - ? "operator_branch" - : executionWorkspaceMode === "agent_default" - ? "adapter_managed" - : "shared_workspace", - strategyType: executionWorkspace.strategy === "git_worktree" ? "git_worktree" : "project_primary", - name: executionWorkspace.branchName ?? issueRef?.identifier ?? `workspace-${agent.id.slice(0, 8)}`, - status: "active", + let persistedExecutionWorkspace = null; + try { + persistedExecutionWorkspace = shouldReuseExisting && existingExecutionWorkspace + ? await executionWorkspacesSvc.update(existingExecutionWorkspace.id, { cwd: executionWorkspace.cwd, repoUrl: executionWorkspace.repoUrl, baseRef: executionWorkspace.repoRef, branchName: executionWorkspace.branchName, providerType: executionWorkspace.strategy === "git_worktree" ? "git_worktree" : "local_fs", providerRef: executionWorkspace.worktreePath, + status: "active", lastUsedAt: new Date(), - openedAt: new Date(), metadata: { + ...(existingExecutionWorkspace.metadata ?? {}), source: executionWorkspace.source, createdByRuntime: executionWorkspace.created, }, }) - : null; + : resolvedProjectId + ? await executionWorkspacesSvc.create({ + companyId: agent.companyId, + projectId: resolvedProjectId, + projectWorkspaceId: resolvedProjectWorkspaceId, + sourceIssueId: issueRef?.id ?? null, + mode: + executionWorkspaceMode === "isolated_workspace" + ? "isolated_workspace" + : executionWorkspaceMode === "operator_branch" + ? "operator_branch" + : executionWorkspaceMode === "agent_default" + ? "adapter_managed" + : "shared_workspace", + strategyType: executionWorkspace.strategy === "git_worktree" ? "git_worktree" : "project_primary", + name: executionWorkspace.branchName ?? issueRef?.identifier ?? `workspace-${agent.id.slice(0, 8)}`, + status: "active", + cwd: executionWorkspace.cwd, + repoUrl: executionWorkspace.repoUrl, + baseRef: executionWorkspace.repoRef, + branchName: executionWorkspace.branchName, + providerType: executionWorkspace.strategy === "git_worktree" ? "git_worktree" : "local_fs", + providerRef: executionWorkspace.worktreePath, + lastUsedAt: new Date(), + openedAt: new Date(), + metadata: { + source: executionWorkspace.source, + createdByRuntime: executionWorkspace.created, + }, + }) + : null; + } catch (error) { + if (executionWorkspace.created) { + try { + await cleanupExecutionWorkspaceArtifacts({ + workspace: { + id: existingExecutionWorkspace?.id ?? `transient-${run.id}`, + cwd: executionWorkspace.cwd, + providerType: executionWorkspace.strategy === "git_worktree" ? "git_worktree" : "local_fs", + providerRef: executionWorkspace.worktreePath, + branchName: executionWorkspace.branchName, + repoUrl: executionWorkspace.repoUrl, + baseRef: executionWorkspace.repoRef, + projectId: resolvedProjectId, + projectWorkspaceId: resolvedProjectWorkspaceId, + sourceIssueId: issueRef?.id ?? null, + metadata: { + createdByRuntime: true, + source: executionWorkspace.source, + }, + }, + projectWorkspace: { + cwd: resolvedWorkspace.cwd, + cleanupCommand: null, + }, + teardownCommand: projectExecutionWorkspacePolicy?.workspaceStrategy?.teardownCommand ?? null, + recorder: workspaceOperationRecorder, + }); + } catch (cleanupError) { + logger.warn( + { + runId: run.id, + issueId, + executionWorkspaceCwd: executionWorkspace.cwd, + cleanupError: cleanupError instanceof Error ? cleanupError.message : String(cleanupError), + }, + "Failed to cleanup realized execution workspace after persistence failure", + ); + } + } + throw error; + } await workspaceOperationRecorder.attachExecutionWorkspaceId(persistedExecutionWorkspace?.id ?? null); + if ( + existingExecutionWorkspace && + persistedExecutionWorkspace && + existingExecutionWorkspace.id !== persistedExecutionWorkspace.id && + existingExecutionWorkspace.status === "active" + ) { + await executionWorkspacesSvc.update(existingExecutionWorkspace.id, { + status: "idle", + cleanupReason: null, + }); + } if (issueId && persistedExecutionWorkspace && issueRef?.executionWorkspaceId !== persistedExecutionWorkspace.id) { await issuesSvc.update(issueId, { executionWorkspaceId: persistedExecutionWorkspace.id, diff --git a/server/src/services/work-products.ts b/server/src/services/work-products.ts index e5e4cc63..6c5365fd 100644 --- a/server/src/services/work-products.ts +++ b/server/src/services/work-products.ts @@ -51,51 +51,61 @@ export function workProductService(db: Db) { }, createForIssue: async (issueId: string, companyId: string, data: Omit) => { - if (data.isPrimary) { - await db - .update(issueWorkProducts) - .set({ isPrimary: false, updatedAt: new Date() }) - .where(and(eq(issueWorkProducts.companyId, companyId), eq(issueWorkProducts.issueId, issueId), eq(issueWorkProducts.type, data.type))); - } - const row = await db - .insert(issueWorkProducts) - .values({ - ...data, - companyId, - issueId, - }) - .returning() - .then((rows) => rows[0] ?? null); + const row = await db.transaction(async (tx) => { + if (data.isPrimary) { + await tx + .update(issueWorkProducts) + .set({ isPrimary: false, updatedAt: new Date() }) + .where( + and( + eq(issueWorkProducts.companyId, companyId), + eq(issueWorkProducts.issueId, issueId), + eq(issueWorkProducts.type, data.type), + ), + ); + } + return await tx + .insert(issueWorkProducts) + .values({ + ...data, + companyId, + issueId, + }) + .returning() + .then((rows) => rows[0] ?? null); + }); return row ? toIssueWorkProduct(row) : null; }, update: async (id: string, patch: Partial) => { - const existing = await db - .select() - .from(issueWorkProducts) - .where(eq(issueWorkProducts.id, id)) - .then((rows) => rows[0] ?? null); - if (!existing) return null; + const row = await db.transaction(async (tx) => { + const existing = await tx + .select() + .from(issueWorkProducts) + .where(eq(issueWorkProducts.id, id)) + .then((rows) => rows[0] ?? null); + if (!existing) return null; - if (patch.isPrimary === true) { - await db + if (patch.isPrimary === true) { + await tx + .update(issueWorkProducts) + .set({ isPrimary: false, updatedAt: new Date() }) + .where( + and( + eq(issueWorkProducts.companyId, existing.companyId), + eq(issueWorkProducts.issueId, existing.issueId), + eq(issueWorkProducts.type, existing.type), + ), + ); + } + + return await tx .update(issueWorkProducts) - .set({ isPrimary: false, updatedAt: new Date() }) - .where( - and( - eq(issueWorkProducts.companyId, existing.companyId), - eq(issueWorkProducts.issueId, existing.issueId), - eq(issueWorkProducts.type, existing.type), - ), - ); - } - - const row = await db - .update(issueWorkProducts) - .set({ ...patch, updatedAt: new Date() }) - .where(eq(issueWorkProducts.id, id)) - .returning() - .then((rows) => rows[0] ?? null); + .set({ ...patch, updatedAt: new Date() }) + .where(eq(issueWorkProducts.id, id)) + .returning() + .then((rows) => rows[0] ?? null); + }); return row ? toIssueWorkProduct(row) : null; },