Address remaining Greptile workspace review

2026-03-17 10:12:44 -05:00
parent 4a5f6ec00a
commit 108792ac51
3 changed files with 252 additions and 93 deletions
--- a/server/src/tests/work-products.test.ts
+++ 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<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");
+  });
+});
--- 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,
--- 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<typeof issueWorkProducts.$inferInsert, "issueId" | "companyId">) => {
-      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<typeof issueWorkProducts.$inferInsert>) => {
-      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;
    },