From dd9498cf1a68aedde276de55119de4ead339a83c Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Tue, 1 Sep 2026 20:32:16 +0200 Subject: [PATCH 1/2] feat(cloud-agent-next): persist and display worktree change summaries --- .../cloud-agent-next/ChatHeader.tsx | 19 +- .../cloud-agent-next/CloudChatPage.tsx | 68 +- .../cloud-agent-next/WorktreeChanges.tsx | 457 ++++++ .../cloud-agent-next/terminal-tabs.test.ts | 102 +- .../cloud-agent-next/terminal-tabs.ts | 12 +- .../cloud-agent-next/worktree-changes.test.ts | 680 +++++++++ .../cloud-agent-next/worktree-changes.ts | 179 +++ apps/web/src/components/ui/sheet.tsx | 12 +- .../cloud-agent-client.test.ts | 107 +- .../cloud-agent-next/cloud-agent-client.ts | 24 + .../routers/cloud-agent-next-router.test.ts | 105 +- .../src/routers/cloud-agent-next-router.ts | 23 + .../routers/cloud-agent-next-schemas.test.ts | 28 + .../src/routers/cloud-agent-next-schemas.ts | 11 + ...ganization-cloud-agent-next-router.test.ts | 197 ++- .../organization-cloud-agent-next-router.ts | 35 + .../e2e/cloud-agent-sandbox-status.spec.ts | 172 +++ packages/worker-utils/package.json | 1 + .../src/cloud-agent-worktree-changes.test.ts | 149 ++ .../src/cloud-agent-worktree-changes.ts | 97 ++ .../src/persistence/SandboxControl.ts | 72 +- services/cloud-agent-next/src/router.ts | 2 + .../handlers/session-worktree-changes.test.ts | 161 +++ .../handlers/session-worktree-changes.ts | 63 + .../cloud-agent-next/src/router/schemas.ts | 7 + .../src/sandbox-control/frames.test.ts | 28 + .../src/sandbox-control/frames.ts | 2 + .../src/sandbox-control/socket.test.ts | 229 ++- .../src/sandbox-control/socket.ts | 55 +- .../src/sandbox-session/SandboxSession.ts | 89 ++ .../session-message-queue.test.ts | 93 ++ .../sandbox-session/worktree-changes.test.ts | 638 +++++++++ .../src/sandbox-session/worktree-changes.ts | 248 ++++ .../src/shared/sandbox-control-protocol.ts | 13 + .../src/shared/worktree-changes-wire.test.ts | 244 ++++ .../src/shared/worktree-changes-wire.ts | 70 + .../test/integration/sandbox-control.test.ts | 1247 ++++++++++++++++- .../test/unit/wrapper/utils.test.ts | 44 + .../wrapper/src/control/main.ts | 9 + .../control/sandbox-control-client.test.ts | 274 +++- .../src/control/sandbox-control-client.ts | 57 +- .../control/sandbox-control-handlers.test.ts | 298 ++++ .../src/control/sandbox-control-handlers.ts | 48 +- .../src/control/session-directories.ts | 10 + .../src/control/standalone-build.test.ts | 50 + .../src/control/worktree-changes.test.ts | 775 ++++++++++ .../wrapper/src/control/worktree-changes.ts | 380 +++++ .../worktree-mutation-notifications.test.ts | 912 ++++++++++++ .../worktree-mutation-notifications.ts | 293 ++++ .../cloud-agent-next/wrapper/src/utils.ts | 16 +- 50 files changed, 8765 insertions(+), 140 deletions(-) create mode 100644 apps/web/src/components/cloud-agent-next/WorktreeChanges.tsx create mode 100644 apps/web/src/components/cloud-agent-next/worktree-changes.test.ts create mode 100644 apps/web/src/components/cloud-agent-next/worktree-changes.ts create mode 100644 packages/worker-utils/src/cloud-agent-worktree-changes.test.ts create mode 100644 packages/worker-utils/src/cloud-agent-worktree-changes.ts create mode 100644 services/cloud-agent-next/src/router/handlers/session-worktree-changes.test.ts create mode 100644 services/cloud-agent-next/src/router/handlers/session-worktree-changes.ts create mode 100644 services/cloud-agent-next/src/sandbox-session/worktree-changes.test.ts create mode 100644 services/cloud-agent-next/src/sandbox-session/worktree-changes.ts create mode 100644 services/cloud-agent-next/src/shared/worktree-changes-wire.test.ts create mode 100644 services/cloud-agent-next/src/shared/worktree-changes-wire.ts create mode 100644 services/cloud-agent-next/wrapper/src/control/standalone-build.test.ts create mode 100644 services/cloud-agent-next/wrapper/src/control/worktree-changes.test.ts create mode 100644 services/cloud-agent-next/wrapper/src/control/worktree-changes.ts create mode 100644 services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts create mode 100644 services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.ts diff --git a/apps/web/src/components/cloud-agent-next/ChatHeader.tsx b/apps/web/src/components/cloud-agent-next/ChatHeader.tsx index 5e02772a67..fd7e92f426 100644 --- a/apps/web/src/components/cloud-agent-next/ChatHeader.tsx +++ b/apps/web/src/components/cloud-agent-next/ChatHeader.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useEffect, useMemo, useRef, useState, type RefObject } from 'react'; +import { useEffect, useMemo, useRef, useState, type MouseEvent, type RefObject } from 'react'; import { useQuery } from '@tanstack/react-query'; import { Button } from '@/components/ui/button'; import { @@ -16,6 +16,7 @@ import type { SessionCostBreakdown } from './session-cost-breakdown'; import { SessionActionsDialog } from './SessionActionsDialog'; import { SoundToggleButton } from '@/components/shared/SoundToggleButton'; import { FeedbackDialog } from './FeedbackDialog'; +import { WorktreeChangesButton } from './WorktreeChanges'; import { buildRepoBrowseUrl, detectGitPlatform } from './utils/git-utils'; import { useTRPC } from '@/lib/trpc/utils'; import { SandboxStatusIndicator } from './SandboxStatusIndicator'; @@ -44,6 +45,8 @@ type ChatHeaderProps = { sessionInfoTriggerRef: RefObject; soundEnabled?: boolean; onToggleSound?: () => void; + changesOpen?: boolean; + onToggleChanges?: (event: MouseEvent) => void; sessionTitle?: string; sessionActive: boolean; sandboxStatusEligible?: boolean; @@ -62,6 +65,8 @@ export function ChatHeader({ sessionInfoTriggerRef, soundEnabled = true, onToggleSound, + changesOpen = false, + onToggleChanges, kiloSessionId, organizationId, sessionTitle, @@ -129,12 +134,22 @@ export function ChatHeader({
{sandboxStatusEligible && ( )} + {onToggleChanges && ( + + )} {onToggleSound && ( )} diff --git a/apps/web/src/components/cloud-agent-next/CloudChatPage.tsx b/apps/web/src/components/cloud-agent-next/CloudChatPage.tsx index 763ccba074..376acfb81b 100644 --- a/apps/web/src/components/cloud-agent-next/CloudChatPage.tsx +++ b/apps/web/src/components/cloud-agent-next/CloudChatPage.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; +import { useCallback, useEffect, useMemo, useRef, useState, type MouseEvent } from 'react'; import { useAtomValue, useSetAtom } from 'jotai'; import { useSearchParams } from 'next/navigation'; import { useMutation, useQueryClient } from '@tanstack/react-query'; @@ -53,6 +53,8 @@ import { import { billingPayerPresentation } from './billing-payer-presentation'; import type { OrganizationRole } from '@/lib/organizations/organization-types'; import { CloudAgentWorkspaceTabs } from './CloudAgentWorkspaceTabs'; +import { WorktreeChangesDrawer } from './WorktreeChanges'; +import { canOpenWorktreeChanges } from './worktree-changes'; import { CHAT_TAB_ID, addTerminalTab, @@ -163,6 +165,8 @@ export default function CloudChatPage({ const childSessionDrawerFocusTargetRef = useRef(null); const [preparationDrawerAttemptId, setPreparationDrawerAttemptId] = useState(null); const preparationDrawerFocusTargetRef = useRef(null); + const [changesDrawerSessionId, setChangesDrawerSessionId] = useState(null); + const changesDrawerFocusTargetRef = useRef(null); // URL-driven session switching const sessionIdFromParams = searchParams?.get('sessionId') ?? null; @@ -171,6 +175,8 @@ export default function CloudChatPage({ setChildSessionStack([]); preparationDrawerFocusTargetRef.current = null; setPreparationDrawerAttemptId(null); + changesDrawerFocusTargetRef.current = null; + setChangesDrawerSessionId(null); if (sessionIdFromParams) { void manager.switchSession(sessionIdFromParams as KiloSessionId); } else { @@ -255,6 +261,12 @@ export default function CloudChatPage({ organizationId, scope: workspaceTabScope, }); + const canOpenChanges = + sessionIdFromParams !== null && + isCurrentSession && + canOpenWorktreeChanges(sessionId, isReadOnly) && + fetchedSessionData?.organizationId === (organizationId ?? null); + const changesDrawerOpen = canOpenChanges && changesDrawerSessionId === sessionId; if ( resolvedWorkspaceScope.currentUserId !== currentUserId || @@ -280,6 +292,11 @@ export default function CloudChatPage({ } }, [sessionIdFromParams]); + useEffect(() => { + changesDrawerFocusTargetRef.current = null; + setChangesDrawerSessionId(null); + }, [sessionId, currentUserId, organizationId, canOpenChanges]); + // -- Session models ------------------------------------------------------- const sessionModels = useSessionModels({ activeSessionType, @@ -693,9 +710,10 @@ export default function CloudChatPage({ const activeElement = document.activeElement; childSessionDrawerFocusTargetRef.current = activeElement instanceof HTMLElement ? activeElement : null; - // The two drawers overlay the same chat pane — only one may be open. preparationDrawerFocusTargetRef.current = null; setPreparationDrawerAttemptId(null); + changesDrawerFocusTargetRef.current = null; + setChangesDrawerSessionId(null); setChildSessionStack([entry]); }, []); @@ -723,9 +741,10 @@ export default function CloudChatPage({ const activeElement = document.activeElement; preparationDrawerFocusTargetRef.current = activeElement instanceof HTMLElement ? activeElement : null; - // The two drawers overlay the same chat pane — only one may be open. childSessionDrawerFocusTargetRef.current = null; setChildSessionStack([]); + changesDrawerFocusTargetRef.current = null; + setChangesDrawerSessionId(null); setPreparationDrawerAttemptId(attemptId); }, []); @@ -741,6 +760,30 @@ export default function CloudChatPage({ focusTarget.focus(); }, []); + const handleToggleChanges = useCallback( + (event: MouseEvent) => { + changesDrawerFocusTargetRef.current = event.currentTarget; + childSessionDrawerFocusTargetRef.current = null; + setChildSessionStack([]); + preparationDrawerFocusTargetRef.current = null; + setPreparationDrawerAttemptId(null); + setChangesDrawerSessionId(current => (current === sessionId ? null : sessionId)); + }, + [sessionId] + ); + + const handleChangesDrawerOpenChange = useCallback((open: boolean) => { + if (!open) setChangesDrawerSessionId(null); + }, []); + + const handleChangesDrawerCloseAutoFocus = useCallback((event: Event) => { + const focusTarget = changesDrawerFocusTargetRef.current; + changesDrawerFocusTargetRef.current = null; + if (!focusTarget?.isConnected) return; + event.preventDefault(); + focusTarget.focus(); + }, []); + // Surface the session's custom agents plus the current visible profile // agents to the chat picker. `runtimeAgents` are the agents active when the // session was created; the profile list enriches those same agents with @@ -950,6 +993,8 @@ export default function CloudChatPage({ sessionInfoTriggerRef={sessionInfoTriggerRef} soundEnabled={soundEnabled} onToggleSound={handleToggleSound} + changesOpen={changesDrawerOpen} + onToggleChanges={canOpenChanges ? handleToggleChanges : undefined} sessionActive={isStreaming || activity.type === 'busy' || activity.type === 'retrying'} sandboxStatusEligible={isSandboxStatusEligible({ currentUserId, @@ -1051,7 +1096,11 @@ export default function CloudChatPage({ className="relative flex min-h-0 flex-1 flex-col" >
0 || preparationDrawerAttemptId !== null} + inert={ + childSessionStack.length > 0 || + preparationDrawerAttemptId !== null || + changesDrawerOpen + } className="flex min-h-0 flex-1 flex-col" >
@@ -1302,6 +1351,17 @@ export default function CloudChatPage({ onCloseAutoFocus={handlePreparationDrawerCloseAutoFocus} portalContainer={childSessionDrawerContainer} /> + {canOpenChanges && sessionId && ( + + )}
diff --git a/apps/web/src/components/cloud-agent-next/WorktreeChanges.tsx b/apps/web/src/components/cloud-agent-next/WorktreeChanges.tsx new file mode 100644 index 0000000000..9c96c7334c --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/WorktreeChanges.tsx @@ -0,0 +1,457 @@ +'use client'; + +import { useEffect, useMemo, useRef, useState, type MouseEvent, type RefObject } from 'react'; +import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; +import { ChevronRight, FileDiff, GitBranch, List, ListTree, RefreshCw } from 'lucide-react'; +import { formatDistanceToNow } from 'date-fns'; +import { Button } from '@/components/ui/button'; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from '@/components/ui/collapsible'; +import { Sheet, SheetContent, SheetDescription, SheetTitle } from '@/components/ui/sheet'; +import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'; +import { useLocalStorage } from '@/hooks/useLocalStorage'; +import { useRawTRPCClient, useTRPC } from '@/lib/trpc/utils'; +import { + buildWorktreeChangesTree, + deserializeWorktreeChangesViewMode, + formatWorktreeChangesBaseBranch, + getWorktreeChangesTotals, + groupWorktreeChangesByDirectory, + preserveNewerWorktreeChanges, + worktreeChangesMessages, + type WorktreeChangesFile, + type WorktreeChangesTreeNode, + type WorktreeChangesViewMode, +} from './worktree-changes'; + +const fileStatusStyles = { + added: { label: 'Added', dot: 'bg-diff-add-text' }, + modified: { label: 'Modified', dot: 'bg-(--status-orange-400)' }, + deleted: { label: 'Deleted', dot: 'bg-diff-delete-text' }, +}; +const compactCountFormatter = new Intl.NumberFormat('en', { + notation: 'compact', + maximumFractionDigits: 0, +}); + +function useSavedWorktreeChanges({ + cloudAgentSessionId, + organizationId, + enabled, + poll = false, + catchUpUntil = 0, +}: { + cloudAgentSessionId: string; + organizationId?: string; + enabled: boolean; + poll?: boolean; + catchUpUntil?: number; +}) { + const trpc = useTRPC(); + const queryOptions = organizationId + ? trpc.organizations.cloudAgentNext.getWorktreeChanges.queryOptions({ + organizationId, + cloudAgentSessionId, + }) + : trpc.cloudAgentNext.getWorktreeChanges.queryOptions({ cloudAgentSessionId }); + const saved = useQuery({ + ...queryOptions, + enabled, + staleTime: 0, + refetchOnMount: 'always', + refetchOnWindowFocus: true, + refetchOnReconnect: true, + refetchInterval: () => (poll || Date.now() < catchUpUntil ? 5_000 : false), + retry: false, + structuralSharing: preserveNewerWorktreeChanges, + }); + return { saved, queryKey: queryOptions.queryKey }; +} + +function ChangeLineCounts({ + additions, + deletions, + countsComplete, + compactOnMobile = false, +}: Pick & { + compactOnMobile?: boolean; +}) { + return ( + + + {compactOnMobile && ( + +{compactCountFormatter.format(additions)} + )} + +{additions} + + + {compactOnMobile && ( + −{compactCountFormatter.format(deletions)} + )} + −{deletions} + + {!countsComplete && ( + + * + + )} + + ); +} + +export function WorktreeChangesButton({ + cloudAgentSessionId, + organizationId, + open, + onToggle, + sessionActive, +}: { + cloudAgentSessionId: string; + organizationId?: string; + open: boolean; + onToggle: (event: MouseEvent) => void; + sessionActive: boolean; +}) { + const [catchUpUntil, setCatchUpUntil] = useState(0); + const { saved } = useSavedWorktreeChanges({ + cloudAgentSessionId, + organizationId, + enabled: true, + poll: sessionActive || open, + catchUpUntil, + }); + const wasSessionActive = useRef(sessionActive); + useEffect(() => { + if (wasSessionActive.current && !sessionActive) { + setCatchUpUntil(Date.now() + 30_000); + void saved.refetch(); + } + wasSessionActive.current = sessionActive; + }, [sessionActive, saved.refetch]); + + const totals = getWorktreeChangesTotals(saved.data?.snapshot); + const summary = totals + ? `${totals.fileCount} changed files, ${totals.additions} additions, ${totals.deletions} deletions${totals.countsComplete ? '' : ' (partial summary)'}` + : 'Changes'; + + return ( + + ); +} + +export function WorktreeChangesDrawer({ + cloudAgentSessionId, + organizationId, + open, + onOpenChange, + onCloseAutoFocus, + portalContainer, +}: { + cloudAgentSessionId: string; + organizationId?: string; + open: boolean; + onOpenChange: (open: boolean) => void; + onCloseAutoFocus: (event: Event) => void; + portalContainer: HTMLElement | null; +}) { + const [viewMode, setViewMode] = useLocalStorage( + 'cloud-agent:worktree-changes-view-mode', + 'flat', + { initializeWithValue: false, deserializer: deserializeWorktreeChangesViewMode } + ); + const activeTabRef = useRef(null); + + useEffect(() => { + if (open) activeTabRef.current?.focus({ preventScroll: true }); + }, [open]); + + return ( + + { + event.preventDefault(); + activeTabRef.current?.focus({ preventScroll: true }); + }} + onCloseAutoFocus={onCloseAutoFocus} + onInteractOutside={event => event.preventDefault()} + > + Changes + + Changed files in the session worktree. + + + + + ); +} + +function ChangedFile({ file }: { file: WorktreeChangesFile }) { + const status = fileStatusStyles[file.status]; + const name = file.path.slice(file.path.lastIndexOf('/') + 1); + + return ( +
  • + + + {name} + + {file.binary ? ( + Binary file; line counts unavailable. + ) : ( + + )} +
    • + ); +} + +function ChangedFileTree({ nodes }: { nodes: WorktreeChangesTreeNode[] }) { + return nodes.map(node => + node.kind === 'file' ? ( + + ) : ( + +
  • + + + + +
      + +
    +     +
    • +     + ) + ); +} + +function WorktreeChanges({ + cloudAgentSessionId, + organizationId, + open, + viewMode, + onViewModeChange, + activeTabRef, +}: { + cloudAgentSessionId: string; + organizationId?: string; + open: boolean; + viewMode: WorktreeChangesViewMode; + onViewModeChange: (value: WorktreeChangesViewMode) => void; + activeTabRef: RefObject; +}) { + const trpcClient = useRawTRPCClient(); + const queryClient = useQueryClient(); + const { saved, queryKey } = useSavedWorktreeChanges({ + cloudAgentSessionId, + organizationId, + enabled: open, + }); + const refresh = useMutation({ + mutationFn: () => + organizationId + ? trpcClient.organizations.cloudAgentNext.refreshWorktreeChanges.mutate({ + organizationId, + cloudAgentSessionId, + }) + : trpcClient.cloudAgentNext.refreshWorktreeChanges.mutate({ cloudAgentSessionId }), + onSuccess: result => { + queryClient.setQueryData(queryKey, previous => + preserveNewerWorktreeChanges(previous, { snapshot: result.snapshot }) + ); + }, + retry: false, + }); + const attemptedOpeningRefresh = useRef(false); + const { mutate } = refresh; + useEffect(() => { + if (!open) { + attemptedOpeningRefresh.current = false; + return; + } + if (!saved.isFetchedAfterMount || saved.isFetching || attemptedOpeningRefresh.current) return; + attemptedOpeningRefresh.current = true; + mutate(); + }, [open, saved.isFetchedAfterMount, saved.isFetching, mutate]); + + const snapshot = saved.data?.snapshot; + const groups = useMemo( + () => groupWorktreeChangesByDirectory(snapshot?.files ?? []), + [snapshot?.files] + ); + const tree = useMemo(() => buildWorktreeChangesTree(snapshot?.files ?? []), [snapshot?.files]); + const messages = worktreeChangesMessages({ + snapshot, + savedReadPending: saved.isPending, + savedReadFailed: saved.isError, + refreshPending: refresh.isPending, + refreshFailed: refresh.isError, + refreshStatus: refresh.data?.status, + }); + + return ( + onViewModeChange(value === 'tree' ? 'tree' : 'flat')} + className="flex min-h-0 flex-1 flex-col" + > +
    + + + + + + + +
    + +
    + {messages.empty && ( +

    + {messages.empty} +

    + )} + + {groups.length > 0 && ( +
      + {groups.map(({ directory, files }) => ( +
    • +

      + + {directory || 'Repository root'} + + {files.length} +

      +
        + {files.map(file => ( + + ))} +
      +
      • + ))} +
    + )} +     + + {tree.length > 0 && ( +
      + +
    + )} +     +
    + + {(snapshot || messages.notice) && ( +
    + {snapshot && ( +
    +
    + )} +
    + {messages.notice &&

    {messages.notice}

    } + {snapshot?.truncated &&

    Partial summary · some files omitted.

    } +
    +
    + )} +     + ); +} diff --git a/apps/web/src/components/cloud-agent-next/terminal-tabs.test.ts b/apps/web/src/components/cloud-agent-next/terminal-tabs.test.ts index 10a2e33466..103cab422e 100644 --- a/apps/web/src/components/cloud-agent-next/terminal-tabs.test.ts +++ b/apps/web/src/components/cloud-agent-next/terminal-tabs.test.ts @@ -3,6 +3,8 @@ import { createRoot, type Root } from 'react-dom/client'; import { createRequire } from 'node:module'; import type { CloudAgentWorkspaceTabs } from './CloudAgentWorkspaceTabs'; import type { CloudChatPage as CloudChatPageComponent } from './CloudChatPage'; +import type { ChatHeader } from './ChatHeader'; +import type { WorktreeChangesDrawer } from './WorktreeChanges'; import type { StoredSession } from './types'; import { CHAT_TAB_ID, @@ -12,6 +14,7 @@ import { getWorkspaceTabScope, resetWorkspaceTabs, selectWorkspaceTab, + terminalIdFromTabId, terminalTabId, } from './terminal-tabs'; @@ -121,7 +124,28 @@ jest.mock('./older-messages-scroll', () => ({ shouldAnnounceOlderMessagesArrival: () => false, })); jest.mock('./MobileSidebarToggle', () => ({ MobileSidebarToggle: () => null })); -jest.mock('./ChatHeader', () => ({ ChatHeader: () => null })); +jest.mock('./ChatHeader', () => ({ + ChatHeader: ({ onToggleChanges, changesOpen }: ComponentProps) => + onToggleChanges + ? createElement('button', { + 'data-changes-trigger': true, + 'aria-expanded': changesOpen, + onClick: onToggleChanges, + }) + : null, +})); +jest.mock('./WorktreeChanges', () => ({ + WorktreeChangesDrawer: ({ + cloudAgentSessionId, + organizationId, + open, + }: ComponentProps) => + createElement('aside', { + 'data-changes-owner': cloudAgentSessionId, + 'data-changes-organization': organizationId ?? 'personal', + hidden: !open, + }), +})); jest.mock('./ChatInput', () => ({ ChatInput: () => null })); jest.mock('./OlderMessagesHeader', () => ({ OlderMessagesHeader: () => null })); jest.mock('./MessageBubble', () => ({ MessageBubble: () => null })); @@ -195,10 +219,17 @@ describe('cloud agent workspace terminal tabs', () => { }); }); + it('distinguishes chat from terminal IDs', () => { + expect(terminalIdFromTabId(CHAT_TAB_ID)).toBeNull(); + expect(terminalIdFromTabId(terminalTabId('tab-a'))).toBe('tab-a'); + }); + it('selects chat and existing terminal tabs only', () => { const state = addTerminalTab(createWorkspaceTabsState(), 'tab-a', 'cloud-agent-session-a'); expect(selectWorkspaceTab(state, CHAT_TAB_ID).activeTabId).toBe(CHAT_TAB_ID); + expect(selectWorkspaceTab(state, 'changes')).toBe(state); + expect(selectWorkspaceTab(state, 'unknown')).toBe(state); expect(selectWorkspaceTab(state, terminalTabId('tab-a')).activeTabId).toBe( terminalTabId('tab-a') ); @@ -334,6 +365,15 @@ describe('CloudChatPage terminal ownership across navigation', () => { return terminal; } + function openChanges() { + const trigger = dom.container.querySelector('[data-changes-trigger]'); + if (!trigger) throw new Error('Missing changes trigger'); + act(() => trigger.click()); + const drawer = dom.container.querySelector('[data-changes-owner]'); + expect(drawer?.hidden).toBe(false); + return drawer; + } + function resolveSession(worktreeId: string | null) { mockWorktreeId = worktreeId; mockAtomValues.fetchedSessionData = { @@ -384,6 +424,66 @@ describe('CloudChatPage terminal ownership across navigation', () => { expect(mockClosedPtys).toEqual([]); }); + it('scopes changes to each sibling control session without replacing its worktree terminal', () => { + render(); + const terminal = openTerminal(); + const activeTabId = mockTabs.activeTabId; + const changes = openChanges(); + expect(changes?.getAttribute('data-changes-owner')).toBe('workspace_recent'); + expect(mockTabs.activeTabId).toBe(activeTabId); + + mockSessionId = 'ses_historical'; + render(); + expect(dom.container.querySelector('[data-changes-owner]')).toBeNull(); + expect(dom.container.querySelector('[data-changes-trigger]')).toBeNull(); + + resolveSession('worktree_shared'); + expect(dom.container.querySelector('[data-changes-owner]')?.hidden).toBe(true); + const siblingChanges = openChanges(); + expect(siblingChanges?.getAttribute('data-changes-owner')).toBe('workspace_ses_historical'); + expect(siblingChanges).not.toBe(changes); + expect(dom.container.querySelector('[data-pty-owner]')).toBe(terminal); + expect(mockClosedPtys).toEqual([]); + }); + + it('clears changes when the last chat closes before its session atoms are cleared', () => { + render(); + const terminal = openTerminal(); + openChanges(); + + mockSessionId = null; + render(); + expect(dom.container.querySelector('[data-changes-owner]')).toBeNull(); + expect(dom.container.querySelector('[data-changes-trigger]')).toBeNull(); + expect(dom.container.querySelector('[data-pty-owner]')).toBe(terminal); + expect(mockClosedPtys).toEqual([]); + + mockSessionId = 'ses_recent'; + render(); + expect(dom.container.querySelector('[data-changes-owner]')?.hidden).toBe(true); + }); + + it('keeps changes hidden across organization navigation until matching session data resolves', () => { + render(); + const personalChanges = openChanges(); + expect(personalChanges?.getAttribute('data-changes-organization')).toBe('personal'); + + render({ organizationId: 'organization-a' }); + expect(dom.container.querySelector('[data-changes-owner]')).toBeNull(); + expect(dom.container.querySelector('[data-changes-trigger]')).toBeNull(); + + mockAtomValues.fetchedSessionData = { + kiloSessionId: mockSessionId, + organizationId: 'organization-a', + worktreeId: mockWorktreeId, + }; + render({ organizationId: 'organization-a' }); + expect(dom.container.querySelector('[data-changes-owner]')?.hidden).toBe(true); + const organizationChanges = openChanges(); + expect(organizationChanges?.getAttribute('data-changes-organization')).toBe('organization-a'); + expect(organizationChanges).not.toBe(personalChanges); + }); + it.each(['worktree_other', null])( 'closes the original PTY only after a different destination resolves to %s', destinationWorktreeId => { diff --git a/apps/web/src/components/cloud-agent-next/terminal-tabs.ts b/apps/web/src/components/cloud-agent-next/terminal-tabs.ts index 53d91cac78..4dacdde709 100644 --- a/apps/web/src/components/cloud-agent-next/terminal-tabs.ts +++ b/apps/web/src/components/cloud-agent-next/terminal-tabs.ts @@ -20,7 +20,7 @@ export function terminalTabId(terminalId: string): TerminalTabId { } export function terminalIdFromTabId(tabId: WorkspaceTabId): string | null { - if (tabId === CHAT_TAB_ID) return null; + if (!tabId.startsWith('terminal:')) return null; return tabId.slice('terminal:'.length); } @@ -61,18 +61,16 @@ export function addTerminalTab( export function selectWorkspaceTab( state: WorkspaceTabsState, - activeTabId: WorkspaceTabId + activeTabId: string ): WorkspaceTabsState { if (activeTabId === CHAT_TAB_ID) { return { ...state, activeTabId }; } - const terminalId = terminalIdFromTabId(activeTabId); - if (!terminalId || !state.terminals.some(tab => tab.id === terminalId)) { - return state; - } + const terminal = state.terminals.find(tab => terminalTabId(tab.id) === activeTabId); + if (!terminal) return state; - return { ...state, activeTabId }; + return { ...state, activeTabId: terminalTabId(terminal.id) }; } export function closeTerminalTab( diff --git a/apps/web/src/components/cloud-agent-next/worktree-changes.test.ts b/apps/web/src/components/cloud-agent-next/worktree-changes.test.ts new file mode 100644 index 0000000000..20c8d29410 --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/worktree-changes.test.ts @@ -0,0 +1,680 @@ +import { QueryClient } from '@tanstack/react-query'; +import type { + GetWorktreeChangesOutput, + WorktreeChangesSnapshot, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; +import { + buildWorktreeChangesTree, + canOpenWorktreeChanges, + deserializeWorktreeChangesViewMode, + formatWorktreeChangesBaseBranch, + getWorktreeChangesTotals, + groupWorktreeChangesByDirectory, + preserveNewerWorktreeChanges, + worktreeChangesMessages, + type WorktreeChangesFile, + type WorktreeChangesTreeNode, +} from './worktree-changes'; + +const snapshot: WorktreeChangesSnapshot = { + schemaVersion: 1, + revision: 3, + capturedAt: '2026-08-26T12:00:00.000Z', + comparison: { baseRef: 'origin/main', mergeBase: 'a'.repeat(40), head: 'b'.repeat(40) }, + files: [ + { + path: 'src/odd\nfile\tname.ts', + status: 'modified', + additions: 2, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, + }, + ], + truncated: false, +}; + +function messages(overrides: Partial[0]> = {}) { + return worktreeChangesMessages({ + snapshot, + savedReadPending: false, + savedReadFailed: false, + refreshPending: false, + refreshFailed: false, + refreshStatus: undefined, + ...overrides, + }); +} + +describe('getWorktreeChangesTotals', () => { + it.each([null, undefined])('returns null for an unavailable snapshot: %p', missing => { + expect(getWorktreeChangesTotals(missing)).toBeNull(); + }); + + it('returns complete zero totals for a saved empty snapshot', () => { + expect(getWorktreeChangesTotals({ ...snapshot, files: [] })).toEqual({ + fileCount: 0, + additions: 0, + deletions: 0, + countsComplete: true, + }); + }); + + it('sums text changes across statuses and tracking states without mutating the snapshot', () => { + const mixedSnapshot: WorktreeChangesSnapshot = { + ...snapshot, + files: [ + { ...snapshot.files[0] }, + { + ...snapshot.files[0], + path: 'src/added.ts', + status: 'added', + additions: 5, + deletions: 0, + }, + { + ...snapshot.files[0], + path: 'src/deleted.ts', + status: 'deleted', + additions: 0, + deletions: 7, + }, + { + ...snapshot.files[0], + path: 'untracked.txt', + status: 'added', + additions: 3, + deletions: 0, + tracked: false, + }, + { + ...snapshot.files[0], + path: 'image.png', + additions: 100, + deletions: 200, + binary: true, + countsComplete: false, + }, + ], + }; + const original = structuredClone(mixedSnapshot); + + expect(getWorktreeChangesTotals(mixedSnapshot)).toEqual({ + fileCount: 5, + additions: 10, + deletions: 8, + countsComplete: true, + }); + expect(mixedSnapshot).toEqual(original); + }); + + it.each([true, false])( + 'ignores binary counters and completeness when binary countsComplete is %s', + countsComplete => { + expect( + getWorktreeChangesTotals({ + ...snapshot, + files: [ + { + ...snapshot.files[0], + path: 'untracked.bin', + status: 'added', + tracked: false, + binary: true, + additions: 100, + deletions: 200, + countsComplete, + }, + ], + }) + ).toEqual({ fileCount: 1, additions: 0, deletions: 0, countsComplete: true }); + } + ); + + it('preserves known subtotals when the snapshot is truncated', () => { + expect(getWorktreeChangesTotals({ ...snapshot, truncated: true })).toEqual({ + fileCount: 1, + additions: 2, + deletions: 1, + countsComplete: false, + }); + }); + + it('does not report omitted entries as complete zero totals', () => { + expect(getWorktreeChangesTotals({ ...snapshot, files: [], truncated: true })).toEqual({ + fileCount: 0, + additions: 0, + deletions: 0, + countsComplete: false, + }); + }); + + it('includes known counts from partial text files alongside complete files', () => { + expect( + getWorktreeChangesTotals({ + ...snapshot, + files: [ + { ...snapshot.files[0], countsComplete: false }, + { + ...snapshot.files[0], + path: 'untracked.txt', + status: 'added', + additions: 7, + deletions: 0, + tracked: false, + }, + ], + }) + ).toEqual({ fileCount: 2, additions: 9, deletions: 1, countsComplete: false }); + }); + + it('does not report unknown text counts as complete zero totals', () => { + expect( + getWorktreeChangesTotals({ + ...snapshot, + files: [{ ...snapshot.files[0], additions: 0, deletions: 0, countsComplete: false }], + }) + ).toEqual({ fileCount: 1, additions: 0, deletions: 0, countsComplete: false }); + }); +}); + +describe('worktree changes capability', () => { + it('uses the control ID rather than a Kilo or legacy session ID', () => { + expect(canOpenWorktreeChanges('workspace_12345678-1234-4234-9234-123456789abc', false)).toBe( + true + ); + expect(canOpenWorktreeChanges('agent_12345678-1234-4234-9234-123456789abc', false)).toBe(false); + expect(canOpenWorktreeChanges('ses_12345678901234567890123456', false)).toBe(false); + expect(canOpenWorktreeChanges(null, false)).toBe(false); + expect(canOpenWorktreeChanges(undefined, false)).toBe(false); + }); + + it('excludes read-only views independently of the session prefix', () => { + expect(canOpenWorktreeChanges('workspace_12345678-1234-4234-9234-123456789abc', true)).toBe( + false + ); + }); +}); + +describe('saved worktree changes cache', () => { + let queryClient: QueryClient; + const queryKey = ['worktree-changes', 'personal', 'workspace-a']; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { structuralSharing: preserveNewerWorktreeChanges, retry: false }, + }, + }); + }); + + afterEach(() => queryClient.clear()); + + it('accepts newer revisions, preserves unusual paths, and ignores clock ordering', () => { + const newer = { ...snapshot, revision: 4, capturedAt: '2026-08-26T11:00:00.000Z' }; + expect(preserveNewerWorktreeChanges({ snapshot }, { snapshot: newer })).toEqual({ + snapshot: newer, + }); + }); + + it.each([ + null, + { ...snapshot, revision: 2 }, + { ...snapshot, capturedAt: '2026-08-26T13:00:00.000Z' }, + ])('does not replace newer or equal saved revisions', incoming => { + expect(preserveNewerWorktreeChanges({ snapshot }, { snapshot: incoming })).toEqual({ + snapshot, + }); + }); + + it('can initialize saved data from an offline or failed refresh response', () => { + expect(preserveNewerWorktreeChanges(undefined, { snapshot })).toEqual({ snapshot }); + expect(preserveNewerWorktreeChanges({ snapshot: null }, { snapshot })).toEqual({ snapshot }); + expect(preserveNewerWorktreeChanges(undefined, { snapshot: null })).toEqual({ snapshot: null }); + }); + + it('does not let an in-flight saved read overwrite a newer refresh', async () => { + queryClient.setQueryData(queryKey, { snapshot }); + const read = Promise.withResolvers(); + const pendingRead = queryClient.fetchQuery({ queryKey, queryFn: () => read.promise }); + const newer = { ...snapshot, revision: 4 }; + queryClient.setQueryData(queryKey, { snapshot: newer }); + read.resolve({ snapshot }); + await pendingRead; + + expect(queryClient.getQueryData(queryKey)).toEqual({ snapshot: newer }); + }); + + it('keeps delayed responses scoped to their session without a global revision floor', () => { + const nextSessionKey = ['worktree-changes', 'personal', 'workspace-b']; + queryClient.setQueryData(queryKey, { snapshot }); + expect(queryClient.getQueryData(nextSessionKey)).toBeUndefined(); + queryClient.setQueryData(nextSessionKey, { snapshot: { ...snapshot, revision: 1 } }); + queryClient.setQueryData(queryKey, { snapshot: { ...snapshot, revision: 4 } }); + + expect(queryClient.getQueryData(nextSessionKey)).toEqual({ + snapshot: { ...snapshot, revision: 1 }, + }); + }); + + it('does not share data between personal and organization scopes', () => { + queryClient.setQueryData(queryKey, { snapshot }); + expect( + queryClient.getQueryData(['worktree-changes', 'organization-a', 'workspace-a']) + ).toBeUndefined(); + }); +}); + +describe('worktree changes messages', () => { + it('keeps saved content while refresh is pending', () => { + expect(messages({ refreshPending: true })).toEqual({ + notice: 'Refreshing…', + empty: null, + }); + }); + + it('labels offline saved data without treating it as empty', () => { + expect(messages({ refreshStatus: 'offline' })).toEqual({ + notice: 'Offline · showing saved changes.', + empty: null, + }); + }); + + it.each([{ refreshStatus: 'failed' as const }, { refreshFailed: true }])( + 'preserves saved content when refresh fails', + failure => { + expect(messages(failure)).toEqual({ + notice: 'Refresh failed · showing saved changes.', + empty: null, + }); + } + ); + + it('preserves saved content when the saved read fails', () => { + expect(messages({ savedReadFailed: true })).toEqual({ + notice: 'Load failed · showing saved changes.', + empty: null, + }); + }); + + it.each([ + { + state: { refreshPending: true, refreshFailed: true, refreshStatus: 'offline' as const }, + notice: 'Refreshing…', + }, + { + state: { refreshFailed: true, refreshStatus: 'offline' as const }, + notice: 'Refresh failed · showing saved changes.', + }, + { + state: { refreshStatus: 'offline' as const }, + notice: 'Offline · showing saved changes.', + }, + ])('prioritizes "$notice" over a saved-read failure', ({ state, notice }) => { + expect(messages({ ...state, savedReadFailed: true })).toEqual({ notice, empty: null }); + }); + + it('distinguishes a missing saved snapshot from an unsuccessful saved read', () => { + expect(messages({ snapshot: null }).empty).toBe('No saved changes yet.'); + expect(messages({ snapshot: undefined, savedReadFailed: true }).empty).toBe( + 'Could not load saved changes.' + ); + expect(messages({ snapshot: undefined, savedReadPending: true }).empty).toBe( + 'Loading saved changes…' + ); + }); + + it('reports an initial read failure even if the subsequent refresh also fails', () => { + expect(messages({ snapshot: undefined, savedReadFailed: true, refreshFailed: true })).toEqual({ + notice: 'Refresh failed.', + empty: 'Could not load saved changes.', + }); + }); + + it('reports an empty successful summary without repeating its timestamp', () => { + expect(messages({ snapshot: { ...snapshot, files: [] } })).toEqual({ + notice: null, + empty: 'No changes.', + }); + }); + + it('does not call a truncated summary clean when all entries were omitted', () => { + expect(messages({ snapshot: { ...snapshot, files: [], truncated: true } }).empty).toBeNull(); + }); + + it('distinguishes an offline sandbox with no saved summary', () => { + expect(messages({ snapshot: null, refreshStatus: 'offline' })).toEqual({ + notice: 'Sandbox offline.', + empty: 'No saved changes yet.', + }); + }); +}); + +describe('groupWorktreeChangesByDirectory', () => { + it('does not create empty folder groups', () => { + expect(groupWorktreeChangesByDirectory([])).toEqual([]); + }); + + it('groups only direct siblings under their complete parent paths', () => { + const files: WorktreeChangesFile[] = [ + 'src/b.ts', + 'src/nested/a.ts', + 'README', + 'src/a.ts', + 'tests/src/a.ts', + ].map(path => ({ ...snapshot.files[0], path })); + + expect(groupWorktreeChangesByDirectory(files)).toEqual([ + { directory: '', files: [files[2]] }, + { directory: 'src', files: [files[3], files[0]] }, + { directory: 'src/nested', files: [files[1]] }, + { directory: 'tests/src', files: [files[4]] }, + ]); + }); + + it('sorts directories and files deterministically with repository-root files first', () => { + const files: WorktreeChangesFile[] = [ + 'z/z.ts', + 'a/z.ts', + 'Z.ts', + 'a/a.ts', + 'a.ts', + 'é/file.ts', + 'A/file.ts', + 'a/Z.ts', + ].map(path => ({ ...snapshot.files[0], path })); + const groups = groupWorktreeChangesByDirectory(files); + + expect(groups.map(group => group.directory)).toEqual(['', 'A', 'a', 'z', 'é']); + expect(groups[0].files.map(file => file.path)).toEqual(['Z.ts', 'a.ts']); + expect(groups[2].files.map(file => file.path)).toEqual(['a/Z.ts', 'a/a.ts', 'a/z.ts']); + expect(groupWorktreeChangesByDirectory(files.toReversed())).toEqual(groups); + }); + + it('preserves unusual paths and metadata without mutating files or treating backslashes as separators', () => { + const binary = Object.freeze({ + ...snapshot.files[0], + path: '__proto__/constructor/odd\\file\t\n雪.ts', + binary: true, + countsComplete: false, + }); + const root = Object.freeze({ + ...snapshot.files[0], + path: 'root\\file\t\n雪.ts', + tracked: false, + }); + const dotfile = Object.freeze({ ...snapshot.files[0], path: '.changeset/fix.md' }); + const files = Object.freeze([binary, root, dotfile]); + const groups = groupWorktreeChangesByDirectory(files); + + expect(groups).toEqual([ + { directory: '', files: [root] }, + { directory: '.changeset', files: [dotfile] }, + { directory: '__proto__/constructor', files: [binary] }, + ]); + expect(groups[2].files[0]).toBe(binary); + expect(files).toEqual([binary, root, dotfile]); + }); + + it('keeps a deleted file distinct from added files in a folder with the same name', () => { + const deleted: WorktreeChangesFile = { + ...snapshot.files[0], + path: 'a', + status: 'deleted', + additions: 0, + }; + const added: WorktreeChangesFile = { + ...snapshot.files[0], + path: 'a/b', + status: 'added', + deletions: 0, + tracked: false, + }; + + expect(groupWorktreeChangesByDirectory([added, deleted])).toEqual([ + { directory: '', files: [deleted] }, + { directory: 'a', files: [added] }, + ]); + }); +}); + +describe('formatWorktreeChangesBaseBranch', () => { + it.each([ + ['refs/remotes/origin/master', 'master'], + ['refs/remotes/origin/main', 'main'], + ['refs/remotes/origin/feature/a', 'feature/a'], + ['refs/remotes/upstream/release/2026', 'release/2026'], + ['refs/remotes/origin/origin/topic', 'origin/topic'], + ['refs/remotes/origin/refs/remotes/upstream/topic', 'refs/remotes/upstream/topic'], + ['main', 'main'], + ['feature/a', 'feature/a'], + ['refs/tags/v1', 'refs/tags/v1'], + ])('formats %s without losing branch path segments', (baseRef, branch) => { + expect(formatWorktreeChangesBaseBranch(baseRef)).toBe(branch); + }); +}); + +describe('deserializeWorktreeChangesViewMode', () => { + it.each(['flat', 'tree'])('restores the JSON-serialized %s preference', mode => { + expect(deserializeWorktreeChangesViewMode(JSON.stringify(mode))).toBe(mode); + }); + + it('accepts JSON whitespace around a valid preference', () => { + expect(deserializeWorktreeChangesViewMode(' \n"tree"\t ')).toBe('tree'); + }); + + it.each([ + '', + 'flat', + 'tree', + 'undefined', + '{', + '"tree', + 'null', + 'true', + '1', + '""', + '"list"', + '"Tree"', + '" tree "', + '["tree"]', + '{"viewMode":"tree"}', + ])('defaults to flat for invalid persisted data: %p', value => { + expect(deserializeWorktreeChangesViewMode(value)).toBe('flat'); + }); +}); + +describe('buildWorktreeChangesTree', () => { + it('returns an empty tree for no changes', () => { + expect(buildWorktreeChangesTree([])).toEqual([]); + }); + + it('groups nested paths while keeping root files and repeated names distinct', () => { + const files: WorktreeChangesFile[] = [ + 'index.ts', + 'src/index.ts', + 'src/src/index.ts', + 'tests/src/index.ts', + ].map(path => ({ ...snapshot.files[0], path })); + + expect(buildWorktreeChangesTree(files)).toEqual([ + { + kind: 'directory', + name: 'src', + path: 'src', + children: [ + { + kind: 'directory', + name: 'src', + path: 'src/src', + children: [ + { kind: 'file', name: 'index.ts', path: 'src/src/index.ts', file: files[2] }, + ], + }, + { kind: 'file', name: 'index.ts', path: 'src/index.ts', file: files[1] }, + ], + }, + { + kind: 'directory', + name: 'tests', + path: 'tests', + children: [ + { + kind: 'directory', + name: 'src', + path: 'tests/src', + children: [ + { kind: 'file', name: 'index.ts', path: 'tests/src/index.ts', file: files[3] }, + ], + }, + ], + }, + { kind: 'file', name: 'index.ts', path: 'index.ts', file: files[0] }, + ]); + }); + + it('sorts directories first and names deterministically at every level', () => { + const files: WorktreeChangesFile[] = [ + 'z.ts', + 'beta/z.ts', + 'alpha/z.ts', + 'beta/a.ts', + 'a.ts', + 'beta/zeta/z.ts', + 'beta/alpha/z.ts', + 'Z.ts', + 'é.ts', + ].map(path => ({ ...snapshot.files[0], path })); + const tree = buildWorktreeChangesTree(files); + + expect(tree.map(({ kind, name }) => [kind, name])).toEqual([ + ['directory', 'alpha'], + ['directory', 'beta'], + ['file', 'Z.ts'], + ['file', 'a.ts'], + ['file', 'z.ts'], + ['file', 'é.ts'], + ]); + expect(tree[1]).toMatchObject({ + children: [ + { kind: 'directory', name: 'alpha' }, + { kind: 'directory', name: 'zeta' }, + { kind: 'file', name: 'a.ts' }, + { kind: 'file', name: 'z.ts' }, + ], + }); + expect(buildWorktreeChangesTree(files.toReversed())).toEqual(tree); + }); + + it('preserves exact paths and metadata without splitting backslashes or mutating input', () => { + const rootFile: WorktreeChangesFile = Object.freeze({ + path: 'root\\file\t\n雪.ts', + status: 'deleted', + additions: 0, + deletions: 17, + tracked: true, + binary: true, + countsComplete: false, + }); + const nestedFile: WorktreeChangesFile = Object.freeze({ + path: ' \tdir\\name\n / e\u0301\\雪\t\n.ts ', + status: 'added', + additions: 12, + deletions: 0, + tracked: false, + binary: false, + countsComplete: true, + }); + const files = Object.freeze([rootFile, nestedFile]); + + expect(buildWorktreeChangesTree(files)).toEqual([ + { + kind: 'directory', + name: ' \tdir\\name\n ', + path: ' \tdir\\name\n ', + children: [ + { + kind: 'file', + name: ' e\u0301\\雪\t\n.ts ', + path: nestedFile.path, + file: nestedFile, + }, + ], + }, + { kind: 'file', name: rootFile.path, path: rootFile.path, file: rootFile }, + ]); + expect(files).toEqual([rootFile, nestedFile]); + }); + + it('keeps a deleted file alongside an added child at the same directory path in either order', () => { + const deleted: WorktreeChangesFile = { + ...snapshot.files[0], + path: 'a', + status: 'deleted', + additions: 0, + }; + const added: WorktreeChangesFile = { + ...snapshot.files[0], + path: 'a/b', + status: 'added', + deletions: 0, + tracked: false, + }; + const expected: WorktreeChangesTreeNode[] = [ + { + kind: 'directory', + name: 'a', + path: 'a', + children: [{ kind: 'file', name: 'b', path: 'a/b', file: added }], + }, + { kind: 'file', name: 'a', path: 'a', file: deleted }, + ]; + + expect(buildWorktreeChangesTree([deleted, added])).toEqual(expected); + expect(buildWorktreeChangesTree([added, deleted])).toEqual(expected); + }); + + it('treats prototype property names as ordinary path segments', () => { + const files: WorktreeChangesFile[] = [ + '__proto__/constructor/toString', + 'constructor/__proto__', + ].map(path => ({ ...snapshot.files[0], path })); + + expect(buildWorktreeChangesTree(files)).toEqual([ + { + kind: 'directory', + name: '__proto__', + path: '__proto__', + children: [ + { + kind: 'directory', + name: 'constructor', + path: '__proto__/constructor', + children: [ + { + kind: 'file', + name: 'toString', + path: '__proto__/constructor/toString', + file: files[0], + }, + ], + }, + ], + }, + { + kind: 'directory', + name: 'constructor', + path: 'constructor', + children: [ + { + kind: 'file', + name: '__proto__', + path: 'constructor/__proto__', + file: files[1], + }, + ], + }, + ]); + }); +}); diff --git a/apps/web/src/components/cloud-agent-next/worktree-changes.ts b/apps/web/src/components/cloud-agent-next/worktree-changes.ts new file mode 100644 index 0000000000..42bc325ff4 --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/worktree-changes.ts @@ -0,0 +1,179 @@ +import { + getWorktreeChangesOutputSchema, + type GetWorktreeChangesOutput, + type RefreshWorktreeChangesOutput, + type WorktreeChangesSnapshot, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; + +export type WorktreeChangesViewMode = 'flat' | 'tree'; + +export type WorktreeChangesFile = WorktreeChangesSnapshot['files'][number]; + +export type WorktreeChangesTreeNode = + | { + kind: 'directory'; + name: string; + path: string; + children: WorktreeChangesTreeNode[]; + } + | { + kind: 'file'; + name: string; + path: string; + file: WorktreeChangesFile; + }; + +export function getWorktreeChangesTotals( + snapshot: WorktreeChangesSnapshot | null | undefined +): { fileCount: number; additions: number; deletions: number; countsComplete: boolean } | null { + if (!snapshot) return null; + + let additions = 0; + let deletions = 0; + let countsComplete = !snapshot.truncated; + + for (const file of snapshot.files) { + if (file.binary) continue; + additions += file.additions; + deletions += file.deletions; + countsComplete = countsComplete && file.countsComplete; + } + + return { fileCount: snapshot.files.length, additions, deletions, countsComplete }; +} + +export function deserializeWorktreeChangesViewMode(value: string): WorktreeChangesViewMode { + try { + const mode: unknown = JSON.parse(value); + return mode === 'flat' || mode === 'tree' ? mode : 'flat'; + } catch { + return 'flat'; + } +} + +export function groupWorktreeChangesByDirectory( + files: readonly WorktreeChangesFile[] +): { directory: string; files: WorktreeChangesFile[] }[] { + const directories = new Map(); + + for (const file of files) { + const separator = file.path.lastIndexOf('/'); + const directory = separator === -1 ? '' : file.path.slice(0, separator); + const group = directories.get(directory); + if (group) group.push(file); + else directories.set(directory, [file]); + } + + return Array.from(directories, ([directory, files]) => ({ + directory, + files: files.sort((a, b) => (a.path < b.path ? -1 : a.path > b.path ? 1 : 0)), + })).sort((a, b) => (a.directory < b.directory ? -1 : a.directory > b.directory ? 1 : 0)); +} + +export function formatWorktreeChangesBaseBranch(baseRef: string): string { + return baseRef.replace(/^refs\/remotes\/[^/]+\//, ''); +} + +export function buildWorktreeChangesTree( + files: readonly WorktreeChangesFile[] +): WorktreeChangesTreeNode[] { + const tree: WorktreeChangesTreeNode[] = []; + const directories = new Map>(); + + for (const file of files) { + const parts = file.path.split('/'); + let siblings = tree; + let path = ''; + + for (const [index, name] of parts.entries()) { + if (index === parts.length - 1) { + siblings.push({ kind: 'file', name, path: file.path, file }); + break; + } + + path = index === 0 ? name : `${path}/${name}`; + let directory = directories.get(path); + if (!directory) { + directory = { kind: 'directory', name, path, children: [] }; + directories.set(path, directory); + siblings.push(directory); + } + siblings = directory.children; + } + } + + const compareNodes = (a: WorktreeChangesTreeNode, b: WorktreeChangesTreeNode): number => { + if (a.kind !== b.kind) return a.kind === 'directory' ? -1 : 1; + return a.name < b.name ? -1 : a.name > b.name ? 1 : 0; + }; + + tree.sort(compareNodes); + for (const directory of directories.values()) { + directory.children.sort(compareNodes); + } + return tree; +} + +export function canOpenWorktreeChanges( + cloudAgentSessionId: string | null | undefined, + isReadOnly: boolean +): boolean { + return !isReadOnly && cloudAgentSessionId?.startsWith('workspace_') === true; +} + +export function preserveNewerWorktreeChanges( + previous: unknown, + incoming: unknown +): GetWorktreeChangesOutput { + const next = getWorktreeChangesOutputSchema.parse(incoming); + const current = + previous === undefined ? undefined : getWorktreeChangesOutputSchema.parse(previous); + if ( + current?.snapshot && + (!next.snapshot || current.snapshot.revision >= next.snapshot.revision) + ) { + return current; + } + return next; +} + +export function worktreeChangesMessages({ + snapshot, + savedReadPending, + savedReadFailed, + refreshPending, + refreshFailed, + refreshStatus, +}: { + snapshot: WorktreeChangesSnapshot | null | undefined; + savedReadPending: boolean; + savedReadFailed: boolean; + refreshPending: boolean; + refreshFailed: boolean; + refreshStatus: RefreshWorktreeChangesOutput['status'] | undefined; +}): { notice: string | null; empty: string | null } { + let notice: string | null = null; + if (refreshPending) { + notice = 'Refreshing…'; + } else if (refreshFailed || refreshStatus === 'failed') { + notice = snapshot ? 'Refresh failed · showing saved changes.' : 'Refresh failed.'; + } else if (refreshStatus === 'offline') { + notice = snapshot ? 'Offline · showing saved changes.' : 'Sandbox offline.'; + } + + if (snapshot) { + return { + notice: notice ?? (savedReadFailed ? 'Load failed · showing saved changes.' : null), + empty: snapshot.files.length === 0 && !snapshot.truncated ? 'No changes.' : null, + }; + } + + return { + notice, + empty: savedReadPending + ? 'Loading saved changes…' + : savedReadFailed + ? 'Could not load saved changes.' + : 'No saved changes yet.', + }; +} diff --git a/apps/web/src/components/ui/sheet.tsx b/apps/web/src/components/ui/sheet.tsx index 2b425adf15..24ea2647fb 100644 --- a/apps/web/src/components/ui/sheet.tsx +++ b/apps/web/src/components/ui/sheet.tsx @@ -53,6 +53,7 @@ function SheetContent({ showOverlay = true, overlayClassName, dismissibleOverlay = false, + showCloseButton = true, ...props }: React.ComponentProps & { side?: 'top' | 'right' | 'bottom' | 'left'; @@ -60,6 +61,7 @@ function SheetContent({ showOverlay?: boolean; overlayClassName?: string; dismissibleOverlay?: boolean; + showCloseButton?: boolean; }) { return ( @@ -93,10 +95,12 @@ function SheetContent({ {...props} > {children} - - + {showCloseButton && ( + + + )} ); 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 20654e9f94..8dc66163eb 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 @@ -1,7 +1,7 @@ import { afterEach, beforeEach, describe, expect, it, jest } from '@jest/globals'; import type * as TrpcClientModule from '@trpc/client'; +import type * as CloudAgentClientModule from './cloud-agent-client'; import type { - CloudAgentNextClient as CloudAgentNextClientType, CreateWorktreeChatInput, CreateWorktreeChatOutput, DeleteWorktreeInput, @@ -9,14 +9,18 @@ import type { PrepareSessionInput, SendMessageInput, } from './cloud-agent-client'; +import type { WorktreeChangesSnapshot } from '@kilocode/worker-utils/cloud-agent-worktree-changes'; const mockCreateTRPCClient = jest.fn(() => ({})); const mockHttpLink = jest.fn<(options: { url: string; headers: () => Record }) => undefined>(); const mockCaptureException = jest.fn(); +const mockGetWorktreeChanges = + jest.fn<(input: { cloudAgentSessionId: string }) => Promise>(); +const mockRefreshWorktreeChanges = + jest.fn<(input: { cloudAgentSessionId: string }) => Promise>(); import type * as SentryModule from '@sentry/nextjs'; -import type * as CloudAgentClientModule from './cloud-agent-client'; import type { SandboxStatusSnapshot } from '@/routers/cloud-agent-next-schemas'; jest.mock('@/lib/dotenvx', () => ({ @@ -84,12 +88,99 @@ beforeEach(() => { // Load the real `closeCloudAgentOrgStreams` (the module mock above does not // expose it) so the test exercises the actual fetch call, not a stub. -const { closeCloudAgentOrgStreams, CloudAgentNextClient } = jest.requireActual( - './cloud-agent-client' -) as { - closeCloudAgentOrgStreams: (userId: string, organizationId: string) => Promise; - CloudAgentNextClient: typeof CloudAgentNextClientType; -}; +const { closeCloudAgentOrgStreams, CloudAgentNextClient } = + jest.requireActual('./cloud-agent-client'); + +describe('CloudAgentNextClient worktree changes', () => { + const cloudAgentSessionId = 'workspace_12345678-1234-4234-9234-123456789abc'; + const snapshot: WorktreeChangesSnapshot = { + schemaVersion: 1, + revision: 1, + capturedAt: '2026-08-26T12:00:00.000Z', + comparison: { baseRef: 'origin/main', mergeBase: 'a'.repeat(40), head: 'b'.repeat(40) }, + files: [ + { + path: 'src/odd\nfile.ts', + status: 'modified', + additions: 2, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, + }, + ], + truncated: false, + }; + + beforeEach(() => { + mockGetWorktreeChanges.mockReset(); + mockRefreshWorktreeChanges.mockReset(); + mockCreateTRPCClient.mockReturnValueOnce({ + getWorktreeChanges: { query: mockGetWorktreeChanges }, + refreshWorktreeChanges: { mutate: mockRefreshWorktreeChanges }, + }); + }); + + it.each([null, snapshot])( + 'validates saved query responses without changing paths', + async saved => { + mockGetWorktreeChanges.mockResolvedValue({ snapshot: saved }); + const client = new CloudAgentNextClient('token'); + + await expect(client.getWorktreeChanges(cloudAgentSessionId)).resolves.toEqual({ + snapshot: saved, + }); + expect(mockGetWorktreeChanges).toHaveBeenCalledWith({ cloudAgentSessionId }); + expect(mockRefreshWorktreeChanges).not.toHaveBeenCalled(); + } + ); + + it.each(['refreshed', 'offline', 'failed'] as const)( + 'validates %s refresh responses', + async status => { + mockRefreshWorktreeChanges.mockResolvedValue({ status, snapshot }); + const client = new CloudAgentNextClient('token'); + + await expect(client.refreshWorktreeChanges(cloudAgentSessionId)).resolves.toEqual({ + status, + snapshot, + }); + expect(mockRefreshWorktreeChanges).toHaveBeenCalledWith({ cloudAgentSessionId }); + expect(mockGetWorktreeChanges).not.toHaveBeenCalled(); + } + ); + + it.each(['offline', 'failed'] as const)('accepts %s without a saved snapshot', async status => { + mockRefreshWorktreeChanges.mockResolvedValue({ status, snapshot: null }); + await expect( + new CloudAgentNextClient('token').refreshWorktreeChanges(cloudAgentSessionId) + ).resolves.toEqual({ status, snapshot: null }); + }); + + it.each([ + { snapshot: { ...snapshot, schemaVersion: 2 } }, + { snapshot: { ...snapshot, revision: 0 } }, + { snapshot: { ...snapshot, capturedAt: 'not-a-date' } }, + { snapshot: { ...snapshot, files: [{ ...snapshot.files[0], patch: 'file content' }] } }, + { snapshot: { ...snapshot, files: [{ ...snapshot.files[0], countsComplete: undefined }] } }, + ])('rejects invalid persisted responses', async response => { + mockGetWorktreeChanges.mockResolvedValue(response); + await expect( + new CloudAgentNextClient('token').getWorktreeChanges(cloudAgentSessionId) + ).rejects.toThrow(); + }); + + it.each([ + { status: 'refreshed', snapshot: null }, + { status: 'unknown', snapshot }, + { status: 'failed', snapshot: { ...snapshot, revision: -1 } }, + ])('rejects invalid refresh responses', async response => { + mockRefreshWorktreeChanges.mockResolvedValue(response); + await expect( + new CloudAgentNextClient('token').refreshWorktreeChanges(cloudAgentSessionId) + ).rejects.toThrow(); + }); +}); describe('createCloudAgentNextClientForModel', () => { it('returns the default client when the model is paid and has no BYOK', () => { 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 25b64da160..293c2d9138 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 @@ -11,6 +11,12 @@ import { captureException } from '@sentry/nextjs'; import { INTERNAL_API_SECRET } from '@/lib/config.server'; import { parseCustomerBillingFailure } from '@kilocode/cloud-agent-sdk'; import type { CloudAgentWorktreeId } from '@kilocode/session-ingest-contracts'; +import { + getWorktreeChangesOutputSchema, + refreshWorktreeChangesOutputSchema, + type GetWorktreeChangesOutput, + type RefreshWorktreeChangesOutput, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; import type { SendMessagePayload } from './types.js'; import { SandboxStatusSnapshotSchema, @@ -573,6 +579,12 @@ type CloudAgentNextTRPCClient = { getSandboxStatus: { query: (input: GetSessionInput) => Promise; }; + getWorktreeChanges: { + query: (input: GetSessionInput) => Promise; + }; + refreshWorktreeChanges: { + mutate: (input: GetSessionInput) => Promise; + }; getComputeBillingStatus: { query: (input: GetSessionInput) => Promise; }; @@ -830,6 +842,18 @@ export class CloudAgentNextClient { } } + async getWorktreeChanges(cloudAgentSessionId: string): Promise { + return getWorktreeChangesOutputSchema.parse( + await this.client.getWorktreeChanges.query({ cloudAgentSessionId }) + ); + } + + async refreshWorktreeChanges(cloudAgentSessionId: string): Promise { + return refreshWorktreeChangesOutputSchema.parse( + await this.client.refreshWorktreeChanges.mutate({ cloudAgentSessionId }) + ); + } + async getComputeBillingStatus(cloudAgentSessionId: string): Promise { return await this.client.getComputeBillingStatus.query({ cloudAgentSessionId }); } diff --git a/apps/web/src/routers/cloud-agent-next-router.test.ts b/apps/web/src/routers/cloud-agent-next-router.test.ts index 1f6374c87c..1af7167ce2 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -1,7 +1,14 @@ import { describe, expect, it, jest, beforeAll, beforeEach } from '@jest/globals'; -import type { User } from '@kilocode/db/schema'; +import { cli_sessions_v2, organizations, type User } from '@kilocode/db/schema'; import type { createWorktreeChat as CreateWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; import type * as MinimumVersionModule from '@/lib/trpc/min-version'; +import { db } from '@/lib/drizzle'; +import { insertTestUser } from '@/tests/helpers/user.helper'; +import type * as SessionOwnership from '@/lib/cloud-agent/session-ownership'; +import type { + GetWorktreeChangesOutput, + RefreshWorktreeChangesOutput, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; import type { z } from 'zod'; import type { personalPrepareSessionNextSchema, @@ -65,6 +72,10 @@ const mockCancelQueuedMessage = jest.fn<(input: { sessionId: string; messageId: string }) => Promise<{ dropped: boolean }>>(); const mockGetSandboxStatus = jest.fn<(cloudAgentSessionId: string) => Promise>(); +const mockGetWorktreeChanges = + jest.fn<(cloudAgentSessionId: string) => Promise>(); +const mockRefreshWorktreeChanges = + jest.fn<(cloudAgentSessionId: string) => Promise>(); const mockCreateCloudAgentNextClient = jest.fn((_authToken: string) => ({ prepareSession: mockPrepareSession, @@ -72,6 +83,8 @@ const mockCreateCloudAgentNextClient = jest.fn((_authToken: string) => ({ getSession: mockGetSession, cancelQueuedMessage: mockCancelQueuedMessage, getSandboxStatus: mockGetSandboxStatus, + getWorktreeChanges: mockGetWorktreeChanges, + refreshWorktreeChanges: mockRefreshWorktreeChanges, })); const mockCreateCloudAgentNextClientForModel = jest.fn( @@ -209,6 +222,10 @@ let createCaller: (ctx: { user: User; headersList?: Headers }) => { getAttachmentDownloadUrl: (input: { messageUuid: string; filename: string }) => Promise; cancelQueuedMessage: (input: { sessionId: string; messageId: string }) => Promise; getSandboxStatus: (input: { cloudAgentSessionId: string }) => Promise; + getWorktreeChanges: (input: { cloudAgentSessionId: string }) => Promise; + refreshWorktreeChanges: (input: { + cloudAgentSessionId: string; + }) => Promise; checkEligibility: () => Promise<{ balance: number; minBalance: number; @@ -225,6 +242,92 @@ beforeAll(async () => { createCaller = createCallerFactory(mod.cloudAgentNextRouter); }); +describe('cloudAgentNextRouter worktree changes access', () => { + const personalSessionId = 'workspace_12345678-1234-4234-9234-123456789abc'; + const orgSessionId = 'workspace_12345678-1234-4234-9234-123456789abd'; + let owner: User; + let otherUser: User; + + beforeAll(async () => { + owner = await insertTestUser({ id: 'oauth/worktree-personal-owner' }); + otherUser = await insertTestUser(); + const [organization] = await db + .insert(organizations) + .values({ + name: 'Personal changes scope test', + created_by_kilo_user_id: owner.id, + }) + .returning(); + await db.insert(cli_sessions_v2).values([ + { + session_id: 'ses_changes_personal', + cloud_agent_session_id: personalSessionId, + kilo_user_id: owner.id, + created_on_platform: 'cloud-agent-web', + }, + { + session_id: 'ses_changes_personal_org', + cloud_agent_session_id: orgSessionId, + organization_id: organization.id, + kilo_user_id: owner.id, + created_on_platform: 'cloud-agent-web', + }, + ]); + }); + + beforeEach(() => { + jest.clearAllMocks(); + mockVerifyUserOwnsSessionV2ByCloudAgentId.mockImplementation( + jest.requireActual('@/lib/cloud-agent/session-ownership') + .verifyUserOwnsSessionV2ByCloudAgentId + ); + mockGetWorktreeChanges.mockResolvedValue({ snapshot: null }); + mockRefreshWorktreeChanges.mockResolvedValue({ status: 'offline', snapshot: null }); + }); + + describe.each(['getWorktreeChanges', 'refreshWorktreeChanges'] as const)('%s', procedure => { + it('allows the creator in personal scope without model or rollout gates', async () => { + const result = await createCaller({ user: owner })[procedure]({ + cloudAgentSessionId: personalSessionId, + }); + expect(result.snapshot).toBeNull(); + expect( + procedure === 'getWorktreeChanges' ? mockGetWorktreeChanges : mockRefreshWorktreeChanges + ).toHaveBeenCalledWith(personalSessionId); + expect(mockComputeCloudAgentNextBalanceCheckEligibility).not.toHaveBeenCalled(); + expect(mockGetBalanceForUser).not.toHaveBeenCalled(); + expect(mockIsFeatureFlagEnabledOrDevelopment).not.toHaveBeenCalled(); + }); + + it('denies another creator before constructing a Worker client', async () => { + await expect( + createCaller({ user: otherUser })[procedure]({ cloudAgentSessionId: personalSessionId }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + expect(mockGetWorktreeChanges).not.toHaveBeenCalled(); + expect(mockRefreshWorktreeChanges).not.toHaveBeenCalled(); + }); + + it('denies the same creator accessing an organization session through personal scope', async () => { + await expect( + createCaller({ user: owner })[procedure]({ cloudAgentSessionId: orgSessionId }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + it.each(['agent_12345678-1234-4234-9234-123456789abc', 'ses_12345678901234567890123456'])( + 'rejects legacy ID %s before ownership or Worker calls', + async cloudAgentSessionId => { + await expect( + createCaller({ user: owner })[procedure]({ cloudAgentSessionId }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(mockVerifyUserOwnsSessionV2ByCloudAgentId).not.toHaveBeenCalled(); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + } + ); + }); +}); + describe('cloudAgentNextRouter attachment forwarding', () => { beforeEach(() => { jest.clearAllMocks(); diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index 13b4147f8b..2d2039f209 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -31,6 +31,7 @@ import { baseGetSessionNextOutputSchema, baseGetSandboxStatusNextSchema, baseGetSandboxStatusNextOutputSchema, + baseWorktreeChangesNextSchema, baseAnswerQuestionNextSchema, baseRejectQuestionNextSchema, baseAnswerPermissionNextSchema, @@ -64,6 +65,10 @@ import { generateMessageId } from '@kilocode/cloud-agent-sdk/message-id'; import { getBalanceForUser } from '@/lib/user/balance'; import { isMobileClient } from '@/lib/trpc/min-version'; import { buildCloudAgentNextEligibility } from './cloud-agent-next-eligibility'; +import { + getWorktreeChangesOutputSchema, + refreshWorktreeChangesOutputSchema, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; function buildTerminalUrl(params: { cloudAgentSessionId: string; @@ -318,6 +323,24 @@ export const cloudAgentNextRouter = createTRPCRouter({ } }), + getWorktreeChanges: baseProcedure + .input(baseWorktreeChangesNextSchema) + .output(getWorktreeChangesOutputSchema) + .query(async ({ ctx, input }) => { + await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); + const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + return await client.getWorktreeChanges(input.cloudAgentSessionId); + }), + + refreshWorktreeChanges: baseProcedure + .input(baseWorktreeChangesNextSchema) + .output(refreshWorktreeChangesOutputSchema) + .mutation(async ({ ctx, input }) => { + await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); + const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + return await client.refreshWorktreeChanges(input.cloudAgentSessionId); + }), + createTerminal: baseProcedure .input(baseCreateTerminalNextSchema) .output(baseCreateTerminalNextOutputSchema) diff --git a/apps/web/src/routers/cloud-agent-next-schemas.test.ts b/apps/web/src/routers/cloud-agent-next-schemas.test.ts index 0724599988..c8811a6cbb 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.test.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.test.ts @@ -8,6 +8,7 @@ import { baseCancelQueuedMessageNextSchema, SANDBOX_STATUS_DETAIL_MESSAGES, type SandboxStatusSnapshot, + baseWorktreeChangesNextSchema, cloudAgentGetAttachmentDownloadUrlSchema, cloudAgentGetAttachmentUploadUrlSchema, cloudAgentRelaxedAttachmentFilenameSchema, @@ -214,6 +215,33 @@ describe('baseGetSandboxStatusNextSchema', () => { }); }); +describe('baseWorktreeChangesNextSchema', () => { + const cloudAgentSessionId = `workspace_${MESSAGE_UUID}`; + + it('accepts only a control-plane session ID', () => { + expect(baseWorktreeChangesNextSchema.parse({ cloudAgentSessionId })).toEqual({ + cloudAgentSessionId, + }); + }); + + it.each([`agent_${MESSAGE_UUID}`, KILO_SESSION_ID, '', 'workspace_', 'workspace_not-a-uuid'])( + 'rejects legacy or malformed session ID %s', + cloudAgentSessionId => { + expect(baseWorktreeChangesNextSchema.safeParse({ cloudAgentSessionId }).success).toBe(false); + } + ); + + it.each(['directory', 'baseRef', 'sandboxId', 'revision'])( + 'rejects client-controlled %s', + field => { + expect( + baseWorktreeChangesNextSchema.safeParse({ cloudAgentSessionId, [field]: 'override' }) + .success + ).toBe(false); + } + ); +}); + describe('cloudAgentGetAttachmentUploadUrlSchema', () => { it('preserves the legacy 9-MIME contract when extension is absent', () => { const result = cloudAgentGetAttachmentUploadUrlSchema.safeParse({ diff --git a/apps/web/src/routers/cloud-agent-next-schemas.ts b/apps/web/src/routers/cloud-agent-next-schemas.ts index b3b06e577a..3cec74b6b3 100644 --- a/apps/web/src/routers/cloud-agent-next-schemas.ts +++ b/apps/web/src/routers/cloud-agent-next-schemas.ts @@ -566,6 +566,17 @@ export const baseGetSandboxStatusNextSchema = z }) .strict(); +export const baseWorktreeChangesNextSchema = z + .object({ + cloudAgentSessionId: z + .string() + .regex( + /^workspace_[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$/, + 'Changes require a control-plane session' + ), + }) + .strict(); + export const cloudAgentTerminalSizeSchema = z.object({ cols: z.number().int().min(2).max(500), rows: z.number().int().min(2).max(200), diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts index 3f74353f4b..91ed51437c 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts @@ -1,7 +1,6 @@ import { describe, expect, it, jest, beforeAll, beforeEach } from '@jest/globals'; import { inspect } from 'node:util'; import { DrizzleQueryError } from 'drizzle-orm'; -import { db } from '@/lib/drizzle'; import type * as TrpcInitModule from '@/lib/trpc/init'; import type { createWorktreeChat as CreateWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; import type * as MinimumVersionModule from '@/lib/trpc/min-version'; @@ -9,7 +8,21 @@ import type * as OrganizationUtilsModule from '@/routers/organizations/utils'; import { fetchRequestHandler } from '@trpc/server/adapters/fetch'; import type * as ZodModule from 'zod'; import type { z } from 'zod'; -import type { User } from '@kilocode/db/schema'; +import { + cli_sessions_v2, + organization_memberships, + organizations, + type Organization, + type User, +} from '@kilocode/db/schema'; +import { db } from '@/lib/drizzle'; +import { insertTestUser } from '@/tests/helpers/user.helper'; +import { and, eq } from 'drizzle-orm'; +import type * as SessionOwnership from '@/lib/cloud-agent/session-ownership'; +import type { + GetWorktreeChangesOutput, + RefreshWorktreeChangesOutput, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; import type * as BitbucketIntegrationHelpers from '@/lib/cloud-agent/bitbucket-integration-helpers'; import type { BitbucketOrganizationRepositoryListResult } from '@/lib/cloud-agent/bitbucket-integration-helpers'; import { TRPCError } from '@trpc/server'; @@ -73,6 +86,10 @@ const mockCancelQueuedMessage = jest.fn<(input: { sessionId: string; messageId: string }) => Promise<{ dropped: boolean }>>(); const mockGetSandboxStatus = jest.fn<(cloudAgentSessionId: string) => Promise>(); +const mockGetWorktreeChanges = + jest.fn<(cloudAgentSessionId: string) => Promise>(); +const mockRefreshWorktreeChanges = + jest.fn<(cloudAgentSessionId: string) => Promise>(); const mockCreateCloudAgentNextClient = jest.fn((_authToken: string) => ({ prepareSession: mockPrepareSession, @@ -80,6 +97,8 @@ const mockCreateCloudAgentNextClient = jest.fn((_authToken: string) => ({ getSession: mockGetSession, cancelQueuedMessage: mockCancelQueuedMessage, getSandboxStatus: mockGetSandboxStatus, + getWorktreeChanges: mockGetWorktreeChanges, + refreshWorktreeChanges: mockRefreshWorktreeChanges, })); const mockCreateCloudAgentNextClientForModel = jest.fn( @@ -144,6 +163,7 @@ const mockOrderRepositoriesByUsage = >(); const mockEnsureOrganizationAccess = jest.fn(); +const mockRequireActiveSubscription = jest.fn<() => void>(); jest.mock('@/lib/tokens', () => ({ generateCloudAgentToken: jest.fn(() => 'cloud-agent-token'), @@ -221,7 +241,10 @@ jest.mock('@/routers/organizations/utils', () => { return { ensureOrganizationAccess: mockEnsureOrganizationAccess, organizationMemberProcedure: organizationProcedure, - organizationMemberMutationProcedure: organizationProcedure, + organizationMemberMutationProcedure: organizationProcedure.use(({ next }) => { + mockRequireActiveSubscription(); + return next(); + }), }; }); @@ -285,6 +308,14 @@ let createCaller: (ctx: { user: User; headersList?: Headers }) => { organizationId: string; forceRefresh: boolean; }) => Promise; + getWorktreeChanges: (input: { + organizationId: string; + cloudAgentSessionId: string; + }) => Promise; + refreshWorktreeChanges: (input: { + organizationId: string; + cloudAgentSessionId: string; + }) => Promise; refreshTerminalTicket: (input: { organizationId: string; cloudAgentSessionId: string; @@ -316,6 +347,7 @@ beforeAll(async () => { beforeEach(() => { mockEnsureOrganizationAccess.mockReset().mockResolvedValue('member'); + mockRequireActiveSubscription.mockReset(); }); describe('organizationCloudAgentNextRouter.getSandboxStatus', () => { @@ -526,6 +558,165 @@ describe('organizationCloudAgentNextRouter.getSandboxStatus', () => { }); }); +describe('organizationCloudAgentNextRouter worktree changes access', () => { + const orgSessionId = 'workspace_12345678-1234-4234-9234-123456789abc'; + const personalSessionId = 'workspace_12345678-1234-4234-9234-123456789abd'; + let owner: User; + let otherMember: User; + let organization: Organization; + let otherOrganization: Organization; + + beforeAll(async () => { + owner = await insertTestUser({ id: 'oauth/worktree-org-owner' }); + otherMember = await insertTestUser(); + [organization, otherOrganization] = await db + .insert(organizations) + .values([ + { name: 'Changes organization', created_by_kilo_user_id: owner.id }, + { name: 'Other changes organization', created_by_kilo_user_id: owner.id }, + ]) + .returning(); + await db.insert(organization_memberships).values([ + { organization_id: organization.id, kilo_user_id: owner.id, role: 'owner' }, + { organization_id: organization.id, kilo_user_id: otherMember.id, role: 'member' }, + { organization_id: otherOrganization.id, kilo_user_id: owner.id, role: 'owner' }, + ]); + await db.insert(cli_sessions_v2).values([ + { + session_id: 'ses_changes_org', + cloud_agent_session_id: orgSessionId, + organization_id: organization.id, + kilo_user_id: owner.id, + created_on_platform: 'cloud-agent-web', + }, + { + session_id: 'ses_changes_org_personal', + cloud_agent_session_id: personalSessionId, + kilo_user_id: owner.id, + created_on_platform: 'cloud-agent-web', + }, + ]); + }); + + beforeEach(async () => { + jest.clearAllMocks(); + mockVerifyOrgOwnsSessionV2ByCloudAgentId.mockImplementation( + jest.requireActual('@/lib/cloud-agent/session-ownership') + .verifyOrgOwnsSessionV2ByCloudAgentId + ); + mockGetWorktreeChanges.mockResolvedValue({ snapshot: null }); + mockRefreshWorktreeChanges.mockResolvedValue({ status: 'offline', snapshot: null }); + await db + .update(organizations) + .set({ deleted_at: null }) + .where(eq(organizations.id, organization.id)); + await db + .insert(organization_memberships) + .values({ + organization_id: organization.id, + kilo_user_id: owner.id, + role: 'owner', + }) + .onConflictDoNothing(); + }); + + describe.each(['getWorktreeChanges', 'refreshWorktreeChanges'] as const)('%s', procedure => { + it('allows the creator without an active subscription, model balance, or rollout gate', async () => { + mockRequireActiveSubscription.mockImplementation(() => { + throw new Error('Subscription inactive'); + }); + const result = await createCaller({ user: owner })[procedure]({ + organizationId: organization.id, + cloudAgentSessionId: orgSessionId, + }); + expect(result.snapshot).toBeNull(); + expect( + procedure === 'getWorktreeChanges' ? mockGetWorktreeChanges : mockRefreshWorktreeChanges + ).toHaveBeenCalledWith(orgSessionId); + expect(mockRequireActiveSubscription).not.toHaveBeenCalled(); + expect(mockComputeCloudAgentNextBalanceCheckEligibility).not.toHaveBeenCalled(); + expect(mockGetBalanceForOrganizationUser).not.toHaveBeenCalled(); + expect(mockIsFeatureFlagEnabledOrDevelopment).not.toHaveBeenCalled(); + }); + + it('denies another member who did not create the session before calling the Worker', async () => { + await expect( + createCaller({ user: otherMember })[procedure]({ + organizationId: organization.id, + cloudAgentSessionId: orgSessionId, + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + it('denies a different organization even when the creator belongs to both', async () => { + await expect( + createCaller({ user: owner })[procedure]({ + organizationId: otherOrganization.id, + cloudAgentSessionId: orgSessionId, + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + it('denies personal sessions through the organization endpoint', async () => { + await expect( + createCaller({ user: owner })[procedure]({ + organizationId: organization.id, + cloudAgentSessionId: personalSessionId, + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + it('denies removed members even when organization middleware grants access', async () => { + await db + .delete(organization_memberships) + .where( + and( + eq(organization_memberships.organization_id, organization.id), + eq(organization_memberships.kilo_user_id, owner.id) + ) + ); + await expect( + createCaller({ user: owner })[procedure]({ + organizationId: organization.id, + cloudAgentSessionId: orgSessionId, + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + it('denies deleted organizations before calling the Worker', async () => { + await db + .update(organizations) + .set({ deleted_at: new Date().toISOString() }) + .where(eq(organizations.id, organization.id)); + await expect( + createCaller({ user: owner })[procedure]({ + organizationId: organization.id, + cloudAgentSessionId: orgSessionId, + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + it.each(['agent_12345678-1234-4234-9234-123456789abc', 'ses_12345678901234567890123456'])( + 'rejects legacy ID %s before ownership or Worker calls', + async cloudAgentSessionId => { + await expect( + createCaller({ user: owner })[procedure]({ + organizationId: organization.id, + cloudAgentSessionId, + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + expect(mockVerifyOrgOwnsSessionV2ByCloudAgentId).not.toHaveBeenCalled(); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + } + ); + }); +}); + describe('organizationCloudAgentNextRouter attachment forwarding', () => { beforeEach(() => { jest.clearAllMocks(); diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index f8fe523c64..4592c9472e 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -40,6 +40,7 @@ import { baseGetSessionNextOutputSchema, baseGetSandboxStatusNextSchema, baseGetSandboxStatusNextOutputSchema, + baseWorktreeChangesNextSchema, baseAnswerQuestionNextSchema, baseRejectQuestionNextSchema, baseAnswerPermissionNextSchema, @@ -71,6 +72,10 @@ import { generateMessageId } from '@kilocode/cloud-agent-sdk/message-id'; import { getBalanceForOrganizationUser } from '@/lib/organizations/organization-usage'; import { isMobileClient } from '@/lib/trpc/min-version'; import { buildCloudAgentNextEligibility } from '../cloud-agent-next-eligibility'; +import { + getWorktreeChangesOutputSchema, + refreshWorktreeChangesOutputSchema, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; function buildTerminalUrl(params: { cloudAgentSessionId: string; @@ -181,6 +186,10 @@ const GetSandboxStatusInput = baseGetSandboxStatusNextSchema.extend({ organizationId: z.uuid(), }); +const WorktreeChangesInput = baseWorktreeChangesNextSchema.extend({ + organizationId: z.uuid(), +}); + const CreateTerminalInput = baseCreateTerminalNextSchema.extend({ organizationId: z.uuid(), }); @@ -458,6 +467,32 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ } }), + getWorktreeChanges: organizationMemberProcedure + .input(WorktreeChangesInput) + .output(getWorktreeChangesOutputSchema) + .query(async ({ ctx, input }) => { + await assertOrganizationOwnsSession({ + organizationId: input.organizationId, + userId: ctx.user.id, + cloudAgentSessionId: input.cloudAgentSessionId, + }); + const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + return await client.getWorktreeChanges(input.cloudAgentSessionId); + }), + + refreshWorktreeChanges: organizationMemberProcedure + .input(WorktreeChangesInput) + .output(refreshWorktreeChangesOutputSchema) + .mutation(async ({ ctx, input }) => { + await assertOrganizationOwnsSession({ + organizationId: input.organizationId, + userId: ctx.user.id, + cloudAgentSessionId: input.cloudAgentSessionId, + }); + const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + return await client.refreshWorktreeChanges(input.cloudAgentSessionId); + }), + createTerminal: organizationMemberMutationProcedure .input(CreateTerminalInput) .output(baseCreateTerminalNextOutputSchema) diff --git a/apps/web/tests/e2e/cloud-agent-sandbox-status.spec.ts b/apps/web/tests/e2e/cloud-agent-sandbox-status.spec.ts index 9b1b1a73e4..cdec9aff5e 100644 --- a/apps/web/tests/e2e/cloud-agent-sandbox-status.spec.ts +++ b/apps/web/tests/e2e/cloud-agent-sandbox-status.spec.ts @@ -5,6 +5,7 @@ import { organization_memberships, organizations } from '@kilocode/db/schema'; import { eq } from 'drizzle-orm'; import { z } from 'zod'; import type { SandboxStatusSnapshot } from '@/routers/cloud-agent-next-schemas'; +import type { WorktreeChangesSnapshot } from '@kilocode/worker-utils/cloud-agent-worktree-changes'; const firstId = 'ses_sandbox_status_first'; const secondId = 'ses_sandbox_status_second'; @@ -76,6 +77,7 @@ async function mountFixtures( const procedures: string[] = []; const sockets = new Map(); let reply: (request: StatusRequest) => RpcResult | Promise = () => success(snapshot()); + let savedChanges: WorktreeChangesSnapshot | null = null; let eventId = 0; function snapshot(overrides: Partial = {}): SandboxStatusSnapshot { @@ -204,6 +206,12 @@ async function mountFixtures( }); if (procedure === 'cliSessionsV2.getSessionMessages') return success({ info: { id: args.session_id }, messages: [] }); + if (procedure.endsWith('.getWorktreeChanges')) return success({ snapshot: savedChanges }); + if (procedure.endsWith('.refreshWorktreeChanges')) + return success({ + status: savedChanges ? 'refreshed' : 'offline', + snapshot: savedChanges, + }); if (procedure.endsWith('.getComputeBillingStatus')) return success({ phase: 'unavailable' }); if (procedure.endsWith('.sendMessage')) @@ -231,6 +239,9 @@ async function mountFixtures( setReply(handler: typeof reply) { reply = handler; }, + setWorktreeChanges(snapshot: WorktreeChangesSnapshot) { + savedChanges = snapshot; + }, async open(id = firstId, organizationId?: string) { await page.goto( `${organizationId ? `/organizations/${organizationId}` : ''}/cloud/chat?sessionId=${id}` @@ -416,6 +427,167 @@ test.describe('control-plane sandbox header', () => { } }); + test('keeps status and changes controls distinct across workspace navigation', async ({ + page, + }) => { + const duplicateKeyErrors: string[] = []; + page.on('console', message => { + if (message.text().includes('Encountered two children with the same key')) { + duplicateKeyErrors.push(message.text()); + } + }); + const fixture = await mountFixtures(page); + fixture.setReply(request => + success( + fixture.snapshot( + request.cloudAgentSessionId === secondWorkspace + ? { status: 'sleeping', detailCode: 'sandbox_stopped', estimatedSleepAt: null } + : {} + ) + ) + ); + await fixture.open(); + for (const [sessionId, status] of [ + [firstId, 'Active'], + [secondId, 'Sleeping'], + [firstId, 'Active'], + ]) { + await fixture.navigate(sessionId); + await expect(indicator(page)).toHaveAccessibleName(`Sandbox status: ${status}`); + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + await expect(changes).toHaveCount(1); + await expect(changes).toBeVisible(); + expect(duplicateKeyErrors).toEqual([]); + } + }); + + test('updates saved file changes during a turn and while the drawer is open without polling captures', async ({ + page, + }) => { + const fixture = await mountFixtures(page); + const readCount = () => + fixture.procedures.filter(procedure => procedure.endsWith('.getWorktreeChanges')).length; + const captureCount = () => + fixture.procedures.filter(procedure => procedure.endsWith('.refreshWorktreeChanges')).length; + function publishRevision(revision: number) { + fixture.setWorktreeChanges({ + schemaVersion: 1, + revision, + capturedAt: new Date(baseTime + revision).toISOString(), + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: [ + { + path: `src/revision-${revision}.ts`, + status: 'modified', + additions: revision, + deletions: 0, + tracked: true, + binary: false, + countsComplete: true, + }, + ], + truncated: false, + }); + } + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + const drawer = page.getByRole('dialog', { name: 'Changes', exact: true }); + const summary = (revision: number) => `1 changed files, ${revision} additions, 0 deletions`; + publishRevision(1); + await fixture.open(); + await expect(changes).toHaveAttribute('aria-description', summary(1)); + const initialReads = readCount(); + await fixture.advance(15_000); + expect(readCount()).toBe(initialReads); + + await fixture.activity(firstWorkspace, 'busy'); + await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toBeVisible(); + publishRevision(2); + await fixture.advance(5_000); + await expect(changes).toHaveAttribute('aria-description', summary(2)); + expect(readCount()).toBeGreaterThan(initialReads); + expect(captureCount()).toBe(0); + + await fixture.activity(firstWorkspace, 'idle'); + await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toHaveCount(0); + publishRevision(3); + await fixture.advance(31_000); + await expect(changes).toHaveAttribute('aria-description', summary(3)); + const idleReads = readCount(); + await fixture.advance(15_000); + expect(readCount()).toBe(idleReads); + + await changes.click(); + await expect(drawer.getByText('revision-3.ts', { exact: true })).toBeVisible(); + await expect.poll(captureCount).toBe(1); + publishRevision(4); + await fixture.advance(5_000); + await expect(drawer.getByText('revision-4.ts', { exact: true })).toBeVisible(); + await expect(changes).toHaveAttribute('aria-description', summary(4)); + expect(captureCount()).toBe(1); + await page.keyboard.press('Escape'); + await expect(drawer).toHaveCount(0); + const closedReads = readCount(); + await fixture.advance(15_000); + expect(readCount()).toBe(closedReads); + expect(captureCount()).toBe(1); + }); + + for (const { width, reducedMotion } of [ + { width: 375, reducedMotion: 'no-preference' }, + { width: 820, reducedMotion: 'no-preference' }, + { width: 1440, reducedMotion: 'no-preference' }, + { width: 1440, reducedMotion: 'reduce' }, + ] as const) { + test(`opens Changes without scrolling the chat at ${width}px with ${reducedMotion} motion`, async ({ + page, + }) => { + await page.setViewportSize({ width, height: 1000 }); + await page.emulateMedia({ reducedMotion }); + const fixture = await mountFixtures(page); + await fixture.open(); + await expect(indicator(page)).toHaveAccessibleName('Sandbox status: Active'); + await page.addStyleTag({ + content: + '#worktree-changes-panel[data-state="open"] { animation-play-state: paused !important; }', + }); + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + const drawer = page.getByRole('dialog', { name: 'Changes', exact: true }); + for (const layout of ['Flat', 'Tree']) { + await changes.click(); + await expect(drawer.getByRole('tab', { name: layout, exact: true })).toBeFocused(); + const opening = await drawer.evaluate(element => { + const scrollOffsets: number[] = []; + for (let parent = element.parentElement; parent; parent = parent.parentElement) { + if (parent.scrollLeft !== 0) scrollOffsets.push(parent.scrollLeft); + } + return { + left: element.getBoundingClientRect().left, + width: element.getBoundingClientRect().width, + scrollOffsets, + }; + }); + expect(opening.scrollOffsets).toEqual([]); + await drawer.evaluate(element => { + for (const animation of element.getAnimations()) animation.finish(); + }); + const openedLeft = await drawer.evaluate(element => element.getBoundingClientRect().left); + expect(opening.left - openedLeft).toBeCloseTo( + reducedMotion === 'reduce' ? 0 : opening.width, + 0 + ); + await expect(drawer).toBeInViewport(); + await drawer.getByRole('tab', { name: 'Tree', exact: true }).click(); + await page.keyboard.press('Escape'); + await expect(drawer).toHaveCount(0); + await expect(changes).toBeFocused(); + } + }); + } + test('renders distinct static lifecycle icons without button text independently of agent progress', async ({ page, }) => { diff --git a/packages/worker-utils/package.json b/packages/worker-utils/package.json index 8b7a77b324..4463ca04d2 100644 --- a/packages/worker-utils/package.json +++ b/packages/worker-utils/package.json @@ -25,6 +25,7 @@ "./cf-access": "./src/cf-access.ts", "./cloud-agent-next-client": "./src/cloud-agent-next-client.ts", "./cloud-agent-session-access": "./src/cloud-agent-session-access.ts", + "./cloud-agent-worktree-changes": "./src/cloud-agent-worktree-changes.ts", "./kilo-model-id": "./src/kilo-model-id.ts", "./extract-bearer-token": "./src/extract-bearer-token.ts", "./cloud-agent-queue-report": "./src/cloud-agent-queue-report.ts", diff --git a/packages/worker-utils/src/cloud-agent-worktree-changes.test.ts b/packages/worker-utils/src/cloud-agent-worktree-changes.test.ts new file mode 100644 index 0000000000..aa893dec9a --- /dev/null +++ b/packages/worker-utils/src/cloud-agent-worktree-changes.test.ts @@ -0,0 +1,149 @@ +import { describe, expect, it } from 'vitest'; +import { + MAX_WORKTREE_CHANGES_BYTES, + MAX_WORKTREE_CHANGES_FILES, + getWorktreeChangesOutputSchema, + refreshWorktreeChangesOutputSchema, + worktreeChangesCaptureRequestSchema, + worktreeChangesCaptureSchema, + worktreeChangesFileSchema, + worktreeChangesSnapshotSchema, + type WorktreeChangesFile, + type WorktreeChangesSnapshot, +} from './cloud-agent-worktree-changes.js'; + +const file: WorktreeChangesFile = { + path: 'src/example.ts', + status: 'modified', + additions: 2, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, +}; +const snapshot: WorktreeChangesSnapshot = { + schemaVersion: 1, + revision: 3, + capturedAt: '2026-08-26T12:00:00.000Z', + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: [file], + truncated: false, +}; + +function capture(files = snapshot.files) { + return { revision: snapshot.revision, comparison: snapshot.comparison, files, truncated: false }; +} + +describe('cloud agent worktree changes contracts', () => { + it('distinguishes no saved snapshot from a successful empty capture', () => { + expect(getWorktreeChangesOutputSchema.parse({ snapshot: null })).toEqual({ snapshot: null }); + const empty = { ...snapshot, files: [] }; + expect( + refreshWorktreeChangesOutputSchema.parse({ status: 'refreshed', snapshot: empty }) + ).toEqual({ + status: 'refreshed', + snapshot: empty, + }); + expect( + refreshWorktreeChangesOutputSchema.safeParse({ status: 'refreshed', snapshot: null }).success + ).toBe(false); + for (const status of ['offline', 'failed']) { + expect(refreshWorktreeChangesOutputSchema.parse({ status, snapshot })).toEqual({ + status, + snapshot, + }); + expect(refreshWorktreeChangesOutputSchema.parse({ status, snapshot: null })).toEqual({ + status, + snapshot: null, + }); + } + }); + + it.each([ + { schemaVersion: 2 }, + { revision: 0 }, + { revision: Number.MAX_SAFE_INTEGER + 1 }, + { capturedAt: 'not a timestamp' }, + { comparison: { ...snapshot.comparison, head: 'HEAD' } }, + { comparison: { ...snapshot.comparison, mergeBase: '' } }, + { patch: 'unexpected file contents' }, + ])('rejects unsupported or malformed persisted records %j', invalid => { + expect(worktreeChangesSnapshotSchema.safeParse({ ...snapshot, ...invalid }).success).toBe( + false + ); + }); + + it('accepts SHA-256 object IDs without weakening commit validation', () => { + expect( + worktreeChangesSnapshotSchema.safeParse({ + ...snapshot, + comparison: { ...snapshot.comparison, head: 'b'.repeat(64), mergeBase: 'a'.repeat(64) }, + }).success + ).toBe(true); + }); + + it.each([ + '/outside', + '../outside', + 'parent/../outside', + 'parent//file', + './file', + 'file\0suffix', + ])('rejects unsafe paths %j', path => { + expect(worktreeChangesFileSchema.safeParse({ ...file, path }).success).toBe(false); + }); + + it('preserves ordinary unusual filenames exactly', () => { + const path = 'parent/ leading\tline\n"back\\slash-é-漢 '; + expect(worktreeChangesFileSchema.parse({ ...file, path }).path).toBe(path); + }); + + it.each([ + { additions: -1 }, + { deletions: 1.5 }, + { additions: Number.POSITIVE_INFINITY }, + { status: 'renamed' }, + { contents: 'not a file summary' }, + ])('rejects invalid file summaries %j', invalid => { + expect(worktreeChangesFileSchema.safeParse({ ...file, ...invalid }).success).toBe(false); + }); + + it('bounds both capture and stored snapshot by UTF-8 bytes rather than string length', () => { + const files = Array.from({ length: 24 }, (_, index) => ({ + ...file, + path: `${index}/${'漢'.repeat(3800)}`, + })); + expect(JSON.stringify(capture(files)).length).toBeLessThan(MAX_WORKTREE_CHANGES_BYTES); + expect(worktreeChangesCaptureSchema.safeParse(capture(files)).success).toBe(false); + expect(worktreeChangesSnapshotSchema.safeParse({ ...snapshot, files }).success).toBe(false); + expect(worktreeChangesCaptureSchema.safeParse(capture(files.slice(0, 1))).success).toBe(true); + }); + + it('bounds file count independently of bytes', () => { + const files = Array.from({ length: MAX_WORKTREE_CHANGES_FILES + 1 }, (_, index) => ({ + ...file, + path: `${index}`, + })); + expect(new TextEncoder().encode(JSON.stringify(capture(files))).byteLength).toBeLessThan( + MAX_WORKTREE_CHANGES_BYTES + ); + expect(worktreeChangesCaptureSchema.safeParse(capture(files)).success).toBe(false); + expect(worktreeChangesSnapshotSchema.safeParse({ ...snapshot, files }).success).toBe(false); + }); + + it('rejects option-like refs and extra routing fields at the capture boundary', () => { + expect(worktreeChangesCaptureRequestSchema.parse({ revision: 1 })).toEqual({ revision: 1 }); + for (const invalid of [ + { revision: 1, baseRef: '--help' }, + { revision: 1, baseRef: 'main\0suffix' }, + { revision: 1, directory: '/another-session' }, + { revision: 1, sessionId: 'another-session' }, + ]) { + expect(worktreeChangesCaptureRequestSchema.safeParse(invalid).success).toBe(false); + } + }); +}); diff --git a/packages/worker-utils/src/cloud-agent-worktree-changes.ts b/packages/worker-utils/src/cloud-agent-worktree-changes.ts new file mode 100644 index 0000000000..6fdfde3584 --- /dev/null +++ b/packages/worker-utils/src/cloud-agent-worktree-changes.ts @@ -0,0 +1,97 @@ +import { z } from 'zod'; + +export const WORKTREE_CHANGES_SCHEMA_VERSION = 1; +export const MAX_WORKTREE_CHANGES_FILES = 1_000; +export const MAX_WORKTREE_CHANGES_BYTES = 256 * 1024; + +const revisionSchema = z.number().int().positive().max(Number.MAX_SAFE_INTEGER); +const commitSchema = z.string().regex(/^(?:[0-9a-f]{40}|[0-9a-f]{64})$/); +const baseRefSchema = z + .string() + .min(1) + .max(1024) + .refine(ref => !ref.startsWith('-') && !ref.includes('\0'), 'Invalid comparison ref'); + +export const worktreeChangesFileSchema = z + .object({ + path: z + .string() + .min(1) + .max(4096) + .refine( + path => + !path.includes('\0') && + path.split('/').every(part => part !== '' && part !== '.' && part !== '..'), + 'Expected a repository-relative path' + ), + status: z.enum(['added', 'modified', 'deleted']), + additions: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + deletions: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + tracked: z.boolean(), + binary: z.boolean(), + countsComplete: z.boolean(), + }) + .strict(); + +export const worktreeChangesCaptureRequestSchema = z + .object({ + revision: revisionSchema, + baseRef: baseRefSchema.optional(), + }) + .strict(); + +const captureFields = { + revision: revisionSchema, + comparison: z + .object({ + baseRef: baseRefSchema, + mergeBase: commitSchema, + head: commitSchema, + }) + .strict(), + files: z.array(worktreeChangesFileSchema).max(MAX_WORKTREE_CHANGES_FILES), + truncated: z.boolean(), +}; + +export const worktreeChangesCaptureSchema = z + .object(captureFields) + .strict() + .refine( + capture => + new TextEncoder().encode(JSON.stringify(capture)).byteLength <= MAX_WORKTREE_CHANGES_BYTES, + 'Worktree summary exceeds the size limit' + ); + +export const worktreeChangesSnapshotSchema = z + .object({ + schemaVersion: z.literal(WORKTREE_CHANGES_SCHEMA_VERSION), + capturedAt: z.string().datetime({ offset: true }), + ...captureFields, + }) + .strict() + .refine( + snapshot => + new TextEncoder().encode(JSON.stringify(snapshot)).byteLength <= MAX_WORKTREE_CHANGES_BYTES, + 'Saved worktree summary exceeds the size limit' + ); + +export const getWorktreeChangesOutputSchema = z + .object({ snapshot: worktreeChangesSnapshotSchema.nullable() }) + .strict(); + +export const refreshWorktreeChangesOutputSchema = z.discriminatedUnion('status', [ + z.object({ status: z.literal('refreshed'), snapshot: worktreeChangesSnapshotSchema }).strict(), + z + .object({ + status: z.enum(['offline', 'failed']), + snapshot: worktreeChangesSnapshotSchema.nullable(), + }) + .strict(), +]); + +export type WorktreeChangesFile = z.infer; +export type WorktreeChangesCaptureRequest = z.infer; +export type WorktreeChangesCapture = z.infer; +export type WorktreeChangesSnapshot = z.infer; +export type GetWorktreeChangesOutput = z.infer; +export type RefreshWorktreeChangesOutput = z.infer; diff --git a/services/cloud-agent-next/src/persistence/SandboxControl.ts b/services/cloud-agent-next/src/persistence/SandboxControl.ts index faa34690ed..dd9a884948 100644 --- a/services/cloud-agent-next/src/persistence/SandboxControl.ts +++ b/services/cloud-agent-next/src/persistence/SandboxControl.ts @@ -34,7 +34,7 @@ import { type SandboxControlOutboundRequest, type SandboxControlSocketHandler, } from '../sandbox-control/socket.js'; -import { parseOperationPayload } from '../sandbox-control/frames.js'; +import { errorResponse, parseOperationPayload } from '../sandbox-control/frames.js'; import { generateSandboxCredential, hashSandboxCredential, @@ -508,6 +508,7 @@ export class SandboxControl extends DurableObject { async request(input: SandboxControlOutboundRequest): Promise { await this.ensureOperationalInitialized(); + if (input.operation === 'session.git.summary') return this.requestWorktreeChanges(input); await this.assertRequestWorktreeAdmission(input); if (input.operation === 'worktree.delete' || input.operation === 'worktree.prepareDeletion') { throw new Error('Worktree cleanup requires the deletion coordinator'); @@ -598,6 +599,75 @@ export class SandboxControl extends DurableObject { return this.socketHandler.sendRequest(input); } + private async requestWorktreeChanges( + input: SandboxControlOutboundRequest + ): Promise { + await this.assertRequestWorktreeAdmission(input); + const session = input.session; + const matchesRoute = (route: SessionRoute | undefined) => + session !== undefined && + route?.kiloSessionId === session.kiloSessionId && + route.directory === session.directory; + const physical = await loadPhysicalRecord(this.ctx.storage); + const routes = await loadRouteTable(this.ctx.storage); + const route = session ? routes.get(session.sessionId) : undefined; + const runtime = this.readyWrapperRuntime(); + const socket = this.socketHandler.getReadySocket(); + this.assertWorktreeAdmission(route?.worktreeId); + if ( + physical.state !== 'running' || + physical.stopTombstone || + !runtime || + physical.providerRef !== runtime.providerInstanceId || + !matchesRoute(route) || + !socket + ) { + return errorResponse(crypto.randomUUID(), 'not_ready', 'Worktree is not attached and ready'); + } + if ( + input.expectedWrapperInstanceId !== undefined && + wrapperInstanceIdSchema.parse(input.expectedWrapperInstanceId) !== runtime.wrapperInstanceId + ) { + return errorResponse( + crypto.randomUUID(), + 'protocol_error', + 'Worktree capture context changed' + ); + } + + let response: ResponseFrame; + try { + response = await this.socketHandler.sendRequest(input); + } catch { + return errorResponse(crypto.randomUUID(), 'protocol_error', 'Worktree capture failed'); + } + const currentPhysical = await loadPhysicalRecord(this.ctx.storage); + const currentRoutes = await loadRouteTable(this.ctx.storage); + const currentRoute = session ? currentRoutes.get(session.sessionId) : undefined; + await this.assertRequestWorktreeAdmission(input); + this.assertWorktreeAdmission(currentRoute?.worktreeId); + const currentRuntime = this.readyWrapperRuntime(); + if ( + currentPhysical.state !== 'running' || + currentPhysical.stopTombstone || + currentPhysical.providerRef !== physical.providerRef || + !sameAllocation(physical, currentPhysical) || + !currentRuntime || + !this.sameConnection(runtime, currentRuntime) || + !matchesRoute(currentRoute) || + currentRoute?.ownerId !== route?.ownerId || + currentRoute?.worktreeId !== route?.worktreeId || + this.socketHandler.getReadySocket() !== socket + ) { + return errorResponse( + response.requestId, + 'protocol_error', + 'Worktree capture context changed' + ); + } + return response; + } + async quarantineRuntime(input: { ownerId: string; sessionId: string; diff --git a/services/cloud-agent-next/src/router.ts b/services/cloud-agent-next/src/router.ts index 3036e0a29f..f0a512863d 100644 --- a/services/cloud-agent-next/src/router.ts +++ b/services/cloud-agent-next/src/router.ts @@ -14,6 +14,7 @@ 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 { createSessionWorktreeChangesHandlers } from './router/handlers/session-worktree-changes.js'; export const appRouter = router({ deleteWorktree, @@ -25,6 +26,7 @@ export const appRouter = router({ ...createSessionStartHandlers(), ...createSessionSendHandlers(), ...createSessionWorktreeHandlers(), + ...createSessionWorktreeChangesHandlers(), }); export type AppRouter = typeof appRouter; diff --git a/services/cloud-agent-next/src/router/handlers/session-worktree-changes.test.ts b/services/cloud-agent-next/src/router/handlers/session-worktree-changes.test.ts new file mode 100644 index 0000000000..61c14e2ffe --- /dev/null +++ b/services/cloud-agent-next/src/router/handlers/session-worktree-changes.test.ts @@ -0,0 +1,161 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { AccessibleCloudAgentSession } from '@kilocode/worker-utils/cloud-agent-session-access'; +import type { WorktreeChangesSnapshot } from '@kilocode/worker-utils/cloud-agent-worktree-changes'; +import type { TRPCContext } from '../../types.js'; + +const { queryAccess } = vi.hoisted(() => ({ queryAccess: vi.fn() })); + +vi.mock('@cloudflare/sandbox', () => ({ getSandbox: vi.fn() })); +vi.mock('@kilocode/db/client', () => ({ getWorkerDb: vi.fn(() => ({})) })); +vi.mock('@kilocode/worker-utils/cloud-agent-session-access', () => ({ + queryAccessibleCloudAgentSession: queryAccess, +})); + +import { appRouter } from '../../router.js'; + +const sessionId = 'workspace_12345678-1234-1234-1234-123456789abc'; +const legacySessionId = 'agent_12345678-1234-1234-1234-123456789abc'; +const access: AccessibleCloudAgentSession = { + kiloSessionId: 'kilo_root', + organizationId: 'org_current', +}; +const snapshot: WorktreeChangesSnapshot = { + schemaVersion: 1, + revision: 1, + capturedAt: '2026-08-20T10:00:00.000Z', + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: [], + truncated: false, +}; + +function setup(userId = 'user_owner') { + const stub = { + getWorktreeChanges: vi.fn().mockResolvedValue({ snapshot }), + refreshWorktreeChanges: vi.fn().mockResolvedValue({ status: 'offline', snapshot }), + }; + const session = { idFromName: vi.fn(name => name), get: vi.fn(() => stub) }; + const legacy = { idFromName: vi.fn(), get: vi.fn() }; + const control = { getByName: vi.fn() }; + const context = { + userId, + authToken: 'test-token', + request: new Request('https://worker.test/trpc'), + env: { + HYPERDRIVE: { connectionString: 'postgresql://test' }, + CONTROL_PLANE_IDS: '', + SANDBOX_SESSION: session, + CLOUD_AGENT_SESSION: legacy, + SANDBOX_CONTROL: control, + }, + } as unknown as TRPCContext; + return { stub, session, legacy, control, context, caller: appRouter.createCaller(context) }; +} + +describe('Worker worktree changes procedures', () => { + beforeEach(() => { + vi.clearAllMocks(); + queryAccess.mockResolvedValue(access); + }); + + it.each(['getWorktreeChanges', 'refreshWorktreeChanges'] as const)( + '%s authorizes existing control sessions without a creation allowlist or runtime work', + async procedure => { + const harness = setup(); + const result = await harness.caller[procedure]({ cloudAgentSessionId: sessionId }); + expect(result).toEqual( + procedure === 'getWorktreeChanges' ? { snapshot } : { status: 'offline', snapshot } + ); + expect(harness.session.idFromName).toHaveBeenCalledWith(`user_owner:${sessionId}`); + expect(harness.legacy.get).not.toHaveBeenCalled(); + expect(harness.control.getByName).not.toHaveBeenCalled(); + expect(queryAccess).toHaveBeenCalledWith(expect.anything(), { + kiloUserId: 'user_owner', + cloudAgentSessionId: sessionId, + }); + } + ); + + describe.each(['getWorktreeChanges', 'refreshWorktreeChanges'] as const)( + '%s access checks', + procedure => { + it('rejects an unauthenticated request before access lookup or DO routing', async () => { + const harness = setup(''); + await expect( + harness.caller[procedure]({ cloudAgentSessionId: sessionId }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + expect(queryAccess).not.toHaveBeenCalled(); + expect(harness.session.get).not.toHaveBeenCalled(); + }); + + it.each(['another owner', 'revoked organization membership', 'a deleted organization'])( + 'rejects current access denied for %s before any DO call', + async () => { + queryAccess.mockResolvedValue(null); + const harness = setup(); + await expect( + harness.caller[procedure]({ cloudAgentSessionId: sessionId }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(harness.session.idFromName).not.toHaveBeenCalled(); + expect(harness.session.get).not.toHaveBeenCalled(); + expect(harness.legacy.get).not.toHaveBeenCalled(); + expect(harness.control.getByName).not.toHaveBeenCalled(); + } + ); + + it('fails closed when the authoritative access lookup fails', async () => { + queryAccess.mockRejectedValueOnce(new Error('database unavailable')); + const harness = setup(); + await expect( + harness.caller[procedure]({ cloudAgentSessionId: sessionId }) + ).rejects.toMatchObject({ code: 'SERVICE_UNAVAILABLE' }); + expect(harness.session.get).not.toHaveBeenCalled(); + }); + + it('does not route to storage until the current access lookup resolves', async () => { + const started = Promise.withResolvers(); + const authorized = Promise.withResolvers(); + queryAccess.mockImplementationOnce(async () => { + started.resolve(); + return authorized.promise; + }); + const harness = setup(); + const result = harness.caller[procedure]({ cloudAgentSessionId: sessionId }); + await started.promise; + expect(harness.session.idFromName).not.toHaveBeenCalled(); + expect(harness.session.get).not.toHaveBeenCalled(); + authorized.resolve(access); + await result; + expect(harness.session.get).toHaveBeenCalledTimes(1); + }); + + it('rejects legacy sessions after authorization without calling either session DO', async () => { + const harness = setup(); + await expect( + harness.caller[procedure]({ cloudAgentSessionId: legacySessionId }) + ).rejects.toMatchObject({ code: 'PRECONDITION_FAILED' }); + expect(queryAccess).toHaveBeenCalledTimes(1); + expect(harness.session.get).not.toHaveBeenCalled(); + expect(harness.legacy.get).not.toHaveBeenCalled(); + queryAccess.mockResolvedValue(null); + await expect( + harness.caller[procedure]({ cloudAgentSessionId: legacySessionId }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + }); + + it('validates the session DO output before exposing cached data', async () => { + const harness = setup(); + harness.stub[procedure].mockResolvedValue({ + status: 'refreshed', + snapshot: { ...snapshot, schemaVersion: 2 }, + }); + await expect( + harness.caller[procedure]({ cloudAgentSessionId: sessionId }) + ).rejects.toMatchObject({ code: 'INTERNAL_SERVER_ERROR' }); + }); + } + ); +}); diff --git a/services/cloud-agent-next/src/router/handlers/session-worktree-changes.ts b/services/cloud-agent-next/src/router/handlers/session-worktree-changes.ts new file mode 100644 index 0000000000..704915684a --- /dev/null +++ b/services/cloud-agent-next/src/router/handlers/session-worktree-changes.ts @@ -0,0 +1,63 @@ +import { TRPCError } from '@trpc/server'; +import { withLogTags } from '../../logger.js'; +import { getSandboxSessionStub } from '../../sandbox-session/session-stub.js'; +import { requireCurrentSessionAccess } from '../../session-access.js'; +import { sessionPlaneFromId } from '../../session-plane.js'; +import { withDORetry } from '../../utils/do-retry.js'; +import { protectedProcedure } from '../auth.js'; +import { + GetWorktreeChangesOutput, + RefreshWorktreeChangesOutput, + WorktreeChangesInput, +} from '../schemas.js'; + +function requireControlSession(sessionId: string): void { + if (sessionPlaneFromId(sessionId) !== 'control') { + throw new TRPCError({ + code: 'PRECONDITION_FAILED', + message: 'Worktree changes are not available for this session', + }); + } +} + +export function createSessionWorktreeChangesHandlers() { + return { + getWorktreeChanges: protectedProcedure + .input(WorktreeChangesInput) + .output(GetWorktreeChangesOutput) + .query(({ input, ctx }) => + withLogTags({ source: 'getWorktreeChanges' }, async () => { + await requireCurrentSessionAccess({ + env: ctx.env, + kiloUserId: ctx.userId, + cloudAgentSessionId: input.cloudAgentSessionId, + }); + requireControlSession(input.cloudAgentSessionId); + return withDORetry( + () => getSandboxSessionStub(ctx.env, ctx.userId, input.cloudAgentSessionId), + session => session.getWorktreeChanges(), + 'getWorktreeChanges' + ); + }) + ), + + refreshWorktreeChanges: protectedProcedure + .input(WorktreeChangesInput) + .output(RefreshWorktreeChangesOutput) + .mutation(({ input, ctx }) => + withLogTags({ source: 'refreshWorktreeChanges' }, async () => { + await requireCurrentSessionAccess({ + env: ctx.env, + kiloUserId: ctx.userId, + cloudAgentSessionId: input.cloudAgentSessionId, + }); + requireControlSession(input.cloudAgentSessionId); + return withDORetry( + () => getSandboxSessionStub(ctx.env, ctx.userId, input.cloudAgentSessionId), + async session => await session.refreshWorktreeChanges(), + 'refreshWorktreeChanges' + ); + }) + ), + }; +} diff --git a/services/cloud-agent-next/src/router/schemas.ts b/services/cloud-agent-next/src/router/schemas.ts index 6684d465bd..35eb78c7db 100644 --- a/services/cloud-agent-next/src/router/schemas.ts +++ b/services/cloud-agent-next/src/router/schemas.ts @@ -1009,6 +1009,13 @@ export const GetSandboxStatusInput = z }) .strict(); +export const WorktreeChangesInput = z.object({ cloudAgentSessionId: sessionIdSchema }).strict(); + +export { + getWorktreeChangesOutputSchema as GetWorktreeChangesOutput, + refreshWorktreeChangesOutputSchema as RefreshWorktreeChangesOutput, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; + /** Customer-safe, no-wake compute billing status for an existing session. */ export const GetComputeBillingStatusOutput = z.object({ payer: z.object({ type: z.enum(['user', 'org']), id: z.string() }), diff --git a/services/cloud-agent-next/src/sandbox-control/frames.test.ts b/services/cloud-agent-next/src/sandbox-control/frames.test.ts index 396ce8325c..d36bf30aaa 100644 --- a/services/cloud-agent-next/src/sandbox-control/frames.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/frames.test.ts @@ -90,6 +90,8 @@ describe('sandbox control frames', () => { it('recognizes known operations', () => { expect(isControlOperation('sandbox.hello')).toBe(true); expect(isControlOperation('session.prompt')).toBe(true); + expect(isControlOperation('session.git.summary')).toBe(true); + expect(isSessionOperation('session.git.summary')).toBe(true); expect(isControlOperation('http.tunnel')).toBe(false); for (const operation of [ 'session.terminal.create', @@ -253,6 +255,32 @@ describe('sandbox control frames', () => { ).toBe(false); }); + it('validates summary requests without accepting a payload directory', () => { + expect(parseOperationPayload('session.git.summary', { revision: 1 })).toEqual({ + ok: true, + payload: { revision: 1 }, + }); + expect( + parseOperationPayload('session.git.summary', { + revision: 2, + baseRef: 'refs/remotes/origin/main', + }).ok + ).toBe(true); + for (const payload of [ + {}, + { revision: 0 }, + { revision: 1.5 }, + { revision: Number.MAX_SAFE_INTEGER + 1 }, + { revision: 1, baseRef: '--help' }, + { revision: 1, directory: '/outside' }, + ]) { + expect(parseOperationPayload('session.git.summary', payload)).toEqual({ + ok: false, + error: { code: 'protocol_error', message: 'Invalid session.git.summary payload' }, + }); + } + }); + it('accepts a full sandbox.heartbeat payload', () => { expect( parseEventPayload('sandbox.heartbeat', { diff --git a/services/cloud-agent-next/src/sandbox-control/frames.ts b/services/cloud-agent-next/src/sandbox-control/frames.ts index 1669698b48..f4f783719a 100644 --- a/services/cloud-agent-next/src/sandbox-control/frames.ts +++ b/services/cloud-agent-next/src/sandbox-control/frames.ts @@ -14,6 +14,7 @@ import { sessionAttachPayloadSchema, sessionDetachPayloadSchema, sessionEventPayloadSchema, + sessionGitSummaryPayloadSchema, sessionPreparingPayloadSchema, sessionPermissionResolvePayloadSchema, sessionPromptPayloadSchema, @@ -50,6 +51,7 @@ const REQUEST_PAYLOAD_SCHEMAS: Record = { 'session.question.resolve': sessionQuestionResolvePayloadSchema, 'session.abort': sessionAbortPayloadSchema, 'session.sync': sessionSyncPayloadSchema, + 'session.git.summary': sessionGitSummaryPayloadSchema, 'session.detach': sessionDetachPayloadSchema, 'session.terminal.create': sessionTerminalCreatePayloadSchema, 'session.terminal.resize': sessionTerminalResizePayloadSchema, diff --git a/services/cloud-agent-next/src/sandbox-control/socket.test.ts b/services/cloud-agent-next/src/sandbox-control/socket.test.ts index 1b887565de..1b2fd782ae 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.test.ts @@ -212,7 +212,10 @@ describe('sandbox control socket handler', () => { ); expect(validateHandshake).toHaveBeenCalledWith('inst_stale'); - expect(incoming.serializeAttachment).not.toHaveBeenCalled(); + expect(incoming.deserializeAttachment()).toMatchObject({ + handshakeComplete: false, + kiloReady: false, + }); expect(incoming.close).toHaveBeenCalledWith(1008, 'invalid_provider_instance'); expect(incoming.send).toHaveBeenCalledWith( JSON.stringify({ @@ -304,6 +307,7 @@ describe('sandbox control socket handler', () => { expect(incoming.serializeAttachment).toHaveBeenCalledWith({ handshakeComplete: true, + kiloReady: false, acceptedAt: expect.any(Number), connectionId: expect.any(String), protocolVersion: 1, @@ -491,6 +495,7 @@ describe('sandbox control socket handler', () => { expect(superseded.deserializeAttachment()).toEqual({ handshakeComplete: false, + kiloReady: false, acceptedAt: expect.any(Number), }); expect(superseded.close).toHaveBeenCalledWith(1008, 'handshake_required'); @@ -1261,6 +1266,177 @@ describe('sandbox control socket handler', () => { ).rejects.toThrow('session identity is required'); expect(ws.send).not.toHaveBeenCalled(); }); + + it('requires readiness on the selected socket only for worktree captures', async () => { + const ws = createFakeWebSocket({ + handshakeComplete: true, + acceptedAt: Date.now(), + providerInstanceId: 'inst_1', + }); + const handler = createSandboxControlSocketHandler(createFakeState([ws]), 'sbx_test'); + const capture = { + operation: 'session.git.summary' as const, + session: { sessionId: 'workspace_1', kiloSessionId: 'kilo_1', directory: '/workspace' }, + payload: { revision: 1 }, + }; + await expect(handler.sendRequest(capture)).resolves.toMatchObject({ + ok: false, + error: { code: 'not_ready' }, + }); + expect(ws.send).not.toHaveBeenCalled(); + expect(handler.getReadySocket()).toBeNull(); + + await handler.handleMessage( + asWs(ws), + JSON.stringify({ + type: 'event', + event: 'sandbox.ready', + payload: { kiloReady: true, globalFeedAttached: true }, + }) + ); + expect(handler.getReadySocket()).toBe(ws); + const pending = handler.sendRequest(capture); + const sent = JSON.parse(ws.send.mock.calls[0]?.[0] as string) as { requestId: string }; + await handler.handleMessage( + asWs(ws), + JSON.stringify({ type: 'response', requestId: sent.requestId, ok: true }) + ); + await expect(pending).resolves.toMatchObject({ ok: true }); + + await handler.handleMessage( + asWs(ws), + JSON.stringify({ + type: 'event', + event: 'sandbox.heartbeat', + payload: { state: 'idle', kilo: { ready: false }, sessions: [] }, + }) + ); + expect(handler.getReadySocket()).toBeNull(); + await expect(handler.sendRequest(capture)).resolves.toMatchObject({ + ok: false, + error: { code: 'not_ready' }, + }); + const status = handler.sendRequest({ operation: 'sandbox.status', payload: {} }); + const statusSent = JSON.parse(ws.send.mock.calls[1]?.[0] as string) as { requestId: string }; + await handler.handleMessage( + asWs(ws), + JSON.stringify({ type: 'response', requestId: statusSent.requestId, ok: true }) + ); + await expect(status).resolves.toMatchObject({ ok: true }); + }); + + it('does not accept a response from a provisional or replaced socket', async () => { + const current = createFakeWebSocket({ + handshakeComplete: true, + kiloReady: true, + acceptedAt: 1, + providerInstanceId: 'inst_1', + }); + const provisional = createFakeWebSocket({ handshakeComplete: false, acceptedAt: Date.now() }); + const sockets = [current, provisional]; + const waiters = createControlRequestWaiters(); + const handler = createSandboxControlSocketHandler( + createFakeState(sockets), + 'sbx_test', + waiters + ); + const request = { + operation: 'session.git.summary' as const, + session: { sessionId: 'workspace_1', kiloSessionId: 'kilo_1', directory: '/workspace' }, + payload: { revision: 1 }, + }; + const pending = handler.sendRequest(request); + const sent = JSON.parse(current.send.mock.calls[0]?.[0] as string) as { requestId: string }; + await handler.handleMessage( + asWs(provisional), + JSON.stringify({ type: 'response', requestId: sent.requestId, ok: true }) + ); + expect(waiters.pendingCount()).toBe(1); + expect(provisional.close).toHaveBeenCalledWith(1008, 'handshake_required'); + await handler.handleMessage( + asWs(current), + JSON.stringify({ type: 'response', requestId: sent.requestId, ok: true }) + ); + await expect(pending).resolves.toMatchObject({ ok: true }); + + const replacement = createFakeWebSocket({ handshakeComplete: false, acceptedAt: Date.now() }); + sockets.push(replacement); + await handler.handleMessage( + asWs(replacement), + JSON.stringify({ + type: 'request', + requestId: 'hello_new', + operation: 'sandbox.hello', + payload: { protocolVersion: 1, providerInstanceId: 'inst_2' }, + }) + ); + expect(handler.getReadySocket()).toBeNull(); + await handler.handleMessage( + asWs(current), + JSON.stringify({ + type: 'event', + event: 'sandbox.ready', + payload: { kiloReady: true, globalFeedAttached: true }, + }) + ); + expect(handler.getReadySocket()).toBeNull(); + await expect(handler.sendRequest(request)).resolves.toMatchObject({ + ok: false, + error: { code: 'not_ready' }, + }); + + await handler.handleMessage( + asWs(replacement), + JSON.stringify({ + type: 'event', + event: 'sandbox.ready', + payload: { kiloReady: true, globalFeedAttached: true }, + }) + ); + const replacementPending = handler.sendRequest(request); + const replacementSent = JSON.parse(replacement.send.mock.calls.at(-1)?.[0] as string) as { + requestId: string; + }; + await handler.handleMessage( + asWs(current), + JSON.stringify({ type: 'response', requestId: replacementSent.requestId, ok: true }) + ); + expect(waiters.pendingCount()).toBe(1); + await handler.handleMessage( + asWs(replacement), + JSON.stringify({ type: 'response', requestId: replacementSent.requestId, ok: true }) + ); + await expect(replacementPending).resolves.toMatchObject({ ok: true }); + }); + + it('rejects a capture if readiness changes before its response is accepted', async () => { + const ws = createFakeWebSocket({ + handshakeComplete: true, + kiloReady: true, + acceptedAt: Date.now(), + providerInstanceId: 'inst_1', + }); + const handler = createSandboxControlSocketHandler(createFakeState([ws]), 'sbx_test'); + const pending = handler.sendRequest({ + operation: 'session.git.summary', + session: { sessionId: 'workspace_1', kiloSessionId: 'kilo_1', directory: '/workspace' }, + payload: { revision: 1 }, + }); + const sent = JSON.parse(ws.send.mock.calls[0]?.[0] as string) as { requestId: string }; + await handler.handleMessage( + asWs(ws), + JSON.stringify({ + type: 'event', + event: 'sandbox.heartbeat', + payload: { state: 'idle', kilo: { ready: false }, sessions: [] }, + }) + ); + await handler.handleMessage( + asWs(ws), + JSON.stringify({ type: 'response', requestId: sent.requestId, ok: true }) + ); + await expect(pending).rejects.toThrow('Worktree capture connection changed'); + }); }); const now = 1_000_000; @@ -1393,6 +1569,57 @@ describe('connection-local sandbox observations', () => { expect(replacement.send).toHaveBeenLastCalledWith(expect.stringContaining('protocol_error')); }); + it.each(['unhealthy', 'closed', 'replaced'] as const)( + 'keeps capture readiness fenced when a delayed heartbeat is overtaken by %s', + async change => { + vi.useFakeTimers({ now }); + const ws = createFakeWebSocket({ ...handshaken, wrapperInstanceId: WRAPPER_INSTANCE_ID }); + const sockets = [ws]; + const state = createFakeState(sockets); + const firstHook = Promise.withResolvers(); + const handler = createSandboxControlSocketHandler(state, 'sbx_test', undefined, { + onHeartbeat: vi.fn().mockImplementationOnce(() => firstHook.promise), + }); + await handler.handleMessage(asWs(ws), readyFrame); + const identity = handler.getConnectionIdentity(); + const earlier = handler.handleMessage(asWs(ws), heartbeatFrame()); + expect(handler.getReadySocket()).toBe(ws); + expect(attachmentOf(ws)).toMatchObject({ ...identity, kiloReady: true }); + const pending = handler.sendRequest({ + operation: 'session.git.summary', + session: { sessionId: 'workspace_1', kiloSessionId: 'kilo_1', directory: '/workspace' }, + payload: { revision: 1 }, + }); + const rejected = expect(pending).rejects.toThrow(); + const sent = JSON.parse(ws.send.mock.calls.at(-1)?.[0] as string) as { requestId: string }; + if (change === 'unhealthy') { + await handler.handleMessage( + asWs(ws), + heartbeatFrame({ + ...idleHeartbeat, + kilo: { ready: false, reason: 'credential_refresh_failed' }, + }) + ); + } else if (change === 'closed') { + handler.closeAll('runtime closed'); + } else { + const replacement = createFakeWebSocket({ handshakeComplete: false, acceptedAt: now }); + sockets.push(replacement); + await handler.handleMessage(asWs(replacement), helloFrame('inst_1', WRAPPER_INSTANCE_ID)); + } + const fencedAttachment = attachmentOf(ws); + firstHook.resolve(); + await earlier; + expect(attachmentOf(ws)).toEqual(fencedAttachment); + expect(handler.getReadySocket()).toBeNull(); + await handler.handleMessage( + asWs(ws), + JSON.stringify({ type: 'response', requestId: sent.requestId, ok: true }) + ); + await rejected; + } + ); + it('does not overwrite newer false evidence when an earlier idle hook completes late', async () => { vi.useFakeTimers({ now }); const ws = createFakeWebSocket(handshaken); diff --git a/services/cloud-agent-next/src/sandbox-control/socket.ts b/services/cloud-agent-next/src/sandbox-control/socket.ts index 4d7e2dfa62..65697a2e8d 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.ts @@ -97,6 +97,7 @@ export type SandboxControlSocketHandler = { sendRequest(input: SandboxControlOutboundRequest): Promise; hasHandshakenSocket(): boolean; getConnectionIdentity(): SandboxControlConnectionIdentity | null; + getReadySocket(): WebSocket | null; closeProvisionalSockets(): void; }; @@ -162,6 +163,10 @@ function sendJson(ws: WebSocket, value: unknown): void { } function closeSocket(ws: WebSocket, code: number, reason: string): void { + const attachment = readAttachment(ws); + if (attachment) { + ws.serializeAttachment({ ...attachment, kiloReady: false }); + } try { ws.close(code, reason); } catch { @@ -283,15 +288,13 @@ export function createSandboxControlSocketHandler( logControlDiagnostic(event, { sandboxId, ...fields }); const observations = new WeakMap(); - function recordObservation( - ws: WebSocket, - attachment: SandboxControlSocketAttachment, - ready: boolean - ): SandboxControlObservation | undefined { + function recordObservation(ws: WebSocket, ready: boolean): SandboxControlObservation | undefined { if (currentHandshakenSocket(state)?.socket !== ws) return undefined; + const attachment = readAttachment(ws); + if (!attachment) return undefined; const observation: SandboxControlObservation = { ready, receivedAt: Date.now(), idle: null }; observations.set(ws, observation); - ws.serializeAttachment({ ...attachment, observation }); + ws.serializeAttachment({ ...attachment, kiloReady: ready, observation }); return observation; } @@ -317,6 +320,11 @@ export function createSandboxControlSocketHandler( return currentHandshakenSocket(state)?.identity ?? null; }, + getReadySocket(): WebSocket | null { + const ws = currentHandshakenSocket(state)?.socket; + return ws && readAttachment(ws)?.kiloReady === true ? ws : null; + }, + closeProvisionalSockets(): void { for (const ws of state.getWebSockets(SANDBOX_CONTROL_WS_TAG)) { const attachment = readAttachment(ws); @@ -465,6 +473,7 @@ export function createSandboxControlSocketHandler( }; const completed: SandboxControlSocketAttachment = { handshakeComplete: true, + kiloReady: false, acceptedAt: attachment.acceptedAt, connectionId: identity.connectionId, protocolVersion: SANDBOX_CONTROL_PROTOCOL_VERSION, @@ -571,16 +580,22 @@ export function createSandboxControlSocketHandler( return; } if (frame.event === 'sandbox.ready') { - recordObservation(ws, attachment, true); + recordObservation(ws, true); await hooks.onReady?.(identity); } else if (frame.event === 'sandbox.heartbeat') { const payload = eventPayload.payload as SandboxHeartbeatPayload; - const observation = recordObservation(ws, attachment, payload.kilo.ready); + const observation = recordObservation(ws, payload.kilo.ready); await hooks.onHeartbeat?.(payload, identity); if (observation) { const idle = await summarizeHeartbeatIdle(payload); if (isCurrentConnection(state, ws, identity) && observations.get(ws) === observation) { - ws.serializeAttachment({ ...attachment, observation: { ...observation, idle } }); + const currentAttachment = readAttachment(ws); + if (currentAttachment) { + ws.serializeAttachment({ + ...currentAttachment, + observation: { ...observation, idle }, + }); + } } } } else if (frame.event === 'session.event') { @@ -638,13 +653,19 @@ export function createSandboxControlSocketHandler( const current = currentHandshakenSocket(state); if (current && current.socket !== ws) return; + const identity = readConnectionIdentity(readAttachment(ws)) ?? undefined; + ws.serializeAttachment({ + ...attachment, + ...identity, + handshakeComplete: false, + kiloReady: false, + }); const remaining = state.getWebSockets(SANDBOX_CONTROL_WS_TAG).some(other => { if (other === ws || other.readyState !== 1) return false; return readAttachment(other)?.handshakeComplete === true; }); if (remaining) return; - const identity = readConnectionIdentity(readAttachment(ws)) ?? undefined; if (identity) activatingConnections.delete(identity.connectionId); waiters.rejectAll('Wrapper socket closed'); await hooks.onSocketClosed?.(true, identity); @@ -670,6 +691,10 @@ export function createSandboxControlSocketHandler( } const current = currentHandshakenSocket(state); + const readyOnly = input.operation === 'session.git.summary'; + if (readyOnly && (!current || readAttachment(current.socket)?.kiloReady !== true)) { + return errorResponse(crypto.randomUUID(), 'not_ready', 'No ready wrapper socket', true); + } if ( !current || current.socket.readyState !== 1 || @@ -695,7 +720,15 @@ export function createSandboxControlSocketHandler( timeoutMs: input.timeoutMs, }); sendJson(current.socket, frame); - return pending; + const response = await pending; + if ( + readyOnly && + (!isCurrentConnection(state, current.socket, current.identity) || + readAttachment(current.socket)?.kiloReady !== true) + ) { + throw new SandboxControlConnectionError('Worktree capture connection changed'); + } + return response; }, }; } diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index e87cd4d23b..85a3b62c94 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -1,4 +1,8 @@ import { DurableObject } from 'cloudflare:workers'; +import type { + GetWorktreeChangesOutput, + RefreshWorktreeChangesOutput, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; import { TRPCError } from '@trpc/server'; import { withTimeout } from '@kilocode/worker-utils'; import { z } from 'zod'; @@ -75,6 +79,12 @@ import { getSandboxControlStub } from '../sandbox-control/stub.js'; import { DEADLINE_MS } from '../sandbox-control/deadlines.js'; import { createMessageId } from '../session/message-id.js'; import { validateControlSessionOptions } from './attach-payload.js'; +import { + createWorktreeChanges, + worktreeChangesContext, + type WorktreeChangesContext, +} from './worktree-changes.js'; +import { WORKTREE_CHANGED_EVENT } from '../shared/worktree-changes-wire.js'; import { createPreparationProgressRecorder } from '../session/preparation-progress.js'; import { finalizeOtherRunningAttemptsForMessage, @@ -192,6 +202,7 @@ export class SandboxSession extends DurableObject { private deletedWorktreeId: CloudAgentWorktreeId | undefined; private readonly activeOperations = new Set>(); private deletionCompletion: Promise | undefined; + private readonly worktreeChanges: ReturnType; constructor(ctx: DurableObjectState, env: Env) { super(ctx, env); @@ -224,6 +235,22 @@ export class SandboxSession extends DurableObject { }); const db = drizzle(ctx.storage, { logger: false }); this.eventQueries = createEventQueries(db, ctx.storage.sql); + this.worktreeChanges = createWorktreeChanges({ + storage: ctx.storage, + readContext: async () => this.worktreeContext(await this.getMetadata()), + requestCapture: (context, payload) => + withDORetry( + () => getSandboxControlStub(this.env, context.sandboxId), + control => + control.request({ + operation: 'session.git.summary', + session: context.session, + payload, + }), + 'captureWorktreeChanges' + ), + waitUntil: promise => this.ctx.waitUntil(promise), + }); void ctx.blockConcurrencyWhile(async () => { await migrate(db, migrations); this.deletedWorktreeId = cloudAgentWorktreeIdSchema @@ -384,6 +411,14 @@ export class SandboxSession extends DurableObject { } if (!this.saveMessages(settled, epoch, 'wrapper_outcome')) return result(false, 'epoch_changed', diagnostic); + if (this.isCurrentEventRuntime(input.wrapperInstanceId)) { + this.worktreeChanges.onEvent( + this.worktreeContext(metadata), + root, + input.payload.type, + outcome.data + ); + } await this.armQueueRetry(); const nextId = nextQueuedMessageId(this.loadMessages()); if (nextId && this.terminalLifecycle.isCurrent(epoch)) { @@ -399,6 +434,17 @@ export class SandboxSession extends DurableObject { (root === undefined || input.identity.rootKiloSessionId !== root) ) return result(false, 'root_mismatch'); + if (input.payload.type === WORKTREE_CHANGED_EVENT) { + if (!root || eventKiloSessionId !== root) return result(false, 'root_mismatch'); + if (!input.wrapperInstanceId) return result(false, 'missing_wrapper_identity'); + this.worktreeChanges.onEvent( + this.worktreeContext(metadata), + eventKiloSessionId, + input.payload.type, + input.payload.properties + ); + return result(true, 'applied'); + } if ( (input.payload.type === 'question.asked' || input.payload.type === 'permission.asked') && !this.loadMessages().some( @@ -415,6 +461,12 @@ export class SandboxSession extends DurableObject { }); const activeMessages = recordAcceptedMessageActivity(this.loadMessages(), Date.now()); if (activeMessages) this.saveMessages(activeMessages, epoch); + this.worktreeChanges.onEvent( + this.worktreeContext(metadata), + eventKiloSessionId, + input.payload.type, + input.payload.properties + ); const ingestItems = controlEventToIngestItems(input.payload.type, input.payload.properties); const rootKiloSessionId = metadata.auth.kiloSessionId; const token = metadata.auth.kilocodeToken; @@ -528,6 +580,7 @@ export class SandboxSession extends DurableObject { async closeOrgStreams(organizationId: string): Promise { const metadata = this.terminalLifecycle.getStoredMetadata(); if (!metadata?.identity.orgId || metadata.identity.orgId !== organizationId) return 0; + this.worktreeChanges.suppress(); const records = this.terminalLifecycle.beginRevocation(metadata); const sockets = this.ctx.getWebSockets('stream'); for (const ws of sockets) ws.close(1000, 'session access revoked'); @@ -590,6 +643,18 @@ export class SandboxSession extends DurableObject { } } + async getWorktreeChanges(): Promise { + if (this.deletedWorktreeId || this.terminalLifecycle.isBlocked()) return { snapshot: null }; + return this.worktreeChanges.get(); + } + + async refreshWorktreeChanges(): Promise { + if (this.deletedWorktreeId || this.terminalLifecycle.isBlocked()) { + return { status: 'offline', snapshot: null }; + } + return this.worktreeChanges.refresh(); + } + async validateKiloGlobalFeedProducer(_params: { kiloSessionId: string; wrapperRunId: string; @@ -668,6 +733,7 @@ export class SandboxSession extends DurableObject { const sandboxId = metadata?.workspace?.sandboxId; const kiloSessionId = metadata?.auth.kiloSessionId; if (!accepted) return { success: true }; + this.worktreeChanges.markInterrupted(this.worktreeContext(metadata)); try { if (!metadata || !sandboxId || !kiloSessionId) throw new Error('Accepted runtime is unavailable'); @@ -822,6 +888,7 @@ export class SandboxSession extends DurableObject { (await this.ctx.storage.get(DELETION_COMPLETED_KEY)) ) return null; + this.worktreeChanges.suppress(); this.ctx.storage.transactionSync(() => { this.ctx.storage.kv.put(DELETED_WORKTREE_KEY, worktreeId); this.terminalLifecycle.beginDeletion(metadata); @@ -922,6 +989,7 @@ export class SandboxSession extends DurableObject { async deleteSession(): Promise { if (this.deletedWorktreeId) throw new Error('worktree_deleting'); + this.worktreeChanges.suppress(); await this.interruptExecution(); if (this.deletedWorktreeId) throw new Error('worktree_deleting'); const metadata = this.terminalLifecycle.getStoredMetadata(); @@ -1528,6 +1596,7 @@ export class SandboxSession extends DurableObject { } let phase: DispatchPhase = 'preparing'; let credentialsPrepared = false; + const preparationGeneration = this.worktreeChanges.beginPreparation(); try { validateControlSessionOptions(metadata); const session = { sessionId, kiloSessionId, directory: this.directory(metadata) }; @@ -1674,6 +1743,7 @@ export class SandboxSession extends DurableObject { throw new Error('Wrapper changed during session attachment'); this.terminalLifecycle.recordAttachment({ metadata, sandboxId, wrapperInstanceId, epoch }); recorder.finalize({ status: 'completed' }); + this.worktreeChanges.attached(preparationGeneration, this.worktreeContext(metadata)); phase = 'prompt'; await dispatch('prompt', async () => { const prompt = await wait(async () => @@ -1736,6 +1806,8 @@ export class SandboxSession extends DurableObject { deadlineAt, error, }); + } finally { + this.worktreeChanges.finishPreparation(preparationGeneration); } } @@ -2278,6 +2350,23 @@ export class SandboxSession extends DurableObject { return { success: true }; } + private worktreeContext(metadata: SessionMetadata | null): WorktreeChangesContext | null { + if ( + this.deletedWorktreeId || + this.terminalLifecycle.isBlocked() || + this.pendingRuntimeCleanup() || + !metadata || + metadata.identity.sessionId !== this.sessionId + ) { + return null; + } + try { + return worktreeChangesContext(metadata, this.directory(metadata)); + } catch { + return null; + } + } + private directory(metadata: SessionMetadata): string { return ( metadata.workspace?.workspacePath ?? diff --git a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts index 9e3d20cd87..130c0745c1 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts @@ -17,6 +17,7 @@ import { SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, sessionPromptPayloadSchema, + sessionGitSummaryPayloadSchema, type ResponseFrame, type SessionAttachPayload, type SessionMessageOutcome, @@ -979,6 +980,8 @@ function sessionFixture(overrides: Partial = {}, sharedControl? orchestrationMocks.signedAttachments.mockResolvedValue([]); const storage = { kv, + get: async (key: string) => kv.get(key), + put: async (key: string, value: T) => kv.put(key, value), sql: {}, transactionSync: (callback: () => T) => callback(), getAlarm: vi.fn(async () => alarmAt), @@ -3165,6 +3168,96 @@ describe('SandboxSession orchestration', () => { expect(await fixture.snapshot()).toMatchObject({ preparationSnapshots: coldPreparation }); }); + it.each(['cloudflare', 'vercel'] as const)( + 'fences captures across warm direct reattachment without clearing ambiguous dispatch on %s', + async sandboxProvider => { + const fixture = sessionFixture({ + repository: { type: 'github', repo: 'acme/repo', upstreamBranch: 'main' }, + workspace: { sandboxId: SANDBOX_ID, workspacePath: DIRECTORY, sandboxProvider }, + }); + fixture.control.ensureReady.mockResolvedValue({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + attachment: { + ...ATTACHMENT, + kilo: { ...ATTACHMENT.kilo, containmentEnabled: false }, + }, + }); + const captureResponse = (input: SandboxControlOutboundRequest) => + controlResponse({ + revision: sessionGitSummaryPayloadSchema.parse(input.payload).revision, + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: [], + truncated: false, + }); + let heldCapture: ReturnType> | undefined; + let heldRequest: SandboxControlOutboundRequest | undefined; + delegateRequest(fixture, 'session.git.summary', async input => { + if (!heldCapture) return captureResponse(input); + heldRequest = input; + return heldCapture.promise; + }); + await fixture.admit('cold'); + await fixture.flush(); + await fixture.outcome('cold', 'completed'); + await fixture.flush(); + const saved = await fixture.session.getWorktreeChanges(); + expect(saved.snapshot).not.toBeNull(); + fixture.reload(); + heldCapture = deferred(); + const staleRefresh = fixture.session.refreshWorktreeChanges(); + await fixture.flush(); + const attached = deferred(); + delegateRequest(fixture, 'session.attach', () => attached.promise); + const lostPrompt = deferred(); + let prompts = 0; + delegateRequest(fixture, 'session.prompt', async () => + ++prompts === 1 ? lostPrompt.promise : controlFailure(true, 'session_busy') + ); + await fixture.admit('warm'); + await fixture.flush(); + if (!heldRequest) throw new Error('Expected in-flight capture'); + heldCapture.resolve(captureResponse(heldRequest)); + heldCapture = undefined; + await expect(staleRefresh).resolves.toEqual({ status: 'failed', snapshot: saved.snapshot }); + await expect(fixture.session.refreshWorktreeChanges()).resolves.toEqual({ + status: 'offline', + snapshot: saved.snapshot, + }); + expect(prompts).toBe(0); + attached.resolve(controlResponse({ attached: true })); + await fixture.flush(); + expect(prompts).toBe(1); + expect(fixture.record('warm')).toMatchObject({ state: 'queued', unresolvedDispatch: true }); + const captured = await fixture.session.getWorktreeChanges(); + expect(captured.snapshot?.revision).toBeGreaterThan(saved.snapshot?.revision ?? 0); + fixture.reload(); + await fixture.fireAlarm(); + await fixture.flush(); + expect(fixture.record('warm')).toMatchObject({ + state: 'queued', + unresolvedDispatch: true, + deliveryRetryScope: 'runtime', + }); + await expect(fixture.session.refreshWorktreeChanges()).resolves.toMatchObject({ + status: 'refreshed', + }); + expect(fixture.record('warm')?.unresolvedDispatch).toBe(true); + await fixture.session.interruptExecution(); + expect(fixture.control.quarantineRuntime).toHaveBeenCalledWith( + expect.objectContaining({ wrapperInstanceId: RUNTIME_ID }) + ); + lostPrompt.resolve(controlResponse({ messageId: 'warm', status: 'accepted' })); + await fixture.flush(); + expect(fixture.record('warm')?.state).toBe('cancelled'); + } + ); + it.each(['stop', 'expiry'] as const)( 'retains ambiguous warm prompt ownership through direct reattachment and %s', async action => { diff --git a/services/cloud-agent-next/src/sandbox-session/worktree-changes.test.ts b/services/cloud-agent-next/src/sandbox-session/worktree-changes.test.ts new file mode 100644 index 0000000000..d2daf4f098 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/worktree-changes.test.ts @@ -0,0 +1,638 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { + WorktreeChangesCapture, + WorktreeChangesCaptureRequest, + WorktreeChangesSnapshot, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; +import type { ResponseFrame } from '../shared/sandbox-control-protocol.js'; +import { WORKTREE_CHANGED_EVENT } from '../shared/worktree-changes-wire.js'; +import { + createWorktreeChanges, + worktreeChangesBaseRef, + worktreeChangesContext, + WORKTREE_CHANGES_KEY, + type WorktreeChangesContext, +} from './worktree-changes.js'; + +const context: WorktreeChangesContext = { + session: { + sessionId: 'workspace_test', + kiloSessionId: 'kilo_root', + directory: '/workspace/test', + }, + ownerId: 'user_test', + sandboxId: 'usr-abc123', + provider: 'cloudflare', + repository: { type: 'github', source: 'acme/demo' }, + baseRef: 'refs/remotes/origin/main', +}; + +function captureResult(revision: number): WorktreeChangesCapture { + return { + revision, + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: [ + { + path: 'src/changed.ts', + status: 'modified', + additions: 3, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, + }, + ], + truncated: false, + }; +} + +const oldSnapshot: WorktreeChangesSnapshot = { + ...captureResult(8), + schemaVersion: 1, + capturedAt: '2026-08-20T10:00:00.000Z', +}; + +function response(result: unknown): ResponseFrame { + return { type: 'response', requestId: 'test', ok: true, result }; +} + +function setup(saved: unknown = oldSnapshot) { + const values = new Map([[WORKTREE_CHANGES_KEY, saved]]); + const storage = { + get: vi.fn(async (key: string) => values.get(key)), + put: vi.fn(async (key: string, value: WorktreeChangesSnapshot) => { + values.set(key, value); + }), + }; + const readContext = vi.fn<() => Promise>(async () => context); + const requestCapture = vi.fn< + ( + context: WorktreeChangesContext, + payload: WorktreeChangesCaptureRequest + ) => Promise + >(async (_context, payload) => response(captureResult(payload.revision))); + const background: Promise[] = []; + const deps = { + storage, + readContext, + requestCapture, + waitUntil: (promise: Promise) => { + background.push(promise); + }, + }; + return { + values, + storage, + readContext, + requestCapture, + background, + deps, + changes: createWorktreeChanges(deps), + }; +} + +function holdCapture(harness: ReturnType) { + const started = Promise.withResolvers(); + const finished = Promise.withResolvers(); + harness.requestCapture.mockImplementationOnce(async (_context, payload) => { + started.resolve(payload); + return finished.promise; + }); + return { started: started.promise, finish: finished.resolve, fail: finished.reject }; +} + +function terminal(harness: ReturnType, currentContext = context): void { + harness.changes.onEvent(currentContext, 'kilo_root', 'session.turn.close', {}); +} + +describe('worktree comparison context', () => { + it.each([ + [undefined, undefined], + ['main', 'refs/remotes/origin/main'], + ['feature/a', 'refs/remotes/origin/feature/a'], + ['refs/heads/release', 'refs/remotes/origin/release'], + ['origin/release', 'refs/remotes/origin/release'], + ['remotes/origin/release', 'refs/remotes/origin/release'], + ['refs/remotes/upstream/release', 'refs/remotes/upstream/release'], + ])('normalizes %s to %s', (branch, expected) => { + expect(worktreeChangesBaseRef(branch)).toBe(expected); + }); + + it('rejects non-branch refs instead of silently changing the comparison', () => { + expect(() => worktreeChangesBaseRef('refs/tags/v1')).toThrow('Unsupported'); + }); + + it('uses selected upstream metadata and never the moving workspace branch or credentials', () => { + const metadata = { + metadataSchemaVersion: 2, + identity: { sessionId: 'workspace_test', userId: 'user_test' }, + auth: { kiloSessionId: 'kilo_root', kilocodeToken: 'private-test-token' }, + repository: { + type: 'github', + repo: 'acme/demo', + upstreamBranch: 'main', + token: 'private-git-token', + }, + workspace: { sandboxId: 'usr-abc123', branchName: 'moving-branch' }, + lifecycle: { version: 1, timestamp: 1 }, + } as SessionMetadata; + expect(worktreeChangesContext(metadata, '/workspace/test')).toEqual(context); + const worktreeId = 'worktree_11111111-1111-4111-8111-111111111111'; + expect( + worktreeChangesContext( + { ...metadata, workspace: { ...metadata.workspace, worktreeId } }, + '/workspace/shared' + ) + ).toEqual({ + ...context, + worktreeId, + session: { ...context.session, directory: '/workspace/shared' }, + }); + expect( + worktreeChangesContext({ ...metadata, repository: undefined }, '/workspace/test') + ).toBeNull(); + expect(worktreeChangesContext({ ...metadata, auth: {} }, '/workspace/test')).toBeNull(); + expect(worktreeChangesContext({ ...metadata, workspace: {} }, '/workspace/test')).toBeNull(); + expect( + worktreeChangesContext( + { ...metadata, repository: { type: 'github', repo: 'acme/demo' } }, + '/workspace/test' + )?.baseRef + ).toBeUndefined(); + }); +}); + +describe('worktree changes capture coordination', () => { + it('reads only persisted, validated storage without resolving runtime context', async () => { + const harness = setup(); + await expect(harness.changes.get()).resolves.toEqual({ snapshot: oldSnapshot }); + harness.values.set(WORKTREE_CHANGES_KEY, { ...oldSnapshot, schemaVersion: 2 }); + await expect(harness.changes.get()).resolves.toEqual({ snapshot: null }); + expect(harness.readContext).not.toHaveBeenCalled(); + expect(harness.requestCapture).not.toHaveBeenCalled(); + expect(harness.storage.put).not.toHaveBeenCalled(); + }); + + it('installs the in-flight promise synchronously and coalesces manual refreshes', async () => { + const harness = setup(); + const held = holdCapture(harness); + const first = harness.changes.refresh(); + const second = harness.changes.refresh(); + expect(second).toBe(first); + expect(await held.started).toEqual({ revision: 9, baseRef: context.baseRef }); + expect(harness.changes.refresh()).toBe(first); + held.finish(response(captureResult(9))); + const result = await first; + expect(result.status).toBe('refreshed'); + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + expect(harness.values.get(WORKTREE_CHANGES_KEY)).toEqual(result.snapshot); + expect(result.snapshot?.capturedAt).not.toBe(oldSnapshot.capturedAt); + }); + + it('keeps one trailing capture for lifecycle events arriving during manual capture', async () => { + const harness = setup(); + const first = holdCapture(harness); + const second = holdCapture(harness); + const refreshed = harness.changes.refresh(); + await first.started; + terminal(harness); + terminal(harness); + terminal(harness); + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + first.finish(response(captureResult(9))); + expect((await second.started).revision).toBe(10); + second.finish(response({ ...captureResult(10), files: [] })); + await expect(refreshed).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 10, files: [] }, + }); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + }); + + it('captures dirty hints immediately and retains one trailing capture during a scan', async () => { + const harness = setup(); + const first = holdCapture(harness); + const second = holdCapture(harness); + harness.changes.onEvent(context, 'kilo_root', WORKTREE_CHANGED_EVENT, {}); + expect(harness.background).toHaveLength(1); + expect((await first.started).revision).toBe(9); + for (let hint = 0; hint < 3; hint++) { + harness.changes.onEvent(context, 'kilo_root', WORKTREE_CHANGED_EVENT, {}); + } + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + first.finish(response(captureResult(9))); + expect((await second.started).revision).toBe(10); + second.finish(response({ ...captureResult(10), files: [] })); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + await expect(harness.changes.get()).resolves.toMatchObject({ + snapshot: { revision: 10, files: [] }, + }); + }); + + it.each(['session.idle', 'session.status'])( + 'preserves pending interruption settlement on %s after a dirty capture', + async type => { + const harness = setup(); + harness.changes.markInterrupted(context); + harness.changes.onEvent(context, 'kilo_root', WORKTREE_CHANGED_EVENT, {}); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + harness.changes.onEvent(context, 'kilo_root', type, { status: { type: 'idle' } }); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + harness.changes.onEvent(context, 'kilo_root', type, { status: { type: 'idle' } }); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + } + ); + + it('recaptures dirty changes at turn close and after finalization', async () => { + const harness = setup(); + harness.changes.onEvent(context, 'kilo_root', WORKTREE_CHANGED_EVENT, {}); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + terminal(harness); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + harness.changes.onEvent(context, 'kilo_root', 'session.message.outcome', { + messageId: 'msg_completed', + status: 'completed', + }); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(3); + await expect(harness.changes.get()).resolves.toMatchObject({ snapshot: { revision: 11 } }); + }); + + it('captures again after finalization when the turn-close capture races a HEAD change', async () => { + const harness = setup(); + const held = holdCapture(harness); + terminal(harness); + await held.started; + harness.changes.onEvent(context, 'kilo_root', 'session.message.outcome', { + messageId: 'msg_completed', + status: 'completed', + }); + held.finish({ + type: 'response', + requestId: 'turn-close', + ok: false, + error: { code: 'capture_failed', message: 'HEAD changed', retryable: true }, + }); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + expect(harness.storage.put).toHaveBeenCalledTimes(1); + await expect(harness.changes.get()).resolves.toMatchObject({ snapshot: { revision: 10 } }); + }); + + it('allows the pending lifecycle slot to refill during a trailing capture without retrying failures', async () => { + const harness = setup(); + const first = holdCapture(harness); + const second = holdCapture(harness); + const third = holdCapture(harness); + const refreshed = harness.changes.refresh(); + await first.started; + terminal(harness); + first.fail(new Error('capture timeout')); + expect((await second.started).revision).toBe(10); + terminal(harness); + terminal(harness); + second.finish(response(captureResult(10))); + expect((await third.started).revision).toBe(11); + third.finish(response(captureResult(11))); + await expect(refreshed).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 11 }, + }); + expect(harness.requestCapture).toHaveBeenCalledTimes(3); + }); + + it.each([ + ['directory', { ...context, session: { ...context.session, directory: '/other' } }], + ['root', { ...context, session: { ...context.session, kiloSessionId: 'other_root' } }], + ['session', { ...context, session: { ...context.session, sessionId: 'workspace_other' } }], + ['sandbox', { ...context, sandboxId: 'usr-other' }], + [ + 'worktree', + { ...context, worktreeId: 'worktree_11111111-1111-4111-8111-111111111111' as const }, + ], + ['provider', { ...context, provider: 'vercel' as const }], + ['repository', { ...context, repository: { type: 'github', source: 'acme/other' } }], + ['base', { ...context, baseRef: 'refs/remotes/origin/other' }], + ['owner', { ...context, ownerId: 'other_owner' }], + ['deleted metadata', null], + ])('discards a result after %s changes during capture', async (_name, nextContext) => { + const harness = setup(); + const held = holdCapture(harness); + const refreshed = harness.changes.refresh(); + await held.started; + harness.readContext.mockResolvedValue(nextContext); + held.finish(response(captureResult(9))); + await expect(refreshed).resolves.toEqual({ status: 'failed', snapshot: oldSnapshot }); + expect(harness.storage.put).not.toHaveBeenCalled(); + await expect(harness.changes.get()).resolves.toEqual({ snapshot: oldSnapshot }); + }); + + it('fences preparation even when replacement metadata is identical and captures after the latest attach', async () => { + const harness = setup(); + const held = holdCapture(harness); + const refreshed = harness.changes.refresh(); + await held.started; + const oldGeneration = harness.changes.beginPreparation(); + const currentGeneration = harness.changes.beginPreparation(); + harness.changes.attached(oldGeneration, context); + expect(harness.background).toHaveLength(0); + held.finish(response(captureResult(9))); + await expect(refreshed).resolves.toEqual({ status: 'failed', snapshot: oldSnapshot }); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'offline', + snapshot: oldSnapshot, + }); + harness.changes.attached(currentGeneration, context); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + await expect(harness.changes.get()).resolves.toMatchObject({ snapshot: { revision: 10 } }); + }); + + it('cleans up only the current preparation without exposing a newer unfinished workspace', async () => { + const harness = setup(); + const obsolete = harness.changes.beginPreparation(); + const current = harness.changes.beginPreparation(); + harness.changes.finishPreparation(obsolete); + terminal(harness); + harness.changes.onEvent(context, 'kilo_root', WORKTREE_CHANGED_EVENT, {}); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'offline', + snapshot: oldSnapshot, + }); + expect(harness.requestCapture).not.toHaveBeenCalled(); + expect(harness.background).toHaveLength(0); + + harness.changes.finishPreparation(current); + expect(harness.background).toHaveLength(0); + await expect(harness.changes.refresh()).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 9 }, + }); + terminal(harness); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + }); + + it('does not restore an obsolete capture fence when preparation ends unsuccessfully', async () => { + const harness = setup(); + const held = holdCapture(harness); + const pending = harness.changes.refresh(); + await held.started; + const generation = harness.changes.beginPreparation(); + harness.changes.finishPreparation(generation); + held.finish(response(captureResult(9))); + await expect(pending).resolves.toEqual({ status: 'failed', snapshot: oldSnapshot }); + expect(harness.storage.put).not.toHaveBeenCalled(); + await expect(harness.changes.refresh()).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 10 }, + }); + }); + + it('does not override deletion suppression when preparation ends', async () => { + const harness = setup(); + const generation = harness.changes.beginPreparation(); + harness.changes.suppress(); + harness.changes.finishPreparation(generation); + terminal(harness); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'offline', + snapshot: oldSnapshot, + }); + expect(harness.requestCapture).not.toHaveBeenCalled(); + expect(harness.background).toHaveLength(0); + }); + + it('uses the latest attached context for a trailing capture after replacement', async () => { + const harness = setup(); + const first = holdCapture(harness); + const trailing = holdCapture(harness); + const refreshed = harness.changes.refresh(); + await first.started; + const replaced = { ...context, sandboxId: 'usr-other', baseRef: 'refs/remotes/origin/release' }; + const generation = harness.changes.beginPreparation(); + harness.readContext.mockResolvedValue(replaced); + harness.changes.attached(generation, replaced); + terminal(harness, replaced); + first.finish(response(captureResult(9))); + expect(await trailing.started).toEqual({ revision: 10, baseRef: replaced.baseRef }); + expect(harness.storage.put).not.toHaveBeenCalled(); + trailing.finish( + response({ + ...captureResult(10), + comparison: { ...captureResult(10).comparison, baseRef: replaced.baseRef }, + }) + ); + await expect(refreshed).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 10, comparison: { baseRef: replaced.baseRef } }, + }); + expect(harness.storage.put).toHaveBeenCalledTimes(1); + expect(harness.requestCapture).toHaveBeenCalledTimes(2); + }); + + it('reports a superseded offline response as failed without overwriting storage', async () => { + const harness = setup(); + const held = holdCapture(harness); + const refreshed = harness.changes.refresh(); + await held.started; + harness.changes.beginPreparation(); + held.finish({ + type: 'response', + requestId: 'test', + ok: false, + error: { code: 'not_ready', message: 'offline', retryable: false }, + }); + await expect(refreshed).resolves.toEqual({ status: 'failed', snapshot: oldSnapshot }); + expect(harness.storage.put).not.toHaveBeenCalled(); + }); + + it('rechecks generation after the final metadata read and immediately before writing', async () => { + const harness = setup(); + harness.readContext.mockResolvedValueOnce(context).mockImplementationOnce(async () => { + harness.changes.beginPreparation(); + return context; + }); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'failed', + snapshot: oldSnapshot, + }); + expect(harness.storage.put).not.toHaveBeenCalled(); + }); + + it('suppresses capture before deletion interrupt and never recreates deleted storage', async () => { + const harness = setup(); + const held = holdCapture(harness); + const refreshed = harness.changes.refresh(); + await held.started; + terminal(harness); + harness.changes.onEvent(context, 'kilo_root', WORKTREE_CHANGED_EVENT, {}); + harness.changes.suppress(); + harness.changes.markInterrupted(context); + terminal(harness); + harness.changes.onEvent(context, 'kilo_root', WORKTREE_CHANGED_EVENT, {}); + harness.changes.onEvent(context, 'kilo_root', 'session.idle', {}); + harness.values.clear(); + held.finish(response(captureResult(9))); + await expect(refreshed).resolves.toEqual({ status: 'failed', snapshot: oldSnapshot }); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + expect(harness.storage.put).not.toHaveBeenCalled(); + await expect(harness.changes.get()).resolves.toEqual({ snapshot: null }); + await expect(harness.changes.refresh()).resolves.toEqual({ status: 'offline', snapshot: null }); + }); + + it.each([ + [ + 'Git failure', + { + type: 'response', + requestId: 'test', + ok: false, + error: { code: 'git_failed', message: 'failed', retryable: false }, + }, + ], + ['malformed data', response({ files: [] })], + ['wrong revision', response(captureResult(10))], + [ + 'wrong comparison', + response({ + ...captureResult(9), + comparison: { ...captureResult(9).comparison, baseRef: 'HEAD' }, + }), + ], + [ + 'oversized file list', + response({ + ...captureResult(9), + files: Array.from({ length: 1001 }, () => captureResult(9).files[0]), + }), + ], + ] satisfies [string, ResponseFrame][])( + 'preserves the exact old snapshot on %s', + async (_name, failed) => { + const harness = setup(); + harness.requestCapture.mockResolvedValue(failed); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'failed', + snapshot: oldSnapshot, + }); + expect(harness.values.get(WORKTREE_CHANGES_KEY)).toEqual(oldSnapshot); + expect(harness.storage.put).not.toHaveBeenCalled(); + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + } + ); + + it('preserves saved data on timeout and persistence failure', async () => { + const harness = setup(); + harness.requestCapture.mockRejectedValueOnce(new Error('timeout')); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'failed', + snapshot: oldSnapshot, + }); + harness.storage.put.mockRejectedValueOnce(new Error('storage unavailable')); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'failed', + snapshot: oldSnapshot, + }); + expect(harness.values.get(WORKTREE_CHANGES_KEY)).toEqual(oldSnapshot); + }); + + it('returns offline without capture when context is unavailable', async () => { + const harness = setup(); + harness.readContext.mockResolvedValue(null); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'offline', + snapshot: oldSnapshot, + }); + expect(harness.requestCapture).not.toHaveBeenCalled(); + expect(harness.storage.put).not.toHaveBeenCalled(); + }); + + it('preserves saved data when the ready-only transport is unavailable', async () => { + const harness = setup(); + harness.requestCapture.mockResolvedValue({ + type: 'response', + requestId: 'test', + ok: false, + error: { code: 'not_ready', message: 'offline', retryable: false }, + }); + await expect(harness.changes.refresh()).resolves.toEqual({ + status: 'offline', + snapshot: oldSnapshot, + }); + expect(harness.storage.put).not.toHaveBeenCalled(); + }); + + it('replaces old files with a valid empty capture and resumes revisions from persisted data', async () => { + const harness = setup(); + harness.requestCapture.mockImplementation(async (_context, payload) => + response({ ...captureResult(payload.revision), files: [] }) + ); + const refreshed = await harness.changes.refresh(); + expect(refreshed).toMatchObject({ status: 'refreshed', snapshot: { files: [], revision: 9 } }); + const restarted = createWorktreeChanges(harness.deps); + await expect(restarted.get()).resolves.toEqual({ snapshot: refreshed.snapshot }); + await expect(restarted.refresh()).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { files: [], revision: 10 }, + }); + }); + + it.each(['session.turn.close', 'session.error', WORKTREE_CHANGED_EVENT])( + 'captures only positively identified root %s events', + async type => { + const harness = setup(); + harness.changes.onEvent(context, 'kilo_child', type, {}); + harness.changes.onEvent(context, undefined, type, {}); + expect(harness.background).toHaveLength(0); + harness.changes.onEvent(context, 'kilo_root', type, {}); + await Promise.all(harness.background); + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + } + ); + + it.each(['session.idle', 'session.status'])( + 'waits for confirmed root %s after interruption', + async type => { + const harness = setup(); + harness.changes.onEvent(context, 'kilo_root', type, { status: { type: 'idle' } }); + expect(harness.background).toHaveLength(0); + harness.changes.markInterrupted(context); + harness.changes.onEvent(context, 'kilo_child', type, { status: { type: 'idle' } }); + harness.changes.onEvent(context, 'kilo_root', 'session.status', { status: { type: 'busy' } }); + harness.changes.onEvent(context, 'kilo_root', 'session.status', { status: 'idle' }); + expect(harness.background).toHaveLength(0); + harness.changes.onEvent(context, 'kilo_root', type, { status: { type: 'idle' } }); + await Promise.all(harness.background); + harness.changes.onEvent(context, 'kilo_root', 'session.idle', {}); + expect(harness.requestCapture).toHaveBeenCalledTimes(1); + } + ); + + it('does not transfer pending interruption settlement into a changed workspace', () => { + const harness = setup(); + harness.changes.markInterrupted(context); + harness.changes.onEvent( + { ...context, sandboxId: 'usr-other' }, + 'kilo_root', + 'session.idle', + {} + ); + harness.changes.beginPreparation(); + harness.changes.onEvent(context, 'kilo_root', 'session.idle', {}); + expect(harness.background).toHaveLength(0); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-session/worktree-changes.ts b/services/cloud-agent-next/src/sandbox-session/worktree-changes.ts new file mode 100644 index 0000000000..a3ed1f718d --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/worktree-changes.ts @@ -0,0 +1,248 @@ +import { z } from 'zod'; +import { + WORKTREE_CHANGES_SCHEMA_VERSION, + worktreeChangesCaptureSchema, + worktreeChangesSnapshotSchema, + type GetWorktreeChangesOutput, + type RefreshWorktreeChangesOutput, + type WorktreeChangesCaptureRequest, + type WorktreeChangesSnapshot, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; +import { getSandboxProvider } from '../persistence/session-metadata.js'; +import type { ResponseFrame, SessionRequestIdentity } from '../shared/sandbox-control-protocol.js'; +import { WORKTREE_CHANGED_EVENT } from '../shared/worktree-changes-wire.js'; + +export const WORKTREE_CHANGES_KEY = 'worktree_changes'; + +export type WorktreeChangesContext = { + session: SessionRequestIdentity; + ownerId: string; + orgId?: string; + sandboxId: string; + worktreeId?: NonNullable['worktreeId']; + provider: 'cloudflare' | 'vercel'; + providerRuntime?: NonNullable['providerRuntime']; + repository: { type: string; source: string }; + baseRef?: string; +}; + +type CaptureTrigger = { generation: number; context?: WorktreeChangesContext }; + +type WorktreeChangesDependencies = { + storage: { + get(key: string): Promise; + put(key: string, value: WorktreeChangesSnapshot): Promise; + }; + readContext(): Promise; + requestCapture( + context: WorktreeChangesContext, + payload: WorktreeChangesCaptureRequest + ): Promise; + waitUntil(promise: Promise): void; +}; + +const idleStatusSchema = z.object({ status: z.object({ type: z.literal('idle') }) }); + +export function worktreeChangesBaseRef(branch: string | undefined): string | undefined { + if (branch === undefined) return undefined; + if (branch.startsWith('refs/remotes/')) return branch; + if (branch.startsWith('remotes/')) return `refs/${branch}`; + if (branch.startsWith('origin/')) return `refs/remotes/${branch}`; + if (branch.startsWith('refs/heads/')) { + return `refs/remotes/origin/${branch.slice('refs/heads/'.length)}`; + } + if (branch.startsWith('refs/')) throw new Error('Unsupported worktree comparison ref'); + return `refs/remotes/origin/${branch}`; +} + +export function worktreeChangesContext( + metadata: SessionMetadata, + directory: string +): WorktreeChangesContext | null { + const sandboxId = metadata.workspace?.sandboxId; + const kiloSessionId = metadata.auth.kiloSessionId; + const repository = metadata.repository; + if (!sandboxId || !kiloSessionId || !repository || !directory) return null; + return { + session: { sessionId: metadata.identity.sessionId, kiloSessionId, directory }, + ownerId: metadata.identity.userId, + orgId: metadata.identity.orgId, + sandboxId, + worktreeId: metadata.workspace?.worktreeId, + provider: getSandboxProvider(metadata), + providerRuntime: metadata.workspace?.providerRuntime, + repository: { + type: repository.type, + source: repository.type === 'github' ? repository.repo : repository.url, + }, + baseRef: worktreeChangesBaseRef(repository.upstreamBranch), + }; +} + +function sameContext(left: WorktreeChangesContext, right: WorktreeChangesContext | null): boolean { + return right !== null && JSON.stringify(left) === JSON.stringify(right); +} + +export function createWorktreeChanges(deps: WorktreeChangesDependencies) { + let generation = 0; + let revision = 0; + let suppressed = false; + let preparing = false; + let inFlight: Promise | undefined; + let pending: CaptureTrigger | undefined; + let pendingInterruption: { generation: number; context: WorktreeChangesContext } | undefined; + + async function readSnapshot(): Promise { + const parsed = worktreeChangesSnapshotSchema.safeParse( + await deps.storage.get(WORKTREE_CHANGES_KEY) + ); + return parsed.success ? parsed.data : null; + } + + async function capture(trigger: CaptureTrigger): Promise { + let snapshot: WorktreeChangesSnapshot | null = null; + try { + snapshot = await readSnapshot(); + if (trigger.generation !== generation) return { status: 'failed', snapshot }; + if (suppressed || preparing) return { status: 'offline', snapshot }; + const context = await deps.readContext(); + if (trigger.generation !== generation) return { status: 'failed', snapshot }; + if (!context) return { status: 'offline', snapshot }; + if (trigger.context && !sameContext(trigger.context, context)) { + return { status: 'failed', snapshot }; + } + revision = Math.max(revision, snapshot?.revision ?? 0) + 1; + if (!Number.isSafeInteger(revision)) return { status: 'failed', snapshot }; + const requestedRevision = revision; + const response = await deps.requestCapture(context, { + revision: requestedRevision, + ...(context.baseRef ? { baseRef: context.baseRef } : {}), + }); + if (trigger.generation !== generation) return { status: 'failed', snapshot }; + if (!response.ok) { + return { status: response.error?.code === 'not_ready' ? 'offline' : 'failed', snapshot }; + } + const parsed = worktreeChangesCaptureSchema.safeParse(response.result); + if ( + !parsed.success || + parsed.data.revision !== requestedRevision || + (context.baseRef !== undefined && parsed.data.comparison.baseRef !== context.baseRef) + ) { + return { status: 'failed', snapshot }; + } + const saved = worktreeChangesSnapshotSchema.safeParse({ + ...parsed.data, + schemaVersion: WORKTREE_CHANGES_SCHEMA_VERSION, + capturedAt: new Date().toISOString(), + }); + if (!saved.success) return { status: 'failed', snapshot }; + const current = await deps.readContext(); + if ( + suppressed || + preparing || + trigger.generation !== generation || + parsed.data.revision !== revision || + !sameContext(context, current) + ) { + return { status: 'failed', snapshot }; + } + await deps.storage.put(WORKTREE_CHANGES_KEY, saved.data); + return { status: 'refreshed', snapshot: saved.data }; + } catch { + return { status: 'failed', snapshot }; + } + } + + function start(trigger: CaptureTrigger): Promise { + if (inFlight) return inFlight; + inFlight = Promise.resolve().then(async () => { + try { + let next = trigger; + while (true) { + const result = await capture(next); + const trailing = pending; + pending = undefined; + if (!trailing) return result; + next = trailing; + } + } finally { + inFlight = undefined; + } + }); + return inFlight; + } + + function schedule(context: WorktreeChangesContext): void { + if (suppressed || preparing) return; + const trigger = { generation, context }; + if (inFlight) pending = trigger; + deps.waitUntil(start(trigger)); + } + + function invalidate(): void { + generation++; + pending = undefined; + pendingInterruption = undefined; + } + + return { + async get(): Promise { + return { snapshot: await readSnapshot() }; + }, + + refresh(): Promise { + return start({ generation }); + }, + + beginPreparation(): number { + invalidate(); + preparing = true; + return generation; + }, + + finishPreparation(preparationGeneration: number): void { + if (preparationGeneration === generation) preparing = false; + }, + + attached(preparationGeneration: number, context: WorktreeChangesContext | null): void { + if (preparationGeneration !== generation || suppressed) return; + preparing = false; + if (context) schedule(context); + }, + + markInterrupted(context: WorktreeChangesContext | null): void { + if (context && !suppressed && !preparing) pendingInterruption = { generation, context }; + }, + + onEvent( + context: WorktreeChangesContext | null, + eventKiloSessionId: string | undefined, + type: string, + properties: Record + ): void { + if (!context || eventKiloSessionId !== context.session.kiloSessionId) return; + if (type === WORKTREE_CHANGED_EVENT) { + schedule(context); + return; + } + const terminal = + type === 'session.turn.close' || + type === 'session.error' || + type === 'session.message.outcome'; + const interruptionSettled = + pendingInterruption?.generation === generation && + sameContext(pendingInterruption.context, context) && + (type === 'session.idle' || + (type === 'session.status' && idleStatusSchema.safeParse(properties).success)); + if (!terminal && !interruptionSettled) return; + pendingInterruption = undefined; + schedule(context); + }, + + suppress(): void { + suppressed = true; + invalidate(); + }, + }; +} diff --git a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts index a8a7ae934e..abd19c3c63 100644 --- a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts +++ b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts @@ -1,6 +1,17 @@ import { z } from 'zod'; import { SandboxRuntimeVersionSchema } from './sandbox-status.js'; +export { + MAX_WORKTREE_CHANGES_BYTES, + MAX_WORKTREE_CHANGES_FILES, + worktreeChangesFileSchema, + sessionGitSummaryPayloadSchema, + sessionGitSummaryResultSchema, + type WorktreeChangesFile, + type SessionGitSummaryPayload, + type SessionGitSummaryResult, +} from './worktree-changes-wire.js'; + export const SANDBOX_CONTROL_PROTOCOL_VERSION = 1; export const MAX_SANDBOX_CONTROL_FRAME_BYTES = 1 * 1024 * 1024; @@ -34,6 +45,7 @@ export const SESSION_OPERATIONS = [ 'session.question.resolve', 'session.abort', 'session.sync', + 'session.git.summary', 'session.detach', 'session.terminal.create', 'session.terminal.resize', @@ -599,6 +611,7 @@ export type SandboxControlObservation = z.infer { + it('keeps input and output types identical to the public contracts', () => { + expectTypeOf>().toEqualTypeOf< + z.input + >(); + expectTypeOf>().toEqualTypeOf< + z.output + >(); + expectTypeOf>().toEqualTypeOf< + z.input + >(); + expectTypeOf>().toEqualTypeOf< + z.output + >(); + expectTypeOf>().toEqualTypeOf< + z.input + >(); + expectTypeOf>().toEqualTypeOf< + z.output + >(); + }); + + it('preserves requests at the revision and base-ref boundaries', () => { + for (const input of [ + { revision: 1 }, + { revision: 2, baseRef: capture.comparison.baseRef }, + { revision: Number.MAX_SAFE_INTEGER, baseRef: 'x'.repeat(1024) }, + ]) { + for (const schema of requestSchemas) expect(schema.parse(input)).toEqual(input); + } + }); + + it('rejects invalid requests and routing fields in both contracts', () => { + for (const input of [ + {}, + { revision: 0 }, + { revision: -1 }, + { revision: 1.5 }, + { revision: Number.MAX_SAFE_INTEGER + 1 }, + { revision: '1' }, + { revision: 1, baseRef: '' }, + { revision: 1, baseRef: null }, + { revision: 1, baseRef: 'x'.repeat(1025) }, + { revision: 1, baseRef: '--help' }, + { revision: 1, baseRef: 'main\0suffix' }, + { revision: 1, directory: '/outside' }, + { revision: 1, sessionId: 'other-session' }, + ]) { + for (const schema of requestSchemas) expect(schema.safeParse(input).success).toBe(false); + } + }); + + it('preserves file statuses, count flags, and unusual paths without normalization', () => { + for (const input of [ + file, + { + ...file, + status: 'added', + tracked: false, + additions: 0, + deletions: 0, + countsComplete: false, + }, + { + ...file, + status: 'deleted', + binary: true, + additions: 0, + deletions: 0, + countsComplete: false, + }, + { ...file, additions: Number.MAX_SAFE_INTEGER, deletions: Number.MAX_SAFE_INTEGER }, + { ...file, path: 'parent/ leading\tline\n"back\\slash-é-漢 ' }, + { ...file, path: 'x'.repeat(4096) }, + ]) { + for (const schema of fileSchemas) expect(schema.parse(input)).toEqual(input); + } + }); + + it('rejects unsafe or oversized paths in both contracts', () => { + for (const path of [ + '', + '/outside', + '../outside', + 'parent/../outside', + './file', + 'parent//file', + 'file\0suffix', + 'x'.repeat(4097), + ]) { + for (const schema of fileSchemas) + expect(schema.safeParse({ ...file, path }).success).toBe(false); + } + }); + + it('rejects invalid counts, missing flags, and file contents in both contracts', () => { + for (const invalid of [ + { additions: -1 }, + { deletions: 1.5 }, + { additions: Number.MAX_SAFE_INTEGER + 1 }, + { deletions: Number.POSITIVE_INFINITY }, + { additions: Number.NaN }, + { tracked: 'true' }, + { binary: undefined }, + { countsComplete: undefined }, + { status: 'renamed' }, + { contents: 'not summary data' }, + { patch: 'not summary data' }, + ]) { + for (const schema of fileSchemas) + expect(schema.safeParse({ ...file, ...invalid }).success).toBe(false); + } + }); + + it('preserves empty, truncated, and SHA-256 capture results', () => { + for (const input of [ + capture, + { ...capture, files: [] }, + { ...capture, truncated: true }, + { + ...capture, + revision: Number.MAX_SAFE_INTEGER, + comparison: { baseRef: 'x'.repeat(1024), mergeBase: 'c'.repeat(64), head: 'd'.repeat(64) }, + }, + ]) { + for (const schema of captureSchemas) expect(schema.parse(input)).toEqual(input); + } + }); + + it('rejects invalid result envelopes and comparison identities in both contracts', () => { + for (const invalid of [ + { revision: 0 }, + { revision: Number.MAX_SAFE_INTEGER + 1 }, + { files: null }, + { truncated: undefined }, + { capturedAt: '2026-08-27T00:00:00.000Z' }, + { schemaVersion: 1 }, + { comparison: { ...capture.comparison, baseRef: '--help' } }, + { comparison: { ...capture.comparison, baseRef: 'x'.repeat(1025) } }, + { comparison: { ...capture.comparison, head: 'HEAD' } }, + { comparison: { ...capture.comparison, mergeBase: 'a'.repeat(39) } }, + { comparison: { ...capture.comparison, head: 'B'.repeat(40) } }, + { comparison: { ...capture.comparison, directory: '/outside' } }, + { files: [{ ...file, contents: 'not summary data' }] }, + { files: [{ ...file, path: 'x'.repeat(4097) }] }, + { files: [{ ...file, additions: -1 }] }, + ]) { + for (const schema of captureSchemas) + expect(schema.safeParse({ ...capture, ...invalid }).success).toBe(false); + } + }); + + it('enforces the same file-count boundary independently of serialized bytes', () => { + expect(MAX_WORKTREE_CHANGES_FILES).toBe(PUBLIC_MAX_FILES); + const input = { + ...capture, + files: Array.from({ length: PUBLIC_MAX_FILES }, (_, index) => ({ + ...file, + path: String(index), + })), + }; + const oversized = { ...input, files: [...input.files, { ...file, path: 'extra' }] }; + expect(Buffer.byteLength(JSON.stringify(oversized))).toBeLessThan(PUBLIC_MAX_BYTES); + for (const schema of captureSchemas) { + expect(schema.parse(input)).toEqual(input); + expect(schema.safeParse(oversized).success).toBe(false); + } + }); + + it('enforces the same inclusive UTF-8 byte limit, not JavaScript string length', () => { + expect(MAX_WORKTREE_CHANGES_BYTES).toBe(PUBLIC_MAX_BYTES); + const input = { + ...capture, + files: Array.from({ length: 80 }, (_, index) => ({ + ...file, + path: `${index}/${'漢'.repeat(1000)}`, + })), + }; + let padding = PUBLIC_MAX_BYTES - Buffer.byteLength(JSON.stringify(input)); + expect(padding).toBeGreaterThan(0); + for (const entry of input.files) { + const length = Math.min(padding, 4096 - entry.path.length); + entry.path += 'x'.repeat(length); + padding -= length; + } + expect(padding).toBe(0); + expect(Buffer.byteLength(JSON.stringify(input))).toBe(PUBLIC_MAX_BYTES); + expect(JSON.stringify(input).length).toBeLessThan(PUBLIC_MAX_BYTES); + const oversized = { + ...input, + files: input.files.map((entry, index) => + index === input.files.length - 1 ? { ...entry, path: `${entry.path}x` } : entry + ), + }; + expect(Buffer.byteLength(JSON.stringify(oversized))).toBe(PUBLIC_MAX_BYTES + 1); + for (const schema of captureSchemas) { + expect(schema.parse(input)).toEqual(input); + expect(schema.safeParse(oversized).success).toBe(false); + } + }); +}); diff --git a/services/cloud-agent-next/src/shared/worktree-changes-wire.ts b/services/cloud-agent-next/src/shared/worktree-changes-wire.ts new file mode 100644 index 0000000000..46922bdf46 --- /dev/null +++ b/services/cloud-agent-next/src/shared/worktree-changes-wire.ts @@ -0,0 +1,70 @@ +import type { + WorktreeChangesCapture, + WorktreeChangesCaptureRequest, + WorktreeChangesFile as PublicWorktreeChangesFile, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; +import { z } from 'zod'; + +export const MAX_WORKTREE_CHANGES_FILES = 1_000; +export const MAX_WORKTREE_CHANGES_BYTES = 256 * 1024; +export const WORKTREE_CHANGED_EVENT = 'session.worktree.changed'; + +const revisionSchema = z.number().int().positive().max(Number.MAX_SAFE_INTEGER); +const commitSchema = z.string().regex(/^(?:[0-9a-f]{40}|[0-9a-f]{64})$/); +const baseRefSchema = z + .string() + .min(1) + .max(1024) + .refine(ref => !ref.startsWith('-') && !ref.includes('\0'), 'Invalid comparison ref'); + +export const worktreeChangesFileSchema = z + .object({ + path: z + .string() + .min(1) + .max(4096) + .refine( + path => + !path.includes('\0') && + path.split('/').every(part => part !== '' && part !== '.' && part !== '..'), + 'Expected a repository-relative path' + ), + status: z.enum(['added', 'modified', 'deleted']), + additions: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + deletions: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + tracked: z.boolean(), + binary: z.boolean(), + countsComplete: z.boolean(), + }) + .strict() satisfies z.ZodType; + +export const sessionGitSummaryPayloadSchema = z + .object({ + revision: revisionSchema, + baseRef: baseRefSchema.optional(), + }) + .strict() satisfies z.ZodType; + +export const sessionGitSummaryResultSchema = z + .object({ + revision: revisionSchema, + comparison: z + .object({ + baseRef: baseRefSchema, + mergeBase: commitSchema, + head: commitSchema, + }) + .strict(), + files: z.array(worktreeChangesFileSchema).max(MAX_WORKTREE_CHANGES_FILES), + truncated: z.boolean(), + }) + .strict() + .refine( + capture => + new TextEncoder().encode(JSON.stringify(capture)).byteLength <= MAX_WORKTREE_CHANGES_BYTES, + 'Worktree summary exceeds the size limit' + ) satisfies z.ZodType; + +export type WorktreeChangesFile = z.infer; +export type SessionGitSummaryPayload = z.infer; +export type SessionGitSummaryResult = z.infer; diff --git a/services/cloud-agent-next/test/integration/sandbox-control.test.ts b/services/cloud-agent-next/test/integration/sandbox-control.test.ts index 94fdc06ccf..e9eaea7446 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -6,6 +6,11 @@ import { runDurableObjectAlarm, runInDurableObject, } from 'cloudflare:test'; +import { + worktreeChangesCaptureRequestSchema, + type WorktreeChangesCapture, + type WorktreeChangesSnapshot, +} from '@kilocode/worker-utils/cloud-agent-worktree-changes'; import { drizzle } from 'drizzle-orm/durable-sqlite'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import type { BillingContext } from '@kilocode/container-usage'; @@ -92,6 +97,7 @@ import { type VercelControlRestClient, } from '../../src/sandbox-control/vercel-provider.js'; import { + ATTACH_FAILURE_LIMIT, createSessionMessageRecord, type SessionMessageRecord, } from '../../src/sandbox-session/session-message-queue.js'; @@ -103,14 +109,17 @@ import { sessionPromptPayloadSchema, SANDBOX_CONTROL_AUTO_PING, SANDBOX_CONTROL_AUTO_PONG, + SANDBOX_CONTROL_WS_TAG, type RequestFrame, type ResponseFrame, type SessionAttachPayload, - SANDBOX_CONTROL_WS_TAG, sandboxControlSocketAttachmentSchema, type SandboxHeartbeatPayload, } from '../../src/shared/sandbox-control-protocol.js'; import { SandboxStatusSnapshotSchema } from '../../src/shared/sandbox-status.js'; +import { WORKTREE_CHANGES_KEY } from '../../src/sandbox-session/worktree-changes.js'; +import { WORKTREE_CHANGED_EVENT } from '../../src/shared/worktree-changes-wire.js'; +import { getWorktreeWorkspacePath } from '../../src/workspace.js'; vi.mock('../../src/session-access.js', () => ({ requireCurrentSessionAccess: vi.fn(), @@ -1325,7 +1334,11 @@ async function deliverWrapperEvent( }); } -function respondToWrapperRequest(ws: WebSocket, request: WrapperRequest, result: unknown): void { +function respondToWrapperRequest( + ws: WebSocket, + request: Pick, + result: unknown +): void { ws.send(JSON.stringify({ type: 'response', requestId: request.requestId, ok: true, result })); } @@ -6644,6 +6657,1027 @@ describe('SandboxSession passive delegation', () => { ); }); +function worktreeCapture(revision: number, empty = false): WorktreeChangesCapture { + return { + revision, + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: empty + ? [] + : [ + { + path: 'changed.ts', + status: 'modified', + additions: 2, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, + }, + ], + truncated: false, + }; +} + +const savedWorktreeSnapshot: WorktreeChangesSnapshot = { + ...worktreeCapture(4), + schemaVersion: 1, + capturedAt: '2026-08-20T10:00:00.000Z', +}; + +async function worktreeFixture() { + const suffix = crypto.randomUUID(); + const userId = `user_worktree_${suffix}`; + const sessionId = `workspace_${suffix}` as const; + const sandboxId = `usr-${suffix.replaceAll('-', '').slice(0, 12)}` as const; + const kiloSessionId = ROOT_ID; + const worktreeId = `worktree_${suffix}` as const; + const directory = getWorktreeWorkspacePath(undefined, userId, worktreeId); + const wrapperInstanceId = crypto.randomUUID(); + const control = env.SANDBOX_CONTROL.getByName(sandboxId); + const session = env.SANDBOX_SESSION.getByName(`${userId}:${sessionId}`); + const credential = generateSandboxCredential(); + const controlTasks: Promise[] = []; + const sessionTasks: Promise[] = []; + await seedRunningCredential(credential, sandboxId); + const { provider } = await installProvider(control, cloudflareRef(sandboxId)); + await runInDurableObject(control, async (instance, state) => { + await instance.initializeOwner(userId); + const waitUntil = state.waitUntil.bind(state); + vi.spyOn(state, 'waitUntil').mockImplementation(promise => { + controlTasks.push(promise); + waitUntil(promise); + }); + }); + await runInDurableObject(session, async (instance, state) => { + await instance.registerSession({ + identity: { sessionId, userId, createdOnPlatform: 'cloud-agent-web' }, + auth: { kiloSessionId, kilocodeToken: KILO_TOKEN }, + agent: { mode: 'code', model: 'test' }, + repository: { + type: 'github', + repo: 'acme/demo', + upstreamBranch: 'main', + }, + workspace: { + sandboxId, + sandboxProvider: 'cloudflare', + worktreeId, + workspacePath: directory, + branchName: 'moving-work-branch', + }, + }); + const waitUntil = state.waitUntil.bind(state); + vi.spyOn(state, 'waitUntil').mockImplementation(promise => { + sessionTasks.push(promise); + waitUntil(promise); + }); + }); + await control.prepareSessionCredentials({ ownerId: userId, sessionId }); + await control.attachSession({ sessionId, kiloSessionId, directory, worktreeId, ownerId: userId }); + let ws = await connect(credential, sandboxId); + await completeHello(ws, `hello_${suffix}`, { wrapperInstanceId }); + const captures: RequestFrame[] = []; + const inbox: RequestFrame[] = []; + const captureWaiters: ((request: RequestFrame) => void)[] = []; + const prompts: RequestFrame[] = []; + const promptSeen = Promise.withResolvers(); + const aborts: RequestFrame[] = []; + let nextAttach: ((request: RequestFrame) => void) | undefined; + + function receive(client: WebSocket): void { + client.addEventListener('message', event => { + const parsed = requestFrameSchema.safeParse(JSON.parse(String(event.data))); + if (!parsed.success) return; + const request = parsed.data; + if (request.operation === 'session.git.summary') { + captures.push(request); + const waiting = captureWaiters.shift(); + if (waiting) waiting(request); + else inbox.push(request); + return; + } + if (request.operation === 'session.attach' && nextAttach) { + const resolve = nextAttach; + nextAttach = undefined; + resolve(request); + return; + } + let result: unknown; + if (request.operation === 'session.attach') result = { attached: true }; + else if (request.operation === 'session.prompt') { + prompts.push(request); + const payload = request.payload as { messageId: string }; + result = { messageId: payload.messageId, status: 'accepted' }; + promptSeen.resolve(); + } else if (request.operation === 'session.abort') { + aborts.push(request); + result = { status: 'aborted' }; + } else if (request.operation === 'session.detach') { + result = { detached: true }; + } else return; + client.send( + JSON.stringify({ type: 'response', requestId: request.requestId, ok: true, result }) + ); + }); + } + receive(ws); + + async function ready(): Promise { + await runInDurableObject(control, async (instance, state) => { + const server = state + .getWebSockets(SANDBOX_CONTROL_WS_TAG) + .find(socket => socket.readyState === 1); + if (!server) throw new Error('Missing test control socket'); + await instance.webSocketMessage( + server, + JSON.stringify({ + type: 'event', + event: 'sandbox.ready', + payload: { kiloReady: true, globalFeedAttached: true }, + }) + ); + }); + } + await ready(); + const noWake = await runInDurableObject(control, instance => { + const prototype = Object.getPrototypeOf(instance) as typeof instance; + return { + ensureReady: vi.spyOn(prototype, 'ensureReady'), + attachSession: vi.spyOn(prototype, 'attachSession'), + claimCreate: vi.spyOn(prototype, 'claimCreate'), + }; + }); + + return { + userId, + sessionId, + sandboxId, + kiloSessionId, + worktreeId, + wrapperInstanceId, + directory, + control, + session, + provider, + captures, + prompts, + aborts, + noWake, + promptSeen: promptSeen.promise, + holdNextAttach(): Promise { + return new Promise(resolve => { + nextAttach = resolve; + }); + }, + async nextCapture(): Promise { + const request = inbox.shift(); + if (request) return request; + return new Promise(resolve => captureWaiters.push(resolve)); + }, + reply(request: RequestFrame, result: unknown): void { + ws.send(JSON.stringify({ type: 'response', requestId: request.requestId, ok: true, result })); + }, + fail( + request: RequestFrame, + retryable = false, + code = retryable ? 'not_ready' : 'git_failed' + ): void { + ws.send( + JSON.stringify({ + type: 'response', + requestId: request.requestId, + ok: false, + error: { + code, + message: 'Fixture request failed', + retryable, + }, + }) + ); + }, + async event( + type: string, + root = kiloSessionId, + properties: Record = {}, + eventDirectory = directory + ): Promise { + await runInDurableObject(control, async (instance, state) => { + const server = state + .getWebSockets(SANDBOX_CONTROL_WS_TAG) + .find(socket => socket.readyState === 1); + if (!server) throw new Error('Missing test control socket'); + await instance.webSocketMessage( + server, + JSON.stringify({ + type: 'event', + event: 'session.event', + session: { + directory: eventDirectory, + kiloSessionId: root, + rootKiloSessionId: kiloSessionId, + }, + payload: { + type, + properties: + type === 'session.message.outcome' || type === WORKTREE_CHANGED_EVENT + ? properties + : { sessionID: root, ...properties }, + }, + }) + ); + }); + await Promise.all(controlTasks); + }, + async settled(): Promise { + await Promise.all(controlTasks); + await Promise.all(sessionTasks); + }, + async rotateSocket(): Promise { + await control.setWrapperCredentialHash(await hashSandboxCredential(credential)); + ws = await connect(credential, sandboxId); + await completeHello(ws, `hello_replacement_${suffix}`, { wrapperInstanceId }); + receive(ws); + }, + ready, + close(): void { + ws.close(); + }, + }; +} + +function captureRevision(request: RequestFrame): number { + return (request.payload as { revision: number }).revision; +} + +describe('SandboxSession worktree changes persistence', () => { + beforeEach(() => { + vi.spyOn(globalThis, 'fetch').mockImplementation(async () => Response.json({ valid: true })); + }); + + afterEach(() => vi.restoreAllMocks()); + + it('captures after attach without a UI request and does not delay prompt delivery', async () => { + const fixture = await worktreeFixture(); + await fixture.session.admitSubmittedMessage({ + userId: fixture.userId, + turn: { type: 'prompt', id: 'msg_worktree_attach', prompt: 'test prompt' }, + }); + const request = await fixture.nextCapture(); + expect(request.session).toEqual({ + sessionId: fixture.sessionId, + kiloSessionId: fixture.kiloSessionId, + directory: fixture.directory, + }); + expect(request.payload).toEqual({ revision: 1, baseRef: 'refs/remotes/origin/main' }); + await fixture.promptSeen; + expect(fixture.prompts).toHaveLength(1); + fixture.reply(request, worktreeCapture(1)); + await fixture.settled(); + await expect(fixture.session.getWorktreeChanges()).resolves.toMatchObject({ + snapshot: { schemaVersion: 1, revision: 1, files: worktreeCapture(1).files }, + }); + await expect(fixture.session.getCurrentMessageWork()).resolves.toMatchObject({ + messageId: 'msg_worktree_attach', + status: 'running', + }); + fixture.close(); + }); + + it('captures dirty hints during an accepted turn without chat events or artificial activity', async () => { + const fixture = await worktreeFixture(); + const messageId = 'msg_worktree_dirty'; + try { + await fixture.session.admitSubmittedMessage({ + userId: fixture.userId, + turn: { type: 'prompt', id: messageId, prompt: 'edit the worktree' }, + }); + const attached = await fixture.nextCapture(); + fixture.reply(attached, worktreeCapture(captureRevision(attached), true)); + await fixture.settled(); + fixture.noWake.ensureReady.mockClear(); + fixture.noWake.attachSession.mockClear(); + fixture.noWake.claimCreate.mockClear(); + const before = await runInDurableObject(fixture.session, async (instance, state) => { + const messages = ( + state.storage.kv.get('session_messages') ?? [] + ).map(message => ({ ...message, lastActivityAt: 1 })); + state.storage.kv.put('session_messages', messages); + const broadcast = vi.fn(instance['broadcastStoredEvent'].bind(instance)); + instance['broadcastStoredEvent'] = broadcast; + return { + broadcast, + messages, + events: createEventQueries(drizzle(state.storage), state.storage.sql).findByFilters({}), + alarm: await state.storage.getAlarm(), + }; + }); + const controlBefore = await runInDurableObject(fixture.control, async (_instance, state) => ({ + records: await state.storage.list(), + alarm: await state.storage.getAlarm(), + })); + + await fixture.event(WORKTREE_CHANGED_EVENT); + await vi.waitFor(() => expect(fixture.captures).toHaveLength(2)); + const dirty = await fixture.nextCapture(); + for (let hint = 0; hint < 3; hint++) await fixture.event(WORKTREE_CHANGED_EVENT); + expect(fixture.captures).toHaveLength(2); + fixture.reply(dirty, worktreeCapture(captureRevision(dirty))); + const trailing = await fixture.nextCapture(); + expect(captureRevision(trailing)).toBe(captureRevision(dirty) + 1); + await expect(fixture.session.getWorktreeChanges()).resolves.toMatchObject({ + snapshot: { revision: 2, files: worktreeCapture(2).files }, + }); + fixture.reply(trailing, worktreeCapture(captureRevision(trailing), true)); + await fixture.settled(); + await expect(fixture.session.getWorktreeChanges()).resolves.toMatchObject({ + snapshot: { revision: 3, files: [] }, + }); + await expect(fixture.session.getCurrentMessageWork()).resolves.toMatchObject({ + messageId, + status: 'running', + }); + expect(fixture.captures).toHaveLength(3); + await runInDurableObject(fixture.session, async (_instance, state) => { + expect(state.storage.kv.get('session_messages')).toEqual(before.messages); + expect( + createEventQueries(drizzle(state.storage), state.storage.sql).findByFilters({}) + ).toEqual(before.events); + expect(await state.storage.getAlarm()).toEqual(before.alarm); + }); + await runInDurableObject(fixture.control, async (_instance, state) => { + expect(await state.storage.list()).toEqual(controlBefore.records); + expect(await state.storage.getAlarm()).toEqual(controlBefore.alarm); + }); + expect(before.broadcast).not.toHaveBeenCalled(); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); + await fixture.event('session.status', fixture.kiloSessionId, { status: { type: 'busy' } }); + expect(before.broadcast).toHaveBeenCalledTimes(1); + } finally { + fixture.close(); + } + }); + + it('rejects dirty hints without positive root and current runtime scope', async () => { + const fixture = await worktreeFixture(); + try { + await runInDurableObject(fixture.session, async (instance, state) => { + const messages: SessionMessageRecord[] = [ + { + messageId: 'msg_worktree_scoped', + state: 'accepted', + wrapperInstanceId: fixture.wrapperInstanceId, + acceptedAt: 1, + lastActivityAt: 2, + }, + ]; + state.storage.kv.put('session_messages', messages); + const identity = { + directory: fixture.directory, + kiloSessionId: fixture.kiloSessionId, + rootKiloSessionId: fixture.kiloSessionId, + }; + const input = { + identity, + wrapperInstanceId: fixture.wrapperInstanceId, + payload: { type: WORKTREE_CHANGED_EVENT, properties: {} }, + }; + for (const invalid of [ + { ...input, identity: { ...identity, directory: '/other' } }, + { ...input, identity: { ...identity, rootKiloSessionId: SECOND_ROOT_ID } }, + { ...input, identity: { ...identity, kiloSessionId: 'kilo_child' } }, + { ...input, identity: { ...identity, kiloSessionId: undefined } }, + { ...input, identity: { directory: fixture.directory } }, + { ...input, wrapperInstanceId: crypto.randomUUID() }, + { ...input, wrapperInstanceId: undefined }, + { ...input, payload: { ...input.payload, properties: { sessionID: 'kilo_child' } } }, + ]) { + await expect(instance.receiveSandboxControlEvent(invalid)).resolves.toEqual({ + applied: false, + }); + } + const metadata = await instance.getMetadata(); + if (!metadata) throw new Error('Missing test metadata'); + state.storage.kv.put('session_metadata', { ...metadata, auth: {} }); + await expect( + instance.receiveSandboxControlEvent({ + ...input, + identity: { directory: fixture.directory }, + }) + ).resolves.toEqual({ applied: false }); + expect(state.storage.kv.get('session_messages')).toEqual(messages); + expect( + createEventQueries(drizzle(state.storage), state.storage.sql).findByFilters({ + eventTypes: ['kilocode'], + }) + ).toEqual([]); + }); + await fixture.settled(); + expect(fixture.captures).toHaveLength(0); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); + } finally { + fixture.close(); + } + }); + + it('recaptures after a current-wrapper finalized outcome without capturing stale or duplicate outcomes', async () => { + const fixture = await worktreeFixture(); + const messageId = 'msg_worktree_finalized'; + const staleMessageId = 'msg_worktree_stale'; + const staleWrapperInstanceId = crypto.randomUUID(); + try { + await expect( + fixture.session.admitSubmittedMessage({ + userId: fixture.userId, + turn: { type: 'prompt', id: messageId, prompt: 'finalize the changes' }, + }) + ).resolves.toMatchObject({ success: true, messageId }); + const attached = await fixture.nextCapture(); + fixture.reply(attached, worktreeCapture(captureRevision(attached))); + await fixture.settled(); + const saved = await fixture.session.getWorktreeChanges(); + + await fixture.event('session.turn.close'); + const early = await fixture.nextCapture(); + fixture.fail(early); + await fixture.settled(); + expect(fixture.captures).toHaveLength(2); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual(saved); + await expect(fixture.session.getMessageResult(messageId)).resolves.toMatchObject({ + result: { status: 'running' }, + }); + + await runInDurableObject(fixture.session, (_instance, state) => { + const messages = state.storage.kv.get('session_messages') ?? []; + state.storage.kv.put('session_messages', [ + { + ...createSessionMessageRecord({ + turn: { type: 'prompt', messageId: staleMessageId, prompt: 'previous wrapper turn' }, + agent: { mode: 'code', model: 'test' }, + }), + state: 'accepted', + wrapperInstanceId: staleWrapperInstanceId, + acceptedAt: Date.now(), + } satisfies SessionMessageRecord, + ...messages, + ]); + }); + await expect( + fixture.session.receiveSandboxControlEvent({ + identity: { directory: fixture.directory, kiloSessionId: fixture.kiloSessionId }, + wrapperInstanceId: staleWrapperInstanceId, + payload: { + type: 'session.message.outcome', + properties: { messageId: staleMessageId, status: 'completed' }, + }, + }) + ).resolves.toEqual({ applied: true }); + await fixture.settled(); + expect(fixture.captures).toHaveLength(2); + await expect(fixture.session.getMessageResult(staleMessageId)).resolves.toMatchObject({ + result: { status: 'completed' }, + }); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual(saved); + + const outcome = { messageId, status: 'completed' }; + await fixture.event('session.message.outcome', fixture.kiloSessionId, outcome); + const finalized = await fixture.nextCapture(); + expect(captureRevision(finalized)).toBe(captureRevision(early) + 1); + expect(finalized.session).toEqual({ + sessionId: fixture.sessionId, + kiloSessionId: fixture.kiloSessionId, + directory: fixture.directory, + }); + const capture = worktreeCapture(captureRevision(finalized)); + capture.comparison.head = 'c'.repeat(40); + fixture.reply(finalized, capture); + await fixture.settled(); + await expect(fixture.session.getMessageResult(messageId)).resolves.toMatchObject({ + result: { status: 'completed' }, + }); + const completed = await fixture.session.getWorktreeChanges(); + expect(completed.snapshot).toMatchObject({ ...capture, schemaVersion: 1 }); + expect(fixture.captures).toHaveLength(3); + + await fixture.event('session.message.outcome', fixture.kiloSessionId, outcome); + await fixture.settled(); + expect(fixture.captures).toHaveLength(3); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual(completed); + } finally { + fixture.close(); + } + }); + + it('captures the shared directory through sibling roots and deletes only the selected chat summary', async () => { + const fixture = await worktreeFixture(); + const siblingId = `workspace_${crypto.randomUUID()}` as const; + const sibling = env.SANDBOX_SESSION.getByName(`${fixture.userId}:${siblingId}`); + const metadata = await runInDurableObject(fixture.session, instance => instance.getMetadata()); + if (!metadata) throw new Error('Missing test metadata'); + await sibling.registerSession({ + identity: { ...metadata.identity, sessionId: siblingId }, + auth: { ...metadata.auth, kiloSessionId: SECOND_ROOT_ID }, + agent: metadata.agent, + repository: metadata.repository, + workspace: metadata.workspace, + }); + await fixture.control.prepareSessionCredentials({ + ownerId: fixture.userId, + sessionId: siblingId, + }); + await fixture.control.attachSession({ + sessionId: siblingId, + kiloSessionId: SECOND_ROOT_ID, + directory: fixture.directory, + worktreeId: fixture.worktreeId, + ownerId: fixture.userId, + }); + fixture.noWake.attachSession.mockClear(); + + for (const [session, sessionId, kiloSessionId] of [ + [fixture.session, fixture.sessionId, fixture.kiloSessionId], + [sibling, siblingId, SECOND_ROOT_ID], + ] as const) { + const pending = session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + expect(request.session).toEqual({ sessionId, kiloSessionId, directory: fixture.directory }); + fixture.reply(request, worktreeCapture(captureRevision(request))); + await expect(pending).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 1 }, + }); + } + const savedSibling = await sibling.getWorktreeChanges(); + await fixture.session.deleteSession(); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ snapshot: null }); + await expect(sibling.getWorktreeChanges()).resolves.toEqual(savedSibling); + await expect(fixture.control.listRoutes()).resolves.toEqual([ + expect.objectContaining({ sessionId: siblingId, worktreeId: fixture.worktreeId }), + ]); + const pending = sibling.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(captureRevision(request), true)); + await expect(pending).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 2, files: [] }, + }); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); + fixture.close(); + }); + + it.each(['cancelled', 'exhausted'] as const)( + 'preserves capture and the healthy runtime after a rejected reattach is %s', + async retry => { + const fixture = await worktreeFixture(); + try { + await fixture.session.admitSubmittedMessage({ + userId: fixture.userId, + turn: { type: 'prompt', id: 'msg_initial_attach', prompt: 'initial prompt' }, + }); + const attachedCapture = await fixture.nextCapture(); + fixture.reply(attachedCapture, worktreeCapture(captureRevision(attachedCapture))); + await fixture.settled(); + await fixture.event('session.turn.close'); + const completedCapture = await fixture.nextCapture(); + fixture.reply(completedCapture, worktreeCapture(captureRevision(completedCapture))); + await fixture.settled(); + await fixture.event('session.message.outcome', fixture.kiloSessionId, { + messageId: 'msg_initial_attach', + status: 'completed', + }); + const finalizedCapture = await fixture.nextCapture(); + fixture.reply(finalizedCapture, worktreeCapture(captureRevision(finalizedCapture))); + await fixture.settled(); + await expect(fixture.session.getMessageResult('msg_initial_attach')).resolves.toMatchObject( + { + result: { status: 'completed' }, + } + ); + await fixture.session.invalidateTerminalRuntime({ + sandboxId: fixture.sandboxId, + wrapperInstanceId: fixture.wrapperInstanceId, + confirmed: true, + }); + const saved = await fixture.session.getWorktreeChanges(); + + const failedAttach = fixture.holdNextAttach(); + await fixture.session.admitSubmittedMessage({ + userId: fixture.userId, + turn: { type: 'prompt', id: 'msg_failed_reattach', prompt: 'follow-up prompt' }, + }); + const request = await failedAttach; + await expect(fixture.session.refreshWorktreeChanges()).resolves.toEqual({ + status: 'offline', + snapshot: saved.snapshot, + }); + fixture.fail(request, true); + await fixture.settled(); + await expect(fixture.control.getStatus()).resolves.toMatchObject({ + physical: 'running', + connection: 'ready', + }); + fixture.noWake.ensureReady.mockClear(); + fixture.noWake.attachSession.mockClear(); + + const refreshed = fixture.session.refreshWorktreeChanges(); + const next = await fixture.nextCapture(); + fixture.reply(next, worktreeCapture(captureRevision(next), true)); + await expect(refreshed).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { files: [] }, + }); + await fixture.event('session.error'); + const terminalCapture = await fixture.nextCapture(); + fixture.reply(terminalCapture, worktreeCapture(captureRevision(terminalCapture))); + await fixture.settled(); + const beforeCleanup = await fixture.session.getWorktreeChanges(); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + + if (retry === 'cancelled') { + await fixture.session.interruptExecution(); + await runInDurableObject(fixture.session, instance => instance.alarm()); + } else { + for (let attempt = 1; attempt < ATTACH_FAILURE_LIMIT; attempt++) { + const failedRetry = fixture.holdNextAttach(); + const alarm = runInDurableObject(fixture.session, instance => instance.alarm()); + fixture.fail(await failedRetry, true); + await alarm; + } + } + await fixture.settled(); + await expect( + fixture.session.getMessageResult('msg_failed_reattach') + ).resolves.toMatchObject({ + type: 'found', + result: { status: retry === 'cancelled' ? 'interrupted' : 'failed' }, + }); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual(beforeCleanup); + await expect(fixture.control.getStatus()).resolves.toMatchObject({ + physical: 'running', + connection: 'ready', + }); + expect(fixture.provider.stop).not.toHaveBeenCalled(); + fixture.noWake.ensureReady.mockClear(); + fixture.noWake.attachSession.mockClear(); + const afterRejection = fixture.session.refreshWorktreeChanges(); + const capture = await fixture.nextCapture(); + fixture.reply(capture, worktreeCapture(captureRevision(capture))); + await expect(afterRejection).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: captureRevision(capture) }, + }); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); + expect(fixture.prompts).toHaveLength(1); + } finally { + fixture.close(); + } + } + ); + + it('keeps saved changes and passive status offline after runtime-unhealthy attachment cleanup', async () => { + const fixture = await worktreeFixture(); + try { + await runInDurableObject(fixture.session, async (_instance, state) => { + await state.storage.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot); + }); + const attached = fixture.holdNextAttach(); + await fixture.session.admitSubmittedMessage({ + userId: fixture.userId, + turn: { type: 'prompt', id: 'msg_unhealthy_attach', prompt: 'follow-up' }, + }); + fixture.fail(await attached, false, 'runtime_unhealthy'); + await fixture.settled(); + expect(fixture.provider.stop).toHaveBeenCalled(); + await expect(fixture.control.getStatus()).resolves.toMatchObject({ + physical: 'stopped', + connection: 'disconnected', + }); + fixture.noWake.ensureReady.mockClear(); + fixture.noWake.attachSession.mockClear(); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ + snapshot: savedWorktreeSnapshot, + }); + await expect(fixture.session.refreshWorktreeChanges()).resolves.toEqual({ + status: 'offline', + snapshot: savedWorktreeSnapshot, + }); + await expect(fixture.session.getSandboxStatus()).resolves.toMatchObject({ + status: 'sleeping', + }); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); + expect(fixture.captures).toHaveLength(0); + } finally { + fixture.close(); + } + }); + + it.each(['session.turn.close', 'session.error', WORKTREE_CHANGED_EVENT])( + 'captures root %s with no accepted queue entry and excludes child events', + async type => { + const fixture = await worktreeFixture(); + await fixture.event(type, 'kilo_child'); + await fixture.settled(); + expect(fixture.captures).toHaveLength(0); + await fixture.event(type); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(1)); + await fixture.settled(); + const saved = await fixture.session.getWorktreeChanges(); + expect(saved.snapshot).toMatchObject({ revision: 1, files: worktreeCapture(1).files }); + expect(saved.snapshot?.capturedAt).toEqual(expect.any(String)); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + fixture.close(); + } + ); + + it('coalesces concurrent manual refreshes and retains one trailing lifecycle capture', async () => { + const fixture = await worktreeFixture(); + const refreshed = runInDurableObject(fixture.session, instance => + Promise.all([instance.refreshWorktreeChanges(), instance.refreshWorktreeChanges()]) + ); + const first = await fixture.nextCapture(); + await fixture.event('session.turn.close'); + await fixture.event('session.error'); + expect(fixture.captures).toHaveLength(1); + fixture.reply(first, worktreeCapture(captureRevision(first))); + const trailing = await fixture.nextCapture(); + expect(captureRevision(trailing)).toBe(captureRevision(first) + 1); + fixture.reply(trailing, worktreeCapture(captureRevision(trailing), true)); + const results = await refreshed; + expect(results[0]).toEqual(results[1]); + expect(results[0]).toMatchObject({ status: 'refreshed', snapshot: { files: [], revision: 2 } }); + await fixture.settled(); + expect(fixture.captures).toHaveLength(2); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + fixture.close(); + }); + + it.each(['session.idle', 'session.status'])( + 'does not capture queue cancellation or abort acknowledgement, only settled root %s', + async type => { + const fixture = await worktreeFixture(); + await runInDurableObject(fixture.session, async (instance, state) => { + await state.storage.put('session_messages', [ + { + messageId: 'msg_interrupted', + state: 'accepted', + acceptedAt: Date.now(), + } satisfies SessionMessageRecord, + ]); + await instance.markAsInterrupted(); + }); + await fixture.event(type, fixture.kiloSessionId, { status: { type: 'idle' } }); + await fixture.settled(); + expect(fixture.captures).toHaveLength(0); + await fixture.session.interruptExecution(); + expect(fixture.aborts).toHaveLength(1); + await fixture.settled(); + expect(fixture.captures).toHaveLength(0); + await fixture.event(type, 'kilo_child', { status: { type: 'idle' } }); + await fixture.event('session.status', fixture.kiloSessionId, { status: { type: 'busy' } }); + expect(fixture.captures).toHaveLength(0); + await fixture.event(WORKTREE_CHANGED_EVENT); + const dirty = await fixture.nextCapture(); + fixture.reply(dirty, worktreeCapture(captureRevision(dirty))); + await fixture.settled(); + expect(fixture.captures).toHaveLength(1); + await fixture.event(type, fixture.kiloSessionId, { status: { type: 'idle' } }); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(captureRevision(request))); + await fixture.settled(); + expect(fixture.captures).toHaveLength(2); + await expect(fixture.session.getWorktreeChanges()).resolves.toMatchObject({ + snapshot: { revision: 2 }, + }); + fixture.close(); + } + ); + + it.each(['session', 'worktree'] as const)( + 'discards capture results after metadata changes or %s deletion', + async deletion => { + const fixture = await worktreeFixture(); + await runInDurableObject(fixture.session, async (_instance, state) => + state.storage.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot) + ); + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + await runInDurableObject(fixture.session, async (instance, state) => { + const metadata = await instance.getMetadata(); + if (!metadata) throw new Error('Missing test metadata'); + await state.storage.put('session_metadata', { + ...metadata, + repository: { ...metadata.repository, upstreamBranch: 'other' }, + }); + }); + fixture.reply(request, worktreeCapture(captureRevision(request))); + await expect(pending).resolves.toEqual({ status: 'failed', snapshot: savedWorktreeSnapshot }); + await runInDurableObject(fixture.session, async (instance, state) => { + const metadata = await instance.getMetadata(); + if (!metadata) throw new Error('Missing test metadata'); + await state.storage.put('session_metadata', { + ...metadata, + repository: { ...metadata.repository, upstreamBranch: 'main' }, + }); + }); + const deletedCapture = fixture.session.refreshWorktreeChanges(); + const lateRequest = await fixture.nextCapture(); + if (deletion === 'worktree') { + await fixture.session.beginWorktreeDeletion({ + worktreeId: fixture.worktreeId, + kiloSessionId: fixture.kiloSessionId, + ownerId: fixture.userId, + }); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ snapshot: null }); + await expect(fixture.session.refreshWorktreeChanges()).resolves.toEqual({ + status: 'offline', + snapshot: null, + }); + await fixture.session.finishWorktreeDeletion(fixture.worktreeId); + } else { + await fixture.session.deleteSession(); + } + fixture.reply(lateRequest, worktreeCapture(captureRevision(lateRequest))); + await expect(deletedCapture).resolves.toEqual({ + status: 'failed', + snapshot: savedWorktreeSnapshot, + }); + await fixture.settled(); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ snapshot: null }); + await expect(fixture.session.getMetadata()).resolves.toBeNull(); + expect(fixture.captures).toHaveLength(2); + expect(fixture.aborts).toHaveLength(0); + fixture.close(); + } + ); + + it('preserves the exact saved snapshot on failed, malformed and wrong-revision results, then accepts empty', async () => { + const fixture = await worktreeFixture(); + await runInDurableObject(fixture.session, async (_instance, state) => + state.storage.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot) + ); + for (const kind of ['git', 'malformed', 'revision'] as const) { + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + if (kind === 'git') fixture.fail(request); + else if (kind === 'malformed') fixture.reply(request, { files: [] }); + else fixture.reply(request, worktreeCapture(captureRevision(request) + 1)); + await expect(pending).resolves.toEqual({ status: 'failed', snapshot: savedWorktreeSnapshot }); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ + snapshot: savedWorktreeSnapshot, + }); + } + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(captureRevision(request), true)); + await expect(pending).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 8, files: [] }, + }); + fixture.close(); + }); + + it('fences credential rotation and requires the new connection to be ready before sending', async () => { + const fixture = await worktreeFixture(); + await runInDurableObject(fixture.session, async (_instance, state) => + state.storage.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot) + ); + const pending = fixture.session.refreshWorktreeChanges(); + await fixture.nextCapture(); + await fixture.rotateSocket(); + await expect(pending).resolves.toEqual({ status: 'failed', snapshot: savedWorktreeSnapshot }); + await expect(fixture.control.getStatus()).resolves.toMatchObject({ + physical: 'running', + connection: 'connected', + }); + await expect(fixture.session.refreshWorktreeChanges()).resolves.toEqual({ + status: 'offline', + snapshot: savedWorktreeSnapshot, + }); + expect(fixture.captures).toHaveLength(1); + await fixture.ready(); + const next = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(captureRevision(request))); + await expect(next).resolves.toMatchObject({ status: 'refreshed' }); + fixture.close(); + }); + + it('requires a running physical sandbox and matching route at the request boundary', async () => { + const fixture = await worktreeFixture(); + for (const identity of [ + { + sessionId: 'workspace_other', + kiloSessionId: fixture.kiloSessionId, + directory: fixture.directory, + }, + { sessionId: fixture.sessionId, kiloSessionId: 'other_root', directory: fixture.directory }, + { sessionId: fixture.sessionId, kiloSessionId: fixture.kiloSessionId, directory: '/other' }, + ]) { + await expect( + fixture.control.request({ + operation: 'session.git.summary', + session: identity, + payload: { revision: 1 }, + }) + ).resolves.toMatchObject({ ok: false, error: { code: 'not_ready' } }); + } + await fixture.control.beginStop('test'); + await expect(fixture.session.refreshWorktreeChanges()).resolves.toEqual({ + status: 'offline', + snapshot: null, + }); + expect(fixture.captures).toHaveLength(0); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + await fixture.control.confirmStopped(); + fixture.close(); + }); + + it.each(['physical stop', 'route detach'] as const)( + 'discards a valid in-flight capture after %s and preserves the saved snapshot', + async change => { + const fixture = await worktreeFixture(); + await runInDurableObject(fixture.session, async (_instance, state) => + state.storage.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot) + ); + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + if (change === 'physical stop') await fixture.control.beginStop('test in-flight capture'); + else await fixture.control.detachSession(fixture.sessionId); + fixture.reply(request, worktreeCapture(captureRevision(request), true)); + await expect(pending).resolves.toEqual({ status: 'failed', snapshot: savedWorktreeSnapshot }); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ + snapshot: savedWorktreeSnapshot, + }); + expect(fixture.captures).toHaveLength(1); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); + if (change === 'physical stop') await fixture.control.confirmStopped(); + fixture.close(); + } + ); + + it('persists through DO eviction and serves offline GET and refresh without starting a sandbox', async () => { + const fixture = await worktreeFixture(); + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(captureRevision(request))); + const saved = await pending; + expect(saved.status).toBe('refreshed'); + await fixture.control.beginStop('test'); + await fixture.control.confirmStopped(); + await fixture.settled(); + let previousInstance: unknown; + await runInDurableObject(fixture.session, instance => { + previousInstance = instance; + }); + await abortAllDurableObjects(); + const freshSession = env.SANDBOX_SESSION.getByName(`${fixture.userId}:${fixture.sessionId}`); + const freshControl = env.SANDBOX_CONTROL.getByName(fixture.sandboxId); + const noWake = await runInDurableObject(freshControl, instance => { + const prototype = Object.getPrototypeOf(instance) as typeof instance; + return { + ensureReady: vi.spyOn(prototype, 'ensureReady'), + attachSession: vi.spyOn(prototype, 'attachSession'), + request: vi.spyOn(prototype, 'request'), + }; + }); + await runInDurableObject(freshSession, async instance => { + expect(instance).not.toBe(previousInstance); + await expect(instance.getWorktreeChanges()).resolves.toEqual({ snapshot: saved.snapshot }); + }); + expect(noWake.request).not.toHaveBeenCalled(); + await expect(freshSession.refreshWorktreeChanges()).resolves.toEqual({ + status: 'offline', + snapshot: saved.snapshot, + }); + await expect(freshControl.getStatus()).resolves.toMatchObject({ + physical: 'stopped', + connection: 'disconnected', + }); + expect(noWake.ensureReady).not.toHaveBeenCalled(); + expect(noWake.attachSession).not.toHaveBeenCalled(); + }); +}); + describe('SandboxSession control-plane regressions', () => { beforeEach(() => { vi.spyOn(globalThis, 'fetch').mockImplementation(async () => Response.json({ valid: true })); @@ -10036,6 +11070,177 @@ describe('SandboxControl worktree activity deadlines', () => { }); }); +function receiveAdmissionPrompt(ws: WebSocket, attach: WrapperRequest) { + expect(attach.operation).toBe('session.attach'); + expect(attach.session).toBeDefined(); + const prompt = Promise.withResolvers(); + let promptCount = 0; + let failure: unknown; + const cleanup = () => { + ws.removeEventListener('message', onMessage); + ws.removeEventListener('error', onError); + ws.removeEventListener('close', onClose); + }; + const fail = (error: unknown) => { + failure = error; + cleanup(); + prompt.reject(error); + }; + const onMessage = (event: MessageEvent) => { + try { + const request = requestFrameSchema.parse(JSON.parse(String(event.data))); + expect(request.session).toEqual(attach.session); + if (request.operation === 'session.git.summary') { + const payload = worktreeChangesCaptureRequestSchema.parse(request.payload); + const capture = worktreeCapture(payload.revision, true); + if (payload.baseRef) capture.comparison.baseRef = payload.baseRef; + respondToWrapperRequest(ws, request, capture); + } else if (request.operation === 'session.prompt') { + promptCount++; + prompt.resolve(request); + } else { + throw new Error(`Unexpected admission request: ${request.operation}`); + } + } catch (error) { + fail(error); + } + }; + const onError = () => fail(new Error('sandbox control websocket error')); + const onClose = (event: CloseEvent) => + fail(new Error(`sandbox control websocket closed: ${event.code}`)); + ws.addEventListener('message', onMessage); + ws.addEventListener('error', onError); + ws.addEventListener('close', onClose); + return { + prompt: prompt.promise, + finish(): void { + cleanup(); + if (failure !== undefined) throw failure; + expect(promptCount).toBe(1); + }, + }; +} + +describe('SandboxSession worktree admission receiver', () => { + const attach: WrapperRequest = { + type: 'request', + requestId: 'attach', + operation: 'session.attach', + session: { + sessionId: GRANT_SESSION_ID, + kiloSessionId: ROOT_ID, + directory: '/workspace/shared', + }, + }; + const capture = { + ...attach, + requestId: 'capture', + operation: 'session.git.summary', + payload: { revision: 7, baseRef: 'refs/remotes/origin/feature/shared-worktree' }, + }; + const prompt = { + ...attach, + requestId: 'prompt', + operation: 'session.prompt', + payload: { messageId: INITIAL_MESSAGE_ID, finalization: { autoCommit: false } }, + }; + + function fixture() { + const pair = new WebSocketPair(); + const client = pair[0]; + const server = pair[1]; + client.accept(); + server.accept(); + const receiver = receiveAdmissionPrompt(client, attach); + return { + client, + server, + receiver, + receive(frame: unknown) { + client.dispatchEvent(new MessageEvent('message', { data: JSON.stringify(frame) })); + }, + close() { + client.close(); + server.close(); + }, + }; + } + + it.each(['capture first', 'prompt first'])( + 'acknowledges capture and retains the prompt before awaiting it: %s', + async order => { + const f = fixture(); + try { + const reply = nextMessage(f.server); + for (const frame of order === 'capture first' ? [capture, prompt] : [prompt, capture]) { + f.receive(frame); + } + await expect(f.receiver.prompt).resolves.toEqual(prompt); + expect(JSON.parse(await reply)).toEqual({ + type: 'response', + requestId: capture.requestId, + ok: true, + result: { + ...worktreeCapture(capture.payload.revision, true), + comparison: { + ...worktreeCapture(capture.payload.revision).comparison, + baseRef: capture.payload.baseRef, + }, + }, + }); + f.receiver.finish(); + } finally { + f.close(); + } + } + ); + + it('does not hide duplicate prompt delivery', async () => { + const f = fixture(); + try { + f.receive(prompt); + await expect(f.receiver.prompt).resolves.toEqual(prompt); + f.receive(prompt); + expect(() => f.receiver.finish()).toThrow(); + } finally { + f.close(); + } + }); + + it.each([ + { name: 'unknown operation', frame: { ...prompt, operation: 'session.unexpected' } }, + { name: 'malformed capture', frame: { ...capture, payload: { revision: 0 } } }, + { + name: 'wrong session', + frame: { ...capture, session: { ...attach.session, sessionId: 'other' } }, + }, + ])('rejects $name instead of skipping it', async ({ frame }) => { + const f = fixture(); + try { + const rejected = expect(f.receiver.prompt).rejects.toThrow(); + f.receive(frame); + await rejected; + expect(() => f.receiver.finish()).toThrow(); + } finally { + f.close(); + } + }); + + it.each(['close', 'error'] as const)('rejects a pending prompt on socket %s', async event => { + const f = fixture(); + try { + const rejected = expect(f.receiver.prompt).rejects.toThrow('sandbox control websocket'); + f.client.dispatchEvent( + event === 'close' ? new CloseEvent('close', { code: 1001 }) : new Event('error') + ); + await rejected; + expect(() => f.receiver.finish()).toThrow('sandbox control websocket'); + } finally { + f.close(); + } + }); +}); + describe('SandboxSession worktree admission', () => { beforeEach(() => { vi.spyOn(globalThis, 'fetch').mockImplementation(async () => Response.json({ valid: true })); @@ -10223,9 +11428,9 @@ describe('SandboxSession worktree admission', () => { ]); }); - const incomingPrompt = nextMessage(ws); + const receiver = receiveAdmissionPrompt(ws, attach); respondToWrapperRequest(ws, attach, { attached: true }); - const prompt = JSON.parse(await incomingPrompt) as WrapperRequest; + const prompt = await receiver.prompt; expect(prompt).toMatchObject({ operation: 'session.prompt', session: { sessionId, kiloSessionId }, @@ -10238,6 +11443,13 @@ describe('SandboxSession worktree admission', () => { messageId: INITIAL_MESSAGE_ID, status: 'accepted', }); + await vi.waitFor(async () => { + expect(await session.getCurrentMessageWork()).toMatchObject({ + messageId: INITIAL_MESSAGE_ID, + status: 'running', + }); + }); + receiver.finish(); ws.close(); } ); @@ -10327,9 +11539,9 @@ describe('SandboxSession worktree admission', () => { }); }); - const incomingCommand = nextMessage(wrapper); + const receiver = receiveAdmissionPrompt(wrapper, attach); respondToWrapperRequest(wrapper, attach, { attached: true }); - const command = JSON.parse(await incomingCommand) as WrapperRequest; + const command = await receiver.prompt; expect(command).toMatchObject({ operation: 'session.prompt', session: { sessionId, kiloSessionId }, @@ -10344,6 +11556,13 @@ describe('SandboxSession worktree admission', () => { messageId: INITIAL_MESSAGE_ID, status: 'accepted', }); + await vi.waitFor(async () => { + expect(await session.getCurrentMessageWork()).toMatchObject({ + messageId: INITIAL_MESSAGE_ID, + status: 'running', + }); + }); + receiver.finish(); wrapper.close(); }); @@ -10418,9 +11637,9 @@ describe('SandboxSession worktree admission', () => { }), ]); }); - const incomingPrompt = nextMessage(wrapper); + const receiver = receiveAdmissionPrompt(wrapper, attach); respondToWrapperRequest(wrapper, attach, { attached: true }); - const prompt = JSON.parse(await incomingPrompt) as WrapperRequest & { + const prompt = (await receiver.prompt) as WrapperRequest & { payload: { attachments: Array<{ mime: string; @@ -10457,6 +11676,13 @@ describe('SandboxSession worktree admission', () => { messageId: INITIAL_MESSAGE_ID, status: 'accepted', }); + await vi.waitFor(async () => { + expect(await session.getCurrentMessageWork()).toMatchObject({ + messageId: INITIAL_MESSAGE_ID, + status: 'running', + }); + }); + receiver.finish(); wrapper.close(); }); @@ -10766,9 +11992,9 @@ describe('SandboxSession worktree admission', () => { finalization: { autoCommit: !expected, condenseOnComplete: true }, }); }); - const incomingPrompt = nextMessage(wrapper); + const receiver = receiveAdmissionPrompt(wrapper, attach); respondToWrapperRequest(wrapper, attach, { attached: true }); - const prompt = JSON.parse(await incomingPrompt) as WrapperRequest; + const prompt = await receiver.prompt; expect(prompt).toMatchObject({ operation: 'session.prompt', payload: { @@ -10789,6 +12015,7 @@ describe('SandboxSession worktree admission', () => { condenseOnComplete: true, }); }); + receiver.finish(); wrapper.close(); } ); diff --git a/services/cloud-agent-next/test/unit/wrapper/utils.test.ts b/services/cloud-agent-next/test/unit/wrapper/utils.test.ts index c2a7a1eef1..7a34ef05b4 100644 --- a/services/cloud-agent-next/test/unit/wrapper/utils.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/utils.test.ts @@ -27,6 +27,50 @@ describe('runProcess', () => { expect(result.elapsedMs).toBeGreaterThanOrEqual(0); }); + it('preserves UTF-8 characters split across stdout and stderr chunks', async () => { + const outputs = { stdout: '', stderr: '' }; + const text = 'é漢字𐐀\tpath\n'; + const result = await runProcess( + process.execPath, + [ + '-e', + `const bytes = Buffer.from(${JSON.stringify(text)}); let i = 0; + const timer = setInterval(() => { + const byte = bytes.subarray(i, ++i); + process.stdout.write(byte); + process.stderr.write(byte); + if (i === bytes.length) clearInterval(timer); + }, 20);`, + ], + { + timeoutMs: 5_000, + onOutput: (stream, output) => { + outputs[stream] += output; + }, + } + ); + + expect(result.exitCode).toBe(0); + expect(result.stdout).toBe(text); + expect(result.stderr).toBe(text); + expect(outputs).toEqual({ stdout: text, stderr: text }); + expect(result.stdoutTruncated).toBeUndefined(); + expect(result.stderrTruncated).toBeUndefined(); + }); + + it('keeps a valid UTF-8 tail within the byte cap', async () => { + const result = await runProcess( + process.execPath, + ['-e', 'process.stdout.write("𐐀".repeat(40) + "end")'], + { timeoutMs: 5_000, maxOutputBytes: 10 } + ); + + expect(result.exitCode).toBe(0); + expect(result.stdout).toBe('𐐀end'); + expect(Buffer.byteLength(result.stdout)).toBeLessThanOrEqual(10); + expect(result.stdoutTruncated).toBe(true); + }); + it('bounds output while retaining the most recent tail', async () => { const result = await runProcess( process.execPath, diff --git a/services/cloud-agent-next/wrapper/src/control/main.ts b/services/cloud-agent-next/wrapper/src/control/main.ts index 1b6c003b94..8390501103 100644 --- a/services/cloud-agent-next/wrapper/src/control/main.ts +++ b/services/cloud-agent-next/wrapper/src/control/main.ts @@ -18,6 +18,7 @@ import { createControlTerminalRuntime } from './terminal-runtime'; import { createWorktreeKiloRuntimes } from './worktree-runtime'; import { createControlDiagnostics, type ControlDiagnostics } from './diagnostics'; import { controlLogWrapperIdSchema } from '../../../src/shared/control-diagnostics.js'; +import { createWorktreeMutationNotifications } from './worktree-mutation-notifications'; const retirementCauses = new Map([ ['Kilo event feed is no longer healthy', 'event_feed_unhealthy'], @@ -53,6 +54,7 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void const kiloRuntimes = createWorktreeKiloRuntimes({ onDiagnostic: diagnostics.onDiagnostic, onEvent: (runtime, event) => { + mutationNotifications.observe(runtime, event); const identity = sessionEventIdentity({ ...event, sessionId: eventKiloSessionId(event.properties), @@ -117,6 +119,13 @@ function main(diagnostics: ControlDiagnostics, wrapperInstanceId: string): void onShutdown: () => shutdown(0, 'Sandbox shutting down'), }; + const mutationNotifications = createWorktreeMutationNotifications({ + sessions: deps.sessions, + kiloRuntimes, + signal: abort.signal, + sendEvent: (event, payload, identity) => control?.sendEvent?.(event, payload, identity), + }); + function withHeartbeatReason(payload: SandboxHeartbeatPayload): SandboxHeartbeatPayload { if (!payload.kilo.ready && heartbeatReason) payload.kilo.reason = heartbeatReason; return payload; diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts index 95e73fa717..b6aa3b048e 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts @@ -417,6 +417,141 @@ describe('createSandboxControlClient', () => { client.close(); }); + it.each([ + { ok: true, result: '漢'.repeat(MAX_SANDBOX_CONTROL_FRAME_BYTES / 2) }, + { + ok: false, + error: { code: 'failure', message: 'private-data'.repeat(100_000), retryable: false }, + }, + ])( + 'replaces oversized response envelopes with a small error and preserves the socket', + async outcome => { + const fake = new FakeWebSocket(); + const client = createSandboxControlClient({ + url: 'wss://example.test/sandbox-control/sbx_1', + credential: 'secret', + providerInstanceId: 'inst_1', + openWebSocket: () => fake as unknown as WebSocket, + onRequest: async operation => + operation === 'sandbox.status' ? { ok: true, result: { healthy: true } } : outcome, + }); + const connecting = client.connect(); + await handshake(fake); + await connecting; + fake.respond( + JSON.stringify({ + type: 'request', + requestId: 'large', + operation: 'session.git.summary', + payload: { revision: 1 }, + }) + ); + await Promise.resolve(); + await Promise.resolve(); + + expect(JSON.parse(fake.sent[2] ?? '{}')).toEqual({ + type: 'response', + requestId: 'large', + ok: false, + error: { + code: 'payload_too_large', + message: 'Response exceeds size limit', + retryable: false, + }, + }); + expect(Buffer.byteLength(fake.sent[2] ?? '')).toBeLessThan(1024); + expect(fake.readyState).toBe(1); + fake.respond( + JSON.stringify({ + type: 'request', + requestId: 'next', + operation: 'sandbox.status', + payload: {}, + }) + ); + await Promise.resolve(); + await Promise.resolve(); + expect(JSON.parse(fake.sent[3] ?? '{}')).toEqual({ + type: 'response', + requestId: 'next', + ok: true, + result: { healthy: true }, + }); + client.close(); + } + ); + + it.each([0, 1])( + 'keeps the full response strictly below the frame limit with %s bytes of headroom', + async headroom => { + const fake = new FakeWebSocket(); + const requestId = 'quote"漢'; + const envelopeBytes = Buffer.byteLength( + JSON.stringify({ type: 'response', requestId, ok: true, result: '' }) + ); + const client = createSandboxControlClient({ + url: 'wss://example.test/sandbox-control/sbx_1', + credential: 'secret', + providerInstanceId: 'inst_1', + openWebSocket: () => fake as unknown as WebSocket, + onRequest: async () => ({ + ok: true, + result: 'x'.repeat(MAX_SANDBOX_CONTROL_FRAME_BYTES - envelopeBytes - headroom), + }), + }); + const connecting = client.connect(); + await handshake(fake); + await connecting; + fake.respond( + JSON.stringify({ + type: 'request', + requestId, + operation: 'session.git.summary', + payload: { revision: 1 }, + }) + ); + await Promise.resolve(); + await Promise.resolve(); + + expect(JSON.parse(fake.sent[2] ?? '{}').ok).toBe(headroom === 1); + expect(Buffer.byteLength(fake.sent[2] ?? '')).toBeLessThan(MAX_SANDBOX_CONTROL_FRAME_BYTES); + client.close(); + } + ); + + it('safely handles unserializable results without closing the shared socket', async () => { + const fake = new FakeWebSocket(); + const client = createSandboxControlClient({ + url: 'wss://example.test/sandbox-control/sbx_1', + credential: 'secret', + providerInstanceId: 'inst_1', + openWebSocket: () => fake as unknown as WebSocket, + onRequest: async () => ({ ok: true, result: 1n }), + }); + const connecting = client.connect(); + await handshake(fake); + await connecting; + fake.respond( + JSON.stringify({ + type: 'request', + requestId: 'invalid-result', + operation: 'session.git.summary', + payload: { revision: 1 }, + }) + ); + await Promise.resolve(); + await Promise.resolve(); + + expect(JSON.parse(fake.sent[2] ?? '{}')).toEqual({ + type: 'response', + requestId: 'invalid-result', + ok: false, + error: { code: 'capture_failed', message: 'Response serialization failed', retryable: false }, + }); + expect(fake.readyState).toBe(1); + client.close(); + }); + it('includes session on sendEvent when provided', async () => { const fake = new FakeWebSocket(); const client = createSandboxControlClient({ @@ -992,69 +1127,86 @@ describe('createSandboxControlClient', () => { } ); - it.each(['event', 'outcome-budget', 'response', 'ping', 'closed-socket'] as const)( - 'retires once on established %s delivery failure', - async failure => { - const timers = spyOn(globalThis, 'setInterval'); - const { client, sockets, openWebSocket, onDisconnected } = createClientFixture({ - onRequest: async () => ({ ok: true }), - }); - try { - const connecting = client.connect(); - await Promise.resolve(); - await handshake(sockets[0]); - await connecting; - const socket = sockets[0]; - if (!socket) throw new Error('missing socket'); - if (failure === 'closed-socket') socket.readyState = 3; - else if (failure !== 'outcome-budget') - socket.send = () => { - throw new Error('send failed'); - }; - if (failure === 'response') { - socket.respond( - JSON.stringify({ - type: 'request', - requestId: 'normal-status', - operation: 'sandbox.status', - payload: {}, - }) - ); - } else if (failure === 'ping') { - const ping = timers.mock.calls[0]?.[0]; - if (typeof ping !== 'function') throw new Error('missing keepalive'); - ping(); - ping(); - } else { - expect( - client.sendEvent?.( - 'session.event', - { - type: 'session.message.outcome', - properties: { messageId: 'msg_1', status: 'completed' }, - }, - { - directory: - failure === 'outcome-budget' - ? 'd'.repeat(MAX_SANDBOX_CONTROL_FRAME_BYTES) - : '/workspace', - } - ) - ).toBe(false); - } - await waitForReconnect(); - socket.error(); - socket.close(); - expect(onDisconnected).toHaveBeenCalledTimes(1); - expect(openWebSocket).toHaveBeenCalledTimes(1); - expect(client.sendEvent?.('sandbox.ready', { kiloReady: true })).toBe(false); - expect(client.connect()).rejects.toThrow('sandbox control client closed'); - } finally { - client.close(); - timers.mockRestore(); + it.each([ + 'event', + 'outcome-budget', + 'response', + 'oversized-response', + 'invalid-response', + 'ping', + 'closed-socket', + ] as const)('retires once on established %s delivery failure', async failure => { + const timers = spyOn(globalThis, 'setInterval'); + const { client, sockets, openWebSocket, onDisconnected } = createClientFixture({ + onRequest: async () => ({ + ok: true, + result: + failure === 'oversized-response' + ? 'x'.repeat(MAX_SANDBOX_CONTROL_FRAME_BYTES) + : failure === 'invalid-response' + ? 1n + : undefined, + }), + }); + try { + const connecting = client.connect(); + await Promise.resolve(); + await handshake(sockets[0]); + await connecting; + const socket = sockets[0]; + if (!socket) throw new Error('missing socket'); + if (failure === 'closed-socket') socket.readyState = 3; + else if (failure !== 'outcome-budget') + socket.send = () => { + throw new Error('send failed'); + }; + if ( + failure === 'response' || + failure === 'oversized-response' || + failure === 'invalid-response' + ) { + socket.respond( + JSON.stringify({ + type: 'request', + requestId: 'normal-status', + operation: 'sandbox.status', + payload: {}, + }) + ); + } else if (failure === 'ping') { + const ping = timers.mock.calls[0]?.[0]; + if (typeof ping !== 'function') throw new Error('missing keepalive'); + ping(); + ping(); + } else { + expect( + client.sendEvent?.( + 'session.event', + { + type: 'session.message.outcome', + properties: { messageId: 'msg_1', status: 'completed' }, + }, + { + directory: + failure === 'outcome-budget' + ? 'd'.repeat(MAX_SANDBOX_CONTROL_FRAME_BYTES) + : '/workspace', + } + ) + ).toBe(false); } + await waitForReconnect(); + socket.error(); + socket.close(); + expect(onDisconnected).toHaveBeenCalledTimes(1); + expect(openWebSocket).toHaveBeenCalledTimes(1); + expect(client.sendEvent?.('sandbox.ready', { kiloReady: true })).toBe(false); + expect(client.connect()).rejects.toThrow('sandbox control client closed'); + } finally { + client.close(); + timers.mockRestore(); } - ); + }); it('fences handler completion and further requests after retirement even if the socket appears open', async () => { const outcome = Promise.withResolvers<{ ok: boolean }>(); diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts index 381f768572..630df7d63f 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts @@ -204,23 +204,48 @@ export function createSandboxControlClient( retireConnection(ws); return; } + let response: string; try { - ws.send( - JSON.stringify({ - type: 'response', - requestId: request.requestId, - ...(outcome.ok - ? { ok: true, ...(outcome.result !== undefined ? { result: outcome.result } : {}) } - : { - ok: false, - error: outcome.error ?? { - code: 'not_ready', - message: 'Request failed', - retryable: true, - }, - }), - }) - ); + response = JSON.stringify({ + type: 'response', + requestId: request.requestId, + ...(outcome.ok + ? { ok: true, ...(outcome.result !== undefined ? { result: outcome.result } : {}) } + : { + ok: false, + error: outcome.error ?? { + code: 'not_ready', + message: 'Request failed', + retryable: true, + }, + }), + }); + } catch { + response = JSON.stringify({ + type: 'response', + requestId: request.requestId, + ok: false, + error: { + code: 'capture_failed', + message: 'Response serialization failed', + retryable: false, + }, + }); + } + if (Buffer.byteLength(response) >= MAX_SANDBOX_CONTROL_FRAME_BYTES) { + response = JSON.stringify({ + type: 'response', + requestId: request.requestId, + ok: false, + error: { + code: 'payload_too_large', + message: 'Response exceeds size limit', + retryable: false, + }, + }); + } + try { + ws.send(response); requestDiagnostic('response_sent', outcome.ok); } catch { requestDiagnostic('response_failed', outcome.ok); diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts index 1a8e196b0b..bcfef1cba4 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts @@ -10,6 +10,7 @@ import { SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, sessionSyncResultSchema, type SessionEventPayload, + type SessionGitSummaryResult, } from '../../../src/shared/sandbox-control-protocol'; import { createWrapperKiloClient, type WrapperKiloClient, type WrapperPty } from '../kilo-api'; import { materializeMessageAttachments } from '../session-bootstrap'; @@ -17,8 +18,10 @@ import { runProcess, withTimeoutAndAbort } from '../utils'; import { applySessionAttach } from './apply-attach'; import { updateSessionSnapshots, unfilteredKiloEvents } from './feed'; import { + forgetAttachedRoot, rememberAttachedRoot, rememberChildSession, + rememberSessionDirectory, resetSessionDirectoryState, } from './session-directories'; import { CONTROL_RUNTIME_RESERVED_ENV_VARS } from '../../../src/shared/runtime-environment.js'; @@ -4806,3 +4809,298 @@ describe('createSessionActivityRegistry', () => { ]); }); }); + +describe('session.git.summary', () => { + const captured: SessionGitSummaryResult = { + revision: 1, + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: [], + truncated: false, + }; + + beforeEach(() => { + resetSessionDirectoryState(); + }); + + it('allows each attached shared-worktree root without waking Kilo or changing activity', async () => { + const sibling = { ...session, sessionId: 'ses_2', kiloSessionId: 'kilo_2' }; + const activity = createSessionActivityRegistry(() => 100); + const directories: string[] = []; + for (const identity of [session, sibling]) { + rememberAttachedRoot(identity.kiloSessionId, identity.directory); + activity.attach(identity.kiloSessionId); + } + activity.markActive(session.kiloSessionId); + const snapshots = activity.snapshots(); + const handlerDeps = deps({ + kiloClient: undefined, + kiloReady: false, + activity, + collectWorktreeChanges: async directory => { + directories.push(directory); + return captured; + }, + }); + + expect(rootForSession(undefined, session.directory)).toBeUndefined(); + for (const identity of [session, sibling]) { + expect( + await handleControlRequest('session.git.summary', identity, { revision: 1 }, handlerDeps) + ).toEqual({ ok: true, result: captured }); + } + expect(directories).toEqual([session.directory, session.directory]); + expect(activity.snapshots()).toEqual(snapshots); + expect(handlerDeps.sessions).toEqual([]); + expect(handlerDeps.tasks.size).toBe(0); + }); + + it('drains in-flight capture before deletion and rejects late results without fencing another worktree', async () => { + const sibling = { sessionId: 'ses_2', kiloSessionId: 'kilo_2', directory: '/other' }; + rememberAttachedRoot(session.kiloSessionId, session.directory); + rememberAttachedRoot(sibling.kiloSessionId, sibling.directory); + const started = Promise.withResolvers(); + const capture = Promise.withResolvers(); + const directories: string[] = []; + const handlerDeps = deps({ + kiloClient: undefined, + collectWorktreeChanges: async directory => { + directories.push(directory); + if (directory !== session.directory) return captured; + started.resolve(); + return capture.promise; + }, + }); + const request = handleControlRequest( + 'session.git.summary', + session, + { revision: 1 }, + handlerDeps + ); + await started.promise; + let fenced = false; + const deletion = fenceDirectoryOperations(session.directory).then(() => { + fenced = true; + }); + try { + await Promise.resolve(); + expect(fenced).toBe(false); + expect( + await handleControlRequest('session.git.summary', session, { revision: 2 }, handlerDeps) + ).toEqual({ + ok: false, + error: { code: 'not_ready', message: 'Worktree is being deleted', retryable: false }, + }); + expect( + await handleControlRequest('session.git.summary', sibling, { revision: 1 }, handlerDeps) + ).toEqual({ ok: true, result: captured }); + expect(directories).toEqual([session.directory, sibling.directory]); + expect(fenced).toBe(false); + capture.resolve(captured); + expect(await request).toEqual({ + ok: false, + error: { code: 'not_ready', message: 'Worktree is being deleted', retryable: false }, + }); + await deletion; + expect(fenced).toBe(true); + expect(handlerDeps.tasks.size).toBe(0); + expect(handlerDeps.sessions).toEqual([]); + } finally { + capture.resolve(captured); + await Promise.all([request, deletion]); + } + }); + + it.each(['detached', 'moved', 'retired'] as const)( + 'rejects capture completed after its root is %s while preserving its sibling', + async change => { + const sibling = { ...session, sessionId: 'ses_2', kiloSessionId: 'kilo_2' }; + rememberAttachedRoot(session.kiloSessionId, session.directory); + rememberAttachedRoot(sibling.kiloSessionId, sibling.directory); + const abort = new AbortController(); + const started = Promise.withResolvers(); + const capture = Promise.withResolvers(); + const handlerDeps = deps({ + kiloClient: undefined, + signal: abort.signal, + collectWorktreeChanges: async (_directory, _request, _runGit, signal) => { + started.resolve(signal); + return capture.promise; + }, + }); + const request = handleControlRequest( + 'session.git.summary', + session, + { revision: 1 }, + handlerDeps + ); + try { + expect(await started.promise).toBe(abort.signal); + if (change === 'detached') { + expect(await handleControlRequest('session.detach', session, {}, handlerDeps)).toEqual({ + ok: true, + result: { detached: true }, + }); + } else if (change === 'moved') { + rememberAttachedRoot(session.kiloSessionId, '/moved'); + } else { + abort.abort(); + } + capture.resolve(captured); + expect(await request).toEqual({ + ok: false, + error: + change === 'retired' + ? { code: 'not_ready', message: 'Kilo is not ready', retryable: true } + : { + code: 'not_ready', + message: 'Session directory is not attached', + retryable: false, + }, + }); + expect(rootForSession(sibling.kiloSessionId, sibling.directory)).toBe( + sibling.kiloSessionId + ); + expect(handlerDeps.tasks.size).toBe(0); + } finally { + capture.resolve(captured); + await request; + } + } + ); + + it('collects only from the attached root without calling Kilo or attaching anything', async () => { + rememberAttachedRoot(session.kiloSessionId, session.directory); + const calls: unknown[] = []; + const capture = { + revision: 12, + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: [], + truncated: false, + }; + const payload = { revision: 12, baseRef: 'refs/remotes/origin/main' }; + const result = await handleControlRequest( + 'session.git.summary', + session, + payload, + deps({ + kiloClient: undefined, + kiloReady: false, + collectWorktreeChanges: async (directory, request) => { + calls.push({ directory, request }); + return capture; + }, + }) + ); + expect(calls).toEqual([{ directory: session.directory, request: payload }]); + expect(result).toEqual({ ok: true, result: capture }); + }); + + it('requires the request envelope identity', async () => { + const result = await handleControlRequest( + 'session.git.summary', + undefined, + { revision: 1 }, + deps() + ); + expect(result).toEqual({ + ok: false, + error: { code: 'protocol_error', message: 'session identity is required', retryable: false }, + }); + }); + + it.each([ + 'unattached', + 'directory-only', + 'wrong-directory', + 'child', + 'unknown-root', + 'detached-root', + ])('rejects %s scope without running capture', async scope => { + if (scope === 'directory-only') + rememberSessionDirectory(session.kiloSessionId, session.directory); + if (scope === 'wrong-directory') rememberAttachedRoot(session.kiloSessionId, '/other'); + if (scope === 'child') { + rememberAttachedRoot('root', session.directory); + rememberChildSession({ + childId: session.kiloSessionId, + parentId: 'root', + directory: session.directory, + }); + } + if (scope === 'unknown-root') rememberAttachedRoot('other', session.directory); + if (scope === 'detached-root') { + rememberAttachedRoot(session.kiloSessionId, session.directory); + rememberAttachedRoot('replacement', session.directory); + forgetAttachedRoot(session.kiloSessionId, session.directory); + } + let called = false; + const result = await handleControlRequest( + 'session.git.summary', + session, + { revision: 1 }, + deps({ + collectWorktreeChanges: async () => { + called = true; + throw new Error('Must not run'); + }, + }) + ); + expect(called).toBe(false); + expect(result).toEqual({ + ok: false, + error: { code: 'not_ready', message: 'Session directory is not attached', retryable: false }, + }); + }); + + it.each([ + { revision: 1, directory: '/outside' }, + { revision: 1, baseRef: '--help' }, + { revision: 0 }, + { revision: Number.MAX_SAFE_INTEGER + 1 }, + ])('rejects invalid payload %j without capture', async payload => { + rememberAttachedRoot(session.kiloSessionId, session.directory); + let called = false; + const result = await handleControlRequest( + 'session.git.summary', + session, + payload, + deps({ + collectWorktreeChanges: async () => { + called = true; + throw new Error('Must not run'); + }, + }) + ); + expect(called).toBe(false); + expect(result).toEqual({ + ok: false, + error: { code: 'protocol_error', message: 'Invalid payload', retryable: false }, + }); + }); + + it('returns a safe failure without exposing subprocess output or file data', async () => { + rememberAttachedRoot(session.kiloSessionId, session.directory); + const result = await handleControlRequest( + 'session.git.summary', + session, + { revision: 1 }, + deps({ + collectWorktreeChanges: async () => { + throw new Error('private stdout, stderr, file contents, token'); + }, + }) + ); + expect(result).toEqual({ + ok: false, + error: { code: 'capture_failed', message: 'Worktree capture failed', retryable: true }, + }); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts index 38bf22a387..4ad1e186b0 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts @@ -15,6 +15,7 @@ import { sessionDetachPayloadSchema, sessionMessageOutcomeSchema, sessionEventPayloadSchema, + sessionGitSummaryPayloadSchema, sessionPermissionResolvePayloadSchema, sessionPromptPayloadSchema, sessionQuestionResolvePayloadSchema, @@ -54,7 +55,11 @@ import { type WorktreeKiloRuntime, type WorktreeKiloRuntimes, } from './worktree-runtime.js'; -import { assertDirectoryActive, fenceDirectoryOperations } from './worktree-operations'; +import { + assertDirectoryActive, + fenceDirectoryOperations, + runDirectoryOperation, +} from './worktree-operations'; import { createWorktreeKiloCleanupClient, deleteWorktree, @@ -62,6 +67,7 @@ import { validateWorktreeDirectory, type WorktreeKiloCleanupClient, } from './delete-worktree'; +import { collectWorktreeChanges } from './worktree-changes'; export type HandlerSessionSnapshot = { kiloSessionId: string; @@ -223,6 +229,7 @@ export type HandlerDeps = { applyAttach?: typeof applySessionAttach; materializeAttachments?: typeof materializeMessageAttachments; runAutoCommit?: typeof runAutoCommit; + collectWorktreeChanges?: typeof collectWorktreeChanges; }; export type ControlHandlerResult = @@ -529,7 +536,7 @@ export async function handleControlRequest( return fail('protocol_error', 'session identity is required', false); } if ( - (!deps.kiloReady || deps.signal?.aborted) && + (deps.signal?.aborted || (!deps.kiloReady && operation !== 'session.git.summary')) && operation !== 'session.abort' && operation !== 'session.detach' ) { @@ -636,6 +643,8 @@ async function handleSessionControlRequest( sessionTerminalConnectResultSchema, (runtime, identity, parsed) => runtime.connect(identity, parsed) ); + case 'session.git.summary': + return handleGitSummary(session, payload, deps); default: return fail('unknown_operation', 'Unknown operation', false); } @@ -775,6 +784,41 @@ async function handleTerminalOperation( } } +async function handleGitSummary( + session: SessionRequestIdentity, + payload: unknown, + deps: HandlerDeps +): Promise { + const parsed = sessionGitSummaryPayloadSchema.safeParse(payload); + if (!parsed.success) return fail('protocol_error', 'Invalid payload', false); + const directory = session.directory; + return runDirectoryOperation(directory, async () => { + if (rootForSession(session.kiloSessionId, directory) !== session.kiloSessionId) { + return fail('not_ready', 'Session directory is not attached', false); + } + if (deps.signal?.aborted) return missingKilo(); + let result: Awaited>; + try { + result = await (deps.collectWorktreeChanges ?? collectWorktreeChanges)( + directory, + parsed.data, + undefined, + deps.signal + ); + } catch { + return deps.signal?.aborted + ? missingKilo() + : fail('capture_failed', 'Worktree capture failed', true); + } + assertDirectoryActive(directory); + if (deps.signal?.aborted) return missingKilo(); + if (rootForSession(session.kiloSessionId, directory) !== session.kiloSessionId) { + return fail('not_ready', 'Session directory is not attached', false); + } + return ok(result); + }); +} + function validAttachmentPaths( session: SessionRequestIdentity, payload: SessionPromptPayload diff --git a/services/cloud-agent-next/wrapper/src/control/session-directories.ts b/services/cloud-agent-next/wrapper/src/control/session-directories.ts index db4ea84b69..0f9c321d17 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-directories.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-directories.ts @@ -2,6 +2,11 @@ const directories = new Map(); const rootBySessionId = new Map(); const rootsByDirectory = new Map>(); const detachedSessionIds = new Set(); +const rootAttachments = new Map(); + +export function rootAttachmentId(rootKiloSessionId: string): symbol | undefined { + return rootAttachments.get(rootKiloSessionId); +} function removeDirectoryRoot(directory: string, rootKiloSessionId: string): void { const roots = rootsByDirectory.get(directory); @@ -16,6 +21,9 @@ export function rememberSessionDirectory(kiloSessionId: string, directory: strin export function rememberAttachedRoot(rootKiloSessionId: string, directory: string): void { const previousDirectory = directories.get(rootKiloSessionId); + if (!rootAttachments.has(rootKiloSessionId) || previousDirectory !== directory) { + rootAttachments.set(rootKiloSessionId, Symbol()); + } if (previousDirectory && previousDirectory !== directory) { removeDirectoryRoot(previousDirectory, rootKiloSessionId); } @@ -37,6 +45,7 @@ export function forgetAttachedRoot(rootKiloSessionId: string, directory?: string return; } if (attachedDirectory) removeDirectoryRoot(attachedDirectory, rootKiloSessionId); + rootAttachments.delete(rootKiloSessionId); detachedSessionIds.add(rootKiloSessionId); for (const [sessionId, root] of rootBySessionId) { if (root !== rootKiloSessionId) continue; @@ -104,6 +113,7 @@ export function resetSessionDirectoryState(): void { rootBySessionId.clear(); rootsByDirectory.clear(); detachedSessionIds.clear(); + rootAttachments.clear(); } export function directoryForSession(kiloSessionId: string | undefined): string | undefined { diff --git a/services/cloud-agent-next/wrapper/src/control/standalone-build.test.ts b/services/cloud-agent-next/wrapper/src/control/standalone-build.test.ts new file mode 100644 index 0000000000..fed861cb2f --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/standalone-build.test.ts @@ -0,0 +1,50 @@ +import { expect, it } from 'bun:test'; +import { cp, mkdir, mkdtemp, rm, symlink } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { dirname, join, resolve } from 'node:path'; + +it('builds the control entrypoint with only standalone wrapper runtime dependencies', async () => { + const wrapper = resolve(import.meta.dir, '../..'); + const fixture = await mkdtemp(join(tmpdir(), 'standalone-control-build-')); + try { + const isolatedWrapper = join(fixture, 'wrapper'); + await mkdir(isolatedWrapper); + await Promise.all([ + cp(join(wrapper, 'src'), join(isolatedWrapper, 'src'), { recursive: true }), + cp(join(wrapper, 'package.json'), join(isolatedWrapper, 'package.json')), + cp(join(wrapper, '../src/shared'), join(fixture, 'src/shared'), { recursive: true }), + ]); + const manifest = (await Bun.file(join(isolatedWrapper, 'package.json')).json()) as { + dependencies: Record; + }; + const nodeModules = join(isolatedWrapper, 'node_modules'); + for (const dependency of Object.keys(manifest.dependencies)) { + const target = join(nodeModules, dependency); + await mkdir(dirname(target), { recursive: true }); + await symlink(join(wrapper, 'node_modules', dependency), target, 'dir'); + } + await symlink(nodeModules, join(fixture, 'node_modules'), 'dir'); + expect(() => + Bun.resolveSync( + '@kilocode/worker-utils/cloud-agent-worktree-changes', + join(fixture, 'src/shared') + ) + ).toThrow(); + + const result = await Bun.build({ + entrypoints: [join(isolatedWrapper, 'src/control/main.ts')], + root: fixture, + target: 'bun', + minify: true, + }); + expect(result.logs.filter(log => log.level === 'error').map(log => log.message)).toEqual([]); + expect(result.success).toBe(true); + expect(result.outputs).toHaveLength(1); + const output = result.outputs[0]; + if (!output) throw new Error('Missing control wrapper bundle'); + expect(output.size).toBeGreaterThan(0); + expect(await output.text()).not.toContain('@kilocode/worker-utils'); + } finally { + await rm(fixture, { recursive: true, force: true }); + } +}); diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-changes.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-changes.test.ts new file mode 100644 index 0000000000..42615b2e18 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/worktree-changes.test.ts @@ -0,0 +1,775 @@ +import { rejects } from 'assert/strict'; +import { afterEach, describe, expect, it, spyOn } from 'bun:test'; +import { chmod, mkdir, mkdtemp, rename, rm, symlink, writeFile } from 'fs/promises'; +import { tmpdir } from 'os'; +import { dirname, join } from 'path'; +import { + MAX_WORKTREE_CHANGES_BYTES, + MAX_WORKTREE_CHANGES_FILES, + sessionGitSummaryResultSchema, +} from '../../../src/shared/sandbox-control-protocol.js'; +import { git, runProcess, type ExecResult } from '../utils.js'; +import { collectWorktreeChanges, parseWorktreeDiff } from './worktree-changes'; + +const directories: string[] = []; +const baseRef = 'refs/remotes/origin/main'; +const hash = 'a'.repeat(40); + +function run(directory: string, args: string[], expectedExitCode = 0): string { + const result = Bun.spawnSync({ + cmd: ['git', '-c', 'commit.gpgsign=false', '-c', 'core.hooksPath=/dev/null', ...args], + cwd: directory, + env: { + ...process.env, + GIT_CONFIG_NOSYSTEM: '1', + GIT_CONFIG_GLOBAL: '/dev/null', + GIT_AUTHOR_NAME: 'Test', + GIT_AUTHOR_EMAIL: 'test@example.com', + GIT_COMMITTER_NAME: 'Test', + GIT_COMMITTER_EMAIL: 'test@example.com', + }, + stdout: 'pipe', + stderr: 'pipe', + }); + if (result.exitCode !== expectedExitCode) + throw new Error(`Fixture git command failed: ${args[0]}`); + return result.stdout.toString('utf8'); +} + +async function write(directory: string, path: string, content: string | Buffer): Promise { + const fullPath = join(directory, path); + await mkdir(dirname(fullPath), { recursive: true }); + await writeFile(fullPath, content); +} + +async function createRepo( + files: Record = { 'seed.txt': 'seed\n' } +): Promise { + const directory = await mkdtemp(join(tmpdir(), 'worktree-changes-')); + directories.push(directory); + run(directory, ['init', '-b', 'main']); + run(directory, ['config', 'core.filemode', 'true']); + for (const [path, content] of Object.entries(files)) await write(directory, path, content); + run(directory, ['add', '--all']); + run(directory, ['commit', '--allow-empty', '-m', 'base']); + run(directory, ['update-ref', baseRef, 'HEAD']); + run(directory, ['symbolic-ref', 'refs/remotes/origin/HEAD', baseRef]); + run(directory, ['switch', '-c', 'feature']); + return directory; +} + +function raw(path: string, status = 'M'): string { + return `:100644 100644 ${hash} ${'0'.repeat(40)} ${status}\0${path}\0`; +} + +function fakeGit(output: { diff?: string; untracked?: string } = {}): typeof git { + return async (args, options) => { + expect(options?.timeoutMs).toBeGreaterThan(0); + expect(options?.timeoutMs).toBeLessThanOrEqual(20_000); + expect(options?.maxOutputBytes).toBe(512 * 1024); + expect(options?.signal).toBeInstanceOf(AbortSignal); + let stdout = ''; + if (args.includes('--show-prefix')) stdout = '\n'; + else if (args.includes('symbolic-ref')) stdout = `${baseRef}\n`; + else if (args.includes('--verify') || args.includes('merge-base')) stdout = `${hash}\n`; + else if (args.includes('diff')) stdout = output.diff ?? ''; + else if (args.includes('ls-files')) stdout = output.untracked ?? ''; + else if (!args.includes('check-ref-format')) throw new Error('Unexpected git command'); + return { stdout, stderr: '', exitCode: 0 }; + }; +} + +afterEach(async () => { + await Promise.all( + directories.splice(0).map(directory => rm(directory, { recursive: true, force: true })) + ); +}); + +describe('collectWorktreeChanges', () => { + const cleanRepositories: Record[] = [{}, { 'seed.txt': 'seed\n' }]; + it.each(cleanRepositories)( + 'captures clean repositories without inventing changes', + async files => { + const directory = await createRepo(files); + const head = run(directory, ['rev-parse', 'HEAD']).slice(0, -1); + expect(await collectWorktreeChanges(directory, { revision: 7 })).toEqual({ + revision: 7, + comparison: { baseRef, mergeBase: head, head }, + files: [], + truncated: false, + }); + } + ); + + it('combines commits, staged edits, unstaged edits, and untracked files without double counting', async () => { + const directory = await createRepo({ 'changed.txt': 'base\n', 'cancelled.txt': 'base\n' }); + await write(directory, 'changed.txt', 'base\ncommitted\n'); + await write(directory, 'committed.txt', 'committed\n'); + await write(directory, 'removed-again.txt', 'temporary\n'); + run(directory, ['add', '--all']); + run(directory, ['commit', '-m', 'feature']); + await write(directory, 'changed.txt', 'base\ncommitted\nstaged\n'); + await write(directory, 'cancelled.txt', 'staged replacement\n'); + await write(directory, 'staged.txt', 'staged\n'); + run(directory, ['add', '--all']); + await write(directory, 'changed.txt', 'base\ncommitted\nstaged\nunstaged\n'); + await write(directory, 'cancelled.txt', 'base\n'); + await rm(join(directory, 'removed-again.txt')); + await write(directory, 'untracked.txt', 'first\nlast'); + + const result = await collectWorktreeChanges(directory, { revision: 1, baseRef }); + expect(result.files).toEqual([ + { + path: 'changed.txt', + status: 'modified', + additions: 3, + deletions: 0, + tracked: true, + binary: false, + countsComplete: true, + }, + { + path: 'committed.txt', + status: 'added', + additions: 1, + deletions: 0, + tracked: true, + binary: false, + countsComplete: true, + }, + { + path: 'staged.txt', + status: 'added', + additions: 1, + deletions: 0, + tracked: true, + binary: false, + countsComplete: true, + }, + { + path: 'untracked.txt', + status: 'added', + additions: 2, + deletions: 0, + tracked: false, + binary: false, + countsComplete: true, + }, + ]); + expect(result.truncated).toBe(false); + }); + + it('uses the merge base rather than a divergent upstream tip', async () => { + const directory = await createRepo(); + const ancestor = run(directory, ['rev-parse', 'HEAD']).slice(0, -1); + run(directory, ['switch', 'main']); + await write(directory, 'upstream.txt', 'upstream\n'); + run(directory, ['add', '--all']); + run(directory, ['commit', '-m', 'upstream']); + run(directory, ['update-ref', baseRef, 'HEAD']); + run(directory, ['switch', 'feature']); + await write(directory, 'feature.txt', 'feature\n'); + run(directory, ['add', '--all']); + run(directory, ['commit', '-m', 'feature']); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.comparison.mergeBase).toBe(ancestor); + expect(result.files.map(file => file.path)).toEqual(['feature.txt']); + }); + + it('respects Git ignores but excludes only the exact untracked bootstrap marker', async () => { + const directory = await createRepo({ '.gitignore': 'ignored.txt\nignored-dir/\n' }); + await write(directory, 'ignored.txt', 'ignored\n'); + await write(directory, 'ignored-dir/file.txt', 'ignored\n'); + await write(directory, '.kilo-bootstrap-complete', 'wrapper\n'); + await write(directory, 'nested/.kilo-bootstrap-complete', 'project\n'); + await write(directory, 'dist/app.log', 'generated\n'); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files.map(file => file.path)).toEqual([ + 'dist/app.log', + 'nested/.kilo-bootstrap-complete', + ]); + run(directory, ['add', '.kilo-bootstrap-complete']); + const trackedMarker = await collectWorktreeChanges(directory, { revision: 2 }); + expect( + trackedMarker.files.find(file => file.path === '.kilo-bootstrap-complete') + ).toMatchObject({ tracked: true, status: 'added' }); + }); + + it('reports deletions and renames as deletion plus addition', async () => { + const directory = await createRepo({ 'old.txt': 'move\n', 'deleted.txt': 'delete\n' }); + await rename(join(directory, 'old.txt'), join(directory, 'new.txt')); + run(directory, ['add', '--all']); + await rm(join(directory, 'deleted.txt')); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect( + result.files.map(({ path, status, additions, deletions }) => ({ + path, + status, + additions, + deletions, + })) + ).toEqual([ + { path: 'deleted.txt', status: 'deleted', additions: 0, deletions: 1 }, + { path: 'new.txt', status: 'added', additions: 1, deletions: 0 }, + { path: 'old.txt', status: 'deleted', additions: 0, deletions: 1 }, + ]); + }); + + it('keeps unresolved conflict contents as a meaningful modification', async () => { + const directory = await createRepo(); + run(directory, ['switch', 'main']); + await write(directory, 'seed.txt', 'main\n'); + run(directory, ['commit', '-am', 'main change']); + run(directory, ['switch', 'feature']); + await write(directory, 'seed.txt', 'feature\n'); + run(directory, ['commit', '-am', 'feature change']); + run(directory, ['merge', '--no-edit', 'main'], 1); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files).toEqual([ + { + path: 'seed.txt', + status: 'modified', + additions: 5, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, + }, + ]); + }); + + it('preserves mode-only and file-type changes', async () => { + const directory = await createRepo({ 'mode.txt': 'same\n', 'type.txt': 'before\n' }); + await chmod(join(directory, 'mode.txt'), 0o755); + await rm(join(directory, 'type.txt')); + await symlink('target', join(directory, 'type.txt')); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files).toEqual([ + { + path: 'mode.txt', + status: 'modified', + additions: 0, + deletions: 0, + tracked: true, + binary: false, + countsComplete: true, + }, + { + path: 'type.txt', + status: 'modified', + additions: 1, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, + }, + ]); + }); + + it('preserves Unicode, tabs, newlines, spaces, quotes, and backslashes in paths', async () => { + const names = [ + ' leading space ', + '\ttab\tname', + 'new\nline', + 'quote"back\\slash', + 'café-漢-𐐀', + 'parent-é/child', + '-dash', + ':(glob)*', + ]; + const directory = await createRepo( + Object.fromEntries(names.map(name => [`tracked/${name}`, 'before\n'])) + ); + for (const name of names) { + await write(directory, `tracked/${name}`, 'after\nsecond\n'); + await write(directory, `untracked/${name}`, 'new\n'); + } + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files).toHaveLength(names.length * 2); + for (const name of names) { + expect(result.files.find(file => file.path === `tracked/${name}`)).toMatchObject({ + additions: 2, + deletions: 1, + tracked: true, + }); + expect(result.files.find(file => file.path === `untracked/${name}`)).toMatchObject({ + additions: 1, + tracked: false, + }); + } + }); + + it('flags tracked and sampled untracked binary files without returning contents', async () => { + const directory = await createRepo({ 'tracked.bin': Buffer.from([0, 1, 2]) }); + await write(directory, 'tracked.bin', Buffer.from([0, 3, 4])); + await write(directory, 'nul.txt', Buffer.from([65, 0, 65])); + await write(directory, 'controls.txt', Buffer.from([1, 2, 65, 65, 65])); + await write( + directory, + 'thirty-percent.txt', + Buffer.from([1, 2, 3, 65, 65, 65, 65, 65, 65, 65]) + ); + await write(directory, 'notes.bin', 'text\n'); + await write(directory, 'after-sample.txt', `${'a'.repeat(8192)}\0`); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + for (const path of ['tracked.bin', 'nul.txt', 'controls.txt']) { + expect(result.files.find(file => file.path === path)).toMatchObject({ + binary: true, + additions: 0, + deletions: 0, + countsComplete: false, + }); + } + for (const path of ['notes.bin', 'thirty-percent.txt', 'after-sample.txt']) { + expect(result.files.find(file => file.path === path)).toMatchObject({ + binary: false, + additions: 1, + countsComplete: true, + }); + } + expect(sessionGitSummaryResultSchema.safeParse(result).success).toBe(true); + }); + + it('counts symlink target text without following existing or dangling targets', async () => { + const directory = await createRepo(); + const outside = await mkdtemp(join(tmpdir(), 'worktree-symlink-target-')); + directories.push(outside); + await write(outside, 'secret.bin', Buffer.alloc(20_000, 0)); + await symlink(join(outside, 'secret.bin'), join(directory, 'link')); + await symlink('missing\nsecond-line', join(directory, 'dangling')); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files).toEqual([ + { + path: 'dangling', + status: 'added', + additions: 2, + deletions: 0, + tracked: false, + binary: false, + countsComplete: true, + }, + { + path: 'link', + status: 'added', + additions: 1, + deletions: 0, + tracked: false, + binary: false, + countsComplete: true, + }, + ]); + }); + + it('caps untracked text reads and still samples oversized binaries', async () => { + const directory = await createRepo(); + await write(directory, 'at-limit.txt', 'x\n'.repeat(500_000)); + await write(directory, 'oversized.txt', 'x'.repeat(1_000_001)); + await write(directory, 'oversized.bin', Buffer.alloc(1_000_001, 0)); + await write(directory, 'empty.txt', ''); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files.find(file => file.path === 'at-limit.txt')).toMatchObject({ + additions: 500_000, + countsComplete: true, + }); + expect(result.files.find(file => file.path === 'oversized.txt')).toMatchObject({ + additions: 0, + binary: false, + countsComplete: false, + }); + expect(result.files.find(file => file.path === 'oversized.bin')).toMatchObject({ + additions: 0, + binary: true, + countsComplete: false, + }); + expect(result.files.find(file => file.path === 'empty.txt')).toMatchObject({ + additions: 0, + binary: false, + countsComplete: true, + }); + expect(result.truncated).toBe(false); + }); + + it('bounds aggregate untracked reads while retaining sampled entries with incomplete counts', async () => { + const directory = await createRepo(); + const content = 'x\n'.repeat(500_000); + for (let index = 0; index < 20; index += 1) + await write(directory, `large-${index}.txt`, content); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files).toHaveLength(20); + const complete = result.files.filter(file => file.countsComplete).length; + expect(complete).toBeGreaterThan(0); + expect(complete).toBeLessThan(20); + expect(complete * 1_000_000 + (20 - complete) * 8192).toBeLessThanOrEqual(16 * 1024 * 1024); + expect(result.files.every(file => !file.binary)).toBe(true); + expect(result.truncated).toBe(false); + }); + + it.each(['missing/ref', '--help', '-c core.worktree=/outside', '', 'main~1', 'main\n'])( + 'rejects invalid explicit base %j without fallback', + async explicitBase => { + const directory = await createRepo(); + await rejects(collectWorktreeChanges(directory, { revision: 1, baseRef: explicitBase })); + } + ); + + it('fails without local origin/HEAD instead of falling back to mutable HEAD', async () => { + const directory = await createRepo(); + run(directory, ['symbolic-ref', '--delete', 'refs/remotes/origin/HEAD']); + await rejects(collectWorktreeChanges(directory, { revision: 1 }), /Worktree capture failed/); + expect((await collectWorktreeChanges(directory, { revision: 1, baseRef })).files).toEqual([]); + }); + + it.each(['head', 'base', 'default-target'])( + 'rejects %s movement during capture', + async movement => { + const directory = await createRepo(); + run(directory, ['commit', '--allow-empty', '-m', 'feature']); + run(directory, ['update-ref', 'refs/remotes/origin/other', baseRef]); + const movingGit: typeof git = async (args, options) => { + const result = await git(args, options); + if (args.includes('diff')) { + if (movement === 'head') run(directory, ['commit', '--allow-empty', '-m', 'moving HEAD']); + if (movement === 'base') run(directory, ['update-ref', baseRef, 'HEAD']); + if (movement === 'default-target') + run(directory, [ + 'symbolic-ref', + 'refs/remotes/origin/HEAD', + 'refs/remotes/origin/other', + ]); + } + return result; + }; + await rejects( + collectWorktreeChanges(directory, { revision: 1 }, movingGit), + /Worktree capture failed/ + ); + } + ); + + it('rejects a subdirectory rather than reading a parent repository', async () => { + const directory = await createRepo(); + await mkdir(join(directory, 'nested')); + await rejects( + collectWorktreeChanges(join(directory, 'nested'), { revision: 1 }), + /Worktree capture failed/ + ); + }); + + it('fails when an untracked file vanishes after enumeration', async () => { + const directory = await createRepo(); + await write(directory, 'vanishing.txt', 'before\n'); + const movingGit: typeof git = async (args, options) => { + const result = await git(args, options); + if (args.includes('ls-files')) await rm(join(directory, 'vanishing.txt')); + return result; + }; + await rejects(collectWorktreeChanges(directory, { revision: 1 }, movingGit)); + }); + + it('rejects symlinked parent directories and special files without reading their contents', async () => { + const directory = await createRepo(); + const outside = await mkdtemp(join(tmpdir(), 'worktree-outside-')); + directories.push(outside); + await write(outside, 'secret.txt', 'private\n'); + await symlink(outside, join(directory, 'parent')); + await rejects( + collectWorktreeChanges( + directory, + { revision: 1 }, + fakeGit({ untracked: 'parent/secret.txt\0' }) + ), + /Worktree capture failed/ + ); + const fifo = Bun.spawnSync(['mkfifo', join(directory, 'pipe')]); + expect(fifo.exitCode).toBe(0); + await rejects( + collectWorktreeChanges(directory, { revision: 1 }, fakeGit({ untracked: 'pipe\0' })), + /Worktree capture failed/ + ); + }); + + it('returns only whole entries up to the file limit from complete output', async () => { + const directory = await createRepo(); + await Promise.all( + Array.from({ length: MAX_WORKTREE_CHANGES_FILES + 5 }, (_, index) => + write(directory, `file-${String(index).padStart(4, '0')}`, '') + ) + ); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files).toHaveLength(MAX_WORKTREE_CHANGES_FILES); + expect(result.files[0]?.path).toBe('file-0000'); + expect(result.files.at(-1)?.path).toBe('file-0999'); + expect(result.truncated).toBe(true); + }); + + it('reserves snapshot space and truncates whole entries by serialized UTF-8 bytes', async () => { + const directory = await createRepo(); + const paths = Array.from( + { length: 300 }, + (_, index) => `nested/${'\t"\\'.repeat(200)}/file-${index}` + ); + const diff = + paths.map(path => raw(path)).join('') + paths.map(path => `1\t0\t${path}\0`).join(''); + expect(Buffer.byteLength(diff)).toBeLessThan(512 * 1024); + + const result = await collectWorktreeChanges(directory, { revision: 1 }, fakeGit({ diff })); + expect(result.truncated).toBe(true); + expect(result.files.length).toBeGreaterThan(0); + expect(result.files.length).toBeLessThan(paths.length); + expect(result.files.map(file => file.path)).toEqual(paths.slice(0, result.files.length)); + expect(Buffer.byteLength(JSON.stringify(result))).toBeLessThanOrEqual( + MAX_WORKTREE_CHANGES_BYTES - 1024 + ); + expect(sessionGitSummaryResultSchema.safeParse(result).success).toBe(true); + }); + + it.each([ + { stdoutTruncated: true }, + { stderrTruncated: true }, + { exitCode: 1, stderr: 'private command failure' }, + { exitCode: 124, terminationReason: 'timeout' }, + { exitCode: 124, terminationReason: 'hard_timeout' }, + ] satisfies Partial[])( + 'rejects valid-looking raw output after unsuccessful execution %j', + async failure => { + const directory = await createRepo(); + const runner = fakeGit({ diff: `${raw('valid.txt')}1\t0\tvalid.txt\0` }); + const failedGit: typeof git = async (args, options) => { + const result = await runner(args, options); + return args.includes('diff') ? { ...result, ...failure } : result; + }; + await rejects( + collectWorktreeChanges(directory, { revision: 1 }, failedGit), + /Worktree capture failed/ + ); + } + ); + + it('shares a single deadline across otherwise successful Git commands', async () => { + const directory = await createRepo(); + let now = Date.now(); + const clock = spyOn(Date, 'now').mockImplementation(() => now); + const runner = fakeGit(); + const timeouts: number[] = []; + let signal: AbortSignal | undefined; + try { + const slowGit: typeof git = async (args, options) => { + const result = await runner(args, options); + if (options?.timeoutMs !== undefined) timeouts.push(options.timeoutMs); + signal = options?.signal; + now += 5_000; + return result; + }; + await rejects( + collectWorktreeChanges(directory, { revision: 1 }, slowGit), + /Worktree capture failed/ + ); + expect(timeouts).toEqual([20_000, 15_000, 10_000, 5_000]); + expect(signal?.aborted).toBe(true); + } finally { + clock.mockRestore(); + } + }); + + it('does not run Git after the wrapper has retired', async () => { + const directory = await createRepo(); + const abort = new AbortController(); + abort.abort(); + const runGit = spyOn({ git }, 'git'); + try { + await rejects( + collectWorktreeChanges(directory, { revision: 1 }, runGit, abort.signal), + /Worktree capture failed/ + ); + expect(runGit).not.toHaveBeenCalled(); + } finally { + runGit.mockRestore(); + } + }); + + it('aborts in-flight Git collection when the wrapper retires', async () => { + const directory = await createRepo(); + const abort = new AbortController(); + const started = Promise.withResolvers(); + const runGit: typeof git = async (_args, options) => { + const signal = options?.signal; + if (!signal) throw new Error('Missing capture cancellation signal'); + started.resolve(signal); + await new Promise(resolve => { + signal.addEventListener('abort', () => resolve(), { once: true }); + }); + return { stdout: '', stderr: '', exitCode: 124, terminationReason: 'abort' }; + }; + const pending = collectWorktreeChanges(directory, { revision: 1 }, runGit, abort.signal); + const failed = rejects(pending, /Worktree capture failed/); + try { + const signal = await started.promise; + expect(signal.aborted).toBe(false); + abort.abort(); + await failed; + expect(signal.aborted).toBe(true); + } finally { + abort.abort(); + await failed; + } + }); + + it('does not inherit wrapper credentials or Git configuration overrides into capture processes', async () => { + const directory = await createRepo(); + const overrides = { + KILOCODE_TOKEN: 'fake-managed-token', + SANDBOX_CONTROL_CREDENTIAL: 'fake-control-credential', + GIT_CONFIG_COUNT: '1', + GIT_CONFIG_KEY_0: 'core.worktree', + GIT_CONFIG_VALUE_0: join(directory, 'missing-worktree'), + }; + const previous = new Map(Object.keys(overrides).map(key => [key, process.env[key]])); + let inspected = false; + const runGit: typeof git = async (args, options) => { + if (!inspected) { + inspected = true; + const result = await runProcess( + process.execPath, + [ + '-e', + 'process.stdout.write(JSON.stringify({ token: !!process.env.KILOCODE_TOKEN, control: !!process.env.SANDBOX_CONTROL_CREDENTIAL, override: !!process.env.GIT_CONFIG_COUNT }))', + ], + options + ); + expect(result.exitCode).toBe(0); + expect(JSON.parse(result.stdout)).toEqual({ + token: false, + control: false, + override: false, + }); + } + return git(args, options); + }; + try { + Object.assign(process.env, overrides); + const result = await collectWorktreeChanges(directory, { revision: 1 }, runGit); + expect(inspected).toBe(true); + expect(result.files).toEqual([]); + } finally { + for (const [key, value] of previous) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + } + }); + + it.each(['unterminated', '\0', '../outside\0', '/absolute\0', 'same\0same\0'])( + 'rejects malformed untracked records %j', + async untracked => { + const directory = await createRepo(); + await rejects( + collectWorktreeChanges(directory, { revision: 1 }, fakeGit({ untracked })), + /Worktree capture failed/ + ); + } + ); + + it('keeps the tracked deletion once when the same path remains untracked', async () => { + const directory = await createRepo(); + run(directory, ['rm', '--cached', 'seed.txt']); + await write(directory, 'seed.txt', 'untracked\nreplacement\n'); + expect(run(directory, ['ls-files', '--others', '--exclude-standard', '-z'])).toBe('seed.txt\0'); + + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files).toEqual([ + { + path: 'seed.txt', + status: 'deleted', + additions: 0, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, + }, + ]); + expect(result.truncated).toBe(false); + }); + + it('disables external diff and textconv hooks', async () => { + const directory = await createRepo({ + '.gitattributes': '*.txt diff=custom\n', + 'seed.txt': 'before\n', + }); + run(directory, ['config', 'diff.external', 'false']); + run(directory, ['config', 'diff.custom.command', 'false']); + run(directory, ['config', 'diff.custom.textconv', 'false']); + await write(directory, 'seed.txt', 'after\n'); + const result = await collectWorktreeChanges(directory, { revision: 1 }); + expect(result.files).toEqual([ + { + path: 'seed.txt', + status: 'modified', + additions: 1, + deletions: 1, + tracked: true, + binary: false, + countsComplete: true, + }, + ]); + }); +}); + +describe('parseWorktreeDiff', () => { + it('joins numstat to exact paths rather than output order', () => { + const first = ' first\t"\\\n漢 '; + const second = 'second'; + const output = `${raw(first)}${raw(second, 'A')}2\t0\t${second}\0` + `1\t3\t${first}\0`; + expect(parseWorktreeDiff(output)).toEqual([ + { + path: first, + status: 'modified', + additions: 1, + deletions: 3, + tracked: true, + binary: false, + countsComplete: true, + }, + { + path: second, + status: 'added', + additions: 2, + deletions: 0, + tracked: true, + binary: false, + countsComplete: true, + }, + ]); + }); + + it.each([ + 'unterminated', + '\0', + raw('missing-counts', 'A'), + raw('missing-type-counts', 'T'), + '1\t0\tmissing-raw\0', + `${raw('one')}1\t0\tother\0`, + `${raw('one')}1\t0\tone\0` + '1\t0\tone\0', + `${raw('one')}${raw('one')}1\t0\tone\0`, + `${raw('one', 'U')}1\t0\tone\0`, + `${raw('one', 'R')}1\t0\tone\0`, + `${raw('one')}-\t0\tone\0`, + `${raw('one')}-1\t0\tone\0`, + `${raw('one')}9007199254740992\t0\tone\0`, + `${raw('one')}1x\t0\tone\0`, + `${raw('../outside')}1\t0\t../outside\0`, + `${raw('one')}1\t0\t\0`, + ])('fails on malformed or unsupported raw/numstat data %j', output => { + expect(() => parseWorktreeDiff(output)).toThrow('Worktree capture failed'); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-changes.ts b/services/cloud-agent-next/wrapper/src/control/worktree-changes.ts new file mode 100644 index 0000000000..5db4589928 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/worktree-changes.ts @@ -0,0 +1,380 @@ +import { constants, type Stats } from 'fs'; +import { lstat, open, readlink, realpath } from 'fs/promises'; +import { dirname, join } from 'path'; +import { + MAX_WORKTREE_CHANGES_BYTES, + MAX_WORKTREE_CHANGES_FILES, + sessionGitSummaryPayloadSchema, + sessionGitSummaryResultSchema, + worktreeChangesFileSchema, + type SessionGitSummaryPayload, + type SessionGitSummaryResult, + type WorktreeChangesFile, +} from '../../../src/shared/sandbox-control-protocol.js'; +import { git, withTimeoutAndAbort } from '../utils.js'; + +const CAPTURE_TIMEOUT_MS = 20_000; +const MAX_RAW_OUTPUT_BYTES = 512 * 1024; +const MAX_UNTRACKED_FILE_BYTES = 1_000_000; +const MAX_UNTRACKED_READ_BYTES = 16 * 1024 * 1024; +const BINARY_SAMPLE_BYTES = 8192; +const SUMMARY_BYTES = MAX_WORKTREE_CHANGES_BYTES - 1024; +const CAPTURE_FAILED = 'Worktree capture failed'; +const DEFAULT_BASE_REF = 'refs/remotes/origin/HEAD'; +const RAW_HEADER = + /^:([0-7]{6}) ([0-7]{6}) (?:[0-9a-f]{40}|[0-9a-f]{64}) (?:[0-9a-f]{40}|[0-9a-f]{64}) ([AMDTU])$/; + +function nulRecords(output: string): string[] { + if (output === '') return []; + if (!output.endsWith('\0')) throw new Error(CAPTURE_FAILED); + return output.slice(0, -1).split('\0'); +} + +function parsePath(value: string | undefined): string { + const parsed = worktreeChangesFileSchema.shape.path.safeParse(value); + if (!parsed.success) throw new Error(CAPTURE_FAILED); + return parsed.data; +} + +function parseCount(value: string): number { + if (!/^(?:0|[1-9][0-9]*)$/.test(value)) throw new Error(CAPTURE_FAILED); + const count = Number(value); + if (!Number.isSafeInteger(count)) throw new Error(CAPTURE_FAILED); + return count; +} + +export function parseWorktreeDiff(output: string): WorktreeChangesFile[] { + const records = nulRecords(output); + const statuses = new Map(); + const possiblyUnchanged = new Set(); + let index = 0; + while (records[index]?.startsWith(':')) { + const header = RAW_HEADER.exec(records[index]); + if (!header || header[3] === 'U') throw new Error(CAPTURE_FAILED); + const path = parsePath(records[index + 1]); + if (statuses.has(path)) throw new Error(CAPTURE_FAILED); + statuses.set(path, header[3] === 'A' ? 'added' : header[3] === 'D' ? 'deleted' : 'modified'); + if (header[3] === 'M' && header[1] === header[2]) possiblyUnchanged.add(path); + index += 2; + } + + const files = new Map(); + for (; index < records.length; index += 1) { + const record = records[index]; + const firstTab = record.indexOf('\t'); + const secondTab = record.indexOf('\t', firstTab + 1); + if (firstTab <= 0 || secondTab <= firstTab + 1) throw new Error(CAPTURE_FAILED); + const path = parsePath(record.slice(secondTab + 1)); + const status = statuses.get(path); + if (!status || files.has(path)) throw new Error(CAPTURE_FAILED); + const added = record.slice(0, firstTab); + const deleted = record.slice(firstTab + 1, secondTab); + const binary = added === '-' && deleted === '-'; + files.set(path, { + path, + status, + additions: binary ? 0 : parseCount(added), + deletions: binary ? 0 : parseCount(deleted), + tracked: true, + binary, + countsComplete: !binary, + }); + } + const result: WorktreeChangesFile[] = []; + for (const path of statuses.keys()) { + const file = files.get(path); + if (file) result.push(file); + else if (!possiblyUnchanged.has(path)) throw new Error(CAPTURE_FAILED); + } + return result; +} + +function isBinary(bytes: Uint8Array): boolean { + let controls = 0; + for (const byte of bytes) { + if (byte === 0) return true; + if (byte < 9 || (byte > 13 && byte < 32)) controls += 1; + } + return bytes.length > 0 && controls / bytes.length > 0.3; +} + +function lineCount(bytes: Uint8Array): number { + let lines = 0; + for (const byte of bytes) { + if (byte === 10) lines += 1; + } + return lines + (bytes.length > 0 && bytes[bytes.length - 1] !== 10 ? 1 : 0); +} + +function sameFile(before: Stats, after: Stats): boolean { + return ( + before.dev === after.dev && + before.ino === after.ino && + before.mode === after.mode && + before.size === after.size && + before.mtimeMs === after.mtimeMs && + before.ctimeMs === after.ctimeMs + ); +} + +async function readUntracked( + directory: string, + path: string, + budget: { remaining: number }, + reservedSampleBytes: number, + checkDeadline: () => number +): Promise { + checkDeadline(); + const fullPath = join(directory, path); + if ((await realpath(dirname(fullPath))) !== dirname(fullPath)) throw new Error(CAPTURE_FAILED); + const before = await lstat(fullPath); + const file: WorktreeChangesFile = { + path, + status: 'added', + additions: 0, + deletions: 0, + tracked: false, + binary: false, + countsComplete: false, + }; + + if (before.isSymbolicLink()) { + const target = await readlink(fullPath, { encoding: 'buffer' }); + budget.remaining -= target.length; + if (budget.remaining < 0 || !sameFile(before, await lstat(fullPath))) { + throw new Error(CAPTURE_FAILED); + } + checkDeadline(); + return { ...file, additions: lineCount(target), countsComplete: true }; + } + if (!before.isFile()) throw new Error(CAPTURE_FAILED); + + const handle = await open( + fullPath, + constants.O_RDONLY | constants.O_NOFOLLOW | constants.O_NONBLOCK + ); + try { + if (!sameFile(before, await handle.stat())) throw new Error(CAPTURE_FAILED); + const sampleBytes = Math.min(before.size, BINARY_SAMPLE_BYTES); + const complete = + before.size <= MAX_UNTRACKED_FILE_BYTES && + before.size <= budget.remaining - reservedSampleBytes; + const bytes = Buffer.alloc(complete ? before.size : sampleBytes); + + async function readUntil(start: number, end: number): Promise { + let offset = start; + while (offset < end) { + checkDeadline(); + if (end - offset > budget.remaining) throw new Error(CAPTURE_FAILED); + const { bytesRead } = await handle.read(bytes, offset, end - offset, offset); + if (bytesRead === 0) throw new Error(CAPTURE_FAILED); + budget.remaining -= bytesRead; + offset += bytesRead; + } + } + + await readUntil(0, sampleBytes); + file.binary = isBinary(bytes.subarray(0, sampleBytes)); + if (!file.binary && complete) { + await readUntil(sampleBytes, bytes.length); + file.additions = lineCount(bytes); + file.countsComplete = true; + } + if (!sameFile(before, await handle.stat()) || !sameFile(before, await lstat(fullPath))) { + throw new Error(CAPTURE_FAILED); + } + checkDeadline(); + return file; + } finally { + await handle.close(); + } +} + +export async function collectWorktreeChanges( + directory: string, + request: SessionGitSummaryPayload, + runGit: typeof git = git, + signal?: AbortSignal +): Promise { + const controller = new AbortController(); + const captureSignal = signal ? AbortSignal.any([signal, controller.signal]) : controller.signal; + const deadline = Date.now() + CAPTURE_TIMEOUT_MS; + + function remainingTime(): number { + const remaining = deadline - Date.now(); + if (captureSignal.aborted || remaining <= 0) throw new Error(CAPTURE_FAILED); + return remaining; + } + + async function capture(): Promise { + remainingTime(); + const payload = sessionGitSummaryPayloadSchema.parse(request); + const root = await realpath(directory); + async function command(args: string[]): Promise { + const result = await runGit( + [ + '--no-pager', + '--no-optional-locks', + '-c', + 'color.ui=false', + '-c', + 'core.quotepath=false', + '-c', + 'core.fsmonitor=false', + '-c', + 'diff.autoRefreshIndex=false', + ...args, + ], + { + cwd: root, + inheritEnv: false, + env: { + PATH: process.env.PATH, + GIT_CONFIG_NOSYSTEM: '1', + GIT_CONFIG_GLOBAL: '/dev/null', + GIT_TERMINAL_PROMPT: '0', + GIT_NO_LAZY_FETCH: '1', + }, + timeoutMs: remainingTime(), + terminationGraceMs: 250, + maxOutputBytes: MAX_RAW_OUTPUT_BYTES, + signal: captureSignal, + } + ); + remainingTime(); + if ( + result.exitCode !== 0 || + result.terminationReason !== undefined || + result.stdoutTruncated || + result.stderrTruncated + ) { + throw new Error(CAPTURE_FAILED); + } + return result.stdout; + } + + async function resolveCommit(ref: string): Promise { + const output = await command([ + 'rev-parse', + '--verify', + '--end-of-options', + `${ref}^{commit}`, + ]); + const commit = /^(?:[0-9a-f]{40}|[0-9a-f]{64})\n$/.test(output) + ? output.slice(0, -1) + : undefined; + if (!commit) throw new Error(CAPTURE_FAILED); + return commit; + } + + async function resolveDefaultBase(): Promise { + const output = await command(['symbolic-ref', '--quiet', DEFAULT_BASE_REF]); + if (!output.endsWith('\n') || output.slice(0, -1).includes('\n')) { + throw new Error(CAPTURE_FAILED); + } + const base = output.slice(0, -1); + if (!base.startsWith('refs/remotes/origin/')) throw new Error(CAPTURE_FAILED); + return base; + } + + if ((await command(['rev-parse', '--show-prefix'])) !== '\n') throw new Error(CAPTURE_FAILED); + const baseRef = payload.baseRef ?? (await resolveDefaultBase()); + await command(['check-ref-format', '--allow-onelevel', baseRef]); + const head = await resolveCommit('HEAD'); + const base = await resolveCommit(baseRef); + const mergeBaseOutput = await command(['merge-base', head, base]); + if (!/^(?:[0-9a-f]{40}|[0-9a-f]{64})\n$/.test(mergeBaseOutput)) { + throw new Error(CAPTURE_FAILED); + } + const mergeBase = mergeBaseOutput.slice(0, -1); + const tracked = parseWorktreeDiff( + await command([ + 'diff', + '--raw', + '--numstat', + '--no-renames', + '--no-ext-diff', + '--no-textconv', + '--no-color', + '--ignore-submodules=none', + '--abbrev=64', + '-z', + mergeBase, + '--', + ]) + ); + const untrackedPaths = nulRecords( + await command(['ls-files', '--others', '--exclude-standard', '-z']) + ).map(parsePath); + if (new Set(untrackedPaths).size !== untrackedPaths.length) throw new Error(CAPTURE_FAILED); + const trackedPaths = new Set(tracked.map(file => file.path)); + const untracked = untrackedPaths.filter( + path => path !== '.kilo-bootstrap-complete' && !trackedPaths.has(path) + ); + + const result: SessionGitSummaryResult = { + revision: payload.revision, + comparison: { baseRef, mergeBase, head }, + files: [], + truncated: false, + }; + let bytes = Buffer.byteLength(JSON.stringify(result)); + function append(file: WorktreeChangesFile): boolean { + const addedBytes = + Buffer.byteLength(JSON.stringify(file)) + (result.files.length > 0 ? 1 : 0); + if (result.files.length >= MAX_WORKTREE_CHANGES_FILES || bytes + addedBytes > SUMMARY_BYTES) { + result.truncated = true; + return false; + } + result.files.push(file); + bytes += addedBytes; + return true; + } + + for (const file of tracked) { + if (!append(file)) break; + } + const budget = { remaining: MAX_UNTRACKED_READ_BYTES }; + if (!result.truncated) { + for (let index = 0; index < untracked.length; index += 1) { + if (result.files.length >= MAX_WORKTREE_CHANGES_FILES) { + result.truncated = true; + break; + } + const remainingFiles = Math.min( + untracked.length - index - 1, + MAX_WORKTREE_CHANGES_FILES - result.files.length - 1 + ); + const file = await readUntracked( + root, + untracked[index], + budget, + remainingFiles * BINARY_SAMPLE_BYTES, + remainingTime + ); + if (!append(file)) break; + } + } + + if ( + (await resolveCommit('HEAD')) !== head || + (await resolveCommit(baseRef)) !== base || + (payload.baseRef === undefined && (await resolveDefaultBase()) !== baseRef) + ) { + throw new Error(CAPTURE_FAILED); + } + remainingTime(); + return sessionGitSummaryResultSchema.parse(result); + } + + try { + return await withTimeoutAndAbort(capture(), { + timeoutMs: CAPTURE_TIMEOUT_MS, + timeoutMessage: CAPTURE_FAILED, + signal: captureSignal, + abortMessage: CAPTURE_FAILED, + }); + } finally { + controller.abort(); + } +} diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts new file mode 100644 index 0000000000..3228b7c3ae --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts @@ -0,0 +1,912 @@ +import { afterEach, beforeEach, describe, expect, it, jest, mock } from 'bun:test'; +import type { + BackgroundProcessInfo, + Event, + EventBackgroundProcessUpdated, + EventInteractiveTerminalUpdated, + EventMessagePartUpdated, + InteractiveTerminalInfo, + Pty, + ToolState, +} from '@kilocode/sdk/v2'; +import { WORKTREE_CHANGED_EVENT } from '../../../src/shared/worktree-changes-wire'; +import type { WrapperKiloClient } from '../kilo-api'; +import { eventKiloSessionId, sessionEventIdentity, unfilteredKiloEvents } from './feed'; +import { + buildHeartbeatPayload, + createSessionActivityRegistry, + type HandlerDeps, + type HandlerSessionSnapshot, +} from './sandbox-control-handlers'; +import { + forgetAttachedRoot, + rememberAttachedRoot, + rememberChildSession, + resetSessionDirectoryState, +} from './session-directories'; +import { fenceDirectoryOperations, resetDirectoryOperationState } from './worktree-operations'; +import { createWorktreeMutationNotifications } from './worktree-mutation-notifications'; +import type { WorktreeKiloRuntime } from './worktree-runtime'; + +const directory = '/worktree'; +const fileEdited = { type: 'file.edited', properties: { file: '/worktree/file.ts' } }; +const nextProperties = { + timestamp: 1, + sessionID: 'root', + assistantMessageID: 'assistant', + callID: 'call', +}; + +function toolEvent(state: ToolState): EventMessagePartUpdated { + return { + id: 'event', + type: 'message.part.updated', + properties: { + sessionID: 'root', + time: 1, + part: { + id: 'part', + sessionID: 'root', + messageID: 'assistant', + type: 'tool', + tool: 'mcp_arbitrary_mutation', + callID: 'call', + state, + }, + }, + }; +} + +function backgroundEvent( + status: BackgroundProcessInfo['status'], + sessionID = 'root' +): EventBackgroundProcessUpdated { + return { + id: 'event', + type: 'background_process.updated', + properties: { + scope: 'session', + info: { + id: 'process-not-a-session', + sessionID, + command: 'git add file.ts', + cwd: directory, + ports: [], + status, + lifetime: 'session', + ready: false, + output: 'process output must not be forwarded', + time: { started: 1, updated: 2 }, + }, + }, + }; +} + +function interactiveEvent( + status: InteractiveTerminalInfo['status'], + sessionID = 'root' +): EventInteractiveTerminalUpdated { + return { + id: 'event', + type: 'interactive_terminal.updated', + properties: { + info: { + id: 'terminal-not-a-session', + sessionID, + pid: 123, + command: 'git commit', + cwd: directory, + status, + cols: 80, + rows: 24, + time: { started: 1, updated: 2 }, + }, + }, + }; +} + +function resourceEvents(sessionID = 'root') { + return [ + ...(['starting', 'running', 'ready', 'exited', 'failed', 'stopping', 'stopped'] as const).map( + status => backgroundEvent(status, sessionID) + ), + { + id: 'event', + type: 'background_process.deleted', + properties: { sessionID, processID: 'process-not-a-session', scope: 'session' }, + }, + interactiveEvent('running', sessionID), + interactiveEvent('closed', sessionID), + { + id: 'event', + type: 'interactive_terminal.data', + properties: { + sessionID, + terminalID: 'terminal-not-a-session', + data: 'output must not be forwarded', + cursor: 10, + }, + }, + { + id: 'event', + type: 'interactive_terminal.deleted', + properties: { sessionID, terminalID: 'terminal-not-a-session' }, + }, + ] satisfies Event[]; +} + +function ptyEvents(sessionID?: Pty['sessionID']) { + const info: Pty = { + id: 'pty-not-a-session', + title: 'Terminal', + command: 'sh', + args: [], + cwd: directory, + status: 'running', + pid: 123, + ...(sessionID !== undefined ? { sessionID } : {}), + }; + return [ + { id: 'event', type: 'pty.created', properties: { info } }, + { id: 'event', type: 'pty.updated', properties: { info } }, + { id: 'event', type: 'pty.updated', properties: { info: { ...info, status: 'exited' } } }, + { id: 'event', type: 'pty.exited', properties: { id: info.id, exitCode: 1 } }, + { id: 'event', type: 'pty.deleted', properties: { id: info.id } }, + ] satisfies Event[]; +} + +const mutations: Event[] = [ + { id: 'event', type: 'file.edited', properties: { file: '/worktree/file.ts' } }, + ...(['add', 'change', 'unlink'] as const).map(event => ({ + id: 'event', + type: 'file.watcher.updated' as const, + properties: { file: '/worktree/file.ts', event }, + })), + { id: 'event', type: 'vcs.branch.updated', properties: { branch: 'feature' } }, + { id: 'event', type: 'vcs.branch.updated', properties: {} }, + { id: 'event', type: 'session.diff', properties: { sessionID: 'root', diff: [] } }, + { + id: 'event', + type: 'message.part.updated', + properties: { + sessionID: 'root', + time: 1, + part: { + id: 'part', + sessionID: 'root', + messageID: 'assistant', + type: 'patch', + hash: 'snapshot', + files: ['file.ts'], + }, + }, + }, + toolEvent({ status: 'running', input: {}, time: { start: 1 } }), + toolEvent({ + status: 'completed', + input: {}, + output: '', + title: '', + metadata: {}, + time: { start: 1, end: 2 }, + }), + toolEvent({ status: 'error', input: {}, error: 'Partially wrote', time: { start: 1, end: 2 } }), + { + id: 'event', + type: 'session.next.tool.called', + properties: { + ...nextProperties, + tool: 'arbitrary_shell_tool', + input: {}, + provider: { executed: false }, + }, + }, + { + id: 'event', + type: 'session.next.tool.progress', + properties: { ...nextProperties, structured: {}, content: [] }, + }, + { + id: 'event', + type: 'session.next.tool.success', + properties: { ...nextProperties, structured: {}, content: [], provider: { executed: false } }, + }, + { + id: 'event', + type: 'session.next.tool.failed', + properties: { + ...nextProperties, + error: { type: 'unknown', message: 'Partially wrote' }, + provider: { executed: false }, + }, + }, + { + id: 'event', + type: 'session.next.shell.started', + properties: { + timestamp: 1, + sessionID: 'root', + messageID: 'message', + callID: 'call', + command: 'sh', + }, + }, + { + id: 'event', + type: 'session.next.shell.ended', + properties: { timestamp: 1, sessionID: 'root', callID: 'call', output: '' }, + }, +]; + +type SendEvent = Parameters[0]['sendEvent']; +const disposers: Array<() => void> = []; + +function setup(deliver?: SendEvent, signal?: AbortSignal) { + const sessions: HandlerSessionSnapshot[] = []; + const runtimes = new Map(); + const abort = new AbortController(); + const sendEvent = mock(deliver ?? (() => true)); + const notifications = createWorktreeMutationNotifications({ + sessions, + kiloRuntimes: { get: directory => runtimes.get(directory) }, + signal: signal ?? abort.signal, + sendEvent, + }); + disposers.push(notifications.dispose); + function addRuntime(directory: string) { + const controller = new AbortController(); + let client = {} as WrapperKiloClient; + const runtime: WorktreeKiloRuntime = { + directory, + scopeId: directory, + env: {}, + signal: controller.signal, + get kiloClient() { + return client; + }, + }; + runtimes.set(directory, runtime); + return { + runtime, + controller, + replaceClient() { + client = {} as WrapperKiloClient; + }, + }; + } + function attach(kiloSessionId: string, worktree = directory) { + rememberAttachedRoot(kiloSessionId, worktree); + const snapshot = { kiloSessionId, lastActivityAt: 123, pendingInputs: new Set(['question']) }; + sessions.push(snapshot); + return snapshot; + } + const source = addRuntime(directory); + const snapshot = attach('root'); + return { + ...source, + notifications, + sessions, + runtimes, + abort, + sendEvent, + addRuntime, + attach, + snapshot, + }; +} + +function expectedHint(kiloSessionId = 'root', worktree = directory): Parameters { + return [ + 'session.event', + { type: WORKTREE_CHANGED_EVENT, properties: {} }, + { directory: worktree, kiloSessionId, rootKiloSessionId: kiloSessionId }, + ]; +} + +beforeEach(() => { + resetSessionDirectoryState(); + resetDirectoryOperationState(); + jest.useFakeTimers(); +}); + +afterEach(() => { + for (const dispose of disposers.splice(0)) dispose(); + jest.useRealTimers(); + resetSessionDirectoryState(); + resetDirectoryOperationState(); +}); + +describe('worktree mutation notifications', () => { + it.each([...mutations, ...ptyEvents(), ...ptyEvents(null)])( + 'recognizes SDK $type mutation signals', + event => { + const h = setup(); + h.notifications.observe(h.runtime, { ...event, directory }); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint()]); + } + ); + + it.each(resourceEvents())( + 'notifies SDK $type activity after its launch tool completes', + event => { + const h = setup(); + h.notifications.observe( + h.runtime, + toolEvent({ + status: 'completed', + input: {}, + output: '', + title: '', + metadata: {}, + time: { start: 1, end: 2 }, + }) + ); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint()]); + h.notifications.observe(h.runtime, { ...event, directory }); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint(), expectedHint()]); + } + ); + + it.each(['root', 'sibling', 'child'])( + 'scopes resource activity for %s without treating resource IDs as sessions', + sessionID => { + const h = setup(); + h.attach('sibling'); + h.addRuntime('/other'); + h.attach('other', '/other'); + rememberChildSession({ childId: 'child', parentId: 'root' }); + const before = structuredClone(h.sessions); + for (const event of [...resourceEvents(sessionID), ...ptyEvents(sessionID)]) { + h.notifications.observe(h.runtime, { ...event, directory }); + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint(), expectedHint('sibling')]); + expect(h.sessions).toEqual(before); + } + ); + + it.each([undefined, null, 123, '', 'foreign'])( + 'does not turn invalid resource session ID %# into a directory hint', + sessionID => { + const h = setup(); + for (const event of resourceEvents()) { + const properties = event.properties; + const malformed = + 'info' in properties + ? { ...properties, info: { ...properties.info, id: 'root', sessionID } } + : { ...properties, sessionID }; + h.notifications.observe(h.runtime, { ...event, properties: malformed, directory }); + if ('info' in properties) { + h.notifications.observe(h.runtime, { + ...event, + properties: { ...malformed, sessionID: 'root' }, + directory, + }); + } + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + } + ); + + it('rejects conflicting resource scope even when both IDs name attached siblings', () => { + const h = setup(); + h.attach('sibling'); + for (const event of [...resourceEvents(), ...ptyEvents('root')]) { + const properties = event.properties; + for (const extra of [{ sessionId: 'sibling' }, { part: { sessionID: 'sibling' } }]) { + h.notifications.observe(h.runtime, { + ...event, + properties: { ...properties, ...extra, sessionID: 'root' }, + directory, + }); + } + if ('info' in properties) { + h.notifications.observe(h.runtime, { + ...event, + properties: { ...properties, sessionID: 'sibling' }, + directory, + }); + } else { + h.notifications.observe(h.runtime, { + ...event, + properties: { ...properties, sessionID: 'root', info: { id: 'resource-not-a-session' } }, + directory, + }); + } + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + }); + + it('rejects resource sessions belonging to another emitting runtime', () => { + const h = setup(); + const other = h.addRuntime('/other'); + h.attach('other', '/other'); + rememberChildSession({ childId: 'other-child', parentId: 'other' }); + for (const sessionID of ['other', 'other-child']) { + for (const event of [...resourceEvents(sessionID), ...ptyEvents(sessionID)]) { + if (event.type === 'pty.exited' || event.type === 'pty.deleted') continue; + h.notifications.observe(h.runtime, { ...event, directory: '/other' }); + h.notifications.observe(h.runtime, { ...event, directory }); + } + } + for (const event of resourceEvents()) { + h.notifications.observe(other.runtime, { ...event, directory }); + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + }); + + it.each([undefined, '/other', '/worktree/', '/worktree/subdirectory'])( + 'requires an exact declared directory for sessionless PTYs %#', + declaredDirectory => { + const h = setup(); + h.attach('sibling'); + for (const event of [...ptyEvents(), ...ptyEvents(null)]) { + h.notifications.observe(h.runtime, { ...event, directory: declaredDirectory }); + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + } + ); + + it.each([false, 123, '', 'foreign'])( + 'rejects malformed or foreign optional PTY session IDs %#', + sessionID => { + const h = setup(); + for (const event of ptyEvents()) { + if (!('info' in event.properties)) continue; + h.notifications.observe(h.runtime, { + ...event, + properties: { info: { ...event.properties.info, sessionID } }, + directory, + }); + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + } + ); + + it('does not infer a PTY session from a resource ID or discard explicit foreign scope', () => { + const h = setup(); + for (const event of [...ptyEvents(), ...ptyEvents(null)]) { + if ('info' in event.properties) { + h.notifications.observe(h.runtime, { + ...event, + properties: { info: { ...event.properties.info, id: 'root' } }, + }); + expect(eventKiloSessionId(event.properties)).toBe('pty-not-a-session'); + } + for (const scope of [ + { sessionID: 'foreign' }, + { sessionID: null }, + { sessionId: 123 }, + { part: { sessionID: 'foreign' } }, + ]) { + h.notifications.observe(h.runtime, { + ...event, + properties: { ...event.properties, ...scope }, + directory, + }); + } + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + }); + + it('ignores malformed resource identity and lifecycle fields', () => { + const h = setup(); + for (const event of [...resourceEvents(), ...ptyEvents()]) { + const properties = event.properties; + const invalid = + 'info' in properties + ? [ + { ...properties, info: null }, + { ...properties, info: { ...properties.info, status: 'unknown' } }, + ...[undefined, null, 123, ''].map(id => ({ + ...properties, + info: { ...properties.info, id }, + })), + ] + : [undefined, null, 123, ''].map(id => ({ + ...properties, + [event.type === 'background_process.deleted' + ? 'processID' + : event.type.startsWith('pty.') + ? 'id' + : 'terminalID']: id, + })); + for (const properties of invalid) { + h.notifications.observe(h.runtime, { ...event, properties, directory }); + } + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + }); + + it('bounds sustained terminal output hints without forwarding output', () => { + const h = setup(); + const event = resourceEvents().find(event => event.type === 'interactive_terminal.data'); + if (!event) throw new Error('Missing terminal data fixture'); + for (let index = 0; index < 20; index += 1) { + h.notifications.observe(h.runtime, { ...event, directory }); + jest.advanceTimersByTime(500); + } + expect(h.sendEvent.mock.calls).toEqual([expectedHint()]); + expect(jest.getTimerCount()).toBe(0); + }); + + it('coalesces a mixed burst after quiet and starts a fresh burst without extra trailing hints', () => { + const h = setup(); + for (const event of [...mutations, ...resourceEvents(), ...ptyEvents()]) { + h.notifications.observe(h.runtime, { ...event, directory }); + jest.advanceTimersByTime(50); + } + jest.advanceTimersByTime(4_949); + expect(h.sendEvent).not.toHaveBeenCalled(); + jest.advanceTimersByTime(1); + expect(h.sendEvent.mock.calls).toEqual([expectedHint()]); + expect(jest.getTimerCount()).toBe(0); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint(), expectedHint()]); + jest.advanceTimersByTime(10_000); + expect(h.sendEvent).toHaveBeenCalledTimes(2); + }); + + it('flushes sustained changes within ten seconds and bounds timers across subsequent bursts', () => { + const h = setup(); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + for (let elapsed = 500; elapsed <= 20_000; elapsed += 500) { + jest.advanceTimersByTime(500); + expect(h.sendEvent).toHaveBeenCalledTimes(Math.floor(elapsed / 10_000)); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + expect(jest.getTimerCount()).toBe(2); + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).toHaveBeenCalledTimes(3); + expect(jest.getTimerCount()).toBe(0); + }); + + it('notifies attached sibling roots for a valid grandchild mutation, not other worktrees', () => { + const h = setup(); + h.attach('sibling'); + h.addRuntime('/other'); + h.attach('other', '/other'); + rememberChildSession({ childId: 'child', parentId: 'root' }); + rememberChildSession({ childId: 'grandchild', parentId: 'child' }); + h.notifications.observe(h.runtime, { + type: 'session.next.tool.failed', + properties: { ...nextProperties, sessionID: 'grandchild' }, + }); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint(), expectedHint('sibling')]); + }); + + it('observes ambiguous sessionless feed events before routing without modifying original events or activity', async () => { + const h = setup(); + h.attach('sibling'); + const activity = createSessionActivityRegistry(); + activity.attach('root'); + activity.attach('sibling'); + const deps = { + sessions: h.sessions, + tasks: new Map(), + activity, + kiloReady: true, + } as HandlerDeps; + const snapshots = structuredClone(h.sessions); + const heartbeat = buildHeartbeatPayload(deps); + const routed = []; + const received = []; + const envelopes = [ + ...mutations.filter(event => + ['file.edited', 'file.watcher.updated', 'vcs.branch.updated'].includes(event.type) + ), + ...ptyEvents(), + ...ptyEvents(null), + ].map(payload => ({ directory, payload })); + for await (const event of unfilteredKiloEvents(envelopes)) { + h.notifications.observe(h.runtime, event); + received.push(event); + const identity = sessionEventIdentity({ + ...event, + sessionId: eventKiloSessionId(event.properties), + runtimeDirectory: h.runtime.directory, + }); + if (identity?.rootKiloSessionId) routed.push(event); + } + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint(), expectedHint('sibling')]); + expect(received).toEqual( + envelopes.map(({ payload, directory }) => ({ + type: payload.type, + properties: payload.properties, + directory, + })) + ); + expect(routed).toEqual([]); + expect(h.sessions).toEqual(snapshots); + expect(buildHeartbeatPayload(deps)).toEqual({ + ...heartbeat, + sessions: heartbeat.sessions?.map(session => ({ + ...session, + idleForMs: session.idleForMs + 5_000, + })), + }); + }); + + it.each([ + { ...fileEdited }, + { ...fileEdited, directory: '/other' }, + { ...fileEdited, directory: '/worktree/' }, + { ...fileEdited, directory: '/worktree/subdirectory' }, + { type: 'vcs.branch.updated', properties: {}, directory: '/other' }, + ...[ + { sessionID: 'foreign' }, + { sessionID: '' }, + { sessionID: null }, + { sessionID: 123 }, + { sessionID: undefined }, + { sessionId: 'foreign' }, + { sessionId: false }, + { info: { id: 'foreign' } }, + { info: { sessionID: null } }, + { part: { sessionID: 'foreign' } }, + { part: { sessionID: 123 } }, + { sessionID: 'root', sessionId: 'foreign' }, + { sessionID: 'root', part: { sessionID: 'foreign' } }, + { sessionID: 'root', info: { id: 'foreign' } }, + ].map(scope => ({ + ...fileEdited, + directory, + properties: { ...fileEdited.properties, ...scope }, + })), + { type: 'session.diff', directory, properties: { diff: [] } }, + { type: 'session.diff', directory, properties: { sessionID: 'foreign', diff: [] } }, + { type: 'session.next.tool.called', directory, properties: { callID: 'call' } }, + { + type: 'message.part.updated', + directory, + properties: { + sessionID: 'root', + part: { type: 'patch', sessionID: 'foreign', files: [], hash: '' }, + }, + }, + ])('rejects missing, foreign, conflicting, or malformed scope %#', event => { + const h = setup(); + h.notifications.observe(h.runtime, event); + jest.advanceTimersByTime(10_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + }); + + it('rejects foreign emitting runtimes and children outside the exact worktree', () => { + const h = setup(); + const other = h.addRuntime('/other'); + h.attach('other', '/other'); + rememberChildSession({ childId: 'external-child', parentId: 'root', directory: '/external' }); + rememberChildSession({ childId: 'other-child', parentId: 'other' }); + for (const [runtime, sessionID, eventDirectory] of [ + [other.runtime, 'root', directory], + [h.runtime, 'other', '/other'], + [h.runtime, 'other-child', '/other'], + [h.runtime, 'external-child', '/external'], + ] as const) { + h.notifications.observe(runtime, { + type: 'session.diff', + directory: eventDirectory, + properties: { sessionID, diff: [] }, + }); + } + jest.advanceTimersByTime(10_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + }); + + it.each([ + 'server.heartbeat', + 'message.part.delta', + 'session.next.text.delta', + 'session.next.reasoning.delta', + 'session.next.tool.input.started', + 'session.next.tool.input.delta', + 'session.next.tool.input.ended', + 'session.status', + 'message.updated', + WORKTREE_CHANGED_EVENT, + ])('ignores nonmutating %s events', type => { + const h = setup(); + h.notifications.observe(h.runtime, { type, properties: nextProperties, directory }); + expect(jest.getTimerCount()).toBe(0); + }); + + it.each([ + toolEvent({ status: 'pending', input: {}, raw: '{}' }), + { + type: 'message.part.updated', + properties: { sessionID: 'root', part: { type: 'text', sessionID: 'root', text: 'hello' } }, + }, + { + type: 'message.part.updated', + properties: { + sessionID: 'root', + part: { type: 'reasoning', sessionID: 'root', text: 'thinking' }, + }, + }, + { + type: 'message.part.updated', + properties: { + sessionID: 'root', + part: { type: 'tool', sessionID: 'root', tool: 'shell', state: { status: ['running'] } }, + }, + }, + { type: 'file.edited', properties: { file: 1 } }, + { type: 'file.watcher.updated', properties: { file: 'file', event: 'read' } }, + { type: 'file.watcher.updated', properties: { file: 'file', event: ['change'] } }, + { type: 'vcs.branch.updated', properties: { branch: null } }, + ])('ignores nonmutating or malformed mutation payload %#', event => { + const h = setup(); + h.notifications.observe(h.runtime, { ...event, directory }); + expect(jest.getTimerCount()).toBe(0); + }); + + it.each([ + 'detach', + 'reattach', + 'move', + 'move-back', + 'snapshot-replacement', + 'snapshot-removal', + 'snapshot-id-change', + ])('drops stale targets after %s', change => { + const h = setup(); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + if (change === 'detach' || change === 'reattach') forgetAttachedRoot('root'); + if (change === 'reattach') rememberAttachedRoot('root', directory); + if (change === 'move' || change === 'move-back') rememberAttachedRoot('root', '/other'); + if (change === 'move-back') rememberAttachedRoot('root', directory); + if (change === 'snapshot-replacement') h.sessions.splice(0, 1, { ...h.snapshot }); + if (change === 'snapshot-removal') h.sessions.splice(0, 1); + if (change === 'snapshot-id-change') h.snapshot.kiloSessionId = 'replacement'; + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + }); + + it('keeps pending hints across idempotent attachment and ignores late nonmutating traffic', () => { + const h = setup(); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + jest.advanceTimersByTime(4_900); + rememberAttachedRoot('root', directory); + h.notifications.observe(h.runtime, { + type: 'session.next.tool.input.delta', + properties: { ...nextProperties, delta: 'input' }, + directory, + }); + jest.advanceTimersByTime(100); + expect(h.sendEvent.mock.calls).toEqual([expectedHint()]); + expect(jest.getTimerCount()).toBe(0); + }); + + it('debounces independent worktrees separately', () => { + const h = setup(); + const other = h.addRuntime('/other'); + h.attach('other', '/other'); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + jest.advanceTimersByTime(500); + h.notifications.observe(other.runtime, { ...fileEdited, directory: '/other' }); + jest.advanceTimersByTime(4_500); + expect(h.sendEvent.mock.calls).toEqual([expectedHint()]); + jest.advanceTimersByTime(500); + expect(h.sendEvent.mock.calls).toEqual([expectedHint(), expectedHint('other', '/other')]); + expect(jest.getTimerCount()).toBe(0); + }); + + it('does not notify a newly attached root or replacement from an older shared-worktree event', () => { + const h = setup(); + h.attach('sibling'); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + forgetAttachedRoot('root'); + h.sessions.splice(0, 1); + h.attach('root'); + h.attach('new'); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint('sibling')]); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls.slice(1)).toEqual([ + expectedHint('sibling'), + expectedHint(), + expectedHint('new'), + ]); + }); + + it.each([ + 'retirement', + 'replacement', + 'missing-runtime', + 'client-replacement', + 'abort', + 'dispose', + 'deletion', + ])('drops queued hints after runtime %s', async change => { + const h = setup(); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + if (change === 'retirement') h.controller.abort(); + if (change === 'replacement') h.addRuntime(directory); + if (change === 'missing-runtime') h.runtimes.clear(); + if (change === 'client-replacement') h.replaceClient(); + if (change === 'abort') h.abort.abort(); + if (change === 'dispose') h.notifications.dispose(); + if (change === 'deletion') await fenceDirectoryOperations(directory); + if (['retirement', 'abort', 'dispose'].includes(change)) expect(jest.getTimerCount()).toBe(0); + jest.advanceTimersByTime(10_000); + expect(h.sendEvent).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + }); + + it('rejects retired runtime callbacks after replacement and allows new runtime mutations', () => { + const h = setup(); + const replacement = h.addRuntime(directory); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + expect(jest.getTimerCount()).toBe(0); + h.notifications.observe(replacement.runtime, { ...fileEdited, directory }); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent.mock.calls).toEqual([expectedHint()]); + }); + + it('never queues while aborted, disposed, deleting, or without an attached snapshot', async () => { + const h = setup(undefined, AbortSignal.abort()); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + const live = setup(); + live.sessions.splice(0); + live.notifications.observe(live.runtime, { ...fileEdited, directory }); + live.attach('root'); + await fenceDirectoryOperations(directory); + live.notifications.observe(live.runtime, { ...fileEdited, directory }); + live.notifications.dispose(); + live.notifications.observe(live.runtime, { ...fileEdited, directory }); + expect(jest.getTimerCount()).toBe(0); + }); + + it.each(['throw', 'reject', 'false'])( + 'contains %s delivery failures without blocking sibling roots or subsequent bursts', + async failure => { + let attempts = 0; + const h = setup(() => { + attempts += 1; + if (attempts === 1) { + if (failure === 'throw') throw new Error('delivery failed'); + if (failure === 'reject') return Promise.reject(new Error('delivery failed')); + return false; + } + return true; + }); + h.attach('sibling'); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + expect(() => jest.advanceTimersByTime(5_000)).not.toThrow(); + await Promise.resolve(); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).toHaveBeenCalledTimes(4); + expect(jest.getTimerCount()).toBe(0); + } + ); + + it('stops the flush when delivery triggers socket-disconnect shutdown', () => { + const h = setup(() => h.abort.abort()); + h.attach('sibling'); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + jest.advanceTimersByTime(5_000); + expect(h.sendEvent).toHaveBeenCalledTimes(1); + h.notifications.observe(h.runtime, { ...fileEdited, directory }); + expect(jest.getTimerCount()).toBe(0); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.ts b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.ts new file mode 100644 index 0000000000..a1b79ceab0 --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.ts @@ -0,0 +1,293 @@ +import type { BackgroundProcessInfo, Event } from '@kilocode/sdk/v2'; +import { WORKTREE_CHANGED_EVENT } from '../../../src/shared/worktree-changes-wire'; +import type { HandlerSessionSnapshot } from './sandbox-control-handlers'; +import { sessionEventIdentity } from './feed'; +import { rootAttachmentId, rootForSession } from './session-directories'; +import { assertDirectoryActive } from './worktree-operations'; +import type { WorktreeKiloRuntime, WorktreeKiloRuntimes } from './worktree-runtime'; + +type KiloEvent = { + type: string; + properties: Record; + directory?: string; +}; + +type NotificationIdentity = { + directory: string; + kiloSessionId: string; + rootKiloSessionId: string; +}; + +type Target = { + snapshot: HandlerSessionSnapshot; + kiloSessionId: string; + attachmentId: symbol; +}; + +type Pending = { + runtime: WorktreeKiloRuntime; + kiloClient: WorktreeKiloRuntime['kiloClient']; + targets: Map; + quietTimer?: ReturnType; + maxTimer?: ReturnType; + onAbort: () => void; +}; + +const executionEvents = new Set([ + 'session.next.tool.called', + 'session.next.tool.progress', + 'session.next.tool.success', + 'session.next.tool.failed', + 'session.next.shell.started', + 'session.next.shell.ended', +] satisfies Event['type'][]); + +const resourceInfoEvents = new Set([ + 'background_process.updated', + 'interactive_terminal.updated', + 'pty.created', + 'pty.updated', +] satisfies Event['type'][]); + +const backgroundStatuses = new Set([ + 'starting', + 'running', + 'ready', + 'exited', + 'failed', + 'stopping', + 'stopped', +] satisfies BackgroundProcessInfo['status'][]); + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function isNonemptyString(value: unknown): value is string { + return typeof value === 'string' && value.length > 0; +} + +function isMutation({ type, properties }: KiloEvent): boolean { + if (type === 'file.edited') return typeof properties.file === 'string'; + if (type === 'file.watcher.updated') { + return ( + typeof properties.file === 'string' && + (properties.event === 'add' || properties.event === 'change' || properties.event === 'unlink') + ); + } + if (type === 'vcs.branch.updated') { + return properties.branch === undefined || typeof properties.branch === 'string'; + } + if (resourceInfoEvents.has(type)) { + const info = properties.info; + if (!isRecord(info) || !isNonemptyString(info.id)) return false; + if (type === 'pty.created' || type === 'pty.updated') { + return info.status === 'running' || info.status === 'exited'; + } + if (!isNonemptyString(info.sessionID)) return false; + if (type === 'background_process.updated') { + return ( + typeof info.status === 'string' && + backgroundStatuses.has(info.status) && + typeof properties.scope === 'string' + ); + } + return info.status === 'running' || info.status === 'closed'; + } + if (type === 'pty.exited' || type === 'pty.deleted') { + return ( + isNonemptyString(properties.id) && + (type === 'pty.deleted' || typeof properties.exitCode === 'number') + ); + } + if (typeof properties.sessionID !== 'string' || !properties.sessionID) return false; + if (type === 'background_process.deleted') { + return isNonemptyString(properties.processID) && typeof properties.scope === 'string'; + } + if (type === 'interactive_terminal.data' || type === 'interactive_terminal.deleted') { + return ( + isNonemptyString(properties.terminalID) && + (type === 'interactive_terminal.deleted' || + (typeof properties.data === 'string' && typeof properties.cursor === 'number')) + ); + } + if (executionEvents.has(type)) return typeof properties.callID === 'string'; + if (type === 'session.diff') return Array.isArray(properties.diff); + if (type !== 'message.part.updated') return false; + const part = properties.part; + if (!isRecord(part) || part.sessionID !== properties.sessionID) return false; + if (part.type === 'patch') { + return ( + typeof part.hash === 'string' && + Array.isArray(part.files) && + part.files.every(file => typeof file === 'string') + ); + } + return ( + part.type === 'tool' && + typeof part.tool === 'string' && + isRecord(part.state) && + (part.state.status === 'running' || + part.state.status === 'completed' || + part.state.status === 'error') + ); +} + +function mutationSessionId({ type, properties }: KiloEvent): string | undefined | null { + const ids: unknown[] = []; + for (const key of ['sessionID', 'sessionId']) { + if (key in properties) ids.push(properties[key]); + } + for (const key of ['info', 'part']) { + if (!(key in properties)) continue; + const nested = properties[key]; + if (!isRecord(nested)) return null; + const sessionlessPty = + key === 'info' && + (type === 'pty.created' || type === 'pty.updated') && + (nested.sessionID === null || nested.sessionID === undefined); + if ('sessionID' in nested && !sessionlessPty) ids.push(nested.sessionID); + if (key === 'info' && 'id' in nested && !resourceInfoEvents.has(type)) ids.push(nested.id); + } + const sessionId = ids[0]; + if (ids.some(id => typeof id !== 'string' || !id || id !== sessionId)) return null; + return typeof sessionId === 'string' ? sessionId : undefined; +} + +export function createWorktreeMutationNotifications(options: { + sessions: readonly HandlerSessionSnapshot[]; + kiloRuntimes: Pick; + signal: AbortSignal; + sendEvent: ( + event: 'session.event', + payload: { type: typeof WORKTREE_CHANGED_EVENT; properties: Record }, + identity: NotificationIdentity + ) => unknown; +}) { + const pending = new Map(); + let disposed = false; + + function isCurrent(runtime: WorktreeKiloRuntime): boolean { + assertDirectoryActive(runtime.directory); + return ( + !disposed && + !options.signal.aborted && + !runtime.signal.aborted && + options.kiloRuntimes.get(runtime.directory) === runtime + ); + } + + function validTarget(target: Target, directory: string): boolean { + return ( + options.sessions.includes(target.snapshot) && + target.snapshot.kiloSessionId === target.kiloSessionId && + rootAttachmentId(target.kiloSessionId) === target.attachmentId && + rootForSession(target.kiloSessionId, directory) === target.kiloSessionId + ); + } + + function remove(entry: Pending): void { + clearTimeout(entry.quietTimer); + clearTimeout(entry.maxTimer); + entry.runtime.signal.removeEventListener('abort', entry.onAbort); + if (pending.get(entry.runtime) === entry) pending.delete(entry.runtime); + } + + function flush(entry: Pending): void { + remove(entry); + for (const target of entry.targets.values()) { + try { + if (!isCurrent(entry.runtime) || entry.runtime.kiloClient !== entry.kiloClient) return; + if (!validTarget(target, entry.runtime.directory)) continue; + void Promise.resolve( + options.sendEvent( + 'session.event', + { type: WORKTREE_CHANGED_EVENT, properties: {} }, + { + directory: entry.runtime.directory, + kiloSessionId: target.kiloSessionId, + rootKiloSessionId: target.kiloSessionId, + } + ) + ).catch(() => {}); + } catch { + continue; + } + } + } + + function dispose(): void { + disposed = true; + options.signal.removeEventListener('abort', dispose); + for (const entry of pending.values()) remove(entry); + } + + options.signal.addEventListener('abort', dispose, { once: true }); + if (options.signal.aborted) dispose(); + + return { + dispose, + observe(runtime: WorktreeKiloRuntime, event: KiloEvent): void { + try { + if (!isRecord(event.properties) || !isMutation(event) || !isCurrent(runtime)) return; + const sessionId = mutationSessionId(event); + if (sessionId === null) return; + if (sessionId !== undefined) { + const identity = sessionEventIdentity({ + ...event, + sessionId, + runtimeDirectory: runtime.directory, + }); + if ( + identity?.directory !== runtime.directory || + !options.sessions.some( + snapshot => snapshot.kiloSessionId === identity.rootKiloSessionId + ) + ) + return; + } else if (event.directory !== runtime.directory) { + return; + } + let entry = pending.get(runtime); + if (entry && entry.kiloClient !== runtime.kiloClient) { + remove(entry); + entry = undefined; + } + const targets = entry?.targets ?? new Map(); + for (const [snapshot, target] of targets) { + if (!validTarget(target, runtime.directory)) targets.delete(snapshot); + } + for (const snapshot of options.sessions) { + const kiloSessionId = snapshot.kiloSessionId; + const attachmentId = rootAttachmentId(kiloSessionId); + if (!attachmentId || rootForSession(kiloSessionId, runtime.directory) !== kiloSessionId) + continue; + targets.set(snapshot, { snapshot, kiloSessionId, attachmentId }); + } + if (!targets.size) { + if (entry) remove(entry); + return; + } + if (!entry) { + const created: Pending = { + runtime, + kiloClient: runtime.kiloClient, + targets, + onAbort: () => remove(created), + }; + entry = created; + pending.set(runtime, entry); + runtime.signal.addEventListener('abort', entry.onAbort, { once: true }); + entry.maxTimer = setTimeout(() => flush(created), 10_000); + entry.maxTimer.unref(); + } + const queued = entry; + clearTimeout(queued.quietTimer); + queued.quietTimer = setTimeout(() => flush(queued), 5_000); + queued.quietTimer.unref(); + } catch { + return; + } + }, + }; +} diff --git a/services/cloud-agent-next/wrapper/src/utils.ts b/services/cloud-agent-next/wrapper/src/utils.ts index c690e7aa57..b2fbe7b965 100644 --- a/services/cloud-agent-next/wrapper/src/utils.ts +++ b/services/cloud-agent-next/wrapper/src/utils.ts @@ -58,10 +58,9 @@ export function isTimeoutTermination(result: ExecResult): boolean { function utf8Tail(value: string, maxBytes: number): string { const bytes = Buffer.from(value); if (bytes.length <= maxBytes) return value; - return bytes - .subarray(bytes.length - maxBytes) - .toString('utf8') - .replace(/^\uFFFD/, ''); + let start = bytes.length - maxBytes; + while (start < bytes.length && (bytes[start] & 0xc0) === 0x80) start += 1; + return bytes.subarray(start).toString('utf8'); } function appendBoundedTail( @@ -233,8 +232,7 @@ export function runProcess( inactivityTimer = setTimeout(() => terminate('inactivity_timeout'), opts.inactivityTimeoutMs); }; - const captureOutput = (stream: ProcessOutputStream, output: Buffer): void => { - const text = output.toString(); + const captureOutput = (stream: ProcessOutputStream, text: string): void => { if (stream === 'stdout') { const bounded = appendBoundedTail(stdout, text, maxOutputBytes); stdout = bounded.value; @@ -260,8 +258,10 @@ export function runProcess( hardTimeoutTimer = setTimeout(() => terminate('hard_timeout'), opts.hardTimeoutMs); } - proc.stdout.on('data', (output: Buffer) => captureOutput('stdout', output)); - proc.stderr.on('data', (output: Buffer) => captureOutput('stderr', output)); + proc.stdout.setEncoding('utf8'); + proc.stderr.setEncoding('utf8'); + proc.stdout.on('data', (output: string) => captureOutput('stdout', output)); + proc.stderr.on('data', (output: string) => captureOutput('stderr', output)); if (opts?.signal) { if (opts.signal.aborted) { From 87fe098b03f4b141dd7feff81d44eacc0c7cfcc3 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Wed, 2 Sep 2026 22:24:06 +0200 Subject: [PATCH 2/2] feat(cloud-agent-next): push saved worktree change updates --- .../cloud-agent-next/ChatHeader.tsx | 1 - .../cloud-agent-next/WorktreeChanges.tsx | 83 ++-- .../cloud-agent-next/worktree-changes.test.ts | 171 +++++++- .../cloud-agent-next/worktree-changes.ts | 50 +++ .../e2e/cloud-agent-sandbox-status.spec.ts | 383 ++++++++++++++---- .../src/__fixtures__/helpers.ts | 10 +- .../src/cloud-agent-transport.test.ts | 54 +++ .../src/cloud-agent-transport.ts | 3 + packages/cloud-agent-sdk/src/index.ts | 1 + .../cloud-agent-sdk/src/normalizer.test.ts | 64 ++- packages/cloud-agent-sdk/src/normalizer.ts | 19 +- packages/cloud-agent-sdk/src/schemas.test.ts | 50 +++ packages/cloud-agent-sdk/src/schemas.ts | 4 + .../src/session-manager.test.ts | 198 +++++++++ .../cloud-agent-sdk/src/session-manager.ts | 30 ++ .../cloud-agent-sdk/src/session-phase.test.ts | 10 +- .../src/session-transport.test.ts | 152 +++++++ .../src/sandbox-session/SandboxSession.ts | 43 +- .../cloud-agent-next/src/shared/protocol.ts | 1 + .../src/shared/worktree-changes-wire.ts | 1 + .../src/websocket/stream.test.ts | 55 +++ .../cloud-agent-next/src/websocket/types.ts | 1 + .../test/integration/sandbox-control.test.ts | 329 ++++++++++++--- 23 files changed, 1547 insertions(+), 166 deletions(-) diff --git a/apps/web/src/components/cloud-agent-next/ChatHeader.tsx b/apps/web/src/components/cloud-agent-next/ChatHeader.tsx index fd7e92f426..67b8389bd2 100644 --- a/apps/web/src/components/cloud-agent-next/ChatHeader.tsx +++ b/apps/web/src/components/cloud-agent-next/ChatHeader.tsx @@ -147,7 +147,6 @@ export function ChatHeader({ organizationId={organizationId} open={changesOpen} onToggle={onToggleChanges} - sessionActive={sessionActive} /> )} {onToggleSound && ( diff --git a/apps/web/src/components/cloud-agent-next/WorktreeChanges.tsx b/apps/web/src/components/cloud-agent-next/WorktreeChanges.tsx index 9c96c7334c..2068b3b6ce 100644 --- a/apps/web/src/components/cloud-agent-next/WorktreeChanges.tsx +++ b/apps/web/src/components/cloud-agent-next/WorktreeChanges.tsx @@ -1,7 +1,8 @@ 'use client'; -import { useEffect, useMemo, useRef, useState, type MouseEvent, type RefObject } from 'react'; +import { useEffect, useMemo, useRef, type MouseEvent, type RefObject } from 'react'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; +import { useAtomValue } from 'jotai'; import { ChevronRight, FileDiff, GitBranch, List, ListTree, RefreshCw } from 'lucide-react'; import { formatDistanceToNow } from 'date-fns'; import { Button } from '@/components/ui/button'; @@ -10,8 +11,10 @@ import { Sheet, SheetContent, SheetDescription, SheetTitle } from '@/components/ import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'; import { useLocalStorage } from '@/hooks/useLocalStorage'; import { useRawTRPCClient, useTRPC } from '@/lib/trpc/utils'; +import { useManager } from './CloudAgentProvider'; import { buildWorktreeChangesTree, + createWorktreeChangesRefresher, deserializeWorktreeChangesViewMode, formatWorktreeChangesBaseBranch, getWorktreeChangesTotals, @@ -37,22 +40,25 @@ function useSavedWorktreeChanges({ cloudAgentSessionId, organizationId, enabled, - poll = false, - catchUpUntil = 0, }: { cloudAgentSessionId: string; organizationId?: string; enabled: boolean; - poll?: boolean; - catchUpUntil?: number; }) { const trpc = useTRPC(); - const queryOptions = organizationId - ? trpc.organizations.cloudAgentNext.getWorktreeChanges.queryOptions({ - organizationId, - cloudAgentSessionId, - }) - : trpc.cloudAgentNext.getWorktreeChanges.queryOptions({ cloudAgentSessionId }); + const queryOptions = useMemo( + () => + organizationId + ? trpc.organizations.cloudAgentNext.getWorktreeChanges.queryOptions( + { organizationId, cloudAgentSessionId }, + { trpc: { abortOnUnmount: true, context: { skipBatch: true } } } + ) + : trpc.cloudAgentNext.getWorktreeChanges.queryOptions( + { cloudAgentSessionId }, + { trpc: { abortOnUnmount: true, context: { skipBatch: true } } } + ), + [trpc, organizationId, cloudAgentSessionId] + ); const saved = useQuery({ ...queryOptions, enabled, @@ -60,8 +66,13 @@ function useSavedWorktreeChanges({ refetchOnMount: 'always', refetchOnWindowFocus: true, refetchOnReconnect: true, - refetchInterval: () => (poll || Date.now() < catchUpUntil ? 5_000 : false), - retry: false, + retry: (failureCount, error) => { + const status = error.data?.httpStatus; + return ( + failureCount < 2 && + (status === undefined || status === 408 || status === 429 || status >= 500) + ); + }, structuralSharing: preserveNewerWorktreeChanges, }); return { saved, queryKey: queryOptions.queryKey }; @@ -106,30 +117,50 @@ export function WorktreeChangesButton({ organizationId, open, onToggle, - sessionActive, }: { cloudAgentSessionId: string; organizationId?: string; open: boolean; onToggle: (event: MouseEvent) => void; - sessionActive: boolean; }) { - const [catchUpUntil, setCatchUpUntil] = useState(0); - const { saved } = useSavedWorktreeChanges({ + const manager = useManager(); + const refreshSignal = useAtomValue(manager.atoms.worktreeChangesRefresh); + const queryClient = useQueryClient(); + const { saved, queryKey } = useSavedWorktreeChanges({ cloudAgentSessionId, organizationId, enabled: true, - poll: sessionActive || open, - catchUpUntil, }); - const wasSessionActive = useRef(sessionActive); + const handledSignal = useRef(refreshSignal); + const refresher = useRef | null>(null); useEffect(() => { - if (wasSessionActive.current && !sessionActive) { - setCatchUpUntil(Date.now() + 30_000); - void saved.refetch(); - } - wasSessionActive.current = sessionActive; - }, [sessionActive, saved.refetch]); + const current = createWorktreeChangesRefresher(queryClient, queryKey); + refresher.current = current; + return () => { + current.dispose(); + refresher.current = null; + }; + }, [queryClient, queryKey]); + useEffect(() => { + const previous = handledSignal.current; + if (previous === refreshSignal) return; + handledSignal.current = refreshSignal; + if (refreshSignal?.cloudSessionId !== cloudAgentSessionId) return; + const reconnected = + previous?.cloudSessionId !== refreshSignal.cloudSessionId || + previous?.connectionVersion !== refreshSignal.connectionVersion; + const cached = queryClient.getQueryData(queryKey); + if ( + !reconnected && + refreshSignal.revision !== undefined && + (cached?.snapshot?.revision ?? 0) >= refreshSignal.revision + ) + return; + void refresher.current?.request({ + revision: refreshSignal.revision ?? 0, + force: reconnected, + }); + }, [refreshSignal, cloudAgentSessionId, queryClient, queryKey]); const totals = getWorktreeChangesTotals(saved.data?.snapshot); const summary = totals diff --git a/apps/web/src/components/cloud-agent-next/worktree-changes.test.ts b/apps/web/src/components/cloud-agent-next/worktree-changes.test.ts index 20c8d29410..26eb853889 100644 --- a/apps/web/src/components/cloud-agent-next/worktree-changes.test.ts +++ b/apps/web/src/components/cloud-agent-next/worktree-changes.test.ts @@ -1,4 +1,4 @@ -import { QueryClient } from '@tanstack/react-query'; +import { QueryClient, QueryObserver } from '@tanstack/react-query'; import type { GetWorktreeChangesOutput, WorktreeChangesSnapshot, @@ -6,6 +6,7 @@ import type { import { buildWorktreeChangesTree, canOpenWorktreeChanges, + createWorktreeChangesRefresher, deserializeWorktreeChangesViewMode, formatWorktreeChangesBaseBranch, getWorktreeChangesTotals, @@ -266,6 +267,174 @@ describe('saved worktree changes cache', () => { }); }); +describe('worktree changes refresh coordination', () => { + const queryKey = ['worktree-changes', 'personal', 'workspace-a']; + const saved = (revision: number): GetWorktreeChangesOutput => ({ + snapshot: { ...snapshot, revision }, + }); + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { structuralSharing: preserveNewerWorktreeChanges, retry: false }, + }, + }); + }); + + afterEach(() => queryClient.clear()); + + it.each([false, true])( + 'supersedes an initial read once (connection refresh: %s)', + async connectionRefresh => { + const initial = Promise.withResolvers(); + const initialStarted = Promise.withResolvers(); + const next = Promise.withResolvers(); + const nextStarted = Promise.withResolvers(); + const signals: AbortSignal[] = []; + const observer = new QueryObserver(queryClient, { + queryKey, + staleTime: Infinity, + initialData: connectionRefresh ? { snapshot: null } : undefined, + queryFn: ({ signal }) => { + signals.push(signal); + if (signals.length === 1) { + initialStarted.resolve(); + return initial.promise; + } + nextStarted.resolve(); + return next.promise; + }, + }); + const unsubscribe = observer.subscribe(() => {}); + const refresher = createWorktreeChangesRefresher(queryClient, queryKey); + try { + if (connectionRefresh) void refresher.request({ revision: 0, force: true }); + await initialStarted.promise; + const refreshing = refresher.request({ revision: 4, force: false }); + await nextStarted.promise; + expect(signals[0].aborted).toBe(true); + expect(signals[1].aborted).toBe(false); + next.resolve(saved(4)); + await refreshing; + initial.resolve(saved(3)); + await initial.promise; + expect(queryClient.getQueryData(queryKey)).toEqual(saved(4)); + expect(signals).toHaveLength(2); + } finally { + initial.resolve(saved(3)); + next.resolve(saved(4)); + refresher.dispose(); + unsubscribe(); + } + } + ); + + it.each([ + { returnedRevision: 4, reconnected: false, expectedReads: 2 }, + { returnedRevision: 7, reconnected: false, expectedReads: 1 }, + { returnedRevision: 7, reconnected: true, expectedReads: 2 }, + ])( + 'coalesces a burst with reply $returnedRevision and reconnect $reconnected into $expectedReads reads', + async ({ returnedRevision, reconnected, expectedReads }) => { + queryClient.setQueryData(queryKey, saved(3)); + const first = Promise.withResolvers(); + const firstStarted = Promise.withResolvers(); + const signals: AbortSignal[] = []; + const observer = new QueryObserver(queryClient, { + queryKey, + staleTime: Infinity, + queryFn: ({ signal }) => { + signals.push(signal); + if (signals.length === 1) { + firstStarted.resolve(); + return first.promise; + } + return Promise.resolve(saved(7)); + }, + }); + const unsubscribe = observer.subscribe(() => {}); + const refresher = createWorktreeChangesRefresher(queryClient, queryKey); + try { + const refreshing = refresher.request({ revision: 4, force: false }); + await firstStarted.promise; + for (const revision of [5, 7, 6]) { + expect(refresher.request({ revision, force: false })).toBe(refreshing); + } + if (reconnected) { + expect(refresher.request({ revision: 6, force: true })).toBe(refreshing); + } + expect(signals).toHaveLength(1); + expect(signals[0].aborted).toBe(false); + first.resolve(saved(returnedRevision)); + await refreshing; + expect(signals).toHaveLength(expectedReads); + expect(signals.every(signal => !signal.aborted)).toBe(true); + expect(queryClient.getQueryData(queryKey)).toEqual(saved(7)); + } finally { + first.resolve(saved(7)); + refresher.dispose(); + unsubscribe(); + } + } + ); + + it('discards queued reads on disposal without touching another session cache', async () => { + queryClient.setQueryData(queryKey, saved(3)); + const otherKey = ['worktree-changes', 'organization-b', 'workspace-b']; + queryClient.setQueryData(otherKey, saved(1)); + const first = Promise.withResolvers(); + const started = Promise.withResolvers(); + const queryFn = jest.fn(() => { + started.resolve(); + return first.promise; + }); + const observer = new QueryObserver(queryClient, { queryKey, staleTime: Infinity, queryFn }); + const unsubscribe = observer.subscribe(() => {}); + const refresher = createWorktreeChangesRefresher(queryClient, queryKey); + try { + const refreshing = refresher.request({ revision: 4, force: false }); + await started.promise; + void refresher.request({ revision: 5, force: true }); + refresher.dispose(); + first.resolve(saved(4)); + await refreshing; + await refresher.request({ revision: 6, force: true }); + expect(queryFn).toHaveBeenCalledTimes(1); + expect(queryClient.getQueryData(otherKey)).toEqual(saved(1)); + } finally { + first.resolve(saved(4)); + refresher.dispose(); + unsubscribe(); + } + }); + + it('lets a read retry recover without another ready notification', async () => { + queryClient.setQueryData(queryKey, saved(3)); + const queryFn = jest + .fn, []>() + .mockRejectedValueOnce(new Error('Transient read failure')) + .mockResolvedValue(saved(4)); + const observer = new QueryObserver(queryClient, { + queryKey, + staleTime: Infinity, + retry: 2, + retryDelay: 0, + queryFn, + }); + const unsubscribe = observer.subscribe(() => {}); + const refresher = createWorktreeChangesRefresher(queryClient, queryKey); + try { + await refresher.request({ revision: 4, force: false }); + expect(queryFn).toHaveBeenCalledTimes(2); + expect(queryClient.getQueryData(queryKey)).toEqual(saved(4)); + } finally { + refresher.dispose(); + unsubscribe(); + } + }); +}); + describe('worktree changes messages', () => { it('keeps saved content while refresh is pending', () => { expect(messages({ refreshPending: true })).toEqual({ diff --git a/apps/web/src/components/cloud-agent-next/worktree-changes.ts b/apps/web/src/components/cloud-agent-next/worktree-changes.ts index 42bc325ff4..76c5cf4078 100644 --- a/apps/web/src/components/cloud-agent-next/worktree-changes.ts +++ b/apps/web/src/components/cloud-agent-next/worktree-changes.ts @@ -1,3 +1,4 @@ +import type { QueryClient, QueryKey } from '@tanstack/react-query'; import { getWorktreeChangesOutputSchema, type GetWorktreeChangesOutput, @@ -137,6 +138,55 @@ export function preserveNewerWorktreeChanges( return next; } +export function createWorktreeChangesRefresher(queryClient: QueryClient, queryKey: QueryKey) { + const filters = { queryKey, exact: true }; + let pending: { revision: number; force: boolean } | undefined; + let inFlight: Promise | undefined; + let activeRevision: number | undefined; + let disposed = false; + + async function drain(): Promise { + try { + await queryClient.cancelQueries(filters, { silent: true, revert: false }); + while (!disposed && pending) { + const request = pending; + pending = undefined; + const cached = queryClient.getQueryData(queryKey); + if (!request.force && (cached?.snapshot?.revision ?? 0) >= request.revision) continue; + activeRevision = request.revision; + await queryClient.invalidateQueries( + { ...filters, refetchType: 'active' }, + { cancelRefetch: false } + ); + activeRevision = undefined; + } + } finally { + activeRevision = undefined; + inFlight = undefined; + } + } + + return { + request(request: { revision: number; force: boolean }): Promise { + if (disposed) return Promise.resolve(); + pending = { + revision: Math.max(pending?.revision ?? 0, request.revision), + force: pending?.force === true || request.force, + }; + if (activeRevision === 0 && request.revision > 0) { + activeRevision = undefined; + void queryClient.cancelQueries(filters, { silent: true, revert: false }); + } + inFlight ??= drain(); + return inFlight; + }, + dispose(): void { + disposed = true; + pending = undefined; + }, + }; +} + export function worktreeChangesMessages({ snapshot, savedReadPending, diff --git a/apps/web/tests/e2e/cloud-agent-sandbox-status.spec.ts b/apps/web/tests/e2e/cloud-agent-sandbox-status.spec.ts index cdec9aff5e..3d7a8bb37c 100644 --- a/apps/web/tests/e2e/cloud-agent-sandbox-status.spec.ts +++ b/apps/web/tests/e2e/cloud-agent-sandbox-status.spec.ts @@ -38,6 +38,7 @@ type SessionFixture = { kind?: 'remote' | 'read-only' | 'unresolved' | 'unrelated'; }; type StatusRequest = { cloudAgentSessionId: string; organizationId?: string; at: number }; +type WorktreeChangesRequest = { cloudAgentSessionId: string; organizationId?: string }; type RpcResult = | { result: { data: unknown } } | { error: { message: string; code: number; data: { code: string; httpStatus: number } } }; @@ -64,6 +65,33 @@ function deferred() { return { promise, resolve }; } +function worktreeSnapshot(revision: number): WorktreeChangesSnapshot { + return { + schemaVersion: 1, + revision, + capturedAt: new Date(baseTime + revision).toISOString(), + comparison: { + baseRef: 'refs/remotes/origin/main', + mergeBase: 'a'.repeat(40), + head: 'b'.repeat(40), + }, + files: [ + { + path: `src/revision-${revision}.ts`, + status: 'modified', + additions: revision, + deletions: 0, + tracked: true, + binary: false, + countsComplete: true, + }, + ], + truncated: false, + }; +} + +const worktreeSummary = (revision: number) => `1 changed files, ${revision} additions, 0 deletions`; + async function mountFixtures( page: Page, sessions: SessionFixture[] = [ @@ -74,10 +102,21 @@ async function mountFixtures( let now = baseTime; await page.clock.install({ time: new Date(baseTime) }); const statusRequests: StatusRequest[] = []; + const worktreeRequests: WorktreeChangesRequest[] = []; const procedures: string[] = []; const sockets = new Map(); + const replayedReadyRevisions = new Map(); + const savedChanges = new Map(); + const worktreeKey = (cloudId: string, organizationId?: string) => + `${organizationId ?? 'personal'}:${cloudId}`; let reply: (request: StatusRequest) => RpcResult | Promise = () => success(snapshot()); - let savedChanges: WorktreeChangesSnapshot | null = null; + let worktreeReply: ( + request: WorktreeChangesRequest + ) => RpcResult | Promise = request => + success({ + snapshot: + savedChanges.get(worktreeKey(request.cloudAgentSessionId, request.organizationId)) ?? null, + }); let eventId = 0; function snapshot(overrides: Partial = {}): SandboxStatusSnapshot { @@ -123,11 +162,11 @@ async function mountFixtures( }; } - function send(cloudId: string, streamEventType: string, data: unknown) { + function send(cloudId: string, streamEventType: string, data: unknown, sessionId = cloudId) { sockets.get(cloudId)?.send( JSON.stringify({ eventId: ++eventId, - sessionId: cloudId, + sessionId, streamEventType, timestamp: new Date(now).toISOString(), data, @@ -145,6 +184,10 @@ async function mountFixtures( if (!cloudId) return; sockets.set(cloudId, socket); send(cloudId, 'connected', { sessionStatus: { type: 'idle' }, cloudStatus: { type: 'ready' } }); + for (const revision of replayedReadyRevisions.get(cloudId) ?? []) { + send(cloudId, 'cloud.worktree.changes.ready', { revision }); + } + replayedReadyRevisions.delete(cloudId); }); await page.route('**/api/cloud-agent-next/sessions/stream-ticket', route => route.fulfill({ @@ -206,12 +249,19 @@ async function mountFixtures( }); if (procedure === 'cliSessionsV2.getSessionMessages') return success({ info: { id: args.session_id }, messages: [] }); - if (procedure.endsWith('.getWorktreeChanges')) return success({ snapshot: savedChanges }); - if (procedure.endsWith('.refreshWorktreeChanges')) - return success({ - status: savedChanges ? 'refreshed' : 'offline', - snapshot: savedChanges, - }); + if (procedure.endsWith('.getWorktreeChanges')) { + const worktreeRequest: WorktreeChangesRequest = { + cloudAgentSessionId: args.cloudAgentSessionId, + organizationId: args.organizationId, + }; + worktreeRequests.push(worktreeRequest); + return worktreeReply(worktreeRequest); + } + if (procedure.endsWith('.refreshWorktreeChanges')) { + const snapshot = + savedChanges.get(worktreeKey(args.cloudAgentSessionId, args.organizationId)) ?? null; + return success({ status: snapshot ? 'refreshed' : 'offline', snapshot }); + } if (procedure.endsWith('.getComputeBillingStatus')) return success({ phase: 'unavailable' }); if (procedure.endsWith('.sendMessage')) @@ -234,13 +284,38 @@ async function mountFixtures( return { statusRequests, + worktreeRequests, procedures, snapshot, setReply(handler: typeof reply) { reply = handler; }, - setWorktreeChanges(snapshot: WorktreeChangesSnapshot) { - savedChanges = snapshot; + setWorktreeReply(handler: typeof worktreeReply) { + worktreeReply = handler; + }, + setWorktreeChanges( + snapshot: WorktreeChangesSnapshot, + cloudId = firstWorkspace, + organizationId?: string + ) { + savedChanges.set(worktreeKey(cloudId, organizationId), snapshot); + }, + async worktreeReady(cloudId: string, revision: number, sessionId = cloudId) { + await expect.poll(() => sockets.has(cloudId)).toBe(true); + send(cloudId, 'cloud.worktree.changes.ready', { revision }, sessionId); + }, + async reconnect(cloudId: string, readyRevisions: number[] = []) { + const previous = sockets.get(cloudId); + if (!previous) throw new Error('Missing stream socket'); + replayedReadyRevisions.set(cloudId, readyRevisions); + await previous.close({ code: 1012, reason: 'Test stream reconnect' }); + await expect + .poll(async () => { + await page.clock.fastForward(1_000); + return sockets.get(cloudId) !== previous; + }) + .toBe(true); + now = await page.evaluate(() => Date.now()); }, async open(id = firstId, organizationId?: string) { await page.goto( @@ -461,79 +536,241 @@ test.describe('control-plane sandbox header', () => { } }); - test('updates saved file changes during a turn and while the drawer is open without polling captures', async ({ + for (const scope of ['personal', 'organization'] as const) { + test(`pushes saved file changes and reconciles reconnects without polling in ${scope} scope`, async ({ + page, + }) => { + async function run(organizationId?: string) { + const fixture = await mountFixtures(page, [ + { id: firstId, cloudId: firstWorkspace, title: 'Push updates', organizationId }, + ]); + const readCount = () => fixture.worktreeRequests.length; + const captureCount = () => + fixture.procedures.filter(procedure => procedure.endsWith('.refreshWorktreeChanges')) + .length; + const publish = (revision: number) => + fixture.setWorktreeChanges(worktreeSnapshot(revision), firstWorkspace, organizationId); + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + const drawer = page.getByRole('dialog', { name: 'Changes', exact: true }); + publish(1); + await fixture.open(firstId, organizationId); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(1)); + await fixture.activity(firstWorkspace, 'busy'); + await expect( + page.getByRole('button', { name: 'Stop response', exact: true }) + ).toBeVisible(); + await fixture.advance(1_000); + const initialReads = readCount(); + publish(2); + await fixture.advance(15_000); + expect(readCount()).toBe(initialReads); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(1)); + + await fixture.worktreeReady(firstWorkspace, 2); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(2)); + expect(readCount()).toBe(initialReads + 1); + expect(captureCount()).toBe(0); + await fixture.worktreeReady(firstWorkspace, 2); + await fixture.worktreeReady(firstWorkspace, 1); + await fixture.activity(firstWorkspace, 'idle'); + await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toHaveCount( + 0 + ); + publish(3); + await fixture.advance(40_000); + expect(readCount()).toBe(initialReads + 1); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(2)); + + await changes.click(); + await expect(drawer.getByText('revision-3.ts', { exact: true })).toBeVisible(); + await expect.poll(captureCount).toBe(1); + const openReads = readCount(); + publish(4); + await fixture.advance(15_000); + expect(readCount()).toBe(openReads); + await expect(drawer.getByText('revision-3.ts', { exact: true })).toBeVisible(); + await fixture.worktreeReady(firstWorkspace, 4); + await expect(drawer.getByText('revision-4.ts', { exact: true })).toBeVisible(); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(4)); + expect(readCount()).toBe(openReads + 1); + expect(captureCount()).toBe(1); + await page.keyboard.press('Escape'); + await expect(drawer).toHaveCount(0); + const closedReads = readCount(); + publish(5); + await fixture.advance(15_000); + expect(readCount()).toBe(closedReads); + await fixture.reconnect(firstWorkspace, [4, 3]); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(5)); + expect(readCount()).toBe(closedReads + 1); + expect(captureCount()).toBe(1); + expect( + fixture.worktreeRequests.every( + request => + request.cloudAgentSessionId === firstWorkspace && + request.organizationId === organizationId + ) + ).toBe(true); + } + if (scope === 'organization') await withOrganization(page, run); + else await run(); + }); + } + + test('retries a transient final saved read without another ready notification', async ({ page, }) => { const fixture = await mountFixtures(page); - const readCount = () => - fixture.procedures.filter(procedure => procedure.endsWith('.getWorktreeChanges')).length; - const captureCount = () => - fixture.procedures.filter(procedure => procedure.endsWith('.refreshWorktreeChanges')).length; - function publishRevision(revision: number) { - fixture.setWorktreeChanges({ - schemaVersion: 1, - revision, - capturedAt: new Date(baseTime + revision).toISOString(), - comparison: { - baseRef: 'refs/remotes/origin/main', - mergeBase: 'a'.repeat(40), - head: 'b'.repeat(40), - }, - files: [ - { - path: `src/revision-${revision}.ts`, - status: 'modified', - additions: revision, - deletions: 0, - tracked: true, - binary: false, - countsComplete: true, - }, - ], - truncated: false, - }); - } - const changes = page.getByRole('button', { name: 'Changes', exact: true }); - const drawer = page.getByRole('dialog', { name: 'Changes', exact: true }); - const summary = (revision: number) => `1 changed files, ${revision} additions, 0 deletions`; - publishRevision(1); + fixture.setWorktreeChanges(worktreeSnapshot(1)); await fixture.open(); - await expect(changes).toHaveAttribute('aria-description', summary(1)); - const initialReads = readCount(); + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(1)); + await fixture.activity(firstWorkspace, 'busy'); + await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toBeVisible(); + await fixture.advance(1_000); + let attempts = 0; + fixture.setWorktreeReply(() => + ++attempts === 1 ? failure() : success({ snapshot: worktreeSnapshot(2) }) + ); + await fixture.worktreeReady(firstWorkspace, 2); + await expect.poll(() => attempts).toBe(1); + await fixture.advance(1_100); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(2)); + expect(attempts).toBe(2); await fixture.advance(15_000); - expect(readCount()).toBe(initialReads); + expect(attempts).toBe(2); + }); + + for (const { code, attempts: expectedAttempts } of [ + { code: 'INTERNAL_SERVER_ERROR', attempts: 3 }, + { code: 'FORBIDDEN', attempts: 1 }, + ]) { + test(`bounds saved-read retries for ${code} and accepts a later ready revision`, async ({ + page, + }) => { + const fixture = await mountFixtures(page); + fixture.setWorktreeChanges(worktreeSnapshot(1)); + await fixture.open(); + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(1)); + await fixture.activity(firstWorkspace, 'busy'); + await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toBeVisible(); + await fixture.advance(1_000); + let attempts = 0; + fixture.setWorktreeReply(() => { + attempts++; + return failure(code); + }); + await fixture.worktreeReady(firstWorkspace, 2); + await expect.poll(() => attempts).toBe(1); + for (let attempt = 2; attempt <= expectedAttempts; attempt++) { + await fixture.advance(2_100); + await expect.poll(() => attempts).toBe(attempt); + } + await fixture.advance(30_000); + expect(attempts).toBe(expectedAttempts); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(1)); + fixture.setWorktreeReply(() => success({ snapshot: worktreeSnapshot(3) })); + await fixture.worktreeReady(firstWorkspace, 3); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(3)); + }); + } + test('coalesces ready bursts and reconnects behind one saved read', async ({ page }) => { + const fixture = await mountFixtures(page); + fixture.setWorktreeChanges(worktreeSnapshot(1)); + await fixture.open(); + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(1)); await fixture.activity(firstWorkspace, 'busy'); await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toBeVisible(); - publishRevision(2); - await fixture.advance(5_000); - await expect(changes).toHaveAttribute('aria-description', summary(2)); - expect(readCount()).toBeGreaterThan(initialReads); - expect(captureCount()).toBe(0); + await fixture.advance(1_000); + const held = deferred(); + let reads = 0; + fixture.setWorktreeReply(() => + ++reads === 1 ? held.promise : success({ snapshot: worktreeSnapshot(5) }) + ); + try { + await fixture.worktreeReady(firstWorkspace, 2); + await expect.poll(() => reads).toBe(1); + for (const revision of [3, 5, 4]) { + await fixture.worktreeReady(firstWorkspace, revision); + await fixture.advance(1); + } + await fixture.reconnect(firstWorkspace, [4, 3]); + await fixture.advance(1_000); + expect(reads).toBe(1); + held.resolve(success({ snapshot: worktreeSnapshot(2) })); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(5)); + expect(reads).toBe(2); + await fixture.advance(15_000); + expect(reads).toBe(2); + } finally { + held.resolve(success({ snapshot: worktreeSnapshot(2) })); + } + }); - await fixture.activity(firstWorkspace, 'idle'); - await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toHaveCount(0); - publishRevision(3); - await fixture.advance(31_000); - await expect(changes).toHaveAttribute('aria-description', summary(3)); - const idleReads = readCount(); - await fixture.advance(15_000); - expect(readCount()).toBe(idleReads); + test('a ready notification supersedes an in-flight initial saved read', async ({ page }) => { + const fixture = await mountFixtures(page); + let abortedReads = 0; + page.on('requestfailed', request => { + if ( + request.url().includes('.getWorktreeChanges') && + request.failure()?.errorText.includes('ABORTED') + ) + abortedReads++; + }); + const held = deferred(); + let ready = false; + fixture.setWorktreeReply(() => + ready ? success({ snapshot: worktreeSnapshot(2) }) : held.promise + ); + try { + await fixture.open(); + await fixture.activity(firstWorkspace, 'busy'); + await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toBeVisible(); + await fixture.advance(1_000); + await expect.poll(() => fixture.worktreeRequests.length).toBeGreaterThan(0); + const initialReads = fixture.worktreeRequests.length; + const initialAbortedReads = abortedReads; + ready = true; + await fixture.worktreeReady(firstWorkspace, 2); + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(2)); + expect(fixture.worktreeRequests.length).toBe(initialReads + 1); + await expect.poll(() => abortedReads).toBeGreaterThan(initialAbortedReads); + held.resolve(success({ snapshot: worktreeSnapshot(1) })); + await fixture.advance(15_000); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(2)); + expect(fixture.worktreeRequests.length).toBe(initialReads + 1); + } finally { + held.resolve(success({ snapshot: worktreeSnapshot(1) })); + } + }); - await changes.click(); - await expect(drawer.getByText('revision-3.ts', { exact: true })).toBeVisible(); - await expect.poll(captureCount).toBe(1); - publishRevision(4); - await fixture.advance(5_000); - await expect(drawer.getByText('revision-4.ts', { exact: true })).toBeVisible(); - await expect(changes).toHaveAttribute('aria-description', summary(4)); - expect(captureCount()).toBe(1); - await page.keyboard.press('Escape'); - await expect(drawer).toHaveCount(0); - const closedReads = readCount(); + test('ignores ready notifications for another session after navigation', async ({ page }) => { + const fixture = await mountFixtures(page); + fixture.setWorktreeChanges(worktreeSnapshot(7)); + fixture.setWorktreeChanges(worktreeSnapshot(1), secondWorkspace); + const changes = page.getByRole('button', { name: 'Changes', exact: true }); + await fixture.open(); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(7)); + await fixture.navigate(secondId); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(1)); + await fixture.activity(secondWorkspace, 'busy'); + await expect(page.getByRole('button', { name: 'Stop response', exact: true })).toBeVisible(); + await fixture.advance(1_000); + const selectedReads = fixture.worktreeRequests.length; + fixture.setWorktreeChanges(worktreeSnapshot(2), secondWorkspace); + await fixture.worktreeReady(secondWorkspace, 8, firstWorkspace); await fixture.advance(15_000); - expect(readCount()).toBe(closedReads); - expect(captureCount()).toBe(1); + expect(fixture.worktreeRequests.length).toBe(selectedReads); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(1)); + await fixture.worktreeReady(secondWorkspace, 2); + await expect(changes).toHaveAttribute('aria-description', worktreeSummary(2)); + expect(fixture.worktreeRequests.slice(selectedReads)).toEqual([ + { cloudAgentSessionId: secondWorkspace, organizationId: undefined }, + ]); }); for (const { width, reducedMotion } of [ diff --git a/packages/cloud-agent-sdk/src/__fixtures__/helpers.ts b/packages/cloud-agent-sdk/src/__fixtures__/helpers.ts index f99c9ea58d..8716832610 100644 --- a/packages/cloud-agent-sdk/src/__fixtures__/helpers.ts +++ b/packages/cloud-agent-sdk/src/__fixtures__/helpers.ts @@ -1,6 +1,6 @@ import type { CloudAgentEvent } from '../event-types'; -function createEventHelpers() { +function createEventHelpers(defaultSessionId = 'ses-1') { let eventCounter = 0; function resetCounter() { @@ -10,7 +10,7 @@ function createEventHelpers() { function createEvent( streamEventType: string, data: unknown, - sessionId = 'ses-1' + sessionId = defaultSessionId ): CloudAgentEvent { return { eventId: ++eventCounter, @@ -22,7 +22,11 @@ function createEventHelpers() { }; } - function kilocode(type: string, properties: unknown, sessionId = 'ses-1'): CloudAgentEvent { + function kilocode( + type: string, + properties: unknown, + sessionId = defaultSessionId + ): CloudAgentEvent { return createEvent('kilocode', { type, properties }, sessionId); } diff --git a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts index c398f2d697..eb9b0aa953 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts @@ -186,6 +186,40 @@ describe('CloudAgentTransport event routing', () => { transport.destroy(); }); + it('routes snapshot-ready events without executionId only to the service sink', async () => { + const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); + transport.connect(); + await flushPromises(); + sendRaw({ + eventId: 5, + sessionId: 'ses-1', + streamEventType: 'cloud.worktree.changes.ready', + timestamp: '2026-09-02T00:00:00.000Z', + data: { revision: 3 }, + }); + expect(serviceEvents.at(-1)).toEqual({ + type: 'worktree.changes.ready', + cloudSessionId: 'ses-1', + revision: 3, + }); + expect(chatEvents).toEqual([]); + transport.destroy(); + }); + + it.each(['cloud.worktree.changes.ready', 'connected'])( + 'drops %s envelopes for another Cloud session', + async streamEventType => { + const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); + transport.connect(); + await flushPromises(); + const previousEvents = [...serviceEvents]; + sendRaw(createEvent(streamEventType, { revision: 3 }, 'other-session')); + expect(serviceEvents).toEqual(previousEvents); + expect(chatEvents).toEqual([]); + transport.destroy(); + } + ); + it('routes mixed events to correct sinks', async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); @@ -414,6 +448,26 @@ describe('CloudAgentTransport lifecycle', () => { expect(mockWs.close).toHaveBeenCalled(); }); + it.each(['disconnect', 'destroy'] as const)( + '%s rejects late ready and connected frames', + async stop => { + const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); + transport.connect(); + await flushPromises(); + const oldOnMessage = mockWs.onmessage; + const previousEvents = [...serviceEvents]; + transport[stop](); + for (const streamEventType of ['cloud.worktree.changes.ready', 'connected']) { + oldOnMessage?.({ + data: JSON.stringify(createEvent(streamEventType, { revision: 3 })), + } as MessageEvent); + } + expect(serviceEvents).toEqual(previousEvents); + expect(chatEvents).toEqual([]); + transport.destroy(); + } + ); + it('stale generation after disconnect prevents connection creation', async () => { const resolveTicket: { resolve?: (value: string) => void } = {}; const getTicket = jest.fn( diff --git a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts index f9be44c921..3c9707f201 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts @@ -201,6 +201,9 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport lifecycleHooks: config.lifecycleHooks, websocketHeaders: config.websocketHeaders, onEvent: raw => { + if (expectedGeneration !== lifecycleGeneration || raw.sessionId !== config.sessionId) { + return; + } // Track high-water mark for reconnect fromId only. Do not filter // by eventId: the DO entity-upserts tool/message parts under a // stable row id and rebroadcasts that same (or older) id with a diff --git a/packages/cloud-agent-sdk/src/index.ts b/packages/cloud-agent-sdk/src/index.ts index 66324e1919..0cb38b9db4 100644 --- a/packages/cloud-agent-sdk/src/index.ts +++ b/packages/cloud-agent-sdk/src/index.ts @@ -12,6 +12,7 @@ export type { SessionManager, SessionManagerConfig, SessionManagerAtoms, + WorktreeChangesRefresh, SessionStatusIndicator, SessionConfig, StandalonePermission, diff --git a/packages/cloud-agent-sdk/src/normalizer.test.ts b/packages/cloud-agent-sdk/src/normalizer.test.ts index 034798a6a9..c7f7f0a125 100644 --- a/packages/cloud-agent-sdk/src/normalizer.test.ts +++ b/packages/cloud-agent-sdk/src/normalizer.test.ts @@ -1344,12 +1344,63 @@ describe('normalize', () => { }); }); + describe('cloud.worktree.changes.ready', () => { + it('normalizes an execution-independent service event with the Cloud envelope identity', () => { + const result = normalize({ + eventId: 1, + sessionId: 'agent-1', + streamEventType: 'cloud.worktree.changes.ready', + timestamp: '2026-09-02T00:00:00.000Z', + data: { revision: 4 }, + }); + expect(result).toEqual({ + type: 'worktree.changes.ready', + cloudSessionId: 'agent-1', + revision: 4, + }); + if (!result) throw new Error('Expected worktree changes event'); + expect(isChatEvent(result)).toBe(false); + }); + + it.each([ + {}, + { revision: 0 }, + { revision: '1' }, + { revision: NaN }, + { revision: 1, extra: true }, + ])('rejects invalid data %p', data => + expect(normalize(createRaw('cloud.worktree.changes.ready', data))).toBeNull() + ); + + it.each(['cloud.worktree.changes.ready', 'worktree.changes.ready'])( + 'does not accept %s as a Kilo event', + type => { + expect(normalize(createKilocode(type, { revision: 1 }))).toBeNull(); + expect(normalizeCliEvent(type, { revision: 1 })).toBeNull(); + } + ); + }); + describe('connected', () => { + it('uses only the Cloud envelope session ID for refresh routing', () => { + expect( + normalize(createRaw('connected', { cloudSessionId: 'agent-other' }, 'agent-1')) + ).toEqual({ + type: 'connected', + cloudSessionId: 'agent-1', + }); + expect(normalizeCliEvent('connected', { cloudSessionId: 'agent-1' })).toEqual({ + type: 'connected', + }); + expect(normalize(createKilocode('connected', {}, 'agent-1'))).toEqual({ type: 'connected' }); + }); + it.each(['active-message', null])( 'retains authoritative active identity %p', activeMessageId => { expect(normalize(createRaw('connected', { activeMessageId }))).toEqual({ type: 'connected', + cloudSessionId: 'ses-1', activeMessageId, }); } @@ -1365,7 +1416,7 @@ describe('normalize', () => { sessionStatus: { type: 'busy' }, }) ) - ).toEqual({ type: 'connected', sessionStatus: { type: 'busy' } }); + ).toEqual({ type: 'connected', cloudSessionId: 'ses-1', sessionStatus: { type: 'busy' } }); } ); @@ -1377,6 +1428,7 @@ describe('normalize', () => { ); expect(result).toEqual({ type: 'connected', + cloudSessionId: 'ses-1', sessionStatus: { type: 'busy' }, }); }); @@ -1390,6 +1442,7 @@ describe('normalize', () => { ); expect(result).toEqual({ type: 'connected', + cloudSessionId: 'ses-1', sessionStatus: { type: 'idle' }, cloudStatus: { type: 'ready' }, }); @@ -1403,6 +1456,7 @@ describe('normalize', () => { ); expect(result).toEqual({ type: 'connected', + cloudSessionId: 'ses-1', cloudStatus: { type: 'preparing' }, }); }); @@ -1417,6 +1471,7 @@ describe('normalize', () => { ); expect(result).toEqual({ type: 'connected', + cloudSessionId: 'ses-1', sessionStatus: { type: 'busy' }, cloudStatus: { type: 'preparing', step: 'cloning' }, }); @@ -1424,13 +1479,13 @@ describe('normalize', () => { it('normalizes without sessionStatus', () => { const result = normalize(createRaw('connected', {})); - expect(result).toEqual({ type: 'connected' }); + expect(result).toEqual({ type: 'connected', cloudSessionId: 'ses-1' }); expect(result).not.toHaveProperty('sessionStatus'); }); it('ignores malformed sessionStatus', () => { const result = normalize(createRaw('connected', { sessionStatus: 'busy' })); - expect(result).toEqual({ type: 'connected' }); + expect(result).toEqual({ type: 'connected', cloudSessionId: 'ses-1' }); expect(result).not.toHaveProperty('sessionStatus'); }); @@ -1443,6 +1498,7 @@ describe('normalize', () => { ); expect(result).toEqual({ type: 'connected', + cloudSessionId: 'ses-1', sessionStatus: { type: 'idle' }, }); }); @@ -1456,6 +1512,7 @@ describe('normalize', () => { ); expect(result).toEqual({ type: 'connected', + cloudSessionId: 'ses-1', sessionStatus: { type: 'idle' }, }); expect(result).not.toHaveProperty('question'); @@ -1477,6 +1534,7 @@ describe('normalize', () => { ); expect(result).toEqual({ type: 'connected', + cloudSessionId: 'ses-1', sessionStatus: { type: 'busy' }, }); expect(result).not.toHaveProperty('permission'); diff --git a/packages/cloud-agent-sdk/src/normalizer.ts b/packages/cloud-agent-sdk/src/normalizer.ts index 7749b5be6f..2ead98c57a 100644 --- a/packages/cloud-agent-sdk/src/normalizer.ts +++ b/packages/cloud-agent-sdk/src/normalizer.ts @@ -14,6 +14,7 @@ import type { } from './types'; import { cloudAgentEventSchema, + cloudWorktreeChangesReadyDataSchema, kilocodePayloadSchema, messageUpdatedDataSchema, messagePartUpdatedDataSchema, @@ -160,11 +161,13 @@ export type ServiceEvent = | { type: 'cloud.status'; cloudStatus: CloudStatus } | { type: 'connected'; + cloudSessionId?: string; sessionStatus?: SessionStatus | undefined; cloudStatus?: CloudStatus | undefined; activeMessageId?: string | null | undefined; } | { type: 'commands.available'; commands: SlashCommandInfo[] } + | { type: 'worktree.changes.ready'; cloudSessionId: string; revision: number } | { type: 'cloud.message.queued'; messageId: string; @@ -605,6 +608,16 @@ function normalizeInnerEvent(eventType: string, data: unknown): NormalizedEvent export function normalize(raw: CloudAgentEvent): NormalizedEvent | null { if (!cloudAgentEventSchema.safeParse(raw).success) return null; + if (raw.streamEventType === 'cloud.worktree.changes.ready') { + const r = cloudWorktreeChangesReadyDataSchema.safeParse(raw.data); + if (!r.success) return null; + return { + type: 'worktree.changes.ready', + cloudSessionId: raw.sessionId, + revision: r.data.revision, + }; + } + let eventType = raw.streamEventType; let data: unknown = raw.data; @@ -614,7 +627,11 @@ export function normalize(raw: CloudAgentEvent): NormalizedEvent | null { data = kilo.data.properties; } - return normalizeInnerEvent(eventType, data); + const event = normalizeInnerEvent(eventType, data); + if (raw.streamEventType === 'connected' && event?.type === 'connected') { + return { ...event, cloudSessionId: raw.sessionId }; + } + return event; } /** diff --git a/packages/cloud-agent-sdk/src/schemas.test.ts b/packages/cloud-agent-sdk/src/schemas.test.ts index 28af0c7bac..135a4cc63c 100644 --- a/packages/cloud-agent-sdk/src/schemas.test.ts +++ b/packages/cloud-agent-sdk/src/schemas.test.ts @@ -1,10 +1,60 @@ import { activeSessionSchema, + cloudAgentEventSchema, + cloudWorktreeChangesReadyDataSchema, parseCustomerBillingFailure, sessionEventPayloadSchema, sessionEventV2RowSchema, } from './schemas'; +describe('cloudWorktreeChangesReadyDataSchema', () => { + it.each([1, 42, Number.MAX_SAFE_INTEGER])('accepts positive safe revision %p', revision => { + expect(cloudWorktreeChangesReadyDataSchema.parse({ revision })).toEqual({ revision }); + }); + + it.each([ + NaN, + Infinity, + -Infinity, + Number.MAX_SAFE_INTEGER + 1, + 0, + -1, + 1.5, + '1', + true, + null, + undefined, + {}, + [], + ])('rejects invalid revision %p', revision => { + expect(cloudWorktreeChangesReadyDataSchema.safeParse({ revision }).success).toBe(false); + }); + + it.each([null, undefined, [], {}, { revision: 1, extra: true }])( + 'rejects invalid payload %p', + data => { + expect(cloudWorktreeChangesReadyDataSchema.safeParse(data).success).toBe(false); + } + ); + + it('keeps envelope event names open and accepts events without executionId', () => { + const envelope = { + eventId: 1, + sessionId: 'agent-1', + streamEventType: 'cloud.worktree.changes.ready', + timestamp: '2026-09-02T00:00:00.000Z', + data: { revision: 1 }, + }; + expect(cloudAgentEventSchema.parse(envelope)).toEqual(envelope); + expect( + cloudAgentEventSchema.safeParse({ + ...envelope, + streamEventType: 'future.event', + }).success + ).toBe(true); + }); +}); + describe('parseCustomerBillingFailure', () => { const failure = { code: 'COMPUTE_STOPPING', diff --git a/packages/cloud-agent-sdk/src/schemas.ts b/packages/cloud-agent-sdk/src/schemas.ts index 54f5cf3807..c46e0c83c0 100644 --- a/packages/cloud-agent-sdk/src/schemas.ts +++ b/packages/cloud-agent-sdk/src/schemas.ts @@ -16,6 +16,10 @@ export const cloudAgentEventSchema = z.object({ }); export type CloudAgentEvent = z.infer; +export const cloudWorktreeChangesReadyDataSchema = z + .object({ revision: z.number().int().positive().max(Number.MAX_SAFE_INTEGER) }) + .strict(); + export const streamErrorSchema = z.object({ type: z.literal('error'), code: z.enum([ diff --git a/packages/cloud-agent-sdk/src/session-manager.test.ts b/packages/cloud-agent-sdk/src/session-manager.test.ts index a528fc6b1d..2ebd92551c 100644 --- a/packages/cloud-agent-sdk/src/session-manager.test.ts +++ b/packages/cloud-agent-sdk/src/session-manager.test.ts @@ -450,6 +450,204 @@ describe('createSessionManager', () => { // switchSession // ------------------------------------------------------------------------- + describe('worktreeChangesRefresh', () => { + const ready = { + type: 'worktree.changes.ready', + cloudSessionId: 'agent-1', + revision: 2, + } satisfies NormalizedEvent; + const connected = { + type: 'connected', + cloudSessionId: 'agent-1', + sessionStatus: { type: 'idle' }, + } satisfies NormalizedEvent; + + it('publishes ready revisions and every idle connected event without changing chat state', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + const { store } = config; + expect(store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + await mgr.switchSession(kiloId('ses-1')); + expect(store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + const state = { + activity: store.get(mgr.atoms.activity), + status: store.get(mgr.atoms.agentStatus), + messages: store.get(mgr.atoms.messagesList), + pending: store.get(mgr.atoms.pendingMessages), + }; + const listener = jest.fn(); + const unsubscribe = store.sub(mgr.atoms.worktreeChangesRefresh, listener); + + mockSessionCallbacks.onEvent?.(ready); + expect(store.get(mgr.atoms.worktreeChangesRefresh)).toEqual({ + cloudSessionId: 'agent-1', + revision: 2, + connectionVersion: 0, + }); + mockSessionCallbacks.onEvent?.(connected); + const firstConnected = store.get(mgr.atoms.worktreeChangesRefresh); + expect(firstConnected).toEqual({ + cloudSessionId: 'agent-1', + revision: 2, + connectionVersion: 1, + }); + mockSessionCallbacks.onEvent?.(connected); + expect(store.get(mgr.atoms.worktreeChangesRefresh)).toEqual({ + cloudSessionId: 'agent-1', + revision: 2, + connectionVersion: 2, + }); + expect(store.get(mgr.atoms.worktreeChangesRefresh)).not.toBe(firstConnected); + expect(listener).toHaveBeenCalledTimes(3); + expect(store.get(mgr.atoms.activity)).toBe(state.activity); + expect(store.get(mgr.atoms.agentStatus)).toBe(state.status); + expect(store.get(mgr.atoms.messagesList)).toBe(state.messages); + expect(store.get(mgr.atoms.pendingMessages)).toBe(state.pending); + unsubscribe(); + mgr.destroy(); + }); + + it('retains the highest ready revision and object when updates are coalesced', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + await mgr.switchSession(kiloId('ses-1')); + const listener = jest.fn(); + const unsubscribe = config.store.sub(mgr.atoms.worktreeChangesRefresh, listener); + try { + mockSessionCallbacks.onEvent?.(ready); + mockSessionCallbacks.onEvent?.({ ...ready, revision: 3 }); + const latestSignal = config.store.get(mgr.atoms.worktreeChangesRefresh); + for (const revision of [2, 3, 1]) { + mockSessionCallbacks.onEvent?.({ ...ready, revision }); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBe(latestSignal); + } + expect(latestSignal).toEqual({ + cloudSessionId: 'agent-1', + revision: 3, + connectionVersion: 0, + }); + expect(listener).toHaveBeenCalledTimes(2); + } finally { + unsubscribe(); + mgr.destroy(); + } + }); + + it.each([1, 2, 3])( + 'preserves the reconnect version when followed by ready revision %s', + async revision => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + await mgr.switchSession(kiloId('ses-1')); + try { + mockSessionCallbacks.onEvent?.(ready); + mockSessionCallbacks.onEvent?.(connected); + const connectedSignal = config.store.get(mgr.atoms.worktreeChangesRefresh); + mockSessionCallbacks.onEvent?.({ ...ready, revision }); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toEqual({ + cloudSessionId: 'agent-1', + revision: Math.max(2, revision), + connectionVersion: 1, + }); + if (revision <= 2) { + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBe(connectedSignal); + } else { + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).not.toBe(connectedSignal); + } + } finally { + mgr.destroy(); + } + } + ); + + it.each([ready, connected])('ignores mismatched Cloud session IDs for $type', async event => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + await mgr.switchSession(kiloId('ses-1')); + for (const cloudSessionId of ['agent-other', 'ses-1', '']) { + mockSessionCallbacks.onEvent?.({ ...event, cloudSessionId }); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + } + mgr.destroy(); + }); + + it('ignores signals without a current Cloud session or envelope identity', async () => { + const config = createMockConfig({ + fetchSession: jest.fn().mockResolvedValue({ + ...defaultFetchedSession, + cloudAgentSessionId: null, + }), + }); + const mgr = createSessionManager(config); + await mgr.switchSession(kiloId('ses-1')); + mockSessionCallbacks.onEvent?.(ready); + mockSessionCallbacks.onEvent?.(connected); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + config.store.set(mgr.atoms.sessionId, cloudAgentId('agent-1')); + mockSessionCallbacks.onEvent?.({ type: 'connected' }); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + mgr.destroy(); + }); + + it.each(['ses-1', 'ses-2'])( + 'resets and rejects stale callbacks when switching to %s', + async nextId => { + const config = createMockConfig({ + fetchSession: jest + .fn() + .mockResolvedValueOnce(defaultFetchedSession) + .mockResolvedValueOnce({ + ...defaultFetchedSession, + kiloSessionId: kiloId(nextId), + cloudAgentSessionId: cloudAgentId('agent-2'), + }), + }); + const mgr = createSessionManager(config); + await mgr.switchSession(kiloId('ses-1')); + const oldOnEvent = mockSessionCallbacks.onEvent; + oldOnEvent?.({ ...ready, revision: 10 }); + oldOnEvent?.(connected); + oldOnEvent?.(connected); + const switching = mgr.switchSession(kiloId(nextId)); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + oldOnEvent?.(connected); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + await switching; + for (const event of [ready, connected]) { + oldOnEvent?.({ ...event, cloudSessionId: 'agent-2' }); + mockSessionCallbacks.onEvent?.(event); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + } + mockSessionCallbacks.onEvent?.({ ...ready, cloudSessionId: 'agent-2' }); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toEqual({ + cloudSessionId: 'agent-2', + revision: 2, + connectionVersion: 0, + }); + mockSessionCallbacks.onEvent?.({ ...connected, cloudSessionId: 'agent-2' }); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toEqual({ + cloudSessionId: 'agent-2', + revision: 2, + connectionVersion: 1, + }); + mgr.destroy(); + } + ); + + it('resets on destroy and rejects retained callbacks', async () => { + const config = createMockConfig(); + const mgr = createSessionManager(config); + await mgr.switchSession(kiloId('ses-1')); + const oldOnEvent = mockSessionCallbacks.onEvent; + oldOnEvent?.(ready); + mgr.destroy(); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + oldOnEvent?.(ready); + oldOnEvent?.(connected); + expect(config.store.get(mgr.atoms.worktreeChangesRefresh)).toBeNull(); + }); + }); + describe('switchSession', () => { it('sets isLoading=true synchronously and clears it after completion', async () => { const config = createMockConfig(); diff --git a/packages/cloud-agent-sdk/src/session-manager.ts b/packages/cloud-agent-sdk/src/session-manager.ts index b473f5fed4..b7fd79fb0a 100644 --- a/packages/cloud-agent-sdk/src/session-manager.ts +++ b/packages/cloud-agent-sdk/src/session-manager.ts @@ -100,6 +100,11 @@ type SessionConfig = { }> | undefined; }; +type WorktreeChangesRefresh = { + cloudSessionId: string; + revision?: number; + connectionVersion: number; +}; type ActiveSessionType = ResolvedSession['type']; type ObservedModelSource = 'session' | 'message' | 'catalog'; type StandaloneQuestion = { requestId: string; questions: QuestionInfo[] }; @@ -369,6 +374,7 @@ type SessionManagerAtoms = { fetchedSessionData: W; /** Slash command catalog reported by the wrapper for the current session. */ availableCommands: W; + worktreeChangesRefresh: W; messagesList: Atom; staticMessages: Atom; dynamicMessages: Atom; @@ -752,6 +758,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { * DO) and on every wrapper push. Empty list = wrapper hasn't reported yet. */ const availableCommandsAtom = atom([]); + const worktreeChangesRefreshAtom = atom(null); const childSessionHydrationStatesAtom = atom>(new Map()); const childSessionErrorsAtom = atom>(new Map()); const hasOlderMessagesAtom = atom(false); @@ -973,6 +980,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { store.set(childSessionErrorsAtom, new Map()); store.set(chatUIAtom, { shouldAutoScroll: true }); store.set(availableCommandsAtom, []); + store.set(worktreeChangesRefreshAtom, null); store.set(hasOlderMessagesAtom, false); store.set(isLoadingOlderMessagesAtom, false); store.set(olderMessagesErrorAtom, null); @@ -1831,6 +1839,26 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { }, onEvent: event => { if (expectedGeneration !== switchGeneration) return; + if (event.type === 'worktree.changes.ready' || event.type === 'connected') { + const cloudSessionId = store.get(sessionIdAtom); + if (!cloudSessionId || event.cloudSessionId !== cloudSessionId) return; + const previous = store.get(worktreeChangesRefreshAtom); + if (event.type === 'worktree.changes.ready') { + if (event.revision <= (previous?.revision ?? 0)) return; + store.set(worktreeChangesRefreshAtom, { + cloudSessionId, + revision: event.revision, + connectionVersion: previous?.connectionVersion ?? 0, + }); + } else { + store.set(worktreeChangesRefreshAtom, { + ...previous, + cloudSessionId, + connectionVersion: (previous?.connectionVersion ?? 0) + 1, + }); + } + return; + } if (event.type === 'commands.available') { // Replace the catalog wholesale. The DO sends the full list on // every connect, so we never need to merge incrementally. @@ -2396,6 +2424,7 @@ function createSessionManager(config: SessionManagerConfig): SessionManager { billingFailure: billingFailureAtom, fetchedSessionData: fetchedSessionDataAtom, availableCommands: availableCommandsAtom, + worktreeChangesRefresh: worktreeChangesRefreshAtom, messagesList: messagesListAtom, staticMessages: staticMessagesAtom, dynamicMessages: dynamicMessagesAtom, @@ -2420,6 +2449,7 @@ export type { SessionManager, SessionManagerConfig, SessionManagerAtoms, + WorktreeChangesRefresh, SessionStatusIndicator, SessionConfig, StandalonePermission, diff --git a/packages/cloud-agent-sdk/src/session-phase.test.ts b/packages/cloud-agent-sdk/src/session-phase.test.ts index 3ba6829ccb..b2f64f53bc 100644 --- a/packages/cloud-agent-sdk/src/session-phase.test.ts +++ b/packages/cloud-agent-sdk/src/session-phase.test.ts @@ -83,7 +83,7 @@ function createSessionWithStateCapture( (_sessionId: string) => 'test-ticket' ) ) { - const { createEvent, kilocode, resetCounter } = createEventHelpers(); + const { createEvent, kilocode, resetCounter } = createEventHelpers(TEST_CLOUD_AGENT_ID); resetCounter(); const errors: string[] = []; @@ -376,7 +376,11 @@ describe('session state transitions', () => { }); it('session.created fires onSessionCreated', async () => { - const { createEvent: _createEvent, kilocode, resetCounter } = createEventHelpers(); + const { + createEvent: _createEvent, + kilocode, + resetCounter, + } = createEventHelpers(TEST_CLOUD_AGENT_ID); resetCounter(); const sessions: unknown[] = []; @@ -491,7 +495,7 @@ describe('authoritative message failure settlement', () => { suggestion: store.get(manager.atoms.suggestion), activeSuggestion: store.get(manager.atoms.activeSuggestion), }); - const events = createEventHelpers(); + const events = createEventHelpers(TEST_CLOUD_AGENT_ID); const requestInput = () => { for (const sessionID of [TEST_KILO_ID, 'child-session']) { sendRaw( diff --git a/packages/cloud-agent-sdk/src/session-transport.test.ts b/packages/cloud-agent-sdk/src/session-transport.test.ts index 164b383d7a..f1ebfb5f9c 100644 --- a/packages/cloud-agent-sdk/src/session-transport.test.ts +++ b/packages/cloud-agent-sdk/src/session-transport.test.ts @@ -1,3 +1,6 @@ +import { createStore } from 'jotai'; +import { assistantMsg, toolPart } from './__fixtures__/helpers'; +import { createSessionManager } from './session-manager'; import { createCloudAgentSession, REMOTE_SESSION_CREATION_NOT_SUPPORTED } from './session'; import type { CloudAgentSession } from './session'; import type { CloudAgentApi } from './transport'; @@ -154,6 +157,155 @@ function emitHeartbeatOwner( // Tests // --------------------------------------------------------------------------- +describe('Cloud Agent worktree refresh event pipeline', () => { + it.each([false, true])( + 'signals every idle reconnect independently of chat replay (cursor: %p)', + async withReplayCursor => { + jest.useFakeTimers(); + const store = createStore(); + const fetchSnapshot = jest.fn(() => Promise.resolve(makeSnapshot({ id: kiloSessionId }))); + const pageshow = { + handler: undefined as ((event: { persisted: boolean }) => void) | undefined, + }; + const manager = createSessionManager({ + store, + resolveSession: async () => ({ type: 'cloud-agent', kiloSessionId, cloudAgentSessionId }), + getTicket: () => 'ticket', + fetchSnapshot, + websocketBaseUrl: 'ws://localhost:9999', + userWebConnection: createUserWebConnection(), + api: createMockApi(), + prepare: jest.fn(), + initiate: jest.fn(), + fetchSession: async () => ({ + kiloSessionId, + cloudAgentSessionId, + title: null, + organizationId: null, + gitUrl: null, + gitBranch: null, + mode: null, + model: null, + variant: null, + repository: null, + isInitiated: true, + needsLegacyPrepare: false, + isPreparingAsync: false, + prompt: null, + initialMessageId: null, + associatedPr: null, + }), + lifecycleHooks: { + onPageshow: handler => { + pageshow.handler = handler; + return jest.fn(); + }, + }, + }); + const sendEnvelope = (streamEventType: string, data: unknown, eventId = 0) => { + mockWs.onmessage?.({ + data: JSON.stringify({ + eventId, + sessionId: cloudAgentSessionId, + streamEventType, + timestamp: '2026-09-02T00:00:00.000Z', + data, + }), + } as MessageEvent); + }; + const listener = jest.fn(); + const unsubscribe = store.sub(manager.atoms.worktreeChangesRefresh, listener); + try { + await manager.switchSession(kiloSessionId); + await jest.advanceTimersByTimeAsync(0); + expect(store.get(manager.atoms.worktreeChangesRefresh)).toBeNull(); + sendEnvelope('connected', { sessionStatus: { type: 'idle' } }); + expect(store.get(manager.atoms.worktreeChangesRefresh)).toEqual({ + cloudSessionId: cloudAgentSessionId, + connectionVersion: 1, + }); + expect(listener).toHaveBeenCalledTimes(1); + if (withReplayCursor) sendEnvelope('heartbeat', {}, 12); + fetchSnapshot.mockImplementation(() => new Promise(() => {})); + + for (let reconnect = 1; reconnect <= 2; reconnect++) { + const previousSignal = store.get(manager.atoms.worktreeChangesRefresh); + pageshow.handler?.({ persisted: true }); + await jest.advanceTimersByTimeAsync(0); + expect(store.get(manager.atoms.activity)).toEqual({ type: 'idle' }); + sendEnvelope('connected', { sessionStatus: { type: 'idle' } }); + expect(store.get(manager.atoms.activity)).toEqual({ type: 'idle' }); + expect(store.get(manager.atoms.worktreeChangesRefresh)).toEqual({ + cloudSessionId: cloudAgentSessionId, + connectionVersion: reconnect + 1, + }); + expect(store.get(manager.atoms.worktreeChangesRefresh)).not.toBe(previousSignal); + expect(listener).toHaveBeenCalledTimes(reconnect + 1); + expect(fetchSnapshot).toHaveBeenCalledTimes(withReplayCursor ? 1 : reconnect + 1); + const url = String(jest.mocked(global.WebSocket).mock.calls.at(-1)?.[0]); + expect(url).toContain(withReplayCursor ? 'fromId=12' : 'replay=false'); + } + + sendEnvelope('kilocode', { + type: 'message.updated', + properties: { info: assistantMsg('msg-assistant', 'msg-user', kiloSessionId) }, + }); + sendEnvelope('kilocode', { + type: 'message.part.updated', + properties: { part: toolPart('part-tool', 'msg-assistant', 'bash', kiloSessionId) }, + }); + sendEnvelope('cloud.message.sent', { messageId: 'msg-user' }); + sendEnvelope('kilocode', { + type: 'session.status', + properties: { sessionID: kiloSessionId, status: { type: 'busy' } }, + }); + expect(store.get(manager.atoms.messagesList)).toHaveLength(1); + expect(store.get(manager.atoms.activity)).toEqual({ type: 'busy' }); + const activity = store.get(manager.atoms.activity); + const status = store.get(manager.atoms.agentStatus); + const cloudStatus = store.get(manager.atoms.cloudStatus); + const messages = store.get(manager.atoms.messagesList); + const pending = store.get(manager.atoms.pendingMessages); + for (const revision of [2, 2, 1, 3]) { + const previousSignal = store.get(manager.atoms.worktreeChangesRefresh); + sendEnvelope('cloud.worktree.changes.ready', { revision }, 13); + expect(store.get(manager.atoms.worktreeChangesRefresh)).toEqual({ + cloudSessionId: cloudAgentSessionId, + revision: Math.max(2, revision), + connectionVersion: 3, + }); + if (revision <= (previousSignal?.revision ?? 0)) { + expect(store.get(manager.atoms.worktreeChangesRefresh)).toBe(previousSignal); + } else { + expect(store.get(manager.atoms.worktreeChangesRefresh)).not.toBe(previousSignal); + } + expect(store.get(manager.atoms.activity)).toBe(activity); + expect(store.get(manager.atoms.agentStatus)).toBe(status); + expect(store.get(manager.atoms.cloudStatus)).toBe(cloudStatus); + expect(store.get(manager.atoms.messagesList)).toBe(messages); + expect(store.get(manager.atoms.pendingMessages)).toBe(pending); + } + const oldOnMessage = mockWs.onmessage; + manager.destroy(); + oldOnMessage?.({ + data: JSON.stringify({ + eventId: 14, + sessionId: cloudAgentSessionId, + streamEventType: 'cloud.worktree.changes.ready', + timestamp: '2026-09-02T00:00:00.000Z', + data: { revision: 4 }, + }), + } as MessageEvent); + expect(store.get(manager.atoms.worktreeChangesRefresh)).toBeNull(); + } finally { + unsubscribe(); + manager.destroy(); + jest.useRealTimers(); + } + } + ); +}); + describe('session transport delegation (cloud agent)', () => { it('session.send() delegates to api.send with resolved cloudAgentSessionId', async () => { const api = createMockApi(); diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index 85a3b62c94..b68a702d27 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -84,7 +84,10 @@ import { worktreeChangesContext, type WorktreeChangesContext, } from './worktree-changes.js'; -import { WORKTREE_CHANGED_EVENT } from '../shared/worktree-changes-wire.js'; +import { + WORKTREE_CHANGED_EVENT, + WORKTREE_CHANGES_READY_EVENT, +} from '../shared/worktree-changes-wire.js'; import { createPreparationProgressRecorder } from '../session/preparation-progress.js'; import { finalizeOtherRunningAttemptsForMessage, @@ -236,7 +239,43 @@ export class SandboxSession extends DurableObject { const db = drizzle(ctx.storage, { logger: false }); this.eventQueries = createEventQueries(db, ctx.storage.sql); this.worktreeChanges = createWorktreeChanges({ - storage: ctx.storage, + storage: { + get: key => ctx.storage.get(key), + put: async (key, snapshot) => { + if (this.deletedWorktreeId || this.terminalLifecycle.isBlocked()) { + throw new Error('Worktree changes persistence is blocked'); + } + const sessionId = this.requireSessionId(); + const payload = JSON.stringify({ revision: snapshot.revision }); + const timestamp = Date.now(); + const id = ctx.storage.transactionSync(() => { + ctx.storage.kv.put(key, snapshot); + return this.eventQueries.insertUnique({ + executionId: '', + sessionId, + streamEventType: WORKTREE_CHANGES_READY_EVENT, + payload, + timestamp, + entityId: `worktree-changes/${snapshot.revision}`, + }); + }); + if (id === null) return; + try { + this.broadcastStoredEvent({ + id, + execution_id: '', + session_id: sessionId, + stream_event_type: WORKTREE_CHANGES_READY_EVENT, + payload, + timestamp, + }); + } catch { + logger + .withFields({ sessionId, eventId: id, revision: snapshot.revision }) + .error('Failed to broadcast saved worktree changes'); + } + }, + }, readContext: async () => this.worktreeContext(await this.getMetadata()), requestCapture: (context, payload) => withDORetry( diff --git a/services/cloud-agent-next/src/shared/protocol.ts b/services/cloud-agent-next/src/shared/protocol.ts index a6e1517203..aea40f9425 100644 --- a/services/cloud-agent-next/src/shared/protocol.ts +++ b/services/cloud-agent-next/src/shared/protocol.ts @@ -35,6 +35,7 @@ export type StreamEventType = | 'preparing' // Lazy workspace preparation step progress // DO -> /stream clients (cloud infrastructure lifecycle) | 'cloud.status' // Cloud infrastructure status (preparing/ready/finalizing/error) + | 'cloud.worktree.changes.ready' // DO -> /stream clients (session message queue) | 'cloud.message.queued' // User message accepted into the pending queue | 'cloud.message.sent' // Queued user message delivered to Kilo diff --git a/services/cloud-agent-next/src/shared/worktree-changes-wire.ts b/services/cloud-agent-next/src/shared/worktree-changes-wire.ts index 46922bdf46..045e6556aa 100644 --- a/services/cloud-agent-next/src/shared/worktree-changes-wire.ts +++ b/services/cloud-agent-next/src/shared/worktree-changes-wire.ts @@ -8,6 +8,7 @@ import { z } from 'zod'; export const MAX_WORKTREE_CHANGES_FILES = 1_000; export const MAX_WORKTREE_CHANGES_BYTES = 256 * 1024; export const WORKTREE_CHANGED_EVENT = 'session.worktree.changed'; +export const WORKTREE_CHANGES_READY_EVENT = 'cloud.worktree.changes.ready'; const revisionSchema = z.number().int().positive().max(Number.MAX_SAFE_INTEGER); const commitSchema = z.string().regex(/^(?:[0-9a-f]{40}|[0-9a-f]{64})$/); diff --git a/services/cloud-agent-next/src/websocket/stream.test.ts b/services/cloud-agent-next/src/websocket/stream.test.ts index ca78410d82..ec47dc695f 100644 --- a/services/cloud-agent-next/src/websocket/stream.test.ts +++ b/services/cloud-agent-next/src/websocket/stream.test.ts @@ -1,5 +1,7 @@ import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest'; import { createStreamHandler, formatStreamEvent } from './stream.js'; +import { parseStreamFilters } from './filters.js'; +import { WORKTREE_CHANGES_READY_EVENT } from '../shared/worktree-changes-wire.js'; import type { StoredEvent, StreamFilters } from './types.js'; import type { SessionId, EventId } from '../types/ids.js'; import type { EventQueries, EventQueryFilters } from '../session/queries/index.js'; @@ -248,6 +250,59 @@ describe('stream handler replayEvents', () => { }); }); +describe('worktree changes ready streaming', () => { + const readyEvents: StoredEvent[] = [5, 7].map((id, index) => ({ + ...makeEvent(id, JSON.stringify({ revision: index + 1 })), + execution_id: '', + stream_event_type: WORKTREE_CHANGES_READY_EVENT, + })); + + it.each([ + { query: '', replayIds: [5, 7], live: true }, + { + query: 'fromId=5&eventTypes=cloud.worktree.changes.ready', + replayIds: [7], + live: true, + }, + { query: 'eventTypes=kilocode', replayIds: [], live: false }, + { query: 'executionIds=exec_1', replayIds: [], live: false }, + { query: 'startTime=1007&endTime=1007', replayIds: [7], live: true }, + { query: 'endTime=1006', replayIds: [5], live: false }, + ])( + 'applies stream filters to replay and live delivery: $query', + async ({ query, replayIds, live }) => { + const state = makeFakeState(); + const handler = createStreamHandler(state, makeFakeEventQueries(readyEvents), SESSION_ID); + const filters = parseStreamFilters( + new URL(`https://example.com/stream?${query}`), + SESSION_ID + ); + const replaySocket = makeFakeWebSocket(); + await handler.replayEvents(replaySocket, filters); + const replay = parseSentMessages(replaySocket); + expect(replay.map(event => event.eventId)).toEqual(replayIds); + for (const event of replay) { + expect(event).toEqual({ + eventId: event.eventId, + sessionId: SESSION_ID, + streamEventType: 'cloud.worktree.changes.ready', + timestamp: new Date(1000 + event.eventId).toISOString(), + data: { revision: event.eventId === 5 ? 1 : 2 }, + }); + } + + const liveSocket = Object.assign(makeFakeWebSocket(), { + deserializeAttachment: () => ({ filters, connectedAt: 0 }), + }); + vi.spyOn(state, 'getWebSockets').mockReturnValue([liveSocket]); + handler.broadcastEvent(readyEvents[1]); + expect(parseSentMessages(liveSocket)).toEqual( + live ? [formatStreamEvent(readyEvents[1], SESSION_ID)] : [] + ); + } + ); +}); + describe('stream handler handleStreamRequest', () => { const OriginalResponse = Response; diff --git a/services/cloud-agent-next/src/websocket/types.ts b/services/cloud-agent-next/src/websocket/types.ts index 0201e0054d..47f91cc6d5 100644 --- a/services/cloud-agent-next/src/websocket/types.ts +++ b/services/cloud-agent-next/src/websocket/types.ts @@ -40,6 +40,7 @@ export type StreamEventType = | 'commands.available' // catalog of kilo slash commands | 'preparing' // lazy workspace preparation step progress | 'cloud.status' // cloud infrastructure status + | 'cloud.worktree.changes.ready' | 'cloud.message.queued' // user message accepted into pending queue | 'cloud.message.sent' // queued user message delivered to Kilo | 'cloud.message.completed' // accepted user message completed execution diff --git a/services/cloud-agent-next/test/integration/sandbox-control.test.ts b/services/cloud-agent-next/test/integration/sandbox-control.test.ts index e9eaea7446..4b032041ce 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -118,8 +118,12 @@ import { } from '../../src/shared/sandbox-control-protocol.js'; import { SandboxStatusSnapshotSchema } from '../../src/shared/sandbox-status.js'; import { WORKTREE_CHANGES_KEY } from '../../src/sandbox-session/worktree-changes.js'; -import { WORKTREE_CHANGED_EVENT } from '../../src/shared/worktree-changes-wire.js'; +import { + WORKTREE_CHANGED_EVENT, + WORKTREE_CHANGES_READY_EVENT, +} from '../../src/shared/worktree-changes-wire.js'; import { getWorktreeWorkspacePath } from '../../src/workspace.js'; +import type { StoredEvent } from '../../src/websocket/types.js'; vi.mock('../../src/session-access.js', () => ({ requireCurrentSessionAccess: vi.fn(), @@ -6702,6 +6706,12 @@ async function worktreeFixture() { const credential = generateSandboxCredential(); const controlTasks: Promise[] = []; const sessionTasks: Promise[] = []; + const readyNotifications: { + event: StoredEvent; + snapshot: unknown; + storedEvent: StoredEvent | null; + inTransaction: boolean; + }[] = []; await seedRunningCredential(credential, sandboxId); const { provider } = await installProvider(control, cloudflareRef(sandboxId)); await runInDurableObject(control, async (instance, state) => { @@ -6712,7 +6722,7 @@ async function worktreeFixture() { waitUntil(promise); }); }); - await runInDurableObject(session, async (instance, state) => { + const restoreObservation = await runInDurableObject(session, async (instance, state) => { await instance.registerSession({ identity: { sessionId, userId, createdOnPlatform: 'cloud-agent-web' }, auth: { kiloSessionId, kilocodeToken: KILO_TOKEN }, @@ -6735,6 +6745,36 @@ async function worktreeFixture() { sessionTasks.push(promise); waitUntil(promise); }); + let inTransaction = false; + const transactionSync = state.storage.transactionSync.bind(state.storage); + const transactionSpy = vi + .spyOn(state.storage, 'transactionSync') + .mockImplementation(callback => { + inTransaction = true; + try { + return transactionSync(callback); + } finally { + inTransaction = false; + } + }); + const broadcast = instance['broadcastStoredEvent'].bind(instance); + instance['broadcastStoredEvent'] = event => { + if (event.stream_event_type === WORKTREE_CHANGES_READY_EVENT) { + readyNotifications.push({ + event, + snapshot: state.storage.kv.get(WORKTREE_CHANGES_KEY), + storedEvent: instance['eventQueries'].findByEntityId( + `worktree-changes/${JSON.parse(event.payload).revision}` + ), + inTransaction, + }); + } + broadcast(event); + }; + return () => { + transactionSpy.mockRestore(); + instance['broadcastStoredEvent'] = broadcast; + }; }); await control.prepareSessionCredentials({ ownerId: userId, sessionId }); await control.attachSession({ sessionId, kiloSessionId, directory, worktreeId, ownerId: userId }); @@ -6824,6 +6864,7 @@ async function worktreeFixture() { session, provider, captures, + readyNotifications, prompts, aborts, noWake, @@ -6904,7 +6945,11 @@ async function worktreeFixture() { }, ready, close(): void { - ws.close(); + try { + ws.close(); + } finally { + restoreObservation(); + } }, }; } @@ -6920,6 +6965,152 @@ describe('SandboxSession worktree changes persistence', () => { afterEach(() => vi.restoreAllMocks()); + it('broadcasts committed snapshots with distinct cursor IDs and idempotent revision events', async () => { + const fixture = await worktreeFixture(); + try { + for (const revision of [1, 2]) { + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + expect(captureRevision(request)).toBe(revision); + expect(fixture.readyNotifications).toHaveLength(revision - 1); + fixture.reply(request, worktreeCapture(revision)); + const saved = await pending; + expect(saved.status).toBe('refreshed'); + expect(fixture.readyNotifications).toHaveLength(revision); + const notification = fixture.readyNotifications[revision - 1]; + expect(notification).toEqual({ + event: { + id: expect.any(Number), + execution_id: '', + session_id: fixture.sessionId, + stream_event_type: WORKTREE_CHANGES_READY_EVENT, + payload: JSON.stringify({ revision }), + timestamp: expect.any(Number), + }, + snapshot: saved.snapshot, + storedEvent: notification.event, + inTransaction: false, + }); + } + const [first, second] = fixture.readyNotifications.map(notification => notification.event); + expect(first.id).toBeGreaterThan(0); + expect(second.id).toBeGreaterThan(first.id); + await runInDurableObject(fixture.session, instance => { + const queries = instance['eventQueries']; + expect(queries.findByFilters({ fromId: first.id })).toEqual([second]); + expect( + queries.insertUnique({ + executionId: '', + sessionId: fixture.sessionId, + streamEventType: WORKTREE_CHANGES_READY_EVENT, + payload: second.payload, + timestamp: second.timestamp, + entityId: 'worktree-changes/2', + }) + ).toBeNull(); + expect(queries.findByFilters({})).toEqual([first, second]); + }); + } finally { + fixture.close(); + } + }); + + it('rolls back the snapshot and ready event when event insertion fails', async () => { + const fixture = await worktreeFixture(); + let restoreInsert: (() => void) | undefined; + try { + restoreInsert = await runInDurableObject(fixture.session, (instance, state) => { + state.storage.kv.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot); + const queries = instance['eventQueries']; + const insert = queries.insertUnique.bind(queries); + const insertSpy = vi.spyOn(queries, 'insertUnique').mockImplementationOnce(params => { + insert(params); + throw new Error('fixture insert failure'); + }); + return () => insertSpy.mockRestore(); + }); + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(captureRevision(request))); + await expect(pending).resolves.toEqual({ status: 'failed', snapshot: savedWorktreeSnapshot }); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ + snapshot: savedWorktreeSnapshot, + }); + await runInDurableObject(fixture.session, instance => { + expect(instance['eventQueries'].findByFilters({})).toEqual([]); + }); + expect(fixture.readyNotifications).toEqual([]); + } finally { + restoreInsert?.(); + fixture.close(); + } + }); + + it('preserves a successful saved result and replay event when broadcast fails', async () => { + const fixture = await worktreeFixture(); + try { + await runInDurableObject(fixture.session, (instance, state) => { + state.storage.kv.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot); + instance['broadcastStoredEvent'] = () => { + throw new Error('fixture broadcast failure'); + }; + }); + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(captureRevision(request), true)); + const saved = await pending; + expect(saved).toMatchObject({ status: 'refreshed', snapshot: { revision: 5, files: [] } }); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ + snapshot: saved.snapshot, + }); + await runInDurableObject(fixture.session, instance => { + expect(instance['eventQueries'].findByFilters({ fromId: 0 })).toEqual([ + { + id: expect.any(Number), + execution_id: '', + session_id: fixture.sessionId, + stream_event_type: WORKTREE_CHANGES_READY_EVENT, + payload: JSON.stringify({ revision: 5 }), + timestamp: expect.any(Number), + }, + ]); + }); + } finally { + fixture.close(); + } + }); + + it('does not rebroadcast an existing revision event', async () => { + const fixture = await worktreeFixture(); + try { + const existingId = await runInDurableObject(fixture.session, instance => + instance['eventQueries'].insertUnique({ + executionId: '', + sessionId: fixture.sessionId, + streamEventType: WORKTREE_CHANGES_READY_EVENT, + payload: JSON.stringify({ revision: 1 }), + timestamp: Date.now(), + entityId: 'worktree-changes/1', + }) + ); + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + fixture.reply(request, worktreeCapture(captureRevision(request))); + await expect(pending).resolves.toMatchObject({ + status: 'refreshed', + snapshot: { revision: 1 }, + }); + expect(fixture.readyNotifications).toEqual([]); + await runInDurableObject(fixture.session, instance => { + expect(instance['eventQueries'].findByFilters({})).toEqual([ + expect.objectContaining({ id: existingId, payload: JSON.stringify({ revision: 1 }) }), + ]); + }); + } finally { + fixture.close(); + } + }); + it('captures after attach without a UI request and does not delay prompt delivery', async () => { const fixture = await worktreeFixture(); await fixture.session.admitSubmittedMessage({ @@ -6985,6 +7176,7 @@ describe('SandboxSession worktree changes persistence', () => { const dirty = await fixture.nextCapture(); for (let hint = 0; hint < 3; hint++) await fixture.event(WORKTREE_CHANGED_EVENT); expect(fixture.captures).toHaveLength(2); + expect(before.broadcast).not.toHaveBeenCalled(); fixture.reply(dirty, worktreeCapture(captureRevision(dirty))); const trailing = await fixture.nextCapture(); expect(captureRevision(trailing)).toBe(captureRevision(dirty) + 1); @@ -7005,19 +7197,29 @@ describe('SandboxSession worktree changes persistence', () => { expect(state.storage.kv.get('session_messages')).toEqual(before.messages); expect( createEventQueries(drizzle(state.storage), state.storage.sql).findByFilters({}) - ).toEqual(before.events); + ).toEqual([ + ...before.events, + ...fixture.readyNotifications.slice(1).map(notification => notification.event), + ]); expect(await state.storage.getAlarm()).toEqual(before.alarm); }); await runInDurableObject(fixture.control, async (_instance, state) => { expect(await state.storage.list()).toEqual(controlBefore.records); expect(await state.storage.getAlarm()).toEqual(controlBefore.alarm); }); - expect(before.broadcast).not.toHaveBeenCalled(); + expect(before.broadcast.mock.calls.map(([event]) => event)).toEqual( + fixture.readyNotifications.slice(1).map(notification => notification.event) + ); + expect(fixture.readyNotifications.map(({ event }) => JSON.parse(event.payload))).toEqual([ + { revision: 1 }, + { revision: 2 }, + { revision: 3 }, + ]); expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); await fixture.event('session.status', fixture.kiloSessionId, { status: { type: 'busy' } }); - expect(before.broadcast).toHaveBeenCalledTimes(1); + expect(before.broadcast).toHaveBeenCalledTimes(3); } finally { fixture.close(); } @@ -7072,12 +7274,11 @@ describe('SandboxSession worktree changes persistence', () => { ).resolves.toEqual({ applied: false }); expect(state.storage.kv.get('session_messages')).toEqual(messages); expect( - createEventQueries(drizzle(state.storage), state.storage.sql).findByFilters({ - eventTypes: ['kilocode'], - }) + createEventQueries(drizzle(state.storage), state.storage.sql).findByFilters({}) ).toEqual([]); }); await fixture.settled(); + expect(fixture.readyNotifications).toEqual([]); expect(fixture.captures).toHaveLength(0); expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); @@ -7523,6 +7724,7 @@ describe('SandboxSession worktree changes persistence', () => { await expect(fixture.session.getMetadata()).resolves.toBeNull(); expect(fixture.captures).toHaveLength(2); expect(fixture.aborts).toHaveLength(0); + expect(fixture.readyNotifications).toEqual([]); fixture.close(); } ); @@ -7542,6 +7744,10 @@ describe('SandboxSession worktree changes persistence', () => { await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ snapshot: savedWorktreeSnapshot, }); + expect(fixture.readyNotifications).toEqual([]); + await runInDurableObject(fixture.session, instance => { + expect(instance['eventQueries'].findByFilters({})).toEqual([]); + }); } const pending = fixture.session.refreshWorktreeChanges(); const request = await fixture.nextCapture(); @@ -7550,6 +7756,9 @@ describe('SandboxSession worktree changes persistence', () => { status: 'refreshed', snapshot: { revision: 8, files: [] }, }); + expect(fixture.readyNotifications.map(({ event }) => JSON.parse(event.payload))).toEqual([ + { revision: 8 }, + ]); fixture.close(); }); @@ -7571,6 +7780,7 @@ describe('SandboxSession worktree changes persistence', () => { snapshot: savedWorktreeSnapshot, }); expect(fixture.captures).toHaveLength(1); + expect(fixture.readyNotifications).toEqual([]); await fixture.ready(); const next = fixture.session.refreshWorktreeChanges(); const request = await fixture.nextCapture(); @@ -7581,57 +7791,69 @@ describe('SandboxSession worktree changes persistence', () => { it('requires a running physical sandbox and matching route at the request boundary', async () => { const fixture = await worktreeFixture(); - for (const identity of [ - { - sessionId: 'workspace_other', - kiloSessionId: fixture.kiloSessionId, - directory: fixture.directory, - }, - { sessionId: fixture.sessionId, kiloSessionId: 'other_root', directory: fixture.directory }, - { sessionId: fixture.sessionId, kiloSessionId: fixture.kiloSessionId, directory: '/other' }, - ]) { - await expect( - fixture.control.request({ - operation: 'session.git.summary', - session: identity, - payload: { revision: 1 }, - }) - ).resolves.toMatchObject({ ok: false, error: { code: 'not_ready' } }); + try { + for (const identity of [ + { + sessionId: 'workspace_other', + kiloSessionId: fixture.kiloSessionId, + directory: fixture.directory, + }, + { sessionId: fixture.sessionId, kiloSessionId: 'other_root', directory: fixture.directory }, + { sessionId: fixture.sessionId, kiloSessionId: fixture.kiloSessionId, directory: '/other' }, + ]) { + await expect( + fixture.control.request({ + operation: 'session.git.summary', + session: identity, + payload: { revision: 1 }, + }) + ).resolves.toMatchObject({ ok: false, error: { code: 'not_ready' } }); + } + await fixture.control.beginStop('test'); + await expect(fixture.session.refreshWorktreeChanges()).resolves.toEqual({ + status: 'offline', + snapshot: null, + }); + expect(fixture.captures).toHaveLength(0); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + await fixture.control.confirmStopped(); + await fixture.settled(); + } finally { + fixture.close(); } - await fixture.control.beginStop('test'); - await expect(fixture.session.refreshWorktreeChanges()).resolves.toEqual({ - status: 'offline', - snapshot: null, - }); - expect(fixture.captures).toHaveLength(0); - expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); - expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); - await fixture.control.confirmStopped(); - fixture.close(); }); it.each(['physical stop', 'route detach'] as const)( 'discards a valid in-flight capture after %s and preserves the saved snapshot', async change => { const fixture = await worktreeFixture(); - await runInDurableObject(fixture.session, async (_instance, state) => - state.storage.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot) - ); - const pending = fixture.session.refreshWorktreeChanges(); - const request = await fixture.nextCapture(); - if (change === 'physical stop') await fixture.control.beginStop('test in-flight capture'); - else await fixture.control.detachSession(fixture.sessionId); - fixture.reply(request, worktreeCapture(captureRevision(request), true)); - await expect(pending).resolves.toEqual({ status: 'failed', snapshot: savedWorktreeSnapshot }); - await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ - snapshot: savedWorktreeSnapshot, - }); - expect(fixture.captures).toHaveLength(1); - expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); - expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); - expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); - if (change === 'physical stop') await fixture.control.confirmStopped(); - fixture.close(); + try { + await runInDurableObject(fixture.session, async (_instance, state) => + state.storage.put(WORKTREE_CHANGES_KEY, savedWorktreeSnapshot) + ); + const pending = fixture.session.refreshWorktreeChanges(); + const request = await fixture.nextCapture(); + if (change === 'physical stop') await fixture.control.beginStop('test in-flight capture'); + else await fixture.control.detachSession(fixture.sessionId); + fixture.reply(request, worktreeCapture(captureRevision(request), true)); + await expect(pending).resolves.toEqual({ + status: 'failed', + snapshot: savedWorktreeSnapshot, + }); + await expect(fixture.session.getWorktreeChanges()).resolves.toEqual({ + snapshot: savedWorktreeSnapshot, + }); + expect(fixture.captures).toHaveLength(1); + expect(fixture.noWake.ensureReady).not.toHaveBeenCalled(); + expect(fixture.noWake.attachSession).not.toHaveBeenCalled(); + expect(fixture.noWake.claimCreate).not.toHaveBeenCalled(); + expect(fixture.readyNotifications).toEqual([]); + if (change === 'physical stop') await fixture.control.confirmStopped(); + await fixture.settled(); + } finally { + fixture.close(); + } } ); @@ -7649,6 +7871,7 @@ describe('SandboxSession worktree changes persistence', () => { await runInDurableObject(fixture.session, instance => { previousInstance = instance; }); + fixture.close(); await abortAllDurableObjects(); const freshSession = env.SANDBOX_SESSION.getByName(`${fixture.userId}:${fixture.sessionId}`); const freshControl = env.SANDBOX_CONTROL.getByName(fixture.sandboxId);