diff --git a/.kilo/skills/kilo-mcp/SKILL.md b/.kilo/skills/kilo-mcp/SKILL.md index 47b548380ed..36d797d5a6d 100644 --- a/.kilo/skills/kilo-mcp/SKILL.md +++ b/.kilo/skills/kilo-mcp/SKILL.md @@ -51,7 +51,7 @@ to send platform feedback and bug reports rather than a support channel. The first path segment names an area. Search to learn an area's vocabulary: each row's `summary` and `tags` are what the search matches. -**836 procedures** — **357 queries**, **479 mutations** — under **49 prefixes**. +**838 procedures** — **359 queries**, **479 mutations** — under **49 prefixes**. | Prefix | Procedures | Queries | Mutations | |---|---:|---:|---:| @@ -63,8 +63,8 @@ each row's `summary` and `tags` are what the search matches. | `securityAgent` | 26 | 15 | 11 | | `kiloPass` | 21 | 11 | 10 | | `githubPrReview` | 20 | 6 | 14 | +| `cliSessionsV2` | 19 | 12 | 7 | | `githubApps` | 19 | 9 | 10 | -| `cliSessionsV2` | 17 | 10 | 7 | | `appBuilder` | 16 | 6 | 10 | | `providerReview` | 16 | 8 | 8 | | `todos` | 16 | 6 | 10 | diff --git a/apps/web/src/components/cloud-agent-next/ChatSidebar.tsx b/apps/web/src/components/cloud-agent-next/ChatSidebar.tsx index edbf66b5ecf..09e202e27ae 100644 --- a/apps/web/src/components/cloud-agent-next/ChatSidebar.tsx +++ b/apps/web/src/components/cloud-agent-next/ChatSidebar.tsx @@ -9,6 +9,7 @@ import { useRef, useState, type DragEvent, + type ReactNode, } from 'react'; import { SquarePen, @@ -110,6 +111,7 @@ type ChatSidebarProps = { onProjectChange?: (gitUrls: string[]) => void; recentProjects?: Array<{ gitUrl: string; displayName: string }>; workspaceFolders?: WorkspaceFolderController; + workspaceCleanup?: ReactNode; }; type WorktreeFolderControls = { @@ -666,6 +668,7 @@ export function ChatSidebar({ onProjectChange, recentProjects = [], workspaceFolders, + workspaceCleanup, }: ChatSidebarProps) { const router = useRouter(); const pathname = usePathname(); @@ -1111,6 +1114,8 @@ export function ChatSidebar({ + {workspaceCleanup} + {/* Collapsible search */} {showSearch && (
diff --git a/apps/web/src/components/cloud-agent-next/CloudSidebarLayout.tsx b/apps/web/src/components/cloud-agent-next/CloudSidebarLayout.tsx index c81a69b8165..f3d95284ad4 100644 --- a/apps/web/src/components/cloud-agent-next/CloudSidebarLayout.tsx +++ b/apps/web/src/components/cloud-agent-next/CloudSidebarLayout.tsx @@ -85,6 +85,8 @@ import { } from './local-sessions/local-session'; import { sendLocalSessionRequest } from './local-sessions/local-session-request'; import { localSessionsStorageKey } from './local-sessions/local-session-storage'; +import { useWorktreeCleanup } from './hooks/useWorktreeCleanup'; +import { MergedWorktreeCleanup } from './MergedWorktreeCleanup'; // Context for children to toggle the mobile sidebar sheet type WorktreeChatTabs = { @@ -803,6 +805,50 @@ function CloudSidebarLayoutContent({ [organizationId, queryClient, renameWorktree, trpc] ); + const handleWorktreeDeleted = useCallback( + async (worktreeId: string, deletedSessionIds: string[]) => { + await Promise.all( + deletedSessionIds.map(sessionId => + removeDeletedSession({ + sessionId, + queryClient, + trpc, + setDbSessions, + deleteSessionFromStore, + }) + ) + ); + forgetWorktreeTabs(worktreeId, deletedSessionIds); + const redirect = getWorktreeDeletionRedirect( + new URL(window.location.href), + worktreeId, + deletedSessionIds, + organizationId + ); + if (redirect) router.push(redirect); + void invalidateSessionQueries({ queryClient, trpc }); + void queryClient.invalidateQueries(trpc.workspaceFolders.list.pathFilter()); + }, + [ + deleteSessionFromStore, + forgetWorktreeTabs, + organizationId, + queryClient, + router, + setDbSessions, + trpc, + ] + ); + const cleanup = useWorktreeCleanup({ organizationId, onDeleted: handleWorktreeDeleted }); + const workspaceCleanup = ( + + ); + const handleConfirmWorktreeDelete = useCallback(() => { const worktreeId = worktreePendingDeletion; if (!worktreeId || deletingWorktreeIds.has(worktreeId)) return; @@ -816,27 +862,7 @@ function CloudSidebarLayoutContent({ worktreeId: cloudAgentWorktreeIdSchema.parse(worktreeId), organizationId: organizationId ?? null, }); - await Promise.all( - deletedSessionIds.map(sessionId => - removeDeletedSession({ - sessionId, - queryClient, - trpc, - setDbSessions, - deleteSessionFromStore, - }) - ) - ); - forgetWorktreeTabs(worktreeId, deletedSessionIds); - const redirect = getWorktreeDeletionRedirect( - new URL(window.location.href), - worktreeId, - deletedSessionIds, - organizationId - ); - if (redirect) router.push(redirect); - void invalidateSessionQueries({ queryClient, trpc }); - void queryClient.invalidateQueries(trpc.workspaceFolders.list.pathFilter()); + await handleWorktreeDeleted(worktreeId, deletedSessionIds); toast.success('Worktree deleted'); } catch (error) { toast.error('Failed to delete worktree', { description: formatSessionError(error) }); @@ -851,19 +877,24 @@ function CloudSidebarLayoutContent({ } })(); }, [ - deleteSessionFromStore, deleteWorktree, deletingWorktreeIds, - forgetWorktreeTabs, + handleWorktreeDeleted, organizationId, queryClient, - router, - setDbSessions, trpc, worktreePendingDeletion, ]); - const deletingWorktreeIdList = useMemo(() => [...deletingWorktreeIds], [deletingWorktreeIds]); + const deletingWorktreeIdList = useMemo( + () => [ + ...deletingWorktreeIds, + ...(cleanup.data?.items + .filter(item => item.status === 'pending' || item.status === 'deleting') + .map(item => item.worktreeId) ?? []), + ], + [cleanup.data, deletingWorktreeIds] + ); const toggleMobileSidebar = useCallback(() => setMobileSheetOpen(prev => !prev), []); const sidebarContext = useMemo( () => ({ @@ -922,6 +953,7 @@ function CloudSidebarLayoutContent({ [0]; +const mockSessions: StoredSession[] = []; +const mockFolders = { + folders: [], + isLoading: false, + refresh: jest.fn().mockResolvedValue(undefined), +}; +const mockTabs = { + closedSessionIds: [], + sessionOrderByWorktree: {}, + openChatTab: jest.fn(), + closeChatTab: jest.fn(), + replaceChatTab: jest.fn(), + forgetWorktreeTabs: mockForgetTabs, +}; +const mockLocalSessions = { + localSessions: [], + createLocalSession: jest.fn(), + retryLocalSession: jest.fn(), + discardLocalSession: jest.fn(), + settleLocalSessions: jest.fn(), +}; +const mockTrpc = { + cliSessionsV2: Object.fromEntries( + [ + 'list', + 'search', + 'recentRepositories', + 'worktreeDetails', + 'delete', + 'rename', + 'renameWorktree', + 'deleteWorktree', + ].map(name => [ + name, + { + queryOptions: () => ({}), + mutationOptions: () => ({}), + mutationKey: () => [name], + pathFilter: () => ({ queryKey: [name] }), + }, + ]) + ), + workspaceFolders: { list: { pathFilter: () => ({ queryKey: ['folders'] }) } }, +}; + +jest.mock('next/navigation', () => ({ + useRouter: () => mockRouter, + usePathname: () => '/cloud', + useSearchParams: () => mockSearchParams, +})); +jest.mock('@tanstack/react-query', () => ({ + skipToken: Symbol('skip'), + useQuery: () => ({ data: undefined }), + useMutation: () => ({ mutate: jest.fn(), mutateAsync: jest.fn() }), + useMutationState: () => [], + useQueryClient: () => mockQueryClient, +})); +jest.mock('@/lib/trpc/utils', () => ({ + useTRPC: () => mockTrpc, + useRawTRPCClient: () => ({}), +})); +jest.mock('jotai', () => ({ useAtomValue: () => null, useSetAtom: () => jest.fn() })); +jest.mock('./store/db-session-atoms', () => ({ + dbSessionsAtom: {}, + deleteSessionFromStoreAtom: {}, +})); +jest.mock('./CloudAgentProvider', () => ({ useManager: () => ({ atoms: {} }) })); +jest.mock('./hooks/useSidebarSessions', () => ({ + useSidebarSessions: () => ({ + sessions: mockSessions, + cachedSessions: mockSessions, + worktreeDetails: {}, + loadingWorktreeIds: [], + isLoading: false, + refetchSessions: jest.fn(), + }), + deriveForegroundSessionStatus: () => null, + mergeWorktreeChatSessions: () => [], +})); +jest.mock('./hooks/useActiveSessions', () => ({ + useActiveSessions: () => ({ activeSessions: [] }), +})); +jest.mock('./hooks/useSidebarPreferences', () => ({ + useSidebarPreferences: () => ({ + platformFilter: [], + projectFilter: ['https://github.com/kilo/filtered'], + }), +})); +jest.mock('./hooks/useWorkspaceFolders', () => ({ useWorkspaceFolders: () => mockFolders })); +jest.mock('./hooks/useClosedWorktreeChatTabs', () => ({ + useClosedWorktreeChatTabs: () => mockTabs, +})); +jest.mock('./local-sessions/LocalSessionsProvider', () => ({ + LocalSessionsProvider: ({ children }: { children: ReactNode }) => children, + useLocalSessions: () => mockLocalSessions, +})); +jest.mock('./local-sessions/local-session-request', () => ({ sendLocalSessionRequest: jest.fn() })); +jest.mock('@/components/SetPageActions', () => ({ SetPageActions: () => null })); +jest.mock('./settings/CloudAgentSettingsButton', () => ({ CloudAgentSettingsButton: () => null })); +jest.mock('@/components/ui/sheet', () => ({ + Sheet: () => null, + SheetContent: () => null, + SheetHeader: () => null, + SheetTitle: () => null, +})); +jest.mock('@/components/ui/alert-dialog', () => ({ + AlertDialog: () => null, + AlertDialogAction: () => null, + AlertDialogCancel: () => null, + AlertDialogContent: () => null, + AlertDialogDescription: () => null, + AlertDialogFooter: () => null, + AlertDialogHeader: () => null, + AlertDialogTitle: () => null, +})); +jest.mock('./hooks/useWorktreeCleanup', () => ({ + useWorktreeCleanup: (options: Parameters[0]) => { + mockCleanupOptions = options; + return mockCleanup; + }, +})); +jest.mock('./MergedWorktreeCleanup', () => ({ + MergedWorktreeCleanup: ({ + organizationId, + cleanup, + statusError, + statusLoading, + }: { + organizationId?: string; + cleanup: WorktreeCleanup | null; + statusError: { message: string } | null; + statusLoading: boolean; + }) => + jest.requireActual('react').createElement( + 'output', + { + 'data-scope': organizationId ?? 'personal', + 'data-pending': String(statusLoading), + }, + statusError?.message ?? cleanup?.id ?? 'No cleanup' + ), +})); +jest.mock('./ChatSidebar', () => ({ + ChatSidebar: ({ + workspaceCleanup, + deletingWorktreeIds, + }: { + workspaceCleanup: ReactNode; + deletingWorktreeIds: string[]; + }) => { + const react = jest.requireActual('react'); + return react.createElement( + 'aside', + null, + workspaceCleanup, + deletingWorktreeIds.map(id => + react.createElement('button', { key: id, disabled: true, 'data-deleting': id }, id) + ) + ); + }, +})); +jest.mock('./session-deletion', () => ({ + ...jest.requireActual('./session-deletion'), + removeDeletedSession: (input: object) => mockRemoveDeletedSession(input), + invalidateSessionQueries: (input: object) => mockInvalidateSessions(input), +})); + +function installDom() { + const requireFromNext = createRequire(createRequire(__filename).resolve('next/package.json')); + const { window, document } = ( + requireFromNext('linkedom') as { + parseHTML: (html: string) => { window: Window & typeof globalThis; document: Document }; + } + ).parseHTML('
'); + const values = { + React, + window, + document, + HTMLElement: window.HTMLElement, + Element: window.Element, + Node: window.Node, + IS_REACT_ACT_ENVIRONMENT: true, + }; + const previous = new Map( + Object.keys(values).map(name => [name, Object.getOwnPropertyDescriptor(globalThis, name)]) + ); + for (const [name, value] of Object.entries(values)) { + Object.defineProperty(globalThis, name, { configurable: true, writable: true, value }); + } + const container = document.getElementById('root'); + if (!container) throw new Error('Sidebar cleanup test root missing'); + return { + container, + cleanup: () => { + for (const [name, descriptor] of previous) { + if (descriptor) Object.defineProperty(globalThis, name, descriptor); + else Reflect.deleteProperty(globalThis, name); + } + }, + }; +} + +describe('CloudSidebarLayout cleanup integration', () => { + let dom: ReturnType; + let root: Root; + let CloudSidebarLayout: typeof SidebarLayout; + + beforeEach(async () => { + jest.clearAllMocks(); + mockSearchParams = new URLSearchParams(); + mockCleanup = { data: null, error: null, isPending: false }; + dom = installDom(); + ({ CloudSidebarLayout } = await import('./CloudSidebarLayout')); + const { createRoot } = await import('react-dom/client'); + root = createRoot(dom.container); + }); + + afterEach(() => { + act(() => root.unmount()); + dom.cleanup(); + }); + + async function render(organizationId?: string) { + await act(async () => + root.render( + React.createElement(CloudSidebarLayout, { + currentUserId: 'user-123', + organizationId, + children: React.createElement('main', null, 'Chat content'), + }) + ) + ); + } + + it('passes scope, status loading and structural errors through independently of sidebar filters', async () => { + mockCleanup.isPending = true; + await render('org-123'); + expect(mockCleanupOptions.organizationId).toBe('org-123'); + expect(dom.container.querySelector('output')?.getAttribute('data-scope')).toBe('org-123'); + expect(dom.container.querySelector('output')?.getAttribute('data-pending')).toBe('true'); + mockCleanup = { data: null, error: { message: 'Progress unavailable' }, isPending: false }; + await render('org-123'); + expect(dom.container.querySelector('output')?.textContent).toBe('Progress unavailable'); + expect(dom.container.textContent).toContain('Chat content'); + }); + + it('marks only pending and deleting workspace rows and releases terminal rows', async () => { + const batch: WorktreeCleanup = { + id: 'batch-1', + createdAt: '2026-10-20T12:00:00.000Z', + items: (['pending', 'deleting', 'deleted', 'failed'] as const).map((status, index) => ({ + worktreeId: `worktree_${index}`, + status, + deletedSessionIds: [], + error: null, + })), + }; + mockCleanup.data = batch; + await render(); + expect( + Array.from(dom.container.querySelectorAll('[data-deleting]')).map(node => + node.getAttribute('data-deleting') + ) + ).toEqual(['worktree_0', 'worktree_1']); + expect(dom.container.querySelector('output')?.textContent).toBe('batch-1'); + mockCleanup.data = { + ...batch, + items: batch.items.map(item => ({ ...item, status: 'deleted' })), + }; + await render(); + expect(dom.container.querySelector('[data-deleting]')).toBeNull(); + }); + + it.each([undefined, 'org-123'])( + 'reconciles all deleted chats, tabs and queries and redirects the selected scope %s', + async organizationId => { + const base = organizationId ? `/organizations/${organizationId}/cloud` : '/cloud'; + Object.defineProperty(window, 'location', { + configurable: true, + value: { href: `https://example.com${base}/chat?sessionId=deleted-chat` }, + }); + await render(organizationId); + await act(async () => + mockCleanupOptions.onDeleted('deleted-worktree', ['deleted-chat', 'deleted-child']) + ); + expect(mockRemoveDeletedSession).toHaveBeenCalledTimes(2); + expect(mockRemoveDeletedSession).toHaveBeenCalledWith( + expect.objectContaining({ + sessionId: 'deleted-chat', + queryClient: mockQueryClient, + trpc: mockTrpc, + }) + ); + expect(mockRemoveDeletedSession).toHaveBeenCalledWith( + expect.objectContaining({ sessionId: 'deleted-child' }) + ); + expect(mockForgetTabs).toHaveBeenCalledWith('deleted-worktree', [ + 'deleted-chat', + 'deleted-child', + ]); + expect(mockPush).toHaveBeenCalledWith(base); + expect(mockInvalidateSessions).toHaveBeenCalledWith({ + queryClient: mockQueryClient, + trpc: mockTrpc, + }); + expect(mockInvalidateQueries).toHaveBeenCalledWith({ queryKey: ['folders'] }); + } + ); + + it('does not navigate away from an unrelated selection when background cleanup completes', async () => { + Object.defineProperty(window, 'location', { + configurable: true, + value: { href: 'https://example.com/cloud/chat?sessionId=surviving-chat' }, + }); + await render(); + await act(async () => mockCleanupOptions.onDeleted('deleted-worktree', ['deleted-chat'])); + expect(mockPush).not.toHaveBeenCalled(); + expect(mockForgetTabs).toHaveBeenCalledWith('deleted-worktree', ['deleted-chat']); + expect(mockInvalidateSessions).toHaveBeenCalledTimes(1); + }); +}); diff --git a/apps/web/src/components/cloud-agent-next/MergedWorktreeCleanup.test.ts b/apps/web/src/components/cloud-agent-next/MergedWorktreeCleanup.test.ts new file mode 100644 index 00000000000..346ce3c0388 --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/MergedWorktreeCleanup.test.ts @@ -0,0 +1,560 @@ +import React, { act, type ReactNode } from 'react'; +import { createRequire } from 'node:module'; +import type { Root } from 'react-dom/client'; +import type { QueryClient, QueryClientProvider as QueryProvider } from '@tanstack/react-query'; +import type { MergedWorktreeCleanup as CleanupComponent } from './MergedWorktreeCleanup'; +import type { + WorktreeCleanup, + useWorktreeCleanup as CleanupHook, +} from './hooks/useWorktreeCleanup'; + +type Scope = { organizationId: string | null }; +type StartInput = Scope & { worktreeIds: string[] }; +const mockPreview = jest.fn, [Scope]>(); +const mockStart = jest.fn, [StartInput]>(); +const mockRetry = jest.fn, [Scope]>(); +const mockGet = jest.fn, [Scope]>(); +const cleanupKey = (scope: Scope) => ['cleanup', scope]; + +jest.mock('@/lib/trpc/utils', () => ({ + useTRPC: () => ({ + cliSessionsV2: { + mergedWorktreeCleanupPreview: { + queryOptions: (scope: Scope) => ({ + queryKey: ['preview', scope], + queryFn: () => mockPreview(scope), + }), + }, + startMergedWorktreeCleanup: { + mutationOptions: (options: object) => ({ + mutationFn: (input: StartInput) => mockStart(input), + ...options, + }), + }, + retryWorktreeCleanup: { + mutationOptions: (options: object) => ({ + mutationFn: (input: Scope) => mockRetry(input), + ...options, + }), + }, + getWorktreeCleanup: { + queryKey: (scope: Scope) => cleanupKey(scope), + queryOptions: (scope: Scope) => ({ + queryKey: cleanupKey(scope), + queryFn: () => mockGet(scope), + }), + }, + }, + }), +})); + +jest.mock('@/components/ui/alert-dialog', () => { + const react = jest.requireActual('react'); + const context = react.createContext<(open: boolean) => void>(() => {}); + const passthrough = ({ children }: { children?: ReactNode }) => + react.createElement(react.Fragment, null, children); + return { + AlertDialog: ({ + open, + onOpenChange, + children, + }: { + open: boolean; + onOpenChange: (open: boolean) => void; + children: ReactNode; + }) => (open ? react.createElement(context.Provider, { value: onOpenChange }, children) : null), + AlertDialogContent: ({ children }: { children: ReactNode }) => + react.createElement('div', { role: 'dialog' }, children), + AlertDialogTitle: passthrough, + AlertDialogDescription: passthrough, + AlertDialogHeader: passthrough, + AlertDialogFooter: passthrough, + AlertDialogAction: ({ children, ...props }: React.ComponentProps<'button'>) => + react.createElement('button', props, children), + AlertDialogCancel: ({ children, ...props }: React.ComponentProps<'button'>) => { + const onOpenChange = react.useContext(context); + return react.createElement( + 'button', + { ...props, onClick: () => onOpenChange(false) }, + children + ); + }, + }; +}); + +function installDom() { + const requireFromNext = createRequire(createRequire(__filename).resolve('next/package.json')); + const { window, document } = ( + requireFromNext('linkedom') as { + parseHTML: (html: string) => { window: Window & typeof globalThis; document: Document }; + } + ).parseHTML('
'); + const values = { + React, + window, + document, + HTMLElement: window.HTMLElement, + Element: window.Element, + Node: window.Node, + IS_REACT_ACT_ENVIRONMENT: true, + }; + const previous = new Map( + Object.keys(values).map(name => [name, Object.getOwnPropertyDescriptor(globalThis, name)]) + ); + for (const [name, value] of Object.entries(values)) { + Object.defineProperty(globalThis, name, { configurable: true, writable: true, value }); + } + const container = document.getElementById('root'); + if (!container) throw new Error('Cleanup test root missing'); + return { + container, + cleanup: () => { + for (const [name, descriptor] of previous) { + if (descriptor) Object.defineProperty(globalThis, name, descriptor); + else Reflect.deleteProperty(globalThis, name); + } + }, + }; +} + +function batch(statuses: WorktreeCleanup['items'][number]['status'][]): WorktreeCleanup { + return { + id: 'batch-1', + createdAt: '2026-10-20T12:00:00.000Z', + items: statuses.map((status, index) => ({ + worktreeId: `worktree_${index}`, + status, + deletedSessionIds: status === 'deleted' ? [`session-${index}`, `child-${index}`] : [], + error: status === 'failed' ? 'Sandbox cleanup failed' : null, + })), + }; +} + +describe('merged workspace cleanup', () => { + let dom: ReturnType; + let root: Root; + let client: QueryClient; + let Provider: typeof QueryProvider; + let MergedWorktreeCleanup: typeof CleanupComponent; + let useWorktreeCleanup: typeof CleanupHook; + + async function advance(milliseconds = 1) { + await act(async () => { + await jest.advanceTimersByTimeAsync(milliseconds); + }); + } + + async function render(element: React.ReactElement) { + await act(async () => { + root.render(React.createElement(Provider, { client }, element)); + }); + await advance(); + await advance(); + } + + function button(label: string) { + const result = Array.from(dom.container.querySelectorAll('button')).find( + candidate => candidate.textContent?.trim() === label + ); + if (!result) throw new Error(`Missing button: ${label}`); + return result; + } + + async function click(label: string) { + await act(async () => { + button(label).dispatchEvent(new window.Event('click', { bubbles: true })); + }); + await advance(); + await advance(); + } + + async function renderCleanup( + overrides: Partial> = {} + ) { + await render( + React.createElement(MergedWorktreeCleanup, { + cleanup: null, + statusError: null, + statusLoading: false, + ...overrides, + }) + ); + } + + beforeEach(async () => { + jest.useFakeTimers(); + dom = installDom(); + const query = await import('@tanstack/react-query'); + expect(query.isServer).toBe(false); + Provider = query.QueryClientProvider; + client = new query.QueryClient({ + defaultOptions: { + queries: { retry: false, gcTime: Infinity }, + mutations: { retry: false }, + }, + }); + const { createRoot } = await import('react-dom/client'); + root = createRoot(dom.container); + ({ MergedWorktreeCleanup } = await import('./MergedWorktreeCleanup')); + ({ useWorktreeCleanup } = await import('./hooks/useWorktreeCleanup')); + mockPreview.mockReset().mockResolvedValue({ worktreeIds: ['eligible-a', 'eligible-b'] }); + mockStart.mockReset().mockResolvedValue(batch(['pending', 'pending'])); + mockRetry.mockReset().mockResolvedValue(batch(['pending'])); + mockGet.mockReset().mockResolvedValue(null); + }); + + afterEach(() => { + act(() => root.unmount()); + client.clear(); + dom.cleanup(); + jest.clearAllTimers(); + jest.useRealTimers(); + }); + + it.each([undefined, 'org-123'])( + 'confirms count, destructive scope and frozen IDs for %s', + async organizationId => { + await renderCleanup({ organizationId }); + expect(mockPreview).not.toHaveBeenCalled(); + await click('Clean up merged workspaces'); + expect(mockPreview).toHaveBeenCalledWith({ organizationId: organizationId ?? null }); + expect(dom.container.textContent).toContain('Permanently delete 2 workspaces'); + expect(dom.container.textContent).toContain('including all chats and files'); + expect(dom.container.textContent).toContain('regardless of sidebar filters'); + expect(dom.container.textContent).toContain('unmerged or unverified pull requests are kept'); + expect(dom.container.textContent).toContain('even if you close this tab'); + expect(mockStart).not.toHaveBeenCalled(); + mockPreview.mockResolvedValue({ worktreeIds: ['newly-eligible'] }); + await click('Delete 2 workspaces'); + expect(mockStart).toHaveBeenCalledTimes(1); + expect(mockStart).toHaveBeenCalledWith({ + organizationId: organizationId ?? null, + worktreeIds: ['eligible-a', 'eligible-b'], + }); + expect(client.getQueryData(cleanupKey({ organizationId: organizationId ?? null }))).toEqual( + batch(['pending', 'pending']) + ); + expect(dom.container.querySelector('[role="dialog"]')).toBeNull(); + } + ); + + it('keeps delete disabled while preview loads and when no workspaces are eligible', async () => { + const preview = Promise.withResolvers<{ worktreeIds: string[] }>(); + mockPreview.mockReturnValue(preview.promise); + await renderCleanup(); + await click('Clean up merged workspaces'); + expect(dom.container.textContent).toContain('Checking your workspaces'); + expect(button('Delete 0 workspaces').disabled).toBe(true); + await act(async () => preview.resolve({ worktreeIds: [] })); + await advance(); + expect(dom.container.textContent).toContain('No workspaces with merged pull requests'); + expect(button('Delete 0 workspaces').disabled).toBe(true); + await click('Delete 0 workspaces'); + expect(mockStart).not.toHaveBeenCalled(); + await click('Cancel'); + expect(dom.container.querySelector('[role="dialog"]')).toBeNull(); + }); + + it('blocks duplicate starts and cancellation while the start mutation is pending', async () => { + const start = Promise.withResolvers(); + mockStart.mockReturnValue(start.promise); + await renderCleanup(); + await click('Clean up merged workspaces'); + await click('Delete 2 workspaces'); + expect(button('Starting cleanup...').disabled).toBe(true); + expect(button('Cancel').disabled).toBe(true); + expect(button('Clean up merged workspaces').disabled).toBe(true); + await click('Starting cleanup...'); + await click('Cancel'); + expect(mockStart).toHaveBeenCalledTimes(1); + expect(dom.container.querySelector('[role="dialog"]')).not.toBeNull(); + await act(async () => start.resolve(batch(['pending']))); + await advance(); + expect(dom.container.querySelector('[role="dialog"]')).toBeNull(); + }); + + it('keeps preview failures cancellable and retries the query on reopening', async () => { + mockPreview.mockRejectedValueOnce(new Error('Preview unavailable')); + await renderCleanup(); + await click('Clean up merged workspaces'); + expect(dom.container.querySelector('[role="alert"]')?.textContent).toBe( + 'Connection failed. Please retry in a moment.' + ); + expect(button('Delete 0 workspaces').disabled).toBe(true); + expect(button('Cancel').disabled).toBe(false); + await click('Cancel'); + await click('Clean up merged workspaces'); + expect(mockPreview).toHaveBeenCalledTimes(2); + expect(button('Delete 2 workspaces').disabled).toBe(false); + expect(mockStart).not.toHaveBeenCalled(); + }); + + it('keeps a failed start open and actionable for another attempt', async () => { + mockStart.mockRejectedValueOnce(new Error('Start unavailable')); + await renderCleanup(); + await click('Clean up merged workspaces'); + await click('Delete 2 workspaces'); + expect(dom.container.querySelector('[role="alert"]')?.textContent).toBe( + 'Connection failed. Please retry in a moment.' + ); + expect(button('Delete 2 workspaces').disabled).toBe(false); + expect(button('Cancel').disabled).toBe(false); + await click('Delete 2 workspaces'); + expect(mockStart).toHaveBeenCalledTimes(2); + expect(dom.container.querySelector('[role="dialog"]')).toBeNull(); + }); + + it('shows structural progress-query errors without stranding the cleanup action', async () => { + await renderCleanup({ statusError: { message: 'Progress unavailable' } }); + expect(dom.container.querySelector('[role="alert"]')?.textContent).toContain( + 'It will refresh automatically' + ); + expect(button('Clean up merged workspaces').disabled).toBe(false); + await click('Clean up merged workspaces'); + expect(button('Delete 2 workspaces').disabled).toBe(false); + }); + + it('disables cleanup during initial status loading and active batches', async () => { + await renderCleanup({ statusLoading: true }); + expect(button('Clean up merged workspaces').disabled).toBe(true); + await renderCleanup({ cleanup: batch(['deleted', 'pending', 'deleting', 'failed']) }); + expect(button('Cleaning up merged workspaces...').disabled).toBe(true); + expect(dom.container.querySelector('[role="status"]')?.textContent).toContain( + '1 of 4 workspaces deleted, 1 failed' + ); + expect(dom.container.textContent).toContain('You can close this tab.'); + expect(button('Retry failed cleanup').disabled).toBe(true); + await renderCleanup({ cleanup: batch(['deleted', 'deleted']) }); + expect(dom.container.querySelector('[role="status"]')?.textContent).toContain( + '2 of 2 workspaces deleted' + ); + expect(dom.container.textContent).not.toContain('You can close this tab.'); + expect(button('Clean up merged workspaces').disabled).toBe(false); + }); + + it('retries failures through the scoped retry endpoint and leaves retry errors actionable', async () => { + const retry = Promise.withResolvers(); + mockRetry + .mockRejectedValueOnce(new Error('Retry unavailable')) + .mockReturnValueOnce(retry.promise); + await renderCleanup({ organizationId: 'org-123', cleanup: batch(['deleted', 'failed']) }); + expect(dom.container.textContent).toContain('1 of 2 workspaces deleted, 1 failed'); + expect(dom.container.querySelector('details')?.textContent).toContain( + 'worktree_1: Sandbox cleanup failed' + ); + await click('Retry failed cleanup'); + expect(dom.container.querySelector('[role="alert"]')?.textContent).toBe( + 'Connection failed. Please retry in a moment.' + ); + expect(button('Retry failed cleanup').disabled).toBe(false); + await click('Retry failed cleanup'); + expect(button('Retrying...').disabled).toBe(true); + await click('Retrying...'); + expect(mockRetry).toHaveBeenCalledTimes(2); + expect(mockRetry).toHaveBeenLastCalledWith({ organizationId: 'org-123' }); + expect(mockStart).not.toHaveBeenCalled(); + expect(mockPreview).not.toHaveBeenCalled(); + await act(async () => retry.resolve(batch(['pending']))); + await advance(); + expect(client.getQueryData(cleanupKey({ organizationId: 'org-123' }))).toEqual( + batch(['pending']) + ); + }); + + it('clears a previous retry error after starting a new batch', async () => { + mockRetry.mockRejectedValueOnce(new Error('Retry unavailable')); + await renderCleanup({ cleanup: batch(['failed']) }); + await click('Retry failed cleanup'); + expect(dom.container.querySelector('[role="alert"]')).not.toBeNull(); + + await click('Clean up merged workspaces'); + await click('Delete 2 workspaces'); + await renderCleanup({ cleanup: { ...batch(['failed']), id: 'batch-2' } }); + expect(dom.container.textContent).toContain('0 of 1 workspaces deleted, 1 failed'); + expect(button('Retry failed cleanup').disabled).toBe(false); + expect(dom.container.querySelector('[role="alert"]')).toBeNull(); + await click('Retry failed cleanup'); + expect(mockRetry).toHaveBeenCalledTimes(2); + }); + + function HookHarness(props: Parameters[0]) { + const result = useWorktreeCleanup(props); + return React.createElement( + 'output', + null, + result.error?.message ?? result.data?.items.map(item => item.status).join(',') + ); + } + + it('polls active work every 2 seconds and terminal batches every 30 seconds', async () => { + mockGet.mockResolvedValue(batch(['pending'])); + const onDeleted = jest.fn().mockResolvedValue(undefined); + await render(React.createElement(HookHarness, { organizationId: 'org-123', onDeleted })); + expect(mockGet).toHaveBeenCalledWith({ organizationId: 'org-123' }); + expect(mockGet).toHaveBeenCalledTimes(1); + mockGet.mockResolvedValue(batch(['deleting'])); + await advance(2000); + await advance(); + expect(mockGet).toHaveBeenCalledTimes(2); + expect(dom.container.textContent).toBe('deleting'); + mockGet.mockResolvedValue(batch(['deleted'])); + await advance(2000); + await advance(); + expect(mockGet).toHaveBeenCalledTimes(3); + expect(dom.container.textContent).toBe('deleted'); + expect(onDeleted).toHaveBeenCalledWith('worktree_0', ['session-0', 'child-0']); + await advance(29_000); + expect(mockGet).toHaveBeenCalledTimes(3); + await advance(1000); + await advance(); + expect(mockGet).toHaveBeenCalledTimes(4); + expect(onDeleted).toHaveBeenCalledTimes(1); + }); + + it('reconciles deleted items once per batch across layout remounts', async () => { + const data = batch(['deleted', 'failed', 'pending']); + mockGet.mockResolvedValue(data); + const onDeleted = jest.fn().mockResolvedValue(undefined); + const element = React.createElement(HookHarness, { onDeleted }); + await render(element); + expect(mockGet).toHaveBeenCalledWith({ organizationId: null }); + expect(onDeleted).toHaveBeenCalledTimes(1); + await act(async () => { + client.setQueryData(cleanupKey({ organizationId: null }), { ...data, createdAt: 'later' }); + }); + await advance(); + expect(onDeleted).toHaveBeenCalledTimes(1); + await render(React.createElement(React.Fragment)); + await render(element); + expect(onDeleted).toHaveBeenCalledTimes(1); + expect(onDeleted).toHaveBeenLastCalledWith('worktree_0', ['session-0', 'child-0']); + await act(async () => { + client.setQueryData(cleanupKey({ organizationId: null }), { ...data, id: 'batch-2' }); + }); + await advance(); + expect(onDeleted).toHaveBeenCalledTimes(2); + }); + + it('reconciles deletions completed while the layout was unmounted', async () => { + mockGet.mockResolvedValue(batch(['deleted', 'pending'])); + const onDeleted = jest.fn().mockResolvedValue(undefined); + const element = React.createElement(HookHarness, { onDeleted }); + await render(element); + expect(onDeleted).toHaveBeenCalledTimes(1); + await render(React.createElement(React.Fragment)); + mockGet.mockResolvedValue(batch(['deleted', 'deleted'])); + await render(element); + expect(onDeleted).toHaveBeenCalledTimes(2); + expect(onDeleted).toHaveBeenLastCalledWith('worktree_1', ['session-1', 'child-1']); + }); + + it('reconciles persisted results again with a fresh query client', async () => { + mockGet.mockResolvedValue(batch(['deleted'])); + const onDeleted = jest.fn().mockResolvedValue(undefined); + const element = React.createElement(HookHarness, { onDeleted }); + await render(element); + expect(onDeleted).toHaveBeenCalledTimes(1); + await render(React.createElement(React.Fragment)); + client.clear(); + const { QueryClient } = await import('@tanstack/react-query'); + client = new QueryClient({ + defaultOptions: { queries: { retry: false, gcTime: Infinity } }, + }); + await render(element); + expect(onDeleted).toHaveBeenCalledTimes(2); + expect(onDeleted).toHaveBeenLastCalledWith('worktree_0', ['session-0', 'child-0']); + }); + + it('keeps reconciliation isolated between personal and organization scopes', async () => { + mockGet.mockResolvedValue(batch(['deleted'])); + const onDeleted = jest.fn().mockResolvedValue(undefined); + await render(React.createElement(HookHarness, { onDeleted })); + await render(React.createElement(HookHarness, { organizationId: 'org-123', onDeleted })); + expect(onDeleted).toHaveBeenCalledTimes(2); + await render(React.createElement(HookHarness, { onDeleted })); + expect(onDeleted).toHaveBeenCalledTimes(2); + }); + + it('continues polling after status-query errors and discovers background cleanup from an empty scope', async () => { + mockGet.mockRejectedValueOnce(new Error('Status temporarily unavailable')); + const onDeleted = jest.fn().mockResolvedValue(undefined); + await render(React.createElement(HookHarness, { onDeleted })); + expect(dom.container.textContent).toBe('Status temporarily unavailable'); + await advance(30_000); + await advance(); + expect(mockGet).toHaveBeenCalledTimes(2); + expect(dom.container.textContent).toBe(''); + mockGet.mockResolvedValue(batch(['pending'])); + await advance(30_000); + await advance(); + expect(mockGet).toHaveBeenCalledTimes(3); + expect(dom.container.textContent).toBe('pending'); + mockGet.mockResolvedValue(batch(['deleted'])); + await advance(2000); + await advance(); + expect(dom.container.textContent).toBe('deleted'); + expect(onDeleted).toHaveBeenCalledTimes(1); + }); + + it('does not duplicate reconciliation while local deletion is still in flight', async () => { + const reconciliation = Promise.withResolvers(); + const onDeleted = jest.fn().mockReturnValue(reconciliation.promise); + const data = batch(['deleted']); + mockGet.mockResolvedValue(data); + await render(React.createElement(HookHarness, { onDeleted })); + expect(onDeleted).toHaveBeenCalledTimes(1); + await act(async () => { + client.setQueryData(cleanupKey({ organizationId: null }), { + ...data, + createdAt: 'new-snapshot', + }); + }); + await advance(); + await advance(30_000); + await advance(); + expect(onDeleted).toHaveBeenCalledTimes(1); + await render(React.createElement(React.Fragment)); + await render(React.createElement(HookHarness, { onDeleted })); + expect(onDeleted).toHaveBeenCalledTimes(1); + await act(async () => reconciliation.resolve()); + await advance(30_000); + expect(onDeleted).toHaveBeenCalledTimes(1); + }); + + it('uses the latest reconciliation callback and allows a failed reconciliation to retry', async () => { + const data = batch(['pending']); + mockGet.mockResolvedValue(data); + const first = jest.fn().mockResolvedValue(undefined); + const latest = jest + .fn() + .mockRejectedValueOnce(new Error('Local store unavailable')) + .mockResolvedValue(undefined); + const consoleError = jest.spyOn(console, 'error').mockImplementation(() => {}); + try { + await render(React.createElement(HookHarness, { onDeleted: first })); + await render(React.createElement(HookHarness, { onDeleted: latest })); + mockGet.mockResolvedValue(batch(['deleted'])); + await advance(2000); + await advance(); + expect(first).not.toHaveBeenCalled(); + expect(latest).toHaveBeenCalledTimes(1); + expect(consoleError).toHaveBeenCalledWith( + 'Failed to reconcile deleted workspace:', + expect.any(Error) + ); + await render(React.createElement(React.Fragment)); + await render(React.createElement(HookHarness, { onDeleted: latest })); + expect(latest).toHaveBeenCalledTimes(2); + await act(async () => { + client.setQueryData(cleanupKey({ organizationId: null }), { + ...batch(['deleted']), + createdAt: 'new-snapshot', + }); + }); + await advance(); + expect(latest).toHaveBeenCalledTimes(2); + } finally { + consoleError.mockRestore(); + } + }); +}); diff --git a/apps/web/src/components/cloud-agent-next/MergedWorktreeCleanup.tsx b/apps/web/src/components/cloud-agent-next/MergedWorktreeCleanup.tsx new file mode 100644 index 00000000000..a2afd3efc0e --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/MergedWorktreeCleanup.tsx @@ -0,0 +1,171 @@ +'use client'; + +import { useState } from 'react'; +import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; +import { BrushCleaning, Check, LoaderCircle } from 'lucide-react'; +import { formatSessionError } from '@kilocode/cloud-agent-sdk'; +import { useTRPC } from '@/lib/trpc/utils'; +import { Button } from '@/components/ui/button'; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from '@/components/ui/alert-dialog'; +import type { WorktreeCleanup } from './hooks/useWorktreeCleanup'; + +export function MergedWorktreeCleanup({ + organizationId, + cleanup, + statusError, + statusLoading, +}: { + organizationId?: string; + cleanup: WorktreeCleanup | null | undefined; + statusError: { message: string } | null; + statusLoading: boolean; +}) { + const trpc = useTRPC(); + const queryClient = useQueryClient(); + const [open, setOpen] = useState(false); + const scope = { organizationId: organizationId ?? null }; + const preview = useQuery({ + ...trpc.cliSessionsV2.mergedWorktreeCleanupPreview.queryOptions(scope), + enabled: open, + staleTime: 0, + }); + const start = useMutation( + trpc.cliSessionsV2.startMergedWorktreeCleanup.mutationOptions({ + onSuccess: batch => { + retry.reset(); + queryClient.setQueryData(trpc.cliSessionsV2.getWorktreeCleanup.queryKey(scope), batch); + setOpen(false); + }, + }) + ); + const retry = useMutation( + trpc.cliSessionsV2.retryWorktreeCleanup.mutationOptions({ + onSuccess: batch => { + queryClient.setQueryData(trpc.cliSessionsV2.getWorktreeCleanup.queryKey(scope), batch); + }, + }) + ); + const active = cleanup?.items.some( + item => item.status === 'pending' || item.status === 'deleting' + ); + const deleted = cleanup?.items.filter(item => item.status === 'deleted').length ?? 0; + const failed = cleanup?.items.filter(item => item.status === 'failed') ?? []; + const count = preview.data?.worktreeIds.length ?? 0; + + return ( +
+ + {cleanup && ( +
+
+ {!active && failed.length === 0 && } + + {deleted} of {cleanup.items.length} workspaces deleted + {failed.length > 0 ? `, ${failed.length} failed` : ''} + +
+ {active &&

You can close this tab.

} + {failed.length > 0 && ( + <> +
+ View cleanup errors +
    + {failed.map(item => ( +
  • + {item.worktreeId}: {item.error} +
    • + ))} +
+
+ + {retry.error && ( +

+ {formatSessionError(retry.error)} +

+ )} + + )} +
+ )} + {statusError && ( +

+ Cleanup progress is unavailable. It will refresh automatically. +

+ )} + !start.isPending && setOpen(value)}> + + + Clean up merged workspaces? + + {preview.isFetching + ? 'Checking your workspaces for merged pull requests...' + : count > 0 + ? `Permanently delete ${count} workspace${count === 1 ? '' : 's'} with merged GitHub pull requests, including all chats and files, and stop any running work.` + : 'No workspaces with merged pull requests are ready for cleanup.'} + + +

