diff --git a/CLAUDE.md b/CLAUDE.md index b64e8f6..af9950a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -38,6 +38,10 @@ Or add to `~/.scrum4me-agent-config.json`: If no repo root is found, `wait_for_job` rolls the claim back to QUEUED and returns an error. +## Manual worktree cleanup + +Run `cleanup_my_worktrees` (no arguments) to scan `~/.scrum4me-agent-worktrees/` and remove worktrees for jobs that are in a terminal state (DONE, FAILED, CANCELLED). Worktrees for active jobs (QUEUED, CLAIMED, RUNNING) are left untouched. Returns `{ removed, kept, skipped }`. + ## Key source files | File | Purpose | @@ -45,6 +49,7 @@ If no repo root is found, `wait_for_job` rolls the claim back to QUEUED and retu | `src/git/worktree.ts` | `createWorktreeForJob` + `removeWorktreeForJob` | | `src/tools/wait-for-job.ts` | `resolveRepoRoot`, `rollbackClaim`, `attachWorktreeToJob` | | `src/tools/update-job-status.ts` | `cleanupWorktreeForTerminalStatus` | +| `src/tools/cleanup-my-worktrees.ts` | `cleanup_my_worktrees` tool — scans + removes stale worktrees | ## Testing diff --git a/__tests__/cleanup-my-worktrees.test.ts b/__tests__/cleanup-my-worktrees.test.ts new file mode 100644 index 0000000..6460157 --- /dev/null +++ b/__tests__/cleanup-my-worktrees.test.ts @@ -0,0 +1,171 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import * as path from 'node:path' +import * as os from 'node:os' + +vi.mock('../src/prisma.js', () => ({ + prisma: { + claudeJob: { findMany: vi.fn() }, + }, +})) + +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() } +}) + +vi.mock('node:fs/promises', () => ({ + readdir: vi.fn(), +})) + +import { prisma } from '../src/prisma.js' +import { removeWorktreeForJob } from '../src/git/worktree.js' +import { resolveRepoRoot } from '../src/tools/wait-for-job.js' +import * as fsPromises from 'node:fs/promises' +import { cleanupWorktrees, listWorktreeJobIds, getWorktreeParent } from '../src/tools/cleanup-my-worktrees.js' + +const mockPrisma = prisma as unknown as { + claudeJob: { findMany: ReturnType } +} +const mockRemove = removeWorktreeForJob as ReturnType +const mockResolve = resolveRepoRoot as ReturnType +const mockReaddir = fsPromises.readdir as ReturnType + +const REPO_ROOT = '/repos/my-project' +const USER_ID = 'user-1' +const PRODUCT_ID = 'product-1' +const WORKTREE_PARENT = '/home/user/.scrum4me-agent-worktrees' + +function makeDirent(name: string, isDir = true) { + return { name, isDirectory: () => isDir } as unknown as import('node:fs').Dirent +} + +beforeEach(() => { + vi.clearAllMocks() + mockResolve.mockResolvedValue(REPO_ROOT) + mockRemove.mockResolvedValue({ removed: true }) +}) + +describe('getWorktreeParent', () => { + it('uses SCRUM4ME_AGENT_WORKTREE_DIR env var when set', async () => { + process.env.SCRUM4ME_AGENT_WORKTREE_DIR = '/custom/dir' + expect(await getWorktreeParent()).toBe('/custom/dir') + delete process.env.SCRUM4ME_AGENT_WORKTREE_DIR + }) + + it('defaults to ~/.scrum4me-agent-worktrees', async () => { + delete process.env.SCRUM4ME_AGENT_WORKTREE_DIR + expect(await getWorktreeParent()).toBe(path.join(os.homedir(), '.scrum4me-agent-worktrees')) + }) +}) + +describe('listWorktreeJobIds', () => { + it('returns directory names from the worktree parent', async () => { + mockReaddir.mockResolvedValue([makeDirent('job-aaa'), makeDirent('job-bbb'), makeDirent('file.txt', false)]) + const ids = await listWorktreeJobIds(WORKTREE_PARENT) + expect(ids).toEqual(['job-aaa', 'job-bbb']) + }) + + it('returns empty array when parent does not exist', async () => { + mockReaddir.mockRejectedValue(Object.assign(new Error('ENOENT'), { code: 'ENOENT' })) + expect(await listWorktreeJobIds(WORKTREE_PARENT)).toEqual([]) + }) +}) + +describe('cleanupWorktrees', () => { + it('removes worktrees for DONE/FAILED/CANCELLED jobs', async () => { + mockReaddir.mockResolvedValue([ + makeDirent('job-done'), + makeDirent('job-failed'), + makeDirent('job-cancelled'), + ]) + mockPrisma.claudeJob.findMany.mockResolvedValue([ + { id: 'job-done', status: 'DONE', product_id: PRODUCT_ID, branch: 'feat/job-done' }, + { id: 'job-failed', status: 'FAILED', product_id: PRODUCT_ID, branch: null }, + { id: 'job-cancelled', status: 'CANCELLED', product_id: PRODUCT_ID, branch: null }, + ]) + + const result = await cleanupWorktrees(WORKTREE_PARENT, USER_ID) + + expect(result.removed).toEqual(expect.arrayContaining(['job-done', 'job-failed', 'job-cancelled'])) + expect(result.kept).toEqual([]) + expect(result.skipped).toEqual([]) + expect(mockRemove).toHaveBeenCalledTimes(3) + }) + + it('keeps worktrees for QUEUED/CLAIMED/RUNNING jobs', async () => { + mockReaddir.mockResolvedValue([ + makeDirent('job-queued'), + makeDirent('job-claimed'), + makeDirent('job-running'), + ]) + mockPrisma.claudeJob.findMany.mockResolvedValue([ + { id: 'job-queued', status: 'QUEUED', product_id: PRODUCT_ID, branch: null }, + { id: 'job-claimed', status: 'CLAIMED', product_id: PRODUCT_ID, branch: null }, + { id: 'job-running', status: 'RUNNING', product_id: PRODUCT_ID, branch: null }, + ]) + + const result = await cleanupWorktrees(WORKTREE_PARENT, USER_ID) + + expect(result.kept).toEqual(expect.arrayContaining(['job-queued', 'job-claimed', 'job-running'])) + expect(result.removed).toEqual([]) + expect(mockRemove).not.toHaveBeenCalled() + }) + + it('calls removeWorktreeForJob with keepBranch=true for DONE with branch', async () => { + mockReaddir.mockResolvedValue([makeDirent('job-done')]) + mockPrisma.claudeJob.findMany.mockResolvedValue([ + { id: 'job-done', status: 'DONE', product_id: PRODUCT_ID, branch: 'feat/job-done' }, + ]) + + await cleanupWorktrees(WORKTREE_PARENT, USER_ID) + + expect(mockRemove).toHaveBeenCalledWith({ repoRoot: REPO_ROOT, jobId: 'job-done', keepBranch: true }) + }) + + it('calls removeWorktreeForJob with keepBranch=false for FAILED jobs', async () => { + mockReaddir.mockResolvedValue([makeDirent('job-failed')]) + mockPrisma.claudeJob.findMany.mockResolvedValue([ + { id: 'job-failed', status: 'FAILED', product_id: PRODUCT_ID, branch: null }, + ]) + + await cleanupWorktrees(WORKTREE_PARENT, USER_ID) + + expect(mockRemove).toHaveBeenCalledWith({ repoRoot: REPO_ROOT, jobId: 'job-failed', keepBranch: false }) + }) + + it('skips orphan worktrees (no DB record)', async () => { + mockReaddir.mockResolvedValue([makeDirent('job-orphan')]) + mockPrisma.claudeJob.findMany.mockResolvedValue([]) + + const result = await cleanupWorktrees(WORKTREE_PARENT, USER_ID) + + expect(result.skipped).toContain('job-orphan') + expect(mockRemove).not.toHaveBeenCalled() + }) + + it('skips worktrees when no repoRoot is configured', async () => { + mockReaddir.mockResolvedValue([makeDirent('job-norepo')]) + mockPrisma.claudeJob.findMany.mockResolvedValue([ + { id: 'job-norepo', status: 'FAILED', product_id: 'unknown-product', branch: null }, + ]) + mockResolve.mockResolvedValue(null) + + const result = await cleanupWorktrees(WORKTREE_PARENT, USER_ID) + + expect(result.skipped).toContain('job-norepo') + expect(mockRemove).not.toHaveBeenCalled() + }) + + it('returns empty result when worktree parent is empty', async () => { + mockReaddir.mockResolvedValue([]) + + const result = await cleanupWorktrees(WORKTREE_PARENT, USER_ID) + + expect(result).toEqual({ removed: [], kept: [], skipped: [] }) + expect(mockPrisma.claudeJob.findMany).not.toHaveBeenCalled() + }) +}) diff --git a/src/index.ts b/src/index.ts index d7c2371..15479e3 100644 --- a/src/index.ts +++ b/src/index.ts @@ -20,6 +20,7 @@ import { registerCancelQuestionTool } from './tools/cancel-question.js' import { registerWaitForJobTool } from './tools/wait-for-job.js' import { registerUpdateJobStatusTool } from './tools/update-job-status.js' import { registerVerifyTaskAgainstPlanTool } from './tools/verify-task-against-plan.js' +import { registerCleanupMyWorktreesTool } from './tools/cleanup-my-worktrees.js' import { registerImplementNextStoryPrompt } from './prompts/implement-next-story.js' const VERSION = '0.1.0' @@ -53,6 +54,7 @@ async function main() { registerWaitForJobTool(server) registerUpdateJobStatusTool(server) registerVerifyTaskAgainstPlanTool(server) + registerCleanupMyWorktreesTool(server) registerImplementNextStoryPrompt(server) const transport = new StdioServerTransport() diff --git a/src/tools/cleanup-my-worktrees.ts b/src/tools/cleanup-my-worktrees.ts new file mode 100644 index 0000000..bfcc444 --- /dev/null +++ b/src/tools/cleanup-my-worktrees.ts @@ -0,0 +1,107 @@ +import { z } from 'zod' +import * as fs from 'node:fs/promises' +import * as path from 'node:path' +import * as os from 'node:os' +import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import { prisma } from '../prisma.js' +import { requireWriteAccess } from '../auth.js' +import { toolJson, withToolErrors } from '../errors.js' +import { removeWorktreeForJob } from '../git/worktree.js' +import { resolveRepoRoot } from './wait-for-job.js' + +const TERMINAL_STATUSES = new Set(['DONE', 'FAILED', 'CANCELLED']) +const ACTIVE_STATUSES = new Set(['QUEUED', 'CLAIMED', 'RUNNING']) + +const inputSchema = z.object({}) + +export async function getWorktreeParent(): Promise { + return ( + process.env.SCRUM4ME_AGENT_WORKTREE_DIR ?? + path.join(os.homedir(), '.scrum4me-agent-worktrees') + ) +} + +export async function listWorktreeJobIds(worktreeParent: string): Promise { + try { + const entries = await fs.readdir(worktreeParent, { withFileTypes: true }) + return entries.filter((e) => e.isDirectory()).map((e) => e.name) + } catch { + return [] + } +} + +export async function cleanupWorktrees( + worktreeParent: string, + userId: string, +): Promise<{ removed: string[]; kept: string[]; skipped: string[] }> { + const jobIds = await listWorktreeJobIds(worktreeParent) + const removed: string[] = [] + const kept: string[] = [] + const skipped: string[] = [] + + if (jobIds.length === 0) return { removed, kept, skipped } + + const jobs = await prisma.claudeJob.findMany({ + where: { id: { in: jobIds }, user_id: userId }, + select: { id: true, status: true, product_id: true, branch: true }, + }) + const jobMap = new Map(jobs.map((j) => [j.id, j])) + + for (const jobId of jobIds) { + const job = jobMap.get(jobId) + + // No DB record for this jobId — orphan worktree, skip safely + if (!job) { + skipped.push(jobId) + continue + } + + if (ACTIVE_STATUSES.has(job.status)) { + kept.push(jobId) + continue + } + + if (TERMINAL_STATUSES.has(job.status)) { + const repoRoot = await resolveRepoRoot(job.product_id) + if (!repoRoot) { + skipped.push(jobId) + continue + } + + // Keep branch for DONE jobs that already pushed (job.branch is set) + const keepBranch = job.status === 'DONE' && job.branch !== null + try { + await removeWorktreeForJob({ repoRoot, jobId, keepBranch }) + removed.push(jobId) + } catch { + skipped.push(jobId) + } + } + } + + return { removed, kept, skipped } +} + +export function registerCleanupMyWorktreesTool(server: McpServer) { + server.registerTool( + 'cleanup_my_worktrees', + { + title: 'Cleanup my worktrees', + description: + 'Remove stale git worktrees left by crashed or cancelled agent runs. ' + + 'Scans ~/.scrum4me-agent-worktrees/ (or SCRUM4ME_AGENT_WORKTREE_DIR) for job directories, ' + + 'looks up each job\'s status, and removes worktrees whose jobs are in a terminal state ' + + '(DONE, FAILED, CANCELLED). Worktrees for active jobs (QUEUED, CLAIMED, RUNNING) are kept. ' + + 'Returns { removed, kept, skipped } for inspection.', + inputSchema, + annotations: { readOnlyHint: false }, + }, + async () => + withToolErrors(async () => { + const auth = await requireWriteAccess() + const worktreeParent = await getWorktreeParent() + const result = await cleanupWorktrees(worktreeParent, auth.userId) + return toolJson(result) + }), + ) +}