Address remaining Greptile workspace review

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<Record<string, unknown>> = {}) {
+  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<unknown>) => 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<unknown>) => 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");
+  });
+});

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
+    const issueRef = issueContext
-      ? await db
+      ? {
-          .select({
+          id: issueContext.id,
-            id: issues.id,
+          identifier: issueContext.identifier,
-            identifier: issues.identifier,
+          title: issueContext.title,
-            title: issues.title,
+          projectId: issueContext.projectId,
-            projectId: issues.projectId,
+          projectWorkspaceId: issueContext.projectWorkspaceId,
-            projectWorkspaceId: issues.projectWorkspaceId,
+          executionWorkspaceId: issueContext.executionWorkspaceId,
-            executionWorkspaceId: issues.executionWorkspaceId,
+          executionWorkspacePreference: issueContext.executionWorkspacePreference,
-            executionWorkspacePreference: issues.executionWorkspacePreference,
+        }
-          })
-          .from(issues)
-          .where(and(eq(issues.id, issueId), eq(issues.companyId, agent.companyId)))
-          .then((rows) => rows[0] ?? null)
       : null;
     const existingExecutionWorkspace =
       issueRef?.executionWorkspaceId ? await executionWorkspacesSvc.getById(issueRef.executionWorkspaceId) : null;
@@ -1748,7 +1748,9 @@ export function heartbeatService(db: Db) {
       issueRef?.executionWorkspacePreference === "reuse_existing" &&
       existingExecutionWorkspace &&
       existingExecutionWorkspace.status !== "archived";
-    const persistedExecutionWorkspace = shouldReuseExisting && existingExecutionWorkspace
+    let persistedExecutionWorkspace = null;
+    try {
+      persistedExecutionWorkspace = shouldReuseExisting && existingExecutionWorkspace
         ? await executionWorkspacesSvc.update(existingExecutionWorkspace.id, {
             cwd: executionWorkspace.cwd,
             repoUrl: executionWorkspace.repoUrl,
@@ -1795,7 +1797,59 @@ export function heartbeatService(db: Db) {
               },
             })
           : 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,

server/src/services/work-products.ts

@@ -51,13 +51,20 @@ export function workProductService(db: Db) {
     },
 
     createForIssue: async (issueId: string, companyId: string, data: Omit<typeof issueWorkProducts.$inferInsert, "issueId" | "companyId">) => {
+      const row = await db.transaction(async (tx) => {
         if (data.isPrimary) {
-        await db
+          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)));
+            .where(
+              and(
+                eq(issueWorkProducts.companyId, companyId),
+                eq(issueWorkProducts.issueId, issueId),
+                eq(issueWorkProducts.type, data.type),
+              ),
+            );
         }
-      const row = await db
+        return await tx
           .insert(issueWorkProducts)
           .values({
             ...data,
@@ -66,11 +73,13 @@ export function workProductService(db: Db) {
           })
           .returning()
           .then((rows) => rows[0] ?? null);
+      });
       return row ? toIssueWorkProduct(row) : null;
     },
 
     update: async (id: string, patch: Partial<typeof issueWorkProducts.$inferInsert>) => {
-      const existing = await db
+      const row = await db.transaction(async (tx) => {
+        const existing = await tx
           .select()
           .from(issueWorkProducts)
           .where(eq(issueWorkProducts.id, id))
@@ -78,7 +87,7 @@ export function workProductService(db: Db) {
         if (!existing) return null;
 
         if (patch.isPrimary === true) {
-        await db
+          await tx
             .update(issueWorkProducts)
             .set({ isPrimary: false, updatedAt: new Date() })
             .where(
@@ -90,12 +99,13 @@ export function workProductService(db: Db) {
             );
         }
 
-      const row = await db
+        return await tx
           .update(issueWorkProducts)
           .set({ ...patch, updatedAt: new Date() })
           .where(eq(issueWorkProducts.id, id))
           .returning()
           .then((rows) => rows[0] ?? null);
+      });
       return row ? toIssueWorkProduct(row) : null;
     },