+ This checks all your workspaces in this account, regardless of sidebar filters. + Workspaces with unmerged or unverified pull requests are kept. Cleanup continues in the + background even if you close this tab. +

+ {(preview.error || start.error) && ( +

+ {formatSessionError(start.error ?? preview.error)} +

+ )} + + Cancel + { + if (preview.data) start.mutate({ ...scope, worktreeIds: preview.data.worktreeIds }); + }} + > + {start.isPending ? : null} + {start.isPending + ? 'Starting cleanup...' + : `Delete ${count} workspace${count === 1 ? '' : 's'}`} + + + + +
+ ); +} diff --git a/apps/web/src/components/cloud-agent-next/hooks/useWorktreeCleanup.ts b/apps/web/src/components/cloud-agent-next/hooks/useWorktreeCleanup.ts new file mode 100644 index 00000000000..cd09df87b39 --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/hooks/useWorktreeCleanup.ts @@ -0,0 +1,54 @@ +'use client'; + +import { useEffect, useEffectEvent } from 'react'; +import { useQuery, useQueryClient, type QueryClient } from '@tanstack/react-query'; +import { useTRPC } from '@/lib/trpc/utils'; +import type { inferRouterOutputs } from '@trpc/server'; +import type { RootRouter } from '@/routers/root-router'; + +export type WorktreeCleanup = NonNullable< + inferRouterOutputs['cliSessionsV2']['getWorktreeCleanup'] +>; + +const reconciledByClient = new WeakMap>(); + +export function useWorktreeCleanup({ + organizationId, + onDeleted, +}: { + organizationId?: string; + onDeleted: (worktreeId: string, deletedSessionIds: string[]) => Promise; +}) { + const trpc = useTRPC(); + const queryClient = useQueryClient(); + const cleanup = useQuery({ + ...trpc.cliSessionsV2.getWorktreeCleanup.queryOptions({ + organizationId: organizationId ?? null, + }), + refetchInterval: query => + query.state.data?.items.some(item => item.status === 'pending' || item.status === 'deleting') + ? 2000 + : 30_000, + }); + const reconcileDeletion = useEffectEvent(onDeleted); + + useEffect(() => { + if (!cleanup.data) return; + let reconciled = reconciledByClient.get(queryClient); + if (!reconciled) { + reconciled = new Set(); + reconciledByClient.set(queryClient, reconciled); + } + for (const item of cleanup.data.items) { + const key = JSON.stringify([organizationId ?? null, cleanup.data.id, item.worktreeId]); + if (item.status !== 'deleted' || reconciled.has(key)) continue; + reconciled.add(key); + void reconcileDeletion(item.worktreeId, item.deletedSessionIds).catch(error => { + reconciled.delete(key); + console.error('Failed to reconcile deleted workspace:', error); + }); + } + }, [cleanup.data, organizationId, queryClient]); + + return cleanup; +} diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts index 948f78f63b3..f3ab07965ae 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts @@ -185,6 +185,63 @@ describe('CloudAgentNextClient review message results', () => { }); }); +describe('CloudAgentNextClient worktree cleanup', () => { + const worktreeId = 'worktree_12345678-1234-4234-9234-123456789abc'; + const batch = { + id: 'cleanup-batch', + createdAt: '2026-10-20T12:00:00.000Z', + items: [{ worktreeId, status: 'pending', deletedSessionIds: [], error: null }], + }; + const mutate = + jest.fn<(input: CloudAgentClientModule.StartWorktreeCleanupInput) => Promise>(); + const query = + jest.fn<(input: CloudAgentClientModule.GetWorktreeCleanupInput) => Promise>(); + + beforeEach(() => { + mutate.mockReset().mockResolvedValue(batch); + query.mockReset().mockResolvedValue(batch); + mockCreateTRPCClient.mockReturnValueOnce({ + startWorktreeCleanup: { mutate }, + getWorktreeCleanup: { query }, + }); + }); + + it('forwards frozen worktree IDs and organization scope to the typed worker mutation', async () => { + const input: CloudAgentClientModule.StartWorktreeCleanupInput = { + worktreeIds: [worktreeId], + kilocodeOrganizationId: '12345678-1234-4234-9234-123456789abc', + }; + await expect(new CloudAgentNextClient('token').startWorktreeCleanup(input)).resolves.toEqual( + batch + ); + expect(mutate).toHaveBeenCalledWith(input); + }); + + it.each([{}, { kilocodeOrganizationId: '12345678-1234-4234-9234-123456789abc' }])( + 'forwards cleanup polling scope %j and preserves null', + async input => { + const client = new CloudAgentNextClient('token'); + await expect(client.getWorktreeCleanup(input)).resolves.toEqual(batch); + expect(query).toHaveBeenCalledWith(input); + query.mockResolvedValueOnce(null); + await expect(client.getWorktreeCleanup(input)).resolves.toBeNull(); + } + ); + + it('rejects malformed worker results rather than exposing unsupported statuses', async () => { + query.mockResolvedValueOnce({ ...batch, items: [{ ...batch.items[0], status: 'done' }] }); + await expect(new CloudAgentNextClient('token').getWorktreeCleanup({})).rejects.toThrow(); + }); + + it('propagates worker failures without storing or replacing them', async () => { + const error = new Error('worker unavailable'); + mutate.mockRejectedValueOnce(error); + await expect( + new CloudAgentNextClient('token').startWorktreeCleanup({ worktreeIds: [worktreeId] }) + ).rejects.toBe(error); + }); +}); + describe('CloudAgentNextClient worktree changes', () => { const cloudAgentSessionId = 'workspace_12345678-1234-4234-9234-123456789abc'; const snapshot: WorktreeChangesSnapshot = { diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts index 2553d43e00c..90abbbef267 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts @@ -38,6 +38,12 @@ import { type SandboxStatusSnapshot, } from '../../../../../services/cloud-agent-next/src/shared/sandbox-status'; import { fetchWithinBudget } from '@/lib/bounded-service-fetch'; +import { + WorktreeCleanup as worktreeCleanupOutputSchema, + type WorktreeCleanup, + type WorktreeCleanupScope, + type StartWorktreeCleanupInput as WorkerStartWorktreeCleanupInput, +} from '../../../../../services/cloud-agent-next/src/worktree-cleanup/schemas'; export type { SendMessagePayload } from './types.js'; const CLOUD_AGENT_NEXT_API_URL = getEnvVariable('CLOUD_AGENT_NEXT_API_URL') || ''; @@ -247,6 +253,10 @@ const deleteWorktreeOutputSchema = z.object({ export type DeleteWorktreeOutput = z.infer; +export type WorktreeCleanupOutput = WorktreeCleanup; +export type GetWorktreeCleanupInput = z.infer; +export type StartWorktreeCleanupInput = z.infer; + /** Input for initiating from a prepared session */ export type InitiateFromPreparedSessionInput = { cloudAgentSessionId: string; @@ -618,6 +628,12 @@ type CloudAgentNextTRPCClient = { deleteWorktree: { mutate: (input: DeleteWorktreeInput) => Promise; }; + startWorktreeCleanup: { + mutate: (input: StartWorktreeCleanupInput) => Promise; + }; + getWorktreeCleanup: { + query: (input: GetWorktreeCleanupInput) => Promise; + }; cleanupSession: { mutate: (input: { sessionId: string }) => Promise<{ success: boolean; message?: string }>; }; @@ -802,6 +818,16 @@ export class CloudAgentNextClient { } } + async startWorktreeCleanup(input: StartWorktreeCleanupInput): Promise { + return worktreeCleanupOutputSchema.parse(await this.client.startWorktreeCleanup.mutate(input)); + } + + async getWorktreeCleanup(input: GetWorktreeCleanupInput): Promise { + return worktreeCleanupOutputSchema + .nullable() + .parse(await this.client.getWorktreeCleanup.query(input)); + } + /** * Clean up a caller-created session and classify its runtime deletion as caller cleanup. */ diff --git a/apps/web/src/routers/cli-sessions-v2-router.ts b/apps/web/src/routers/cli-sessions-v2-router.ts index 202ca010b29..a784dbfca7a 100644 --- a/apps/web/src/routers/cli-sessions-v2-router.ts +++ b/apps/web/src/routers/cli-sessions-v2-router.ts @@ -30,7 +30,10 @@ import { organization_memberships, type CliSessionV2, } from '@kilocode/db/schema'; -import { createCloudAgentNextClient } from '@/lib/cloud-agent-next/cloud-agent-client'; +import { + createCloudAgentNextClient, + rethrowCloudAgentError, +} from '@/lib/cloud-agent-next/cloud-agent-client'; import { createControlTokenForRequest } from '@/lib/auth/resource-delegation'; import { readFirstPageStreamState, @@ -52,6 +55,7 @@ import { projectPrivateWorktreePaths, validateKiloSdkMessagesCursor, type SessionDiffRead, + type CloudAgentWorktreeId, } from '@kilocode/session-ingest-contracts'; import { baseGetSessionNextOutputSchema } from './cloud-agent-next-schemas'; import { projectSessionGoal } from '@kilocode/cloud-agent-sdk'; @@ -683,6 +687,9 @@ const WorktreeInputSchema = z.object({ organizationId: z.uuid().nullable(), }); +const WorktreeCleanupScopeSchema = z.object({ organizationId: z.uuid().nullable() }); +const MAX_WORKTREE_CLEANUP_SIZE = 1000; + const WorktreeDetailsInputSchema = z.object({ worktreeIds: z .array(cloudAgentWorktreeIdSchema) @@ -954,6 +961,70 @@ function joinWithAnd(fragments: SQL[]): SQL { return sql.join(fragments, sql` AND `); } +async function mergedWorktreeCleanupIds( + ctx: Parameters[0], + organizationId: string | null, + worktreeIds?: string[] +): Promise { + const conditions: SQL[] = [eq(cli_sessions_v2.kilo_user_id, ctx.user.id)]; + await addOrganizationCondition(conditions, ctx, organizationId); + if (worktreeIds) { + conditions.push(inArray(cli_sessions_v2.cloud_agent_worktree_id, worktreeIds)); + } + const mergedPr = sql`( + ${cli_sessions_v2.platform} = 'github' + AND ${cli_sessions_v2.pr_url} IS NOT NULL + AND ${cli_sessions_v2.pr_link_verified_at} IS NOT NULL + AND ${github_branch_pull_requests.pr_state} = 'merged' + )`; + const { rows } = await db.execute<{ worktree_id: string }>(sql` + WITH RECURSIVE worktree_sessions AS ( + SELECT ${cli_sessions_v2.session_id} AS session_id, + ${cli_sessions_v2.cloud_agent_worktree_id} AS worktree_id, + ${cli_sessions_v2.cloud_agent_session_id} AS cloud_agent_session_id + FROM ${cli_sessions_v2} + WHERE ${and(...conditions)} + AND ${cli_sessions_v2.cloud_agent_worktree_id} IS NOT NULL + UNION + SELECT child.session_id, parent.worktree_id, child.cloud_agent_session_id + FROM ${cli_sessions_v2} child + JOIN worktree_sessions parent ON child.parent_session_id = parent.session_id + OR child.cloud_agent_session_scope_id = parent.cloud_agent_session_id + WHERE child.kilo_user_id = ${ctx.user.id} + AND child.organization_id IS NOT DISTINCT FROM ${organizationId}::uuid + ) + SELECT worktree_sessions.worktree_id + FROM worktree_sessions + JOIN ${cli_sessions_v2} ON ${cli_sessions_v2.session_id} = worktree_sessions.session_id + AND ${cli_sessions_v2.kilo_user_id} = ${ctx.user.id} + LEFT JOIN ${github_branch_pull_requests} ON ${sessionPrJoinPredicate} + LEFT JOIN ${cloud_agent_worktrees} + ON ${cloud_agent_worktrees.worktree_id} = worktree_sessions.worktree_id + WHERE ${worktreeMembershipCondition(ctx.user.id, organizationId) ?? sql`true`} + AND (${cloud_agent_worktrees.worktree_id} IS NULL OR ( + ${cloud_agent_worktrees.kilo_user_id} = ${ctx.user.id} + AND ${cloud_agent_worktrees.organization_id} IS NOT DISTINCT FROM ${organizationId}::uuid + )) + AND ${cloud_agent_worktrees.deletion_completed_at} IS NULL + GROUP BY worktree_sessions.worktree_id + HAVING bool_or(COALESCE(${mergedPr}, false)) + AND NOT bool_or( + (${cli_sessions_v2.pr_url} IS NOT NULL OR ${cli_sessions_v2.pr_number} IS NOT NULL) + AND NOT COALESCE(${mergedPr}, false) + ) + ORDER BY worktree_sessions.worktree_id + LIMIT ${MAX_WORKTREE_CLEANUP_SIZE + 1} + `); + if (rows.length > MAX_WORKTREE_CLEANUP_SIZE) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: + 'More than 1000 merged worktrees are eligible for cleanup. Clean up some worktrees individually first.', + }); + } + return z.array(cloudAgentWorktreeIdSchema).parse(rows.map(row => row.worktree_id)); +} + /** * Router for cli_sessions_v2 table operations. * Used by cloud-agent-next for session storage and retrieval. @@ -962,6 +1033,109 @@ function joinWithAnd(fragments: SQL[]): SQL { * This router only queries the data. */ export const cliSessionsV2Router = createTRPCRouter({ + mergedWorktreeCleanupPreview: baseProcedure + .input(WorktreeCleanupScopeSchema) + .output(z.object({ worktreeIds: z.array(cloudAgentWorktreeIdSchema) })) + .query(async ({ ctx, input }) => ({ + worktreeIds: await mergedWorktreeCleanupIds(ctx, input.organizationId), + })), + + startMergedWorktreeCleanup: baseProcedure + .input( + WorktreeCleanupScopeSchema.extend({ + worktreeIds: z.array(cloudAgentWorktreeIdSchema).max(MAX_WORKTREE_CLEANUP_SIZE), + }) + ) + .mutation(async ({ ctx, input }) => { + const worktreeIds = await mergedWorktreeCleanupIds( + ctx, + input.organizationId, + input.worktreeIds + ); + if (worktreeIds.length === 0) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'None of the selected worktrees are still eligible for merged-worktree cleanup.', + }); + } + const client = createCloudAgentNextClient( + ( + await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + organizationId: input.organizationId ?? undefined, + tokenSource: 'cloud-agent', + }) + ).token + ); + try { + return await client.startWorktreeCleanup({ + worktreeIds, + ...(input.organizationId === null + ? {} + : { kilocodeOrganizationId: input.organizationId }), + }); + } catch (error) { + rethrowCloudAgentError(error); + } + }), + + retryWorktreeCleanup: baseProcedure + .input(WorktreeCleanupScopeSchema.strict()) + .mutation(async ({ ctx, input }) => { + if (input.organizationId !== null) { + await ensureOrganizationAccess(ctx, input.organizationId); + } + const client = createCloudAgentNextClient( + ( + await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + organizationId: input.organizationId ?? undefined, + tokenSource: 'cloud-agent', + }) + ).token + ); + const scope = + input.organizationId === null ? {} : { kilocodeOrganizationId: input.organizationId }; + try { + const batch = await client.getWorktreeCleanup(scope); + const worktreeIds = + batch?.items.filter(item => item.status === 'failed').map(item => item.worktreeId) ?? []; + if (worktreeIds.length === 0) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'There are no failed worktrees to retry in this cleanup batch.', + }); + } + return await client.startWorktreeCleanup({ ...scope, worktreeIds }); + } catch (error) { + rethrowCloudAgentError(error); + } + }), + + getWorktreeCleanup: baseProcedure + .input(WorktreeCleanupScopeSchema) + .query(async ({ ctx, input }) => { + if (input.organizationId !== null) { + await ensureOrganizationAccess(ctx, input.organizationId); + } + const client = createCloudAgentNextClient( + ( + await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + organizationId: input.organizationId ?? undefined, + tokenSource: 'cloud-agent', + }) + ).token + ); + try { + return await client.getWorktreeCleanup( + input.organizationId === null ? {} : { kilocodeOrganizationId: input.organizationId } + ); + } catch (error) { + rethrowCloudAgentError(error); + } + }), + worktreeDetails: baseProcedure.input(WorktreeDetailsInputSchema).query(async ({ ctx, input }) => { const worktrees: Record = {}; if (input.worktreeIds.length === 0) { diff --git a/apps/web/src/routers/cli-sessions-v2-worktree.test.ts b/apps/web/src/routers/cli-sessions-v2-worktree.test.ts index 00f45a1658c..71f30d0d819 100644 --- a/apps/web/src/routers/cli-sessions-v2-worktree.test.ts +++ b/apps/web/src/routers/cli-sessions-v2-worktree.test.ts @@ -11,7 +11,10 @@ import { type NewCliSessionV2, type User, } from '@kilocode/db/schema'; -import type { CloudAgentWorktreeId } from '@kilocode/session-ingest-contracts'; +import { + cloudAgentWorktreeIdSchema, + type CloudAgentWorktreeId, +} from '@kilocode/session-ingest-contracts'; import { TRPCClientError } from '@trpc/client'; import type * as TrpcServerModule from '@trpc/server'; import { and, eq, inArray, isNotNull, or } from 'drizzle-orm'; @@ -20,6 +23,9 @@ import type { cliSessionsV2Router } from './cli-sessions-v2-router'; import type { DeleteWorktreeInput, DeleteWorktreeOutput, + StartWorktreeCleanupInput, + WorktreeCleanupOutput, + GetWorktreeCleanupInput, } from '@/lib/cloud-agent-next/cloud-agent-client'; const UUID = '12345678-1234-4234-9234-123456789abc'; @@ -38,6 +44,10 @@ const PR_URL = `${GIT_URL}/pull/42`; const mockRuntimeDelete = jest.fn<(sessionId: string) => Promise<{ success: boolean }>>(); const mockWorktreeDelete = jest.fn<(input: DeleteWorktreeInput) => Promise>(); +const mockStartWorktreeCleanup = + jest.fn<(input: StartWorktreeCleanupInput) => Promise>(); +const mockGetWorktreeCleanup = + jest.fn<(input: GetWorktreeCleanupInput) => Promise>(); const mockGetRuntimeSession = jest.fn<(sessionId: string) => Promise>>(); const mockDeleteSessionIngest = jest.fn<(sessionId: string, userId: string) => Promise>(); const mockFetchSessionSnapshot = jest.fn<() => Promise>(); @@ -50,6 +60,8 @@ const mockCreateCloudAgentNextClient = jest.fn((_token: string) => ({ deleteSession: mockRuntimeDelete, deleteWorktree: mockWorktreeDelete, getSession: mockGetRuntimeSession, + startWorktreeCleanup: mockStartWorktreeCleanup, + getWorktreeCleanup: mockGetWorktreeCleanup, })); jest.mock('@kilocode/web-shared/lib/trpc/init', () => { @@ -69,6 +81,9 @@ jest.mock('@sentry/nextjs', () => ({ jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ createCloudAgentNextClient: mockCreateCloudAgentNextClient, + rethrowCloudAgentError: (error: unknown): never => { + throw error; + }, })); jest.mock('@kilocode/web-shared/lib/tokens', () => ({ @@ -255,6 +270,521 @@ beforeEach(async () => { ); }); mockFetchSessionSnapshot.mockReset().mockResolvedValue(null); + mockStartWorktreeCleanup.mockReset().mockImplementation(async input => ({ + id: 'cleanup-batch', + createdAt: INITIAL_TIME, + items: input.worktreeIds.map(worktreeId => ({ + worktreeId, + status: 'pending', + deletedSessionIds: [], + error: null, + })), + })); + mockGetWorktreeCleanup.mockReset().mockResolvedValue(null); +}); + +describe('merged worktree cleanup', () => { + async function mergedSession(overrides: Partial = {}, state = 'merged') { + const session = await insertSession({ + cloud_agent_worktree_id: newWorktreeId(), + platform: 'github', + git_url: GIT_URL, + pr_url: PR_URL, + pr_number: 42, + pr_link_verified_at: INITIAL_TIME, + ...overrides, + }); + await db.insert(github_branch_pull_requests).values({ + git_url: GIT_URL, + git_branch: `branch-${crypto.randomUUID()}`, + owned_by_user_id: session.organization_id === null ? session.kilo_user_id : null, + owned_by_organization_id: session.organization_id, + pr_url: session.pr_url ?? PR_URL, + pr_number: session.pr_number ?? 42, + pr_state: state, + }); + return { + ...session, + cloud_agent_worktree_id: cloudAgentWorktreeIdSchema.parse(session.cloud_agent_worktree_id), + }; + } + + it('enumerates all eligible worktrees regardless of list filters, dates, and pagination', async () => { + const sessions = await Promise.all( + Array.from({ length: 12 }, (_, index) => + mergedSession({ + pr_number: index + 1, + pr_url: `${GIT_URL}/pull/${index + 1}`, + created_at: '2020-01-01T00:00:00.000Z', + created_on_platform: 'app-builder', + }) + ) + ); + const result = await createCaller({ user }).mergedWorktreeCleanupPreview({ + organizationId: null, + }); + expect(result.worktreeIds.sort()).toEqual( + sessions.map(session => session.cloud_agent_worktree_id).sort() + ); + }); + + it.each(['open', 'closed', 'unknown', 'unverified', 'gitlab', 'missing-cache'])( + 'rejects a worktree with a merged PR and another %s PR, including descendants without a worktree ID', + async state => { + const root = await mergedSession(); + const child = await insertSession({ + parent_session_id: root.session_id, + cloud_agent_worktree_id: null, + }); + await insertSession({ + parent_session_id: child.session_id, + cloud_agent_worktree_id: null, + platform: state === 'gitlab' ? 'gitlab' : 'github', + git_url: GIT_URL, + pr_url: `${GIT_URL}/pull/99`, + pr_number: 99, + pr_link_verified_at: state === 'unverified' ? null : INITIAL_TIME, + }); + if (state !== 'missing-cache') { + await db.insert(github_branch_pull_requests).values({ + git_url: GIT_URL, + git_branch: 'other-pr', + owned_by_user_id: USER_ID, + pr_url: `${GIT_URL}/pull/99`, + pr_number: 99, + pr_state: ['open', 'closed', 'unknown'].includes(state) ? state : 'merged', + }); + } + await expect( + createCaller({ user }).mergedWorktreeCleanupPreview({ organizationId: null }) + ).resolves.toEqual({ worktreeIds: [] }); + } + ); + + it('accepts multiple verified merged PRs in the same worktree', async () => { + const first = await mergedSession(); + await mergedSession({ + cloud_agent_worktree_id: first.cloud_agent_worktree_id, + pr_url: `${GIT_URL}/pull/99`, + pr_number: 99, + }); + await expect( + createCaller({ user }).mergedWorktreeCleanupPreview({ organizationId: null }) + ).resolves.toEqual({ worktreeIds: [first.cloud_agent_worktree_id] }); + }); + + it('rejects mixed merged and open PRs across root sessions in the same worktree', async () => { + const first = await mergedSession(); + await mergedSession( + { + cloud_agent_worktree_id: first.cloud_agent_worktree_id, + pr_url: `${GIT_URL}/pull/99`, + pr_number: 99, + }, + 'open' + ); + await expect( + createCaller({ user }).mergedWorktreeCleanupPreview({ organizationId: null }) + ).resolves.toEqual({ worktreeIds: [] }); + }); + + it('rejects worktrees whose metadata belongs to another owner or tenant', async () => { + const otherOwner = await mergedSession(); + const otherTenant = await mergedSession({ pr_number: 99 }); + await db.insert(cloud_agent_worktrees).values([ + { worktree_id: otherOwner.cloud_agent_worktree_id, kilo_user_id: OTHER_USER_ID }, + { + worktree_id: otherTenant.cloud_agent_worktree_id, + kilo_user_id: USER_ID, + organization_id: ORGANIZATION_ID, + }, + ]); + await expect( + createCaller({ user }).mergedWorktreeCleanupPreview({ organizationId: null }) + ).resolves.toEqual({ worktreeIds: [] }); + await expect( + createCaller({ user }).startMergedWorktreeCleanup({ + organizationId: null, + worktreeIds: [otherOwner.cloud_agent_worktree_id, otherTenant.cloud_agent_worktree_id], + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(mockStartWorktreeCleanup).not.toHaveBeenCalled(); + }); + + it('does not trust unverified, branch-only, non-GitHub, or absent PR links', async () => { + await mergedSession({ pr_link_verified_at: null }); + await mergedSession({ pr_number: 43, pr_url: null }); + await mergedSession({ pr_number: 44, platform: 'gitlab' }); + await insertSession({ git_url: GIT_URL, git_branch: BRANCH }); + await expect( + createCaller({ user }).mergedWorktreeCleanupPreview({ organizationId: null }) + ).resolves.toEqual({ worktreeIds: [] }); + }); + + it('isolates owners and organization scopes, including PR cache ownership', async () => { + const personal = await mergedSession(); + const org = await mergedSession({ organization_id: ORGANIZATION_ID }); + await mergedSession({ organization_id: OTHER_ORGANIZATION_ID }); + await mergedSession({ kilo_user_id: OTHER_USER_ID }); + await insertSession({ + cloud_agent_worktree_id: newWorktreeId(), + platform: 'github', + git_url: GIT_URL, + pr_url: PR_URL, + pr_number: 42, + pr_link_verified_at: INITIAL_TIME, + kilo_user_id: OTHER_USER_ID, + organization_id: ORGANIZATION_ID, + }); + const caller = createCaller({ user }); + await expect(caller.mergedWorktreeCleanupPreview({ organizationId: null })).resolves.toEqual({ + worktreeIds: [personal.cloud_agent_worktree_id], + }); + await expect( + caller.mergedWorktreeCleanupPreview({ organizationId: ORGANIZATION_ID }) + ).resolves.toEqual({ worktreeIds: [org.cloud_agent_worktree_id] }); + await db + .delete(github_branch_pull_requests) + .where(eq(github_branch_pull_requests.owned_by_user_id, USER_ID)); + await expect(caller.mergedWorktreeCleanupPreview({ organizationId: null })).resolves.toEqual({ + worktreeIds: [], + }); + }); + + it.each([false, true])( + 'ignores colliding foreign session PRs when owner eligibility is %s', + async eligible => { + const owner = eligible + ? await mergedSession() + : await insertSession({ cloud_agent_worktree_id: newWorktreeId() }); + await mergedSession( + { session_id: owner.session_id, kilo_user_id: OTHER_USER_ID }, + eligible ? 'open' : 'merged' + ); + const caller = createCaller({ user }); + await expect(caller.mergedWorktreeCleanupPreview({ organizationId: null })).resolves.toEqual({ + worktreeIds: eligible ? [owner.cloud_agent_worktree_id] : [], + }); + const start = caller.startMergedWorktreeCleanup({ + organizationId: null, + worktreeIds: [cloudAgentWorktreeIdSchema.parse(owner.cloud_agent_worktree_id)], + }); + if (eligible) { + await expect(start).resolves.toMatchObject({ id: 'cleanup-batch' }); + expect(mockStartWorktreeCleanup).toHaveBeenCalledTimes(1); + } else { + await expect(start).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(mockStartWorktreeCleanup).not.toHaveBeenCalled(); + } + } + ); + + it('allows restarting eligible unfinished deletions without clearing the deletion fence', async () => { + const session = await mergedSession(); + await db.insert(cloud_agent_worktrees).values({ + worktree_id: session.cloud_agent_worktree_id, + kilo_user_id: USER_ID, + deletion_started_at: INITIAL_TIME, + }); + const caller = createCaller({ user }); + await expect(caller.mergedWorktreeCleanupPreview({ organizationId: null })).resolves.toEqual({ + worktreeIds: [session.cloud_agent_worktree_id], + }); + await caller.startMergedWorktreeCleanup({ + organizationId: null, + worktreeIds: [session.cloud_agent_worktree_id], + }); + expect(mockStartWorktreeCleanup).toHaveBeenCalledWith({ + worktreeIds: [session.cloud_agent_worktree_id], + }); + const [worktree] = await db + .select() + .from(cloud_agent_worktrees) + .where(eq(cloud_agent_worktrees.worktree_id, session.cloud_agent_worktree_id)); + expect(new Date(worktree.deletion_started_at!).toISOString()).toBe(INITIAL_TIME); + }); + + it('continues to exclude completed deletions', async () => { + const session = await mergedSession(); + await db.insert(cloud_agent_worktrees).values({ + worktree_id: session.cloud_agent_worktree_id, + kilo_user_id: USER_ID, + deletion_started_at: INITIAL_TIME, + deletion_completed_at: INITIAL_TIME, + }); + await expect( + createCaller({ user }).mergedWorktreeCleanupPreview({ organizationId: null }) + ).resolves.toEqual({ worktreeIds: [] }); + }); + + it.each(['open', 'unverified'])( + 'rejects a %s PR on a scope-linked session without parent or worktree identity', + async state => { + const root = await mergedSession(); + await insertSession({ + parent_session_id: null, + cloud_agent_worktree_id: null, + cloud_agent_session_id: null, + cloud_agent_session_scope_id: root.cloud_agent_session_id, + platform: 'github', + git_url: GIT_URL, + pr_url: `${GIT_URL}/pull/99`, + pr_number: 99, + pr_link_verified_at: state === 'unverified' ? null : INITIAL_TIME, + }); + await db.insert(github_branch_pull_requests).values({ + git_url: GIT_URL, + git_branch: 'scope-linked', + owned_by_user_id: USER_ID, + pr_url: `${GIT_URL}/pull/99`, + pr_number: 99, + pr_state: state === 'open' ? 'open' : 'merged', + }); + const caller = createCaller({ user }); + await expect(caller.mergedWorktreeCleanupPreview({ organizationId: null })).resolves.toEqual({ + worktreeIds: [], + }); + await expect( + caller.startMergedWorktreeCleanup({ + organizationId: null, + worktreeIds: [root.cloud_agent_worktree_id], + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + } + ); + + it('keeps scope-linked membership inside the authenticated owner and organization', async () => { + const root = await mergedSession(); + for (const overrides of [ + { kilo_user_id: OTHER_USER_ID }, + { organization_id: ORGANIZATION_ID }, + ]) { + await insertSession({ + cloud_agent_worktree_id: null, + cloud_agent_session_id: null, + cloud_agent_session_scope_id: root.cloud_agent_session_id, + pr_url: `${GIT_URL}/pull/99`, + pr_number: 99, + ...overrides, + }); + } + await expect( + createCaller({ user }).mergedWorktreeCleanupPreview({ organizationId: null }) + ).resolves.toEqual({ worktreeIds: [root.cloud_agent_worktree_id] }); + }); + + it('rechecks and intersects frozen IDs without adding newly eligible worktrees', async () => { + const kept = await mergedSession({ pr_number: 1 }); + const changed = await mergedSession({ pr_number: 2 }); + await mergedSession({ pr_number: 3 }); + await db + .update(github_branch_pull_requests) + .set({ pr_state: 'open' }) + .where(eq(github_branch_pull_requests.pr_number, 2)); + const result = await createCaller({ user }).startMergedWorktreeCleanup({ + organizationId: null, + worktreeIds: [ + kept.cloud_agent_worktree_id, + changed.cloud_agent_worktree_id, + WORKTREE_ID, + kept.cloud_agent_worktree_id, + ], + }); + expect(mockStartWorktreeCleanup).toHaveBeenCalledWith({ + worktreeIds: [kept.cloud_agent_worktree_id], + }); + expect(result.items).toEqual([ + expect.objectContaining({ worktreeId: kept.cloud_agent_worktree_id }), + ]); + expect(mockGenerateCloudAgentToken).toHaveBeenCalledWith(user); + expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('cloud-agent-token'); + }); + + it('rejects empty or no-longer-eligible selections before creating a worker client', async () => { + const caller = createCaller({ user }); + for (const worktreeIds of [[], [WORKTREE_ID]]) { + await expect( + caller.startMergedWorktreeCleanup({ organizationId: null, worktreeIds }) + ).rejects.toMatchObject({ + code: 'BAD_REQUEST', + message: expect.stringContaining('still eligible'), + }); + } + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + it('forwards organization scope on start and get and returns the worker batch unchanged', async () => { + const session = await mergedSession({ organization_id: ORGANIZATION_ID }); + const caller = createCaller({ user }); + const batch = await caller.startMergedWorktreeCleanup({ + organizationId: ORGANIZATION_ID, + worktreeIds: [session.cloud_agent_worktree_id], + }); + expect(mockStartWorktreeCleanup).toHaveBeenCalledWith({ + worktreeIds: [session.cloud_agent_worktree_id], + kilocodeOrganizationId: ORGANIZATION_ID, + }); + mockGetWorktreeCleanup.mockResolvedValue(batch); + await expect(caller.getWorktreeCleanup({ organizationId: ORGANIZATION_ID })).resolves.toEqual( + batch + ); + expect(mockGetWorktreeCleanup).toHaveBeenCalledWith({ + kilocodeOrganizationId: ORGANIZATION_ID, + }); + await caller.getWorktreeCleanup({ organizationId: null }); + expect(mockGetWorktreeCleanup).toHaveBeenLastCalledWith({}); + }); + + it('denies removed organization members on preview, start, and get', async () => { + const session = await mergedSession({ organization_id: ORGANIZATION_ID }); + await db + .delete(organization_memberships) + .where( + and( + eq(organization_memberships.kilo_user_id, USER_ID), + eq(organization_memberships.organization_id, ORGANIZATION_ID) + ) + ); + const caller = createCaller({ user }); + await expect( + caller.mergedWorktreeCleanupPreview({ organizationId: ORGANIZATION_ID }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + await expect( + caller.startMergedWorktreeCleanup({ + organizationId: ORGANIZATION_ID, + worktreeIds: [session.cloud_agent_worktree_id], + }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + await expect( + caller.getWorktreeCleanup({ organizationId: ORGANIZATION_ID }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + it('rejects previews above the 1000-worktree batch limit instead of truncating', async () => { + await mergedSession(); + await db.insert(cli_sessions_v2).values( + Array.from({ length: 1000 }, () => + makeSession({ + cloud_agent_worktree_id: newWorktreeId(), + platform: 'github', + git_url: GIT_URL, + pr_url: PR_URL, + pr_number: 42, + pr_link_verified_at: INITIAL_TIME, + }) + ) + ); + await expect( + createCaller({ user }).mergedWorktreeCleanupPreview({ organizationId: null }) + ).rejects.toMatchObject({ + code: 'BAD_REQUEST', + message: expect.stringContaining('More than 1000'), + }); + }); + + it('rejects input larger than 1000 worktrees', async () => { + await expect( + createCaller({ user }).startMergedWorktreeCleanup({ + organizationId: null, + worktreeIds: Array.from({ length: 1001 }, newWorktreeId), + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(mockStartWorktreeCleanup).not.toHaveBeenCalled(); + }); +}); + +describe('retryWorktreeCleanup', () => { + function cleanupBatch( + statuses: WorktreeCleanupOutput['items'][number]['status'][] + ): WorktreeCleanupOutput { + return { + id: 'persisted-cleanup', + createdAt: INITIAL_TIME, + items: statuses.map(status => ({ + worktreeId: newWorktreeId(), + status, + deletedSessionIds: [], + error: status === 'failed' ? 'Cleanup attempts exhausted' : null, + })), + }; + } + + it.each([null, ORGANIZATION_ID])( + 'retrieves the persisted batch in scope %s and submits only its failed IDs', + async organizationId => { + const batch = cleanupBatch(['failed', 'deleted', 'pending', 'deleting', 'failed']); + mockGetWorktreeCleanup.mockResolvedValue(batch); + const result = await createCaller({ user }).retryWorktreeCleanup({ organizationId }); + const scope = organizationId === null ? {} : { kilocodeOrganizationId: organizationId }; + const worktreeIds = [batch.items[0].worktreeId, batch.items[4].worktreeId]; + expect(mockGetWorktreeCleanup).toHaveBeenCalledWith(scope); + expect(mockStartWorktreeCleanup).toHaveBeenCalledWith({ ...scope, worktreeIds }); + expect(result.items.map(item => item.worktreeId)).toEqual(worktreeIds); + expect(mockGenerateCloudAgentToken).toHaveBeenCalledWith(user); + expect(mockCreateCloudAgentNextClient).toHaveBeenCalledTimes(1); + expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('cloud-agent-token'); + } + ); + + it('cannot submit caller-provided worktree IDs', async () => { + mockGetWorktreeCleanup.mockResolvedValue(cleanupBatch(['failed'])); + const input = { organizationId: null, worktreeIds: [WORKTREE_ID] }; + await expect(createCaller({ user }).retryWorktreeCleanup(input)).rejects.toMatchObject({ + code: 'BAD_REQUEST', + }); + expect(mockGetWorktreeCleanup).not.toHaveBeenCalled(); + expect(mockStartWorktreeCleanup).not.toHaveBeenCalled(); + }); + + it.each([null, cleanupBatch([]), cleanupBatch(['deleted', 'pending', 'deleting'])])( + 'rejects a batch without failures: %j', + async batch => { + mockGetWorktreeCleanup.mockResolvedValue(batch); + await expect( + createCaller({ user }).retryWorktreeCleanup({ organizationId: null }) + ).rejects.toMatchObject({ + code: 'BAD_REQUEST', + message: expect.stringContaining('no failed worktrees'), + }); + expect(mockStartWorktreeCleanup).not.toHaveBeenCalled(); + } + ); + + it('returns the active worker batch unchanged when a retry races another start', async () => { + mockGetWorktreeCleanup.mockResolvedValue(cleanupBatch(['failed'])); + const active = cleanupBatch(['pending', 'deleting']); + mockStartWorktreeCleanup.mockResolvedValue(active); + await expect( + createCaller({ user }).retryWorktreeCleanup({ organizationId: null }) + ).resolves.toEqual(active); + }); + + it('denies removed membership before reading any persisted batch', async () => { + await db + .delete(organization_memberships) + .where( + and( + eq(organization_memberships.kilo_user_id, USER_ID), + eq(organization_memberships.organization_id, ORGANIZATION_ID) + ) + ); + await expect( + createCaller({ user }).retryWorktreeCleanup({ organizationId: ORGANIZATION_ID }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + expect(mockGetWorktreeCleanup).not.toHaveBeenCalled(); + expect(mockStartWorktreeCleanup).not.toHaveBeenCalled(); + }); + + it('does not start a retry if reading the authoritative batch fails', async () => { + const error = new Error('Batch unavailable'); + mockGetWorktreeCleanup.mockRejectedValue(error); + await expect( + createCaller({ user }).retryWorktreeCleanup({ organizationId: null }) + ).rejects.toMatchObject({ message: error.message }); + expect(mockStartWorktreeCleanup).not.toHaveBeenCalled(); + }); }); describe('cliSessionsV2 worktree projections and session deletion', () => { diff --git a/apps/web/src/scripts/mcp-catalog/catalog.test.ts b/apps/web/src/scripts/mcp-catalog/catalog.test.ts index 3c87b068ba3..84ee2645bfe 100644 --- a/apps/web/src/scripts/mcp-catalog/catalog.test.ts +++ b/apps/web/src/scripts/mcp-catalog/catalog.test.ts @@ -135,6 +135,10 @@ describe('mcp-catalog catalog', () => { expect(isDenylistedPath('agentBriefing.generateSummary')).toBe(true); expect(isDenylistedPath('agentBriefing.retrySummary')).toBe(true); expect(isDenylistedPath('agentBriefing.updateContextPreferences')).toBe(true); + expect(isDenylistedPath('cliSessionsV2.startMergedWorktreeCleanup')).toBe(true); + expect(isDenylistedPath('cliSessionsV2.retryWorktreeCleanup')).toBe(true); + expect(isDenylistedPath('cliSessionsV2.mergedWorktreeCleanupPreview')).toBe(false); + expect(isDenylistedPath('cliSessionsV2.getWorktreeCleanup')).toBe(false); expect(isDenylistedPath('agentBriefing.estimateSummary')).toBe(false); expect(isDenylistedPath('agentBriefing.dismiss')).toBe(false); expect(isDenylistedPath('activeSessions.getToken')).toBe(false); diff --git a/apps/web/src/scripts/mcp-catalog/catalog.ts b/apps/web/src/scripts/mcp-catalog/catalog.ts index bb417f13fcf..6bb53a8c71b 100644 --- a/apps/web/src/scripts/mcp-catalog/catalog.ts +++ b/apps/web/src/scripts/mcp-catalog/catalog.ts @@ -44,12 +44,16 @@ export const ROOT_ROUTER_PATH = join(__dirname, '..', '..', 'routers', 'root-rou * `agentBriefing` generation and consent mutations spend the context's balance * behind a cost disclosure that only a person in the app can acknowledge; an MCP * agent could otherwise mint the estimate token itself and bypass that consent. + * Bulk worktree cleanup permanently deletes chats and files behind an in-app + * confirmation, so starting and retrying it stay out of MCP. */ const NATIVE_ONLY_PATHS: Record = { 'activeSessions.widgetCredential': true, 'agentBriefing.generateSummary': true, 'agentBriefing.retrySummary': true, 'agentBriefing.updateContextPreferences': true, + 'cliSessionsV2.startMergedWorktreeCleanup': true, + 'cliSessionsV2.retryWorktreeCleanup': true, }; /** diff --git a/packages/session-ingest-contracts/src/rpc-contract.ts b/packages/session-ingest-contracts/src/rpc-contract.ts index 6d2adf20b81..6a41dc1e0d3 100644 --- a/packages/session-ingest-contracts/src/rpc-contract.ts +++ b/packages/session-ingest-contracts/src/rpc-contract.ts @@ -13,6 +13,7 @@ export const cloudAgentWorktreeLocationSchema = z .strict(); export type CloudAgentWorktreeLocation = z.infer; export const WORKTREE_RUNTIME_HISTORY_UNAVAILABLE = 'worktree_runtime_history_unavailable'; +export const WORKTREE_MERGED_PULL_REQUESTS_REQUIRED = 'worktree_merged_pull_requests_required'; export const cloudAgentWorktreeDeletionManifestSchema = z .object({ @@ -36,6 +37,7 @@ export const cloudAgentWorktreeDeletionSchema = z worktreeId: cloudAgentWorktreeIdSchema, kiloUserId: z.string().min(1), organizationId: z.uuid().optional(), + requireMergedPullRequests: z.literal(true).optional(), }) .strict(); export type CloudAgentWorktreeDeletionParams = z.infer; diff --git a/services/cloud-agent-next/src/e2e-entry.ts b/services/cloud-agent-next/src/e2e-entry.ts index 85790461033..c53d0606fdf 100644 --- a/services/cloud-agent-next/src/e2e-entry.ts +++ b/services/cloud-agent-next/src/e2e-entry.ts @@ -27,6 +27,7 @@ export { CloudAgentSession } from './persistence/CloudAgentSession.js'; export { SandboxControlV2 as SandboxControl } from './control-plane/sandbox/sandbox-do.js'; export { SandboxSessionV2 as SandboxSession } from './control-plane/session/session-do.js'; export { StreamTicketNonceDO } from './persistence/StreamTicketNonceDO.js'; +export { WorktreeCleanupCoordinator } from './worktree-cleanup/WorktreeCleanupCoordinator.js'; export { UserKiloFacade } from './kilo-facade/user-kilo-facade.js'; export { E2eCallbackSink } from './persistence/E2eCallbackSink.js'; diff --git a/services/cloud-agent-next/src/index.ts b/services/cloud-agent-next/src/index.ts index 759269b2f06..188ef5c65f2 100644 --- a/services/cloud-agent-next/src/index.ts +++ b/services/cloud-agent-next/src/index.ts @@ -17,4 +17,5 @@ export { SandboxControlV2 as SandboxControl } from './control-plane/sandbox/sand export { SandboxSessionV2 as SandboxSession } from './control-plane/session/session-do.js'; export { SandboxContainers } from './sandbox-containers/SandboxContainers.js'; export { StreamTicketNonceDO } from './persistence/StreamTicketNonceDO.js'; +export { WorktreeCleanupCoordinator } from './worktree-cleanup/WorktreeCleanupCoordinator.js'; export { UserKiloFacade } from './kilo-facade/user-kilo-facade.js'; diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index f153a45faef..6360b8166ad 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/services/cloud-agent-next/src/router.test.ts @@ -393,6 +393,7 @@ describe('router sessionId validation', () => { } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], + WORKTREE_CLEANUP: {} as TRPCContext['env']['WORKTREE_CLEANUP'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), @@ -850,6 +851,7 @@ describe('router sessionId validation', () => { } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], + WORKTREE_CLEANUP: {} as TRPCContext['env']['WORKTREE_CLEANUP'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), @@ -1003,6 +1005,7 @@ describe('router sessionId validation', () => { } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], + WORKTREE_CLEANUP: {} as TRPCContext['env']['WORKTREE_CLEANUP'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), @@ -1523,6 +1526,7 @@ describe('router sessionId validation', () => { } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], + WORKTREE_CLEANUP: {} as TRPCContext['env']['WORKTREE_CLEANUP'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn() } as unknown as TRPCContext['env']['SESSION_INGEST'], CONTAINER_USAGE_METER: {} as TRPCContext['env']['CONTAINER_USAGE_METER'], @@ -1643,6 +1647,7 @@ describe('router sessionId validation', () => { } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], + WORKTREE_CLEANUP: {} as TRPCContext['env']['WORKTREE_CLEANUP'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), @@ -1953,6 +1958,7 @@ describe('router sessionId validation', () => { } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], + WORKTREE_CLEANUP: {} as TRPCContext['env']['WORKTREE_CLEANUP'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), diff --git a/services/cloud-agent-next/src/router.ts b/services/cloud-agent-next/src/router.ts index 716d86c7726..39bab110153 100644 --- a/services/cloud-agent-next/src/router.ts +++ b/services/cloud-agent-next/src/router.ts @@ -14,12 +14,15 @@ import { createSessionStartHandlers } from './router/handlers/session-start.js'; import { createSessionSendHandlers } from './router/handlers/session-send.js'; import { createSessionWorktreeHandlers } from './router/handlers/session-worktree.js'; import { deleteWorktree } from './router/handlers/worktree-deletion.js'; +import { startWorktreeCleanup, getWorktreeCleanup } from './router/handlers/worktree-cleanup.js'; import { createSessionWorktreeChangesHandlers } from './router/handlers/session-worktree-changes.js'; import { getSandboxSelectionOptions } from './router/handlers/sandbox-selection.js'; import { getPreviewOptions } from './router/handlers/preview-options.js'; export const appRouter = router({ deleteWorktree, + startWorktreeCleanup, + getWorktreeCleanup, getSandboxSelectionOptions, getPreviewOptions, ...createSessionManagementHandlers(), diff --git a/services/cloud-agent-next/src/router/handlers/worktree-cleanup.test.ts b/services/cloud-agent-next/src/router/handlers/worktree-cleanup.test.ts new file mode 100644 index 00000000000..0ffd034a63e --- /dev/null +++ b/services/cloud-agent-next/src/router/handlers/worktree-cleanup.test.ts @@ -0,0 +1,112 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type * as WorkerUtils from '@kilocode/worker-utils'; +import type { TRPCContext } from '../../types.js'; +import { router } from '../auth.js'; +import { getWorktreeCleanup, startWorktreeCleanup } from './worktree-cleanup.js'; + +const mocks = vi.hoisted(() => ({ + access: vi.fn(), + getDb: vi.fn(), + getStub: vi.fn(), + start: vi.fn(), + get: vi.fn(), +})); +vi.mock('@kilocode/worker-utils', async importOriginal => ({ + ...(await importOriginal()), + hasOrganizationAccess: mocks.access, +})); +vi.mock('../../db/pg.js', () => ({ getPgDb: mocks.getDb })); +vi.mock('../../worktree-cleanup/stub.js', () => ({ getWorktreeCleanupStub: mocks.getStub })); + +const worktreeId = 'worktree_11111111-1111-4111-8111-111111111111'; +const organizationId = '22222222-2222-4222-8222-222222222222'; +const batch = { + id: 'batch', + createdAt: '2026-10-20T12:00:00.000Z', + items: [{ worktreeId, status: 'pending', deletedSessionIds: [], error: null }], +}; + +function fixture(override: Partial = {}) { + const ctx: TRPCContext = { + userId: 'oauth/google:owner', + authToken: 'not-persisted', + request: new Request('https://worker.test'), + env: {} as never, + ...override, + }; + return { ctx, caller: router({ startWorktreeCleanup, getWorktreeCleanup }).createCaller(ctx) }; +} + +beforeEach(() => { + vi.resetAllMocks(); + mocks.access.mockResolvedValue(true); + mocks.getStub.mockReturnValue({ start: mocks.start, get: mocks.get }); + mocks.start.mockResolvedValue(batch); + mocks.get.mockResolvedValue(null); +}); + +describe('worktree cleanup endpoints', () => { + it('passes only authenticated identity and scope to the durable coordinator', async () => { + const { caller, ctx } = fixture(); + expect(await caller.startWorktreeCleanup({ worktreeIds: [worktreeId] })).toEqual(batch); + expect(mocks.start).toHaveBeenCalledWith({ userId: ctx.userId }, { worktreeIds: [worktreeId] }); + expect(await caller.getWorktreeCleanup({})).toBeNull(); + expect(mocks.get).toHaveBeenCalledWith({ userId: ctx.userId }); + expect(mocks.access).not.toHaveBeenCalled(); + }); + + it.each(['start', 'get'] as const)( + 'rejects bots and unauthenticated callers for %s', + async operation => { + for (const override of [{ botId: 'bot' }, { authToken: '' }]) { + const { caller } = fixture(override); + const result = + operation === 'start' + ? caller.startWorktreeCleanup({ worktreeIds: [worktreeId] }) + : caller.getWorktreeCleanup({}); + await expect(result).rejects.toMatchObject({ + code: override.botId ? 'FORBIDDEN' : 'UNAUTHORIZED', + }); + } + expect(mocks.getStub).not.toHaveBeenCalled(); + expect(mocks.access).not.toHaveBeenCalled(); + } + ); + + it.each(['start', 'get'] as const)('reauthorizes organization access on %s', async operation => { + const { caller, ctx } = fixture(); + const execute = () => + operation === 'start' + ? caller.startWorktreeCleanup({ + worktreeIds: [worktreeId], + kilocodeOrganizationId: organizationId, + }) + : caller.getWorktreeCleanup({ kilocodeOrganizationId: organizationId }); + await execute(); + expect(mocks.access).toHaveBeenCalledWith(undefined, { + kiloUserId: ctx.userId, + organizationId, + }); + expect(mocks.getStub).toHaveBeenCalledWith(ctx.env, { + userId: ctx.userId, + kilocodeOrganizationId: organizationId, + }); + mocks.access.mockResolvedValue(false); + mocks.getStub.mockClear(); + await expect(execute()).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mocks.getStub).not.toHaveBeenCalled(); + }); + + it.each( + [[], ['invalid'], Array.from({ length: 1001 }, () => worktreeId)].map(worktreeIds => ({ + worktreeIds, + })) + )('rejects invalid or oversized frozen IDs', async ({ worktreeIds }) => { + await expect( + fixture().caller.startWorktreeCleanup({ worktreeIds: worktreeIds as `worktree_${string}`[] }) + ).rejects.toMatchObject({ + code: 'BAD_REQUEST', + }); + expect(mocks.getStub).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/router/handlers/worktree-cleanup.ts b/services/cloud-agent-next/src/router/handlers/worktree-cleanup.ts new file mode 100644 index 00000000000..13f6c7b18e2 --- /dev/null +++ b/services/cloud-agent-next/src/router/handlers/worktree-cleanup.ts @@ -0,0 +1,50 @@ +import { hasOrganizationAccess } from '@kilocode/worker-utils'; +import { TRPCError } from '@trpc/server'; +import { getPgDb } from '../../db/pg.js'; +import type { TRPCContext } from '../../types.js'; +import { withDORetry } from '../../utils/do-retry.js'; +import { + StartWorktreeCleanupInput, + WorktreeCleanup, + WorktreeCleanupScope, +} from '../../worktree-cleanup/schemas.js'; +import { getWorktreeCleanupStub } from '../../worktree-cleanup/stub.js'; +import { protectedProcedure } from '../auth.js'; + +async function authorize(ctx: TRPCContext, kilocodeOrganizationId?: string) { + if ( + ctx.botId !== undefined || + (kilocodeOrganizationId && + !(await hasOrganizationAccess(getPgDb(ctx.env), { + kiloUserId: ctx.userId, + organizationId: kilocodeOrganizationId, + }))) + ) { + throw new TRPCError({ code: 'FORBIDDEN', message: 'Worktree access denied' }); + } + return { userId: ctx.userId, ...(kilocodeOrganizationId ? { kilocodeOrganizationId } : {}) }; +} + +export const startWorktreeCleanup = protectedProcedure + .input(StartWorktreeCleanupInput) + .output(WorktreeCleanup) + .mutation(async ({ input, ctx }) => { + const owner = await authorize(ctx, input.kilocodeOrganizationId); + return withDORetry( + () => getWorktreeCleanupStub(ctx.env, owner), + stub => stub.start(owner, input), + 'startWorktreeCleanup' + ); + }); + +export const getWorktreeCleanup = protectedProcedure + .input(WorktreeCleanupScope) + .output(WorktreeCleanup.nullable()) + .query(async ({ input, ctx }) => { + const owner = await authorize(ctx, input.kilocodeOrganizationId); + return withDORetry( + () => getWorktreeCleanupStub(ctx.env, owner), + stub => stub.get(owner), + 'getWorktreeCleanup' + ); + }); diff --git a/services/cloud-agent-next/src/router/handlers/worktree-deletion.test.ts b/services/cloud-agent-next/src/router/handlers/worktree-deletion.test.ts index 465d25cabca..9e979c7cf79 100644 --- a/services/cloud-agent-next/src/router/handlers/worktree-deletion.test.ts +++ b/services/cloud-agent-next/src/router/handlers/worktree-deletion.test.ts @@ -148,6 +148,82 @@ function fixture(rootCount = 1, childCount = 0) { beforeEach(() => vi.resetAllMocks()); describe('deleteWorktree authorization and completion', () => { + it('checks merged-only eligibility before destruction, not when recording its result', async () => { + const f = fixture(); + await f.caller.deleteWorktree({ worktreeId, requireMergedPullRequests: true }); + expect(f.begin).toHaveBeenCalledWith({ + worktreeId, + kiloUserId: userId, + requireMergedPullRequests: true, + }); + expect(f.record).toHaveBeenCalledTimes(2); + expect(f.record.mock.calls[0][0].requireMergedPullRequests).toBe(true); + expect(f.record.mock.calls[1][0]).not.toHaveProperty('requireMergedPullRequests'); + expect(f.complete).toHaveBeenCalledWith({ + worktreeId, + kiloUserId: userId, + requireMergedPullRequests: true, + }); + }); + + it.each([undefined, organizationId])( + 'records runtime-discovered sessions without rechecking merged eligibility in scope %s', + async kilocodeOrganizationId => { + const f = fixture(); + const discoveredSessionId = kiloId(1); + const record = f.record.getMockImplementation(); + if (!record) throw new Error('Missing record fixture'); + f.cleanup.mockResolvedValue({ + deleted: true, + sessionIds: [kiloId(0), discoveredSessionId], + }); + f.record.mockImplementation(async input => { + if (input.requireMergedPullRequests && input.sessionIds?.includes(discoveredSessionId)) { + throw new Error('worktree_merged_pull_requests_required'); + } + return record(input); + }); + + await expect( + f.caller.deleteWorktree({ + worktreeId, + kilocodeOrganizationId, + requireMergedPullRequests: true, + }) + ).resolves.toEqual({ + success: true, + deletedSessionIds: [kiloId(0), discoveredSessionId], + }); + expect(f.record.mock.calls[1][0]).toEqual({ + worktreeId, + kiloUserId: userId, + ...(kilocodeOrganizationId ? { organizationId: kilocodeOrganizationId } : {}), + directory: getWorktreeWorkspacePath(kilocodeOrganizationId, userId, worktreeId), + sessionIds: [kiloId(0), discoveredSessionId], + }); + expect(f.finishSession).toHaveBeenCalledTimes(1); + expect(f.complete).toHaveBeenCalledTimes(1); + } + ); + + it.each(['begin', 'record'] as const)( + 'reports permanent merged-only rejection at %s without wiping resources', + async phase => { + const f = fixture(); + f[phase].mockRejectedValue(new Error('worktree_merged_pull_requests_required')); + await expect( + f.caller.deleteWorktree({ worktreeId, requireMergedPullRequests: true }) + ).rejects.toMatchObject({ + code: 'PRECONDITION_FAILED', + cause: { error: 'WORKTREE_MERGED_PULL_REQUESTS_REQUIRED', retryable: false }, + }); + expect(f.cleanup).not.toHaveBeenCalled(); + expect(f.finishSession).not.toHaveBeenCalled(); + expect(f.complete).not.toHaveBeenCalled(); + if (phase === 'begin') expect(f.beginSession).not.toHaveBeenCalled(); + } + ); + it.each([false, true])('reports aggregate phase timings when cleanup fails=%s', async fail => { const f = fixture(2); let now = 1000; diff --git a/services/cloud-agent-next/src/router/handlers/worktree-deletion.ts b/services/cloud-agent-next/src/router/handlers/worktree-deletion.ts index 2dc5b30af46..cb29b9763ad 100644 --- a/services/cloud-agent-next/src/router/handlers/worktree-deletion.ts +++ b/services/cloud-agent-next/src/router/handlers/worktree-deletion.ts @@ -7,6 +7,7 @@ import { cloudAgentChildSessionLineageSchema, sessionIdSchema, WORKTREE_RUNTIME_HISTORY_UNAVAILABLE, + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED, type CloudAgentWorktreeDeletionParams, type RecordCloudAgentWorktreeCleanupParams, } from '@kilocode/session-ingest-contracts'; @@ -47,6 +48,7 @@ export const DeleteWorktreeInput = z .object({ worktreeId: cloudAgentWorktreeIdSchema, kilocodeOrganizationId: z.uuid().optional(), + requireMergedPullRequests: z.literal(true).optional(), }) .strict(); @@ -59,13 +61,14 @@ export const DeleteWorktreeOutput = z export async function deleteWorktreeResources( input: z.infer, - ctx: TRPCContext + ctx: Pick ): Promise> { if (ctx.botId !== undefined) throw new TRPCError({ code: 'FORBIDDEN', message: 'Worktree access denied' }); const params: CloudAgentWorktreeDeletionParams = { worktreeId: input.worktreeId, kiloUserId: ctx.userId, + ...(input.requireMergedPullRequests ? { requireMergedPullRequests: true } : {}), ...(input.kilocodeOrganizationId ? { organizationId: input.kilocodeOrganizationId } : {}), }; const startedAt = Date.now(); @@ -201,7 +204,9 @@ export async function deleteWorktreeResources( for (const directory of recordDirectories) { state = cloudAgentWorktreeDeletionStateSchema.parse( await ctx.env.SESSION_INGEST.recordCloudAgentWorktreeCleanup({ - ...params, + worktreeId: params.worktreeId, + kiloUserId: params.kiloUserId, + ...(params.organizationId ? { organizationId: params.organizationId } : {}), directory, sessionIds: cleanup.sessionIds, }) @@ -236,6 +241,17 @@ export async function deleteWorktreeResources( errorCode = error.code; throw error; } + if (error instanceof Error && error.message.includes(WORKTREE_MERGED_PULL_REQUESTS_REQUIRED)) { + outcome = 'rejected'; + errorCode = 'WORKTREE_MERGED_PULL_REQUESTS_REQUIRED'; + retryable = false; + throw new TRPCError({ + code: 'PRECONDITION_FAILED', + message: + 'Worktree cleanup requires a verified merged GitHub pull request and no other linked pull requests', + cause: { error: 'WORKTREE_MERGED_PULL_REQUESTS_REQUIRED', retryable: false }, + }); + } if (error instanceof Error && error.message.includes(WORKTREE_RUNTIME_HISTORY_UNAVAILABLE)) { outcome = 'history_unavailable'; errorCode = 'WORKTREE_RUNTIME_HISTORY_UNAVAILABLE'; diff --git a/services/cloud-agent-next/src/types.ts b/services/cloud-agent-next/src/types.ts index abceae7982a..1f58eea4f87 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -547,6 +547,7 @@ export type Env = { USER_KILO_FACADE: DurableObjectNamespace; /** Durable Object namespace for one-time stream/terminal ticket nonce consumption */ STREAM_TICKET_NONCE_DO: DurableObjectNamespace; + WORKTREE_CLEANUP: Cloudflare.Env['WORKTREE_CLEANUP']; /** e2e-only callback sink; bound only by the rendered e2e configs. */ E2E_CALLBACK_SINK?: DurableObjectNamespace; /** One-way shared sandbox failover overrides keyed by shared identity */ diff --git a/services/cloud-agent-next/src/worktree-cleanup/WorktreeCleanupCoordinator.test.ts b/services/cloud-agent-next/src/worktree-cleanup/WorktreeCleanupCoordinator.test.ts new file mode 100644 index 00000000000..f90c10d1591 --- /dev/null +++ b/services/cloud-agent-next/src/worktree-cleanup/WorktreeCleanupCoordinator.test.ts @@ -0,0 +1,311 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { TRPCError } from '@trpc/server'; +import type { Env } from '../types.js'; +import { StoredWorktreeCleanupItem } from './schemas.js'; +import { WorktreeCleanupCoordinator } from './WorktreeCleanupCoordinator.js'; +import { getWorktreeCleanupStub } from './stub.js'; + +const { deletion } = vi.hoisted(() => ({ deletion: vi.fn() })); +vi.mock('../router/handlers/worktree-deletion.js', () => ({ deleteWorktreeResources: deletion })); +vi.mock('cloudflare:workers', () => ({ + DurableObject: class { + constructor( + protected ctx: DurableObjectState, + protected env: Env + ) {} + }, +})); + +const owner = { userId: 'oauth/google:cleanup-owner' }; +const worktreeId = 'worktree_11111111-1111-4111-8111-111111111111'; +const otherWorktreeId = 'worktree_22222222-2222-4222-8222-222222222222'; +const sessionId = `ses_${'1'.repeat(26)}`; +const childId = `ses_${'2'.repeat(26)}`; + +function fixture() { + const records = new Map(); + let alarm: number | null = null; + let failResultWrite = false; + const storage = { + get: vi.fn(async (key: string) => structuredClone(records.get(key))), + put: vi.fn(async (key: string, value: unknown) => { + if ( + failResultWrite && + StoredWorktreeCleanupItem.safeParse(value).data?.status === 'deleted' + ) { + failResultWrite = false; + throw new Error('isolate terminated before result was persisted'); + } + records.set(key, structuredClone(value)); + }), + delete: vi.fn(async (key: string) => records.delete(key)), + list: vi.fn( + async ({ prefix }: { prefix: string }) => + new Map( + [...records] + .filter(([key]) => key.startsWith(prefix)) + .sort(([a], [b]) => a.localeCompare(b)) + ) + ), + setAlarm: vi.fn(async (time: number) => { + alarm = time; + }), + deleteAlarm: vi.fn(async () => { + alarm = null; + }), + transaction: vi.fn(async (operation: (txn: DurableObjectTransaction) => Promise) => + operation(storage as unknown as DurableObjectTransaction) + ), + }; + const env = { WORKTREE_CLEANUP: {} } as Env; + const create = () => + new WorktreeCleanupCoordinator({ storage } as unknown as DurableObjectState, env); + return { + records, + storage, + create, + coordinator: create(), + alarm: () => alarm, + interruptResultWrite: () => { + failResultWrite = true; + }, + }; +} + +beforeEach(() => { + vi.resetAllMocks(); + deletion.mockResolvedValue({ success: true, deletedSessionIds: [sessionId, childId] }); +}); + +describe('WorktreeCleanupCoordinator', () => { + it('freezes deduplicated IDs, persists before returning, and processes one worktree per alarm', async () => { + const f = fixture(); + expect(await f.coordinator.get(owner)).toBeNull(); + const batch = await f.coordinator.start(owner, { + worktreeIds: [worktreeId, otherWorktreeId, worktreeId], + }); + expect(batch.items.map(item => item.worktreeId)).toEqual([worktreeId, otherWorktreeId]); + expect(f.alarm()).toBeGreaterThan(Date.now()); + expect(deletion).not.toHaveBeenCalled(); + expect(await f.create().get(owner)).toEqual(batch); + expect(await f.coordinator.start(owner, { worktreeIds: [otherWorktreeId] })).toEqual(batch); + + await f.create().alarm(); + expect(deletion).toHaveBeenCalledTimes(1); + expect(deletion).toHaveBeenCalledWith( + { worktreeId, kilocodeOrganizationId: undefined, requireMergedPullRequests: true }, + { env: expect.any(Object), userId: owner.userId } + ); + expect((await f.create().get(owner))?.items).toEqual([ + { worktreeId, status: 'deleted', deletedSessionIds: [sessionId, childId], error: null }, + { worktreeId: otherWorktreeId, status: 'pending', deletedSessionIds: [], error: null }, + ]); + await f.create().alarm(); + expect(deletion).toHaveBeenCalledTimes(2); + expect(f.alarm()).toBeNull(); + const replacement = await f.coordinator.start(owner, { worktreeIds: [otherWorktreeId] }); + expect(replacement.id).not.toBe(batch.id); + expect(replacement.items).toHaveLength(1); + }); + + it('keeps starts deduplicated while deletion is in flight and exposes deleting durably', async () => { + const f = fixture(); + const completed = Promise.withResolvers<{ success: true; deletedSessionIds: string[] }>(); + deletion.mockReturnValue(completed.promise); + const batch = await f.coordinator.start(owner, { worktreeIds: [worktreeId] }); + const running = f.coordinator.alarm(); + await vi.waitFor(() => expect(deletion).toHaveBeenCalledTimes(1)); + expect(f.alarm()).toBeGreaterThan(Date.now() + 20_000); + expect(StoredWorktreeCleanupItem.parse(f.records.get('item:0000'))).toMatchObject({ + status: 'deleting', + attempts: 1, + }); + expect((await f.coordinator.start(owner, { worktreeIds: [otherWorktreeId] })).id).toBe( + batch.id + ); + completed.resolve({ success: true, deletedSessionIds: [sessionId] }); + await running; + expect((await f.coordinator.get(owner))?.items[0].status).toBe('deleted'); + }); + + it('replays an interrupted deletion on a fresh coordinator using its pre-armed alarm', async () => { + const f = fixture(); + await f.coordinator.start(owner, { worktreeIds: [worktreeId] }); + f.interruptResultWrite(); + await expect(f.coordinator.alarm()).rejects.toThrow('isolate terminated'); + expect(f.alarm()).not.toBeNull(); + expect((await f.create().get(owner))?.items[0].status).toBe('deleting'); + await f.create().alarm(); + expect(deletion).toHaveBeenCalledTimes(2); + expect((await f.create().get(owner))?.items[0]).toMatchObject({ + status: 'deleted', + deletedSessionIds: [sessionId, childId], + }); + expect(f.alarm()).toBeNull(); + }); + + it('retries explicit transient failures with backoff and stops after three attempts', async () => { + const f = fixture(); + deletion.mockRejectedValue( + new TRPCError({ + code: 'INTERNAL_SERVER_ERROR', + message: 'Worktree deletion is incomplete', + cause: { retryable: true }, + }) + ); + await f.coordinator.start(owner, { worktreeIds: [worktreeId, otherWorktreeId] }); + await f.create().alarm(); + expect((await f.create().get(owner))?.items[0]).toMatchObject({ + status: 'pending', + error: expect.stringContaining('retrying automatically'), + }); + expect(f.alarm()).toBeGreaterThanOrEqual(Date.now() + 4900); + await f.create().alarm(); + expect(f.alarm()).toBeGreaterThanOrEqual(Date.now() + 9900); + await f.create().alarm(); + expect(deletion).toHaveBeenCalledTimes(3); + expect((await f.create().get(owner))?.items[0]).toMatchObject({ + status: 'failed', + error: expect.stringContaining('stopped after 3 attempts'), + }); + deletion.mockResolvedValue({ success: true, deletedSessionIds: [sessionId] }); + await f.create().alarm(); + expect((await f.create().get(owner))?.items[1].status).toBe('deleted'); + expect(f.alarm()).toBeNull(); + }); + + it('clears the transient error after a successful automatic retry', async () => { + const f = fixture(); + deletion.mockRejectedValueOnce( + new TRPCError({ code: 'SERVICE_UNAVAILABLE', message: 'Cleanup service unavailable' }) + ); + await f.coordinator.start(owner, { worktreeIds: [worktreeId] }); + await f.create().alarm(); + expect((await f.create().get(owner))?.items[0]).toMatchObject({ + status: 'pending', + error: expect.stringContaining('retrying automatically'), + }); + await f.create().alarm(); + expect((await f.create().get(owner))?.items[0]).toMatchObject({ + status: 'deleted', + error: null, + deletedSessionIds: [sessionId, childId], + }); + expect(f.alarm()).toBeNull(); + }); + + it('retries unexpected platform errors only when explicitly transient, within the same bound', async () => { + const f = fixture(); + deletion.mockRejectedValue( + Object.assign(new Error('platform error with secret credential'), { retryable: true }) + ); + await f.coordinator.start(owner, { worktreeIds: [worktreeId] }); + await f.create().alarm(); + expect((await f.create().get(owner))?.items[0].status).toBe('pending'); + await f.create().alarm(); + await f.create().alarm(); + expect(deletion).toHaveBeenCalledTimes(3); + expect((await f.create().get(owner))?.items[0]).toMatchObject({ + status: 'failed', + error: expect.stringContaining('stopped after 3 attempts'), + }); + expect(JSON.stringify([...f.records.values()])).not.toContain('secret credential'); + expect(f.alarm()).toBeNull(); + }); + + it('fails a queued worktree permanently when it becomes ineligible while another is deleting', async () => { + const f = fixture(); + await f.coordinator.start(owner, { worktreeIds: [worktreeId, otherWorktreeId] }); + await f.create().alarm(); + deletion.mockRejectedValue( + new TRPCError({ + code: 'PRECONDITION_FAILED', + message: 'Worktree now includes an open pull request', + cause: { error: 'WORKTREE_MERGED_PULL_REQUESTS_REQUIRED', retryable: false }, + }) + ); + await f.create().alarm(); + expect(deletion).toHaveBeenLastCalledWith( + { + worktreeId: otherWorktreeId, + kilocodeOrganizationId: undefined, + requireMergedPullRequests: true, + }, + expect.any(Object) + ); + expect((await f.create().get(owner))?.items.map(item => item.status)).toEqual([ + 'deleted', + 'failed', + ]); + expect(f.alarm()).toBeNull(); + await f.create().alarm(); + expect(deletion).toHaveBeenCalledTimes(2); + }); + + it.each([ + new TRPCError({ code: 'FORBIDDEN', message: 'Worktree access denied' }), + new TRPCError({ + code: 'INTERNAL_SERVER_ERROR', + message: 'Worktree runtime history is unavailable; recovery is required', + cause: { retryable: false }, + }), + new Error('unexpected error with secret credential'), + ])('does not retry permanent failures: %s', async error => { + const f = fixture(); + deletion.mockRejectedValue(error); + await f.coordinator.start(owner, { worktreeIds: [worktreeId] }); + await f.create().alarm(); + expect((await f.create().get(owner))?.items[0]).toMatchObject({ + status: 'failed', + error: expect.any(String), + }); + expect(JSON.stringify([...f.records.values()])).not.toContain('secret credential'); + expect(f.alarm()).toBeNull(); + await f.create().alarm(); + expect(deletion).toHaveBeenCalledTimes(1); + }); + + it('bounds repeated interruptions and advances to the next worktree', async () => { + const f = fixture(); + await f.coordinator.start(owner, { worktreeIds: [worktreeId, otherWorktreeId] }); + for (let attempt = 0; attempt < 3; attempt++) { + f.interruptResultWrite(); + await expect(f.create().alarm()).rejects.toThrow('isolate terminated'); + } + await f.create().alarm(); + expect(deletion).toHaveBeenCalledTimes(3); + expect((await f.create().get(owner))?.items[0]).toMatchObject({ + status: 'failed', + error: expect.stringContaining('interrupted after 3 attempts'), + }); + await f.create().alarm(); + expect((await f.create().get(owner))?.items[1].status).toBe('deleted'); + expect(f.alarm()).toBeNull(); + }); + + it('isolates user and organization namespaces and rejects mismatched stored ownership', async () => { + const names: string[] = []; + const env = { + WORKTREE_CLEANUP: { + getByName: (name: string) => { + names.push(name); + return {}; + }, + }, + } as unknown as Env; + getWorktreeCleanupStub(env, owner); + getWorktreeCleanupStub(env, { userId: `${owner.userId}:other` }); + getWorktreeCleanupStub(env, { + ...owner, + kilocodeOrganizationId: '22222222-2222-4222-8222-222222222222', + }); + expect(new Set(names).size).toBe(3); + const f = fixture(); + await f.coordinator.start(owner, { worktreeIds: [worktreeId] }); + await expect(f.coordinator.get({ userId: 'other' })).rejects.toThrow('Worktree access denied'); + await expect( + f.coordinator.start({ userId: 'other' }, { worktreeIds: [worktreeId] }) + ).rejects.toThrow('Worktree access denied'); + expect(JSON.stringify([...f.records.values()])).not.toContain('authToken'); + }); +}); diff --git a/services/cloud-agent-next/src/worktree-cleanup/WorktreeCleanupCoordinator.ts b/services/cloud-agent-next/src/worktree-cleanup/WorktreeCleanupCoordinator.ts new file mode 100644 index 00000000000..be44b70db02 --- /dev/null +++ b/services/cloud-agent-next/src/worktree-cleanup/WorktreeCleanupCoordinator.ts @@ -0,0 +1,176 @@ +import { DurableObject } from 'cloudflare:workers'; +import { TRPCError } from '@trpc/server'; +import type { z } from 'zod'; +import { deleteWorktreeResources } from '../router/handlers/worktree-deletion.js'; +import type { Env } from '../types.js'; +import { + StartWorktreeCleanupInput, + StoredWorktreeCleanupBatch, + StoredWorktreeCleanupItem, + WorktreeCleanup, + WorktreeCleanupOwner, +} from './schemas.js'; + +const MAX_ATTEMPTS = 3; +const RECOVERY_DELAY_MS = 30_000; +const ITEM_PREFIX = 'item:'; + +async function readBatch(storage: DurableObjectTransaction) { + const raw = await storage.get('batch'); + if (raw === undefined) return null; + const batch = StoredWorktreeCleanupBatch.parse(raw); + const records = await storage.list({ prefix: ITEM_PREFIX }); + const entries = [...records].map(([key, value]) => ({ + key, + item: StoredWorktreeCleanupItem.parse(value), + })); + return { batch, entries }; +} + +function isActive(item: z.infer): boolean { + return item.status === 'pending' || item.status === 'deleting'; +} + +function cleanupFailure(error: unknown): { retryable: boolean; message: string } { + if (error instanceof TRPCError) { + const cause = error.cause; + return { + retryable: + cause && 'retryable' in cause && typeof cause.retryable === 'boolean' + ? cause.retryable + : ['TIMEOUT', 'SERVICE_UNAVAILABLE', 'TOO_MANY_REQUESTS'].includes(error.code), + message: error.message, + }; + } + return { + retryable: + typeof error === 'object' && + error !== null && + 'retryable' in error && + error.retryable === true, + message: 'Worktree cleanup failed unexpectedly', + }; +} + +export class WorktreeCleanupCoordinator extends DurableObject { + async start(ownerInput: WorktreeCleanupOwner, input: z.infer) { + const owner = WorktreeCleanupOwner.parse(ownerInput); + const { worktreeIds } = StartWorktreeCleanupInput.parse(input); + return this.ctx.storage.transaction(async storage => { + const current = await readBatch(storage); + if (current) { + this.assertOwner(current.batch.owner, owner); + if (current.entries.some(({ item }) => isActive(item))) { + return WorktreeCleanup.parse({ + ...current.batch, + items: current.entries.map(({ item }) => item), + }); + } + for (const { key } of current.entries) await storage.delete(key); + } + const batch = { id: crypto.randomUUID(), createdAt: new Date().toISOString(), owner }; + const items = [...new Set(worktreeIds)].map(worktreeId => ({ + worktreeId, + status: 'pending' as const, + attempts: 0, + deletedSessionIds: [], + error: null, + })); + await storage.put('batch', batch); + for (const [index, item] of items.entries()) { + await storage.put(`${ITEM_PREFIX}${String(index).padStart(4, '0')}`, item); + } + await storage.setAlarm(Date.now() + 1000); + return WorktreeCleanup.parse({ ...batch, items }); + }); + } + + async get(ownerInput: WorktreeCleanupOwner): Promise { + const owner = WorktreeCleanupOwner.parse(ownerInput); + return this.ctx.storage.transaction(async storage => { + const current = await readBatch(storage); + if (!current) return null; + this.assertOwner(current.batch.owner, owner); + return WorktreeCleanup.parse({ + ...current.batch, + items: current.entries.map(({ item }) => item), + }); + }); + } + + async alarm(): Promise { + const next = await this.ctx.storage.transaction(async storage => { + const current = await readBatch(storage); + const entry = current?.entries.find(({ item }) => isActive(item)); + if (!current || !entry) { + await storage.deleteAlarm(); + return null; + } + if (entry.item.attempts >= MAX_ATTEMPTS) { + await storage.put(entry.key, { + ...entry.item, + status: 'failed', + error: `Worktree cleanup was interrupted after ${MAX_ATTEMPTS} attempts; start cleanup again to retry`, + }); + await storage.setAlarm(Date.now() + 1000); + return null; + } + const item = { + ...entry.item, + status: 'deleting' as const, + attempts: entry.item.attempts + 1, + error: null, + }; + await storage.put(entry.key, item); + // This alarm survives an eviction between the external deletion and its result write. + await storage.setAlarm(Date.now() + RECOVERY_DELAY_MS); + return { batch: current.batch, key: entry.key, item }; + }); + if (!next) return; + + let item: z.infer = next.item; + let retryDelayMs = 1000; + try { + const result = await deleteWorktreeResources( + { + worktreeId: item.worktreeId, + kilocodeOrganizationId: next.batch.owner.kilocodeOrganizationId, + requireMergedPullRequests: true, + }, + { env: this.env, userId: next.batch.owner.userId } + ); + item = { ...item, status: 'deleted', deletedSessionIds: result.deletedSessionIds }; + } catch (error) { + const failure = cleanupFailure(error); + const retry = failure.retryable && item.attempts < MAX_ATTEMPTS; + item = { + ...item, + status: retry ? 'pending' : 'failed', + error: retry + ? `${failure.message}; retrying automatically` + : failure.retryable + ? `${failure.message}; cleanup stopped after ${MAX_ATTEMPTS} attempts` + : failure.message, + }; + if (retry) retryDelayMs = 5000 * 2 ** (item.attempts - 1); + } + await this.ctx.storage.transaction(async storage => { + await storage.put(next.key, item); + const current = await readBatch(storage); + if (current?.entries.some(({ item: stored }) => isActive(stored))) { + await storage.setAlarm(Date.now() + retryDelayMs); + } else { + await storage.deleteAlarm(); + } + }); + } + + private assertOwner(stored: WorktreeCleanupOwner, owner: WorktreeCleanupOwner): void { + if ( + stored.userId !== owner.userId || + stored.kilocodeOrganizationId !== owner.kilocodeOrganizationId + ) { + throw new TRPCError({ code: 'FORBIDDEN', message: 'Worktree access denied' }); + } + } +} diff --git a/services/cloud-agent-next/src/worktree-cleanup/schemas.ts b/services/cloud-agent-next/src/worktree-cleanup/schemas.ts new file mode 100644 index 00000000000..7da275b38cc --- /dev/null +++ b/services/cloud-agent-next/src/worktree-cleanup/schemas.ts @@ -0,0 +1,38 @@ +import { cloudAgentWorktreeIdSchema, sessionIdSchema } from '@kilocode/session-ingest-contracts'; +import { z } from 'zod'; + +export const WorktreeCleanupScope = z + .object({ kilocodeOrganizationId: z.uuid().optional() }) + .strict(); + +export const StartWorktreeCleanupInput = WorktreeCleanupScope.extend({ + worktreeIds: z.array(cloudAgentWorktreeIdSchema).min(1).max(1000), +}); + +export const WorktreeCleanupItem = z.object({ + worktreeId: cloudAgentWorktreeIdSchema, + status: z.enum(['pending', 'deleting', 'deleted', 'failed']), + deletedSessionIds: z.array(sessionIdSchema), + error: z.string().nullable(), +}); + +export const WorktreeCleanup = z.object({ + id: z.string(), + items: z.array(WorktreeCleanupItem), + createdAt: z.iso.datetime(), +}); + +export type WorktreeCleanup = z.infer; + +export const WorktreeCleanupOwner = WorktreeCleanupScope.extend({ userId: z.string().min(1) }); +export type WorktreeCleanupOwner = z.infer; + +export const StoredWorktreeCleanupItem = WorktreeCleanupItem.extend({ + attempts: z.number().int().nonnegative(), +}); + +export const StoredWorktreeCleanupBatch = z.object({ + id: z.string(), + createdAt: z.iso.datetime(), + owner: WorktreeCleanupOwner, +}); diff --git a/services/cloud-agent-next/src/worktree-cleanup/stub.ts b/services/cloud-agent-next/src/worktree-cleanup/stub.ts new file mode 100644 index 00000000000..95d0a036dd8 --- /dev/null +++ b/services/cloud-agent-next/src/worktree-cleanup/stub.ts @@ -0,0 +1,11 @@ +import type { Env } from '../types.js'; +import type { WorktreeCleanupOwner } from './schemas.js'; + +export function getWorktreeCleanupStub( + env: Pick, + owner: WorktreeCleanupOwner +) { + return env.WORKTREE_CLEANUP.getByName( + JSON.stringify([owner.userId, owner.kilocodeOrganizationId ?? null]) + ); +} diff --git a/services/cloud-agent-next/test/integration/worktree-cleanup.test.ts b/services/cloud-agent-next/test/integration/worktree-cleanup.test.ts new file mode 100644 index 00000000000..573b1d1818c --- /dev/null +++ b/services/cloud-agent-next/test/integration/worktree-cleanup.test.ts @@ -0,0 +1,78 @@ +import { env, reset, runDurableObjectAlarm, runInDurableObject } from 'cloudflare:test'; +import { afterEach, describe, expect, it } from 'vitest'; +import { getWorktreeCleanupStub } from '../../src/worktree-cleanup/stub.js'; +import { StoredWorktreeCleanupItem } from '../../src/worktree-cleanup/schemas.js'; + +const owner = { userId: 'oauth/google:background-cleanup' }; +const worktreeId = 'worktree_11111111-1111-4111-8111-111111111111'; +const otherWorktreeId = 'worktree_22222222-2222-4222-8222-222222222222'; + +afterEach(() => reset()); + +describe('worktree cleanup durable storage', () => { + it('atomically accepts one frozen batch for concurrent starts and isolates tenants', async () => { + const stub = getWorktreeCleanupStub(env, owner); + const batches = await Promise.all([ + stub.start(owner, { worktreeIds: [worktreeId] }), + stub.start(owner, { worktreeIds: [otherWorktreeId] }), + ]); + expect(batches[0]).toEqual(batches[1]); + expect(await getWorktreeCleanupStub(env, owner).get(owner)).toEqual(batches[0]); + await runInDurableObject(stub, async (_instance, state) => { + expect(await state.storage.getAlarm()).not.toBeNull(); + const batch = await state.storage.get('batch'); + expect(batch).toMatchObject({ owner }); + expect(JSON.stringify(batch)).not.toContain('authToken'); + await state.storage.deleteAlarm(); + }); + const otherUser = { userId: 'other-user' }; + const organization = { + ...owner, + kilocodeOrganizationId: '22222222-2222-4222-8222-222222222222', + }; + expect(await getWorktreeCleanupStub(env, otherUser).get(otherUser)).toBeNull(); + expect(await getWorktreeCleanupStub(env, organization).get(organization)).toBeNull(); + }); + + it('stores the maximum 1000-item batch as separate durable records', async () => { + const stub = getWorktreeCleanupStub(env, owner); + const worktreeIds = Array.from( + { length: 1000 }, + (_, index): `worktree_${string}` => + `worktree_${String(index).padStart(8, '0')}-1111-4111-8111-111111111111` + ); + const batch = await stub.start(owner, { worktreeIds }); + expect(batch.items.map(item => item.worktreeId)).toEqual(worktreeIds); + await runInDurableObject(stub, async (_instance, state) => { + expect((await state.storage.list({ prefix: 'item:' })).size).toBe(1000); + await state.storage.deleteAlarm(); + }); + expect(await stub.get(owner)).toEqual(batch); + }); + + it('recovers a deleting item and bounds retries without another browser request', async () => { + const stub = getWorktreeCleanupStub(env, owner); + await stub.start(owner, { worktreeIds: [worktreeId, otherWorktreeId] }); + await runInDurableObject(stub, async (_instance, state) => { + const item = StoredWorktreeCleanupItem.parse(await state.storage.get('item:0000')); + await state.storage.put('item:0000', { ...item, status: 'deleting', attempts: 1 }); + }); + expect(await runDurableObjectAlarm(stub)).toBe(true); + await runInDurableObject(stub, async (_instance, state) => { + expect(StoredWorktreeCleanupItem.parse(await state.storage.get('item:0000'))).toMatchObject({ + attempts: 2, + status: 'pending', + }); + expect(await state.storage.getAlarm()).not.toBeNull(); + }); + expect(await runDurableObjectAlarm(stub)).toBe(true); + expect((await stub.get(owner))?.items[0]).toMatchObject({ + status: 'failed', + error: expect.stringContaining('stopped after 3 attempts'), + }); + expect((await stub.get(owner))?.items[1].status).toBe('pending'); + await runInDurableObject(stub, async (_instance, state) => { + await state.storage.deleteAlarm(); + }); + }); +}); diff --git a/services/cloud-agent-next/test/test-worker.ts b/services/cloud-agent-next/test/test-worker.ts index e498ce005df..abae17a3e4e 100644 --- a/services/cloud-agent-next/test/test-worker.ts +++ b/services/cloud-agent-next/test/test-worker.ts @@ -72,6 +72,7 @@ export class CloudAgentSession extends RealCloudAgentSession { } export { UserKiloFacade } from '../src/kilo-facade/user-kilo-facade.js'; +export { WorktreeCleanupCoordinator } from '../src/worktree-cleanup/WorktreeCleanupCoordinator.js'; // The test Worker binds the production `SANDBOX_CONTROL`/`SANDBOX_SESSION` // names to the control-plane V2 classes, matching production after the C1c flip. export { SandboxControlV2, SandboxSessionV2 }; diff --git a/services/cloud-agent-next/test/unit/render-e2e-worker-config.test.ts b/services/cloud-agent-next/test/unit/render-e2e-worker-config.test.ts index 70d6ff8246a..26e29b04875 100644 --- a/services/cloud-agent-next/test/unit/render-e2e-worker-config.test.ts +++ b/services/cloud-agent-next/test/unit/render-e2e-worker-config.test.ts @@ -11,6 +11,7 @@ import { tmpdir } from 'node:os'; import { join, resolve } from 'node:path'; import { parse } from 'jsonc-parser'; +import { createSourceFile, isExportDeclaration, isNamedExports, ScriptTarget } from 'typescript'; import { afterEach, describe, expect, it } from 'vitest'; import { @@ -47,6 +48,7 @@ const KEPT_DO_CLASSES = [ 'UserKiloFacade', 'SandboxControl', 'SandboxSession', + 'WorktreeCleanupCoordinator', ]; function readSourceConfig() { @@ -114,12 +116,34 @@ describe('buildE2eWorkerConfig', () => { const migratedClasses = config.migrations.flatMap( (migration: { new_sqlite_classes: string[] }) => migration.new_sqlite_classes ); + expect([...bindingClasses].sort()).toEqual( + ['Sandbox', 'E2eCallbackSink', ...KEPT_DO_CLASSES].sort() + ); for (const className of KEPT_DO_CLASSES) { expect(bindingClasses).toContain(className); expect(migratedClasses).toContain(className); } }); + it('exports every rendered Durable Object binding from the e2e entrypoint', () => { + const { config } = render(); + const entrypoint = createSourceFile( + 'e2e-entry.ts', + readFileSync(join(process.cwd(), 'src/e2e-entry.ts'), 'utf8'), + ScriptTarget.Latest + ); + const exportedNames = entrypoint.statements.flatMap(statement => + isExportDeclaration(statement) && + statement.exportClause && + isNamedExports(statement.exportClause) + ? statement.exportClause.elements.map(element => element.name.text) + : [] + ); + for (const binding of config.durable_objects.bindings) { + expect(exportedNames).toContain(binding.class_name); + } + }); + it('keeps each surviving production class on its original production migration tag', () => { const { source, config } = render(); diff --git a/services/cloud-agent-next/worker-configuration.d.ts b/services/cloud-agent-next/worker-configuration.d.ts index 066668c69c4..f8d806ab4f1 100644 --- a/services/cloud-agent-next/worker-configuration.d.ts +++ b/services/cloud-agent-next/worker-configuration.d.ts @@ -1,5 +1,5 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: ee7a2a6eac8fdf8d011ae1608ca24507) +// Generated by Wrangler by running `wrangler types` (hash: a5002d85da7f17091d5b7911fab63743) // Runtime types generated with workerd@1.20260918.1 2026-06-03 nodejs_compat interface __BaseEnv_Env { SHARED_SANDBOX_OVERRIDES: KVNamespace; @@ -52,6 +52,7 @@ interface __BaseEnv_Env { SandboxCodeReviewContainment: DurableObjectNamespace; CLOUD_AGENT_SESSION: DurableObjectNamespace; STREAM_TICKET_NONCE_DO: DurableObjectNamespace; + WORKTREE_CLEANUP: DurableObjectNamespace; USER_KILO_FACADE: DurableObjectNamespace; SANDBOX_CONTROL: DurableObjectNamespace; SANDBOX_SESSION: DurableObjectNamespace; @@ -69,7 +70,7 @@ interface __BaseEnv_Env { declare namespace Cloudflare { interface GlobalProps { mainModule: typeof import("./src/index"); - durableNamespaces: "Sandbox" | "CloudAgentSession" | "SandboxSmall" | "SandboxDIND" | "UserKiloFacade" | "SandboxCodeReview" | "SandboxContainment" | "SandboxSmallContainment" | "SandboxCodeReviewContainment" | "StreamTicketNonceDO" | "SandboxControl" | "SandboxSession" | "SandboxContainers"; + durableNamespaces: "Sandbox" | "CloudAgentSession" | "SandboxSmall" | "SandboxDIND" | "UserKiloFacade" | "SandboxCodeReview" | "SandboxContainment" | "SandboxSmallContainment" | "SandboxCodeReviewContainment" | "StreamTicketNonceDO" | "SandboxControl" | "SandboxSession" | "SandboxContainers" | "WorktreeCleanupCoordinator"; } interface DevEnv { SHARED_SANDBOX_OVERRIDES: KVNamespace; @@ -122,6 +123,7 @@ declare namespace Cloudflare { SandboxCodeReviewContainment: DurableObjectNamespace; CLOUD_AGENT_SESSION: DurableObjectNamespace; STREAM_TICKET_NONCE_DO: DurableObjectNamespace; + WORKTREE_CLEANUP: DurableObjectNamespace; USER_KILO_FACADE: DurableObjectNamespace; SANDBOX_CONTROL: DurableObjectNamespace; SANDBOX_SESSION: DurableObjectNamespace; diff --git a/services/cloud-agent-next/wrangler.jsonc b/services/cloud-agent-next/wrangler.jsonc index ec9e691c512..74baba78395 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -333,6 +333,10 @@ "class_name": "StreamTicketNonceDO", "name": "STREAM_TICKET_NONCE_DO", }, + { + "class_name": "WorktreeCleanupCoordinator", + "name": "WORKTREE_CLEANUP", + }, { "class_name": "UserKiloFacade", "name": "USER_KILO_FACADE", @@ -400,6 +404,10 @@ "new_sqlite_classes": ["SandboxContainers"], "tag": "v11", }, + { + "new_sqlite_classes": ["WorktreeCleanupCoordinator"], + "tag": "v12", + }, ], /** * Hyperdrive Bindings @@ -683,6 +691,10 @@ "class_name": "StreamTicketNonceDO", "name": "STREAM_TICKET_NONCE_DO", }, + { + "class_name": "WorktreeCleanupCoordinator", + "name": "WORKTREE_CLEANUP", + }, { "class_name": "UserKiloFacade", "name": "USER_KILO_FACADE", diff --git a/services/cloud-agent-next/wrangler.test.jsonc b/services/cloud-agent-next/wrangler.test.jsonc index 11852b97aa3..d68e1f3b4de 100644 --- a/services/cloud-agent-next/wrangler.test.jsonc +++ b/services/cloud-agent-next/wrangler.test.jsonc @@ -20,6 +20,10 @@ "class_name": "StreamTicketNonceDO", "name": "STREAM_TICKET_NONCE_DO", }, + { + "class_name": "WorktreeCleanupCoordinator", + "name": "WORKTREE_CLEANUP", + }, { "class_name": "UserKiloFacade", "name": "USER_KILO_FACADE", @@ -63,6 +67,10 @@ "new_classes": ["PreviewPortTestDO"], "tag": "v11", }, + { + "new_sqlite_classes": ["WorktreeCleanupCoordinator"], + "tag": "v12", + }, ], "kv_namespaces": [{ "binding": "SHARED_SANDBOX_OVERRIDES" }], "r2_buckets": [ diff --git a/services/kilo-mcp/catalog.json b/services/kilo-mcp/catalog.json index 6acb5f051bb..99eebebd260 100644 --- a/services/kilo-mcp/catalog.json +++ b/services/kilo-mcp/catalog.json @@ -5286,6 +5286,38 @@ ], "searchBlob": "cliSessionsV2.getWithRuntimeState Get a session record together with its current cloud agent runtime state and any associated GitHub pull request metadata for the session's branch. clisessionsv2 getwithruntimestate session_id session_id" }, + "cliSessionsV2.getWorktreeCleanup": { + "path": "cliSessionsV2.getWorktreeCleanup", + "kind": "query", + "summary": "Get the status of the current worktree cleanup batch, including each worktree's cleanup result and any failures. Use to check progress or see which merged worktrees failed to clean up.", + "inputSchema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "organizationId": { + "anyOf": [ + { + "type": "string", + "format": "uuid", + "pattern": "^([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-8][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}|00000000-0000-0000-0000-000000000000|ffffffff-ffff-ffff-ffff-ffffffffffff)$" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "organizationId" + ] + }, + "tags": [ + "clisessionsv2", + "getworktreecleanup", + "organizationid" + ], + "searchBlob": "cliSessionsV2.getWorktreeCleanup Get the status of the current worktree cleanup batch, including each worktree's cleanup result and any failures. Use to check progress or see which merged worktrees failed to clean up. clisessionsv2 getworktreecleanup organizationid organizationId" + }, "cliSessionsV2.list": { "path": "cliSessionsV2.list", "kind": "query", @@ -5401,6 +5433,38 @@ ], "searchBlob": "cliSessionsV2.list List the current user's CLI/agent sessions with pagination, filtering by organization, worktree, git URL, platform, shared/public status, version, or recent updates. clisessionsv2 list cursor limit orderby includechildren sharedonly createdonplatform organizationid worktreeid giturl updatedsince version fetchreviewdecision cursor limit orderBy includeChildren sharedOnly createdOnPlatform organizationId worktreeId gitUrl updatedSince version fetchReviewDecision" }, + "cliSessionsV2.mergedWorktreeCleanupPreview": { + "path": "cliSessionsV2.mergedWorktreeCleanupPreview", + "kind": "query", + "summary": "List which merged worktrees are eligible for git worktree cleanup, so you can preview what would be removed before starting a cleanup. Use to check candidates for cleaning up branches and worktrees after merges.", + "inputSchema": { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": { + "organizationId": { + "anyOf": [ + { + "type": "string", + "format": "uuid", + "pattern": "^([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[1-8][0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}|00000000-0000-0000-0000-000000000000|ffffffff-ffff-ffff-ffff-ffffffffffff)$" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "organizationId" + ] + }, + "tags": [ + "clisessionsv2", + "mergedworktreecleanuppreview", + "organizationid" + ], + "searchBlob": "cliSessionsV2.mergedWorktreeCleanupPreview List which merged worktrees are eligible for git worktree cleanup, so you can preview what would be removed before starting a cleanup. Use to check candidates for cleaning up branches and worktrees after merges. clisessionsv2 mergedworktreecleanuppreview organizationid organizationId" + }, "cliSessionsV2.recentRepositories": { "path": "cliSessionsV2.recentRepositories", "kind": "query", diff --git a/services/session-ingest/src/services/worktree-deletion.ts b/services/session-ingest/src/services/worktree-deletion.ts index ec78ed979f7..47b3537f0a3 100644 --- a/services/session-ingest/src/services/worktree-deletion.ts +++ b/services/session-ingest/src/services/worktree-deletion.ts @@ -4,6 +4,7 @@ import { getWorkerDb, type WorkerDb } from '@kilocode/db/client'; import { cli_sessions_v2, cloud_agent_worktrees, + github_branch_pull_requests, type CloudAgentWorktree, } from '@kilocode/db/schema'; import { @@ -13,6 +14,7 @@ import { containedKiloSessionIdSchema, kiloSdkSessionSnapshotOutcomeSchema, WORKTREE_RUNTIME_HISTORY_UNAVAILABLE, + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED, type CanDestroyCloudAgentWorktreeSandboxResult, type UnresolvedCloudAgentSandboxOwner, type CloudAgentWorktreeDeletionParams, @@ -272,6 +274,73 @@ async function lockRoots(db: WorktreeDb, params: CloudAgentWorktreeDeletionParam .for('update'); } +async function requireMergedPullRequests( + db: WorktreeDb, + params: CloudAgentWorktreeDeletionParams, + sessions: CloudAgentWorktreeDeletionState['manifest']['sessions'] +): Promise { + if (!params.requireMergedPullRequests) return; + if (sessions.length === 0) throw new Error(WORKTREE_MERGED_PULL_REQUESTS_REQUIRED); + const rows = await db + .select({ + organizationId: cli_sessions_v2.organization_id, + platform: cli_sessions_v2.platform, + prUrl: cli_sessions_v2.pr_url, + prNumber: cli_sessions_v2.pr_number, + verifiedAt: cli_sessions_v2.pr_link_verified_at, + prState: github_branch_pull_requests.pr_state, + }) + .from(cli_sessions_v2) + .leftJoin( + github_branch_pull_requests, + and( + eq(github_branch_pull_requests.git_url, cli_sessions_v2.git_url), + eq(github_branch_pull_requests.pr_number, cli_sessions_v2.pr_number), + or( + and( + isNotNull(cli_sessions_v2.organization_id), + eq( + github_branch_pull_requests.owned_by_organization_id, + cli_sessions_v2.organization_id + ) + ), + and( + isNull(cli_sessions_v2.organization_id), + eq(github_branch_pull_requests.owned_by_user_id, cli_sessions_v2.kilo_user_id) + ) + ), + isNotNull(cli_sessions_v2.pr_url), + isNotNull(cli_sessions_v2.pr_link_verified_at) + ) + ) + .where( + and( + eq(cli_sessions_v2.kilo_user_id, params.kiloUserId), + inArray( + cli_sessions_v2.session_id, + sessions.map(session => session.sessionId) + ) + ) + ) + .orderBy(cli_sessions_v2.session_id) + .for('update', { of: cli_sessions_v2 }); + const merged = (row: (typeof rows)[number]) => + row.platform === 'github' && + row.prUrl !== null && + row.verifiedAt !== null && + row.prState === 'merged'; + if ( + !rows.some(merged) || + rows.some( + row => + row.organizationId !== (params.organizationId ?? null) || + ((row.prUrl !== null || row.prNumber !== null) && !merged(row)) + ) + ) { + throw new Error(WORKTREE_MERGED_PULL_REQUESTS_REQUIRED); + } +} + /** * The whole deletion-journal body, run on a caller-owned transaction so a * caller can hold the same worktree row lock (`lockWorktree`) across its own @@ -323,6 +392,7 @@ async function beginWorktreeDeletionOn( sessions: await discoverMembers(tx, params), }); } + await requireMergedPullRequests(tx, params, manifest.sessions); const runtimeLocations = locations(row); if (runtimeLocations.length === 0) { const cloudAgentSessionId = firstWorktreeSessionId(params.worktreeId); @@ -692,6 +762,10 @@ export async function recordWorktreeCleanup( const sessions = new Map( state.manifest.sessions.map(session => [session.sessionId, session]) ); + await requireMergedPullRequests(tx, params, [ + ...sessions.values(), + ...childSessions.map(child => ({ sessionId: child.sessionId, cloudAgentSessionId: null })), + ]); await recoverChildSessionLineage( tx, { ...params, childSessions }, @@ -779,6 +853,7 @@ export async function recordWorktreeCleanup( version: 1, sessions: [...sessions.values()], } satisfies typeof state.manifest; + await requireMergedPullRequests(tx, params, manifest.sessions); await tx .update(cloud_agent_worktrees) .set({ runtime_locations: runtimeLocations, deletion_manifest: manifest }) diff --git a/services/session-ingest/test/postgres/worktree-deletion.test.ts b/services/session-ingest/test/postgres/worktree-deletion.test.ts new file mode 100644 index 00000000000..ceb6afb66f5 --- /dev/null +++ b/services/session-ingest/test/postgres/worktree-deletion.test.ts @@ -0,0 +1,314 @@ +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; +import { createDrizzleClient } from '@kilocode/db/client'; +import { + cli_sessions_v2, + cloud_agent_worktrees, + github_branch_pull_requests, +} from '@kilocode/db/schema'; +import { eq, sql } from 'drizzle-orm'; +import { + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED, + type CloudAgentWorktreeDeletionParams, +} from '@kilocode/session-ingest-contracts'; +import { beginWorktreeDeletion, recordWorktreeCleanup } from '../../src/services/worktree-deletion'; +import type { Env } from '../../src/env'; + +const mocks = vi.hoisted(() => ({ getDb: vi.fn(), snapshot: vi.fn() })); +vi.mock('cloudflare:workers', () => ({ DurableObject: class {} })); +vi.mock(import('@kilocode/db/client'), async importOriginal => ({ + ...(await importOriginal()), + getWorkerDb: mocks.getDb, +})); +vi.mock(import('@kilocode/worker-utils'), async importOriginal => ({ + ...(await importOriginal()), + hasOrganizationAccess: async () => true, +})); +vi.mock('../../src/dos/SessionIngestDO', () => ({ + getSessionIngestDO: () => ({ readKiloSdkSessionSnapshot: mocks.snapshot }), +})); +vi.mock('../../src/dos/SessionAccessCacheDO', () => ({ getSessionAccessCacheDO: vi.fn() })); +vi.mock('../../src/dos/UserConnectionDO', () => ({ getUserConnectionDO: vi.fn() })); + +const userId = 'oauth/github:merged-cleanup'; +const organizationId = '33333333-3333-4333-8333-333333333333'; +const worktreeId = 'worktree_11111111-1111-4111-8111-111111111111'; +const cloudAgentSessionId = 'workspace_11111111-1111-4111-8111-111111111111'; +const rootId = `ses_${'0'.repeat(26)}`; +const childId = `ses_${'1'.repeat(26)}`; +const gitUrl = 'https://github.com/example/repository'; +const directory = `/workspace/owner/worktrees/${worktreeId}`; +const params = { + kiloUserId: userId, + worktreeId, + requireMergedPullRequests: true, +} satisfies CloudAgentWorktreeDeletionParams; +const env = { HYPERDRIVE: { connectionString: 'unused' } } as Env; +let client: ReturnType; + +beforeAll(async () => { + const connectionString = process.env.POSTGRES_URL; + if (!connectionString) throw new Error('Set POSTGRES_URL to a local test database'); + if (!['localhost', '127.0.0.1', '[::1]'].includes(new URL(connectionString).hostname)) { + throw new Error('The deletion tests require a local PostgreSQL database'); + } + client = createDrizzleClient({ connectionString, poolConfig: { max: 1 }, ssl: false }); + for (const table of [cli_sessions_v2, cloud_agent_worktrees, github_branch_pull_requests]) { + await client.db.execute(sql`CREATE TEMP TABLE ${table} (LIKE public.${table} INCLUDING ALL)`); + } + mocks.getDb.mockReturnValue(client.db); +}); + +afterAll(async () => { + if (client) await client.pool.end(); +}); + +beforeEach(async () => { + await client.db.execute(sql`TRUNCATE pg_temp.cli_sessions_v2, + pg_temp.cloud_agent_worktrees, pg_temp.github_branch_pull_requests`); + mocks.snapshot.mockReset(); + mocks.snapshot.mockResolvedValue({ kind: 'pending' }); + await client.db.insert(cloud_agent_worktrees).values({ + worktree_id: worktreeId, + kilo_user_id: userId, + runtime_locations: [{ sandboxId: 'sandbox', provider: 'cloudflare' }], + }); + await client.db.insert(cli_sessions_v2).values({ + session_id: rootId, + kilo_user_id: userId, + cloud_agent_session_id: cloudAgentSessionId, + cloud_agent_worktree_id: worktreeId, + git_url: gitUrl, + platform: 'github', + pr_url: `${gitUrl}/pull/1`, + pr_number: 1, + pr_link_verified_at: '2026-10-20 01:00:00+00', + }); + await client.db.insert(github_branch_pull_requests).values({ + git_url: gitUrl, + git_branch: 'old-reused-branch', + owned_by_user_id: userId, + pr_number: 1, + pr_state: 'merged', + }); +}); + +async function worktree() { + const [row] = await client.db.select().from(cloud_agent_worktrees); + return row; +} + +async function addChild(values: Partial = {}) { + await client.db.insert(cli_sessions_v2).values({ + session_id: childId, + kilo_user_id: userId, + git_url: gitUrl, + platform: 'github', + pr_url: `${gitUrl}/pull/2`, + pr_number: 2, + pr_link_verified_at: '2026-10-20 01:00:00+00', + ...values, + }); + await client.db.insert(github_branch_pull_requests).values({ + git_url: gitUrl, + git_branch: 'new-branch', + owned_by_user_id: userId, + pr_number: 2, + pr_state: 'open', + }); +} + +describe('merged-only deletion transactions', () => { + it('rejects a queued worktree that acquired an open-PR root before deletion, without fencing it', async () => { + await addChild({ + cloud_agent_worktree_id: worktreeId, + cloud_agent_session_id: 'workspace_22222222-2222-4222-8222-222222222222', + }); + await expect(beginWorktreeDeletion(env, params)).rejects.toThrow( + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED + ); + expect(await worktree()).toMatchObject({ deletion_started_at: null, deletion_manifest: null }); + }); + + it.each(['parent', 'scope'])('rejects an open-PR descendant linked only by %s', async linkage => { + await addChild( + linkage === 'parent' + ? { parent_session_id: rootId } + : { cloud_agent_session_scope_id: cloudAgentSessionId } + ); + await expect(beginWorktreeDeletion(env, params)).rejects.toThrow( + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED + ); + expect((await worktree()).deletion_started_at).toBeNull(); + }); + + it('revalidates persisted members on retry even when their parent link is no longer discoverable', async () => { + await addChild({ parent_session_id: rootId, pr_url: null, pr_number: null }); + await beginWorktreeDeletion(env, params); + const journal = await worktree(); + await client.db + .update(cli_sessions_v2) + .set({ parent_session_id: null, pr_url: `${gitUrl}/pull/2`, pr_number: 2 }) + .where(eq(cli_sessions_v2.session_id, childId)); + await expect(beginWorktreeDeletion(env, params)).rejects.toThrow( + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED + ); + expect(await worktree()).toEqual(journal); + }); + + it.each(['scope', 'snapshot', 'retained'])( + 'rejects newly discovered %s members before runtime cleanup and preserves recovery state', + async linkage => { + await beginWorktreeDeletion(env, params); + const journal = await worktree(); + await addChild( + linkage === 'scope' ? { cloud_agent_session_scope_id: cloudAgentSessionId } : {} + ); + if (linkage === 'snapshot') { + mocks.snapshot.mockResolvedValue({ + kind: 'value', + info: { + id: childId, + parentID: rootId, + directory, + slug: 'child', + projectID: 'project', + title: 'Orphan', + version: '7.4.20', + time: { created: 1, updated: 1 }, + }, + byteLength: 100, + }); + } + await expect( + recordWorktreeCleanup(env, { + ...params, + directory, + ...(linkage === 'retained' + ? { + childSessions: [ + { sessionId: childId, parentSessionId: rootId, cloudAgentSessionId }, + ], + } + : {}), + }) + ).rejects.toThrow(WORKTREE_MERGED_PULL_REQUESTS_REQUIRED); + expect(await worktree()).toEqual(journal); + const [child] = await client.db + .select() + .from(cli_sessions_v2) + .where(eq(cli_sessions_v2.session_id, childId)); + expect(child.cloud_agent_worktree_id).toBeNull(); + expect(child.parent_session_id).toBeNull(); + } + ); + + it.each([ + 'open', + 'closed', + 'unknown', + 'unverified', + 'nonGithub', + 'number-only', + 'url-only', + 'no-linked-pr', + ])('fails closed for %s links', async problem => { + if (['open', 'closed', 'unknown'].includes(problem)) { + await client.db.update(github_branch_pull_requests).set({ pr_state: problem }); + } else { + await client.db.update(cli_sessions_v2).set({ + ...(problem === 'unverified' ? { pr_link_verified_at: null } : {}), + ...(problem === 'nonGithub' ? { platform: 'gitlab' } : {}), + ...(problem === 'number-only' ? { pr_url: null } : {}), + ...(problem === 'url-only' ? { pr_number: null } : {}), + ...(problem === 'no-linked-pr' ? { pr_url: null, pr_number: null } : {}), + }); + } + await expect(beginWorktreeDeletion(env, params)).rejects.toThrow( + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED + ); + expect((await worktree()).deletion_started_at).toBeNull(); + }); + + it.each(['user', 'organization', 'git-url', 'pr-number'])( + 'does not borrow a merged cache entry with mismatched %s identity', + async identity => { + await client.db.update(github_branch_pull_requests).set({ + ...(identity === 'user' ? { owned_by_user_id: 'another-user' } : {}), + ...(identity === 'organization' + ? { owned_by_user_id: null, owned_by_organization_id: organizationId } + : {}), + ...(identity === 'git-url' ? { git_url: 'https://github.com/other/repo' } : {}), + ...(identity === 'pr-number' ? { pr_number: 2 } : {}), + }); + await expect(beginWorktreeDeletion(env, params)).rejects.toThrow( + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED + ); + } + ); + + it('uses the organization cache rather than personal state for organization sessions', async () => { + await client.db.update(cli_sessions_v2).set({ organization_id: organizationId }); + await client.db.update(cloud_agent_worktrees).set({ organization_id: organizationId }); + const scoped = { ...params, organizationId }; + await expect(beginWorktreeDeletion(env, scoped)).rejects.toThrow( + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED + ); + await client.db + .update(github_branch_pull_requests) + .set({ owned_by_user_id: null, owned_by_organization_id: organizationId }); + await expect(beginWorktreeDeletion(env, scoped)).resolves.toMatchObject({ completed: false }); + }); + + it('checks recovered tenant identity before lineage repair could move a personal PR into an organization', async () => { + await client.db.update(cli_sessions_v2).set({ organization_id: organizationId }); + await client.db.update(cloud_agent_worktrees).set({ organization_id: organizationId }); + await client.db.update(github_branch_pull_requests).set({ + owned_by_user_id: null, + owned_by_organization_id: organizationId, + }); + const scoped = { ...params, organizationId }; + await beginWorktreeDeletion(env, scoped); + await addChild(); + await client.db.insert(github_branch_pull_requests).values({ + git_url: gitUrl, + git_branch: 'organization-branch', + owned_by_organization_id: organizationId, + pr_number: 2, + pr_state: 'merged', + }); + await expect( + recordWorktreeCleanup(env, { + ...scoped, + childSessions: [{ sessionId: childId, parentSessionId: rootId, cloudAgentSessionId }], + }) + ).rejects.toThrow(WORKTREE_MERGED_PULL_REQUESTS_REQUIRED); + const [child] = await client.db + .select() + .from(cli_sessions_v2) + .where(eq(cli_sessions_v2.session_id, childId)); + expect(child.organization_id).toBeNull(); + expect(child.cloud_agent_worktree_id).toBeNull(); + }); + + it('allows verified merged links independent of branch names and includes unlinked descendants', async () => { + await addChild({ + cloud_agent_session_scope_id: cloudAgentSessionId, + pr_url: null, + pr_number: null, + }); + await client.db.update(cli_sessions_v2).set({ git_branch: 'different-new-branch' }); + const state = await beginWorktreeDeletion(env, params); + expect(state.manifest.sessions.map(session => session.sessionId)).toEqual([rootId, childId]); + await expect(recordWorktreeCleanup(env, params)).resolves.toMatchObject({ completed: false }); + }); + + it('keeps individual deletion unrestricted even after a rejected merged-only attempt', async () => { + await client.db.update(github_branch_pull_requests).set({ pr_state: 'open' }); + await expect(beginWorktreeDeletion(env, params)).rejects.toThrow( + WORKTREE_MERGED_PULL_REQUESTS_REQUIRED + ); + await expect( + beginWorktreeDeletion(env, { worktreeId, kiloUserId: userId }) + ).resolves.toMatchObject({ completed: false }); + }); +});