Fix execution workspace runtime lifecycle

2026-03-14 09:35:35 -05:00
parent 7a06a577ce
commit 3b25268c0b
7 changed files with 553 additions and 10 deletions
--- a/server/src/services/heartbeat.ts
+++ b/server/src/services/heartbeat.ts
@@ -139,6 +139,20 @@ export type ResolvedWorkspaceForRun = {
  warnings: string[];
 };

+type ProjectWorkspaceCandidate = {
+  id: string;
+};
+
+export function prioritizeProjectWorkspaceCandidatesForRun<T extends ProjectWorkspaceCandidate>(
+  rows: T[],
+  preferredWorkspaceId: string | null | undefined,
+): T[] {
+  if (!preferredWorkspaceId) return rows;
+  const preferredIndex = rows.findIndex((row) => row.id === preferredWorkspaceId);
+  if (preferredIndex <= 0) return rows;
+  return [rows[preferredIndex]!, ...rows.slice(0, preferredIndex), ...rows.slice(preferredIndex + 1)];
+}
+
 function readNonEmptyString(value: unknown): string | null {
  return typeof value === "string" && value.trim().length > 0 ? value : null;
 }
@@ -537,18 +551,25 @@ export function heartbeatService(db: Db) {
  ): Promise<ResolvedWorkspaceForRun> {
    const issueId = readNonEmptyString(context.issueId);
    const contextProjectId = readNonEmptyString(context.projectId);
-    const issueProjectId = issueId
+    const contextProjectWorkspaceId = readNonEmptyString(context.projectWorkspaceId);
+    const issueProjectRef = issueId
      ? await db
-          .select({ projectId: issues.projectId })
+          .select({
+            projectId: issues.projectId,
+            projectWorkspaceId: issues.projectWorkspaceId,
+          })
          .from(issues)
          .where(and(eq(issues.id, issueId), eq(issues.companyId, agent.companyId)))
-          .then((rows) => rows[0]?.projectId ?? null)
+          .then((rows) => rows[0] ?? null)
      : null;
+    const issueProjectId = issueProjectRef?.projectId ?? null;
+    const preferredProjectWorkspaceId =
+      issueProjectRef?.projectWorkspaceId ?? contextProjectWorkspaceId ?? null;
    const resolvedProjectId = issueProjectId ?? contextProjectId;
    const useProjectWorkspace = opts?.useProjectWorkspace !== false;
    const workspaceProjectId = useProjectWorkspace ? resolvedProjectId : null;

-    const projectWorkspaceRows = workspaceProjectId
+    const unorderedProjectWorkspaceRows = workspaceProjectId
      ? await db
          .select()
          .from(projectWorkspaces)
@@ -560,6 +581,10 @@ export function heartbeatService(db: Db) {
          )
          .orderBy(asc(projectWorkspaces.createdAt), asc(projectWorkspaces.id))
      : [];
+    const projectWorkspaceRows = prioritizeProjectWorkspaceCandidatesForRun(
+      unorderedProjectWorkspaceRows,
+      preferredProjectWorkspaceId,
+    );

    const workspaceHints = projectWorkspaceRows.map((workspace) => ({
      workspaceId: workspace.id,
@@ -569,11 +594,22 @@ export function heartbeatService(db: Db) {
    }));

    if (projectWorkspaceRows.length > 0) {
+      const preferredWorkspace = preferredProjectWorkspaceId
+        ? projectWorkspaceRows.find((workspace) => workspace.id === preferredProjectWorkspaceId) ?? null
+        : null;
      const missingProjectCwds: string[] = [];
      let hasConfiguredProjectCwd = false;
+      let preferredWorkspaceWarning: string | null = null;
+      if (preferredProjectWorkspaceId && !preferredWorkspace) {
+        preferredWorkspaceWarning =
+          `Selected project workspace "${preferredProjectWorkspaceId}" is not available on this project.`;
+      }
      for (const workspace of projectWorkspaceRows) {
        const projectCwd = readNonEmptyString(workspace.cwd);
        if (!projectCwd || projectCwd === REPO_ONLY_CWD_SENTINEL) {
+          if (preferredWorkspace?.id === workspace.id) {
+            preferredWorkspaceWarning = `Selected project workspace "${workspace.name}" has no local cwd configured.`;
+          }
          continue;
        }
        hasConfiguredProjectCwd = true;
@@ -590,15 +626,22 @@ export function heartbeatService(db: Db) {
            repoUrl: workspace.repoUrl,
            repoRef: workspace.repoRef,
            workspaceHints,
-            warnings: [],
+            warnings: preferredWorkspaceWarning ? [preferredWorkspaceWarning] : [],
          };
        }
+        if (preferredWorkspace?.id === workspace.id) {
+          preferredWorkspaceWarning =
+            `Selected project workspace path "${projectCwd}" is not available yet.`;
+        }
        missingProjectCwds.push(projectCwd);
      }

      const fallbackCwd = resolveDefaultAgentWorkspaceDir(agent.id);
      await fs.mkdir(fallbackCwd, { recursive: true });
      const warnings: string[] = [];
+      if (preferredWorkspaceWarning) {
+        warnings.push(preferredWorkspaceWarning);
+      }
      if (missingProjectCwds.length > 0) {
        const firstMissing = missingProjectCwds[0];
        const extraMissingCount = Math.max(0, missingProjectCwds.length - 1);
@@ -1464,6 +1507,7 @@ export function heartbeatService(db: Db) {
        },
        issue: issueRef,
        workspace: executionWorkspace,
+        executionWorkspaceId: persistedExecutionWorkspace?.id ?? issueRef?.executionWorkspaceId ?? null,
        config: resolvedConfig,
        adapterEnv,
        onLog,
--- a/server/src/services/workspace-runtime.ts
+++ b/server/src/services/workspace-runtime.ts
@@ -46,6 +46,7 @@ export interface RuntimeServiceRef {
  companyId: string;
  projectId: string | null;
  projectWorkspaceId: string | null;
+  executionWorkspaceId: string | null;
  issueId: string | null;
  serviceName: string;
  status: "starting" | "running" | "stopped" | "failed";
@@ -92,6 +93,17 @@ function stableStringify(value: unknown): string {
  return JSON.stringify(value);
 }

+function sanitizeRuntimeServiceBaseEnv(baseEnv: NodeJS.ProcessEnv): NodeJS.ProcessEnv {
+  const env: NodeJS.ProcessEnv = { ...baseEnv };
+  for (const key of Object.keys(env)) {
+    if (key.startsWith("PAPERCLIP_")) {
+      delete env[key];
+    }
+  }
+  delete env.DATABASE_URL;
+  return env;
+}
+
 function stableRuntimeServiceId(input: {
  adapterType: string;
  runId: string;
@@ -126,6 +138,7 @@ function toRuntimeServiceRef(record: RuntimeServiceRecord, overrides?: Partial<R
    companyId: record.companyId,
    projectId: record.projectId,
    projectWorkspaceId: record.projectWorkspaceId,
+    executionWorkspaceId: record.executionWorkspaceId,
    issueId: record.issueId,
    serviceName: record.serviceName,
    status: record.status,
@@ -330,6 +343,55 @@ async function provisionExecutionWorktree(input: {
  });
 }

+function buildExecutionWorkspaceCleanupEnv(input: {
+  workspace: {
+    cwd: string | null;
+    providerRef: string | null;
+    branchName: string | null;
+    repoUrl: string | null;
+    baseRef: string | null;
+    projectId: string | null;
+    projectWorkspaceId: string | null;
+    sourceIssueId: string | null;
+  };
+  projectWorkspaceCwd?: string | null;
+}) {
+  const env: NodeJS.ProcessEnv = { ...process.env };
+  env.PAPERCLIP_WORKSPACE_CWD = input.workspace.cwd ?? "";
+  env.PAPERCLIP_WORKSPACE_PATH = input.workspace.cwd ?? "";
+  env.PAPERCLIP_WORKSPACE_WORKTREE_PATH =
+    input.workspace.providerRef ?? input.workspace.cwd ?? "";
+  env.PAPERCLIP_WORKSPACE_BRANCH = input.workspace.branchName ?? "";
+  env.PAPERCLIP_WORKSPACE_BASE_CWD = input.projectWorkspaceCwd ?? "";
+  env.PAPERCLIP_WORKSPACE_REPO_ROOT = input.projectWorkspaceCwd ?? "";
+  env.PAPERCLIP_WORKSPACE_REPO_URL = input.workspace.repoUrl ?? "";
+  env.PAPERCLIP_WORKSPACE_REPO_REF = input.workspace.baseRef ?? "";
+  env.PAPERCLIP_PROJECT_ID = input.workspace.projectId ?? "";
+  env.PAPERCLIP_PROJECT_WORKSPACE_ID = input.workspace.projectWorkspaceId ?? "";
+  env.PAPERCLIP_ISSUE_ID = input.workspace.sourceIssueId ?? "";
+  return env;
+}
+
+async function resolveGitRepoRootForWorkspaceCleanup(
+  worktreePath: string,
+  projectWorkspaceCwd: string | null,
+): Promise<string | null> {
+  if (projectWorkspaceCwd) {
+    const resolvedProjectWorkspaceCwd = path.resolve(projectWorkspaceCwd);
+    const gitDir = await runGit(["rev-parse", "--git-common-dir"], resolvedProjectWorkspaceCwd)
+      .catch(() => null);
+    if (gitDir) {
+      const resolvedGitDir = path.resolve(resolvedProjectWorkspaceCwd, gitDir);
+      return path.dirname(resolvedGitDir);
+    }
+  }
+
+  const gitDir = await runGit(["rev-parse", "--git-common-dir"], worktreePath).catch(() => null);
+  if (!gitDir) return null;
+  const resolvedGitDir = path.resolve(worktreePath, gitDir);
+  return path.dirname(resolvedGitDir);
+}
+
 export async function realizeExecutionWorkspace(input: {
  base: ExecutionWorkspaceInput;
  config: Record<string, unknown>;
@@ -418,6 +480,98 @@ export async function realizeExecutionWorkspace(input: {
  };
 }

+export async function cleanupExecutionWorkspaceArtifacts(input: {
+  workspace: {
+    id: string;
+    cwd: string | null;
+    providerType: string;
+    providerRef: string | null;
+    branchName: string | null;
+    repoUrl: string | null;
+    baseRef: string | null;
+    projectId: string | null;
+    projectWorkspaceId: string | null;
+    sourceIssueId: string | null;
+    metadata?: Record<string, unknown> | null;
+  };
+  projectWorkspace?: {
+    cwd: string | null;
+    cleanupCommand: string | null;
+  } | null;
+  teardownCommand?: string | null;
+}) {
+  const warnings: string[] = [];
+  const workspacePath = input.workspace.providerRef ?? input.workspace.cwd;
+  const cleanupEnv = buildExecutionWorkspaceCleanupEnv({
+    workspace: input.workspace,
+    projectWorkspaceCwd: input.projectWorkspace?.cwd ?? null,
+  });
+  const createdByRuntime = input.workspace.metadata?.createdByRuntime === true;
+  const cleanupCommands = [
+    input.projectWorkspace?.cleanupCommand ?? null,
+    input.teardownCommand ?? null,
+  ]
+    .map((value) => asString(value, "").trim())
+    .filter(Boolean);
+
+  for (const command of cleanupCommands) {
+    try {
+      await runWorkspaceCommand({
+        command,
+        cwd: workspacePath ?? input.projectWorkspace?.cwd ?? process.cwd(),
+        env: cleanupEnv,
+        label: `Execution workspace cleanup command "${command}"`,
+      });
+    } catch (err) {
+      warnings.push(err instanceof Error ? err.message : String(err));
+    }
+  }
+
+  if (input.workspace.providerType === "git_worktree" && workspacePath) {
+    const worktreeExists = await directoryExists(workspacePath);
+    if (worktreeExists) {
+      const repoRoot = await resolveGitRepoRootForWorkspaceCleanup(
+        workspacePath,
+        input.projectWorkspace?.cwd ?? null,
+      );
+      if (!repoRoot) {
+        warnings.push(`Could not resolve git repo root for "${workspacePath}".`);
+      } else {
+        try {
+          await runGit(["worktree", "remove", "--force", workspacePath], repoRoot);
+        } catch (err) {
+          warnings.push(err instanceof Error ? err.message : String(err));
+        }
+        if (createdByRuntime && input.workspace.branchName) {
+          try {
+            await runGit(["branch", "-D", input.workspace.branchName], repoRoot);
+          } catch (err) {
+            warnings.push(err instanceof Error ? err.message : String(err));
+          }
+        }
+      }
+    }
+  } else if (input.workspace.providerType === "local_fs" && createdByRuntime && workspacePath) {
+    const projectWorkspaceCwd = input.projectWorkspace?.cwd ? path.resolve(input.projectWorkspace.cwd) : null;
+    const resolvedWorkspacePath = path.resolve(workspacePath);
+    if (projectWorkspaceCwd && resolvedWorkspacePath === projectWorkspaceCwd) {
+      warnings.push(`Refusing to remove shared project workspace "${workspacePath}".`);
+    } else {
+      await fs.rm(resolvedWorkspacePath, { recursive: true, force: true });
+    }
+  }
+
+  const cleaned =
+    !workspacePath ||
+    !(await directoryExists(workspacePath));
+
+  return {
+    cleanedPath: workspacePath,
+    cleaned,
+    warnings,
+  };
+}
+
 async function allocatePort(): Promise<number> {
  return await new Promise<number>((resolve, reject) => {
    const server = net.createServer();
@@ -521,6 +675,7 @@ function toPersistedWorkspaceRuntimeService(record: RuntimeServiceRecord): typeo
    companyId: record.companyId,
    projectId: record.projectId,
    projectWorkspaceId: record.projectWorkspaceId,
+    executionWorkspaceId: record.executionWorkspaceId,
    issueId: record.issueId,
    scopeType: record.scopeType,
    scopeId: record.scopeId,
@@ -556,6 +711,7 @@ async function persistRuntimeServiceRecord(db: Db | undefined, record: RuntimeSe
      set: {
        projectId: values.projectId,
        projectWorkspaceId: values.projectWorkspaceId,
+        executionWorkspaceId: values.executionWorkspaceId,
        issueId: values.issueId,
        scopeType: values.scopeType,
        scopeId: values.scopeId,
@@ -593,6 +749,7 @@ export function normalizeAdapterManagedRuntimeServices(input: {
  agent: ExecutionWorkspaceAgentRef;
  issue: ExecutionWorkspaceIssueRef | null;
  workspace: RealizedExecutionWorkspace;
+  executionWorkspaceId?: string | null;
  reports: AdapterRuntimeServiceReport[];
  now?: Date;
 }): RuntimeServiceRef[] {
@@ -629,6 +786,7 @@ export function normalizeAdapterManagedRuntimeServices(input: {
      companyId: input.agent.companyId,
      projectId: report.projectId ?? input.workspace.projectId,
      projectWorkspaceId: report.projectWorkspaceId ?? input.workspace.workspaceId,
+      executionWorkspaceId: input.executionWorkspaceId ?? null,
      issueId: report.issueId ?? input.issue?.id ?? null,
      serviceName,
      status,
@@ -660,6 +818,7 @@ async function startLocalRuntimeService(input: {
  agent: ExecutionWorkspaceAgentRef;
  issue: ExecutionWorkspaceIssueRef | null;
  workspace: RealizedExecutionWorkspace;
+  executionWorkspaceId?: string | null;
  adapterEnv: Record<string, string>;
  service: Record<string, unknown>;
  onLog?: (stream: "stdout" | "stderr", chunk: string) => Promise<void>;
@@ -683,7 +842,10 @@ async function startLocalRuntimeService(input: {
    port,
  });
  const serviceCwd = resolveConfiguredPath(renderTemplate(serviceCwdTemplate, templateData), input.workspace.cwd);
-  const env: Record<string, string> = { ...process.env, ...input.adapterEnv } as Record<string, string>;
+  const env: Record<string, string> = {
+    ...sanitizeRuntimeServiceBaseEnv(process.env),
+    ...input.adapterEnv,
+  } as Record<string, string>;
  for (const [key, value] of Object.entries(envConfig)) {
    if (typeof value === "string") {
      env[key] = renderTemplate(value, templateData);
@@ -735,6 +897,7 @@ async function startLocalRuntimeService(input: {
    companyId: input.agent.companyId,
    projectId: input.workspace.projectId,
    projectWorkspaceId: input.workspace.workspaceId,
+    executionWorkspaceId: input.executionWorkspaceId ?? null,
    issueId: input.issue?.id ?? null,
    serviceName,
    status: "running",
@@ -791,6 +954,28 @@ async function stopRuntimeService(serviceId: string) {
  await persistRuntimeServiceRecord(record.db, record);
 }

+async function markPersistedRuntimeServicesStoppedForExecutionWorkspace(input: {
+  db: Db;
+  executionWorkspaceId: string;
+}) {
+  const now = new Date();
+  await input.db
+    .update(workspaceRuntimeServices)
+    .set({
+      status: "stopped",
+      healthStatus: "unknown",
+      stoppedAt: now,
+      lastUsedAt: now,
+      updatedAt: now,
+    })
+    .where(
+      and(
+        eq(workspaceRuntimeServices.executionWorkspaceId, input.executionWorkspaceId),
+        inArray(workspaceRuntimeServices.status, ["starting", "running"]),
+      ),
+    );
+}
+
 function registerRuntimeService(db: Db | undefined, record: RuntimeServiceRecord) {
  record.db = db;
  runtimeServicesById.set(record.id, record);
@@ -820,6 +1005,7 @@ export async function ensureRuntimeServicesForRun(input: {
  agent: ExecutionWorkspaceAgentRef;
  issue: ExecutionWorkspaceIssueRef | null;
  workspace: RealizedExecutionWorkspace;
+  executionWorkspaceId?: string | null;
  config: Record<string, unknown>;
  adapterEnv: Record<string, string>;
  onLog?: (stream: "stdout" | "stderr", chunk: string) => Promise<void>;
@@ -871,6 +1057,7 @@ export async function ensureRuntimeServicesForRun(input: {
        agent: input.agent,
        issue: input.issue,
        workspace: input.workspace,
+        executionWorkspaceId: input.executionWorkspaceId,
        adapterEnv: input.adapterEnv,
        service,
        onLog: input.onLog,
@@ -911,6 +1098,32 @@ export async function releaseRuntimeServicesForRun(runId: string) {
  }
 }

+export async function stopRuntimeServicesForExecutionWorkspace(input: {
+  db?: Db;
+  executionWorkspaceId: string;
+  workspaceCwd?: string | null;
+}) {
+  const normalizedWorkspaceCwd = input.workspaceCwd ? path.resolve(input.workspaceCwd) : null;
+  const matchingServiceIds = Array.from(runtimeServicesById.values())
+    .filter((record) => {
+      if (record.executionWorkspaceId === input.executionWorkspaceId) return true;
+      if (!normalizedWorkspaceCwd || !record.cwd) return false;
+      return path.resolve(record.cwd).startsWith(normalizedWorkspaceCwd);
+    })
+    .map((record) => record.id);
+
+  for (const serviceId of matchingServiceIds) {
+    await stopRuntimeService(serviceId);
+  }
+
+  if (input.db) {
+    await markPersistedRuntimeServicesStoppedForExecutionWorkspace({
+      db: input.db,
+      executionWorkspaceId: input.executionWorkspaceId,
+    });
+  }
+}
+
 export async function listWorkspaceRuntimeServicesForProjectWorkspaces(
  db: Db,
  companyId: string,
@@ -978,6 +1191,7 @@ export async function persistAdapterManagedRuntimeServices(input: {
  agent: ExecutionWorkspaceAgentRef;
  issue: ExecutionWorkspaceIssueRef | null;
  workspace: RealizedExecutionWorkspace;
+  executionWorkspaceId?: string | null;
  reports: AdapterRuntimeServiceReport[];
 }) {
  const refs = normalizeAdapterManagedRuntimeServices(input);
@@ -1000,6 +1214,7 @@ export async function persistAdapterManagedRuntimeServices(input: {
        companyId: ref.companyId,
        projectId: ref.projectId,
        projectWorkspaceId: ref.projectWorkspaceId,
+        executionWorkspaceId: ref.executionWorkspaceId,
        issueId: ref.issueId,
        scopeType: ref.scopeType,
        scopeId: ref.scopeId,
@@ -1028,6 +1243,7 @@ export async function persistAdapterManagedRuntimeServices(input: {
        set: {
          projectId: ref.projectId,
          projectWorkspaceId: ref.projectWorkspaceId,
+          executionWorkspaceId: ref.executionWorkspaceId,
          issueId: ref.issueId,
          scopeType: ref.scopeType,
          scopeId: ref.scopeId,