feat: cleanup worktree in update_job_status on terminal transitions

On DONE/FAILED, resolves repoRoot and calls removeWorktreeForJob (best-effort).
keepBranch=true when status=done and agent reported a branch (push assumed);
false otherwise. Cleanup failures are logged as warnings — DB status is preserved.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Janpeter Visser 2026-05-01 11:52:16 +02:00
parent 6ee55e79b6
commit ce4afa1928
2 changed files with 109 additions and 0 deletions

View file

@ -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<typeof import('../src/tools/wait-for-job.js')>()
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<typeof vi.fn>
const mockResolve = resolveRepoRoot as ReturnType<typeof vi.fn>
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()
})
})

View file

@ -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<void> {
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,