diff --git a/__tests__/update-job-status-worktree.test.ts b/__tests__/update-job-status-worktree.test.ts new file mode 100644 index 0000000..5b084b4 --- /dev/null +++ b/__tests__/update-job-status-worktree.test.ts @@ -0,0 +1,84 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' + +vi.mock('../src/git/worktree.js', () => ({ + removeWorktreeForJob: vi.fn(), +})) + +vi.mock('../src/tools/wait-for-job.js', async (importOriginal) => { + const original = await importOriginal() + return { + ...original, + resolveRepoRoot: vi.fn(), + } +}) + +import { removeWorktreeForJob } from '../src/git/worktree.js' +import { resolveRepoRoot } from '../src/tools/wait-for-job.js' +import { cleanupWorktreeForTerminalStatus } from '../src/tools/update-job-status.js' + +const mockRemove = removeWorktreeForJob as ReturnType +const mockResolve = resolveRepoRoot as ReturnType + +beforeEach(() => { + vi.clearAllMocks() +}) + +describe('cleanupWorktreeForTerminalStatus', () => { + it('calls removeWorktreeForJob with keepBranch=true when done and branch set', async () => { + mockResolve.mockResolvedValue('/repos/my-project') + mockRemove.mockResolvedValue({ removed: true }) + + await cleanupWorktreeForTerminalStatus('prod-001', 'job-abc', 'done', 'feat/job-abc') + + expect(mockRemove).toHaveBeenCalledWith({ + repoRoot: '/repos/my-project', + jobId: 'job-abc', + keepBranch: true, + }) + }) + + it('calls removeWorktreeForJob with keepBranch=false when done but no branch', async () => { + mockResolve.mockResolvedValue('/repos/my-project') + mockRemove.mockResolvedValue({ removed: true }) + + await cleanupWorktreeForTerminalStatus('prod-001', 'job-abc', 'done', undefined) + + expect(mockRemove).toHaveBeenCalledWith({ + repoRoot: '/repos/my-project', + jobId: 'job-abc', + keepBranch: false, + }) + }) + + it('calls removeWorktreeForJob with keepBranch=false when failed', async () => { + mockResolve.mockResolvedValue('/repos/my-project') + mockRemove.mockResolvedValue({ removed: true }) + + await cleanupWorktreeForTerminalStatus('prod-001', 'job-abc', 'failed', 'feat/job-abc') + + expect(mockRemove).toHaveBeenCalledWith({ + repoRoot: '/repos/my-project', + jobId: 'job-abc', + keepBranch: false, + }) + }) + + it('skips cleanup and does not throw when no repoRoot configured', async () => { + mockResolve.mockResolvedValue(null) + + await expect( + cleanupWorktreeForTerminalStatus('prod-no-root', 'job-abc', 'done', undefined), + ).resolves.toBeUndefined() + + expect(mockRemove).not.toHaveBeenCalled() + }) + + it('does not throw when removeWorktreeForJob fails (best-effort)', async () => { + mockResolve.mockResolvedValue('/repos/my-project') + mockRemove.mockRejectedValue(new Error('git error')) + + await expect( + cleanupWorktreeForTerminalStatus('prod-001', 'job-abc', 'done', 'feat/job-abc'), + ).resolves.toBeUndefined() + }) +}) diff --git a/src/tools/update-job-status.ts b/src/tools/update-job-status.ts index 21ec566..05e5a95 100644 --- a/src/tools/update-job-status.ts +++ b/src/tools/update-job-status.ts @@ -8,6 +8,8 @@ import { Client } from 'pg' import { prisma } from '../prisma.js' import { requireWriteAccess } from '../auth.js' import { toolJson, toolError, withToolErrors } from '../errors.js' +import { removeWorktreeForJob } from '../git/worktree.js' +import { resolveRepoRoot } from './wait-for-job.js' const inputSchema = z.object({ job_id: z.string().min(1), @@ -17,6 +19,24 @@ const inputSchema = z.object({ error: z.string().max(2_000).optional(), }) +export async function cleanupWorktreeForTerminalStatus( + productId: string, + jobId: string, + status: 'done' | 'failed', + branch: string | undefined, +): Promise { + const repoRoot = await resolveRepoRoot(productId) + if (!repoRoot) return + + // Keep branch when job is done and a branch was reported (agent pushed) + const keepBranch = status === 'done' && branch !== undefined + try { + await removeWorktreeForJob({ repoRoot, jobId, keepBranch }) + } catch (err) { + console.warn(`[update_job_status] Worktree cleanup failed for job ${jobId}:`, err) + } +} + const DB_STATUS_MAP = { running: 'RUNNING', done: 'DONE', @@ -108,6 +128,11 @@ export function registerUpdateJobStatusTool(server: McpServer) { // non-fatal — status is already persisted } + // Best-effort worktree cleanup on terminal transitions + if (status === 'done' || status === 'failed') { + await cleanupWorktreeForTerminalStatus(job.product_id, job_id, status, branch) + } + return toolJson({ job_id: updated.id, status,