diff --git a/apps/mobile/src/components/agents/mobile-session-page-adapter.test.ts b/apps/mobile/src/components/agents/mobile-session-page-adapter.test.ts index 2241f569403..d688429c285 100644 --- a/apps/mobile/src/components/agents/mobile-session-page-adapter.test.ts +++ b/apps/mobile/src/components/agents/mobile-session-page-adapter.test.ts @@ -93,6 +93,7 @@ describe('fetchMobileSessionSnapshotPage', () => { mocks.getSessionMessagesPageQuery.mockResolvedValueOnce({ kiloSessionId: 'ses_123', history: page, + watermarkEventId: 42, }); const { fetchMobileSessionSnapshotPage } = await importAdapter(); @@ -104,6 +105,7 @@ describe('fetchMobileSessionSnapshotPage', () => { messages: page.messages, nextCursor: 'opaque-cursor', omittedItemCount: 2, + watermarkEventId: 42, }); // No `cursor` was supplied, so the tRPC layer must see only `session_id` // and let its input schema default the bounded page size to 50. @@ -239,4 +241,30 @@ describe('fetchMobileSessionSnapshotPage', () => { expect(call).toBeDefined(); expect(call.limit).toBeUndefined(); }); + + it('carries watermarkEventId through when the web router includes it', async () => { + mocks.getSessionMessagesPageQuery.mockResolvedValueOnce({ + kiloSessionId: 'ses_123', + history: historyPage({ nextCursor: null }), + watermarkEventId: 77, + }); + + const { fetchMobileSessionSnapshotPage } = await importAdapter(); + const result = await fetchMobileSessionSnapshotPage(kiloSessionId('ses_123'), {}); + + expect(result).toMatchObject({ kind: 'success', watermarkEventId: 77 }); + }); + + it('omits watermarkEventId from the page when the web router does not include it', async () => { + mocks.getSessionMessagesPageQuery.mockResolvedValueOnce({ + kiloSessionId: 'ses_123', + history: historyPage({ nextCursor: null }), + }); + + const { fetchMobileSessionSnapshotPage } = await importAdapter(); + const result = await fetchMobileSessionSnapshotPage(kiloSessionId('ses_123'), {}); + + expect(result).toMatchObject({ kind: 'success' }); + expect((result as Record).watermarkEventId).toBeUndefined(); + }); }); diff --git a/apps/mobile/src/components/agents/mobile-session-page-adapter.ts b/apps/mobile/src/components/agents/mobile-session-page-adapter.ts index 002ba646c87..21221728944 100644 --- a/apps/mobile/src/components/agents/mobile-session-page-adapter.ts +++ b/apps/mobile/src/components/agents/mobile-session-page-adapter.ts @@ -20,7 +20,8 @@ function isHistoryPage(history: KiloSdkMessageHistory): history is KiloSdkMessag * * Maps the shared `KiloSdkMessageHistory` union to the SDK's * `SessionSnapshotPageOutcome`: - * - Page variant → success outcome with messages, cursor, and omitted count + * - Page variant → success outcome with messages, cursor, omitted count, + * and the Cloud Agent event-log watermark * - Failure variants → passed through verbatim (retryable, too_large, invalid_data) * - Null history → empty success page (existing session with zero messages). A genuine * NOT_FOUND is returned as a thrown/rejected TRPCError, never as `history === null`. @@ -45,6 +46,7 @@ export async function fetchMobileSessionSnapshotPage( messages: [], nextCursor: null, omittedItemCount: 0, + watermarkEventId: result.watermarkEventId ?? undefined, }; } @@ -55,6 +57,7 @@ export async function fetchMobileSessionSnapshotPage( messages: history.messages as SessionSnapshotPage['messages'], nextCursor: history.nextCursor, omittedItemCount: history.omittedItemCount, + watermarkEventId: result.watermarkEventId ?? undefined, }; } diff --git a/apps/mobile/src/components/kilo-chat/conversation-list-screen.tsx b/apps/mobile/src/components/kilo-chat/conversation-list-screen.tsx index ad16395a0fa..7855c4ac111 100644 --- a/apps/mobile/src/components/kilo-chat/conversation-list-screen.tsx +++ b/apps/mobile/src/components/kilo-chat/conversation-list-screen.tsx @@ -38,6 +38,7 @@ import { } from './hooks/use-conversations'; import { useInstancePresence } from './hooks/use-instance-presence'; import { useNowTicker } from './hooks/use-now-ticker'; +import { useAppActiveAndFocused } from './hooks/use-app-active-and-focused'; type Props = { sandboxId: string; @@ -105,6 +106,7 @@ export function ConversationListScreen({ sandboxId, sandboxLabel }: Props) { const { fontScale } = useWindowDimensions(); const client = useKiloChatClient(); const eventClient = useEventServiceClient(); + const activeAndFocused = useAppActiveAndFocused(); const listQuery = useConversations(client, sandboxId); const createConversation = useCreateConversation(client); const leaveConversation = useLeaveConversation(client); @@ -137,7 +139,7 @@ export function ConversationListScreen({ sandboxId, sandboxLabel }: Props) { ); useInstancePresence(sandboxId); - useBotStatus(client, eventClient, sandboxId); + useBotStatus(client, eventClient, sandboxId, activeAndFocused); function handleRowPress(conversationId: string) { void Haptics.selectionAsync(); diff --git a/apps/mobile/src/components/kilo-chat/conversation-screen.tsx b/apps/mobile/src/components/kilo-chat/conversation-screen.tsx index 152d6158c85..b9392183d23 100644 --- a/apps/mobile/src/components/kilo-chat/conversation-screen.tsx +++ b/apps/mobile/src/components/kilo-chat/conversation-screen.tsx @@ -29,6 +29,7 @@ import { useMessageCacheUpdater, useMessages } from './hooks/use-messages'; import { useNowTicker } from './hooks/use-now-ticker'; import { useCurrentUserId } from './hooks/use-current-user-id'; import { useKiloChatTokenError } from './kilo-chat-provider'; +import { useAppActiveAndFocused } from './hooks/use-app-active-and-focused'; import { instanceOrgId, useAllKiloClawInstances, @@ -56,6 +57,7 @@ export function ConversationScreen({ }: Props) { const client = useKiloChatClient(); const eventClient = useEventServiceClient(); + const activeAndFocused = useAppActiveAndFocused(); const router = useRouter(); const currentUserId = useCurrentUserId(); const tokenError = useKiloChatTokenError(); @@ -67,7 +69,7 @@ export function ConversationScreen({ const { data: instances } = useAllKiloClawInstances(); const currentInstance = instances?.find(instance => instance.sandboxId === sandboxId); const instanceStatus = instanceStatusQuery.data?.status ?? currentInstance?.status ?? null; - const botStatus = useBotStatus(client, eventClient, sandboxId); + const botStatus = useBotStatus(client, eventClient, sandboxId, activeAndFocused); const botPresence = botStatus ? { online: botStatus.online, lastAt: botStatus.at } : undefined; const hasAttachmentsCapability = botStatus?.capabilities?.includes('attachments') ?? false; const now = useNowTicker(10_000); diff --git a/apps/mobile/src/components/kilo-chat/hooks/use-conversation-event-subscription.ts b/apps/mobile/src/components/kilo-chat/hooks/use-conversation-event-subscription.ts index 65a8e2248d3..036643730f2 100644 --- a/apps/mobile/src/components/kilo-chat/hooks/use-conversation-event-subscription.ts +++ b/apps/mobile/src/components/kilo-chat/hooks/use-conversation-event-subscription.ts @@ -29,7 +29,7 @@ export function useConversationEventSubscription( if (!conversationId) { return undefined; } - return eventService.onReconnect(() => { + return eventService.onResync(() => { void queryClient.invalidateQueries({ queryKey: messagesKey(conversationId) }); }); }, [eventService, queryClient, conversationId]); diff --git a/apps/mobile/src/components/kilo-chat/kilo-chat-provider.mounted.test.tsx b/apps/mobile/src/components/kilo-chat/kilo-chat-provider.mounted.test.tsx new file mode 100644 index 00000000000..d97e6f35193 --- /dev/null +++ b/apps/mobile/src/components/kilo-chat/kilo-chat-provider.mounted.test.tsx @@ -0,0 +1,166 @@ +/* eslint-disable typescript-eslint/no-deprecated -- react-test-renderer is the DOM-free renderer used to mount React trees under vitest (same pattern as src/test/render-with-providers.tsx) */ +/* eslint-disable eslint/max-classes-per-file -- hoisted mock classes for two different packages */ +import { createElement, type ReactElement, type ReactNode } from 'react'; +import TestRenderer, { act } from 'react-test-renderer'; +import { describe, expect, it, vi } from 'vitest'; + +import { KiloChatProvider } from './kilo-chat-provider'; + +// ── Hoisted mocks ────────────────────────────────────────────────────────── + +const mocks = vi.hoisted(() => { + const acquire = vi.fn<() => Promise>().mockResolvedValue(undefined); + const release = vi.fn(); + + // Must be a real class (not vi.fn()) because it is instantiated with `new`. + class MockEventServiceClient { + acquire = acquire; + release = release; + } + + // eslint-disable-next-line typescript-eslint/no-extraneous-class -- hoisted mock stub for constructor + class MockKiloChatClient {} + + const subscribeToKiloChatTokenResponses = vi.fn(() => vi.fn()); + const clearKiloChatTokenCache = vi.fn(); + const useKiloChatTokenGetter = vi.fn(() => vi.fn().mockResolvedValue('mock-token')); + const useKiloChatTokenResponseGetter = vi.fn(() => + vi.fn().mockResolvedValue({ userId: 'test-user', token: 'mock-token' }) + ); + + const useAppActiveAndFocused = vi.fn<() => boolean>(); + + return { + acquire, + release, + EventServiceClient: MockEventServiceClient, + KiloChatClient: MockKiloChatClient, + subscribeToKiloChatTokenResponses, + clearKiloChatTokenCache, + useKiloChatTokenGetter, + useKiloChatTokenResponseGetter, + useAppActiveAndFocused, + }; +}); + +// ── Module mocks ─────────────────────────────────────────────────────────── + +vi.mock('@kilocode/event-service', () => ({ + EventServiceClient: mocks.EventServiceClient, +})); + +vi.mock('@kilocode/kilo-chat', () => ({ + KiloChatClient: mocks.KiloChatClient, +})); + +vi.mock('@kilocode/kilo-chat-hooks', () => ({ + KiloChatHooksProvider: ({ children }: { children: ReactNode; value: unknown }) => + children as ReactElement, +})); + +vi.mock('@/lib/config', () => ({ + EVENT_SERVICE_URL: 'https://events.test', + KILO_CHAT_URL: 'https://chat.test', +})); + +vi.mock('./hooks/use-app-active-and-focused', () => ({ + useAppActiveAndFocused: mocks.useAppActiveAndFocused, +})); + +vi.mock('./hooks/use-kilo-chat-token', () => ({ + subscribeToKiloChatTokenResponses: mocks.subscribeToKiloChatTokenResponses, + clearKiloChatTokenCache: mocks.clearKiloChatTokenCache, + useKiloChatTokenGetter: mocks.useKiloChatTokenGetter, + useKiloChatTokenResponseGetter: mocks.useKiloChatTokenResponseGetter, +})); + +// ── Flat render helper ───────────────────────────────────────────────────── + +async function mountProvider(): Promise { + const ref: { current: TestRenderer.ReactTestRenderer | undefined } = { current: undefined }; + await act(async () => { + ref.current = TestRenderer.create(createElement(KiloChatProvider, null)); + await Promise.resolve(); + }); + const created = ref.current; + if (!created) { + throw new Error('mountProvider: renderer was not created'); + } + return created; +} + +async function unmountProvider(renderer: TestRenderer.ReactTestRenderer): Promise { + await act(async () => { + renderer.unmount(); + await Promise.resolve(); + }); +} + +// ── Tests ────────────────────────────────────────────────────────────────── + +describe('KiloChatProvider lifecycle', () => { + it('mounts active, calls acquire, unmounts and releases the hold', async () => { + mocks.useAppActiveAndFocused.mockReturnValue(true); + vi.clearAllMocks(); + + const renderer = await mountProvider(); + + expect(mocks.acquire).toHaveBeenCalledTimes(1); + expect(mocks.release).not.toHaveBeenCalled(); + + await unmountProvider(renderer); + + // Active mount → cleanup must release the hold. + expect(mocks.release).toHaveBeenCalledTimes(1); + }); + + it('mounts inactive, calls neither acquire nor release', async () => { + mocks.useAppActiveAndFocused.mockReturnValue(false); + vi.clearAllMocks(); + + const renderer = await mountProvider(); + + // Inactive: acquires nothing, releases nothing — no unpaired release. + expect(mocks.acquire).not.toHaveBeenCalled(); + expect(mocks.release).not.toHaveBeenCalled(); + + await unmountProvider(renderer); + + // Cleanup: holding was false, so no release on unmount either. + expect(mocks.release).not.toHaveBeenCalled(); + }); + + it('handles active → inactive → active transition without unmounting', async () => { + mocks.useAppActiveAndFocused.mockReturnValue(true); + vi.clearAllMocks(); + + const renderer = await mountProvider(); + expect(mocks.acquire).toHaveBeenCalledTimes(1); + expect(mocks.release).not.toHaveBeenCalled(); + + // Transition to inactive — cleanup from previous active effect releases the hold. + mocks.useAppActiveAndFocused.mockReturnValue(false); + vi.clearAllMocks(); + await act(async () => { + renderer.update(createElement(KiloChatProvider, null)); + await Promise.resolve(); + }); + expect(mocks.release).toHaveBeenCalledTimes(1); + expect(mocks.acquire).not.toHaveBeenCalled(); + + // Transition back to active — reacquires. + mocks.useAppActiveAndFocused.mockReturnValue(true); + vi.clearAllMocks(); + await act(async () => { + renderer.update(createElement(KiloChatProvider, null)); + await Promise.resolve(); + }); + expect(mocks.acquire).toHaveBeenCalledTimes(1); + expect(mocks.release).not.toHaveBeenCalled(); + + // Unmount releases the reacquired hold. + vi.clearAllMocks(); + await unmountProvider(renderer); + expect(mocks.release).toHaveBeenCalledTimes(1); + }); +}); diff --git a/apps/mobile/src/components/kilo-chat/kilo-chat-provider.tsx b/apps/mobile/src/components/kilo-chat/kilo-chat-provider.tsx index 1fd0d32b7b3..21ce21d3e24 100644 --- a/apps/mobile/src/components/kilo-chat/kilo-chat-provider.tsx +++ b/apps/mobile/src/components/kilo-chat/kilo-chat-provider.tsx @@ -6,6 +6,7 @@ import { KiloChatHooksProvider } from '@kilocode/kilo-chat-hooks'; import { EVENT_SERVICE_URL, KILO_CHAT_URL } from '@/lib/config'; +import { useAppActiveAndFocused } from './hooks/use-app-active-and-focused'; import { clearKiloChatTokenCache, subscribeToKiloChatTokenResponses, @@ -41,6 +42,7 @@ export function KiloChatProvider({ children }: KiloChatProviderProps) { const [currentUserId, setCurrentUserId] = useState(null); const [tokenError, setTokenError] = useState(false); const [retryCount, setRetryCount] = useState(0); + const activeAndFocused = useAppActiveAndFocused(); const [value] = useState(() => { const eventService = new EventServiceClient({ @@ -64,11 +66,17 @@ export function KiloChatProvider({ children }: KiloChatProviderProps) { }); useEffect(() => { - void value.eventService.connect(); + let holding = false; + if (activeAndFocused) { + void value.eventService.acquire(); + holding = true; + } return () => { - value.eventService.disconnect(); + if (holding) { + value.eventService.release(); + } }; - }, [value]); + }, [value, activeAndFocused]); useEffect(() => { let cancelled = false; diff --git a/apps/mobile/src/lib/trpc.test.ts b/apps/mobile/src/lib/trpc.test.ts index 1334ec98f28..ceb874ce252 100644 --- a/apps/mobile/src/lib/trpc.test.ts +++ b/apps/mobile/src/lib/trpc.test.ts @@ -1,4 +1,4 @@ -import { afterEach, describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; const httpLinkMock = vi.hoisted(() => vi.fn()); const httpBatchLinkMock = vi.hoisted(() => vi.fn()); @@ -78,3 +78,119 @@ describe('tRPC client link options', () => { expect(httpBatchLinkOpts).toHaveProperty('methodOverride', 'POST'); }); }); + +const mockFetch = vi.fn(); + +// Suppress Node.js 24 unhandledRejection from AbortController.abort() +// with a non-DOMException reason during fake-timer tests. +// eslint-disable-next-line @typescript-eslint/no-empty-function +function swallowUnhandledRejection(): void {} + +describe('deadlineFetch', () => { + beforeEach(() => { + vi.useRealTimers(); + mockFetch.mockReset(); + vi.stubGlobal('fetch', mockFetch); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + vi.resetModules(); + }); + + it('returns the response when fetch completes before the deadline', async () => { + mockFetch.mockResolvedValue(new Response('ok', { status: 200 })); + const { deadlineFetch } = await import('./trpc'); + const response = await deadlineFetch('https://api.example.com/api/trpc'); + expect(response.status).toBe(200); + expect(await response.text()).toBe('ok'); + expect(mockFetch).toHaveBeenCalledTimes(1); + }); + + it('rejects with RequestDeadlineError when the deadline expires', async () => { + vi.useFakeTimers(); + process.on('unhandledRejection', swallowUnhandledRejection); + try { + // Fetch never resolves unless its signal aborts — simulates a hanging + // backend where the only thing that cancels the request is the deadline. + mockFetch.mockImplementation( + // eslint-disable-next-line typescript-eslint/promise-function-async + (_url, init?: RequestInit) => + new Promise((resolve, reject) => { + const signal = init?.signal; + if (signal?.aborted) { + reject(signal.reason as Error); + return; + } + const id = setTimeout(resolve, 60_000, new Response('never')); + signal?.addEventListener('abort', () => { + clearTimeout(id); + reject(signal.reason as Error); + }); + }) + ); + const { deadlineFetch } = await import('./trpc'); + const promise = deadlineFetch('https://api.example.com/api/trpc'); + + vi.advanceTimersByTime(15_001); + // eslint-disable-next-line typescript-eslint/await-thenable + await vi.runAllTicks(); + + await expect(promise).rejects.toThrow('timed out after 15000ms'); + } finally { + process.off('unhandledRejection', swallowUnhandledRejection); + vi.useRealTimers(); + } + }); + + it('rejects immediately when the caller signal is already aborted', async () => { + const { deadlineFetch } = await import('./trpc'); + const controller = new AbortController(); + controller.abort(new Error('caller cancelled')); + + const promise = deadlineFetch('https://api.example.com/api/trpc', { + signal: controller.signal, + }); + + await expect(promise).rejects.toThrow('caller cancelled'); + expect(mockFetch).not.toHaveBeenCalled(); + }); + + it('rejects with the caller abort reason when the signal is aborted during fetch', async () => { + vi.useFakeTimers(); + try { + // Fetch hangs unless its signal aborts. + mockFetch.mockImplementation( + // eslint-disable-next-line typescript-eslint/promise-function-async + (_url, init?: RequestInit) => + new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal?.aborted) { + reject(signal.reason as Error); + return; + } + signal?.addEventListener('abort', () => { + reject(signal.reason as Error); + }); + }) + ); + + const { deadlineFetch } = await import('./trpc'); + const controller = new AbortController(); + const promise = deadlineFetch('https://api.example.com/api/trpc', { + signal: controller.signal, + }); + + // Let the fetch start, then abort from the caller side. + await vi.advanceTimersByTimeAsync(0); + controller.abort(new Error('caller cancelled mid-flight')); + + // eslint-disable-next-line typescript-eslint/await-thenable + await vi.runAllTicks(); + + await expect(promise).rejects.toThrow('caller cancelled mid-flight'); + } finally { + vi.useRealTimers(); + } + }); +}); diff --git a/apps/mobile/src/lib/trpc.ts b/apps/mobile/src/lib/trpc.ts index a1d7e8fd2de..0bfd6569602 100644 --- a/apps/mobile/src/lib/trpc.ts +++ b/apps/mobile/src/lib/trpc.ts @@ -1,6 +1,7 @@ import { type MobileRouter } from '@kilocode/trpc/mobile'; import { createTRPCClient, httpBatchLink, httpLink, splitLink } from '@trpc/client'; import { createTRPCContext } from '@trpc/tanstack-react-query'; +import { CONTROL_PLANE_DEADLINE_MS, withDeadline } from '@kilocode/event-service'; import * as SecureStore from 'expo-secure-store'; import { API_BASE_URL, E2E_LATENCY_MESSAGES_MS, E2E_LATENCY_SESSION_MS } from '@/lib/config'; @@ -58,6 +59,25 @@ const e2eLatencyFetch: typeof fetch = async (url, init) => { const e2eFetch = E2E_LATENCY_SESSION_MS > 0 || E2E_LATENCY_MESSAGES_MS > 0 ? e2eLatencyFetch : fetch; +/** + * Fetch wrapper that adds a control-plane deadline (15 s). When E2E latency + * values are set the deadline is extended by the applicable delay so the + * synthetic latency does not eat into the real request budget. + */ +export const deadlineFetch: typeof fetch = async (url, init) => { + const delayMs = e2eLatencyForUrl(requestUrlString(url)); + const totalDeadline = CONTROL_PLANE_DEADLINE_MS + delayMs; + const response = await withDeadline( + totalDeadline, + async signal => { + const res = await e2eFetch(url, { ...init, signal }); + return res; + }, + init?.signal ?? undefined + ); + return response; +}; + async function getAuthHeaders() { const token = await SecureStore.getItemAsync(AUTH_TOKEN_KEY); if (!token) { @@ -83,14 +103,14 @@ async function getAuthHeaders() { const singleLink = httpLink({ url: trpcUrl, headers: getAuthHeaders, - fetch: e2eFetch, + fetch: deadlineFetch, methodOverride: 'POST', }); const batchLink = httpBatchLink({ url: trpcUrl, headers: getAuthHeaders, - fetch: e2eFetch, + fetch: deadlineFetch, methodOverride: 'POST', }); diff --git a/apps/web/src/app/(app)/claw/kilo-chat/components/kiloChatLayoutCache.test.ts b/apps/web/src/app/(app)/claw/kilo-chat/components/kiloChatLayoutCache.test.ts index c460bb46610..50871a9a138 100644 --- a/apps/web/src/app/(app)/claw/kilo-chat/components/kiloChatLayoutCache.test.ts +++ b/apps/web/src/app/(app)/claw/kilo-chat/components/kiloChatLayoutCache.test.ts @@ -72,18 +72,19 @@ function createFakeKiloChatClient(): FakeKiloChatClient & { } function createFakeEventService() { - let reconnectHandler: (() => void) | null = null; + let resyncHandler: (() => void) | null = null; return { - onReconnect: (handler: () => void) => { - reconnectHandler = handler; + onReconnect: (_handler: () => void) => () => {}, + emitReconnect: () => { + resyncHandler?.(); + }, + onResync: (handler: () => void) => { + resyncHandler = handler; return () => { - reconnectHandler = null; + resyncHandler = null; }; }, - emitReconnect: () => { - reconnectHandler?.(); - }, }; } 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 e04c84fba07..0b14fe34e62 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 @@ -272,7 +272,12 @@ export type GetSessionOutput = { upstreamBranch?: string; /** Custom agents stored on this session (slug + name, plus optional model and thinking-effort overrides). */ - runtimeAgents?: Array<{ slug: string; name: string; model?: string; variant?: string }>; + runtimeAgents?: Array<{ + slug: string; + name: string; + model?: string; + variant?: string; + }>; // Execution status (grouped for cleaner API) execution: ExecutionStatus; @@ -292,6 +297,11 @@ export type GetSessionOutput = { // Versioning timestamp: number; version: number; + + // Latest persisted event ID from the Cloud Agent event log. Numeric and + // present after at least one event has been persisted. Null when no events + // exist, and absent when the DO cannot be reached (fails open). + latestEventId?: number | null; }; /** @@ -395,7 +405,10 @@ function isInsufficientCreditsError(err: unknown): boolean { if (httpStatus === 402) { return true; } - const cause = err.cause as { error?: { status?: number }; suppressed?: { status?: number } }; + const cause = err.cause as { + error?: { status?: number }; + suppressed?: { status?: number }; + }; if (cause?.error?.status === 402 || cause?.suppressed?.status === 402) { return true; } @@ -583,7 +596,10 @@ export class CloudAgentNextClient { } catch (error) { console.error(`Error interrupting session ${sessionId}:`, error); captureException(error, { - tags: { source: 'cloud-agent-next-client', endpoint: 'interruptSession' }, + tags: { + source: 'cloud-agent-next-client', + endpoint: 'interruptSession', + }, extra: { sessionId }, }); throw error; @@ -686,7 +702,10 @@ export class CloudAgentNextClient { } captureException(normalizedError, { - tags: { source: 'cloud-agent-next-client', endpoint: 'initiateFromPreparedSession' }, + tags: { + source: 'cloud-agent-next-client', + endpoint: 'initiateFromPreparedSession', + }, extra: { input }, }); throw normalizedError; @@ -736,7 +755,10 @@ export class CloudAgentNextClient { } catch (error) { captureException(error, { tags: { source: 'cloud-agent-next-client', endpoint: 'resizeTerminal' }, - extra: { cloudAgentSessionId: input.cloudAgentSessionId, ptyId: input.ptyId }, + extra: { + cloudAgentSessionId: input.cloudAgentSessionId, + ptyId: input.ptyId, + }, }); throw error; } @@ -748,7 +770,10 @@ export class CloudAgentNextClient { } catch (error) { captureException(error, { tags: { source: 'cloud-agent-next-client', endpoint: 'closeTerminal' }, - extra: { cloudAgentSessionId: input.cloudAgentSessionId, ptyId: input.ptyId }, + extra: { + cloudAgentSessionId: input.cloudAgentSessionId, + ptyId: input.ptyId, + }, }); throw error; } @@ -783,7 +808,10 @@ export class CloudAgentNextClient { return await this.client.answerPermission.mutate(input); } catch (error) { captureException(error, { - tags: { source: 'cloud-agent-next-client', endpoint: 'answerPermission' }, + tags: { + source: 'cloud-agent-next-client', + endpoint: 'answerPermission', + }, extra: { sessionId: input.sessionId, permissionId: input.permissionId }, }); throw error; diff --git a/apps/web/src/lib/session-ingest-client.ts b/apps/web/src/lib/session-ingest-client.ts index 223e268e1e1..dd4681e0c6e 100644 --- a/apps/web/src/lib/session-ingest-client.ts +++ b/apps/web/src/lib/session-ingest-client.ts @@ -133,6 +133,7 @@ export type SessionMessagesPageOptions = { export type SessionMessagesPageResult = { kiloSessionId: string; history: KiloSdkMessageHistory | null; + watermarkEventId?: number | null; }; /** diff --git a/apps/web/src/routers/cli-sessions-v2-router.test.ts b/apps/web/src/routers/cli-sessions-v2-router.test.ts index 7096f436b3f..28410adca38 100644 --- a/apps/web/src/routers/cli-sessions-v2-router.test.ts +++ b/apps/web/src/routers/cli-sessions-v2-router.test.ts @@ -19,6 +19,31 @@ import type { fetchSessionMessagesPage as FetchSessionMessagesPageType } from '@ import { notifyCliSessionRenamed } from '@/lib/cloud-agent/session-events'; import { captureException } from '@sentry/nextjs'; +// Mock the cloud agent client for watermark and runtime-state tests. +const mockGetSession = jest.fn().mockRejectedValue(new Error('not mocked')); + +jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ + createCloudAgentNextClient: jest.fn(() => ({ + getSession: mockGetSession, + })), + rethrowAsPaymentRequired: jest.fn(), + InsufficientCreditsError: class InsufficientCreditsError extends Error { + constructor(message = 'Insufficient credits') { + super(message); + this.name = 'InsufficientCreditsError'; + } + }, +})); + +jest.mock('@/lib/tokens', () => { + const actual: Record = jest.requireActual('@/lib/tokens'); + return { + ...actual, + generateApiToken: jest.fn(() => 'test-api-token'), + generateInternalServiceToken: jest.fn(() => 'test-internal-token'), + }; +}); + jest.mock('@/lib/config.server', () => { const actual: Record = jest.requireActual('@/lib/config.server'); return { @@ -167,7 +192,9 @@ describe('cli-sessions-v2-router', () => { it('returns validated snapshot info together with messages', async () => { const caller = await createCallerForUser(regularUser.id); - const result = await caller.cliSessionsV2.getSessionMessages({ session_id: sessionId }); + const result = await caller.cliSessionsV2.getSessionMessages({ + session_id: sessionId, + }); expect(result).toEqual({ info: { @@ -225,7 +252,10 @@ describe('cli-sessions-v2-router', () => { role: 'user' as const, time: { created: 1761000000100 }, agent: 'build', - model: { providerID: 'openrouter', modelID: 'anthropic/claude-sonnet-4' }, + model: { + providerID: 'openrouter', + modelID: 'anthropic/claude-sonnet-4', + }, }, parts: [ { @@ -241,7 +271,10 @@ describe('cli-sessions-v2-router', () => { nextCursor: 'opaque-cursor', omittedItemCount: 0, }; - fetchSessionMessagesPage.mockResolvedValueOnce({ kiloSessionId: sessionId, history }); + fetchSessionMessagesPage.mockResolvedValueOnce({ + kiloSessionId: sessionId, + history, + }); const caller = await createCallerForUser(regularUser.id); const result = await caller.cliSessionsV2.getSessionMessagesPage({ @@ -249,7 +282,11 @@ describe('cli-sessions-v2-router', () => { limit: 50, }); - expect(result).toEqual({ kiloSessionId: sessionId, history }); + expect(result).toEqual({ + kiloSessionId: sessionId, + history, + watermarkEventId: null, + }); expect(fetchSessionMessagesPage).toHaveBeenCalledWith(sessionId, regularUser.id, { limit: 50, }); @@ -285,7 +322,9 @@ describe('cli-sessions-v2-router', () => { }); const caller = await createCallerForUser(regularUser.id); - await caller.cliSessionsV2.getSessionMessagesPage({ session_id: sessionId }); + await caller.cliSessionsV2.getSessionMessagesPage({ + session_id: sessionId, + }); // The tRPC input schema must fill in the shared default before the // client is called so the worker's bounded reader always sees a @@ -318,8 +357,14 @@ describe('cli-sessions-v2-router', () => { }); it('preserves retryable_failure so the UI can offer Retry', async () => { - const history = { kind: 'retryable_failure' as const, phase: 'page_parts' as const }; - fetchSessionMessagesPage.mockResolvedValueOnce({ kiloSessionId: sessionId, history }); + const history = { + kind: 'retryable_failure' as const, + phase: 'page_parts' as const, + }; + fetchSessionMessagesPage.mockResolvedValueOnce({ + kiloSessionId: sessionId, + history, + }); const caller = await createCallerForUser(regularUser.id); const result = await caller.cliSessionsV2.getSessionMessagesPage({ @@ -327,7 +372,11 @@ describe('cli-sessions-v2-router', () => { limit: 10, }); - expect(result).toEqual({ kiloSessionId: sessionId, history }); + expect(result).toEqual({ + kiloSessionId: sessionId, + history, + watermarkEventId: null, + }); }); it('preserves too_large and invalid_data as non-retryable outcomes', async () => { @@ -340,13 +389,20 @@ describe('cli-sessions-v2-router', () => { { kind: 'invalid_data' as const }, ]) { fetchSessionMessagesPage.mockReset(); - fetchSessionMessagesPage.mockResolvedValueOnce({ kiloSessionId: sessionId, history }); + fetchSessionMessagesPage.mockResolvedValueOnce({ + kiloSessionId: sessionId, + history, + }); const caller = await createCallerForUser(regularUser.id); const result = await caller.cliSessionsV2.getSessionMessagesPage({ session_id: sessionId, limit: 10, }); - expect(result).toEqual({ kiloSessionId: sessionId, history }); + expect(result).toEqual({ + kiloSessionId: sessionId, + history, + watermarkEventId: null, + }); } }); @@ -365,6 +421,7 @@ describe('cli-sessions-v2-router', () => { expect(result).toEqual({ kiloSessionId: sessionId, history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + watermarkEventId: null, }); }); @@ -372,7 +429,10 @@ describe('cli-sessions-v2-router', () => { const caller = await createCallerForUser(otherUser.id); await expect( - caller.cliSessionsV2.getSessionMessagesPage({ session_id: sessionId, limit: 50 }) + caller.cliSessionsV2.getSessionMessagesPage({ + session_id: sessionId, + limit: 50, + }) ).rejects.toMatchObject({ code: 'NOT_FOUND' }); expect(fetchSessionMessagesPage).not.toHaveBeenCalled(); }); @@ -404,7 +464,10 @@ describe('cli-sessions-v2-router', () => { const caller = await createCallerForUser(regularUser.id); await expect( - caller.cliSessionsV2.getSessionMessagesPage({ session_id: orgSessionId, limit: 50 }) + caller.cliSessionsV2.getSessionMessagesPage({ + session_id: orgSessionId, + limit: 50, + }) ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); expect(fetchSessionMessagesPage).not.toHaveBeenCalled(); } finally { @@ -416,7 +479,10 @@ describe('cli-sessions-v2-router', () => { const caller = await createCallerForUser(regularUser.id); await expect( - caller.cliSessionsV2.getSessionMessagesPage({ session_id: sessionId, limit: 101 }) + caller.cliSessionsV2.getSessionMessagesPage({ + session_id: sessionId, + limit: 101, + }) ).rejects.toThrow(); expect(fetchSessionMessagesPage).not.toHaveBeenCalled(); }); @@ -425,7 +491,10 @@ describe('cli-sessions-v2-router', () => { const caller = await createCallerForUser(regularUser.id); await expect( - caller.cliSessionsV2.getSessionMessagesPage({ session_id: sessionId, limit: 0 }) + caller.cliSessionsV2.getSessionMessagesPage({ + session_id: sessionId, + limit: 0, + }) ).rejects.toThrow(); expect(fetchSessionMessagesPage).not.toHaveBeenCalled(); }); @@ -472,7 +541,10 @@ describe('cli-sessions-v2-router', () => { const caller = await createCallerForUser(regularUser.id); await expect( - caller.cliSessionsV2.getSessionMessagesPage({ session_id: sessionId, cursor: 'bad' }) + caller.cliSessionsV2.getSessionMessagesPage({ + session_id: sessionId, + cursor: 'bad', + }) ).rejects.toThrow(); expect(fetchSessionMessagesPage).not.toHaveBeenCalled(); }); @@ -494,10 +566,191 @@ describe('cli-sessions-v2-router', () => { const caller = await createCallerForUser(regularUser.id); await expect( - caller.cliSessionsV2.getSessionMessagesPage({ session_id: 'not-a-session', limit: 50 }) + caller.cliSessionsV2.getSessionMessagesPage({ + session_id: 'not-a-session', + limit: 50, + }) ).rejects.toThrow(); expect(fetchSessionMessagesPage).not.toHaveBeenCalled(); }); + + describe('watermark', () => { + const cloudAgentSessionId = 'agent_wm123456-1234-1234-1234-123456789abc'; + const watermarkSessionId = 'ses_watermark_test_123456'; + + beforeEach(async () => { + mockGetSession.mockReset().mockRejectedValue(new Error('not mocked')); + await db.insert(cli_sessions_v2).values({ + session_id: watermarkSessionId, + cloud_agent_session_id: cloudAgentSessionId, + kilo_user_id: regularUser.id, + created_on_platform: 'cloud-agent', + }); + }); + + afterEach(async () => { + await db.delete(cli_sessions_v2).where(eq(cli_sessions_v2.session_id, watermarkSessionId)); + }); + + it('carries a present watermark on the initial page response', async () => { + mockGetSession.mockResolvedValue({ latestEventId: 42 }); + fetchSessionMessagesPage.mockResolvedValueOnce({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + }); + + const caller = await createCallerForUser(regularUser.id); + const result = await caller.cliSessionsV2.getSessionMessagesPage({ + session_id: watermarkSessionId, + limit: 50, + }); + + expect(result).toEqual({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + watermarkEventId: 42, + }); + expect(mockGetSession).toHaveBeenCalledWith(cloudAgentSessionId); + // Verify both the watermark and the page were fetched. + expect(fetchSessionMessagesPage).toHaveBeenCalled(); + }); + + it('carries null watermark when the DO has no events', async () => { + mockGetSession.mockResolvedValue({ latestEventId: null }); + fetchSessionMessagesPage.mockResolvedValueOnce({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + }); + + const caller = await createCallerForUser(regularUser.id); + const result = await caller.cliSessionsV2.getSessionMessagesPage({ + session_id: watermarkSessionId, + limit: 50, + }); + + expect(result).toEqual({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + watermarkEventId: null, + }); + }); + + it('fails open: a watermark read failure returns null and page still succeeds', async () => { + mockGetSession.mockRejectedValue(new Error('DO unreachable')); + fetchSessionMessagesPage.mockResolvedValueOnce({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + }); + + const caller = await createCallerForUser(regularUser.id); + const result = await caller.cliSessionsV2.getSessionMessagesPage({ + session_id: watermarkSessionId, + limit: 50, + }); + + // The page must succeed with null watermark — no INTERNAL_SERVER_ERROR. + expect(result).toEqual({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + watermarkEventId: null, + }); + }); + + it('skips the watermark read when the session has no cloud_agent_session_id', async () => { + // Remove the CA session ID and test that no watermark fetch happens. + await db + .update(cli_sessions_v2) + .set({ cloud_agent_session_id: null }) + .where(eq(cli_sessions_v2.session_id, watermarkSessionId)); + + fetchSessionMessagesPage.mockResolvedValueOnce({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + }); + + const caller = await createCallerForUser(regularUser.id); + const result = await caller.cliSessionsV2.getSessionMessagesPage({ + session_id: watermarkSessionId, + limit: 50, + }); + + expect(mockGetSession).not.toHaveBeenCalled(); + expect(result).toEqual({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + watermarkEventId: null, + }); + }); + + it('skips the watermark read on cursor pages (continuation reads)', async () => { + mockGetSession.mockRejectedValue(new Error('must not be called')); + fetchSessionMessagesPage.mockResolvedValueOnce({ + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + }); + + const caller = await createCallerForUser(regularUser.id); + const validCursor = btoa( + JSON.stringify({ id: 'msg_user_01', time: 1761000000100 }) + ).replace(/=+$/, ''); + const result = await caller.cliSessionsV2.getSessionMessagesPage({ + session_id: watermarkSessionId, + limit: 50, + cursor: validCursor, + }); + + // Cursor pages must skip the Cloud Agent read entirely. + expect(mockGetSession).not.toHaveBeenCalled(); + expect(result.watermarkEventId).toBeNull(); + }); + + it('resolves getSession before fetchSessionMessagesPage starts (deferred-promise order proof)', async () => { + // Deferred promise proves the router awaits getSession before calling + // fetchSessionMessagesPage. Without this ordering, the watermark + // read could race with the page fetch. + let getSessionResolved = false; + let resolveGetSession!: (value: { latestEventId: number }) => void; + + const getSessionPromise = new Promise<{ latestEventId: number }>(resolve => { + resolveGetSession = resolve; + }); + mockGetSession.mockReturnValue(getSessionPromise); + + let fetchPageCalled = false; + let pageResolved = false; + fetchSessionMessagesPage.mockImplementationOnce(async () => { + fetchPageCalled = true; + // If getSession hasn't resolved yet, the ordering is broken. + expect(getSessionResolved).toBe(true); + pageResolved = true; + return { + kiloSessionId: watermarkSessionId, + history: { messages: [], nextCursor: null, omittedItemCount: 0 }, + }; + }); + + const caller = await createCallerForUser(regularUser.id); + const resultPromise = caller.cliSessionsV2.getSessionMessagesPage({ + session_id: watermarkSessionId, + limit: 50, + }); + + // Let the router reach the getSession call. + await new Promise(r => setTimeout(r, 0)); + + // Neither getSession nor fetchPage has resolved yet. + expect(getSessionResolved).toBe(false); + expect(fetchPageCalled).toBe(false); + + // Resolve getSession — the router must then call fetchSessionMessagesPage. + getSessionResolved = true; + resolveGetSession({ latestEventId: 99 }); + + const result = await resultPromise; + expect(result.watermarkEventId).toBe(99); + expect(pageResolved).toBe(true); + }); + }); }); describe('shareForWebhookTrigger', () => { @@ -1009,7 +1262,9 @@ describe('cli-sessions-v2-router', () => { const caller = await createCallerForUser(regularUser.id); await expect( - caller.cliSessionsV2.getSessionMessages({ session_id: organizationSessionId }) + caller.cliSessionsV2.getSessionMessages({ + session_id: organizationSessionId, + }) ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); expect(fetchSpy).not.toHaveBeenCalled(); } finally { @@ -1175,7 +1430,9 @@ describe('cli-sessions-v2-router', () => { try { const caller = await createCallerForUser(regularUser.id); await expect( - caller.cliSessionsV2.getWithRuntimeState({ session_id: sessionWithPr }) + caller.cliSessionsV2.getWithRuntimeState({ + session_id: sessionWithPr, + }) ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); } finally { await db @@ -1373,7 +1630,9 @@ describe('cli-sessions-v2-router', () => { it('search returns associatedPr per row matching the search string', async () => { const caller = await createCallerForUser(regularUser.id); - const result = await caller.cliSessionsV2.search({ search_string: 'session' }); + const result = await caller.cliSessionsV2.search({ + search_string: 'session', + }); const withPr = result.results.find(s => s.session_id === sessionWithPr); const withoutPr = result.results.find(s => s.session_id === sessionWithoutPr); @@ -1459,7 +1718,9 @@ describe('cli-sessions-v2-router', () => { }); const caller = await createCallerForUser(regularUser.id); - const result = await caller.cliSessionsV2.refreshAssociatedPullRequest({ sessionId }); + const result = await caller.cliSessionsV2.refreshAssociatedPullRequest({ + sessionId, + }); expect(mockedFetchPullRequestForBranch).toHaveBeenCalledTimes(1); expect(mockedFetchPullRequestForBranch).toHaveBeenCalledWith({ @@ -1505,7 +1766,9 @@ describe('cli-sessions-v2-router', () => { mockedFetchPullRequestForBranch.mockResolvedValue(null); const caller = await createCallerForUser(regularUser.id); - const result = await caller.cliSessionsV2.refreshAssociatedPullRequest({ sessionId }); + const result = await caller.cliSessionsV2.refreshAssociatedPullRequest({ + sessionId, + }); expect(result.associatedPr).toBeNull(); const rows = await readCacheRows(); @@ -1537,7 +1800,9 @@ describe('cli-sessions-v2-router', () => { expect(mockedFetchPullRequestForBranch).toHaveBeenCalledTimes(1); // Second call within the throttle window short-circuits. - const second = await caller.cliSessionsV2.refreshAssociatedPullRequest({ sessionId }); + const second = await caller.cliSessionsV2.refreshAssociatedPullRequest({ + sessionId, + }); expect(mockedFetchPullRequestForBranch).toHaveBeenCalledTimes(1); expect(second.associatedPr).toBeNull(); }); @@ -1555,7 +1820,9 @@ describe('cli-sessions-v2-router', () => { }); const caller = await createCallerForUser(regularUser.id); - const result = await caller.cliSessionsV2.refreshAssociatedPullRequest({ sessionId }); + const result = await caller.cliSessionsV2.refreshAssociatedPullRequest({ + sessionId, + }); expect(mockedFetchPullRequestForBranch).not.toHaveBeenCalled(); expect(result.associatedPr).toMatchObject({ number: 99, state: 'open' }); @@ -1764,7 +2031,9 @@ describe('cli-sessions-v2-router', () => { it('defaults search to updated_at descending', async () => { const caller = await createCallerForUser(regularUser.id); - const result = await caller.cliSessionsV2.search({ search_string: searchableTitle }); + const result = await caller.cliSessionsV2.search({ + search_string: searchableTitle, + }); expect(result.results.map(session => session.session_id)).toEqual([ olderCreatedNewerUpdated, @@ -2008,7 +2277,10 @@ describe('cli-sessions-v2-router', () => { const caller = await createCallerForUser(regularUser.id); // Fixtures are ordered by created_at; placeholders sit between visible // rows. limit=2 over created_at should page only visible rows. - const page1 = await caller.cliSessionsV2.list({ limit: 2, orderBy: 'created_at' }); + const page1 = await caller.cliSessionsV2.list({ + limit: 2, + orderBy: 'created_at', + }); const page1Ids = page1.cliSessions.map(session => session.session_id); expect(page1Ids).toEqual([normalCliId, costOnlyZeroId]); @@ -2028,7 +2300,9 @@ describe('cli-sessions-v2-router', () => { it('search by exact session_id does not return a bare placeholder', async () => { const caller = await createCallerForUser(regularUser.id); - const result = await caller.cliSessionsV2.search({ search_string: placeholderId }); + const result = await caller.cliSessionsV2.search({ + search_string: placeholderId, + }); expect(result.results.map(session => session.session_id)).not.toContain(placeholderId); expect(result.total).toBe(0); diff --git a/apps/web/src/routers/cli-sessions-v2-router.ts b/apps/web/src/routers/cli-sessions-v2-router.ts index ab20ce0d285..6fbae9f9858 100644 --- a/apps/web/src/routers/cli-sessions-v2-router.ts +++ b/apps/web/src/routers/cli-sessions-v2-router.ts @@ -713,7 +713,10 @@ export const cliSessionsV2Router = createTRPCRouter({ await addOrganizationCondition(whereConditions, ctx, input.organizationId); - const { rows } = await db.execute<{ git_url: string; last_used_at: string }>(sql` + const { rows } = await db.execute<{ + git_url: string; + last_used_at: string; + }>(sql` SELECT ${cli_sessions_v2.git_url} AS git_url, MAX(${cli_sessions_v2.updated_at}) AS last_used_at FROM ${cli_sessions_v2} WHERE ${joinWithAnd(whereConditions)} @@ -829,7 +832,28 @@ export const cliSessionsV2Router = createTRPCRouter({ getSessionMessagesPage: baseProcedure .input(GetSessionMessagesPageInputSchema) .query(async ({ ctx, input }) => { - await getSessionWithAccessCheck(input.session_id, ctx); + const session = await getSessionWithAccessCheck(input.session_id, ctx); + + // Read the event-log watermark from the existing getSession response + // before the initial history page, so the transport can use `fromId` + // on its first WebSocket connect instead of `replay=false`. Cursor + // pages skip the Cloud Agent read — the watermark is only seeded once. + // Failures are swallowed and return null so the page endpoint is + // never blocked on an optional watermark read. + let watermarkEventId: number | null = null; + if (!input.cursor && session.cloud_agent_session_id) { + try { + const authToken = generateApiToken(ctx.user); + const client = createCloudAgentNextClient(authToken); + const sessionState = await client.getSession(session.cloud_agent_session_id); + watermarkEventId = sessionState.latestEventId ?? null; + } catch (error) { + console.warn( + `Failed to fetch watermark for session ${input.session_id}:`, + error instanceof Error ? error.message : error + ); + } + } let result; try { @@ -866,7 +890,7 @@ export const cliSessionsV2Router = createTRPCRouter({ }); } - return result; + return { ...result, watermarkEventId }; }), /** @@ -1074,7 +1098,10 @@ export const cliSessionsV2Router = createTRPCRouter({ .limit(1); if (!row) { - throw new TRPCError({ code: 'NOT_FOUND', message: 'Session not found' }); + throw new TRPCError({ + code: 'NOT_FOUND', + message: 'Session not found', + }); } const { session } = row; @@ -1217,7 +1244,10 @@ export const cliSessionsV2Router = createTRPCRouter({ const normalizedGitUrl = normalizeGitUrl(gitUrl); const ownerValues = session.organization_id - ? { owned_by_organization_id: session.organization_id, owned_by_user_id: null } + ? { + owned_by_organization_id: session.organization_id, + owned_by_user_id: null, + } : { owned_by_organization_id: null, owned_by_user_id: ctx.user.id }; const conflictTarget = session.organization_id @@ -1369,7 +1399,10 @@ export const cliSessionsV2Router = createTRPCRouter({ userId: ctx.user.id, }).catch(error => { captureException(error, { - tags: { source: 'cli-sessions-v2-router', endpoint: 'rename-notify' }, + tags: { + source: 'cli-sessions-v2-router', + endpoint: 'rename-notify', + }, extra: { sessionId: session_id }, }); }) diff --git a/dev/local/scripts/start-tunnel.test.ts b/dev/local/scripts/start-tunnel.test.ts new file mode 100644 index 00000000000..1733ec6d083 --- /dev/null +++ b/dev/local/scripts/start-tunnel.test.ts @@ -0,0 +1,44 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { buildDockerLocalEnvValues } from './start-tunnel'; + +test('returns BACKEND_API_URL with localhost when using default port', () => { + const env = buildDockerLocalEnvValues('3000', '8795', '8808'); + + assert.equal(env.BACKEND_API_URL, 'http://localhost:3000'); + assert.equal(env.KILOCODE_API_BASE_URL, 'http://host.docker.internal:3000/api/gateway/'); + assert.equal(env.KILOCLAW_CHECKIN_URL, 'http://host.docker.internal:8795/api/controller/checkin'); + assert.equal(env.KILOCHAT_BASE_URL, 'http://host.docker.internal:8808'); +}); + +test('writes BACKEND_API_URL and KILOCODE values with port 4900', () => { + const env = buildDockerLocalEnvValues('4900', '8795', '8808'); + + assert.equal(env.BACKEND_API_URL, 'http://localhost:4900'); + assert.equal(env.KILOCODE_API_BASE_URL, 'http://host.docker.internal:4900/api/gateway/'); + assert.equal(env.KILOCLAW_CHECKIN_URL, 'http://host.docker.internal:8795/api/controller/checkin'); + assert.equal(env.KILOCHAT_BASE_URL, 'http://host.docker.internal:8808'); +}); + +test('BACKEND_API_URL uses localhost, not host.docker.internal', () => { + const env = buildDockerLocalEnvValues('3000', '8795', '8808'); + + assert.ok(env.BACKEND_API_URL.startsWith('http://localhost:')); + assert.ok( + env.KILOCODE_API_BASE_URL.startsWith('http://host.docker.internal:'), + 'KILOCODE_API_BASE_URL must use host.docker.internal for Docker containers' + ); +}); + +test('includes all four required keys', () => { + const env = buildDockerLocalEnvValues('3000', '8795', '8808'); + + const keys = Object.keys(env).sort(); + assert.deepEqual(keys, [ + 'BACKEND_API_URL', + 'KILOCHAT_BASE_URL', + 'KILOCLAW_CHECKIN_URL', + 'KILOCODE_API_BASE_URL', + ]); +}); diff --git a/dev/local/scripts/start-tunnel.ts b/dev/local/scripts/start-tunnel.ts index b417c40c210..889bc9ffd10 100644 --- a/dev/local/scripts/start-tunnel.ts +++ b/dev/local/scripts/start-tunnel.ts @@ -2,6 +2,7 @@ import { spawn, spawnSync } from 'node:child_process'; import * as fs from 'node:fs'; import * as os from 'node:os'; import * as path from 'node:path'; +import { fileURLToPath } from 'node:url'; const repoRoot = path.resolve(import.meta.dirname, '../../..'); const devVarsPath = path.join(repoRoot, 'services/kiloclaw/.dev.vars'); @@ -155,190 +156,211 @@ function prefixAndWrite(label: string, chunk: Buffer): void { } } -const port = process.argv[2] ?? '3000'; -const controllerPort = process.argv[3] ?? '8795'; -const kiloChatPort = process.argv[4] ?? '8808'; -const config = loadTunnelConfig(); -const provider = loadKiloClawProvider(); +export function buildDockerLocalEnvValues( + port: string, + controllerPort: string, + kiloChatPort: string +): Record { + return { + BACKEND_API_URL: `http://localhost:${port}`, + KILOCODE_API_BASE_URL: `http://${DOCKER_HOST_INTERNAL}:${port}/api/gateway/`, + KILOCLAW_CHECKIN_URL: `http://${DOCKER_HOST_INTERNAL}:${controllerPort}/api/controller/checkin`, + KILOCHAT_BASE_URL: `http://${DOCKER_HOST_INTERNAL}:${kiloChatPort}`, + }; +} -const children: Array<{ label: string; child: ReturnType }> = []; +function main(): void { + const port = process.argv[2] ?? '3000'; + const controllerPort = process.argv[3] ?? '8795'; + const kiloChatPort = process.argv[4] ?? '8808'; + const config = loadTunnelConfig(); + const provider = loadKiloClawProvider(); -function trackChild(label: string, child: ReturnType): void { - children.push({ label, child }); -} + const children: Array<{ label: string; child: ReturnType }> = []; -function stopAllChildren(signal: NodeJS.Signals): void { - for (const { child } of children) { - child.kill(signal); + function trackChild(label: string, child: ReturnType): void { + children.push({ label, child }); } -} -let exiting = false; - -function exitAndStopOthers(originLabel: string, code: number | null): void { - if (exiting) return; - exiting = true; - for (const { label, child } of children) { - if (label !== originLabel) { - child.kill('SIGTERM'); + function stopAllChildren(signal: NodeJS.Signals): void { + for (const { child } of children) { + child.kill(signal); } } - process.exit(code ?? 1); -} -function startQuickTunnel(options: { - label: string; - localPort: string; - onUrl: (url: string) => void; -}): void { - const { label, localPort, onUrl } = options; - const child = spawn('cloudflared', ['tunnel', '--url', `http://localhost:${localPort}`], { - stdio: ['ignore', 'pipe', 'pipe'], - }); - trackChild(label, child); - - console.log(`Starting quick tunnel (${label}) -> http://localhost:${localPort}...`); - - let captured = false; - const urlPattern = /https:\/\/[a-z0-9-]+\.trycloudflare\.com/; - const handleOutput = (data: Buffer) => { - prefixAndWrite(label, data); - - if (captured) return; - const match = data.toString().match(urlPattern); - if (!match) return; - - captured = true; - onUrl(match[0]); - }; + let exiting = false; - child.stdout.on('data', handleOutput); - child.stderr.on('data', handleOutput); - child.on('close', code => exitAndStopOthers(label, code)); -} + function exitAndStopOthers(originLabel: string, code: number | null): void { + if (exiting) return; + exiting = true; + for (const { label, child } of children) { + if (label !== originLabel) { + child.kill('SIGTERM'); + } + } + process.exit(code ?? 1); + } -if (provider === DOCKER_LOCAL_PROVIDER) { - const apiUrl = `http://${DOCKER_HOST_INTERNAL}:${port}/api/gateway/`; - const checkinUrl = `http://${DOCKER_HOST_INTERNAL}:${controllerPort}/api/controller/checkin`; - const kiloChatUrl = `http://${DOCKER_HOST_INTERNAL}:${kiloChatPort}`; + function startQuickTunnel(options: { + label: string; + localPort: string; + onUrl: (url: string) => void; + }): void { + const { label, localPort, onUrl } = options; + const child = spawn('cloudflared', ['tunnel', '--url', `http://localhost:${localPort}`], { + stdio: ['ignore', 'pipe', 'pipe'], + }); + trackChild(label, child); + + console.log(`Starting quick tunnel (${label}) -> http://localhost:${localPort}...`); + + let captured = false; + const urlPattern = /https:\/\/[a-z0-9-]+\.trycloudflare\.com/; + const handleOutput = (data: Buffer) => { + prefixAndWrite(label, data); + + if (captured) return; + const match = data.toString().match(urlPattern); + if (!match) return; + + captured = true; + onUrl(match[0]); + }; + + child.stdout.on('data', handleOutput); + child.stderr.on('data', handleOutput); + child.on('close', code => exitAndStopOthers(label, code)); + } - updateEnvValue(devVarsPath, 'KILOCODE_API_BASE_URL', apiUrl); - updateEnvValue(devVarsPath, 'KILOCLAW_CHECKIN_URL', checkinUrl); - updateEnvValue(devVarsPath, 'KILOCHAT_BASE_URL', kiloChatUrl); + if (provider === DOCKER_LOCAL_PROVIDER) { + const envValues = buildDockerLocalEnvValues(port, controllerPort, kiloChatPort); - console.log('Docker-local provider detected; skipping Cloudflare quick tunnels.'); - console.log(`Set KILOCODE_API_BASE_URL=${apiUrl}`); - console.log(`Set KILOCLAW_CHECKIN_URL=${checkinUrl}`); - console.log(`Set KILOCHAT_BASE_URL=${kiloChatUrl}`); + for (const [key, value] of Object.entries(envValues)) { + updateEnvValue(devVarsPath, key, value); + } - setInterval(() => undefined, 60_000); -} else if (spawnSync('cloudflared', ['version'], { stdio: 'ignore' }).error) { - console.error( - 'cloudflared not found on PATH. Install it:\n https://developers.cloudflare.com/cloudflare-one/connections/connect-networks/downloads/\n brew install cloudflared' - ); - process.exit(1); -} + console.log('Docker-local provider detected; skipping Cloudflare quick tunnels.'); + console.log(`Set BACKEND_API_URL=${envValues.BACKEND_API_URL}`); + console.log(`Set KILOCODE_API_BASE_URL=${envValues.KILOCODE_API_BASE_URL}`); + console.log(`Set KILOCLAW_CHECKIN_URL=${envValues.KILOCLAW_CHECKIN_URL}`); + console.log(`Set KILOCHAT_BASE_URL=${envValues.KILOCHAT_BASE_URL}`); + + setInterval(() => undefined, 60_000); + } else if (spawnSync('cloudflared', ['version'], { stdio: 'ignore' }).error) { + console.error( + 'cloudflared not found on PATH. Install it:\n https://developers.cloudflare.com/cloudflare-one/connections/connect-networks/downloads/\n brew install cloudflared' + ); + process.exit(1); + } -if (provider !== DOCKER_LOCAL_PROVIDER && (config.tunnelName || config.tunnelConfig)) { - const label = 'kiloclaw-tunnel'; - const args = config.tunnelConfig - ? ['tunnel', '--config', config.tunnelConfig, 'run'] - : ['tunnel', 'run', config.tunnelName]; - const child = spawn('cloudflared', args, { - stdio: ['ignore', 'pipe', 'pipe'], - }); - trackChild(label, child); - - const appOrigin = originFromHostname(config.appHostname); - const kiloclawOrigin = originFromHostname(config.kiloclawHostname); - const kiloChatOrigin = originFromHostname(config.kiloChatHostname); - - console.log( - `Named tunnel: ${config.tunnelConfig || config.tunnelName}` + - `${appOrigin ? `\n app -> ${appOrigin}` : ''}` + - `${kiloclawOrigin ? `\n kiloclaw -> ${kiloclawOrigin}` : ''}` + - `${kiloChatOrigin ? `\n kilochat -> ${kiloChatOrigin}` : ''}` - ); + if (provider !== DOCKER_LOCAL_PROVIDER && (config.tunnelName || config.tunnelConfig)) { + const label = 'kiloclaw-tunnel'; + const args = config.tunnelConfig + ? ['tunnel', '--config', config.tunnelConfig, 'run'] + : ['tunnel', 'run', config.tunnelName]; + const child = spawn('cloudflared', args, { + stdio: ['ignore', 'pipe', 'pipe'], + }); + trackChild(label, child); + + const appOrigin = originFromHostname(config.appHostname); + const kiloclawOrigin = originFromHostname(config.kiloclawHostname); + const kiloChatOrigin = originFromHostname(config.kiloChatHostname); + + console.log( + `Named tunnel: ${config.tunnelConfig || config.tunnelName}` + + `${appOrigin ? `\n app -> ${appOrigin}` : ''}` + + `${kiloclawOrigin ? `\n kiloclaw -> ${kiloclawOrigin}` : ''}` + + `${kiloChatOrigin ? `\n kilochat -> ${kiloChatOrigin}` : ''}` + ); + + if (appOrigin) { + const apiUrl = `${appOrigin}/api/gateway/`; + updateEnvValue(devVarsPath, 'BACKEND_API_URL', appOrigin); + updateEnvValue(devVarsPath, 'KILOCODE_API_BASE_URL', apiUrl); + console.log(`Set BACKEND_API_URL=${appOrigin}`); + console.log(`Set KILOCODE_API_BASE_URL=${apiUrl}`); - if (appOrigin) { - const apiUrl = `${appOrigin}/api/gateway/`; - updateEnvValue(devVarsPath, 'BACKEND_API_URL', appOrigin); - updateEnvValue(devVarsPath, 'KILOCODE_API_BASE_URL', apiUrl); - console.log(`Set BACKEND_API_URL=${appOrigin}`); - console.log(`Set KILOCODE_API_BASE_URL=${apiUrl}`); - - if (config.updateAppEnv) { - updateEnvValue(envLocalPath, 'APP_URL_OVERRIDE', appOrigin); - updateEnvValue(envLocalPath, 'NEXTAUTH_URL', appOrigin); - console.log(`Set APP_URL_OVERRIDE=${appOrigin}`); - console.log(`Set NEXTAUTH_URL=${appOrigin}`); + if (config.updateAppEnv) { + updateEnvValue(envLocalPath, 'APP_URL_OVERRIDE', appOrigin); + updateEnvValue(envLocalPath, 'NEXTAUTH_URL', appOrigin); + console.log(`Set APP_URL_OVERRIDE=${appOrigin}`); + console.log(`Set NEXTAUTH_URL=${appOrigin}`); + } } - } - if (kiloclawOrigin) { - const checkinUrl = `${kiloclawOrigin}/api/controller/checkin`; - updateEnvValue(devVarsPath, 'KILOCLAW_CHECKIN_URL', checkinUrl); - console.log(`Set KILOCLAW_CHECKIN_URL=${checkinUrl}`); + if (kiloclawOrigin) { + const checkinUrl = `${kiloclawOrigin}/api/controller/checkin`; + updateEnvValue(devVarsPath, 'KILOCLAW_CHECKIN_URL', checkinUrl); + console.log(`Set KILOCLAW_CHECKIN_URL=${checkinUrl}`); - if (config.updateAppEnv) { - updateEnvValue(envLocalPath, 'KILOCLAW_API_URL', kiloclawOrigin); - console.log(`Set KILOCLAW_API_URL=${kiloclawOrigin}`); + if (config.updateAppEnv) { + updateEnvValue(envLocalPath, 'KILOCLAW_API_URL', kiloclawOrigin); + console.log(`Set KILOCLAW_API_URL=${kiloclawOrigin}`); + } } - } - if (kiloChatOrigin) { - updateEnvValue(devVarsPath, 'KILOCHAT_BASE_URL', kiloChatOrigin); - console.log(`Set KILOCHAT_BASE_URL=${kiloChatOrigin}`); - } + if (kiloChatOrigin) { + updateEnvValue(devVarsPath, 'KILOCHAT_BASE_URL', kiloChatOrigin); + console.log(`Set KILOCHAT_BASE_URL=${kiloChatOrigin}`); + } - appendEnvListValues( - devVarsPath, - 'OPENCLAW_ALLOWED_ORIGINS', - [appOrigin, kiloclawOrigin, kiloChatOrigin].filter((origin): origin is string => !!origin) - ); + appendEnvListValues( + devVarsPath, + 'OPENCLAW_ALLOWED_ORIGINS', + [appOrigin, kiloclawOrigin, kiloChatOrigin].filter((origin): origin is string => !!origin) + ); + + child.stdout.on('data', data => prefixAndWrite(label, data)); + child.stderr.on('data', data => prefixAndWrite(label, data)); + child.on('close', code => exitAndStopOthers(label, code)); + } else if (provider !== DOCKER_LOCAL_PROVIDER) { + startQuickTunnel({ + label: 'gateway', + localPort: port, + onUrl: url => { + const apiUrl = `${url}/api/gateway/`; + updateEnvValue(devVarsPath, 'KILOCODE_API_BASE_URL', apiUrl); + console.log(`\nGateway tunnel URL: ${url}`); + console.log(`Set KILOCODE_API_BASE_URL=${apiUrl}`); + }, + }); + + startQuickTunnel({ + label: 'controller', + localPort: controllerPort, + onUrl: url => { + const checkinUrl = `${url}/api/controller/checkin`; + updateEnvValue(devVarsPath, 'KILOCLAW_CHECKIN_URL', checkinUrl); + console.log(`\nController tunnel URL: ${url}`); + console.log(`Set KILOCLAW_CHECKIN_URL=${checkinUrl}`); + }, + }); + + startQuickTunnel({ + label: 'kilo-chat', + localPort: kiloChatPort, + onUrl: url => { + updateEnvValue(devVarsPath, 'KILOCHAT_BASE_URL', url); + console.log(`\nKilo-chat tunnel URL: ${url}`); + console.log(`Set KILOCHAT_BASE_URL=${url}`); + }, + }); + } - child.stdout.on('data', data => prefixAndWrite(label, data)); - child.stderr.on('data', data => prefixAndWrite(label, data)); - child.on('close', code => exitAndStopOthers(label, code)); -} else if (provider !== DOCKER_LOCAL_PROVIDER) { - startQuickTunnel({ - label: 'gateway', - localPort: port, - onUrl: url => { - const apiUrl = `${url}/api/gateway/`; - updateEnvValue(devVarsPath, 'KILOCODE_API_BASE_URL', apiUrl); - console.log(`\nGateway tunnel URL: ${url}`); - console.log(`Set KILOCODE_API_BASE_URL=${apiUrl}`); - }, - }); - - startQuickTunnel({ - label: 'controller', - localPort: controllerPort, - onUrl: url => { - const checkinUrl = `${url}/api/controller/checkin`; - updateEnvValue(devVarsPath, 'KILOCLAW_CHECKIN_URL', checkinUrl); - console.log(`\nController tunnel URL: ${url}`); - console.log(`Set KILOCLAW_CHECKIN_URL=${checkinUrl}`); - }, - }); - - startQuickTunnel({ - label: 'kilo-chat', - localPort: kiloChatPort, - onUrl: url => { - updateEnvValue(devVarsPath, 'KILOCHAT_BASE_URL', url); - console.log(`\nKilo-chat tunnel URL: ${url}`); - console.log(`Set KILOCHAT_BASE_URL=${url}`); - }, - }); + for (const signal of ['SIGINT', 'SIGTERM'] as const) { + process.on(signal, () => { + stopAllChildren(signal); + if (children.length === 0) { + process.exit(0); + } + }); + } } -for (const signal of ['SIGINT', 'SIGTERM'] as const) { - process.on(signal, () => { - stopAllChildren(signal); - if (children.length === 0) { - process.exit(0); - } - }); +if ( + process.argv[1] !== undefined && + fileURLToPath(import.meta.url) === path.resolve(process.argv[1]) +) { + main(); } diff --git a/dev/local/tmux.test.ts b/dev/local/tmux.test.ts index 927f575e9c8..b5bd0face36 100644 --- a/dev/local/tmux.test.ts +++ b/dev/local/tmux.test.ts @@ -61,6 +61,9 @@ test( const tmux = (...args: string[]) => execFileSync('tmux', args, { stdio: 'ignore' }); const tmuxOutput = (...args: string[]) => execFileSync('tmux', args, { encoding: 'utf8' }).trim(); + // set-environment -g needs a running server; hold one open for the test. + const keepAliveSession = `${sessionName}-keepalive`; + tmux('new-session', '-d', '-s', keepAliveSession); const savedGlobalEnvironment = new Map(); for (const key of ['PWD', 'OLDPWD']) { try { @@ -120,6 +123,11 @@ test( // The variable may not exist in this tmux version's environment. } } + try { + tmux('kill-session', '-t', keepAliveSession); + } catch { + // The keep-alive session may already be gone. + } } } ); diff --git a/packages/cloud-agent-sdk/src/cli-live-transport.test.ts b/packages/cloud-agent-sdk/src/cli-live-transport.test.ts index 428ff65f5df..b6d35b5ea36 100644 --- a/packages/cloud-agent-sdk/src/cli-live-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cli-live-transport.test.ts @@ -2359,7 +2359,8 @@ describe('CliLiveTransport createSession', () => { KILO_SESSION_ID, 'create_session', { protocolVersion: 1 }, - 'owner' + 'owner', + expect.any(String) as string ); expect(userWebConnection.sendCommandToConnection).not.toHaveBeenCalled(); transport.destroy(); @@ -2486,7 +2487,8 @@ describe('CliLiveTransport createSession', () => { KILO_SESSION_ID, 'create_session', { protocolVersion: 1 }, - 'owner-a' + 'owner-a', + expect.any(String) as string ); expect(userWebConnection.sendCommandToConnection).not.toHaveBeenCalled(); transport.destroy(); @@ -2545,7 +2547,8 @@ describe('CliLiveTransport createSession', () => { model: { providerID: 'anthropic', modelID: 'claude-sonnet-4', variant: 'high' }, orgId: 'aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee', }, - 'owner' + 'owner', + expect.any(String) as string ); transport.destroy(); }); @@ -2652,6 +2655,39 @@ describe('CliLiveTransport createSession', () => { ).toHaveLength(2); transport.destroy(); }); + + it('uses a fresh mutationId for the bare retry so the DO treats it as a new command', async () => { + const connection = createConnection(); + const { transport, userWebConnection } = createTransportWithSinks({ connection }); + + transport.connect(); + emitOwner(connection); + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); + + jest.mocked(userWebConnection.sendCommand).mockClear(); + jest + .mocked(userWebConnection.sendCommand) + .mockRejectedValueOnce(new CommandDeliveredError('invalid create_session command')) + .mockResolvedValueOnce({ protocolVersion: 1, sessionID: NEW_KILO_SESSION_ID }); + + await transport.createSession?.({ agent: 'code' }); + + const createCalls = jest + .mocked(userWebConnection.sendCommand) + .mock.calls.filter(([, command]) => command === 'create_session'); + expect(createCalls).toHaveLength(2); + + // The second call (bare retry) must have a mutationId that differs + // from the first call's mutationId. + const firstMutationId = createCalls[0]?.[4]; + const secondMutationId = createCalls[1]?.[4]; + expect(firstMutationId).toBeDefined(); + expect(secondMutationId).toBeDefined(); + expect(secondMutationId).not.toBe(firstMutationId); + transport.destroy(); + }); }); describe('CliLiveTransport exitSession', () => { diff --git a/packages/cloud-agent-sdk/src/cli-live-transport.ts b/packages/cloud-agent-sdk/src/cli-live-transport.ts index a674799a6cb..ebda25fc70f 100644 --- a/packages/cloud-agent-sdk/src/cli-live-transport.ts +++ b/packages/cloud-agent-sdk/src/cli-live-transport.ts @@ -138,7 +138,10 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor config.onCapabilitiesChange?.(next); } let catalogRequestGeneration = 0; - let catalogRequestInFlight: { ownerConnectionId: string; generation: number } | null = null; + let catalogRequestInFlight: { + ownerConnectionId: string; + generation: number; + } | null = null; // Command catalog discovery runs on its own generation so it stays // independent of model discovery: a model refresh that completes or // fails must not drop an in-flight command catalog request. @@ -168,7 +171,10 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor } function publishRemoteCommandState(next: RemoteCommandState): void { - remoteCommandState = { ...next, commands: deepCopyCommands(next.commands) }; + remoteCommandState = { + ...next, + commands: deepCopyCommands(next.commands), + }; config.onRemoteCommandStateChange?.(next); } @@ -184,7 +190,10 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // state callback each hold an independent snapshot. const snapshotted = deepCopyCommands(commands); lastValidCommands = snapshotted; - sink.onServiceEvent({ type: 'commands.available', commands: deepCopyCommands(snapshotted) }); + sink.onServiceEvent({ + type: 'commands.available', + commands: deepCopyCommands(snapshotted), + }); } function snapshotCommands(): SlashCommandInfo[] { @@ -589,7 +598,11 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor } } - async function sendCommand(command: string, data: unknown): Promise { + async function sendCommand( + command: string, + data: unknown, + mutationId?: string + ): Promise { const expectedOwnerConnectionId = ownerConnectionId; if (!expectedOwnerConnectionId) throw new Error('Remote session has no connected owner'); @@ -598,7 +611,8 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor config.kiloSessionId, command, data, - expectedOwnerConnectionId + expectedOwnerConnectionId, + ...(mutationId !== undefined ? [mutationId] : []) ); } catch (error) { if (error instanceof UserWebCommandError && error.code === 'SESSION_OWNER_CHANGED') { @@ -861,6 +875,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // snapshotted before the await and SESSION_OWNER_CHANGED is handled. // Extended fields degrade via a single bare retry on the exact // old-CLI delivered error; other failures are hard rejects. + const mutationId = input?.mutationId ?? crypto.randomUUID(); const wireData = buildCreateSessionWireData(input); const hasExtendedFields = wireData.agent !== undefined || @@ -868,7 +883,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor wireData.orgId !== undefined; let result: unknown; try { - result = await sendCommand('create_session', wireData); + result = await sendCommand('create_session', wireData, mutationId); } catch (error) { // Only bare-retry when extended fields made the original wire differ // from `{ protocolVersion: 1 }`; otherwise the retry is identical. @@ -877,7 +892,11 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor error instanceof CommandDeliveredError && error.message === INVALID_CREATE_SESSION_COMMAND ) { - result = await sendCommand('create_session', { protocolVersion: 1 }); + result = await sendCommand( + 'create_session', + { protocolVersion: 1 }, + crypto.randomUUID() + ); } else { throw error; } 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 b0983f8f5f2..e4df5261118 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts @@ -5,6 +5,7 @@ import type { CloudAgentEvent } from './event-types'; import { createEventHelpers } from './__fixtures__/helpers'; import type { ChatEvent, ServiceEvent } from './normalizer'; +import type { SessionStatus } from './schemas'; import { createCloudAgentTransport } from './cloud-agent-transport'; import type { SessionSnapshotPageOutcome } from './types'; import { kiloId, cloudAgentId, makeSnapshot } from './test-helpers'; @@ -115,7 +116,12 @@ describe('CloudAgentTransport event routing', () => { await flushPromises(); sendRaw( kilocode('message.updated', { - info: { id: 'msg-1', sessionID: 'ses-1', role: 'assistant', time: { created: 1 } }, + info: { + id: 'msg-1', + sessionID: 'ses-1', + role: 'assistant', + time: { created: 1 }, + }, }) ); @@ -133,7 +139,12 @@ describe('CloudAgentTransport event routing', () => { transport.connect(); await flushPromises(); - sendRaw(kilocode('session.status', { sessionID: 'ses-1', status: { type: 'busy' } })); + sendRaw( + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) + ); // session.created from snapshot replay + session.status from sendRaw expect(serviceEvents).toHaveLength(2); @@ -166,7 +177,10 @@ describe('CloudAgentTransport event routing', () => { data: { commands }, }); - expect(serviceEvents).toContainEqual({ type: 'commands.available', commands }); + expect(serviceEvents).toContainEqual({ + type: 'commands.available', + commands, + }); transport.destroy(); }); @@ -180,12 +194,22 @@ describe('CloudAgentTransport event routing', () => { // Chat event sendRaw( kilocode('message.updated', { - info: { id: 'msg-1', sessionID: 'ses-1', role: 'assistant', time: { created: 1 } }, + info: { + id: 'msg-1', + sessionID: 'ses-1', + role: 'assistant', + time: { created: 1 }, + }, }) ); // Service event - sendRaw(kilocode('session.status', { sessionID: 'ses-1', status: { type: 'busy' } })); + sendRaw( + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) + ); // Chat event (delta) sendRaw( @@ -253,14 +277,26 @@ describe('CloudAgentTransport unexpected disconnect', () => { transport.connect(); await flushPromises(); - sendRaw(kilocode('session.status', { sessionID: 'ses-1', status: { type: 'busy' } })); + sendRaw( + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) + ); // Non-auth close triggers onUnexpectedDisconnect in connection.ts - mockWs.onclose?.({ code: 1011, reason: 'network dropped', wasClean: false } as CloseEvent); + mockWs.onclose?.({ + code: 1011, + reason: 'network dropped', + wasClean: false, + } as CloseEvent); const stoppedEvents = serviceEvents.filter(e => e.type === 'stopped'); expect(stoppedEvents).toHaveLength(1); - expect(stoppedEvents[0]).toEqual({ type: 'stopped', reason: 'transport-disconnected' }); + expect(stoppedEvents[0]).toEqual({ + type: 'stopped', + reason: 'transport-disconnected', + }); transport.destroy(); }); @@ -269,7 +305,12 @@ describe('CloudAgentTransport unexpected disconnect', () => { const { transport, serviceEvents } = createTransportWithSinks(); transport.connect(); - sendRaw(kilocode('session.status', { sessionID: 'ses-1', status: { type: 'busy' } })); + sendRaw( + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) + ); // complete → stopped(complete) through normal pipeline sendRaw(createEvent('complete', { currentBranch: 'main' })); @@ -277,7 +318,11 @@ describe('CloudAgentTransport unexpected disconnect', () => { const stoppedBefore = serviceEvents.filter(e => e.type === 'stopped').length; // Now close unexpectedly — should NOT generate another stopped - mockWs.onclose?.({ code: 1011, reason: 'network dropped', wasClean: false } as CloseEvent); + mockWs.onclose?.({ + code: 1011, + reason: 'network dropped', + wasClean: false, + } as CloseEvent); const stoppedAfter = serviceEvents.filter(e => e.type === 'stopped').length; expect(stoppedAfter).toBe(stoppedBefore); @@ -309,7 +354,12 @@ describe('CloudAgentTransport ticket handling', () => { expect(webSocketConstructor).toHaveBeenCalled(); // Events still route correctly (session.created from replay + session.status from sendRaw) - sendRaw(kilocode('session.status', { sessionID: 'ses-1', status: { type: 'busy' } })); + sendRaw( + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) + ); expect(serviceEvents).toHaveLength(2); transport.destroy(); @@ -408,7 +458,12 @@ describe('CloudAgentTransport command delegation', () => { expect(api.send).toHaveBeenCalledWith({ sessionId: 'ses-1', - payload: { type: 'prompt', prompt: 'hello', mode: 'code', model: 'gpt-4' }, + payload: { + type: 'prompt', + prompt: 'hello', + mode: 'code', + model: 'gpt-4', + }, }); transport.destroy(); @@ -453,7 +508,12 @@ describe('CloudAgentTransport command delegation', () => { expect(api.send).toHaveBeenCalledWith({ sessionId: 'ses-1', - payload: { type: 'prompt', prompt: 'read it', mode: 'code', model: 'gpt-4' }, + payload: { + type: 'prompt', + prompt: 'read it', + mode: 'code', + model: 'gpt-4', + }, attachments, }); @@ -504,7 +564,10 @@ describe('CloudAgentTransport command delegation', () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); - void transport.respondToPermission!({ requestId: 'req-3', response: 'once' }); + void transport.respondToPermission!({ + requestId: 'req-3', + response: 'once', + }); expect(api.respondToPermission).toHaveBeenCalledWith({ sessionId: 'ses-1', @@ -585,7 +648,11 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { newMockWs.onopen?.(new Event('open')); sendRawOn( newMockWs, - establishEvent ?? kilocode('session.status', { sessionID: 'ses-1', status: { type: 'busy' } }) + establishEvent ?? + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) ); return newMockWs; @@ -626,7 +693,11 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { expect(replayedCreated).toHaveLength(0); transport.destroy(); - newMockWs.onclose?.({ code: 1000, reason: '', wasClean: true } as CloseEvent); + newMockWs.onclose?.({ + code: 1000, + reason: '', + wasClean: true, + } as CloseEvent); } finally { jest.useRealTimers(); } @@ -680,7 +751,11 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { expect(chatEvents.filter(e => e.type === 'message.part.updated')).toHaveLength(2); transport.destroy(); - newMockWs.onclose?.({ code: 1000, reason: '', wasClean: true } as CloseEvent); + newMockWs.onclose?.({ + code: 1000, + reason: '', + wasClean: true, + } as CloseEvent); } finally { jest.useRealTimers(); } @@ -728,7 +803,13 @@ describe('CloudAgentTransport page-seam', () => { onServiceEvent: event => serviceEvents.push(event), }); - return { transport, chatEvents, serviceEvents, fetchSnapshotPage, onInitialPageLoaded }; + return { + transport, + chatEvents, + serviceEvents, + fetchSnapshotPage, + onInitialPageLoaded, + }; } it('uses fetchSnapshotPage for the initial bounded read and reports the page via onInitialPageLoaded', async () => { @@ -895,7 +976,11 @@ describe('CloudAgentTransport page-seam', () => { sendRaw(sentinel); // Trigger an unexpected disconnect on the first socket. - mockWs.onclose?.({ code: 1006, reason: '', wasClean: false } as CloseEvent); + mockWs.onclose?.({ + code: 1006, + reason: '', + wasClean: false, + } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); @@ -919,4 +1004,568 @@ describe('CloudAgentTransport page-seam', () => { jest.useRealTimers(); } }); + + describe('watermark cursor', () => { + it('uses fromId=0 on first connect when the page carries a watermark (closes materialization gap)', async () => { + const page = { + kind: 'success' as const, + info: { id: 'ses-1' } as const, + messages: [], + nextCursor: null, + omittedItemCount: 0, + watermarkEventId: 42, + }; + const { transport } = createTransportWithPageFetch(page); + + transport.connect(); + await flushPromises(); + + expect(webSocketConstructor).toHaveBeenCalledTimes(1); + const wsUrl = String(webSocketConstructor.mock.calls[0]?.[0]); + // A present watermark seeds lastEventId=0 so the DO replays every + // stored event from the beginning — not fromId=42, which could skip + // events not yet materialized into the SessionIngest page. + expect(wsUrl).toContain('fromId=0'); + expect(wsUrl).not.toContain('replay=false'); + + transport.destroy(); + }); + + it('uses replay=false on first connect when the page has no watermark', async () => { + const page = { + kind: 'success' as const, + info: { id: 'ses-1' } as const, + messages: [], + nextCursor: null, + omittedItemCount: 0, + }; + const { transport } = createTransportWithPageFetch(page); + + transport.connect(); + await flushPromises(); + + expect(webSocketConstructor).toHaveBeenCalledTimes(1); + const wsUrl = String(webSocketConstructor.mock.calls[0]?.[0]); + expect(wsUrl).toContain('replay=false'); + expect(wsUrl).not.toContain('fromId'); + + transport.destroy(); + }); + + it('uses fromId with watermark null on reconnect when wire events have set lastEventId', async () => { + jest.useFakeTimers(); + try { + async function flushMicrotasks(): Promise { + for (let i = 0; i < 10; i++) { + await Promise.resolve(); + } + } + + const page = { + kind: 'success' as const, + info: { id: 'ses-1' } as const, + messages: [], + nextCursor: null, + omittedItemCount: 0, + watermarkEventId: 10, + }; + const fetchSnapshotPage = jest.fn().mockResolvedValue(page); + + const chatEvents: ChatEvent[] = []; + const serviceEvents: ServiceEvent[] = []; + const factory = createCloudAgentTransport({ + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), + api: createMockApi(), + getTicket: () => 'test-ticket', + fetchSnapshot: () => Promise.reject(new Error('should not be called')), + fetchSnapshotPage, + websocketBaseUrl: 'ws://localhost:9999', + }); + + const transport = factory({ + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), + }); + + transport.connect(); + await flushMicrotasks(); + + // Verify first connect used fromId=0 (the watermark seeds + // lastEventId=0 so the DO replays all stored events). + const firstUrl = String(webSocketConstructor.mock.calls[0]?.[0]); + expect(firstUrl).toContain('fromId=0'); + + // Send a wire event with eventId > 0 to advance the live cursor + const establish = { + eventId: 55, + executionId: null, + sessionId: 'ses-1', + streamEventType: 'kilocode', + timestamp: new Date().toISOString(), + data: { + type: 'session.status', + properties: { sessionID: 'ses-1', status: { type: 'busy' } }, + }, + }; + sendRaw(establish); + + // Disconnect and reconnect + mockWs.onclose?.({ + code: 1006, + reason: '', + wasClean: false, + } as CloseEvent); + jest.advanceTimersByTime(2000); + await flushMicrotasks(); + + const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); + // Reconnect must use the live cursor (55), not the stale watermark (10) + expect(reconnectUrl).toContain('fromId=55'); + // fetchSnapshotPage must NOT be called on reconnect when a cursor exists + // because the socket replays missed events via fromId + expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); + + transport.destroy(); + } finally { + jest.useRealTimers(); + } + }); + + it('reconnect on a watermarked page with only sentinel events uses fromId=0 to close the materialization gap', async () => { + jest.useFakeTimers(); + try { + async function flushMicrotasks(): Promise { + for (let i = 0; i < 10; i++) { + await Promise.resolve(); + } + } + + const page = { + kind: 'success' as const, + info: { id: 'ses-1' } as const, + messages: [], + nextCursor: null, + omittedItemCount: 0, + watermarkEventId: 42, + }; + const fetchSnapshotPage = jest.fn().mockResolvedValue(page); + + const chatEvents: ChatEvent[] = []; + const serviceEvents: ServiceEvent[] = []; + const factory = createCloudAgentTransport({ + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), + api: createMockApi(), + getTicket: () => 'test-ticket', + fetchSnapshot: () => Promise.reject(new Error('should not be called')), + fetchSnapshotPage, + websocketBaseUrl: 'ws://localhost:9999', + }); + + const transport = factory({ + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), + }); + + transport.connect(); + await flushMicrotasks(); + + // Verify first connect used fromId=0 (the watermark seeds + // lastEventId=0 so the DO replays all stored events, closing + // the materialization gap). + const firstUrl = String(webSocketConstructor.mock.calls[0]?.[0]); + expect(firstUrl).toContain('fromId=0'); + expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); + + // Send only sentinel events (eventId: 0) — they do NOT advance the + // live cursor, so lastEventId stays at 0. + const sentinel: CloudAgentEvent = { + eventId: 0, + executionId: null, + sessionId: 'ses-1', + streamEventType: 'kilocode', + timestamp: new Date().toISOString(), + data: { + type: 'session.status', + properties: { sessionID: 'ses-1', status: { type: 'busy' } }, + }, + }; + sendRaw(sentinel); + + // Disconnect and reconnect. + mockWs.onclose?.({ + code: 1006, + reason: '', + wasClean: false, + } as CloseEvent); + jest.advanceTimersByTime(2000); + await flushMicrotasks(); + + const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); + // Reconnect uses fromId=0 because sentinel events never advance + // lastEventId past 0. The DO replays all stored events again. + expect(reconnectUrl).toContain('fromId=0'); + expect(reconnectUrl).not.toContain('replay=false'); + + // No second page fetch — lastEventId (0) is not null, so + // onReconnected skips the snapshot fallback path entirely. + expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); + + transport.destroy(); + } finally { + jest.useRealTimers(); + } + }); + + it('delivers all events when the watermark is ahead of the page materialization', async () => { + // Regression: when the Cloud Agent DO event-log watermark leads the + // SessionIngest history page, events between the page's last + // materialized event and the watermark must not be skipped. + // + // Setup: page with watermark=10, messages that represent content + // from events 1–5 only. The DO replays all events via fromId=0. + // Events 6–10 (the gap) must reach the sink. + const page = { + kind: 'success' as const, + info: { id: 'ses-1' } as const, + messages: [], + nextCursor: null, + omittedItemCount: 0, + watermarkEventId: 10, + }; + const { transport, serviceEvents } = createTransportWithPageFetch(page); + + transport.connect(); + await flushPromises(); + + // Verify the first connect uses fromId=0 so the DO replays all + // stored events from the beginning, covering the materialization gap. + const wsUrl = String(webSocketConstructor.mock.calls[0]?.[0]); + expect(wsUrl).toContain('fromId=0'); + + const serviceCountBefore = serviceEvents.length; + + // Simulate the DO replaying events 1–10 (fromId=0 replays everything). + // Each event has a unique observable tag so we can identify which + // events reached the sink. + for (let id = 1; id <= 10; id++) { + sendRaw({ + eventId: id, + executionId: null, + sessionId: 'ses-1', + streamEventType: 'kilocode', + timestamp: new Date().toISOString(), + data: { + type: 'session.status', + properties: { + sessionID: 'ses-1', + status: { type: 'retry', attempt: id, message: `event-${id}`, next: id + 1 }, + }, + }, + }); + } + + // Collect the observable IDs from session.status retry events. + const deliveredIds = serviceEvents + .slice(serviceCountBefore) + .filter( + ( + e + ): e is Extract & { + status: Extract; + } => { + if (e.type !== 'session.status') return false; + return e.status.type === 'retry'; + } + ) + .map(e => e.status.attempt) + .sort((a, b) => a - b); + + // Every event 1–10 must be delivered. The dedupe with lastEventId=0 + // must not drop any of them. If the old watermark-as-cursor behavior + // were still in place, events 6–10 would be dropped (eventId ≤ 10). + expect(deliveredIds).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); + + transport.destroy(); + }); + }); +}); + +// --------------------------------------------------------------------------- +// Exactly-once event replay +// --------------------------------------------------------------------------- + +describe('CloudAgentTransport exactly-once event replay', () => { + /** + * Create a session.status event with an explicit eventId and a valid + * session status shape. Count-based tests verify which IDs were delivered. + */ + function eventWithId(eventId: number): CloudAgentEvent { + return { + eventId, + executionId: null, + sessionId: 'ses-1', + streamEventType: 'kilocode', + timestamp: new Date().toISOString(), + data: { + type: 'session.status', + properties: { sessionID: 'ses-1', status: { type: 'busy' } }, + }, + }; + } + + it('drops a duplicate event ID at the cursor', async () => { + const { transport, serviceEvents } = createTransportWithSinks(); + + transport.connect(); + await flushPromises(); + + // snapshot replay puts session.created at serviceEvents[0] + const serviceCountBefore = serviceEvents.length; + + // First delivery sets the cursor. + sendRaw(eventWithId(5)); + // Duplicate — same eventId as the cursor — must be dropped. + sendRaw(eventWithId(5)); + + // Only one new service event (the first delivery) reached the sink. + expect(serviceEvents.length).toBe(serviceCountBefore + 1); + + transport.destroy(); + }); + + it('drops a positive event ID below the cursor', async () => { + const { transport, serviceEvents } = createTransportWithSinks(); + + transport.connect(); + await flushPromises(); + + const serviceCountBefore = serviceEvents.length; + + // Advance the cursor to 10. + sendRaw(eventWithId(10)); + // eventId 5 is below the cursor — must be dropped. + sendRaw(eventWithId(5)); + + // Only the eventId=10 delivery reached the sink. + expect(serviceEvents.length).toBe(serviceCountBefore + 1); + + transport.destroy(); + }); + + it('delivers a higher positive event ID and advances the cursor', async () => { + const { transport, serviceEvents } = createTransportWithSinks(); + + transport.connect(); + await flushPromises(); + + const serviceCountBefore = serviceEvents.length; + + sendRaw(eventWithId(3)); + sendRaw(eventWithId(8)); + + // Both events delivered — cursor advanced from 3 to 8. + expect(serviceEvents.length).toBe(serviceCountBefore + 2); + + transport.destroy(); + }); + + it('always delivers event ID zero (synthetic sentinel)', async () => { + const { transport, serviceEvents } = createTransportWithSinks(); + + transport.connect(); + await flushPromises(); + + const serviceCountBefore = serviceEvents.length; + + // Sentinel events always pass through — even repeated ones. + sendRaw(eventWithId(0)); + sendRaw(eventWithId(0)); + sendRaw(eventWithId(0)); + + expect(serviceEvents.length).toBe(serviceCountBefore + 3); + + transport.destroy(); + }); + + it('replays 3 through 8 after 1 through 5 on resume, delivering each once', async () => { + jest.useFakeTimers(); + try { + async function flushMicrotasks(): Promise { + for (let i = 0; i < 10; i++) { + await Promise.resolve(); + } + } + + // Create events with observable unique IDs embedded in the status payload. + // The retry status type carries an `attempt` field we use as a tag. + function eventWithObservedId(eventId: number): CloudAgentEvent { + return { + eventId, + executionId: null, + sessionId: 'ses-1', + streamEventType: 'kilocode', + timestamp: new Date().toISOString(), + data: { + type: 'session.status', + properties: { + sessionID: 'ses-1', + status: { + type: 'retry', + attempt: eventId, + message: `exactly-once-${eventId}`, + next: eventId + 1, + }, + }, + }, + }; + } + + const { transport, serviceEvents } = createTransportWithSinks(); + + transport.connect(); + await flushMicrotasks(); + + // Phase 1: deliver events 1 through 5. + for (let id = 1; id <= 5; id++) { + sendRaw(eventWithObservedId(id)); + } + + // Disconnect (triggers reconnect via connection.ts reconnection logic). + mockWs.onclose?.({ + code: 1006, + reason: '', + wasClean: false, + } as CloseEvent); + jest.advanceTimersByTime(2000); + await flushMicrotasks(); + + // Phase 2: the DO replays 3 through 8 on the new socket. + // Events 3, 4, 5 must be dropped (≤ cursor at 5). + // Events 6, 7, 8 must be delivered. + const newMockWs = webSocketConstructor.mock.results.at(-1)?.value as MockWebSocket; + newMockWs.onopen?.(new Event('open')); + for (let id = 3; id <= 8; id++) { + newMockWs.onmessage?.({ + data: JSON.stringify(eventWithObservedId(id)), + } as MessageEvent); + } + await flushMicrotasks(); + + // Collect IDs delivered via session.status events tagged with retry. + const deliveredIds = serviceEvents + .filter( + ( + e + ): e is Extract & { + status: Extract; + } => { + if (e.type !== 'session.status') return false; + return e.status.type === 'retry'; + } + ) + .map(e => e.status.attempt) + .sort((a, b) => a - b); + + // Each ID from 1 through 8 must appear exactly once. + // Duplicates, missing IDs, or extra deliveries will fail the assertion. + expect(deliveredIds).toEqual([1, 2, 3, 4, 5, 6, 7, 8]); + + transport.destroy(); + newMockWs.onclose?.({ + code: 1000, + reason: '', + wasClean: true, + } as CloseEvent); + } finally { + jest.useRealTimers(); + } + }); + + it('cursor never moves backward after delivering a higher ID', async () => { + jest.useFakeTimers(); + try { + async function flushMicrotasks(): Promise { + for (let i = 0; i < 10; i++) { + await Promise.resolve(); + } + } + + const { transport, serviceEvents } = createTransportWithSinks(); + + transport.connect(); + await flushMicrotasks(); + + const serviceCountBefore = serviceEvents.length; + + // Advance the cursor to 10. + sendRaw(eventWithId(10)); + expect(serviceEvents.length).toBe(serviceCountBefore + 1); + + // Send eventId 5 — dropped, cursor must stay at 10. + sendRaw(eventWithId(5)); + expect(serviceEvents.length).toBe(serviceCountBefore + 1); + + // Disconnect and reconnect to inspect the replay cursor. + mockWs.onclose?.({ + code: 1006, + reason: '', + wasClean: false, + } as CloseEvent); + jest.advanceTimersByTime(2000); + await flushMicrotasks(); + + const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); + // Cursor must still be 10, not 5. + expect(reconnectUrl).toContain('fromId=10'); + + transport.destroy(); + } finally { + jest.useRealTimers(); + } + }); + + it('passes through sentinel (ID 0) events without affecting the cursor for replay URLs', async () => { + jest.useFakeTimers(); + try { + async function flushMicrotasks(): Promise { + for (let i = 0; i < 10; i++) { + await Promise.resolve(); + } + } + + const { transport, serviceEvents } = createTransportWithSinks(); + + transport.connect(); + await flushMicrotasks(); + + const serviceCountBefore = serviceEvents.length; + + // Advance cursor to 7. + sendRaw(eventWithId(7)); + expect(serviceEvents.length).toBe(serviceCountBefore + 1); + + // Sentinels flow through but do not change the cursor. + sendRaw(eventWithId(0)); + sendRaw(eventWithId(0)); + expect(serviceEvents.length).toBe(serviceCountBefore + 3); + + // Disconnect and reconnect — cursor must still be 7. + mockWs.onclose?.({ + code: 1006, + reason: '', + wasClean: false, + } as CloseEvent); + jest.advanceTimersByTime(2000); + await flushMicrotasks(); + + const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); + expect(reconnectUrl).toContain('fromId=7'); + + transport.destroy(); + } finally { + jest.useRealTimers(); + } + }); }); diff --git a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts index 77b4b8fbee9..8dce2fea12c 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts @@ -85,11 +85,15 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport function buildWebsocketUrl(): string { const url = new URL('/stream', websocketBaseUrl); url.searchParams.set('cloudAgentSessionId', config.sessionId); - if (lastEventId === null) { - // First connect: messages are pre-loaded via REST, skip the full replay. - url.searchParams.set('replay', 'false'); - } else { + if (lastEventId !== null) { + // Reconnect cursor or initial watermark: the DO replays everything + // after this id — either a live cursor from the wire, or the + // event-log watermark from the initial bounded page. url.searchParams.set('fromId', String(lastEventId)); + } else { + // No cursor and no watermark: messages are pre-loaded via REST, + // skip the DO replay. + url.searchParams.set('replay', 'false'); } return url.toString(); } @@ -145,6 +149,14 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport return null; } if (page.kind === 'success') { + // Seed lastEventId from the page's event-log watermark. A + // present watermark sets the cursor to 0 so the first + // WebSocket connect uses `fromId=0` — the DO replays every + // stored event, closing the gap when SessionIngest + // materialization lags behind the event-log high-water mark. + // On reconnect, wire events advance lastEventId past 0 and + // the live cursor takes over. + lastEventId = page.watermarkEventId != null ? 0 : null; config.onInitialPageLoaded?.(page); replayPage(page); return { ticket }; @@ -189,6 +201,13 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport lifecycleHooks: config.lifecycleHooks, websocketHeaders: config.websocketHeaders, onEvent: raw => { + // Drop replayed duplicates: positive IDs at or below the cursor + // are redundant — the sink saw them on the first delivery. + // eventId 0 (synthetic sentinel) is always delivered. + if (raw.eventId > 0 && lastEventId !== null && raw.eventId <= lastEventId) { + return; + } + if (raw.eventId > 0) { lastEventId = raw.eventId; } diff --git a/packages/cloud-agent-sdk/src/session.test.ts b/packages/cloud-agent-sdk/src/session.test.ts index 97fb0d4dd1a..0d7b8686f3a 100644 --- a/packages/cloud-agent-sdk/src/session.test.ts +++ b/packages/cloud-agent-sdk/src/session.test.ts @@ -510,7 +510,8 @@ describe('remote session create and retry commands', () => { kiloId('ses-remote'), 'create_session', { protocolVersion: 1 }, - 'owner' + 'owner', + expect.any(String) as string ); expect(userWebConnection.sendCommandToConnection).not.toHaveBeenCalled(); session.destroy(); @@ -544,7 +545,8 @@ describe('remote session create and retry commands', () => { model: { providerID: 'kilo', modelID: 'kilo-auto' }, orgId: 'aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee', }, - 'owner' + 'owner', + expect.any(String) as string ); session.destroy(); }); diff --git a/packages/cloud-agent-sdk/src/transport.ts b/packages/cloud-agent-sdk/src/transport.ts index 5aed52e3753..7f188b28a71 100644 --- a/packages/cloud-agent-sdk/src/transport.ts +++ b/packages/cloud-agent-sdk/src/transport.ts @@ -16,6 +16,8 @@ import type { ModelRef, RemoteModelOverride } from './remote-model-catalog'; * connection-scoped spawn helper. */ type CreateRemoteSessionInput = { + /** Reuse for a caller retry of the same create intent. */ + mutationId?: string; agent?: string; model?: { providerID: string; diff --git a/packages/cloud-agent-sdk/src/types.ts b/packages/cloud-agent-sdk/src/types.ts index 25aaf2fb1f3..84c2034c183 100644 --- a/packages/cloud-agent-sdk/src/types.ts +++ b/packages/cloud-agent-sdk/src/types.ts @@ -252,12 +252,19 @@ export type SessionSnapshot = { * page (or `null` when the history has been fully read); `omittedItemCount` * reports how many individual items the worker filtered out before the page * left the DO so the UI can faithfully report omissions. + * + * `watermarkEventId` carries the latest persisted Cloud Agent event-log + * event ID at the time the page was fetched. A present watermark makes the + * transport use `fromId` on its first WebSocket connect, closing the gap + * between page-snapshot and live-stream events; a null watermark keeps the + * existing `replay=false` behaviour. */ export type SessionSnapshotPage = { info: SessionInfo; messages: SessionSnapshot['messages']; nextCursor: string | null; omittedItemCount: number; + watermarkEventId?: number | null; }; /** diff --git a/packages/cloud-agent-sdk/src/user-web-connection.ts b/packages/cloud-agent-sdk/src/user-web-connection.ts index 5b81fbda2ee..7973d438f8b 100644 --- a/packages/cloud-agent-sdk/src/user-web-connection.ts +++ b/packages/cloud-agent-sdk/src/user-web-connection.ts @@ -66,6 +66,8 @@ type SendCommandToConnectionInput = { command: string; data: unknown; expectedConnectionId: string; + // Stable intent id forwarded to the relay for durable dedupe (D8). + mutationId?: string; }; type UserWebConnection = { @@ -94,7 +96,10 @@ type UserWebConnection = { sessionId: string, command: string, data: unknown, - expectedOwnerConnectionId?: string + expectedOwnerConnectionId?: string, + // Stable intent id forwarded to the relay for durable dedupe. + // Absent on legacy paths that do not re-issue (D5). + mutationId?: string ) => Promise; /** * Send a viewer command that is scoped to a specific CLI connection and has @@ -537,6 +542,10 @@ function createUserWebConnection( data: unknown; connectionId?: string; sessionId?: string; + // Stable intent id forwarded to the relay for durable dedupe (D8). + // Absent on legacy paths that do not re-issue; the relay falls back to + // a per-wire random correlation id. + mutationId?: string; }; /** @@ -591,6 +600,7 @@ function createUserWebConnection( command: wire.command, ...(wire.sessionId ? { sessionId: wire.sessionId } : {}), ...(wire.connectionId ? { connectionId: wire.connectionId } : {}), + ...(wire.mutationId ? { mutationId: wire.mutationId } : {}), data: wire.data, }) ); @@ -653,7 +663,7 @@ function createUserWebConnection( releaseConnection(); }; }, - sendCommand(sessionId, command, data, expectedOwnerConnectionId) { + sendCommand(sessionId, command, data, expectedOwnerConnectionId, mutationId) { return sendRawCommand({ command, // `expectedOwnerConnectionId` undefined still flows through `connectionId` @@ -662,6 +672,7 @@ function createUserWebConnection( ...(expectedOwnerConnectionId ? { connectionId: expectedOwnerConnectionId } : {}), sessionId, data, + ...(mutationId ? { mutationId } : {}), }); }, sendCommandToConnection(input) { @@ -669,6 +680,7 @@ function createUserWebConnection( command: input.command, connectionId: input.expectedConnectionId, data: input.data, + ...(input.mutationId ? { mutationId: input.mutationId } : {}), }); }, onCliEvent(sessionId, listener) { diff --git a/packages/event-service/src/__tests__/client.test.ts b/packages/event-service/src/__tests__/client.test.ts index f5383219ca1..901b7fbfaae 100644 --- a/packages/event-service/src/__tests__/client.test.ts +++ b/packages/event-service/src/__tests__/client.test.ts @@ -116,6 +116,7 @@ describe('EventServiceClient', () => { expect(fetch).toHaveBeenCalledWith('http://localhost:8080/connect-ticket', { method: 'POST', headers: { authorization: 'Bearer header.payload.sig' }, + signal: expect.any(AbortSignal) as unknown, }); expect(lastMockWs.url).toBe('ws://localhost:8080/connect?ticket=ticket-1'); expectNonSecretProtocol(lastMockWs.protocols); @@ -273,6 +274,7 @@ describe('EventServiceClient', () => { expect(fetch).toHaveBeenLastCalledWith('http://localhost:8080/connect-ticket', { method: 'POST', headers: { authorization: 'Bearer fresh.token.sig' }, + signal: expect.any(AbortSignal) as unknown, }); expect(allMockWs[1]?.url).toBe('ws://localhost:8080/connect?ticket=ticket-2'); expectNonSecretProtocol(allMockWs[1]?.protocols); @@ -572,6 +574,65 @@ describe('EventServiceClient', () => { } }); + it('concurrent connects leave one open socket and no orphan delivery', async () => { + const client = makeClient(); + + const received: unknown[] = []; + client.on('test.event', (_, payload) => received.push(payload)); + + // Start the first connect. It will create a socket inside connectOnce + // after getToken + fetchConnectionTicket, then wait for the open event. + const p1 = client.connect(); + + // Drain all pending microtasks so the first connect runs past the + // generation guard, creates its WebSocket, and the open event fires. + // Without this delay, a second connect() call synchronously bumps + // this.generation and the first connect bails before socket creation. + await new Promise(r => setTimeout(r, 0)); + expect(allMockWs.length).toBe(1); + + // Start the second connect. connectOnce() closes the first socket + // (line 154) before fetching a new ticket, producing an actual orphan. + const p2 = client.connect(); + + const [r1, r2] = await Promise.allSettled([p1, p2]); + + // Both must resolve. + expect(r1.status).toBe('fulfilled'); + expect(r2.status).toBe('fulfilled'); + + // Exactly one socket stays open — the last one created. + const connectedCount = allMockWs.filter(ws => ws.readyState === 1).length; + expect(connectedCount).toBe(1); + + // All earlier sockets are closed. + const closedCount = allMockWs.filter(ws => ws.readyState === 3).length; + expect(closedCount).toBe(allMockWs.length - 1); + + // The live socket delivers events. + const liveWs = allMockWs.find(ws => ws.readyState === 1)!; + liveWs.triggerMessage({ + type: 'event', + context: 'room:1', + event: 'test.event', + payload: 'live', + }); + expect(received).toEqual(['live']); + + // Orphan sockets cannot call handlers. + received.length = 0; + const stale = allMockWs.filter(ws => ws.readyState === 3); + for (const ws of stale) { + ws.triggerMessage({ + type: 'event', + context: 'room:1', + event: 'test.event', + payload: 'orphan', + }); + } + expect(received).toHaveLength(0); + }); + it('exports HandshakeTimeoutError', () => { // Sanity check on the public error type. const err = new HandshakeTimeoutError(); @@ -796,4 +857,561 @@ describe('EventServiceClient', () => { } }); }); + + describe('onResync and sequence gap detection', () => { + function eventWithSeq(seq: number, event = 'message.created') { + return { type: 'event', context: 'room:1', event, payload: { id: seq }, seq }; + } + + function eventWithoutSeq(event = 'message.created') { + return { type: 'event', context: 'room:1', event, payload: {} }; + } + + it('in-order events emit no resync', async () => { + const client = makeClient(); + await client.connect(); + + const resyncCalls: number[] = []; + client.onResync(() => resyncCalls.push(1)); + + const received: unknown[] = []; + client.on('message.created', (_, payload) => received.push(payload)); + + lastMockWs.triggerMessage(eventWithSeq(1)); + lastMockWs.triggerMessage(eventWithSeq(2)); + lastMockWs.triggerMessage(eventWithSeq(3)); + + expect(resyncCalls).toHaveLength(0); + expect(received).toHaveLength(3); + }); + + it('a gap emits exactly one resync and delivers the triggering event', async () => { + const client = makeClient(); + await client.connect(); + + const resyncCalls: number[] = []; + client.onResync(() => resyncCalls.push(1)); + + const received: unknown[] = []; + client.on('message.created', (_, payload) => received.push(payload)); + + // Establish baseline: seq 1, 2, 3 + lastMockWs.triggerMessage(eventWithSeq(1)); + lastMockWs.triggerMessage(eventWithSeq(2)); + lastMockWs.triggerMessage(eventWithSeq(3)); + expect(resyncCalls).toHaveLength(0); + expect(received).toHaveLength(3); + + // Gap: seq 10 arrives instead of expected seq 4 + lastMockWs.triggerMessage(eventWithSeq(10)); + expect(resyncCalls).toHaveLength(1); // one resync + expect(received).toHaveLength(4); // gap event still delivered + expect(received[3]).toEqual({ id: 10 }); + + // Back to in-order: seq 11 + lastMockWs.triggerMessage(eventWithSeq(11)); + expect(resyncCalls).toHaveLength(1); // no additional resync + expect(received).toHaveLength(5); + }); + + it('throwing resync handler does not prevent seq state advance', async () => { + vi.useFakeTimers(); + try { + const client = makeClient(); + await client.connect(); + // Clear pings + lastMockWs.sent = []; + + // Register a resync handler that throws on first call + let resyncCallCount = 0; + client.onResync(() => { + resyncCallCount++; + throw new Error('resync throws'); + }); + + // Baseline: in-order events + lastMockWs.triggerMessage(eventWithSeq(1)); + lastMockWs.triggerMessage(eventWithSeq(2)); + expect(resyncCallCount).toBe(0); + + // Trigger a gap (10 > 2 + 1), which fires the throwing resync handler. + // The seq state must still be advanced despite the throw. + expect(() => { + lastMockWs.triggerMessage(eventWithSeq(10)); + }).toThrow('resync throws'); + + expect(resyncCallCount).toBe(1); + + // Advance past the ping interval. The client should have advanced + // highestAckedSeq to 10 despite the throw. + await vi.advanceTimersByTimeAsync(15_000); + + const messages = lastMockWs.sent.map(s => { + if (s === 'ping') return s; + return JSON.parse(s) as unknown; + }); + expect(messages).toContain('ping'); + expect(messages).toContainEqual({ type: 'ack', seq: 10 }); + } finally { + vi.useRealTimers(); + } + }); + it('multiple gaps each trigger their own resync', async () => { + const client = makeClient(); + await client.connect(); + + const resyncCalls: number[] = []; + client.onResync(() => resyncCalls.push(1)); + + const received: unknown[] = []; + client.on('message.created', (_, payload) => received.push(payload)); + + lastMockWs.triggerMessage(eventWithSeq(1)); + expect(resyncCalls).toHaveLength(0); + + // First gap + lastMockWs.triggerMessage(eventWithSeq(5)); + expect(resyncCalls).toHaveLength(1); + + // Second gap + lastMockWs.triggerMessage(eventWithSeq(20)); + expect(resyncCalls).toHaveLength(2); + + expect(received).toHaveLength(3); + }); + + it('events without seq do not trigger gap logic', async () => { + const client = makeClient(); + await client.connect(); + + const resyncCalls: number[] = []; + client.onResync(() => resyncCalls.push(1)); + + const received: unknown[] = []; + client.on('message.created', (_, payload) => received.push(payload)); + + // Old-server traffic without seq + lastMockWs.triggerMessage(eventWithoutSeq()); + lastMockWs.triggerMessage(eventWithoutSeq()); + lastMockWs.triggerMessage(eventWithoutSeq()); + + expect(resyncCalls).toHaveLength(0); + expect(received).toHaveLength(3); + + // Then a seq-bearing event still works correctly + lastMockWs.triggerMessage(eventWithSeq(5)); + expect(resyncCalls).toHaveLength(1); // 5 > 0 + 1, so gap detected + expect(received).toHaveLength(4); + }); + + it('duplicate seq does not trigger resync', async () => { + const client = makeClient(); + await client.connect(); + + const resyncCalls: number[] = []; + client.onResync(() => resyncCalls.push(1)); + + const received: unknown[] = []; + client.on('message.created', (_, payload) => received.push(payload)); + + lastMockWs.triggerMessage(eventWithSeq(1)); + lastMockWs.triggerMessage(eventWithSeq(2)); + lastMockWs.triggerMessage(eventWithSeq(3)); + expect(resyncCalls).toHaveLength(0); + + // Duplicate: seq 2 again (already seen) + lastMockWs.triggerMessage(eventWithSeq(2)); + expect(resyncCalls).toHaveLength(0); // no resync for duplicate + expect(received).toHaveLength(4); // event still delivered + }); + + it('reconnect invokes resync handlers', async () => { + vi.useFakeTimers(); + try { + const client = makeClient(); + const resyncCalls: number[] = []; + client.onResync(() => resyncCalls.push(1)); + + await client.connect(); + // First connect — resync does NOT fire (not a reconnect) + expect(resyncCalls).toHaveLength(0); + + // Drop and reconnect + lastMockWs.triggerClose(); + await vi.advanceTimersByTimeAsync(2000); + await vi.advanceTimersByTimeAsync(0); + + expect(allMockWs).toHaveLength(2); + expect(resyncCalls).toHaveLength(1); + } finally { + vi.useRealTimers(); + } + }); + + it('resync unsubscribe stops further firings', async () => { + const client = makeClient(); + await client.connect(); + + const resyncCalls: number[] = []; + const off = client.onResync(() => resyncCalls.push(1)); + + // First event (no gap) + lastMockWs.triggerMessage(eventWithSeq(1)); + + // Gap triggers resync + lastMockWs.triggerMessage(eventWithSeq(10)); + expect(resyncCalls).toHaveLength(1); + + off(); + + // Another gap — no resync + lastMockWs.triggerMessage(eventWithSeq(20)); + expect(resyncCalls).toHaveLength(1); + }); + + it('onReconnect and onResync are independent', async () => { + vi.useFakeTimers(); + try { + const client = makeClient(); + const reconnectCalls: number[] = []; + const resyncCalls: number[] = []; + + client.onReconnect(() => reconnectCalls.push(1)); + client.onResync(() => resyncCalls.push(1)); + + await client.connect(); + expect(reconnectCalls).toHaveLength(0); + expect(resyncCalls).toHaveLength(0); + + // Gap triggers resync but not reconnect + lastMockWs.triggerMessage(eventWithSeq(1)); + lastMockWs.triggerMessage(eventWithSeq(10)); + expect(reconnectCalls).toHaveLength(0); + expect(resyncCalls).toHaveLength(1); + + // Reconnect triggers both + lastMockWs.triggerClose(); + await vi.advanceTimersByTimeAsync(2000); + await vi.advanceTimersByTimeAsync(0); + + expect(allMockWs).toHaveLength(2); + expect(reconnectCalls).toHaveLength(1); + expect(resyncCalls).toHaveLength(2); + } finally { + vi.useRealTimers(); + } + }); + + it('sequence state resets per socket so same seq on new socket is fine', async () => { + vi.useFakeTimers(); + try { + const client = makeClient(); + const resyncCalls: number[] = []; + client.onResync(() => resyncCalls.push(1)); + + await client.connect(); + + // First socket: receive events + lastMockWs.triggerMessage(eventWithSeq(1)); + lastMockWs.triggerMessage(eventWithSeq(2)); + expect(resyncCalls).toHaveLength(0); + + // Reconnect: sequence state reset + lastMockWs.triggerClose(); + await vi.advanceTimersByTimeAsync(2000); + await vi.advanceTimersByTimeAsync(0); + + // resync from reconnect: 1 call + expect(resyncCalls).toHaveLength(1); + + // New socket starts fresh — seq 1 is not a gap + allMockWs[1]!.triggerMessage(eventWithSeq(1)); + allMockWs[1]!.triggerMessage(eventWithSeq(2)); + expect(resyncCalls).toHaveLength(1); // no additional gap resync + } finally { + vi.useRealTimers(); + } + }); + + it('sends ack during ping interval', async () => { + vi.useFakeTimers(); + try { + const client = makeClient(); + await client.connect(); + // Clear pings from connect open + subscribe + lastMockWs.sent = []; + + // Send events with seq to populate highestAckedSeq + lastMockWs.triggerMessage(eventWithSeq(1)); + lastMockWs.triggerMessage(eventWithSeq(5)); + + // Advance past one ping interval (15s) + await vi.advanceTimersByTimeAsync(15_000); + + const messages = lastMockWs.sent.map(s => { + if (s === 'ping') return s; + return JSON.parse(s) as unknown; + }); + + expect(messages).toContain('ping'); + expect(messages).toContainEqual({ type: 'ack', seq: 5 }); + } finally { + vi.useRealTimers(); + } + }); + + it('does not send ack before any seq-bearing event', async () => { + vi.useFakeTimers(); + try { + const client = makeClient(); + await client.connect(); + // Clear pings from connect open + subscribe + lastMockWs.sent = []; + + // No seq events — highestAckedSeq stays at -1 + await vi.advanceTimersByTimeAsync(15_000); + + const messages = lastMockWs.sent.map(s => { + if (s === 'ping') return s; + return JSON.parse(s) as unknown; + }); + + expect(messages).toContain('ping'); + expect( + messages.filter(m => typeof m === 'object' && (m as { type?: string }).type === 'ack') + ).toHaveLength(0); + } finally { + vi.useRealTimers(); + } + }); + }); + + describe('acquire / release refcounting', () => { + it('two acquires then one release holds the socket open', async () => { + const client = makeClient(); + client.subscribe(['room:1']); + + await client.acquire(); + expect(client.isConnected()).toBe(true); + const ws1 = lastMockWs; + + // Second acquire: no new WebSocket, refcount is 2. + await client.acquire(); + expect(allMockWs).toHaveLength(1); + expect(client.isConnected()).toBe(true); + + // One release: refcount drops to 1, socket stays open. + client.release(); + expect(client.isConnected()).toBe(true); + expect(ws1.readyState).toBe(1); + + // Verify the socket still delivers events. + const received: unknown[] = []; + client.on('test.event', (_, payload) => received.push(payload)); + ws1.triggerMessage({ + type: 'event', + context: 'room:1', + event: 'test.event', + payload: 'hello', + }); + expect(received).toEqual(['hello']); + }); + + it('the last release closes the socket', async () => { + const client = makeClient(); + + await client.acquire(); + expect(client.isConnected()).toBe(true); + + client.release(); + expect(client.isConnected()).toBe(false); + }); + + it('next acquire after full release reconnects', async () => { + const client = makeClient(); + + await client.acquire(); + expect(client.isConnected()).toBe(true); + expect(allMockWs).toHaveLength(1); + + client.release(); + expect(client.isConnected()).toBe(false); + + await client.acquire(); + expect(client.isConnected()).toBe(true); + expect(allMockWs).toHaveLength(2); + }); + + it('extra releases are no-ops', async () => { + const client = makeClient(); + + await client.acquire(); + expect(client.isConnected()).toBe(true); + expect(allMockWs).toHaveLength(1); + + client.release(); + expect(client.isConnected()).toBe(false); + + // Extra release: no crash, no side effects. + client.release(); + expect(client.isConnected()).toBe(false); + expect(allMockWs).toHaveLength(1); + }); + + it('acquire after extra release reconnects normally', async () => { + const client = makeClient(); + + await client.acquire(); + client.release(); + client.release(); // extra + + await client.acquire(); + expect(client.isConnected()).toBe(true); + expect(allMockWs).toHaveLength(2); + }); + + it('reacquire after release triggers resync handlers', async () => { + const client = makeClient(); + const resyncCalls: number[] = []; + client.onResync(() => resyncCalls.push(1)); + + // First acquire: connect() does NOT fire resync (not a reconnect). + await client.acquire(); + expect(client.isConnected()).toBe(true); + expect(resyncCalls).toHaveLength(0); + + // Release closes the socket. + client.release(); + expect(client.isConnected()).toBe(false); + + // Reacquire: connect() fires resync on reconnect. + await client.acquire(); + expect(client.isConnected()).toBe(true); + expect(resyncCalls).toHaveLength(1); + }); + }); + + describe('fetchConnectionTicket deadline', () => { + it('times out after CONTROL_PLANE_DEADLINE_MS when the ticket endpoint hangs', async () => { + vi.useFakeTimers(); + try { + // Override the global fetch to never resolve — simulates a hanging + // /connect-ticket endpoint. The mock respects the abort signal so + // the withDeadline timeout actually propagates. + vi.stubGlobal( + 'fetch', + vi.fn((_url, init?: RequestInit) => { + return new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal) { + if (signal.aborted) { + reject(signal.reason); + return; + } + signal.addEventListener('abort', () => reject(signal.reason)); + } + }); + }) + ); + + const client = makeClient(); + const connectPromise = client.connect(); + + // Advance past the 15s control-plane deadline. + await vi.advanceTimersByTimeAsync(15_001); + + // connect() catches the error and schedules a reconnect. The promise + // resolves (connect doesn't propagate the error). + await connectPromise; + + // No WebSocket was created because the ticket fetch never returned. + expect(allMockWs).toHaveLength(0); + expect(client.isConnected()).toBe(false); + + // A reconnect should be scheduled. + await vi.advanceTimersByTimeAsync(2_000); + expect(fetch).toHaveBeenCalledTimes(2); // first call timed out, second is reconnect + } finally { + vi.useRealTimers(); + } + }); + + it('ticket fetch rejects on deadline, handshake timer stays independent', async () => { + vi.useFakeTimers(); + try { + // Fetch hangs forever (with signal respect for deadline). + vi.stubGlobal( + 'fetch', + vi.fn((_url, init?: RequestInit) => { + return new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal) { + if (signal.aborted) { + reject(signal.reason); + return; + } + signal.addEventListener('abort', () => reject(signal.reason)); + } + }); + }) + ); + + const client = makeClient(); + const connectPromise = client.connect(); + + // Advance to just before the 15s deadline. + await vi.advanceTimersByTimeAsync(14_000); + + // Ticket has not yet timed out — no WebSocket created. + expect(allMockWs).toHaveLength(0); + + // Advance past the 15s deadline. + await vi.advanceTimersByTimeAsync(2_000); + await connectPromise; + + // Still no WebSocket — ticket fetch timed out before WS creation. + expect(allMockWs).toHaveLength(0); + expect(client.isConnected()).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it('successful ticket fetch creates WebSocket and clears deadline', async () => { + vi.useFakeTimers(); + try { + // Fetch resolves just under the 15s deadline. + vi.stubGlobal( + 'fetch', + vi.fn( + () => + new Promise(resolve => { + setTimeout(() => { + resolve( + new Response(JSON.stringify({ ticket: 'fast-ticket' }), { + headers: { 'content-type': 'application/json' }, + }) + ); + }, 5_000); + }) + ) + ); + + const client = makeClient(); + client.subscribe(['room:1']); + const connectPromise = client.connect(); + + // Advance past the 5s fetch resolve. + await vi.advanceTimersByTimeAsync(6_000); + await connectPromise; + + // WebSocket was created and opened successfully. + expect(allMockWs.length).toBeGreaterThanOrEqual(1); + expect(client.isConnected()).toBe(true); + expect(fetch).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + }); }); diff --git a/packages/event-service/src/__tests__/deadline.test.ts b/packages/event-service/src/__tests__/deadline.test.ts new file mode 100644 index 00000000000..014a25eb2f9 --- /dev/null +++ b/packages/event-service/src/__tests__/deadline.test.ts @@ -0,0 +1,308 @@ +import { describe, it, expect, vi } from 'vitest'; +import { + RequestDeadlineError, + CONTROL_PLANE_DEADLINE_MS, + SEND_DEADLINE_MS, + withDeadline, +} from '../deadline'; + +/** + * Hermes (React Native) whatwg-fetch rejects with a generic `AbortError` + * DOMException after the signal aborts, discarding `signal.reason`. + * This constant simulates that behavior in tests. + */ +const HERMES_ABORT_ERROR = new DOMException('Aborted', 'AbortError'); + +describe('withDeadline', () => { + it('resolves with the return value when fn completes before the deadline', async () => { + const result = await withDeadline(5_000, async () => 'ok'); + expect(result).toBe('ok'); + }); + + it('rejects with RequestDeadlineError when the deadline expires', async () => { + vi.useFakeTimers(); + try { + const promise = withDeadline(100, signal => { + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => reject(signal.reason)); + }); + }); + + vi.advanceTimersByTime(200); + vi.runAllTicks(); + + await expect(promise).rejects.toThrow(RequestDeadlineError); + await expect(promise).rejects.toThrow('timed out after 100ms'); + } finally { + vi.useRealTimers(); + } + }); + + it('forwards the signal to fn', async () => { + const signalRef: AbortSignal[] = []; + + await withDeadline(5_000, signal => { + signalRef.push(signal); + return Promise.resolve('ok'); + }); + + expect(signalRef).toHaveLength(1); + expect(signalRef[0]).toBeInstanceOf(AbortSignal); + expect(signalRef[0].aborted).toBe(false); + }); + + it('clears the timer after success (no late aborts)', async () => { + vi.useFakeTimers(); + + const fn = vi.fn(async () => 'ok'); + const promise = withDeadline(10_000, fn); + + await vi.advanceTimersByTimeAsync(0); + const result = await promise; + + expect(result).toBe('ok'); + + // Advance past the deadline — the timer must be cleared, so fn is not + // called again and no unhandled rejection is emitted. + await vi.advanceTimersByTimeAsync(20_000); + expect(fn).toHaveBeenCalledTimes(1); + + vi.useRealTimers(); + }); + + it('caller abort wins over the deadline', async () => { + vi.useFakeTimers(); + try { + const callerController = new AbortController(); + + const promise = withDeadline( + 60_000, + signal => { + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => reject(signal.reason)); + }); + }, + callerController.signal + ); + + // Abort from the caller side before either settles. + callerController.abort(new Error('caller cancelled')); + + vi.runAllTicks(); + + await expect(promise).rejects.toThrow('caller cancelled'); + } finally { + vi.useRealTimers(); + } + }); + + it('already-aborted caller signal rejects immediately', async () => { + const callerController = new AbortController(); + callerController.abort(new Error('already aborted')); + + const fn = vi.fn(async () => 'never'); + + await expect(withDeadline(60_000, fn, callerController.signal)).rejects.toThrow( + 'already aborted' + ); + + expect(fn).not.toHaveBeenCalled(); + }); + + it('removes the caller abort listener after fn settles', async () => { + const callerController = new AbortController(); + const removeSpy = vi.spyOn(callerController.signal, 'removeEventListener'); + + await withDeadline(5_000, async () => 'ok', callerController.signal); + + expect(removeSpy).toHaveBeenCalledWith('abort', expect.any(Function)); + }); + + it('does not add abort listener when no caller signal is provided', async () => { + const addSpy = vi.spyOn(AbortSignal.prototype, 'addEventListener'); + const removeSpy = vi.spyOn(AbortSignal.prototype, 'removeEventListener'); + + await withDeadline(5_000, async () => 'ok'); + + // Neither addEventListener nor removeEventListener should be called + // because there is no caller signal. + expect(addSpy).not.toHaveBeenCalled(); + expect(removeSpy).not.toHaveBeenCalled(); + }); +}); + +describe('Hermes AbortError resilience', () => { + it('rejects with RequestDeadlineError when fn rejects generic AbortError after deadline', async () => { + vi.useFakeTimers(); + try { + const promise = withDeadline(100, signal => { + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => { + // Hermes whatwg-fetch: discards signal.reason, rejects with + // generic DOMException('Aborted', 'AbortError'). + reject(HERMES_ABORT_ERROR); + }); + }); + }); + + vi.advanceTimersByTime(200); + vi.runAllTicks(); + + await expect(promise).rejects.toThrow(RequestDeadlineError); + await expect(promise).rejects.toThrow('timed out after 100ms'); + } finally { + vi.useRealTimers(); + } + }); + + it('preserves caller abort reason when fn rejects generic AbortError', async () => { + const callerController = new AbortController(); + + const promise = withDeadline( + 60_000, + signal => { + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => { + // Hermes whatwg-fetch: discards signal.reason. + reject(HERMES_ABORT_ERROR); + }); + }); + }, + callerController.signal + ); + + callerController.abort(new Error('caller cancelled')); + + await expect(promise).rejects.toThrow('caller cancelled'); + // Must not be the generic AbortError. + await expect(promise).rejects.not.toThrow('Aborted'); + }); + + it('does not leak unhandled rejection when inner promise rejects after outer settles', async () => { + vi.useFakeTimers(); + try { + // The inner promise (simulating Hermes whatwg-fetch) rejects with a + // generic AbortError after the deadline timer fires. The settle gate + // must suppress this rejection so it does not propagate to the outer + // promise, which already rejected with RequestDeadlineError. + const promise = withDeadline(100, signal => { + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => { + reject(HERMES_ABORT_ERROR); + }); + }); + }); + + // Timer fires, outer rejects with RequestDeadlineError. + vi.advanceTimersByTime(200); + vi.runAllTicks(); + + // The outer promise must have rejected with the deadline error, + // not the generic AbortError. + await expect(promise).rejects.toThrow(RequestDeadlineError); + await expect(promise).rejects.not.toThrow('Aborted'); + } finally { + vi.useRealTimers(); + } + }); + + it('does not leak timer or listener when fn rejects generic AbortError after deadline', async () => { + vi.useFakeTimers(); + try { + const clearTimeoutSpy = vi.spyOn(globalThis, 'clearTimeout'); + + const promise = withDeadline(100, signal => { + return new Promise((_resolve, reject) => { + signal.addEventListener('abort', () => { + reject(HERMES_ABORT_ERROR); + }); + }); + }); + + vi.advanceTimersByTime(200); + vi.runAllTicks(); + + await expect(promise).rejects.toThrow(RequestDeadlineError); + + // clearTimeout must have been called (even though the timer already + // fired, the finally block still clears it for safety). + expect(clearTimeoutSpy).toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + it('successful fn still resolves normally alongside Hermes sim', async () => { + // Regression: normal resolution must work even with the new + // independent-settlement pattern. + const result = await withDeadline(5_000, async () => { + // Return a value — no abort involved. + return 42; + }); + + expect(result).toBe(42); + }); +}); + +describe('RequestDeadlineError', () => { + it('is an instance of Error', () => { + const err = new RequestDeadlineError(10_000); + expect(err).toBeInstanceOf(Error); + }); + + it('has name RequestDeadlineError', () => { + const err = new RequestDeadlineError(10_000); + expect(err.name).toBe('RequestDeadlineError'); + }); + + it('includes the deadline in its message', () => { + const err = new RequestDeadlineError(15_000); + expect(err.message).toContain('15000'); + }); +}); + +describe('CONTROL_PLANE_DEADLINE_MS', () => { + it('is 15 seconds', () => { + expect(CONTROL_PLANE_DEADLINE_MS).toBe(15_000); + }); +}); + +describe('SEND_DEADLINE_MS', () => { + it('is 30 seconds', () => { + expect(SEND_DEADLINE_MS).toBe(30_000); + }); +}); + +describe('forbidden APIs', () => { + it('never calls AbortSignal.timeout or AbortSignal.any', async () => { + // withDeadline must use manual AbortController + setTimeout, not the + // newer AbortSignal.timeout / AbortSignal.any static methods, because + // those are absent on Hermes (React Native iOS/Android). + const timeoutSpy = vi.spyOn(AbortSignal, 'timeout'); + const anySpy = vi.spyOn(AbortSignal, 'any'); + + try { + await withDeadline(5_000, async () => 'ok'); + expect(timeoutSpy).not.toHaveBeenCalled(); + expect(anySpy).not.toHaveBeenCalled(); + } finally { + timeoutSpy.mockRestore(); + anySpy.mockRestore(); + } + }); + + it('never calls AbortSignal.timeout or AbortSignal.any with a caller signal', async () => { + const callerController = new AbortController(); + const timeoutSpy = vi.spyOn(AbortSignal, 'timeout'); + const anySpy = vi.spyOn(AbortSignal, 'any'); + + try { + await withDeadline(5_000, async () => 'ok', callerController.signal); + expect(timeoutSpy).not.toHaveBeenCalled(); + expect(anySpy).not.toHaveBeenCalled(); + } finally { + timeoutSpy.mockRestore(); + anySpy.mockRestore(); + } + }); +}); diff --git a/packages/event-service/src/__tests__/schemas.test.ts b/packages/event-service/src/__tests__/schemas.test.ts new file mode 100644 index 00000000000..d8fd63b9062 --- /dev/null +++ b/packages/event-service/src/__tests__/schemas.test.ts @@ -0,0 +1,198 @@ +import { describe, it, expect } from 'vitest'; +import { + eventMessageSchema, + errorMessageSchema, + ackMessageSchema, + clientMessageSchema, + serverMessageSchema, +} from '../schemas'; + +describe('eventMessageSchema', () => { + it('accepts an event with a seq field', () => { + const result = eventMessageSchema.safeParse({ + type: 'event', + context: 'project:abc', + event: 'task.created', + payload: { taskId: '1' }, + seq: 5, + }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.seq).toBe(5); + } + }); + + it('accepts an event without a seq field (backward compat)', () => { + const result = eventMessageSchema.safeParse({ + type: 'event', + context: 'project:abc', + event: 'task.created', + payload: {}, + }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.seq).toBeUndefined(); + } + }); + + it('rejects an event with a negative seq', () => { + const result = eventMessageSchema.safeParse({ + type: 'event', + context: 'project:abc', + event: 'task.created', + payload: {}, + seq: -1, + }); + expect(result.success).toBe(false); + }); + + it('rejects an event with a non-integer seq', () => { + const result = eventMessageSchema.safeParse({ + type: 'event', + context: 'project:abc', + event: 'task.created', + payload: {}, + seq: 1.5, + }); + expect(result.success).toBe(false); + }); + + it('rejects an event with a string seq', () => { + const result = eventMessageSchema.safeParse({ + type: 'event', + context: 'project:abc', + event: 'task.created', + payload: {}, + seq: '5', + }); + expect(result.success).toBe(false); + }); + + it('accepts an event with seq 0', () => { + const result = eventMessageSchema.safeParse({ + type: 'event', + context: 'project:abc', + event: 'task.created', + payload: {}, + seq: 0, + }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.seq).toBe(0); + } + }); +}); + +describe('errorMessageSchema', () => { + it('accepts an error with a seq field', () => { + const result = errorMessageSchema.safeParse({ + type: 'error', + code: 'too_many_contexts', + max: 200, + seq: 10, + }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.seq).toBe(10); + } + }); + + it('accepts an error without a seq field (backward compat)', () => { + const result = errorMessageSchema.safeParse({ + type: 'error', + code: 'too_many_contexts', + max: 200, + }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.seq).toBeUndefined(); + } + }); +}); + +describe('ackMessageSchema', () => { + it('accepts a valid ack message', () => { + const result = ackMessageSchema.safeParse({ type: 'ack', seq: 42 }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.seq).toBe(42); + } + }); + + it('accepts an ack with seq 0', () => { + const result = ackMessageSchema.safeParse({ type: 'ack', seq: 0 }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.seq).toBe(0); + } + }); + + it('rejects an ack with a negative seq', () => { + const result = ackMessageSchema.safeParse({ type: 'ack', seq: -1 }); + expect(result.success).toBe(false); + }); + + it('rejects an ack with a non-integer seq', () => { + const result = ackMessageSchema.safeParse({ type: 'ack', seq: 1.5 }); + expect(result.success).toBe(false); + }); + + it('rejects an ack with a missing seq', () => { + const result = ackMessageSchema.safeParse({ type: 'ack' }); + expect(result.success).toBe(false); + }); +}); + +describe('clientMessageSchema (ack in discriminant union)', () => { + it('parses an ack message as a client message', () => { + const result = clientMessageSchema.safeParse({ type: 'ack', seq: 7 }); + expect(result.success).toBe(true); + if (result.success && result.data.type === 'ack') { + expect(result.data.seq).toBe(7); + } + }); + + it('still parses context.subscribe', () => { + const result = clientMessageSchema.safeParse({ + type: 'context.subscribe', + contexts: ['ctx:a'], + }); + expect(result.success).toBe(true); + }); + + it('still parses context.unsubscribe', () => { + const result = clientMessageSchema.safeParse({ + type: 'context.unsubscribe', + contexts: ['ctx:a'], + }); + expect(result.success).toBe(true); + }); + + it('rejects an unknown message type', () => { + const result = clientMessageSchema.safeParse({ type: 'unknown' }); + expect(result.success).toBe(false); + }); +}); + +describe('serverMessageSchema', () => { + it('parses an event as a server message', () => { + const result = serverMessageSchema.safeParse({ + type: 'event', + context: 'project:a', + event: 'x', + payload: null, + seq: 1, + }); + expect(result.success).toBe(true); + }); + + it('parses an error as a server message', () => { + const result = serverMessageSchema.safeParse({ + type: 'error', + code: 'too_many_contexts', + max: 200, + seq: 1, + }); + expect(result.success).toBe(true); + }); +}); diff --git a/packages/event-service/src/client.ts b/packages/event-service/src/client.ts index 14230344640..e009b332c4a 100644 --- a/packages/event-service/src/client.ts +++ b/packages/event-service/src/client.ts @@ -1,5 +1,6 @@ import type { ClientMessage, ConnectTicketQuery, EventServiceConfig } from './types'; import { connectTicketResponseSchema, serverMessageSchema } from './schemas'; +import { CONTROL_PLANE_DEADLINE_MS, withDeadline } from './deadline'; const WEBSOCKET_PROTOCOL = 'kilo.events.v1'; @@ -53,19 +54,22 @@ function connectUrlFor(wsBase: string, ticket: string): string { } async function fetchConnectionTicket(wsBase: string, token: string): Promise { - const response = await fetch(ticketEndpointFor(wsBase), { - method: 'POST', - headers: { authorization: `Bearer ${token}` }, - }); - if (response.status === 401 || response.status === 403) { - throw new WebSocketAuthError(); - } - if (!response.ok) { - throw new Error('Failed to mint WebSocket connection ticket'); - } + return withDeadline(CONTROL_PLANE_DEADLINE_MS, async signal => { + const response = await fetch(ticketEndpointFor(wsBase), { + method: 'POST', + headers: { authorization: `Bearer ${token}` }, + signal, + }); + if (response.status === 401 || response.status === 403) { + throw new WebSocketAuthError(); + } + if (!response.ok) { + throw new Error('Failed to mint WebSocket connection ticket'); + } - const body = connectTicketResponseSchema.parse(await response.json()); - return body.ticket; + const body = connectTicketResponseSchema.parse(await response.json()); + return body.ticket; + }); } export class EventServiceClient { @@ -75,6 +79,7 @@ export class EventServiceClient { private ws: WebSocket | null = null; private connected = false; + private generation = 0; private eventHandlers = new Map void>>(); // Refcounted so multiple consumers can independently subscribe to and // unsubscribe from the same context without trampling each other. The wire @@ -88,6 +93,19 @@ export class EventServiceClient { private hasConnectedBefore = false; private reconnectHandlers = new Set<() => void>(); private connectedHandlers = new Set<() => void>(); + private resyncHandlers = new Set<() => void>(); + // Acquire/release lifecycle refcount. Multiple consumers can independently + // hold the socket open. The 0→1 transition calls connect(); the 1→0 + // transition calls disconnect(). Acquire after release reconnects by + // clearing destroyed through connect() (which resets the flag on entry). + private acquireCount = 0; + // Per-socket sequence state. The server emits seq starting at 1. + // Initialize to 0 so the first in-order seq:1 event does not trigger + // a gap (1 > 0 + 1 is false). + private lastKnownSeq = 0; + // Highest acknowledged seq; sent as an { type: 'ack' } message during the + // regular ping interval, not for every event. + private highestAckedSeq = -1; private pingTimer: ReturnType | null = null; private handshakeTimer: ReturnType | null = null; private abortHandshake: ((err: Error) => void) | null = null; @@ -99,6 +117,7 @@ export class EventServiceClient { } async connect(): Promise { + this.generation++; this.destroyed = false; this.reconnectAttempts = 0; this.authRecoveryAttempts = 0; @@ -146,6 +165,8 @@ export class EventServiceClient { } private async connectOnce(): Promise { + const gen = this.generation; + // Close any existing socket to avoid leaking connections. if (this.ws) { const oldWs = this.ws; @@ -154,12 +175,12 @@ export class EventServiceClient { } const token = await this.getToken(); - if (this.destroyed) return; + if (gen !== this.generation || this.destroyed) return; const ticket = await fetchConnectionTicket(this.url, token); // disconnect() may have run while we were awaiting the token. Bail before // creating the socket so we don't leak a WebSocket + ping timer past // provider unmount (e.g. sign-out, navigation, strict-mode remount). - if (this.destroyed) return; + if (gen !== this.generation || this.destroyed) return; return new Promise((resolve, reject) => { const ws = new WebSocket(connectUrlFor(this.url, ticket), [WEBSOCKET_PROTOCOL]); @@ -197,16 +218,23 @@ export class EventServiceClient { }, HANDSHAKE_TIMEOUT_MS); ws.addEventListener('open', () => { + if (this.ws !== ws) return; const isReconnect = this.hasConnectedBefore; this.connected = true; this.hasConnectedBefore = true; this.reconnectAttempts = 0; this.authRecoveryAttempts = 0; + // Reset per-socket sequence state on every new WebSocket. + this.lastKnownSeq = 0; + this.highestAckedSeq = -1; this.resubscribeContexts(); if (isReconnect) { for (const handler of this.reconnectHandlers) { handler(); } + for (const handler of this.resyncHandlers) { + handler(); + } } for (const handler of this.connectedHandlers) { handler(); @@ -216,6 +244,7 @@ export class EventServiceClient { }); ws.addEventListener('message', (event: MessageEvent) => { + if (this.ws !== ws) return; this.handleMessage(event.data as string); }); @@ -250,6 +279,7 @@ export class EventServiceClient { } disconnect(): void { + this.generation++; this.destroyed = true; if (this.reconnectTimer !== null) { clearTimeout(this.reconnectTimer); @@ -269,6 +299,33 @@ export class EventServiceClient { this.connected = false; } + /** + * Holds the socket open for this caller. The zero-to-one transition calls + * connect(); every subsequent call increments an internal refcount without + * creating a second socket. Call release() when the caller no longer needs + * the connection. Callers must pair every acquire with exactly one release. + */ + async acquire(): Promise { + this.acquireCount++; + if (this.acquireCount === 1) { + await this.connect(); + } + } + + /** + * Releases one refcount. The one-to-zero transition calls disconnect() and + * tears down the socket. Safe to call more times than acquire (extra calls + * are no-ops). After the last release, the next acquire reconnects because + * connect() clears the destroyed flag. + */ + release(): void { + if (this.acquireCount <= 0) return; + this.acquireCount--; + if (this.acquireCount === 0) { + this.disconnect(); + } + } + isConnected(): boolean { return this.connected && this.ws !== null && this.ws.readyState === WebSocket.OPEN; } @@ -337,6 +394,21 @@ export class EventServiceClient { }; } + /** + * Registers a handler that fires when the connection state may be stale: + * on every reconnect and on every detected sequence-number gap in server + * events. Consumers use this to refetch or invalidate caches that may + * have missed events while the socket was disconnected or out of sync. + * + * Returns an unsubscribe function. Calling it removes the handler. + */ + onResync(handler: () => void): () => void { + this.resyncHandlers.add(handler); + return () => { + this.resyncHandlers.delete(handler); + }; + } + on(event: string, handler: (context: string, payload: unknown) => void): () => void { let handlers = this.eventHandlers.get(event); if (!handlers) { @@ -371,6 +443,26 @@ export class EventServiceClient { if (!result.success) return; const message = result.data; + // Sequence gap detection: fire resync once when a seq skips ahead of + // the expected next value, then continue from the received seq. Events + // without seq are passed through unchanged (old-server compatibility). + // Advance lastKnownSeq and highestAckedSeq *before* invoking resync + // handlers so a throwing handler cannot stop state progress. + if (message.seq !== undefined) { + const oldLastKnownSeq = this.lastKnownSeq; + if (message.seq > this.lastKnownSeq) { + this.lastKnownSeq = message.seq; + } + if (message.seq > this.highestAckedSeq) { + this.highestAckedSeq = message.seq; + } + if (message.seq > oldLastKnownSeq + 1) { + for (const handler of this.resyncHandlers) { + handler(); + } + } + } + if (message.type === 'event') { const handlers = this.eventHandlers.get(message.event); if (handlers) { @@ -394,6 +486,13 @@ export class EventServiceClient { this.pingTimer = setInterval(() => { if (this.ws && this.ws.readyState === WebSocket.OPEN) { this.ws.send('ping'); + if (this.highestAckedSeq >= 0) { + const ack = { + type: 'ack', + seq: this.highestAckedSeq, + } satisfies ClientMessage; + this.send(ack); + } } }, 15000); } @@ -422,6 +521,7 @@ export class EventServiceClient { this.reconnectAttempts++; this.reconnectTimer = setTimeout(() => { this.reconnectTimer = null; + this.generation++; this.connectOnce().catch(err => { void this.handleReconnectFailure(err); }); diff --git a/packages/event-service/src/deadline.ts b/packages/event-service/src/deadline.ts new file mode 100644 index 00000000000..e7b387eacc9 --- /dev/null +++ b/packages/event-service/src/deadline.ts @@ -0,0 +1,114 @@ +/** + * Thrown when a {@link withDeadline}-wrapped request exceeds its time budget. + * The caller must treat the outcome as uncertain: the server may have accepted + * the request after the client gave up. Callers must reconcile before retry. + */ +export class RequestDeadlineError extends Error { + constructor(deadlineMs: number) { + super(`Request timed out after ${deadlineMs}ms`); + this.name = 'RequestDeadlineError'; + } +} + +/** + * Deadline for control-plane requests: tRPC, Kilo Chat non-send HTTP calls, + * and Event Service connection-ticket fetches. + */ +export const CONTROL_PLANE_DEADLINE_MS = 15_000; + +/** + * Deadline for Kilo Chat message send (`POST /v1/messages`) only. + * Longer than the control-plane deadline because the server may be + * busy with agent processing and the outcome is always uncertain + * for timed-out sends. + */ +export const SEND_DEADLINE_MS = 30_000; + +/** + * Run `fn` with a deadline. If the deadline expires before `fn` settles, + * the returned promise rejects with a {@link RequestDeadlineError}. + * + * This helper is safe for Hermes (React Native). It uses `setTimeout` and + * manual `AbortController` forwarding; it never calls `AbortSignal.timeout` + * or `AbortSignal.any`, which are absent on Hermes. + * + * When `callerSignal` is provided, an abort on the caller signal forwards + * to the deadline controller. Caller abort wins: the returned promise + * rejects with `callerSignal.reason` instead of a deadline error. + * + * The outer promise settles independently of the inner `fn` result. On + * Hermes, `whatwg-fetch` rejects with a generic `AbortError` after abort, + * discarding `signal.reason`. This helper preserves the correct error type + * (deadline or caller reason) regardless of what `fn` rejects. + * + * Cleanup (timer clear, listener removal) runs in a `finally` block so + * the deadline is always disarmed after the promise settles. + */ +export function withDeadline( + timeoutMs: number, + fn: (signal: AbortSignal) => Promise, + callerSignal?: AbortSignal +): Promise { + const controller = new AbortController(); + let settled = false; + let timer!: ReturnType; + let onCallerAbort: (() => void) | undefined; + + return new Promise((resolve, reject) => { + timer = setTimeout(() => { + // Settle the outer promise first, then abort the inner controller. + // This preserves the deadline error even when fn's fetch discards + // signal.reason (Hermes whatwg-fetch). + if (!settled) { + settled = true; + reject(new RequestDeadlineError(timeoutMs)); + } + controller.abort(new RequestDeadlineError(timeoutMs)); + }, timeoutMs); + + onCallerAbort = () => { + // Settle the outer promise first, then abort the inner controller. + // This preserves the caller's reason even when fn's fetch discards + // signal.reason (Hermes whatwg-fetch). + if (!settled) { + settled = true; + reject(callerSignal?.reason); + } + controller.abort(callerSignal?.reason); + }; + + if (callerSignal) { + if (callerSignal.aborted) { + clearTimeout(timer); + // Already-aborted: reject immediately, never call fn. + reject(callerSignal.reason); + settled = true; + controller.abort(callerSignal.reason); + return; + } + callerSignal.addEventListener('abort', onCallerAbort); + } + + fn(controller.signal).then( + value => { + if (!settled) { + settled = true; + resolve(value); + } + }, + reason => { + // Forward inner rejection only when neither timer nor caller abort + // has settled the outer promise yet. + if (!settled) { + settled = true; + reject(reason); + } + } + ); + }).finally(() => { + clearTimeout(timer); + if (callerSignal && onCallerAbort) { + callerSignal.removeEventListener('abort', onCallerAbort); + } + }); +} diff --git a/packages/event-service/src/index.ts b/packages/event-service/src/index.ts index c2e238f1dd2..7ef81b4f315 100644 --- a/packages/event-service/src/index.ts +++ b/packages/event-service/src/index.ts @@ -1,4 +1,10 @@ export { EventServiceClient, WebSocketAuthError, HandshakeTimeoutError } from './client'; +export { + RequestDeadlineError, + CONTROL_PLANE_DEADLINE_MS, + SEND_DEADLINE_MS, + withDeadline, +} from './deadline'; export * from './presence'; export * from './kiloclaw-contexts'; export * from './schemas'; diff --git a/packages/event-service/src/schemas.ts b/packages/event-service/src/schemas.ts index a3676d09b1b..4193e7d3940 100644 --- a/packages/event-service/src/schemas.ts +++ b/packages/event-service/src/schemas.ts @@ -16,9 +16,15 @@ export const contextUnsubscribeMessageSchema = z.object({ contexts: z.array(contextSchema).max(MAX_CONTEXTS), }); +export const ackMessageSchema = z.object({ + type: z.literal('ack'), + seq: z.number().int().nonnegative(), +}); + export const clientMessageSchema = z.discriminatedUnion('type', [ contextSubscribeMessageSchema, contextUnsubscribeMessageSchema, + ackMessageSchema, ]); // ── Server → Client ──────────────────────────────────────────────── @@ -28,12 +34,14 @@ export const eventMessageSchema = z.object({ context: z.string(), event: z.string(), payload: z.unknown(), + seq: z.number().int().nonnegative().optional(), }); export const errorMessageSchema = z.object({ type: z.literal('error'), code: z.enum(['too_many_contexts']), max: z.number().int().positive(), + seq: z.number().int().nonnegative().optional(), }); export const serverMessageSchema = z.discriminatedUnion('type', [ diff --git a/packages/event-service/src/types.ts b/packages/event-service/src/types.ts index 3bb378ca3cd..397e71a14fa 100644 --- a/packages/event-service/src/types.ts +++ b/packages/event-service/src/types.ts @@ -2,6 +2,7 @@ import type { z } from 'zod'; import type { contextSubscribeMessageSchema, contextUnsubscribeMessageSchema, + ackMessageSchema, clientMessageSchema, errorMessageSchema, eventMessageSchema, @@ -14,6 +15,7 @@ import type { export type ContextSubscribeMessage = z.infer; export type ContextUnsubscribeMessage = z.infer; +export type AckMessage = z.infer; export type ClientMessage = z.infer; // ── Server → Client ──────────────────────────────────────────────── diff --git a/packages/kilo-chat-hooks/src/use-bot-status.test.ts b/packages/kilo-chat-hooks/src/use-bot-status.test.ts index 10c9f5e945f..928f89f39f7 100644 --- a/packages/kilo-chat-hooks/src/use-bot-status.test.ts +++ b/packages/kilo-chat-hooks/src/use-bot-status.test.ts @@ -18,7 +18,12 @@ const testState = vi.hoisted(() => ({ // tanstack query mocks queryData: new Map(), invalidateCalls: [] as Array<{ queryKey: unknown }>, - queryCalls: [] as Array<{ queryKey: unknown; queryFn: () => Promise; enabled: boolean }>, + queryCalls: [] as Array<{ + queryKey: unknown; + queryFn: () => Promise; + enabled: boolean; + refetchInterval?: number | false; + }>, })); /** Call before each simulated render to reset the slot cursor. */ @@ -62,7 +67,12 @@ vi.mock('@tanstack/react-query', () => ({ testState.invalidateCalls.push(opts); }, }), - useQuery: (opts: { queryKey: unknown; queryFn: () => Promise; enabled: boolean }) => { + useQuery: (opts: { + queryKey: unknown; + queryFn: () => Promise; + enabled: boolean; + refetchInterval?: number | false; + }) => { testState.queryCalls.push(opts); const key = JSON.stringify(opts.queryKey); return { data: testState.queryData.get(key) }; @@ -91,16 +101,16 @@ function record(overrides: Partial = {}): BotStatusRecord { } function makeClients({ connected = false } = {}) { - const capturedOnConnected: Array<() => void> = []; + const capturedOnResync: Array<() => void> = []; const capturedBotStatus: Array<(ctx: string, ev: BotStatusEvent) => void> = []; const eventClient = { - onConnected: vi.fn((handler: () => void) => { - capturedOnConnected.push(handler); + onResync: vi.fn((handler: () => void) => { + capturedOnResync.push(handler); if (connected) handler(); return () => { - const i = capturedOnConnected.indexOf(handler); - if (i >= 0) capturedOnConnected.splice(i, 1); + const i = capturedOnResync.indexOf(handler); + if (i >= 0) capturedOnResync.splice(i, 1); }; }), }; @@ -123,7 +133,7 @@ function makeClients({ connected = false } = {}) { return { eventClient: eventClient as unknown as EventServiceClient, kiloChatClient: kiloChatClient as unknown as KiloChatClient, - fireConnected: () => capturedOnConnected.forEach(h => h()), + fireResync: () => capturedOnResync.forEach(h => h()), fireBotStatus: (ev: BotStatusEvent) => capturedBotStatus.forEach(h => h('ctx', ev)), mockRequestBotStatus, }; @@ -210,36 +220,15 @@ describe('useBotStatus', () => { } }); - it('query is disabled until onConnected fires (wsReady gate)', () => { + it('query is enabled when sandboxId is set (no wsReady gate)', () => { const { eventClient, kiloChatClient } = makeClients({ connected: false }); beginRender(); useBotStatus(kiloChatClient, eventClient, 'sb-1'); - // wsReady starts as false — query should be disabled + // Query should be enabled even before WS connects const queryCall = testState.queryCalls[0]; expect(queryCall).toBeDefined(); - expect(queryCall?.enabled).toBe(false); - }); - - it('query becomes enabled once WS connects', () => { - const { eventClient, kiloChatClient } = makeClients({ connected: true }); - - beginRender(); - useBotStatus(kiloChatClient, eventClient, 'sb-1'); - - // onConnected fires synchronously when connected=true, which calls setWsReady(true). - // In real React this triggers a re-render; verify the slot was updated. - expect(testState.stateSlots[0]).toBe(true); - - // Simulate the re-render: reset cursor + effects/queries for the second pass. - testState.queryCalls = []; - // (Don't reset cleanups — effects shouldn't re-run on a state-only re-render.) - - beginRender(); - useBotStatus(kiloChatClient, eventClient, 'sb-1'); - - const queryCall = testState.queryCalls[0]; expect(queryCall?.enabled).toBe(true); }); @@ -352,15 +341,15 @@ describe('useBotStatus', () => { expect(testState.queryData.get(key)).toBe(existing); }); - it('reconnect triggers invalidateQueries for the sandbox', () => { - const { eventClient, kiloChatClient, fireConnected } = makeClients({ connected: false }); + it('resync triggers invalidateQueries for the sandbox', () => { + const { eventClient, kiloChatClient, fireResync } = makeClients({ connected: false }); beginRender(); useBotStatus(kiloChatClient, eventClient, 'sb-1'); expect(testState.invalidateCalls).toHaveLength(0); - fireConnected(); + fireResync(); expect(testState.invalidateCalls).toHaveLength(1); expect(testState.invalidateCalls[0]?.queryKey).toEqual(botStatusKey('sb-1')); @@ -373,7 +362,28 @@ describe('useBotStatus', () => { useBotStatus(kiloChatClient, eventClient, null); const queryCall = testState.queryCalls[0]; - // enabled = sandboxId !== null && wsReady → false expect(queryCall?.enabled).toBe(false); }); + + it('disables polling when active is false', () => { + const { eventClient, kiloChatClient } = makeClients({ connected: false }); + + beginRender(); + useBotStatus(kiloChatClient, eventClient, 'sb-1', false); + + const queryCall = testState.queryCalls[0]; + expect(queryCall?.enabled).toBe(true); + expect(queryCall?.refetchInterval).toBeUndefined(); + }); + + it('enables polling by default (active defaults to true)', () => { + const { eventClient, kiloChatClient } = makeClients({ connected: false }); + + beginRender(); + useBotStatus(kiloChatClient, eventClient, 'sb-1'); + + const queryCall = testState.queryCalls[0]; + expect(queryCall?.enabled).toBe(true); + expect(queryCall?.refetchInterval).toBe(15_000); + }); }); diff --git a/packages/kilo-chat-hooks/src/use-bot-status.ts b/packages/kilo-chat-hooks/src/use-bot-status.ts index 76ae71ff2c8..238601691a6 100644 --- a/packages/kilo-chat-hooks/src/use-bot-status.ts +++ b/packages/kilo-chat-hooks/src/use-bot-status.ts @@ -1,4 +1,4 @@ -import { useEffect, useState } from 'react'; +import { useEffect } from 'react'; import { useQuery, useQueryClient } from '@tanstack/react-query'; import type { EventServiceClient } from '@kilocode/event-service'; import { @@ -28,16 +28,14 @@ export function reduceBotStatusOnEvent( export function useBotStatus( client: KiloChatClient, eventClient: EventServiceClient, - sandboxId: string | null + sandboxId: string | null, + active: boolean = true ): BotStatusRecord | null { const queryClient = useQueryClient(); - // WS-ready gate. onConnected fires synchronously if already connected, and on every reconnect. - const [wsReady, setWsReady] = useState(false); + // On resync (reconnect or sequence gap), refetch to catch up on missed events. useEffect(() => { - return eventClient.onConnected(() => { - setWsReady(true); - // On reconnect, refetch to catch up on anything we missed while disconnected. + return eventClient.onResync(() => { if (sandboxId) { void queryClient.invalidateQueries({ queryKey: botStatusKey(sandboxId) }); } @@ -69,8 +67,8 @@ export function useBotStatus( }); return queryClient.getQueryData(botStatusKey(sandboxId)) ?? null; }, - enabled: sandboxId !== null && wsReady, - refetchInterval: POLL_INTERVAL_MS, + enabled: sandboxId !== null, + refetchInterval: active ? POLL_INTERVAL_MS : undefined, staleTime: 0, }); diff --git a/packages/kilo-chat-hooks/src/use-conversations.test.ts b/packages/kilo-chat-hooks/src/use-conversations.test.ts index 771addee366..18251de9180 100644 --- a/packages/kilo-chat-hooks/src/use-conversations.test.ts +++ b/packages/kilo-chat-hooks/src/use-conversations.test.ts @@ -427,6 +427,7 @@ describe('registerConversationListCacheHandlers', () => { const listKey = conversationsKey(sandboxId); const messageKey = messagesKey(conversationId); let activityHandler: EventHandler | undefined; + let resyncHandler: (() => void) | undefined; const off = () => {}; const kiloChatClient: HandlerOptions['kiloChatClient'] = { @@ -440,7 +441,10 @@ describe('registerConversationListCacheHandlers', () => { }, }; const eventService: HandlerOptions['eventService'] = { - onReconnect: () => off, + onResync: handler => { + resyncHandler = handler; + return off; + }, }; queryClient.setQueryData( @@ -458,6 +462,13 @@ describe('registerConversationListCacheHandlers', () => { sandboxId, }); + if (!resyncHandler) throw new Error('resync handler was not registered'); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(false); + + // Resync invalidates the conversation list + resyncHandler(); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(messageKey)?.isInvalidated).toBe(false); if (!activityHandler) throw new Error('activity handler was not registered'); @@ -476,6 +487,7 @@ describe('registerConversationListCacheHandlers', () => { const listKey = conversationsKey(sandboxId); const messageKey = messagesKey(conversationId); let activityHandler: EventHandler | undefined; + let resyncHandler: (() => void) | undefined; const off = () => {}; const kiloChatClient: HandlerOptions['kiloChatClient'] = { @@ -489,7 +501,10 @@ describe('registerConversationListCacheHandlers', () => { }, }; const eventService: HandlerOptions['eventService'] = { - onReconnect: () => off, + onResync: handler => { + resyncHandler = handler; + return off; + }, }; queryClient.setQueryData( @@ -508,6 +523,13 @@ describe('registerConversationListCacheHandlers', () => { sandboxId, }); + if (!resyncHandler) throw new Error('resync handler was not registered'); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(false); + + // Resync invalidates the conversation list + resyncHandler(); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(true); + if (!activityHandler) throw new Error('activity handler was not registered'); activityHandler(kiloclawInstanceContext(sandboxId), { diff --git a/packages/kilo-chat-hooks/src/use-conversations.ts b/packages/kilo-chat-hooks/src/use-conversations.ts index a7a0ee730d3..3ebb8da86e0 100644 --- a/packages/kilo-chat-hooks/src/use-conversations.ts +++ b/packages/kilo-chat-hooks/src/use-conversations.ts @@ -46,14 +46,14 @@ type KiloChatConversationEventClient = Pick< | 'onConversationActivity' >; -type ReconnectEventService = { - onReconnect: (handler: () => void) => () => void; +type ResyncEventService = { + onResync: (handler: () => void) => () => void; }; type RegisterConversationListCacheHandlersOptions = { activeConversationId?: string | null; currentUserId: string | null; - eventService: ReconnectEventService; + eventService: ResyncEventService; kiloChatClient: KiloChatConversationEventClient; queryClient: QueryClient; queryKey: QueryKey; @@ -577,7 +577,7 @@ export function registerConversationListCacheHandlers({ if (!matchesContext(ctx)) return; patchActivity(event); }), - eventService.onReconnect(() => { + eventService.onResync(() => { invalidateConversationListQuery(queryClient, queryKey); }), ]; diff --git a/packages/kilo-chat/src/client.ts b/packages/kilo-chat/src/client.ts index 4a8e9c9738b..d502995c1d1 100644 --- a/packages/kilo-chat/src/client.ts +++ b/packages/kilo-chat/src/client.ts @@ -1,4 +1,5 @@ import { z } from 'zod'; +import { withDeadline, CONTROL_PLANE_DEADLINE_MS, SEND_DEADLINE_MS } from '@kilocode/event-service'; import { KiloChatApiError } from './errors'; import { conversationListResponseSchema, @@ -98,6 +99,14 @@ export class KiloChatClient { // ── Mutations via HTTP ──────────────────────────────────────────────────── + /** + * Send a message into a conversation. The send is queued per conversation so + * concurrent callers cannot race ahead and get a lower server-assigned ULID. + * + * **Timeout is uncertain.** When the deadline fires the server may already + * have accepted the message. Callers must reconcile the conversation state + * via `listMessages` (or event-service events) before retrying the send. + */ async sendMessage(req: CreateMessageRequest): Promise { const body = req satisfies CreateMessageRequest; const prev = this.sendQueues.get(req.conversationId) ?? Promise.resolve(); @@ -107,6 +116,7 @@ export class KiloChatClient { method: 'POST', body, schema: createMessageResponseSchema, + deadlineMs: SEND_DEADLINE_MS, }), // A failed prior send must not block subsequent sends — swallow the // rejection on the chain; the original caller already received it. @@ -115,6 +125,7 @@ export class KiloChatClient { method: 'POST', body, schema: createMessageResponseSchema, + deadlineMs: SEND_DEADLINE_MS, }) ); this.sendQueues.set(req.conversationId, next); @@ -435,21 +446,31 @@ export class KiloChatClient { body?: unknown; query?: Record; schema: z.ZodType; + signal?: AbortSignal; + deadlineMs?: number; } ): Promise { - try { - return await this.httpRequestOnce(path, opts); - } catch (err) { - const onUnauthorized = this.onUnauthorized; - if (!this.shouldRecoverFromUnauthorized(err) || onUnauthorized === undefined) { - throw err; - } - const decision = await onUnauthorized(); - if (decision !== 'retry') { - throw err; - } - return this.httpRequestOnce(path, opts); - } + const deadline = opts.deadlineMs ?? CONTROL_PLANE_DEADLINE_MS; + return withDeadline( + deadline, + async signal => { + try { + return await this.httpRequestOnce(path, { ...opts, signal }); + } catch (err) { + // Deadline/timeout errors must NOT trigger unauthorized recovery. + const onUnauthorized = this.onUnauthorized; + if (!this.shouldRecoverFromUnauthorized(err) || onUnauthorized === undefined) { + throw err; + } + const decision = await onUnauthorized(); + if (decision !== 'retry') { + throw err; + } + return this.httpRequestOnce(path, { ...opts, signal }); + } + }, + opts.signal + ); } private shouldRecoverFromUnauthorized(err: unknown): err is KiloChatApiError { @@ -467,6 +488,7 @@ export class KiloChatClient { body?: unknown; query?: Record; schema: z.ZodType; + signal?: AbortSignal; } ): Promise { const token = await this.getToken(); @@ -489,6 +511,7 @@ export class KiloChatClient { method: opts.method ?? 'GET', headers, body: opts.body !== undefined ? JSON.stringify(opts.body) : undefined, + signal: opts.signal, }); if (!res.ok) { diff --git a/packages/kilo-chat/src/events.ts b/packages/kilo-chat/src/events.ts index 15633a3160b..fd81eb7241f 100644 --- a/packages/kilo-chat/src/events.ts +++ b/packages/kilo-chat/src/events.ts @@ -1,7 +1,7 @@ import { z } from 'zod'; import { actionGroupIdSchema, - capabilitySchema, + capabilityListSchema, conversationListItemSchema, conversationTitleSchema, contentBlockSchema, @@ -106,7 +106,7 @@ export const botStatusEventSchema = z.object({ sandboxId: sandboxIdSchema, online: z.boolean(), at: nonNegativeIntegerSchema, - capabilities: z.array(capabilitySchema).optional(), + capabilities: capabilityListSchema.optional(), }); export const conversationStatusEventSchema = z.object({ diff --git a/packages/kilo-chat/src/schemas.ts b/packages/kilo-chat/src/schemas.ts index bb4036bae51..10c8d8aa5b0 100644 --- a/packages/kilo-chat/src/schemas.ts +++ b/packages/kilo-chat/src/schemas.ts @@ -165,6 +165,10 @@ export const memberKindSchema = z.enum(['user', 'bot']); export const capabilitySchema = z.enum(['attachments']); export type Capability = z.infer; +export const capabilityListSchema = z + .array(z.string()) + .transform(arr => arr.filter((c): c is Capability => capabilitySchema.safeParse(c).success)); + export const conversationMemberSchema = z.object({ id: z.string(), kind: memberKindSchema, @@ -335,7 +339,7 @@ export const listMessagesQuerySchema = z.object({ export const botStatusRequestSchema = z.object({ online: z.boolean(), at: nonNegativeIntegerSchema, - capabilities: z.array(capabilitySchema).optional(), + capabilities: capabilityListSchema.optional(), }); export const conversationStatusRequestSchema = z.object({ @@ -350,7 +354,7 @@ export const botStatusRecordSchema = z.object({ online: z.boolean(), at: nonNegativeIntegerSchema, updatedAt: nonNegativeIntegerSchema, - capabilities: z.array(capabilitySchema).optional(), + capabilities: capabilityListSchema.optional(), }); export const conversationStatusRecordSchema = z.object({ diff --git a/packages/kilo-chat/test/client.test.ts b/packages/kilo-chat/test/client.test.ts index e3ee3468e6e..87a8fead349 100644 --- a/packages/kilo-chat/test/client.test.ts +++ b/packages/kilo-chat/test/client.test.ts @@ -1,6 +1,7 @@ import { describe, it, expect, vi } from 'vitest'; import { KiloChatClient } from '../src/client'; import { KiloChatApiError } from '../src/errors'; +import { RequestDeadlineError } from '@kilocode/event-service'; import type { KiloChatClientConfig } from '../src/types'; function createMockConfig(fetchFn: typeof globalThis.fetch): KiloChatClientConfig { @@ -375,4 +376,273 @@ describe('KiloChatClient', () => { await expect(client.listConversations()).rejects.toThrow(); }); }); + + describe('request deadlines', () => { + const sentMessage = { + id: 'm1', + senderId: 'user-1', + content: [{ type: 'text' as const, text: 'hi' }], + inReplyToMessageId: null, + replyTo: null, + updatedAt: null, + clientUpdatedAt: null, + deleted: false, + deliveryFailed: false, + reactions: [], + }; + + it('timed-out send releases the conversation queue for the next send', async () => { + vi.useFakeTimers(); + try { + // First fetch hangs and respects the abort signal. Second fetch resolves. + const fetch = vi + .fn() + .mockImplementationOnce((_url, init) => { + return new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal) { + signal.addEventListener('abort', () => { + reject(signal.reason); + }); + } + }); + }) + .mockResolvedValueOnce( + new Response( + JSON.stringify({ messageId: 'm2', message: { ...sentMessage, id: 'm2' } }), + { + status: 201, + } + ) + ); + + const client = new KiloChatClient(createMockConfig(fetch)); + + // Queue send 1 — it will hang on fetch. + const send1Promise = client.sendMessage({ + conversationId: 'c1', + content: [{ type: 'text', text: 'first' }], + }); + + // Queue send 2 immediately — chained after send 1's promise. + const send2Promise = client.sendMessage({ + conversationId: 'c1', + content: [{ type: 'text', text: 'second' }], + }); + + // Suppress unhandled rejections during fake-timer advance. + send1Promise.catch(() => {}); + + // Advance past the 30s send deadline. + await vi.advanceTimersByTimeAsync(30_001); + + // Send 1 rejects with RequestDeadlineError. + await expect(send1Promise).rejects.toThrow(RequestDeadlineError); + + // Send 2 goes through — the queue was not blocked. + await expect(send2Promise).resolves.toEqual({ + messageId: 'm2', + message: { ...sentMessage, id: 'm2' }, + }); + + expect(fetch).toHaveBeenCalledTimes(2); + // First fetch: timed-out POST + expect(fetch).toHaveBeenNthCalledWith( + 1, + 'https://chat.example.com/v1/messages', + expect.objectContaining({ method: 'POST' }) + ); + // Second fetch: successful POST after queue release + expect(fetch).toHaveBeenNthCalledWith( + 2, + 'https://chat.example.com/v1/messages', + expect.objectContaining({ method: 'POST' }) + ); + } finally { + vi.useRealTimers(); + } + }); + + it('control-plane deadline applies to listConversations', async () => { + vi.useFakeTimers(); + try { + const fetch = vi.fn().mockImplementationOnce((_url, init) => { + return new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal) { + signal.addEventListener('abort', () => { + reject(signal.reason); + }); + } + }); + }); + + const client = new KiloChatClient(createMockConfig(fetch)); + const promise = client.listConversations(); + promise.catch(() => {}); // suppress unhandled rejection + + // Advance past 15s control-plane deadline. + await vi.advanceTimersByTimeAsync(15_001); + + await expect(promise).rejects.toThrow(RequestDeadlineError); + await expect(promise).rejects.toThrow('timed out after 15000ms'); + } finally { + vi.useRealTimers(); + } + }); + + it('successful send completes within deadline', async () => { + vi.useFakeTimers(); + try { + const fetch = vi + .fn() + .mockResolvedValueOnce( + new Response(JSON.stringify({ messageId: 'm1', message: sentMessage }), { status: 201 }) + ); + + const client = new KiloChatClient(createMockConfig(fetch)); + const promise = client.sendMessage({ + conversationId: 'c1', + content: [{ type: 'text', text: 'hi' }], + }); + + // Advance a small amount — well under the 30s deadline. + await vi.advanceTimersByTimeAsync(100); + await expect(promise).resolves.toEqual({ messageId: 'm1', message: sentMessage }); + expect(fetch).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it('timeout does not trigger unauthorized recovery', async () => { + vi.useFakeTimers(); + try { + const fetch = vi.fn().mockImplementationOnce((_url, init) => { + return new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal) { + signal.addEventListener('abort', () => { + reject(signal.reason); + }); + } + }); + }); + + const onUnauthorized = vi.fn<() => 'retry'>(() => 'retry'); + const client = new KiloChatClient({ + ...createMockConfig(fetch), + onUnauthorized, + }); + + const promise = client.listConversations(); + promise.catch(() => {}); // suppress unhandled rejection + + await vi.advanceTimersByTimeAsync(15_001); + + // Rejects with a deadline error, NOT an auth error. + await expect(promise).rejects.toThrow(RequestDeadlineError); + + // onUnauthorized must NOT be called — timeouts are not auth failures. + expect(onUnauthorized).not.toHaveBeenCalled(); + expect(fetch).toHaveBeenCalledTimes(1); // no retry + } finally { + vi.useRealTimers(); + } + }); + + it('ABORT: no automatic resend after a timed-out send', async () => { + vi.useFakeTimers(); + try { + const fetch = vi.fn().mockImplementationOnce((_url, init) => { + return new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal) { + signal.addEventListener('abort', () => { + reject(signal.reason); + }); + } + }); + }); + + const client = new KiloChatClient(createMockConfig(fetch)); + const promise = client.sendMessage({ + conversationId: 'c1', + content: [{ type: 'text', text: 'timeout' }], + }); + promise.catch(() => {}); // suppress unhandled rejection + + // Advance past the 30s send deadline. + await vi.advanceTimersByTimeAsync(30_001); + + await expect(promise).rejects.toThrow(RequestDeadlineError); + + // No automatic resend: fetch called exactly once (not retried). + expect(fetch).toHaveBeenCalledTimes(1); + + // Advance further — still no additional fetch calls. + await vi.advanceTimersByTimeAsync(60_000); + expect(fetch).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it('ABORT: send deadline is 30s', async () => { + vi.useFakeTimers(); + try { + const fetch = vi.fn().mockImplementationOnce((_url, init) => { + return new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal) { + signal.addEventListener('abort', () => { + reject(signal.reason); + }); + } + }); + }); + + const client = new KiloChatClient(createMockConfig(fetch)); + const sendPromise = client.sendMessage({ + conversationId: 'c1', + content: [{ type: 'text', text: 'hi' }], + }); + sendPromise.catch(() => {}); // suppress unhandled rejection + + // At 15s, the send should still be alive (send deadline is 30s). + await vi.advanceTimersByTimeAsync(15_000); + + // At 30s, it should reject. + await vi.advanceTimersByTimeAsync(15_001); + await expect(sendPromise).rejects.toThrow('timed out after 30000ms'); + } finally { + vi.useRealTimers(); + } + }); + + it('ABORT: control-plane deadline is 15s', async () => { + vi.useFakeTimers(); + try { + const fetch = vi.fn().mockImplementationOnce((_url, init) => { + return new Promise((_resolve, reject) => { + const signal = init?.signal; + if (signal) { + signal.addEventListener('abort', () => { + reject(signal.reason); + }); + } + }); + }); + + const client = new KiloChatClient(createMockConfig(fetch)); + const promise = client.listConversations(); + promise.catch(() => {}); // suppress unhandled rejection + + await vi.advanceTimersByTimeAsync(15_001); + await expect(promise).rejects.toThrow('timed out after 15000ms'); + } finally { + vi.useRealTimers(); + } + }); + }); }); diff --git a/packages/kilo-chat/test/events.bot-status.test.ts b/packages/kilo-chat/test/events.bot-status.test.ts index 9d51c44292b..53844f77719 100644 --- a/packages/kilo-chat/test/events.bot-status.test.ts +++ b/packages/kilo-chat/test/events.bot-status.test.ts @@ -11,6 +11,7 @@ describe('botStatusEventSchema', () => { }); expect(parsed.capabilities).toEqual(['attachments']); }); + it('parses legacy event without capabilities', () => { const parsed = botStatusEventSchema.parse({ sandboxId: 'sbx-abc', @@ -19,4 +20,24 @@ describe('botStatusEventSchema', () => { }); expect(parsed.capabilities).toBeUndefined(); }); + + it('filters unknown capabilities and retains known ones', () => { + const parsed = botStatusEventSchema.parse({ + sandboxId: 'sbx-abc', + online: true, + at: 1, + capabilities: ['attachments', 'future-feature', 'another-unknown'], + }); + expect(parsed.capabilities).toEqual(['attachments']); + }); + + it('parses event with only unknown capabilities (empty array result)', () => { + const parsed = botStatusEventSchema.parse({ + sandboxId: 'sbx-abc', + online: true, + at: 1, + capabilities: ['future-feature'], + }); + expect(parsed.capabilities).toEqual([]); + }); }); diff --git a/packages/kilo-chat/test/schemas.bot-status.test.ts b/packages/kilo-chat/test/schemas.bot-status.test.ts index e6e21493bf0..6d1a12a8284 100644 --- a/packages/kilo-chat/test/schemas.bot-status.test.ts +++ b/packages/kilo-chat/test/schemas.bot-status.test.ts @@ -14,13 +14,28 @@ describe('botStatusRequestSchema with capabilities', () => { const parsed = botStatusRequestSchema.parse({ online: true, at: 1000 }); expect(parsed.capabilities).toBeUndefined(); }); - it('rejects unknown capability', () => { + it('filters unknown capabilities, keeping known ones', () => { const r = botStatusRequestSchema.safeParse({ online: true, at: 1, capabilities: ['foo'], }); - expect(r.success).toBe(false); + expect(r.success).toBe(true); + if (r.success) { + expect(r.data.capabilities).toEqual([]); + } + }); + + it('filters mixed known and unknown capabilities', () => { + const r = botStatusRequestSchema.safeParse({ + online: true, + at: 1, + capabilities: ['attachments', 'foo'], + }); + expect(r.success).toBe(true); + if (r.success) { + expect(r.data.capabilities).toEqual(['attachments']); + } }); }); @@ -34,4 +49,14 @@ describe('botStatusRecordSchema with capabilities', () => { }); expect(parsed.capabilities).toEqual(['attachments']); }); + + it('filters unknown capabilities from record', () => { + const parsed = botStatusRecordSchema.parse({ + online: true, + at: 1, + updatedAt: 2, + capabilities: ['attachments', 'unknown'], + }); + expect(parsed.capabilities).toEqual(['attachments']); + }); }); diff --git a/packages/kilo-chat/test/schemas.capability.test.ts b/packages/kilo-chat/test/schemas.capability.test.ts index 7c33ef47bf4..15e1bbb3ec8 100644 --- a/packages/kilo-chat/test/schemas.capability.test.ts +++ b/packages/kilo-chat/test/schemas.capability.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect } from 'vitest'; -import { capabilitySchema } from '../src/schemas'; +import { capabilitySchema, capabilityListSchema } from '../src/schemas'; describe('capabilitySchema', () => { it('accepts "attachments"', () => { @@ -12,3 +12,37 @@ describe('capabilitySchema', () => { expect(capabilitySchema.safeParse(42).success).toBe(false); }); }); + +describe('capabilityListSchema', () => { + it('retains only known capabilities and drops unknowns', () => { + const result = capabilityListSchema.safeParse(['attachments', 'foo', 'bar']); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data).toEqual(['attachments']); + } + }); + + it('returns empty array when all capabilities are unknown', () => { + const result = capabilityListSchema.safeParse(['foo', 'bar']); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data).toEqual([]); + } + }); + + it('returns single known capability', () => { + const result = capabilityListSchema.safeParse(['attachments']); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data).toEqual(['attachments']); + } + }); + + it('returns empty array for empty input', () => { + const result = capabilityListSchema.safeParse([]); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data).toEqual([]); + } + }); +}); diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 410381ba543..e6ae60ae797 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -1347,6 +1347,17 @@ export class CloudAgentSession extends DurableObject { return this.eventQueries.getLatestAssistantMessage(sessionId, metadata.auth.kiloSessionId); } + /** + * Get the latest persisted event ID from the event log. + * + * Returns `null` when the DO has a pending deletion intent or when no + * events have been persisted yet. + */ + async getLatestEventId(): Promise { + if (await this.hasDeletionIntent()) return null; + return this.eventQueries.getLatestEventId(); + } + async getMessageResult(messageId: string): Promise { const metadata = await this.getMetadata(); if (!metadata) return { type: 'session-not-found' }; diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index d2b69e1cc91..66135d15215 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/services/cloud-agent-next/src/router.test.ts @@ -913,6 +913,7 @@ describe('router sessionId validation', () => { getMetadata: mockGetMetadata, getCurrentRuntimeExecution: mockGetCurrentRuntimeExecution, getCurrentMessageWork: mockGetCurrentMessageWork, + getLatestEventId: vi.fn().mockResolvedValue(null), })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], @@ -1183,6 +1184,115 @@ describe('router sessionId validation', () => { ).rejects.toThrow('Authentication required'); }); }); + + describe('latestEventId', () => { + let mockContext: TRPCContext; + let caller: ReturnType; + let mockGetMetadata: ReturnType; + let mockGetLatestEventId: ReturnType; + + beforeEach(() => { + vi.clearAllMocks(); + + mockGetMetadata = vi.fn(); + mockGetLatestEventId = vi.fn().mockResolvedValue(null); + + mockContext = { + userId: 'test-user-123', + authToken: 'test-token', + botId: undefined, + request: {} as Request, + env: { + Sandbox: {} as TRPCContext['env']['Sandbox'], + SandboxSmall: {} as TRPCContext['env']['SandboxSmall'], + SandboxDIND: {} as TRPCContext['env']['SandboxDIND'], + SandboxCodeReview: {} as TRPCContext['env']['SandboxCodeReview'], + SandboxContainment: {} as TRPCContext['env']['SandboxContainment'], + SandboxSmallContainment: {} as TRPCContext['env']['SandboxSmallContainment'], + SandboxCodeReviewContainment: + {} as TRPCContext['env']['SandboxCodeReviewContainment'], + CLOUD_AGENT_SESSION: { + idFromName: vi.fn((id: string) => ({ id })), + get: vi.fn(() => ({ + getMetadata: mockGetMetadata, + getCurrentRuntimeExecution: vi.fn().mockResolvedValue(null), + getCurrentMessageWork: vi.fn().mockResolvedValue(null), + getLatestEventId: mockGetLatestEventId, + })), + } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], + USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], + SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], + SESSION_INGEST: { fetch: vi.fn() } as unknown as TRPCContext['env']['SESSION_INGEST'], + CONTAINER_USAGE_METER: {} as TRPCContext['env']['CONTAINER_USAGE_METER'], + R2_BUCKET: {} as TRPCContext['env']['R2_BUCKET'], + CLOUD_AGENT_REPORT_QUEUE: {} as TRPCContext['env']['CLOUD_AGENT_REPORT_QUEUE'], + GIT_TOKEN_SERVICE: {} as Env['GIT_TOKEN_SERVICE'], + NEXTAUTH_SECRET: 'test-secret', + INTERNAL_API_SECRET_PROD: { + get: vi.fn().mockResolvedValue('test-secret'), + } as unknown as TRPCContext['env']['INTERNAL_API_SECRET_PROD'], + HYPERDRIVE: { + connectionString: 'postgresql://test', + } as unknown as TRPCContext['env']['HYPERDRIVE'], + NOTIFICATIONS: {} as TRPCContext['env']['NOTIFICATIONS'], + }, + }; + caller = appRouter.createCaller(mockContext); + }); + + it('returns latestEventId when events have been persisted', async () => { + const sessionId: SessionId = 'agent_a1b2c3d4-0000-0000-0000-000000000001'; + mockGetMetadata.mockResolvedValue( + legacySessionMetadata({ + version: 123456789, + sessionId, + userId: 'test-user-123', + timestamp: 123456789, + }) + ); + mockGetLatestEventId.mockResolvedValue(42); + + const result = await caller.getSession({ cloudAgentSessionId: sessionId }); + + expect(result.latestEventId).toBe(42); + }); + + it('returns null latestEventId when no events exist', async () => { + const sessionId: SessionId = 'agent_a1b2c3d4-0000-0000-0000-000000000002'; + mockGetMetadata.mockResolvedValue( + legacySessionMetadata({ + version: 123456789, + sessionId, + userId: 'test-user-123', + timestamp: 123456789, + }) + ); + mockGetLatestEventId.mockResolvedValue(null); + + const result = await caller.getSession({ cloudAgentSessionId: sessionId }); + + expect(result.latestEventId).toBeNull(); + }); + + it('fails open: swallows getLatestEventId failure and returns null', async () => { + const sessionId: SessionId = 'agent_a1b2c3d4-0000-0000-0000-000000000003'; + mockGetMetadata.mockResolvedValue( + legacySessionMetadata({ + version: 123456789, + sessionId, + userId: 'test-user-123', + timestamp: 123456789, + }) + ); + mockGetLatestEventId.mockRejectedValue(new Error('DO unreachable')); + + const result = await caller.getSession({ cloudAgentSessionId: sessionId }); + + // Session must still resolve successfully with null latestEventId. + expect(result.sessionId).toBe(sessionId); + expect(result.latestEventId).toBeNull(); + }); + }); }); describe('getSessionHealth procedure', () => { @@ -1641,6 +1751,17 @@ describe('router sessionId validation', () => { }); }); }); + + describe('getWatermarkEventId procedure', () => { + // Removed — the standalone getWatermarkEventId tRPC procedure has been + // replaced by getSession(...).latestEventId. The watermark is now an + // optional field on the existing getSession response. + it('has been replaced by getSession(...).latestEventId', () => { + // Placeholder to keep the describe block from being empty and to + // document the removal. Actual watermark coverage lives in the + // getSession procedure tests below. + }); + }); }); describe('getMessageResult procedure', () => { diff --git a/services/cloud-agent-next/src/router/handlers/session-management.ts b/services/cloud-agent-next/src/router/handlers/session-management.ts index ae666927aad..4f90a80458f 100644 --- a/services/cloud-agent-next/src/router/handlers/session-management.ts +++ b/services/cloud-agent-next/src/router/handlers/session-management.ts @@ -275,6 +275,25 @@ export function createSessionManagementHandlers() { 'getCurrentMessageWork' ); + // Fetch the latest persisted event ID for the durability watermark. + // Failures are swallowed so an optional watermark read never blocks + // the session response. + let latestEventId: number | null = null; + try { + latestEventId = await withDORetry( + getStub, + s => s.getLatestEventId(), + 'getLatestEventId' + ); + } catch (error) { + logger + .withFields({ + sessionId, + error: error instanceof Error ? error.message : String(error), + }) + .warn('Failed to fetch latest event ID for getSession'); + } + // Compute sandboxId for log correlation const sessionMetadata = metadata; const metadataProfile = readProfileBundle(sessionMetadata); @@ -350,6 +369,7 @@ export function createSessionManagementHandlers() { timestamp: sessionMetadata.lifecycle.timestamp, version: sessionMetadata.lifecycle.version, + latestEventId, }; }); }), diff --git a/services/cloud-agent-next/src/router/schemas.ts b/services/cloud-agent-next/src/router/schemas.ts index 00de4c12ea9..c5829b5ada4 100644 --- a/services/cloud-agent-next/src/router/schemas.ts +++ b/services/cloud-agent-next/src/router/schemas.ts @@ -976,6 +976,17 @@ export const GetSessionOutput = z.object({ // Versioning timestamp: z.number().describe('Last update timestamp'), version: z.number().describe('Metadata version for cache invalidation'), + + // Latest persisted event ID from the Cloud Agent event log. Used as a + // durability watermark: numeric and present after at least one event has + // been persisted. Null when no events exist or the session has no events yet. + latestEventId: z + .number() + .int() + .positive() + .nullable() + .optional() + .describe('Latest persisted event ID, or null when no events exist'), }); export type GetSessionResponse = z.infer; diff --git a/services/event-service/src/__tests__/user-session-do.test.ts b/services/event-service/src/__tests__/user-session-do.test.ts index 8f51cdba971..fcd89001d55 100644 --- a/services/event-service/src/__tests__/user-session-do.test.ts +++ b/services/event-service/src/__tests__/user-session-do.test.ts @@ -1,4 +1,4 @@ -import { env } from 'cloudflare:test'; +import { env, runInDurableObject } from 'cloudflare:test'; import { describe, it, expect } from 'vitest'; // Helper: connect a WebSocket to a DO stub for a given userId @@ -65,6 +65,7 @@ describe('UserSessionDO', () => { context: 'project:abc', event: 'task.created', payload: { taskId: '1' }, + seq: 1, }); ws.close(); @@ -149,6 +150,168 @@ describe('UserSessionDO', () => { ws.close(); }); + it('includes an incrementing seq in each event', async () => { + const userId = 'user-seq-increment'; + const { ws, stub } = await connectWs(userId); + + ws.send(JSON.stringify({ type: 'context.subscribe', contexts: ['project:seq'] })); + await new Promise(r => setTimeout(r, 50)); + + const msg1Promise = nextMessage(ws); + await stub.pushEvent('project:seq', 'task.created', { taskId: '1' }); + const msg1 = await msg1Promise; + expect(msg1).toMatchObject({ type: 'event', seq: 1 }); + + const msg2Promise = nextMessage(ws); + await stub.pushEvent('project:seq', 'task.updated', { taskId: '2' }); + const msg2 = await msg2Promise; + expect(msg2).toMatchObject({ type: 'event', seq: 2 }); + + const msg3Promise = nextMessage(ws); + await stub.pushEvent('project:seq', 'task.deleted', { taskId: '3' }); + const msg3 = await msg3Promise; + expect(msg3).toMatchObject({ type: 'event', seq: 3 }); + + ws.close(); + }); + + it('updates acked on receipt of an ack message', async () => { + const userId = 'user-ack-update'; + const { ws, stub } = await connectWs(userId); + + ws.send(JSON.stringify({ type: 'context.subscribe', contexts: ['project:ack'] })); + await new Promise(r => setTimeout(r, 100)); + + // Push two events and drain messages to ensure seq is built + const msg1Promise = nextMessage(ws, 2000); + await stub.pushEvent('project:ack', 'e1', {}); + await msg1Promise; + + const msg2Promise = nextMessage(ws, 2000); + await stub.pushEvent('project:ack', 'e2', {}); + await msg2Promise; + + // Send ack for seq 2 + ws.send(JSON.stringify({ type: 'ack', seq: 2 })); + await new Promise(r => setTimeout(r, 100)); + + // Push one more event — socket should still be open (gap is 1) + const msg3Promise = nextMessage(ws, 2000); + await stub.pushEvent('project:ack', 'e3', {}); + const msg3 = await msg3Promise; + expect(msg3).toMatchObject({ type: 'event', seq: 3 }); + + ws.close(); + }); + + it( + 'closes socket with code 4002 when unacked gap exceeds threshold', + { timeout: 30000 }, + async () => { + const userId = 'user-zombie-close'; + const { ws, stub } = await connectWs(userId); + + ws.send(JSON.stringify({ type: 'context.subscribe', contexts: ['project:z'] })); + await new Promise(r => setTimeout(r, 100)); + + // Ack seq 0 — the socket is now eligible for zombie close + ws.send(JSON.stringify({ type: 'ack', seq: 0 })); + await new Promise(r => setTimeout(r, 100)); + + const closePromise = new Promise(resolve => { + ws.addEventListener('close', resolve); + }); + + // Push enough events to exceed UNACKED_THRESHOLD (1000). + // seq starts at 0, increments each push. After N pushes seq = N. + // Gap = seq - acked = N - 0. Need N > 1000. + // Push 1100 events for margin. + for (let i = 0; i < 1100; i++) { + await stub.pushEvent('project:z', 'fill', {}); + if (i % 100 === 0) { + // Periodically yield to let the DO process the close + await new Promise(r => setTimeout(r, 10)); + } + if (ws.readyState === 3 /* CLOSED */) break; + } + + const closeEvent = await closePromise; + expect(closeEvent.code).toBe(4002); + } + ); + + it( + 'never closes a socket that has not acknowledged (old client compat)', + { timeout: 25000 }, + async () => { + const userId = 'user-no-ack-no-close'; + const { ws, stub } = await connectWs(userId); + + ws.send(JSON.stringify({ type: 'context.subscribe', contexts: ['project:noack'] })); + await new Promise(r => setTimeout(r, 100)); + + let closed = false; + ws.addEventListener('close', () => { + closed = true; + }); + + // Push events beyond the threshold without any ack — should never + // trigger zombie close because the socket has never acknowledged. + for (let i = 0; i < 1100; i++) { + await stub.pushEvent('project:noack', 'fill', {}); + if (i % 100 === 0) { + await new Promise(r => setTimeout(r, 10)); + } + if (closed) break; + } + + expect(closed).toBe(false); + ws.close(); + } + ); + + it( + 'healthy high-volume clients survive with periodic acks in a 15s cadence', + { timeout: 30000 }, + async () => { + const userId = 'user-high-volume-ack'; + const { ws, stub } = await connectWs(userId); + + ws.send(JSON.stringify({ type: 'context.subscribe', contexts: ['project:hv'] })); + await new Promise(r => setTimeout(r, 100)); + + // Ack seq 0 — the socket is now eligible for zombie close + ws.send(JSON.stringify({ type: 'ack', seq: 0 })); + await new Promise(r => setTimeout(r, 100)); + + let closed = false; + ws.addEventListener('close', () => { + closed = true; + }); + + // Push 200 events — far more than the old threshold of 50 — with + // periodic acks that simulate a healthy 15s ack cadence. The high + // threshold (1000) must not close the socket. + for (let i = 0; i < 200; i++) { + await stub.pushEvent('project:hv', 'fill', {}); + // Ack every 25 events to keep the gap well under 1000 + if (i > 0 && i % 25 === 0) { + ws.send(JSON.stringify({ type: 'ack', seq: i })); + await new Promise(r => setTimeout(r, 10)); + } + await new Promise(r => setTimeout(r, 2)); + if (closed) break; + } + + // Final ack to cover all remaining events + ws.send(JSON.stringify({ type: 'ack', seq: 200 })); + await new Promise(r => setTimeout(r, 50)); + + expect(closed).toBe(false); + ws.close(); + } + ); + it('closes sockets that exceed the per-socket context limit', async () => { const userId = 'user-too-many-contexts'; const { ws, stub } = await connectWs(userId); @@ -160,4 +323,63 @@ describe('UserSessionDO', () => { expect(await stub.pushEvent('project:200', 'test', {})).toBe(false); ws.close(); }); + + it( + 'defaults counters to zero and emits sequence one with old attachment', + { timeout: 25000 }, + async () => { + // Load a WebSocket whose server-side attachment has only the old + // shape: `{ contexts: ['project:old'] }`. The DO fetches a new + // attachment; we overwrite it with the old shape inside the DO. + const userId = 'user-old-attachment'; + const id = env.USER_SESSION_DO.idFromName(userId); + const stub = env.USER_SESSION_DO.get(id); + + const res = await stub.fetch('https://do/connect', { + headers: { Upgrade: 'websocket' }, + }); + const ws = res.webSocket!; + ws.accept(); + + const rawState = await runInDurableObject(stub, async (_instance, state) => { + const sockets = state.getWebSockets(); + const serverSocket = sockets[0]!; + serverSocket.serializeAttachment({ contexts: ['project:old'] }); + return serverSocket.deserializeAttachment(); + }); + + // The old attachment lacks seq, acked, and ackReceived. + // getState defaults them to 0, 0, false. + expect(rawState).toEqual({ contexts: ['project:old'] }); + + // First event: the missing seq defaults to 0 → pushEvent emits seq=1. + const msg1Promise = nextMessage(ws, 2000); + await stub.pushEvent('project:old', 'test', {}); + const msg1 = await msg1Promise; + expect(msg1).toMatchObject({ type: 'event', seq: 1 }); + + // Second event: increments to seq=2. + const msg2Promise = nextMessage(ws, 2000); + await stub.pushEvent('project:old', 'test', {}); + const msg2 = await msg2Promise; + expect(msg2).toMatchObject({ type: 'event', seq: 2 }); + + // Without an ack the socket is never zombie-closed. + let closed = false; + ws.addEventListener('close', () => { + closed = true; + }); + + for (let i = 0; i < 1100; i++) { + await stub.pushEvent('project:old', 'fill', {}); + if (i % 100 === 0) { + await new Promise(r => setTimeout(r, 10)); + } + if (closed) break; + } + + expect(closed).toBe(false); + ws.close(); + } + ); }); diff --git a/services/event-service/src/do/user-session-do.ts b/services/event-service/src/do/user-session-do.ts index d8e595fc1a1..e39ef0ee906 100644 --- a/services/event-service/src/do/user-session-do.ts +++ b/services/event-service/src/do/user-session-do.ts @@ -3,7 +3,13 @@ import { clientMessageSchema, MAX_CONTEXTS } from '@kilocode/event-service'; import { configureDevLogging, logger, withLogTags } from '../util/logger'; import type { ServerMessage } from '../types'; -type SerializedState = { contexts: string[] }; +/** Number of unacked events before a socket is closed as a zombie. */ +const UNACKED_THRESHOLD = 1000; + +/** Close code for a socket that fell too far behind on acknowledgements. */ +const ZOMBIE_CLOSE_CODE = 4002; + +type SerializedState = { contexts: string[]; seq: number; acked: number; ackReceived: boolean }; export class UserSessionDO extends DurableObject { constructor(ctx: DurableObjectState, env: Env) { @@ -24,7 +30,12 @@ export class UserSessionDO extends DurableObject { const [client, server] = [pair[0], pair[1]]; this.ctx.acceptWebSocket(server); - server.serializeAttachment({ contexts: [] } satisfies SerializedState); + server.serializeAttachment({ + contexts: [], + seq: 0, + acked: 0, + ackReceived: false, + } satisfies SerializedState); logger.debug('fetch: accepted websocket'); return new Response(null, { status: 101, webSocket: client }); @@ -81,6 +92,15 @@ export class UserSessionDO extends DurableObject { logger.debug('unsubscribed', { contexts: msg.contexts }); break; } + case 'ack': { + const state = this.getState(ws); + state.acked = Math.max(state.acked, msg.seq); + state.ackReceived = true; + this.saveState(ws, state); + logger.debug('acked', { seq: msg.seq, acked: state.acked }); + this.checkZombieClose(ws, state); + break; + } } }); } @@ -106,19 +126,32 @@ export class UserSessionDO extends DurableObject { logger.setTags({ userId: this.ctx.id.name, context, event }); const sockets = this.ctx.getWebSockets(); - const message = { type: 'event', context, event, payload } satisfies ServerMessage; - const text = JSON.stringify(message); let delivered = false; for (const ws of sockets) { const state = this.getState(ws); if (!state.contexts.has(context)) continue; + + state.seq += 1; + this.saveState(ws, state); + + const message = { + type: 'event', + context, + event, + payload, + seq: state.seq, + } satisfies ServerMessage; + const text = JSON.stringify(message); + try { ws.send(text); delivered = true; } catch { // Connection dead — hibernation will clean up } + + this.checkZombieClose(ws, state); } logger.debug('pushEvent: delivered', { sockets: sockets.length, delivered }); return delivered; @@ -144,14 +177,46 @@ export class UserSessionDO extends DurableObject { // ── Helpers ──────────────────────────────────────────────────────── - private getState(ws: WebSocket): { contexts: Set } { + /** + * Close a socket that has acknowledged at least once but is + * too far behind (seq - acked > UNACKED_THRESHOLD). + * A socket that has never acknowledged is never closed by this rule. + */ + private checkZombieClose( + ws: WebSocket, + state: { seq: number; acked: number; ackReceived: boolean } + ): void { + if (!state.ackReceived) return; + if (state.seq - state.acked > UNACKED_THRESHOLD) { + logger.warn('zombie close', { seq: state.seq, acked: state.acked }); + ws.close(ZOMBIE_CLOSE_CODE, 'too many unacked events'); + } + } + + private getState(ws: WebSocket): { + contexts: Set; + seq: number; + acked: number; + ackReceived: boolean; + } { const raw = ws.deserializeAttachment() as SerializedState | null; - return { contexts: new Set(raw?.contexts ?? []) }; + return { + contexts: new Set(raw?.contexts ?? []), + seq: raw?.seq ?? 0, + acked: raw?.acked ?? 0, + ackReceived: raw?.ackReceived ?? false, + }; } - private saveState(ws: WebSocket, state: { contexts: Set }): void { + private saveState( + ws: WebSocket, + state: { contexts: Set; seq: number; acked: number; ackReceived: boolean } + ): void { ws.serializeAttachment({ contexts: [...state.contexts], + seq: state.seq, + acked: state.acked, + ackReceived: state.ackReceived, } satisfies SerializedState); } } diff --git a/services/kiloclaw/plugins/kilo-chat/src/synced/events.ts b/services/kiloclaw/plugins/kilo-chat/src/synced/events.ts index 15633a3160b..fd81eb7241f 100644 --- a/services/kiloclaw/plugins/kilo-chat/src/synced/events.ts +++ b/services/kiloclaw/plugins/kilo-chat/src/synced/events.ts @@ -1,7 +1,7 @@ import { z } from 'zod'; import { actionGroupIdSchema, - capabilitySchema, + capabilityListSchema, conversationListItemSchema, conversationTitleSchema, contentBlockSchema, @@ -106,7 +106,7 @@ export const botStatusEventSchema = z.object({ sandboxId: sandboxIdSchema, online: z.boolean(), at: nonNegativeIntegerSchema, - capabilities: z.array(capabilitySchema).optional(), + capabilities: capabilityListSchema.optional(), }); export const conversationStatusEventSchema = z.object({ diff --git a/services/kiloclaw/plugins/kilo-chat/src/synced/schemas.ts b/services/kiloclaw/plugins/kilo-chat/src/synced/schemas.ts index bb4036bae51..10c8d8aa5b0 100644 --- a/services/kiloclaw/plugins/kilo-chat/src/synced/schemas.ts +++ b/services/kiloclaw/plugins/kilo-chat/src/synced/schemas.ts @@ -165,6 +165,10 @@ export const memberKindSchema = z.enum(['user', 'bot']); export const capabilitySchema = z.enum(['attachments']); export type Capability = z.infer; +export const capabilityListSchema = z + .array(z.string()) + .transform(arr => arr.filter((c): c is Capability => capabilitySchema.safeParse(c).success)); + export const conversationMemberSchema = z.object({ id: z.string(), kind: memberKindSchema, @@ -335,7 +339,7 @@ export const listMessagesQuerySchema = z.object({ export const botStatusRequestSchema = z.object({ online: z.boolean(), at: nonNegativeIntegerSchema, - capabilities: z.array(capabilitySchema).optional(), + capabilities: capabilityListSchema.optional(), }); export const conversationStatusRequestSchema = z.object({ @@ -350,7 +354,7 @@ export const botStatusRecordSchema = z.object({ online: z.boolean(), at: nonNegativeIntegerSchema, updatedAt: nonNegativeIntegerSchema, - capabilities: z.array(capabilitySchema).optional(), + capabilities: capabilityListSchema.optional(), }); export const conversationStatusRecordSchema = z.object({ diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 8bfd31904b5..a169a76b685 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -22,7 +22,11 @@ vi.mock('./SessionIngestDO', () => ({ getSessionIngestDO: sessionIngestMocks.getSessionIngestDO, })); -import { MAX_CATALOG_RESULT_BYTES, UserConnectionDO } from './UserConnectionDO'; +import { + MAX_CATALOG_RESULT_BYTES, + MAX_DURABLE_RESULT_BYTES, + UserConnectionDO, +} from './UserConnectionDO'; // --------------------------------------------------------------------------- // Mock WebSocket @@ -319,27 +323,41 @@ function sendPing(doInstance: UserConnectionDO, webWs: MockWS, nonce: string) { doInstance.webSocketMessage(webWs as never, msg); } -/** Send a command from a web ws */ -function sendCommand( +/** Send a command from a web ws. Auto-flushes so durable-before-send + * dispatch completes before callers inspect CLI state. */ +async function sendCommand( doInstance: UserConnectionDO, webWs: MockWS, - opts: { id: string; command: string; sessionId?: string; connectionId?: string; data?: unknown } -) { + opts: { + id: string; + command: string; + sessionId?: string; + connectionId?: string; + data?: unknown; + mutationId?: string; + } +): Promise { const msg = JSON.stringify({ type: 'command', ...opts }); doInstance.webSocketMessage(webWs as never, msg); + await flushAsync(); } -/** Send a response from a CLI ws */ -function sendCliResponse( +/** Send a response from a CLI ws. Durable-before-send means the live + * response arrives after storage.put; flushAsync settles the waitUntil + * promise in the test harness. */ +async function sendCliResponse( doInstance: UserConnectionDO, cliWs: MockWS, opts: { id: string; result?: unknown; error?: unknown } -) { +): Promise { const msg = JSON.stringify({ type: 'response', ...opts }); doInstance.webSocketMessage(cliWs as never, msg); + await flushAsync(); } -function createResultWithSerializedBytes(targetBytes: number): { padding: string } { +function createResultWithSerializedBytes(targetBytes: number): { + padding: string; +} { const framingBytes = new TextEncoder().encode(JSON.stringify({ padding: '' })).byteLength; const result = { padding: 'x'.repeat(targetBytes - framingBytes) }; if (new TextEncoder().encode(JSON.stringify(result)).byteLength !== targetBytes) { @@ -432,7 +450,10 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx); await expect( - doInstance.notifySessionEvent({ type: 'session.created', data: { source: 'v1' } } as never) + doInstance.notifySessionEvent({ + type: 'session.created', + data: { source: 'v1' }, + } as never) ).rejects.toThrow(); expect(webWs.send).not.toHaveBeenCalled(); }); @@ -455,7 +476,7 @@ describe('UserConnectionDO', () => { expect(doInstance.hasActiveCliSession('ses_1')).toBe(false); }); - it('reconstructs live session ownership from a hibernated CLI attachment', () => { + it('reconstructs live session ownership from a hibernated CLI attachment', async () => { const { doInstance, mockCtx } = setup(); addCliSocket(mockCtx, 'cli-1', [makeSession('ses_1')]); @@ -468,7 +489,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('heartbeat processing', () => { - it('updates session ownership and persists attachment', () => { + it('updates session ownership and persists attachment', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); addWebSocket(mockCtx, 'web-1'); @@ -481,7 +502,7 @@ describe('UserConnectionDO', () => { expect(att.sessions).toEqual(sessions); }); - it('removes session ownership when session disappears from heartbeat', () => { + it('removes session ownership when session disappears from heartbeat', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); addWebSocket(mockCtx, 'web-1'); @@ -494,7 +515,11 @@ describe('UserConnectionDO', () => { // Verify via command routing: command to s2 should fail (no owner) const webWs2 = addWebSocket(mockCtx, 'web-2'); - sendCommand(doInstance, webWs2, { id: 'cmd-1', command: 'send_message', sessionId: 's2' }); + await sendCommand(doInstance, webWs2, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's2', + }); const resp = parseSent(webWs2); expect(resp).toMatchObject({ type: 'response', @@ -503,7 +528,7 @@ describe('UserConnectionDO', () => { }); }); - it('fails an in-flight command when the session owner changes', () => { + it('fails an in-flight command when the session owner changes', async () => { const { doInstance, mockCtx } = setup(); const firstOwner = addCliSocket(mockCtx, 'cli-1'); const nextOwner = addCliSocket(mockCtx, 'cli-2'); @@ -513,7 +538,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, nextOwner, []); firstOwner.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -523,6 +548,9 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + // Settle the asynchronous durable sweep (finishDurablePendingCommands + // runs inside waitUntil) before asserting the response count. + await flushAsync(); // The owner-change heartbeat broadcasts sessions.heartbeat and also fires // the SESSION_OWNER_CHANGED error response for the in-flight command. The @@ -536,7 +564,10 @@ describe('UserConnectionDO', () => { message: 'Session owner changed', }, }); - sendCliResponse(doInstance, firstOwner, { id: correlationId, result: 'late' }); + await sendCliResponse(doInstance, firstOwner, { + id: correlationId, + result: 'late', + }); // The late sendCliResponse is filtered (the pending entry was already // removed by the owner-change path) so webWs sees exactly the two // messages produced by the owner-change heartbeat itself: the broadcast @@ -544,7 +575,7 @@ describe('UserConnectionDO', () => { expect(webWs.send).toHaveBeenCalledTimes(2); }); - it('replays existing web subscriptions when a session gets a new CLI owner', () => { + it('replays existing web subscriptions when a session gets a new CLI owner', async () => { const { doInstance, mockCtx } = setup(); const cli1 = addCliSocket(mockCtx, 'cli-1'); const cli2 = addCliSocket(mockCtx, 'cli-2'); @@ -568,7 +599,7 @@ describe('UserConnectionDO', () => { expect(cli2Msgs).toContainEqual({ type: 'subscribe', sessionId: 's1' }); }); - it('broadcasts heartbeat to every web socket regardless of subscription', () => { + it('broadcasts heartbeat to every web socket regardless of subscription', async () => { const { doInstance, mockCtx } = setup(); const cli1 = addCliSocket(mockCtx, 'cli-1'); const cli2 = addCliSocket(mockCtx, 'cli-2'); @@ -603,7 +634,7 @@ describe('UserConnectionDO', () => { }); }); - it('delivers one heartbeat per web socket (delivery count equals ws count)', () => { + it('delivers one heartbeat per web socket (delivery count equals ws count)', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const web1 = addWebSocket(mockCtx, 'web-1'); @@ -624,7 +655,9 @@ describe('UserConnectionDO', () => { // The delivered payload is the same shape (including the connectionId). for (const ws of [web1, web2, web3]) { - const sent = parseSent(ws) as { data: { connectionId: string; sessions: unknown[] } }; + const sent = parseSent(ws) as { + data: { connectionId: string; sessions: unknown[] }; + }; expect(sent).toMatchObject({ type: 'system', event: 'sessions.heartbeat', @@ -633,16 +666,20 @@ describe('UserConnectionDO', () => { } }); - it('forwards the CLI-reported protocolVersion to every web socket', () => { + it('forwards the CLI-reported protocolVersion to every web socket', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { protocolVersion: '1' }); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + protocolVersion: '1', + }); sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { protocolVersion: '1' }); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + protocolVersion: '1', + }); expect(parseSent(webWs)).toMatchObject({ type: 'system', @@ -651,7 +688,7 @@ describe('UserConnectionDO', () => { }); }); - it('omits protocolVersion for a legacy CLI that never reports one', () => { + it('omits protocolVersion for a legacy CLI that never reports one', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -666,7 +703,7 @@ describe('UserConnectionDO', () => { expect(sent.data).not.toHaveProperty('protocolVersion'); }); - it('broadcasts removed-session information to every web socket (no subscriber special-case)', () => { + it('broadcasts removed-session information to every web socket (no subscriber special-case)', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); // subWeb subscribed to s1, otherWeb is unrelated — both must learn s1 is gone. @@ -696,7 +733,7 @@ describe('UserConnectionDO', () => { }); }); - it('delivers heartbeat to web sockets that are not subscribed to anything', () => { + it('delivers heartbeat to web sockets that are not subscribed to anything', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -712,7 +749,7 @@ describe('UserConnectionDO', () => { }); }); - it('schedules stale alarm on heartbeat', () => { + it('schedules stale alarm on heartbeat', async () => { const { doInstance, mockCtx, ctx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -720,7 +757,7 @@ describe('UserConnectionDO', () => { expect(ctx.storage.setAlarm).toHaveBeenCalled(); }); - it('sends heartbeat_ack to CLI socket', () => { + it('sends heartbeat_ack to CLI socket', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -738,7 +775,7 @@ describe('UserConnectionDO', () => { // capabilities. // ----------------------------------------------------------------------- - it('projects capabilities.attachments=true on every aggregateSessions row when the owning CLI advertises it', () => { + it('projects capabilities.attachments=true on every aggregateSessions row when the owning CLI advertises it', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -753,7 +790,7 @@ describe('UserConnectionDO', () => { } }); - it('projects the latest connection capabilities onto every sessions.heartbeat row', () => { + it('projects the latest connection capabilities onto every sessions.heartbeat row', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -800,7 +837,7 @@ describe('UserConnectionDO', () => { }); }); - it('omits capabilities from aggregateSessions rows when the latest heartbeat omits the field (legacy CLI)', () => { + it('omits capabilities from aggregateSessions rows when the latest heartbeat omits the field (legacy CLI)', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -816,7 +853,7 @@ describe('UserConnectionDO', () => { expect(rows[0]).not.toHaveProperty('capabilities'); }); - it('omits capabilities from sessions.heartbeat event envelope when the latest heartbeat omits the field', () => { + it('omits capabilities from sessions.heartbeat event envelope when the latest heartbeat omits the field', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -834,14 +871,16 @@ describe('UserConnectionDO', () => { expect(sent.data).not.toHaveProperty('capabilities'); }); - it('flips capabilities.attachments from true to false on the next heartbeat (CLI revocation)', () => { + it('flips capabilities.attachments from true to false on the next heartbeat (CLI revocation)', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: true }, }); - expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ attachments: true }); + expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ + attachments: true, + }); // CLI advertises attachments=false (e.g. feature gated, profile change) sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { @@ -853,7 +892,7 @@ describe('UserConnectionDO', () => { expect(rows[0].capabilities).toEqual({ attachments: false }); }); - it('flips capabilities.attachments from absent to true when a legacy CLI starts advertising it', () => { + it('flips capabilities.attachments from absent to true when a legacy CLI starts advertising it', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -871,23 +910,27 @@ describe('UserConnectionDO', () => { expect(upgraded[0].capabilities).toEqual({ attachments: true }); }); - it('flips capabilities.attachments from false to true on the next heartbeat', () => { + it('flips capabilities.attachments from false to true on the next heartbeat', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: false }, }); - expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ attachments: false }); + expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ + attachments: false, + }); sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: true }, }); - expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ attachments: true }); + expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ + attachments: true, + }); }); - it('projects the same owning-connection capabilities on every session row of a multi-session heartbeat', () => { + it('projects the same owning-connection capabilities on every session row of a multi-session heartbeat', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -901,7 +944,7 @@ describe('UserConnectionDO', () => { expect(rows[1].capabilities).toEqual({ attachments: false }); }); - it('reconstructs capabilities from a hibernated CLI attachment', () => { + it('reconstructs capabilities from a hibernated CLI attachment', async () => { const { doInstance, mockCtx } = setup(); // Pre-existing attachment with capabilities — simulates a socket that // was accepted before the DO was evicted. @@ -993,7 +1036,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('subscribe/unsubscribe', () => { - it('sends subscribe to owning CLI when web subscribes', () => { + it('sends subscribe to owning CLI when web subscribes', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1009,7 +1052,7 @@ describe('UserConnectionDO', () => { expect(parseSent(cliWs)).toEqual({ type: 'subscribe', sessionId: 's1' }); }); - it('sends the active session list when web subscribes after the socket is open', () => { + it('sends the active session list when web subscribes after the socket is open', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1021,11 +1064,20 @@ describe('UserConnectionDO', () => { expect(parseSent(webWs)).toEqual({ type: 'system', event: 'sessions.list', - data: { sessions: [{ id: 's1', status: 'busy', title: 'Fix bug', connectionId: 'cli-1' }] }, + data: { + sessions: [ + { + id: 's1', + status: 'busy', + title: 'Fix bug', + connectionId: 'cli-1', + }, + ], + }, }); }); - it('broadcasts subscribe to all CLIs when no owner found', () => { + it('broadcasts subscribe to all CLIs when no owner found', async () => { const { doInstance, mockCtx } = setup(); const cli1 = addCliSocket(mockCtx, 'cli-1'); const cli2 = addCliSocket(mockCtx, 'cli-2'); @@ -1042,7 +1094,7 @@ describe('UserConnectionDO', () => { expect(parseSent(cli2)).toEqual({ type: 'subscribe', sessionId: 's1' }); }); - it('duplicate subscribe is idempotent for attachment', () => { + it('duplicate subscribe is idempotent for attachment', async () => { const { doInstance, mockCtx } = setup(); addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1050,11 +1102,13 @@ describe('UserConnectionDO', () => { sendSubscribe(doInstance, webWs, 's1'); sendSubscribe(doInstance, webWs, 's1'); - const att = webWs.deserializeAttachment() as { subscribedSessions: string[] }; + const att = webWs.deserializeAttachment() as { + subscribedSessions: string[]; + }; expect(att.subscribedSessions).toEqual(['s1']); }); - it('unsubscribe sends to CLI when last subscriber leaves', () => { + it('unsubscribe sends to CLI when last subscriber leaves', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1068,10 +1122,13 @@ describe('UserConnectionDO', () => { sendUnsubscribe(doInstance, webWs, 's1'); expect(cliWs.send).toHaveBeenCalledTimes(1); - expect(parseSent(cliWs)).toEqual({ type: 'unsubscribe', sessionId: 's1' }); + expect(parseSent(cliWs)).toEqual({ + type: 'unsubscribe', + sessionId: 's1', + }); }); - it('unsubscribe does not send to CLI when other subscribers remain', () => { + it('unsubscribe does not send to CLI when other subscribers remain', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const web1 = addWebSocket(mockCtx, 'web-1'); @@ -1091,7 +1148,10 @@ describe('UserConnectionDO', () => { // Unsubscribe second — CLI SHOULD get unsubscribe sendUnsubscribe(doInstance, web2, 's1'); expect(cliWs.send).toHaveBeenCalledTimes(1); - expect(parseSent(cliWs)).toEqual({ type: 'unsubscribe', sessionId: 's1' }); + expect(parseSent(cliWs)).toEqual({ + type: 'unsubscribe', + sessionId: 's1', + }); }); }); @@ -1100,7 +1160,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('viewer liveness', () => { - it('replies to a viewer ping with the matching nonce only', () => { + it('replies to a viewer ping with the matching nonce only', async () => { const { doInstance, mockCtx, ctx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'viewer-1'); @@ -1157,7 +1217,7 @@ describe('UserConnectionDO', () => { expect(mockCtx.sockets.filter(socket => socket._tags.includes('web'))).toHaveLength(2); }); - it('does not migrate old subscriptions when replacing a viewer', () => { + it('does not migrate old subscriptions when replacing a viewer', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); @@ -1170,7 +1230,10 @@ describe('UserConnectionDO', () => { const newWeb = connectWebSocket(doInstance, 'viewer-1'); expect(cliWs.send).toHaveBeenCalledTimes(1); - expect(parseSent(cliWs)).toEqual({ type: 'unsubscribe', sessionId: 's1' }); + expect(parseSent(cliWs)).toEqual({ + type: 'unsubscribe', + sessionId: 's1', + }); expect(newWeb.deserializeAttachment()).toEqual({ role: 'web', connectionId: 'viewer-1', @@ -1178,7 +1241,7 @@ describe('UserConnectionDO', () => { }); }); - it('ignores messages from a viewer that has been replaced', () => { + it('ignores messages from a viewer that has been replaced', async () => { const { doInstance } = setup(); const oldWeb = connectWebSocket(doInstance, 'viewer-1'); connectWebSocket(doInstance, 'viewer-1'); @@ -1215,7 +1278,7 @@ describe('UserConnectionDO', () => { expect(secondWeb.send).toHaveBeenCalledTimes(1); }); - it('does not replace a CLI socket when a viewer connectionId collides', () => { + it('does not replace a CLI socket when a viewer connectionId collides', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'shared-id'); @@ -1256,8 +1319,15 @@ describe('UserConnectionDO', () => { // Session no longer routable const web2 = addWebSocket(mockCtx, 'web-2'); - sendCommand(doInstance, web2, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); - expect(parseSent(web2)).toMatchObject({ type: 'response', error: 'Session owner not found' }); + await sendCommand(doInstance, web2, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + expect(parseSent(web2)).toMatchObject({ + type: 'response', + error: 'Session owner not found', + }); }); it('sends error responses for pending commands on disconnect', async () => { @@ -1268,7 +1338,11 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // Send command from web - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); webWs.send.mockClear(); // CLI disconnects @@ -1280,7 +1354,11 @@ describe('UserConnectionDO', () => { const errorResp = msgs.find( (m: Record) => m.type === 'response' && m.id === 'cmd-1' ); - expect(errorResp).toMatchObject({ type: 'response', id: 'cmd-1', error: 'CLI disconnected' }); + expect(errorResp).toMatchObject({ + type: 'response', + id: 'cmd-1', + error: 'CLI disconnected', + }); }); it('reports owner change when an owner-fenced command target disconnects', async () => { @@ -1289,7 +1367,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -1311,14 +1389,14 @@ describe('UserConnectionDO', () => { }); }); - it('fails pending commands as soon as their target socket is replaced', () => { + it('fails pending commands as soon as their target socket is replaced', async () => { const { doInstance, mockCtx } = setup(); const firstCli = connectCliSocket(doInstance, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, firstCli, [makeSession('s1')]); firstCli.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -1328,6 +1406,11 @@ describe('UserConnectionDO', () => { connectCliSocket(doInstance, 'cli-1'); + // failPendingCommandsForSocket runs inside ctx.waitUntil from + // closeStaleSocket. Drain microtasks so the durable write and + // live response settle before the assertion. + await flushAsync(); + expect(firstCli.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); expect(parseSent(webWs)).toEqual({ type: 'response', @@ -1348,7 +1431,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, []); // Send command routed by connectionId (no sessionId) - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-conn', command: 'send_message', connectionId: 'cli-1', @@ -1378,7 +1461,10 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, []); // Send command with no sessionId or connectionId (fallback routing) - sendCommand(doInstance, webWs, { id: 'cmd-fallback', command: 'send_message' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-fallback', + command: 'send_message', + }); webWs.send.mockClear(); mockCtx.removeSocket(cliWs); @@ -1438,7 +1524,11 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); // Web sends a command targeting s1 — should route to cli2 (the replacement) - sendCommand(doInstance, webWs, { id: 'cmd-new', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-new', + command: 'send_message', + sessionId: 's1', + }); expect(cli2.send).toHaveBeenCalled(); const correlationId = getCorrelationId(cli2); @@ -1455,7 +1545,7 @@ describe('UserConnectionDO', () => { // cli2 responds successfully webWs.send.mockClear(); - sendCliResponse(doInstance, cli2, { id: correlationId, result: 'ok' }); + await sendCliResponse(doInstance, cli2, { id: correlationId, result: 'ok' }); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ @@ -1473,7 +1563,11 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cli1, [makeSession('s1')]); // Web sends a command that gets forwarded to cli1 - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); webWs.send.mockClear(); // CLI2 connects with the same connectionId (reconnect) @@ -1667,7 +1761,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('web disconnect', () => { - it('removes from all subscription sets', () => { + it('removes from all subscription sets', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1690,7 +1784,7 @@ describe('UserConnectionDO', () => { // No web sockets to receive the event — no crash = success }); - it('sends unsubscribe to CLI when last subscriber leaves', () => { + it('sends unsubscribe to CLI when last subscriber leaves', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1708,21 +1802,25 @@ describe('UserConnectionDO', () => { expect(unsub).toEqual({ type: 'unsubscribe', sessionId: 's1' }); }); - it('cleans up pending commands from disconnecting web socket', () => { + it('cleans up pending commands from disconnecting web socket', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); const correlationId = getCorrelationId(cliWs); mockCtx.removeSocket(webWs); disconnectWeb(doInstance, webWs); // CLI sends response with correlationId, but the pending command is gone — no crash - sendCliResponse(doInstance, cliWs, { id: correlationId, result: 'ok' }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: 'ok' }); }); }); @@ -1731,7 +1829,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('command routing', () => { - it('routes web command to correct CLI by sessionId', () => { + it('routes web command to correct CLI by sessionId', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1739,7 +1837,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -1758,19 +1856,26 @@ describe('UserConnectionDO', () => { expect(sent.id).not.toBe('cmd-1'); }); - it('routes CLI response to correct web socket with original id', () => { + it('routes CLI response to correct web socket with original id', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { id: correlationId, result: { success: true } }); + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { success: true }, + }); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ @@ -1780,18 +1885,22 @@ describe('UserConnectionDO', () => { }); }); - it('sanitizes a relay-shaped CLI error before forwarding it to web', () => { + it('sanitizes a relay-shaped CLI error before forwarding it to web', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: { source: 'relay', @@ -1810,18 +1919,22 @@ describe('UserConnectionDO', () => { }); }); - it('preserves CLI string errors for old-CLI compatibility', () => { + it('preserves CLI string errors for old-CLI compatibility', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_models', }); @@ -1833,7 +1946,7 @@ describe('UserConnectionDO', () => { }); }); - it('accepts a pending response only from the targeted CLI socket', () => { + it('accepts a pending response only from the targeted CLI socket', async () => { const { doInstance, mockCtx } = setup(); const targetCli = addCliSocket(mockCtx, 'cli-1'); const otherCli = addCliSocket(mockCtx, 'cli-2'); @@ -1842,18 +1955,32 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, targetCli, [makeSession('s1')]); sendHeartbeat(doInstance, otherCli, []); targetCli.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); const correlationId = getCorrelationId(targetCli); webWs.send.mockClear(); - sendCliResponse(doInstance, otherCli, { id: correlationId, result: 'wrong-owner' }); + await sendCliResponse(doInstance, otherCli, { + id: correlationId, + result: 'wrong-owner', + }); expect(webWs.send).not.toHaveBeenCalled(); - sendCliResponse(doInstance, targetCli, { id: correlationId, result: 'ok' }); - expect(parseSent(webWs)).toEqual({ type: 'response', id: 'cmd-1', result: 'ok' }); + await sendCliResponse(doInstance, targetCli, { + id: correlationId, + result: 'ok', + }); + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + result: 'ok', + }); }); - it('rejects a duplicate in-flight list_models request for the same viewer session and owner', () => { + it('rejects a duplicate in-flight list_models request for the same viewer session and owner', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1861,13 +1988,13 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', connectionId: 'cli-1', }); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-2', command: 'list_models', sessionId: 's1', @@ -1886,7 +2013,7 @@ describe('UserConnectionDO', () => { expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(1); }); - it('expires pending commands before handling another command', () => { + it('expires pending commands before handling another command', async () => { const now = 1_000_000; vi.spyOn(Date, 'now').mockReturnValue(now); const { doInstance, mockCtx } = setup(); @@ -1896,7 +2023,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -1904,7 +2031,7 @@ describe('UserConnectionDO', () => { }); vi.mocked(Date.now).mockReturnValue(now + 35_001); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-2', command: 'list_models', sessionId: 's1', @@ -1923,7 +2050,7 @@ describe('UserConnectionDO', () => { expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(2); }); - it('does not postpone pending-command expiry when heartbeats reschedule the alarm', () => { + it('does not postpone pending-command expiry when heartbeats reschedule the alarm', async () => { const now = 1_000_000; vi.spyOn(Date, 'now').mockReturnValue(now); const { doInstance, mockCtx, ctx } = setup(); @@ -1931,7 +2058,11 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); ctx.storage.setAlarm.mockClear(); vi.mocked(Date.now).mockReturnValue(now + 20_000); @@ -1949,7 +2080,11 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); const correlationId = getCorrelationId(cliWs); vi.mocked(Date.now).mockReturnValue(now + 34_000); @@ -1968,11 +2103,11 @@ describe('UserConnectionDO', () => { message: 'Command expired', }, }); - sendCliResponse(doInstance, cliWs, { id: correlationId, result: 'late' }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: 'late' }); expect(webWs.send).toHaveBeenCalledTimes(1); }); - it('rejects commands after reaching the global pending-command cap', () => { + it('rejects commands after reaching the global pending-command cap', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -1981,14 +2116,14 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); webWs.send.mockClear(); for (let index = 0; index < 128; index++) { - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: `cmd-${index}`, command: 'send_message', sessionId: 's1', }); } - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-over-cap', command: 'send_message', sessionId: 's1', @@ -2006,14 +2141,14 @@ describe('UserConnectionDO', () => { expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(128); }); - it('accepts a list_models result at exactly 512 KiB', () => { + it('accepts a list_models result at exactly 512 KiB', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -2023,19 +2158,23 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES); - sendCliResponse(doInstance, cliWs, { id: correlationId, result }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result }); - expect(parseSent(webWs)).toEqual({ type: 'response', id: 'cmd-1', result }); + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + result, + }); }); - it('rejects a list_models result one byte over 512 KiB', () => { + it('rejects a list_models result one byte over 512 KiB', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -2044,7 +2183,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1), }); @@ -2060,14 +2199,14 @@ describe('UserConnectionDO', () => { }); }); - it('rejects a multibyte list_models result over 512 KiB', () => { + it('rejects a multibyte list_models result over 512 KiB', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -2076,7 +2215,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: createUtf8OversizedResult(), }); @@ -2092,11 +2231,11 @@ describe('UserConnectionDO', () => { }); }); - it('returns error when CLI not found for session', () => { + it('returns error when CLI not found for session', async () => { const { doInstance, mockCtx } = setup(); const webWs = addWebSocket(mockCtx, 'web-1'); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 'unknown-session', @@ -2110,7 +2249,7 @@ describe('UserConnectionDO', () => { }); }); - it('rejects a stale expected session owner without forwarding', () => { + it('rejects a stale expected session owner without forwarding', async () => { const { doInstance, mockCtx } = setup(); const currentOwner = addCliSocket(mockCtx, 'cli-1'); const staleOwner = addCliSocket(mockCtx, 'cli-2'); @@ -2122,7 +2261,7 @@ describe('UserConnectionDO', () => { staleOwner.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -2142,7 +2281,7 @@ describe('UserConnectionDO', () => { expect(staleOwner.send).not.toHaveBeenCalled(); }); - it('routes command by connectionId to specific CLI', () => { + it('routes command by connectionId to specific CLI', async () => { const { doInstance, mockCtx } = setup(); const cli1 = addCliSocket(mockCtx, 'cli-1'); const cli2 = addCliSocket(mockCtx, 'cli-2'); @@ -2154,7 +2293,7 @@ describe('UserConnectionDO', () => { cli1.send.mockClear(); cli2.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', connectionId: 'cli-2', @@ -2164,7 +2303,7 @@ describe('UserConnectionDO', () => { expect(cli2.send).toHaveBeenCalledTimes(1); }); - it('two web sockets with the same command id each get the correct response', () => { + it('two web sockets with the same command id each get the correct response', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const web1 = addWebSocket(mockCtx, 'web-1'); @@ -2174,10 +2313,18 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); // Both web sockets send commands with the same id - sendCommand(doInstance, web1, { id: 'dup-id', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, web1, { + id: 'dup-id', + command: 'send_message', + sessionId: 's1', + }); const corr1 = getCorrelationId(cliWs, 0); - sendCommand(doInstance, web2, { id: 'dup-id', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, web2, { + id: 'dup-id', + command: 'send_message', + sessionId: 's1', + }); const corr2 = getCorrelationId(cliWs, 1); expect(corr1).not.toBe(corr2); @@ -2185,14 +2332,22 @@ describe('UserConnectionDO', () => { web1.send.mockClear(); web2.send.mockClear(); - sendCliResponse(doInstance, cliWs, { id: corr1, result: 'result-1' }); - sendCliResponse(doInstance, cliWs, { id: corr2, result: 'result-2' }); + await sendCliResponse(doInstance, cliWs, { id: corr1, result: 'result-1' }); + await sendCliResponse(doInstance, cliWs, { id: corr2, result: 'result-2' }); - expect(parseSent(web1)).toEqual({ type: 'response', id: 'dup-id', result: 'result-1' }); - expect(parseSent(web2)).toEqual({ type: 'response', id: 'dup-id', result: 'result-2' }); + expect(parseSent(web1)).toEqual({ + type: 'response', + id: 'dup-id', + result: 'result-1', + }); + expect(parseSent(web2)).toEqual({ + type: 'response', + id: 'dup-id', + result: 'result-2', + }); }); - it('routes to first CLI when no sessionId or connectionId given', () => { + it('routes to first CLI when no sessionId or connectionId given', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -2200,7 +2355,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message' }); + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message' }); expect(cliWs.send).toHaveBeenCalledTimes(1); }); }); @@ -2225,7 +2380,7 @@ describe('UserConnectionDO', () => { 'exit_cli', ]; - it('forwards every allowed viewer command to the owning CLI', () => { + it('forwards every allowed viewer command to the owning CLI', async () => { for (const command of ALLOWED) { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -2234,7 +2389,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command, sessionId: 's1', @@ -2243,13 +2398,17 @@ describe('UserConnectionDO', () => { expect(cliWs.send).toHaveBeenCalledTimes(1); const sent = parseSent(cliWs) as Record; - expect(sent).toMatchObject({ type: 'command', command, sessionId: 's1' }); + expect(sent).toMatchObject({ + type: 'command', + command, + sessionId: 's1', + }); expect(typeof sent.id).toBe('string'); expect(sent.id).not.toBe('cmd-1'); } }); - it('rejects a non-allowlisted command with structured COMMAND_NOT_ALLOWED', () => { + it('rejects a non-allowlisted command with structured COMMAND_NOT_ALLOWED', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -2258,7 +2417,11 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'eval', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'eval', + sessionId: 's1', + }); expect(parseSent(webWs)).toEqual({ type: 'response', @@ -2272,7 +2435,7 @@ describe('UserConnectionDO', () => { expect(cliWs.send).not.toHaveBeenCalled(); }); - it('rejects a non-allowlisted command even when targeting a known session owner via connectionId', () => { + it('rejects a non-allowlisted command even when targeting a known session owner via connectionId', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -2281,7 +2444,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'shell', sessionId: 's1', @@ -2301,11 +2464,11 @@ describe('UserConnectionDO', () => { expect(cliWs.send).not.toHaveBeenCalled(); }); - it('rejects a non-allowlisted command with an unknown session before owner resolution', () => { + it('rejects a non-allowlisted command with an unknown session before owner resolution', async () => { const { doInstance, mockCtx } = setup(); const webWs = addWebSocket(mockCtx, 'web-1'); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'eval', sessionId: 'unknown-session', @@ -2323,7 +2486,7 @@ describe('UserConnectionDO', () => { }); }); - it('does not allocate a pending entry or forward a disallowed command', () => { + it('does not allocate a pending entry or forward a disallowed command', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -2332,7 +2495,11 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'eval', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'eval', + sessionId: 's1', + }); const sent = parseSent(webWs) as Record; expect(sent).toMatchObject({ type: 'response', id: 'cmd-1' }); expect(sent.id).toBe('cmd-1'); @@ -2341,11 +2508,11 @@ describe('UserConnectionDO', () => { // not have allocated a pending slot, so a follow-up CLI response for a // fabricated correlation id is a no-op. expect(cliWs.send).not.toHaveBeenCalled(); - sendCliResponse(doInstance, cliWs, { id: 'fabricated', result: 'noop' }); + await sendCliResponse(doInstance, cliWs, { id: 'fabricated', result: 'noop' }); expect(webWs.send).toHaveBeenCalledTimes(1); }); - it('still rejects an owner-fenced allowed command with SESSION_OWNER_CHANGED', () => { + it('still rejects an owner-fenced allowed command with SESSION_OWNER_CHANGED', async () => { const { doInstance, mockCtx } = setup(); const currentOwner = addCliSocket(mockCtx, 'cli-1'); const staleOwner = addCliSocket(mockCtx, 'cli-2'); @@ -2357,7 +2524,7 @@ describe('UserConnectionDO', () => { staleOwner.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -2383,7 +2550,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('list_commands dedupe and size cap', () => { - it('rejects a duplicate in-flight list_commands request for the same viewer session and owner', () => { + it('rejects a duplicate in-flight list_commands request for the same viewer session and owner', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -2391,13 +2558,13 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', connectionId: 'cli-1', }); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-2', command: 'list_commands', sessionId: 's1', @@ -2416,20 +2583,20 @@ describe('UserConnectionDO', () => { expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(1); }); - it('treats list_models and list_commands as distinct for dedupe purposes', () => { + it('treats list_models and list_commands as distinct for dedupe purposes', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', connectionId: 'cli-1', }); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-2', command: 'list_commands', sessionId: 's1', @@ -2439,14 +2606,14 @@ describe('UserConnectionDO', () => { expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(2); }); - it('accepts a list_commands result at exactly 512 KiB', () => { + it('accepts a list_commands result at exactly 512 KiB', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2456,19 +2623,23 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES); - sendCliResponse(doInstance, cliWs, { id: correlationId, result }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result }); - expect(parseSent(webWs)).toEqual({ type: 'response', id: 'cmd-1', result }); + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + result, + }); }); - it('rejects a list_commands result one byte over 512 KiB', () => { + it('rejects a list_commands result one byte over 512 KiB', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2477,7 +2648,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1), }); @@ -2493,14 +2664,14 @@ describe('UserConnectionDO', () => { }); }); - it('rejects a multibyte list_commands result over 512 KiB', () => { + it('rejects a multibyte list_commands result over 512 KiB', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2509,7 +2680,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: createUtf8OversizedResult(), }); @@ -2531,14 +2702,14 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('old CLI upgrade-required mapping', () => { - it('maps "unknown command: list_commands" to CLI_UPGRADE_REQUIRED with slash message', () => { + it('maps "unknown command: list_commands" to CLI_UPGRADE_REQUIRED with slash message', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2547,7 +2718,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_commands', }); @@ -2563,14 +2734,14 @@ describe('UserConnectionDO', () => { }); }); - it('maps "unknown command: send_command" to CLI_UPGRADE_REQUIRED with slash message', () => { + it('maps "unknown command: send_command" to CLI_UPGRADE_REQUIRED with slash message', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_command', sessionId: 's1', @@ -2580,7 +2751,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: send_command', }); @@ -2596,14 +2767,14 @@ describe('UserConnectionDO', () => { }); }); - it('maps "unknown command: exit_cli" to CLI_UPGRADE_REQUIRED with slash message', () => { + it('maps "unknown command: exit_cli" to CLI_UPGRADE_REQUIRED with slash message', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -2613,7 +2784,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: exit_cli', }); @@ -2629,14 +2800,14 @@ describe('UserConnectionDO', () => { }); }); - it('maps "unknown command: create_session" to CLI_UPGRADE_REQUIRED with create_session message', () => { + it('maps "unknown command: create_session" to CLI_UPGRADE_REQUIRED with create_session message', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'create_session', connectionId: 'cli-1', @@ -2645,7 +2816,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: create_session', }); @@ -2662,18 +2833,22 @@ describe('UserConnectionDO', () => { }); }); - it('preserves "unknown command: list_models" because list_models is not in the upgrade-required set', () => { + it('preserves "unknown command: list_models" because list_models is not in the upgrade-required set', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1' }); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_models', }); @@ -2685,14 +2860,14 @@ describe('UserConnectionDO', () => { }); }); - it('preserves an unrelated CLI string error for send_command', () => { + it('preserves an unrelated CLI string error for send_command', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_command', sessionId: 's1', @@ -2701,7 +2876,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'session not ready', }); @@ -2713,14 +2888,14 @@ describe('UserConnectionDO', () => { }); }); - it('does not match a longer error that merely starts with "unknown command: list_commands"', () => { + it('does not match a longer error that merely starts with "unknown command: list_commands"', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2729,7 +2904,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_commands: try again', }); @@ -2742,14 +2917,14 @@ describe('UserConnectionDO', () => { }); }); - it('preserves a longer exit_cli unknown-command error', () => { + it('preserves a longer exit_cli unknown-command error', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -2759,7 +2934,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: exit_cli: session not ready', }); @@ -2779,7 +2954,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('send_command / create_session negative coverage', () => { - it('forwards two in-flight same-owner/same-session send_command requests without deduping', () => { + it('forwards two in-flight same-owner/same-session send_command requests without deduping', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -2787,14 +2962,14 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_command', sessionId: 's1', connectionId: 'cli-1', data: { command: 'init' }, }); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-2', command: 'send_command', sessionId: 's1', @@ -2820,7 +2995,7 @@ describe('UserConnectionDO', () => { expect(webWs.send).not.toHaveBeenCalled(); }); - it('forwards two in-flight create_session requests without deduping', () => { + it('forwards two in-flight create_session requests without deduping', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -2828,13 +3003,13 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'create_session', connectionId: 'cli-1', data: { title: 'First session' }, }); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-2', command: 'create_session', connectionId: 'cli-1', @@ -2857,14 +3032,14 @@ describe('UserConnectionDO', () => { expect(webWs.send).not.toHaveBeenCalled(); }); - it('relays a send_command result over 512 KiB unchanged', () => { + it('relays a send_command result over 512 KiB unchanged', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_command', sessionId: 's1', @@ -2875,7 +3050,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); - sendCliResponse(doInstance, cliWs, { id: correlationId, result }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ type: 'response', @@ -2884,14 +3059,14 @@ describe('UserConnectionDO', () => { }); }); - it('relays a create_session result over 512 KiB unchanged', () => { + it('relays a create_session result over 512 KiB unchanged', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'create_session', connectionId: 'cli-1', @@ -2901,7 +3076,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); - sendCliResponse(doInstance, cliWs, { id: correlationId, result }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ type: 'response', @@ -2919,15 +3094,24 @@ describe('UserConnectionDO', () => { it.each([ { label: 'missing sessionId', input: { data: { protocolVersion: 1 } } }, { label: 'missing data', input: { sessionId: 's1' } }, - { label: 'wrong protocol version', input: { sessionId: 's1', data: { protocolVersion: 2 } } }, + { + label: 'wrong protocol version', + input: { sessionId: 's1', data: { protocolVersion: 2 } }, + }, { label: 'extra data field', input: { sessionId: 's1', data: { protocolVersion: 1, extra: true } }, }, { label: 'null data', input: { sessionId: 's1', data: null } }, - { label: 'array data', input: { sessionId: 's1', data: [{ protocolVersion: 1 }] } }, - { label: 'primitive data', input: { sessionId: 's1', data: 'protocolVersion=1' } }, - ])('rejects $label before routing or pending allocation', ({ input }) => { + { + label: 'array data', + input: { sessionId: 's1', data: [{ protocolVersion: 1 }] }, + }, + { + label: 'primitive data', + input: { sessionId: 's1', data: 'protocolVersion=1' }, + }, + ])('rejects $label before routing or pending allocation', async ({ input }) => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -2935,7 +3119,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', ...input, @@ -2954,7 +3138,7 @@ describe('UserConnectionDO', () => { expect(Reflect.get(doInstance, 'pendingCommands')).toEqual(new Map()); }); - it('routes exit_cli to the selected session owner with its data unchanged', () => { + it('routes exit_cli to the selected session owner with its data unchanged', async () => { const { doInstance, mockCtx } = setup(); const selectedOwner = addCliSocket(mockCtx, 'cli-1'); const otherCli = addCliSocket(mockCtx, 'cli-2'); @@ -2965,7 +3149,7 @@ describe('UserConnectionDO', () => { selectedOwner.send.mockClear(); otherCli.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -2984,7 +3168,7 @@ describe('UserConnectionDO', () => { expect(otherCli.send).not.toHaveBeenCalled(); }); - it('rejects exit_cli when the selected owner snapshot is stale', () => { + it('rejects exit_cli when the selected owner snapshot is stale', async () => { const { doInstance, mockCtx } = setup(); const currentOwner = addCliSocket(mockCtx, 'cli-1'); const staleOwner = addCliSocket(mockCtx, 'cli-2'); @@ -2996,7 +3180,7 @@ describe('UserConnectionDO', () => { staleOwner.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3017,11 +3201,11 @@ describe('UserConnectionDO', () => { expect(staleOwner.send).not.toHaveBeenCalled(); }); - it('rejects exit_cli when the session has no owner', () => { + it('rejects exit_cli when the session has no owner', async () => { const { doInstance, mockCtx } = setup(); const webWs = addWebSocket(mockCtx, 'web-1'); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3035,7 +3219,7 @@ describe('UserConnectionDO', () => { }); }); - it('does not dedupe concurrent exit_cli requests', () => { + it('does not dedupe concurrent exit_cli requests', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3043,14 +3227,14 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', connectionId: 'cli-1', data: { protocolVersion: 1 }, }); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-2', command: 'exit_cli', sessionId: 's1', @@ -3064,7 +3248,7 @@ describe('UserConnectionDO', () => { expect(webWs.send).not.toHaveBeenCalled(); }); - it('relays an exit_cli result over 512 KiB unchanged', () => { + it('relays an exit_cli result over 512 KiB unchanged', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3072,7 +3256,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3082,12 +3266,16 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); - sendCliResponse(doInstance, cliWs, { id: correlationId, result }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result }); - expect(parseSent(webWs)).toEqual({ type: 'response', id: 'cmd-1', result }); + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + result, + }); }); - it('resolves exit_cli successfully when heartbeat drops the session', () => { + it('resolves exit_cli successfully when heartbeat drops the session', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3095,7 +3283,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3108,6 +3296,10 @@ describe('UserConnectionDO', () => { // Exit's own effect: CLI heartbeat no longer lists the session. sendHeartbeat(doInstance, cliWs, []); + // failPendingCommandsForOwnerChange runs inside ctx.waitUntil. + // Drain microtasks so the durable write and live response settle. + await flushAsync(); + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ type: 'response', id: 'cmd-1', @@ -3127,7 +3319,7 @@ describe('UserConnectionDO', () => { const responsesBeforeLateAck = allSent(webWs).filter( m => m.type === 'response' && m.id === 'cmd-1' ).length; - sendCliResponse(doInstance, cliWs, { id: correlationId, result: {} }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: {} }); expect(allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-1')).toHaveLength( responsesBeforeLateAck ); @@ -3139,7 +3331,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3160,7 +3352,7 @@ describe('UserConnectionDO', () => { it.each(['list_models', 'send_message'] as const)( 'still fails %s with SESSION_OWNER_CHANGED when heartbeat drops the session', - command => { + async command => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3168,7 +3360,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command, sessionId: 's1', @@ -3178,6 +3370,9 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, []); + // failPendingCommandsForOwnerChange runs inside ctx.waitUntil. + await flushAsync(); + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ type: 'response', id: 'cmd-1', @@ -3190,7 +3385,7 @@ describe('UserConnectionDO', () => { } ); - it('still fails exit_cli with SESSION_OWNER_CHANGED on genuine takeover', () => { + it('still fails exit_cli with SESSION_OWNER_CHANGED on genuine takeover', async () => { const { doInstance, mockCtx } = setup(); const firstOwner = addCliSocket(mockCtx, 'cli-1'); const nextOwner = addCliSocket(mockCtx, 'cli-2'); @@ -3200,7 +3395,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, nextOwner, []); firstOwner.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3211,6 +3406,9 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + // failPendingCommandsForOwnerChange runs inside ctx.waitUntil. + await flushAsync(); + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ type: 'response', id: 'cmd-1', @@ -3222,14 +3420,14 @@ describe('UserConnectionDO', () => { }); }); - it('still fails exit_cli with SESSION_OWNER_CHANGED when socket is replaced by reconnect', () => { + it('still fails exit_cli with SESSION_OWNER_CHANGED when socket is replaced by reconnect', async () => { const { doInstance, mockCtx } = setup(); const firstCli = connectCliSocket(doInstance, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, firstCli, [makeSession('s1')]); firstCli.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3240,6 +3438,10 @@ describe('UserConnectionDO', () => { connectCliSocket(doInstance, 'cli-1'); + // failPendingCommandsForSocket runs inside ctx.waitUntil from + // closeStaleSocket. Drain microtasks before asserting. + await flushAsync(); + expect(firstCli.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); expect(parseSent(webWs)).toEqual({ type: 'response', @@ -3258,7 +3460,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('CLI event forwarding', () => { - it('forwards events to subscribed web sockets only', () => { + it('forwards events to subscribed web sockets only', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const subWeb = addWebSocket(mockCtx, 'web-sub'); @@ -3288,7 +3490,7 @@ describe('UserConnectionDO', () => { expect(otherWeb.send).not.toHaveBeenCalled(); }); - it('sends child events to both direct child subscribers and parent subscribers', () => { + it('sends child events to both direct child subscribers and parent subscribers', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const parentWeb = addWebSocket(mockCtx, 'web-parent'); @@ -3322,7 +3524,7 @@ describe('UserConnectionDO', () => { expect(parseSent(childWeb)).toEqual(expected); }); - it('deduplicates when same socket subscribes to both child and parent', () => { + it('deduplicates when same socket subscribes to both child and parent', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3345,7 +3547,7 @@ describe('UserConnectionDO', () => { expect(webWs.send).toHaveBeenCalledTimes(1); }); - it('routes child event to parent session subscribers via parentSessionId', () => { + it('routes child event to parent session subscribers via parentSessionId', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3374,7 +3576,7 @@ describe('UserConnectionDO', () => { }); }); - it('drops child event when neither sessionId nor parentSessionId has subscribers', () => { + it('drops child event when neither sessionId nor parentSessionId has subscribers', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3396,7 +3598,7 @@ describe('UserConnectionDO', () => { expect(webWs.send).not.toHaveBeenCalled(); }); - it('events without parentSessionId still route normally (backward compat)', () => { + it('events without parentSessionId still route normally (backward compat)', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3422,7 +3624,7 @@ describe('UserConnectionDO', () => { }); }); - it('child event does not include parentSessionId when not set', () => { + it('child event does not include parentSessionId when not set', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -3449,7 +3651,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('broadcast resilience', () => { - it('one closed socket does not abort send to other web sockets', () => { + it('one closed socket does not abort send to other web sockets', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const failWeb = addWebSocket(mockCtx, 'web-fail'); @@ -3480,7 +3682,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('ensureState (hibernation recovery)', () => { - it('reconstructs sessionOwners and connectionSessions from CLI attachments', () => { + it('reconstructs sessionOwners and connectionSessions from CLI attachments', async () => { const { doInstance, mockCtx } = setup(); // Simulate hibernation: sockets exist with pre-set attachments @@ -3493,17 +3695,24 @@ describe('UserConnectionDO', () => { // Verify state was reconstructed by routing a command const web2 = addWebSocket(mockCtx, 'web-2'); - sendCommand(doInstance, web2, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + await sendCommand(doInstance, web2, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); // Should route to cli-1 (not "Session owner not found") const cliWs = mockCtx.sockets.find(s => s._tags.includes('cli')); expect(cliWs?.send).toHaveBeenCalled(); const cliMsgs = allSent(cliWs!); const cmdMsg = cliMsgs.find((m: Record) => m.type === 'command'); - expect(cmdMsg).toMatchObject({ type: 'command', command: 'send_message' }); + expect(cmdMsg).toMatchObject({ + type: 'command', + command: 'send_message', + }); }); - it('reconstructs webSubscriptions from web attachments', () => { + it('reconstructs webSubscriptions from web attachments', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1', [makeSession('s1')]); @@ -3527,7 +3736,7 @@ describe('UserConnectionDO', () => { }); }); - it('does not restore subscriptions from a viewer already replaced before hibernation', () => { + it('does not restore subscriptions from a viewer already replaced before hibernation', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1', [makeSession('s1')]); const replacedWeb = addWebSocket(mockCtx, 'web-old', ['s1']); @@ -3540,7 +3749,12 @@ describe('UserConnectionDO', () => { doInstance.webSocketMessage( cliWs as never, - JSON.stringify({ type: 'event', sessionId: 's1', event: 'test', data: {} }) + JSON.stringify({ + type: 'event', + sessionId: 's1', + event: 'test', + data: {}, + }) ); expect(replacedWeb.send).not.toHaveBeenCalled(); @@ -3552,7 +3766,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('getActiveSessions', () => { - it('returns sessions from live CLI connections', () => { + it('returns sessions from live CLI connections', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -3568,7 +3782,7 @@ describe('UserConnectionDO', () => { ]); }); - it('includes the CLI-reported protocolVersion on each session row', () => { + it('includes the CLI-reported protocolVersion on each session row', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -3578,11 +3792,17 @@ describe('UserConnectionDO', () => { const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: 's1', status: 'busy', title: 'Fix bug', connectionId: 'cli-1', protocolVersion: '1' }, + { + id: 's1', + status: 'busy', + title: 'Fix bug', + connectionId: 'cli-1', + protocolVersion: '1', + }, ]); }); - it('excludes sessions from stale connections without live sockets', () => { + it('excludes sessions from stale connections without live sockets', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); @@ -3594,7 +3814,7 @@ describe('UserConnectionDO', () => { expect(result).toEqual([]); }); - it('excludes child sessions reported with parentSessionId in heartbeat', () => { + it('excludes child sessions reported with parentSessionId in heartbeat', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -3605,11 +3825,16 @@ describe('UserConnectionDO', () => { const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: 'root-1', status: 'busy', title: 'Root session', connectionId: 'cli-1' }, + { + id: 'root-1', + status: 'busy', + title: 'Root session', + connectionId: 'cli-1', + }, ]); }); - it('cleans up child tracking when session disappears from heartbeat', () => { + it('cleans up child tracking when session disappears from heartbeat', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -3624,11 +3849,16 @@ describe('UserConnectionDO', () => { const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: 'root-1', status: 'idle', title: 'Root session', connectionId: 'cli-1' }, + { + id: 'root-1', + status: 'idle', + title: 'Root session', + connectionId: 'cli-1', + }, ]); }); - it('forwards the per-session platform when the CLI reports it (newer CLIs)', () => { + it('forwards the per-session platform when the CLI reports it (newer CLIs)', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -3639,12 +3869,18 @@ describe('UserConnectionDO', () => { const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: 's1', status: 'busy', title: 'On a Mac', connectionId: 'cli-1', platform: 'darwin' }, + { + id: 's1', + status: 'busy', + title: 'On a Mac', + connectionId: 'cli-1', + platform: 'darwin', + }, { id: 's2', status: 'idle', title: 'Other', connectionId: 'cli-1' }, ]); }); - it('omits the platform key entirely for legacy CLIs (byte-identical response)', () => { + it('omits the platform key entirely for legacy CLIs (byte-identical response)', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -3664,7 +3900,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('getConnectedInstances', () => { - it('returns one row per CLI socket that has an `instance` attachment', () => { + it('returns one row per CLI socket that has an `instance` attachment', async () => { const { doInstance, mockCtx } = setup(); // Use the hibernated-attachment pattern (no heartbeat) — the live // scan reads the `instance` directly from the attachment, which is @@ -3674,29 +3910,40 @@ describe('UserConnectionDO', () => { projectName: 'kilo', version: '0.1.2', }); - addCliSocket(mockCtx, 'cli-B', [], { name: 'laptop-B', projectName: 'kilo' }); + addCliSocket(mockCtx, 'cli-B', [], { + name: 'laptop-B', + projectName: 'kilo', + }); addWebSocket(mockCtx); const { instances } = doInstance.getConnectedInstances(); expect(instances).toHaveLength(2); expect(instances).toEqual( expect.arrayContaining([ - { connectionId: 'cli-A', name: 'laptop-A', projectName: 'kilo', version: '0.1.2' }, + { + connectionId: 'cli-A', + name: 'laptop-A', + projectName: 'kilo', + version: '0.1.2', + }, { connectionId: 'cli-B', name: 'laptop-B', projectName: 'kilo' }, ]) ); }); - it('omits the `version` key when the CLI did not report one', () => { + it('omits the `version` key when the CLI did not report one', async () => { const { doInstance, mockCtx } = setup(); - addCliSocket(mockCtx, 'cli-1', [], { name: 'laptop-1', projectName: 'kilo' }); + addCliSocket(mockCtx, 'cli-1', [], { + name: 'laptop-1', + projectName: 'kilo', + }); const { instances } = doInstance.getConnectedInstances(); expect(instances).toEqual([{ connectionId: 'cli-1', name: 'laptop-1', projectName: 'kilo' }]); expect(instances[0]).not.toHaveProperty('version'); }); - it('excludes legacy CLIs that never reported an `instance`', () => { + it('excludes legacy CLIs that never reported an `instance`', async () => { const { doInstance, mockCtx } = setup(); // Legacy CLI: pre-spawner heartbeat has no `instance`. const cliWs = addCliSocket(mockCtx, 'legacy-1'); @@ -3706,7 +3953,7 @@ describe('UserConnectionDO', () => { expect(instances).toEqual([]); }); - it('excludes web sockets', () => { + it('excludes web sockets', async () => { const { doInstance, mockCtx } = setup(); addWebSocket(mockCtx); // A web socket with an `instance`-shaped attachment must still be skipped. @@ -3721,7 +3968,7 @@ describe('UserConnectionDO', () => { expect(instances).toEqual([]); }); - it('reads `instance` directly from the live socket (no in-memory map)', () => { + it('reads `instance` directly from the live socket (no in-memory map)', async () => { const { doInstance, mockCtx } = setup(); // Simulate a hibernated attach: socket exists, attachment has `instance` // set, but no heartbeat has been processed through the in-memory state. @@ -3733,11 +3980,16 @@ describe('UserConnectionDO', () => { const { instances } = doInstance.getConnectedInstances(); expect(instances).toEqual([ - { connectionId: 'cli-h', name: 'laptop-h', projectName: 'kilo', version: '1.0.0' }, + { + connectionId: 'cli-h', + name: 'laptop-h', + projectName: 'kilo', + version: '1.0.0', + }, ]); }); - it('includes capabilities when the CLI attachment advertises them', () => { + it('includes capabilities when the CLI attachment advertises them', async () => { const { doInstance, mockCtx } = setup(); // Hibernated attachment carries capabilities — same source // getConnectedInstances already uses for instance/version. @@ -3761,7 +4013,7 @@ describe('UserConnectionDO', () => { ]); }); - it('omits capabilities when the CLI attachment has none (legacy CLI)', () => { + it('omits capabilities when the CLI attachment has none (legacy CLI)', async () => { const { doInstance, mockCtx } = setup(); addCliSocket(mockCtx, 'cli-legacy-cap', [], { name: 'laptop-legacy', @@ -3770,12 +4022,16 @@ describe('UserConnectionDO', () => { const { instances } = doInstance.getConnectedInstances(); expect(instances).toEqual([ - { connectionId: 'cli-legacy-cap', name: 'laptop-legacy', projectName: 'kilo' }, + { + connectionId: 'cli-legacy-cap', + name: 'laptop-legacy', + projectName: 'kilo', + }, ]); expect(instances[0]).not.toHaveProperty('capabilities'); }); - it('persists `instance` in the WS attachment across heartbeats', () => { + it('persists `instance` in the WS attachment across heartbeats', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [], { @@ -3783,11 +4039,17 @@ describe('UserConnectionDO', () => { instance: { name: 'laptop-1', projectName: 'kilo', version: '0.1.0' }, }); - const att = cliWs.deserializeAttachment() as { instance?: { name: string } }; - expect(att.instance).toEqual({ name: 'laptop-1', projectName: 'kilo', version: '0.1.0' }); + const att = cliWs.deserializeAttachment() as { + instance?: { name: string }; + }; + expect(att.instance).toEqual({ + name: 'laptop-1', + projectName: 'kilo', + version: '0.1.0', + }); }); - it('drops `instance` from the attachment on a subsequent heartbeat that omits it', () => { + it('drops `instance` from the attachment on a subsequent heartbeat that omits it', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); // First heartbeat: with instance. @@ -3819,7 +4081,7 @@ describe('UserConnectionDO', () => { // attachment budget. const INSTANCE_HEADROOM = 250; - it('keeps the combined CLI attachment comfortably under 2 KiB with a worst-case instance', () => { + it('keeps the combined CLI attachment comfortably under 2 KiB with a worst-case instance', async () => { const worstCaseInstance = { name: 'x'.repeat(64), projectName: 'x'.repeat(64), @@ -3860,7 +4122,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('owner-unique active sessions', () => { - it('emits owner-unique rows: ownership transfer with both CLIs live yields exactly one row under the new owner', () => { + it('emits owner-unique rows: ownership transfer with both CLIs live yields exactly one row under the new owner', async () => { const { doInstance, ctx, mockCtx } = setup(); const oldOwner = addCliSocket(mockCtx, 'cli-old'); const newOwner = addCliSocket(mockCtx, 'cli-new'); @@ -3897,7 +4159,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- describe('edge cases', () => { - it('ignores non-JSON messages', () => { + it('ignores non-JSON messages', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); @@ -3905,7 +4167,7 @@ describe('UserConnectionDO', () => { doInstance.webSocketMessage(cliWs as never, 'not-json'); }); - it('logs invalid CLI JSON metadata without raw payload content', () => { + it('logs invalid CLI JSON metadata without raw payload content', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); @@ -3921,7 +4183,7 @@ describe('UserConnectionDO', () => { expect(JSON.stringify(warn.mock.calls)).not.toContain('raw-secret-must-not-be-logged'); }); - it('ignores messages from socket with no attachment', () => { + it('ignores messages from socket with no attachment', async () => { const { doInstance, mockCtx } = setup(); const ws = createMockWs(['cli'], null); mockCtx.addSocket(ws); @@ -3931,7 +4193,7 @@ describe('UserConnectionDO', () => { // Should not throw }); - it('ignores messages that fail Zod validation', () => { + it('ignores messages that fail Zod validation', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, []); // trigger ensureState @@ -3947,7 +4209,7 @@ describe('UserConnectionDO', () => { // Should not throw }); - it('logs malformed CLI message metadata without raw payload content', () => { + it('logs malformed CLI message metadata without raw payload content', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); @@ -3986,13 +4248,13 @@ describe('UserConnectionDO', () => { expect(msgs.some((m: Record) => m.event === 'cli.disconnected')).toBe(true); }); - it('CLI response for unknown correlation ID is a no-op', () => { + it('CLI response for unknown correlation ID is a no-op', async () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, []); // Should not throw - sendCliResponse(doInstance, cliWs, { id: 'nonexistent', result: 'ok' }); + await sendCliResponse(doInstance, cliWs, { id: 'nonexistent', result: 'ok' }); }); }); @@ -4018,7 +4280,12 @@ describe('UserConnectionDO', () => { connectionId: string, kiloUserId: string ): MockWS { - const attachment = { role: 'cli' as const, connectionId, sessions: [], kiloUserId }; + const attachment = { + role: 'cli' as const, + connectionId, + sessions: [], + kiloUserId, + }; const ws = createMockWs(['cli'], attachment); mockCtx.addSocket(ws); return ws; @@ -4070,7 +4337,9 @@ describe('UserConnectionDO', () => { fireAt: originalFireAt, attempts: 2, }); - const internal = doInstance as unknown as { readyPushFireAt: Map }; + const internal = doInstance as unknown as { + readyPushFireAt: Map; + }; internal.readyPushFireAt.set('ses_reconnect', originalFireAt); // Disconnect clears sessionOwners; reconnect is a new "first sight". @@ -4239,7 +4508,9 @@ describe('UserConnectionDO', () => { const now = 1_700_000_000_000; vi.spyOn(Date, 'now').mockReturnValue(now); const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); - const internal = doInstance as unknown as { readyPushFireAt: Map }; + const internal = doInstance as unknown as { + readyPushFireAt: Map; + }; // Establish ownership for both sessions first (first-sight schedules entries). sendHeartbeat(doInstance, cliWs, [ @@ -4282,7 +4553,9 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses_retry')]); await flushAsync(); - const internal = doInstance as unknown as { readyPushFireAt: Map }; + const internal = doInstance as unknown as { + readyPushFireAt: Map; + }; internal.readyPushFireAt.set('ses_retry', now - 1); await ctx.storage.put('readyPush:ses_retry', { kiloUserId: 'usr_1', @@ -4328,7 +4601,9 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses_backoff')]); await flushAsync(); - const internal = doInstance as unknown as { readyPushFireAt: Map }; + const internal = doInstance as unknown as { + readyPushFireAt: Map; + }; internal.readyPushFireAt.set('ses_backoff', now - 1); await ctx.storage.put('readyPush:ses_backoff', { kiloUserId: 'usr_1', @@ -4366,7 +4641,9 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses_ok')]); await flushAsync(); - const internal = doInstance as unknown as { readyPushFireAt: Map }; + const internal = doInstance as unknown as { + readyPushFireAt: Map; + }; internal.readyPushFireAt.set('ses_ok', now - 1); await ctx.storage.put('readyPush:ses_ok', { kiloUserId: 'usr_1', @@ -4397,7 +4674,9 @@ describe('UserConnectionDO', () => { attempts: 0, }); // Mirror intentionally empty - const internal = doInstance as unknown as { readyPushFireAt: Map }; + const internal = doInstance as unknown as { + readyPushFireAt: Map; + }; internal.readyPushFireAt.clear(); await doInstance.alarm(); @@ -4405,7 +4684,7 @@ describe('UserConnectionDO', () => { expect(ctx.storage.store.has('readyPush:ses_kv_only')).toBe(false); }); - it('stores the kiloUserId from the connection URL on the attachment', () => { + it('stores the kiloUserId from the connection URL on the attachment', async () => { const { doInstance } = setupWithIngestDO(); const client = createMockWs(); const server = createMockWs(); @@ -4429,7 +4708,10 @@ describe('UserConnectionDO', () => { }) ); - expect(server.deserializeAttachment()).toMatchObject({ role: 'cli', kiloUserId: 'usr_1' }); + expect(server.deserializeAttachment()).toMatchObject({ + role: 'cli', + kiloUserId: 'usr_1', + }); }); }); @@ -4446,13 +4728,18 @@ describe('UserConnectionDO', () => { const result = await doInstance.notifySessionRenamed('ses_r', 'Renamed Title'); expect(result).toEqual({ delivered: true }); - expect(ctx.storage.store.get('rename:ses_r')).toMatchObject({ title: 'Renamed Title' }); + expect(ctx.storage.store.get('rename:ses_r')).toMatchObject({ + title: 'Renamed Title', + }); const systemMsgs = allSent(cliWs).filter( m => m.type === 'system' && m.event === 'session.renamed' ); expect(systemMsgs).toHaveLength(1); - expect(systemMsgs[0].data).toEqual({ sessionId: 'ses_r', title: 'Renamed Title' }); + expect(systemMsgs[0].data).toEqual({ + sessionId: 'ses_r', + title: 'Renamed Title', + }); }); it('returns delivered:false when no owner but still persists KV', async () => { @@ -4479,7 +4766,10 @@ describe('UserConnectionDO', () => { m => m.type === 'system' && m.event === 'session.renamed' ); expect(reEmits).toHaveLength(1); - expect(reEmits[0].data).toEqual({ sessionId: 'ses_catch', title: 'New Title' }); + expect(reEmits[0].data).toEqual({ + sessionId: 'ses_catch', + title: 'New Title', + }); expect(ctx.storage.store.has('rename:ses_catch')).toBe(true); // Match → delete entry @@ -4517,7 +4807,10 @@ describe('UserConnectionDO', () => { m => m.type === 'system' && m.event === 'session.renamed' ); expect(reEmits).toHaveLength(1); - expect(reEmits[0].data).toEqual({ sessionId: 'ses_wu', title: 'Catch Up Title' }); + expect(reEmits[0].data).toEqual({ + sessionId: 'ses_wu', + title: 'Catch Up Title', + }); }); it('prunes rename entries older than TTL and does not re-emit them', async () => { @@ -4564,8 +4857,2769 @@ describe('UserConnectionDO', () => { m => m.type === 'system' && m.event === 'session.renamed' ); expect(reEmits).toHaveLength(1); - expect(reEmits[0].data).toEqual({ sessionId: 'ses_fresh', title: 'Within TTL' }); + expect(reEmits[0].data).toEqual({ + sessionId: 'ses_fresh', + title: 'Within TTL', + }); expect(ctx.storage.store.has('rename:ses_fresh')).toBe(true); }); }); + + // --------------------------------------------------------------------------- + // D8: durable remote CLI correlations + // + // Force-hibernate mechanism: these tests simulate hibernation by + // pre-populating the storage fake (ctx.storage.put) with durable entries, + // then re-instantiating the DO and calling ensureState(). Cloudflare + // Durable Object hibernation is not directly triggerable from the Vitest + // harness; the test covers the reconstruction path instead. + // --------------------------------------------------------------------------- + + describe('durable pending commands', () => { + it('rehydrates a CLI reply on wake and routes it to the live web socket (D8 case 1)', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + // Pre-populate a durable pending entry. + const now = Date.now(); + const correlationId = 'mut-1'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-a', + originalId: 'original-req-1', + command: 'send_message', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-1', + expiresAt: now + 35_000, + webConnectionId: 'web-1', + state: 'pending' as const, + }); + + // Set up a live CLI socket with the matching connectionId. + const cliWs = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-a', 'busy', 'Session A')]); + + // Set up a live web socket with the matching connectionId. + const webWs = addWebSocket(mockCtx, 'web-1'); + + // Send the CLI response with the correlation id. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { ok: true }, + }); + + // The response must be routed to the web socket. + const responses = allSent(webWs).filter(m => m.type === 'response'); + expect(responses).toHaveLength(1); + expect(responses[0].id).toBe('original-req-1'); + expect(responses[0].result).toEqual({ ok: true }); + + // The durable entry must be marked 'done' with the result. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).result).toEqual({ ok: true }); + }); + + it('keeps the durable entry and skips send when the originating web socket is gone (D8 case 2)', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'mut-2'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-b', + originalId: 'original-req-2', + command: 'send_message', + targetConnectionId: 'cli-2', + expiresAt: now + 35_000, + webConnectionId: 'web-gone', + state: 'pending' as const, + }); + + // Set up a live CLI but NO web socket. + const cliWs = addCliSocket(mockCtx, 'cli-2'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-b', 'busy', 'Session B')]); + + // Send the CLI response. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { ok: true }, + }); + + // The durable entry must be marked 'done'. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + expect((entry as Record).state).toBe('done'); + }); + + it('deduplicates a pending mutationId and returns an idempotent response for a non-catalog command', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const mutationId = 'mut-dedup-pending'; + await ctx.storage.put(`pendingCommand/${mutationId}`, { + sessionId: 'ses-c', + originalId: 'first-req', + command: 'send_message', + targetConnectionId: 'cli-3', + expiresAt: now + 35_000, + webConnectionId: 'web-3', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-3'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-c', 'busy', 'Session C')]); + const webWs = addWebSocket(mockCtx, 'web-3'); + + // Clear CLI sends after heartbeat setup. + cliWs.send.mockClear(); + + // Send a second command with the same mutationId. + await sendCommand(doInstance, webWs, { + id: 'second-req', + command: 'send_message', + sessionId: 'ses-c', + connectionId: 'cli-3', + data: { ok: true }, + mutationId, + }); + await flushAsync(); + + // The web socket must receive a dedupe error, not forward to CLI. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'second-req'); + expect(responses).toHaveLength(1); + expect(responses[0].error).toBeDefined(); + expect((responses[0].error as Record).code).toBe('COMMAND_ALREADY_PENDING'); + + // No new command must be sent to the CLI. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(0); + }); + + it('returns the stored result for a done mutationId under the new request id (D8 response identity)', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const mutationId = 'mut-done-1'; + await ctx.storage.put(`pendingCommand/${mutationId}`, { + sessionId: 'ses-d', + originalId: 'first-req-done', + command: 'send_message', + targetConnectionId: 'cli-4', + expiresAt: now + 35_000, + webConnectionId: 'web-4', + state: 'done' as const, + result: { stored: true }, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-4'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-d', 'busy', 'Session D')]); + const webWs = addWebSocket(mockCtx, 'web-4'); + cliWs.send.mockClear(); + + // Send a retry with the same mutationId but new wire id. + await sendCommand(doInstance, webWs, { + id: 'retry-req-id', + command: 'send_message', + sessionId: 'ses-d', + connectionId: 'cli-4', + mutationId, + }); + await flushAsync(); + + // The response must come under the new request's id. + const responses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'retry-req-id' + ); + expect(responses).toHaveLength(1); + expect(responses[0].result).toEqual({ stored: true }); + + // The durable entry's originalId must be updated for future retries. + const entry = await ctx.storage.get(`pendingCommand/${mutationId}`); + expect((entry as Record).originalId).toBe('retry-req-id'); + + // No command must reach the CLI. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(0); + }); + + it('returns both stored result and error for a done mutationId with combined outcome', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const mutationId = 'mut-done-combined'; + await ctx.storage.put(`pendingCommand/${mutationId}`, { + sessionId: 'ses-dc', + originalId: 'first-req-combined', + command: 'send_message', + targetConnectionId: 'cli-dc', + expiresAt: now + 35_000, + webConnectionId: 'web-dc', + state: 'done' as const, + result: { partial: 'data' }, + error: 'partial error', + }); + + const cliWs = addCliSocket(mockCtx, 'cli-dc'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-dc', 'busy', 'Session DC')]); + const webWs = addWebSocket(mockCtx, 'web-dc'); + cliWs.send.mockClear(); + + // Retry with the same mutationId but new wire id. + await sendCommand(doInstance, webWs, { + id: 'retry-combined-id', + command: 'send_message', + sessionId: 'ses-dc', + connectionId: 'cli-dc', + mutationId, + }); + await flushAsync(); + + const responses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'retry-combined-id' + ); + expect(responses).toHaveLength(1); + expect(responses[0].result).toEqual({ partial: 'data' }); + expect(responses[0].error).toBe('partial error'); + + // No command must reach the CLI. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(0); + }); + + it('behaves identically without a mutationId (per-send random correlation id)', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const cliWs = addCliSocket(mockCtx, 'cli-5'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-e', 'busy', 'Session E')]); + const webWs = addWebSocket(mockCtx, 'web-5'); + + cliWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'no-mut-req', + command: 'send_message', + sessionId: 'ses-e', + // No mutationId. + }); + await flushAsync(); + + // The command must be forwarded to the CLI with a random correlationId. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(1); + expect(cliCommands[0].id).toBeTruthy(); + expect(cliCommands[0].mutationId).toBeUndefined(); + + // A durable entry must exist (persisted at creation time). + const correlationId = cliCommands[0].id as string; + await flushAsync(); + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + + // Send CLI response to verify normal resolution. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { normal: true }, + }); + + const responses = allSent(webWs).filter(m => m.type === 'response'); + expect(responses.some(r => r.id === 'no-mut-req' && r.result)).toBe(true); + }); + + it('does not delete the durable entry on web disconnect (step 26b)', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const cliWs = addCliSocket(mockCtx, 'cli-6'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-f', 'busy', 'Session F')]); + const webWs = addWebSocket(mockCtx, 'web-6'); + + cliWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'disconnect-req', + command: 'send_message', + sessionId: 'ses-f', + connectionId: 'cli-6', + }); + await flushAsync(); + + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(1); + const correlationId = cliCommands[0].id as string; + + // Verify durable entry exists. + await flushAsync(); + let entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + + // Disconnect the web socket. + await doInstance.webSocketClose(webWs as never, 1000, '', true); + await flushAsync(); + + // The durable entry must still exist. + entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + }); + + it('expires a durable pending entry and marks it done with COMMAND_EXPIRED_ERROR', async () => { + const { mockCtx, ctx } = setup(); + + const pastTime = Date.now() - 10_000; + const correlationId = 'mut-expired'; + + // Set expiredAt in the past. + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-g', + originalId: 'expired-req', + command: 'send_message', + targetConnectionId: 'cli-7', + expiresAt: pastTime, + webConnectionId: 'web-7', + state: 'pending' as const, + }); + + // Create a fresh DO to simulate a wake with the durable entry. + const doInstance2 = new UserConnectionDO(ctx as never, {} as never); + + // Set up sockets so the DO has something to work with. + const cliWs = addCliSocket(mockCtx, 'cli-7'); + sendHeartbeat(doInstance2, cliWs, [makeSession('ses-g', 'busy', 'Session G')]); + + // Trigger alarm which calls expirePendingCommands. + await doInstance2.alarm(); + await flushAsync(); + + // The entry must be marked done with COMMAND_EXPIRED_ERROR. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'COMMAND_EXPIRED', + message: 'Command expired', + }); + }); + + it('counts durable entries toward the pending command cap (mutationId path)', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const cliWs = addCliSocket(mockCtx, 'cli-8'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-h', 'busy', 'Session H')]); + const webWs = addWebSocket(mockCtx, 'web-8'); + + // Pre-fill storage with many durable entries to approach the cap. + const now = Date.now(); + for (let i = 0; i < 127; i++) { + await ctx.storage.put(`pendingCommand/prefill-${i}`, { + sessionId: 'ses-h', + originalId: `prefill-${i}`, + command: 'send_message', + targetConnectionId: 'cli-8', + expiresAt: now + 35_000, + webConnectionId: 'web-8', + state: 'pending' as const, + }); + } + + // The first command (no mutationId) is just within the cap: 0 in-memory + 1 new = 1. + await sendCommand(doInstance, webWs, { + id: 'at-cap', + command: 'send_message', + sessionId: 'ses-h', + connectionId: 'cli-8', + }); + await flushAsync(); + const commandsAtCap = allSent(cliWs).filter(m => m.type === 'command'); + expect(commandsAtCap.length).toBeGreaterThanOrEqual(1); + + // The second command uses a fresh mutationId. The mutationId path + // counts durable entries: 127 prefill + 1 from the first + // command's durable write + 1 in-memory = 129 ≥ 128. Rejected. + cliWs.send.mockClear(); + webWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'over-cap', + command: 'send_message', + sessionId: 'ses-h', + connectionId: 'cli-8', + mutationId: 'fresh-cap-check', + }); + await flushAsync(); + + const rejection = allSent(webWs).filter(m => m.type === 'response' && m.id === 'over-cap'); + expect(rejection).toHaveLength(1); + expect((rejection[0].error as Record)?.code).toBe('PENDING_COMMAND_LIMIT'); + + // No second command to CLI. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(0); + }); + + // ------------------------------------------------------------------------- + // Fix 2: durable entry persists the exact terminal error sent live + // ------------------------------------------------------------------------- + it('persists the exact CLI string error in the durable entry, not CLI_COMMAND_ERROR', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'err-1', + command: 'list_models', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + error: 'unknown command: list_models', + }); + + // Live response carries the original string. + const live = parseSent(webWs) as { error: unknown }; + expect(live.error).toBe('unknown command: list_models'); + + // Durable entry carries the same string. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toBe('unknown command: list_models'); + }); + + it('persists the exact structured CLI_UPGRADE_REQUIRED error in the durable entry', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'err-2', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + error: 'unknown command: list_commands', + }); + + // Live response carries the structured error. + const live = parseSent(webWs) as { error: unknown }; + expect(live.error).toEqual({ + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', + message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', + }); + + // Durable entry carries the same structured error. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', + message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', + }); + }); + + // ------------------------------------------------------------------------- + // Fix 3: reply durability — durable write is extended via ctx.waitUntil + // ------------------------------------------------------------------------- + it('completes the durable write for a CLI response via ctx.waitUntil', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'dur-1', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + + // Send CLI response — the fix uses ctx.waitUntil so the durable + // write is extended past the handler's return. flushAsync settles + // the waitUntil promise in the test harness. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { durable: true }, + }); + + // The durable entry must now be written. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).result).toEqual({ + durable: true, + }); + }); + + // ------------------------------------------------------------------------- + // Fix 5: rehydrated oversized catalog sends to live web socket + // ------------------------------------------------------------------------- + it('sends the oversized-catalog error to a live web socket after rehydration', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'cat-rehydrated'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-cat', + originalId: 'original-cat', + command: 'list_models', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-cat', + expiresAt: now + 35_000, + webConnectionId: 'web-cat', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-cat'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-cat', 'busy', 'Session Cat')]); + const webWs = addWebSocket(mockCtx, 'web-cat'); + webWs.send.mockClear(); + + const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: oversized, + }); + + // The live web socket must receive the oversized error. + const responses = allSent(webWs).filter(m => m.type === 'response'); + expect(responses).toHaveLength(1); + expect(responses[0].error).toEqual({ + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', + }); + + // The durable entry must be marked done. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', + }); + }); + + // ------------------------------------------------------------------------- + // Fix 6: bare 'CLI disconnected' string for live sends + // ------------------------------------------------------------------------- + it('sends a bare string "CLI disconnected" to the live web socket on disconnect', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'cmd-disco', + command: 'send_message', + sessionId: 's1', + }); + webWs.send.mockClear(); + + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + + const msgs = allSent(webWs); + const errorResp = msgs.find(m => m.type === 'response' && m.id === 'cmd-disco'); + expect(errorResp).toBeDefined(); + // Live wire must carry the bare string, not a structured object. + expect(errorResp!.error).toBe('CLI disconnected'); + }); + + it('stores the live CLI-disconnected error in the durable entry for retries', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'cmd-dur-disco', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + await flushAsync(); + + // Durable entry must carry structured error for typed retries. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toBe('CLI disconnected'); + }); + + // ------------------------------------------------------------------------- + // Fix: D8 case 2 shaped terminal outcome + // ------------------------------------------------------------------------- + it('shapes the error for a no-web D8 case 2 retry (CLI_UPGRADE_REQUIRED mapping)', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'mut-no-web-shaped'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-nw', + originalId: 'original-nw', + command: 'list_commands', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-nw', + expiresAt: now + 35_000, + webConnectionId: 'web-gone-shaped', + state: 'pending' as const, + }); + + // Set up a live CLI but NO web socket. + const cliWs = addCliSocket(mockCtx, 'cli-nw'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-nw', 'busy', 'Session NW')]); + + // Send the CLI response with an upgrade-required error. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + error: 'unknown command: list_commands', + }); + + // The durable entry must be marked 'done' with the shaped structured error. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + expect((entry as Record).state).toBe('done'); + // Must be the shaped CLI_UPGRADE_REQUIRED error, not the raw string. + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', + message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', + }); + }); + + it('shapes a non-allowlist CLI string error correctly for a no-web D8 case 2 retry', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'mut-no-web-string'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-nws', + originalId: 'original-nws', + command: 'list_models', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-nws', + expiresAt: now + 35_000, + webConnectionId: 'web-gone-string', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-nws'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-nws', 'busy', 'Session NWS')]); + + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + error: 'unknown command: list_models', + }); + + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + // list_models is not in CLI_UPGRADE_REQUIRED_COMMANDS, so the raw + // string is preserved verbatim. + expect((entry as Record).error).toBe('unknown command: list_models'); + }); + + it('does not persist raw oversized catalog data for a no-web D8 case 2 retry', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'mut-no-web-oversized'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-nwo', + originalId: 'original-nwo', + command: 'list_models', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-nwo', + expiresAt: now + 35_000, + webConnectionId: 'web-gone-oversized', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-nwo'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-nwo', 'busy', 'Session NWO')]); + + const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: oversized, + }); + + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + // Must store the CATALOG_TOO_LARGE error, not the raw oversized result. + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', + }); + // Must NOT store the raw result. + expect((entry as Record).result).toBeUndefined(); + }); + + it('shapes a relay-object CLI error to CLI_COMMAND_ERROR for a no-web D8 case 2 retry', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'mut-no-web-relay'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-nwr', + originalId: 'original-nwr', + command: 'send_message', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-nwr', + expiresAt: now + 35_000, + webConnectionId: 'web-gone-relay', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-nwr'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-nwr', 'busy', 'Session NWR')]); + + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + error: { + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }, + }); + + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + // Relay-shaped objects from the CLI are sanitized to CLI_COMMAND_ERROR. + expect((entry as Record).error).toEqual({ + source: 'cli', + message: 'Command failed', + }); + }); + + it('persists both result and error for a no-web D8 case 2 combined response', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'mut-no-web-combined'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-nwc', + originalId: 'original-nwc', + command: 'send_message', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-nwc', + expiresAt: now + 35_000, + webConnectionId: 'web-gone-combined', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-nwc'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-nwc', 'busy', 'Session NWC')]); + + // CLI sends both result and error. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { partial: 'data' }, + error: 'something went wrong', + }); + + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).result).toEqual({ partial: 'data' }); + expect((entry as Record).error).toBe('something went wrong'); + }); + + it('persists the terminal catalog-too-large outcome without a second durable read', async () => { + // Regression test: the rehydrated catalog-too-large branch must not + // do a second getDurablePendingCommand. It must use the value + // captured at the top of handleCliResponse instead. + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'cat-no-second-read'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-cat2', + originalId: 'original-cat2', + command: 'list_models', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-cat2', + expiresAt: now + 35_000, + webConnectionId: 'web-cat2', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-cat2'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-cat2', 'busy', 'Session Cat2')]); + const webWs = addWebSocket(mockCtx, 'web-cat2'); + webWs.send.mockClear(); + + const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: oversized, + }); + + // The live web socket must receive the oversized error. + const responses = allSent(webWs).filter(m => m.type === 'response'); + expect(responses).toHaveLength(1); + expect(responses[0].error).toEqual({ + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', + }); + + // The durable entry must be written with the terminal outcome. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect(entry).toBeDefined(); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', + }); + + // Must not store the raw oversized result. + expect((entry as Record).result).toBeUndefined(); + }); + + // ------------------------------------------------------------------------- + // Fix: Durable result bounding (non-catalog commands) + // ------------------------------------------------------------------------- + it('bounds non-catalog results over the durable limit before the durable write', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'big-result', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Create a result just over the durable limit. + const oversized = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES + 1); + + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); + + // Live response must carry the full result unchanged. + const live = parseSent(webWs) as { id: string; result: unknown }; + expect(live.result).toEqual(oversized); + + // Durable entry must carry an error, not the truncated marker as a result. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).result).toBeUndefined(); + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'DURABLE_RESULT_TOO_LARGE', + message: 'Result is too large to store for retries', + }); + }); + + it('stores non-catalog results at exactly the durable limit unchanged', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'exact-result', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + const exact = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES); + + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: exact }); + + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).result).toEqual(exact); + }); + + it('bounds the result in the D8 case 2 (no-web) durable write', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'mut-no-web-big'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-big', + originalId: 'original-big', + command: 'send_message', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-big', + expiresAt: now + 35_000, + webConnectionId: 'web-gone-big', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-big'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-big', 'busy', 'Session Big')]); + + const oversized = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES + 1); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); + + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + // Durable entry must carry an error, not the truncated marker as a result. + expect((entry as Record).result).toBeUndefined(); + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'DURABLE_RESULT_TOO_LARGE', + message: 'Result is too large to store for retries', + }); + }); + + // ------------------------------------------------------------------------- + // Fix: Duplicate CLI reply is idempotent + // ------------------------------------------------------------------------- + it('delivers a CLI response exactly once when the same response arrives twice', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'dedup-cmd', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // First response — delivered live. + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true } }); + + const firstResponses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'dedup-cmd' + ); + expect(firstResponses).toHaveLength(1); + expect(firstResponses[0].result).toEqual({ ok: true }); + + // Second response with same correlationId — must NOT deliver again. + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true } }); + + const allResponses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'dedup-cmd' + ); + expect(allResponses).toHaveLength(1); + }); + + // ------------------------------------------------------------------------- + // Fix: Late response cannot override expiry + // ------------------------------------------------------------------------- + it('ignores a live CLI response that has already expired', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + const baseTime = 1_000_000; + vi.spyOn(Date, 'now').mockReturnValue(baseTime); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'expiry-cmd', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Advance time past the 35-second TTL. + vi.spyOn(Date, 'now').mockReturnValue(baseTime + 35_001); + + // CLI sends response — should be dropped because it's expired. + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { late: true } }); + + // No response for this command ID. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'expiry-cmd'); + expect(responses).toHaveLength(0); + }); + + // ------------------------------------------------------------------------- + // Fix: Oversized mutationId rejected + // ------------------------------------------------------------------------- + it('rejects a mutationId longer than 128 characters via schema validation', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + + const longMutationId = 'x'.repeat(129); + await sendCommand(doInstance, webWs, { + id: 'oversized-mut', + command: 'send_message', + sessionId: 's1', + mutationId: longMutationId, + }); + + // The Zod schema rejects the oversized mutationId. No response is sent + // (the invalid message is dropped with a warn log), and no command + // reaches the CLI. + const responses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'oversized-mut' + ); + expect(responses).toHaveLength(0); + + // No command must reach the CLI. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(0); + }); + + it('accepts a mutationId at exactly 128 characters', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + + const exactMutationId = 'x'.repeat(128); + await sendCommand(doInstance, webWs, { + id: 'exact-mut', + command: 'send_message', + sessionId: 's1', + mutationId: exactMutationId, + }); + await flushAsync(); + + // Command must be forwarded to CLI. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(1); + expect(cliCommands[0].mutationId).toBe(exactMutationId); + }); + + it('rejects a mutationId longer than 128 characters in the web message schema', async () => { + const { doInstance, mockCtx } = setup(); + addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + const longMutationId = 'x'.repeat(129); + const msg = JSON.stringify({ + type: 'command', + id: 'schema-mut', + command: 'send_message', + mutationId: longMutationId, + }); + + // Send raw — schema validation must reject it. + doInstance.webSocketMessage(webWs as never, msg); + + // No response forwarded to CLI. + const cliCommands = allSent(mockCtx.sockets.find(s => s._tags.includes('cli'))!).filter( + m => m.type === 'command' + ); + expect(cliCommands).toHaveLength(0); + }); + + // ------------------------------------------------------------------------- + // Fix: Rehydrated reply after expiry cannot deliver or persist a success outcome + // ------------------------------------------------------------------------- + it('rejects a rehydrated CLI reply when the durable entry has expired', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const pastTime = Date.now() - 10_000; + const correlationId = 'rehydrated-expired'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-re', + originalId: 'original-re', + command: 'send_message', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-re', + expiresAt: pastTime, + webConnectionId: 'web-re', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-re'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-re', 'busy', 'Session RE')]); + const webWs = addWebSocket(mockCtx, 'web-re'); + webWs.send.mockClear(); + + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { late: true }, + }); + + // No response must reach the web socket. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'original-re'); + expect(responses).toHaveLength(0); + + // The durable entry must be marked done with COMMAND_EXPIRED_ERROR. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toEqual({ + source: 'relay', + code: 'COMMAND_EXPIRED', + message: 'Command expired', + }); + // Must not store the late result. + expect((entry as Record).result).toBeUndefined(); + }); + + // ------------------------------------------------------------------------- + // Fix: Total durable entry with result at limit stays under 128 KiB + // ------------------------------------------------------------------------- + it('keeps the total serialized durable entry under 128 KiB when the result is at the bound', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'safety-check', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + const result = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result }); + + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).result).toEqual(result); + + // The total serialized entry must stay under 128 KiB (131,072 bytes). + const serialized = new TextEncoder().encode(JSON.stringify(entry)).byteLength; + expect(serialized).toBeLessThan(131_072); + }); + // ------------------------------------------------------------------------- + // Fix: duplicate rehydration yields at most one terminal delivery + // ------------------------------------------------------------------------- + it('delivers at most one terminal result when two rehydrated replies race', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'rehydrated-race'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-race', + originalId: 'original-race', + command: 'send_message', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-race', + expiresAt: now + 35_000, + webConnectionId: 'web-race', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-race'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-race', 'busy', 'Session Race')]); + const webWs = addWebSocket(mockCtx, 'web-race'); + webWs.send.mockClear(); + + // Deferred promise: holds the first reply's terminal storage.put so + // the second reply enters the handler while the first is blocked on + // the durable write. This ractests the marker+write window, not the + // read. + let release: () => void; + const deferred = new Promise(resolve => { + release = resolve; + }); + let deferredFired = false; + ctx.storage.put.mockImplementation(async (key, value) => { + // Defer only the first terminal 'done' write for the rehydrated + // correlation id. + if ( + !deferredFired && + key === `pendingCommand/${correlationId}` && + (value as Record)?.state === 'done' + ) { + deferredFired = true; + await deferred; + } + ctx.storage.store.set(key as string, value); + }); + + // Start the first CLI reply — the handler rehydrates the durable + // entry, builds the in-memory entry, then blocks on the deferred + // storage.put before sending the live response. + doInstance.webSocketMessage( + cliWs as never, + JSON.stringify({ type: 'response', id: correlationId, result: { first: true } }) + ); + await flushAsync(); + + // Start the second CLI reply while the first is held on storage.put. + // The completedCorrelationIds reservation is still set, so the + // second reply returns early without delivering. + doInstance.webSocketMessage( + cliWs as never, + JSON.stringify({ type: 'response', id: correlationId, result: { second: true } }) + ); + await flushAsync(); + + // Release the deferred put so the first reply completes. + release!(); + await flushAsync(); + + // Exactly one terminal response must reach the web socket. + const responses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'original-race' + ); + expect(responses).toHaveLength(1); + }); + + // ------------------------------------------------------------------------- + // Fix: no live send and clean marker when the durable write fails + // ------------------------------------------------------------------------- + it('sends no live response and clears the marker when the terminal durable write fails', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'cmd-fail-write', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Make the terminal durable write fail. The mock harness fires + // storage.put multiple times: (1) dispatchWebCommandSync creates a + // pending entry in waitUntil, (2) the done write inside + // handleCliResponse. Clear history and make the next call reject to + // target the terminal write. + ctx.storage.put.mockClear(); + ctx.storage.put.mockRejectedValueOnce(new Error('durable write failed')); + + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { ok: true }, + }); + + // No live response must have been sent — durable write failed first. + const responses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'cmd-fail-write' + ); + expect(responses).toHaveLength(0); + + // The completedCorrelationIds marker must be cleared on failure. + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId)).toBe(false); + }); + + // ------------------------------------------------------------------------- + // Fix: no live send and clean marker when catalog durable write fails + // ------------------------------------------------------------------------- + it('sends no live response and clears the marker when the catalog durable write fails', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'cmd-catalog', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Make the storage.put inside the catalog-too-large path reject. + ctx.storage.put.mockClear(); + ctx.storage.put.mockRejectedValueOnce(new Error('catalog write failed')); + + const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: oversized, + }); + + // No live response — the durable write must succeed first. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-catalog'); + expect(responses).toHaveLength(0); + + // The completedCorrelationIds marker must be cleared on failure. + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId)).toBe(false); + }); + + // ------------------------------------------------------------------------- + // Fix: re-send after failed durable write succeeds (no duplicate) + // ------------------------------------------------------------------------- + it('re-sends after a failed durable write (the CLI re-send succeeds because nothing was sent first time)', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'retry-after-fail', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Make the terminal durable write fail. + ctx.storage.put.mockClear(); + ctx.storage.put.mockRejectedValueOnce(new Error('durable write failed')); + + // First CLI response: durable write fails, no live send, marker cleared. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { first: true }, + }); + + // Zero responses — durable write failed first, so nothing was sent. + const responses1 = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'retry-after-fail' + ); + expect(responses1).toHaveLength(0); + + // Marker must be clear (reservation was cleaned on failure). + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId)).toBe(false); + + // CLI re-sends the response (simulates a real retry or DO wake). + // The durable write succeeds this time because the mock no longer rejects. + webWs.send.mockClear(); + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { second: true }, + }); + + // The second send must deliver (the rehydration path processes it from + // the durable pending state). + const responses2 = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'retry-after-fail' + ); + expect(responses2).toHaveLength(1); + expect(responses2[0].result).toEqual({ second: true }); + }); + + // ------------------------------------------------------------------------- + // Fix: durable-before-send — marker cleared after successful write + // ------------------------------------------------------------------------- + it('clears the completedCorrelationIds marker after a successful durable write and live send', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'cmd-clear', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { ok: true }, + }); + + // Live response must have been delivered. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-clear'); + expect(responses).toHaveLength(1); + expect(responses[0].result).toEqual({ ok: true }); + + // Marker must be cleared — the durable 'done' state is the dedupe guard now. + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId)).toBe(false); + + // Verify durable entry is 'done' (confirms the marker is replaced). + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + }); + + // ------------------------------------------------------------------------- + // Fix: durable-before-send — durable persisted before live send + // ------------------------------------------------------------------------- + it('persists the durable entry before sending the live response', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'order-cmd', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Track call order of storage.put vs ws.send. + const callOrder: string[] = []; + ctx.storage.put.mockImplementation(async (key: string, value: unknown): Promise => { + callOrder.push('put'); + ctx.storage.store.set(key, value); + }); + webWs.send.mockImplementation((..._args: unknown[]) => { + callOrder.push('send'); + }); + + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { ok: true }, + }); + + // The durable write must happen before the web socket send. + const putIdx = callOrder.indexOf('put'); + const sendIdx = callOrder.indexOf('send'); + expect(putIdx).toBeGreaterThanOrEqual(0); + expect(sendIdx).toBeGreaterThanOrEqual(0); + expect(putIdx).toBeLessThan(sendIdx); + }); + + it('retries an in-memory disconnect terminal write before one live response', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'disconnect-write-fail', + command: 'send_message', + sessionId: 's1', + }); + webWs.send.mockClear(); + ctx.storage.put.mockRejectedValueOnce(new Error('durable write failed')); + + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + await flushAsync(); + + const responses = allSent(webWs).filter( + message => message.type === 'response' && message.id === 'disconnect-write-fail' + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toBe('CLI disconnected'); + expect(markers.size).toBe(0); + const entry = await ctx.storage.get(`pendingCommand/${getCorrelationId(cliWs)}`); + expect(entry).toMatchObject({ state: 'done', error: 'CLI disconnected' }); + }); + + it('retries an in-memory owner-change terminal write before one live response', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const firstOwner = addCliSocket(mockCtx, 'cli-1'); + const nextOwner = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, firstOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, nextOwner, []); + await sendCommand(doInstance, webWs, { + id: 'owner-write-fail', + command: 'send_message', + sessionId: 's1', + }); + webWs.send.mockClear(); + ctx.storage.put.mockRejectedValueOnce(new Error('durable write failed')); + + sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + await flushAsync(); + + const responses = allSent(webWs).filter( + message => message.type === 'response' && message.id === 'owner-write-fail' + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toEqual({ + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }); + expect(markers.size).toBe(0); + const entry = await ctx.storage.get(`pendingCommand/${getCorrelationId(firstOwner)}`); + expect(entry).toMatchObject({ + state: 'done', + error: { + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }, + }); + }); + + it('fences a disconnect terminal outcome while its durable write is pending', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'disconnect-fence', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + let resolvePut: (() => void) | undefined; + ctx.storage.put.mockImplementationOnce( + (key, value) => + new Promise(resolve => { + resolvePut = () => { + ctx.storage.store.set(key, value); + resolve(); + }; + }) + ); + + mockCtx.removeSocket(cliWs); + const closing = disconnectCli(doInstance, cliWs); + await Promise.resolve(); + + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId)).toBe(true); + + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { wrong: true } }); + expect( + allSent(webWs).filter( + message => message.type === 'response' && message.id === 'disconnect-fence' + ) + ).toHaveLength(0); + + resolvePut?.(); + await closing; + + const responses = allSent(webWs).filter( + message => message.type === 'response' && message.id === 'disconnect-fence' + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toBe('CLI disconnected'); + expect(markers.has(correlationId)).toBe(false); + }); + + it('fences an owner-change terminal outcome while its durable write is pending', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const firstOwner = addCliSocket(mockCtx, 'cli-1'); + const nextOwner = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, firstOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, nextOwner, []); + await sendCommand(doInstance, webWs, { + id: 'owner-fence', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(firstOwner); + webWs.send.mockClear(); + + let resolvePut: (() => void) | undefined; + ctx.storage.put.mockImplementationOnce( + (key, value) => + new Promise(resolve => { + resolvePut = () => { + ctx.storage.store.set(key, value); + resolve(); + }; + }) + ); + + sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + await Promise.resolve(); + + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId)).toBe(true); + + await sendCliResponse(doInstance, firstOwner, { id: correlationId, result: { wrong: true } }); + expect( + allSent(webWs).filter( + message => message.type === 'response' && message.id === 'owner-fence' + ) + ).toHaveLength(0); + + resolvePut?.(); + await flushAsync(); + + const responses = allSent(webWs).filter( + message => message.type === 'response' && message.id === 'owner-fence' + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toEqual({ + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }); + expect(markers.has(correlationId)).toBe(false); + }); + + it('continues disconnect cleanup after one terminal write retries', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'first-disconnect', + command: 'send_message', + sessionId: 's1', + }); + await sendCommand(doInstance, webWs, { + id: 'second-disconnect', + command: 'send_message', + sessionId: 's1', + }); + webWs.send.mockClear(); + ctx.storage.put.mockRejectedValueOnce(new Error('durable write failed')); + + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + await flushAsync(); + + for (const id of ['first-disconnect', 'second-disconnect']) { + const responses = allSent(webWs).filter( + message => message.type === 'response' && message.id === id + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toBe('CLI disconnected'); + } + expect( + allSent(webWs).some( + message => message.type === 'system' && message.event === 'cli.disconnected' + ) + ).toBe(true); + }); + + it('fences every matching disconnect command before the first terminal write settles', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'fence-first', + command: 'send_message', + sessionId: 's1', + }); + await sendCommand(doInstance, webWs, { + id: 'fence-second', + command: 'send_message', + sessionId: 's1', + }); + const [firstCorrelationId, secondCorrelationId] = [ + getCorrelationId(cliWs, 0), + getCorrelationId(cliWs, 1), + ]; + webWs.send.mockClear(); + + let resolvePut: (() => void) | undefined; + ctx.storage.put.mockImplementationOnce( + (key, value) => + new Promise(resolve => { + resolvePut = () => { + ctx.storage.store.set(key, value); + resolve(); + }; + }) + ); + + mockCtx.removeSocket(cliWs); + const closing = disconnectCli(doInstance, cliWs); + await Promise.resolve(); + + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(firstCorrelationId)).toBe(true); + expect(markers.has(secondCorrelationId)).toBe(true); + + await sendCliResponse(doInstance, cliWs, { + id: secondCorrelationId, + result: { wrong: true }, + }); + expect( + allSent(webWs).filter( + message => message.type === 'response' && message.id === 'fence-second' + ) + ).toHaveLength(0); + + resolvePut?.(); + await closing; + await flushAsync(); + + for (const id of ['fence-first', 'fence-second']) { + const responses = allSent(webWs).filter( + message => message.type === 'response' && message.id === id + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toBe('CLI disconnected'); + } + for (const correlationId of [firstCorrelationId, secondCorrelationId]) { + await expect(ctx.storage.get(`pendingCommand/${correlationId}`)).resolves.toMatchObject({ + state: 'done', + error: 'CLI disconnected', + }); + } + }); + + it('keeps a disconnect terminal outcome after the initial pending write settles', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + let resolveInitialPut: (() => void) | undefined; + ctx.storage.put.mockImplementationOnce( + (key, value) => + new Promise(resolve => { + resolveInitialPut = () => { + ctx.storage.store.set(key, value); + resolve(); + }; + }) + ); + + doInstance.webSocketMessage( + webWs as never, + JSON.stringify({ + type: 'command', + id: 'initial-write-disconnect', + command: 'send_message', + sessionId: 's1', + }) + ); + const correlationId = [ + ...(Reflect.get(doInstance, 'pendingInitialCommandWrites') as Set), + ][0]; + expect(correlationId).toBeDefined(); + webWs.send.mockClear(); + + mockCtx.removeSocket(cliWs); + const closing = disconnectCli(doInstance, cliWs); + await Promise.resolve(); + resolveInitialPut?.(); + await Promise.resolve(); + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId!)).toBe(true); + doInstance.webSocketMessage( + cliWs as never, + JSON.stringify({ type: 'response', id: correlationId, result: { wrong: true } }) + ); + await closing; + await flushAsync(); + + await expect(ctx.storage.get(`pendingCommand/${correlationId!}`)).resolves.toMatchObject({ + state: 'done', + error: 'CLI disconnected', + }); + expect( + allSent(webWs).filter( + message => message.type === 'response' && message.id === 'initial-write-disconnect' + ) + ).toEqual([{ type: 'response', id: 'initial-write-disconnect', error: 'CLI disconnected' }]); + + await sendCliResponse(doInstance, cliWs, { id: correlationId!, result: { wrong: true } }); + expect( + allSent(webWs).filter( + message => message.type === 'response' && message.id === 'initial-write-disconnect' + ) + ).toHaveLength(1); + }); + + it('keeps an owner-change terminal outcome after the initial pending write settles', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const firstOwner = addCliSocket(mockCtx, 'cli-1'); + const nextOwner = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, firstOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, nextOwner, []); + let resolveInitialPut: (() => void) | undefined; + ctx.storage.put.mockImplementationOnce( + (key, value) => + new Promise(resolve => { + resolveInitialPut = () => { + ctx.storage.store.set(key, value); + resolve(); + }; + }) + ); + + doInstance.webSocketMessage( + webWs as never, + JSON.stringify({ + type: 'command', + id: 'initial-write-owner-change', + command: 'send_message', + sessionId: 's1', + }) + ); + const correlationId = [ + ...(Reflect.get(doInstance, 'pendingInitialCommandWrites') as Set), + ][0]; + expect(correlationId).toBeDefined(); + webWs.send.mockClear(); + + sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + await Promise.resolve(); + resolveInitialPut?.(); + await Promise.resolve(); + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId!)).toBe(true); + doInstance.webSocketMessage( + firstOwner as never, + JSON.stringify({ type: 'response', id: correlationId, result: { wrong: true } }) + ); + await flushAsync(); + + const error = { + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }; + await expect(ctx.storage.get(`pendingCommand/${correlationId!}`)).resolves.toMatchObject({ + state: 'done', + error, + }); + expect( + allSent(webWs).filter( + message => message.type === 'response' && message.id === 'initial-write-owner-change' + ) + ).toEqual([{ type: 'response', id: 'initial-write-owner-change', error }]); + + await sendCliResponse(doInstance, firstOwner, { + id: correlationId!, + result: { wrong: true }, + }); + expect( + allSent(webWs).filter( + message => message.type === 'response' && message.id === 'initial-write-owner-change' + ) + ).toHaveLength(1); + }); + + // ------------------------------------------------------------------------- + // Fix: handleCliResponse failure caught at waitUntil boundary + // ------------------------------------------------------------------------- + it('catches a handleCliResponse storage-read failure at the waitUntil boundary', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + // Pre-populate a durable pending entry with a valid CLI socket. + const now = Date.now(); + const correlationId = 'rehydrated-read-fail'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-rrf', + originalId: 'original-rrf', + command: 'send_message', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-rrf', + expiresAt: now + 35_000, + webConnectionId: 'web-rrf', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-rrf'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-rrf', 'busy', 'Session RRF')]); + addWebSocket(mockCtx, 'web-rrf'); + + // Make the first durable read reject (simulates storage failure). + ctx.storage.get.mockRejectedValueOnce(new Error('storage read failure')); + + // Track console.error calls to verify the catch handler logged. + const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => undefined); + + // The CLI response enters the rehydration path, getDurablePendingCommand + // reads durable state and rejects. The waitUntil catch logs the error. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { ok: true }, + }); + + // The error must have been caught and logged, not a rejected waitUntil. + expect(errorSpy).toHaveBeenCalledWith( + 'Failed to handle CLI response (non-fatal)', + expect.objectContaining({ + correlationId, + error: 'storage read failure', + }) + ); + + // The completedCorrelationIds marker must be clean (reservation was + // set before the read and cleared on failure — the .catch at + // waitUntil is the outer guard, but the internal marker reservation + // must still clean up). + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId)).toBe(false); + }); + + it('clears the rehydration reservation on a post-acquisition throw so a later retry is not blocked', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'mut-throw-retry'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-throw', + originalId: 'original-throw', + command: 'send_message', + expectedOwnerConnectionId: undefined, + targetConnectionId: 'cli-throw', + expiresAt: now + 35_000, + webConnectionId: 'web-throw', + state: 'pending' as const, + }); + + const cliWs = addCliSocket(mockCtx, 'cli-throw'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-throw', 'busy', 'Session Throw')]); + const webWs = addWebSocket(mockCtx, 'web-throw'); + + // Force boundDurableResult to throw on the first call so the outer + // finally path exercises the reservation cleanup. + const doAny = doInstance as unknown as { boundDurableResult: (r: unknown) => unknown }; + const originalBound = doAny.boundDurableResult; + doAny.boundDurableResult = vi.fn().mockImplementationOnce(() => { + throw new Error('forced boundDurableResult throw'); + }); + + // First CLI response: enters the rehydration path, builds the in-memory + // entry, then throws inside boundDurableResult. The outer finally must + // clear the completedCorrelationIds reservation. + webWs.send.mockClear(); + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { ok: true }, + }); + + // The reservation marker must be clean after the throw. + const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; + expect(markers.has(correlationId)).toBe(false); + + // Restore boundDurableResult so a retry can succeed. + doAny.boundDurableResult = originalBound; + + // Second CLI response with the same correlationId: the rehydration path + // must not be blocked by a stale marker. The entry is still 'pending' + // so it must be processed normally. + await sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { second: true }, + }); + + const responses = allSent(webWs).filter(m => m.type === 'response'); + expect(responses).toHaveLength(1); + expect(responses[0].id).toBe('original-throw'); + expect(responses[0].result).toEqual({ second: true }); + + // The durable entry must be marked 'done'. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).result).toEqual({ second: true }); + }); + + // ------------------------------------------------------------------------- + // Fix: Terminal delivery to reattached web sockets + // ------------------------------------------------------------------------- + it('delivers COMMAND_EXPIRED_ERROR to a live web socket when a durable entry expires', async () => { + const { mockCtx, ctx } = setup(); + + const pastTime = Date.now() - 10_000; + const correlationId = 'expired-delivery'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-exp', + originalId: 'expired-original', + command: 'send_message', + targetConnectionId: 'cli-exp', + expiresAt: pastTime, + webConnectionId: 'web-exp', + state: 'pending' as const, + }); + + // Create a fresh DO to simulate a wake with the durable entry. + const doInstance2 = new UserConnectionDO(ctx as never, {} as never); + + // Set up a live CLI and a live web socket with the matching webConnectionId. + const cliWs = addCliSocket(mockCtx, 'cli-exp'); + sendHeartbeat(doInstance2, cliWs, [makeSession('ses-exp', 'busy', 'Session Exp')]); + // Web socket connects after the entry was created (simulates reattach). + const webWs = addWebSocket(mockCtx, 'web-exp'); + + // Trigger alarm which calls expirePendingCommands. + await doInstance2.alarm(); + await flushAsync(); + + // The reattached web socket must receive the expired error via originalId. + const responses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'expired-original' + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toEqual({ + source: 'relay', + code: 'COMMAND_EXPIRED', + message: 'Command expired', + }); + + // The durable entry must be marked done. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + }); + + it('delivers CLI-disconnect error to a reattached web socket via finishDurablePendingCommands', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'disco-delivery'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-disco', + originalId: 'disco-original', + command: 'send_message', + targetConnectionId: 'cli-disco', + expiresAt: now + 35_000, + webConnectionId: 'web-disco', + state: 'pending' as const, + }); + + // Set up a live CLI. + const cliWs = addCliSocket(mockCtx, 'cli-disco'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-disco', 'busy', 'Session Disco')]); + + // Web socket connects after the entry was created (simulates reattach). + const webWs = addWebSocket(mockCtx, 'web-disco'); + + // Disconnect the CLI — finishDurablePendingCommands handles the + // durable-only entries that were never in-memory. + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + await flushAsync(); + + // The reattached web socket must receive the disconnect error. + const responses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'disco-original' + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toBe('CLI disconnected'); + + // The durable entry must be marked done. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toBe('CLI disconnected'); + }); + + it('does not deliver terminal response when no live web socket matches webConnectionId', async () => { + const { mockCtx, ctx } = setup(); + + const pastTime = Date.now() - 10_000; + const correlationId = 'no-web-delivery'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-nwd', + originalId: 'no-web-original', + command: 'send_message', + targetConnectionId: 'cli-nwd', + expiresAt: pastTime, + webConnectionId: 'web-missing', + state: 'pending' as const, + }); + + const doInstance2 = new UserConnectionDO(ctx as never, {} as never); + + const cliWs = addCliSocket(mockCtx, 'cli-nwd'); + sendHeartbeat(doInstance2, cliWs, [makeSession('ses-nwd', 'busy', 'Session NWD')]); + // No web socket with connectionId 'web-missing'. + + // Alarm must not throw — storage-only behavior when no socket matches. + await doInstance2.alarm(); + await flushAsync(); + + // The durable entry must be marked done (storage-only). + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + }); + + // ------------------------------------------------------------------------- + // Fix: Oversized durable result returns error on mutationId retry + // ------------------------------------------------------------------------- + it('returns DURABLE_RESULT_TOO_LARGE on a mutationId retry of an oversized result', async () => { + const { doInstance, mockCtx } = setup(); + + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'oversized-mut', + command: 'send_message', + sessionId: 's1', + mutationId: 'mut-oversized', + data: { hello: 'world' }, + }); + await flushAsync(); + + // CLI sends an oversized result. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(1); + const correlationId = cliCommands[0].id as string; + webWs.send.mockClear(); + + const oversized = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES + 1); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); + + // Live response must carry the full result. + const live = parseSent(webWs) as { id: string; result: unknown }; + expect(live.id).toBe('oversized-mut'); + expect(live.result).toEqual(oversized); + + // Now retry with the same mutationId. + webWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'retry-mut', + command: 'send_message', + sessionId: 's1', + mutationId: 'mut-oversized', + }); + await flushAsync(); + + // The retry must receive the error, not the truncated marker as a result. + const retry = parseSent(webWs) as { id: string; error?: unknown; result?: unknown }; + expect(retry.id).toBe('retry-mut'); + expect(retry.result).toBeUndefined(); + expect(retry.error).toEqual({ + source: 'relay', + code: 'DURABLE_RESULT_TOO_LARGE', + message: 'Result is too large to store for retries', + }); + }); + + // ------------------------------------------------------------------------- + // Fix: Default dispatch awaits durable persistence before forwarding + // ------------------------------------------------------------------------- + it('persists durably before forwarding a default (non-mutationId) command to the CLI', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + + // Track the order of storage.put vs send calls. + const callOrder: string[] = []; + const originalPut = ctx.storage.put; + ctx.storage.put = vi.fn(async (...args: unknown[]) => { + callOrder.push('put'); + return (originalPut as any)(...args); + }); + // Spy on cliWs.send to track order. + const originalSend = cliWs.send; + cliWs.send = vi.fn((...args: unknown[]) => { + callOrder.push('send'); + return (originalSend as any).apply(cliWs, args); + }); + + await sendCommand(doInstance, webWs, { + id: 'order-cmd', + command: 'send_message', + sessionId: 's1', + }); + + // Storage.put must be called before sendToCli. + // Since sendToCli is now chained via .then() inside waitUntil, + // the put call is registered first (promise creation), but the + // actual send happens after put resolves. + await flushAsync(); + + // Storage.put must have been called. + expect(ctx.storage.put).toHaveBeenCalled(); + + // The command must have been forwarded to the CLI. + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(1); + expect(cliCommands[0]).toMatchObject({ + type: 'command', + command: 'send_message', + sessionId: 's1', + }); + + // Verify that put was called before send by checking that both + // operations completed and the call order is correct. + const putIdx = callOrder.indexOf('put'); + const sendIdx = callOrder.indexOf('send'); + expect(putIdx).toBeGreaterThanOrEqual(0); + expect(sendIdx).toBeGreaterThanOrEqual(0); + expect(putIdx).toBeLessThan(sendIdx); + }); + + // ------------------------------------------------------------------------- + // Race repair: at-most-once terminal delivery on CLI disconnect + // ------------------------------------------------------------------------- + it('delivers exactly one CLI-disconnect error when entry exists in both memory and durable storage', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + await sendCommand(doInstance, webWs, { + id: 'race-disco', + command: 'send_message', + sessionId: 's1', + }); + webWs.send.mockClear(); + + // Disconnect CLI — failPendingCommandsForSocket processes the in-memory + // entry, then finishDurablePendingCommands scans durable entries. + // The fix passes handledIds to skip the duplicate. + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + await flushAsync(); + + // Exactly one CLI-disconnect error, not two. + const errors = allSent(webWs).filter(m => m.type === 'response' && m.id === 'race-disco'); + expect(errors).toHaveLength(1); + expect(errors[0].error).toBe('CLI disconnected'); + }); + + // ------------------------------------------------------------------------- + // Race repair: at-most-once terminal delivery on owner change + // ------------------------------------------------------------------------- + it('delivers exactly one SESSION_OWNER_CHANGED error when entry exists in both memory and durable storage', async () => { + const { doInstance, mockCtx } = setup(); + const firstOwner = addCliSocket(mockCtx, 'cli-1'); + const nextOwner = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, firstOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, nextOwner, []); + firstOwner.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'race-owner', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', + }); + webWs.send.mockClear(); + + // Ownership change via heartbeat — failPendingCommandsForOwnerChange + // processes the in-memory entry, then finishDurablePendingCommands + // scans durable entries. The fix passes handledIds to skip the duplicate. + // failPendingCommandsForOwnerChange runs inside ctx.waitUntil; drain + // microtasks so both the in-memory and durable sweeps settle. + sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + await flushAsync(); + + // Exactly one SESSION_OWNER_CHANGED error. + const errors = allSent(webWs).filter(m => m.type === 'response' && m.id === 'race-owner'); + expect(errors).toHaveLength(1); + expect(errors[0].error).toEqual({ + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }); + }); + + // ------------------------------------------------------------------------- + // Race repair: durable persisted before live send in finishDurablePendingCommands + // ------------------------------------------------------------------------- + it('persists the durable entry before delivering to a reattached web socket in finishDurablePendingCommands', async () => { + const { doInstance, mockCtx, ctx } = setup(); + + const now = Date.now(); + const correlationId = 'disco-reorder'; + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 'ses-reorder', + originalId: 'reorder-original', + command: 'send_message', + targetConnectionId: 'cli-reorder', + expiresAt: now + 35_000, + webConnectionId: 'web-reorder', + state: 'pending' as const, + }); + + // Set up a live CLI and a reattached web socket. + const cliWs = addCliSocket(mockCtx, 'cli-reorder'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-reorder', 'busy', 'Session Reorder')]); + const webWs = addWebSocket(mockCtx, 'web-reorder'); + + // Track call order: storage.put vs ws.send inside finishDurablePendingCommands. + const callOrder: string[] = []; + const originalPut = ctx.storage.put; + ctx.storage.put = vi.fn(async (...args: unknown[]) => { + callOrder.push('put'); + return (originalPut as any)(...args); + }); + const originalSend = webWs.send; + webWs.send = vi.fn((...args: unknown[]) => { + callOrder.push('send'); + return (originalSend as any).apply(webWs, args); + }); + + // Disconnect the CLI — finishDurablePendingCommands handles the + // durable-only entry. + mockCtx.removeSocket(cliWs); + await disconnectCli(doInstance, cliWs); + await flushAsync(); + + // The reattached web socket must receive the error. + const responses = allSent(webWs).filter( + m => m.type === 'response' && m.id === 'reorder-original' + ); + expect(responses).toHaveLength(1); + expect(responses[0].error).toBe('CLI disconnected'); + + // The durable entry must be marked done. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + + // The durable put must happen before the web socket send. + const putIdx = callOrder.indexOf('put'); + const sendIdx = callOrder.indexOf('send'); + expect(putIdx).toBeGreaterThanOrEqual(0); + expect(sendIdx).toBeGreaterThanOrEqual(0); + expect(putIdx).toBeLessThan(sendIdx); + }); + }); + + // ------------------------------------------------------------------------- + // Concurrent durable sweep fence (Fix: at-most-once delivery) + // ------------------------------------------------------------------------- + + describe('concurrent durable sweep fence', () => { + it('expirePendingCommands sync sweep adds to completedCorrelationIds', async () => { + const now = 1_000_000; + vi.spyOn(Date, 'now').mockReturnValue(now); + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Advance past expiry. + vi.mocked(Date.now).mockReturnValue(now + 35_001); + (doInstance as unknown as { expirePendingCommands(n: number): void }).expirePendingCommands( + now + 35_001 + ); + + // The sync sweep delivered the expiry error. + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + error: { source: 'relay', code: 'COMMAND_EXPIRED', message: 'Command expired' }, + }); + + // The correlationId must be in completedCorrelationIds so a concurrent + // finishDurablePendingCommands does not double-deliver. + const completed = (doInstance as unknown as { completedCorrelationIds: Set }) + .completedCorrelationIds; + expect(completed.has(correlationId)).toBe(true); + }); + + it('expirePendingCommands async sweep adds to completedCorrelationIds before delivery', async () => { + const now = 1_000_000; + vi.spyOn(Date, 'now').mockReturnValue(now); + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Advance past expiry. The sync sweep delivers; the async sweep will run + // during flushAsync. + vi.mocked(Date.now).mockReturnValue(now + 35_001); + (doInstance as unknown as { expirePendingCommands(n: number): void }).expirePendingCommands( + now + 35_001 + ); + + // Settle the waitUntil durable sweep. + await flushAsync(); + + // The durable entry must be marked done. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record)?.state).toBe('done'); + }); + + it('finishDurablePendingCommands skips entries already in completedCorrelationIds', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Manually add to completedCorrelationIds (simulating an + // expirePendingCommands sync sweep that ran first). + ( + doInstance as unknown as { completedCorrelationIds: Set } + ).completedCorrelationIds.add(correlationId); + + // finishDurablePendingCommands must not deliver a second response. + await ( + doInstance as unknown as { + finishDurablePendingCommands( + matches: (e: unknown) => boolean, + error: unknown, + skipIds?: ReadonlySet + ): Promise; + } + ).finishDurablePendingCommands( + (e: unknown) => (e as Record).state === 'pending', + 'CLI disconnected', + new Set() + ); + + // No additional delivery. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-1'); + expect(responses).toHaveLength(0); + }); + + it('expirePendingCommands sync sweep skips entry already reserved in completedCorrelationIds', async () => { + const now = 1_000_000; + vi.spyOn(Date, 'now').mockReturnValue(now); + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Manually add to completedCorrelationIds (simulating a concurrent + // finishDurablePendingCommands that reserved the id first). + ( + doInstance as unknown as { completedCorrelationIds: Set } + ).completedCorrelationIds.add(correlationId); + + // Advance past expiry and call expirePendingCommands sync sweep. + vi.mocked(Date.now).mockReturnValue(now + 35_001); + (doInstance as unknown as { expirePendingCommands(n: number): void }).expirePendingCommands( + now + 35_001 + ); + + // The sync sweep must NOT deliver a second response. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-1'); + expect(responses).toHaveLength(0); + }); + + it('expirePendingCommands async sweep skips durable entry already reserved in completedCorrelationIds', async () => { + const now = 1_000_000; + vi.spyOn(Date, 'now').mockReturnValue(now); + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // Manually add to completedCorrelationIds (simulating a concurrent + // finishDurablePendingCommands that reserved the id first). + ( + doInstance as unknown as { completedCorrelationIds: Set } + ).completedCorrelationIds.add(correlationId); + + // Advance past expiry. The sync sweep will see the entry in + // pendingCommands but skip it because completedCorrelationIds already + // has it. Then the async sweep scans durable storage. + vi.mocked(Date.now).mockReturnValue(now + 35_001); + (doInstance as unknown as { expirePendingCommands(n: number): void }).expirePendingCommands( + now + 35_001 + ); + + // Settle the waitUntil durable sweep. + await flushAsync(); + + // The async sweep must NOT deliver a second response. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-1'); + expect(responses).toHaveLength(0); + }); + }); + + // ------------------------------------------------------------------------- + // Rejected initial write cleanup (Fix: no stale stash leak) + // ------------------------------------------------------------------------- + + describe('rejected initial write cleanup', () => { + it('cleans terminalDuringInitialWrite and completedCorrelationIds when initial write rejects', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + webWs.send.mockClear(); + + // Make the first storage.put reject — this is the initial pending write. + const originalPut = ctx.storage.put; + let initialWriteRejected = false; + ctx.storage.put = vi.fn(async (...args: unknown[]) => { + if (!initialWriteRejected) { + initialWriteRejected = true; + throw new Error('simulated write failure'); + } + return (originalPut as any)(...args); + }); + + // Send a command — dispatchWebCommandSync triggers the initial write + // inside waitUntil. Do NOT auto-flush yet. + const msg = JSON.stringify({ + type: 'command', + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + doInstance.webSocketMessage(webWs as never, msg); + + // Flush enough for the initial write to fail and the .catch + .finally + // to execute. The command was never forwarded to the CLI. + await flushAsync(); + + // No command was sent to CLI (the write failed before forwarding). + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); + expect(cliCommands).toHaveLength(0); + + // The web received no response yet (command was never forwarded). + expect(webWs.send).not.toHaveBeenCalled(); + }); + + it('cleans terminalDuringInitialWrite when initial write fails after a terminal stash', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + webWs.send.mockClear(); + + // Use a controllable promise so the initial write stays pending. + let _resolveInitialWrite: (() => void) | undefined; + let rejectInitialWrite: ((err: Error) => void) | undefined; + const initialWritePromise = new Promise((resolve, reject) => { + _resolveInitialWrite = resolve; + rejectInitialWrite = reject; + }); + + let initialWriteCalled = false; + const originalPut = ctx.storage.put; + ctx.storage.put = vi.fn(async (...args: unknown[]) => { + if (!initialWriteCalled) { + initialWriteCalled = true; + return initialWritePromise; + } + return (originalPut as any)(...args); + }); + + // Send command — initial write is now pending. + const msg = JSON.stringify({ + type: 'command', + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + doInstance.webSocketMessage(webWs as never, msg); + + // Drain microtasks so the waitUntil launches and the initial write + // promise is awaited (but not yet resolved). + await flushAsync(); + + // Now disconnect the CLI. failPendingCommandsForSocket finds the + // entry in pendingCommands and stashes a terminal in + // terminalDuringInitialWrite because the initial write is still pending. + mockCtx.removeSocket(cliWs); + const disconnectPromise = disconnectCli(doInstance, cliWs); + + // Drain microtasks so failPendingCommandsForSocket processes the entry. + await flushAsync(); + + // The terminal stash must exist before the initial write rejects. + const terminalStash = ( + doInstance as unknown as { + terminalDuringInitialWrite: Map; + } + ).terminalDuringInitialWrite; + expect(terminalStash.size).toBe(1); + + // Now reject the initial write. + rejectInitialWrite!(new Error('simulated write failure')); + + // Let the disconnect finish (including its own durable write retry). + await disconnectPromise; + await flushAsync(); + + // After the initial write rejection, the terminal stash and + // completedCorrelationIds must be cleaned. + expect(terminalStash.size).toBe(0); + }); + + it('.finally does not clear completedCorrelationIds when expirePendingCommands already removed the entry', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + webWs.send.mockClear(); + + // Use a controllable promise for the initial write. + let resolveInitialWrite: (() => void) | undefined; + const initialWritePromise = new Promise(resolve => { + resolveInitialWrite = resolve; + }); + + let initialWriteCalled = false; + const originalPut = ctx.storage.put; + ctx.storage.put = vi.fn(async (...args: unknown[]) => { + if (!initialWriteCalled) { + initialWriteCalled = true; + return initialWritePromise; + } + return (originalPut as any)(...args); + }); + + // Send command — initial write is pending. + const msg = JSON.stringify({ + type: 'command', + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + doInstance.webSocketMessage(webWs as never, msg); + await flushAsync(); + + // Get the correlationId that was assigned. + const completed = (doInstance as unknown as { completedCorrelationIds: Set }) + .completedCorrelationIds; + const pendingCommands = ( + doInstance as unknown as { + pendingCommands: Map; + } + ).pendingCommands; + const correlationIds = [...pendingCommands.keys()]; + expect(correlationIds).toHaveLength(1); + const correlationId = correlationIds[0]!; + + // Simulate expiry: remove from pendingCommands and add to completedCorrelationIds. + pendingCommands.delete(correlationId); + completed.add(correlationId); + + // Now resolve the initial write. The .then runs (no terminal stash), + // then .finally runs. The .finally must NOT clear completedCorrelationIds + // because the entry is no longer in pendingCommands (was handled by expiry). + resolveInitialWrite!(); + await flushAsync(); + + // completedCorrelationIds must still contain the correlationId because + // the expiry path owns cleanup (via the async durable sweep). + expect(completed.has(correlationId)).toBe(true); + }); + }); + + // ------------------------------------------------------------------------- + // Concurrent sweep claim gate (terminal sweep final repair) + // ------------------------------------------------------------------------- + + describe('concurrent sweep stale-snapshot claim gate', () => { + it('a stale list snapshot cannot pass through an awaited put after another sweep delivered', async () => { + const { doInstance, mockCtx, ctx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + + const correlationId = 'corr-sweep-gate'; + const now = 3_000_000; + vi.spyOn(Date, 'now').mockReturnValue(now); + + // Insert an expired durable-only pending entry (not in-memory). + await ctx.storage.put(`pendingCommand/${correlationId}`, { + sessionId: 's1', + originalId: 'cmd-gate', + command: 'send_message', + expectedOwnerConnectionId: 'cli-1', + targetConnectionId: 'cli-1', + expiresAt: now - 1, + webConnectionId: 'web-1', + state: 'pending', + }); + + cliWs.send.mockClear(); + webWs.send.mockClear(); + + // Override storage.list: the second call returns a stale snapshot that + // still shows the entry as 'pending', simulating a snapshot captured + // before the first sweep's storage.put completed. + const realList = ctx.storage.list; + let callCount = 0; + ctx.storage.list = vi.fn(async (opts?: { prefix?: string }) => { + callCount++; + if (callCount === 1) { + return realList(opts); + } + // Stale snapshot: the entry still appears pending. + const staleMap = new Map(); + staleMap.set(`pendingCommand/${correlationId}`, { + sessionId: 's1', + originalId: 'cmd-gate', + command: 'send_message', + expectedOwnerConnectionId: 'cli-1', + targetConnectionId: 'cli-1', + expiresAt: now - 1, + webConnectionId: 'web-1', + state: 'pending', + }); + return staleMap as Awaited>; + }); + + // First sweep: expirePendingCommands (via alarm) claims and delivers. + await doInstance.alarm(); + await flushAsync(); + + // Second sweep: finishDurablePendingCommands (via owner change) sees + // the stale snapshot but the claim gate prevents double delivery. + const cli2 = addCliSocket(mockCtx, 'cli-2'); + sendHeartbeat(doInstance, cli2, [makeSession('s1')]); + await flushAsync(); + + // Exactly one terminal response must be delivered to the web socket. + const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-gate'); + expect(responses).toHaveLength(1); + expect(responses[0]).toEqual({ + type: 'response', + id: 'cmd-gate', + error: { + source: 'relay', + code: 'COMMAND_EXPIRED', + message: 'Command expired', + }, + }); + }); + }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 1c092be5e1a..d6bc20ac929 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -53,7 +53,12 @@ type WSAttachment = // fresh and avoids stale instance rows on restart. instance?: Instance; } - | { role: 'web'; connectionId: string; subscribedSessions: string[]; replaced?: true }; + | { + role: 'web'; + connectionId: string; + subscribedSessions: string[]; + replaced?: true; + }; // Type re-export so test files and other internal callers can reference the // connection-row shape from a single place. @@ -70,6 +75,16 @@ export type ConnectedInstanceRow = { export const MAX_CATALOG_RESULT_BYTES = 512 * 1024; +// Maximum durable result size before truncation. Results larger than this are +// replaced with a size-exceeded marker before the storage.put so the DO value +// limit (128 KiB) is never hit. 120 KiB leaves generous headroom for the +// serialized PendingCommandEntry framing overhead while preserving enough +// payload for most retry scenarios. +export const MAX_DURABLE_RESULT_BYTES = 120 * 1024; + +// Maximum allowed mutationId length used as a durable storage key prefix. +const MAX_MUTATION_ID_LENGTH = 128; + // Viewer command allowlist. Anything outside this set is rejected by the relay // before owner resolution, pending allocation, or CLI forwarding. export const ALLOWED_VIEWER_COMMANDS: ReadonlySet = new Set([ @@ -160,6 +175,18 @@ const CLI_COMMAND_ERROR = { message: 'Command failed', }; +const DURABLE_RESULT_TOO_LARGE_ERROR = { + source: 'relay', + code: 'DURABLE_RESULT_TOO_LARGE', + message: 'Result is too large to store for retries', +}; + +const MUTATION_ID_TOO_LONG_ERROR = { + source: 'relay', + code: 'MUTATION_ID_TOO_LONG', + message: 'Mutation ID is too long', +}; + type ReadyPushEntry = { kiloUserId: string; title: string; @@ -172,8 +199,25 @@ type RenameEntry = { at: number; }; +type PendingCommandEntry = { + sessionId?: string; + originalId: string; + command: string; + expectedOwnerConnectionId?: string; + targetConnectionId: string; + expiresAt: number; + // The originating web socket's connectionId (from attachment). + // Stable across hibernation; used to resolve the socket on wake. + webConnectionId: string; + // D8 state: 'pending' or 'done'. When 'done', result or error holds the outcome. + state: 'pending' | 'done'; + result?: unknown; + error?: unknown; +}; + const READY_PUSH_KEY_PREFIX = 'readyPush:'; const RENAME_KEY_PREFIX = 'rename:'; +const PENDING_COMMAND_KEY_PREFIX = 'pendingCommand/'; const SESSION_READY_PUSH_DELAY_MS = 5_000; /** Backoff between ready-push claim retries so the 3-attempt bound spans real time. */ const READY_PUSH_RETRY_BACKOFF_MS = 5_000; @@ -210,6 +254,8 @@ export class UserConnectionDO extends DurableObject { targetCliWs: WebSocket; } >(); + private pendingInitialCommandWrites = new Set(); + private terminalDuringInitialWrite = new Map(); // Last heartbeat timestamp per CLI connectionId (for staleness eviction) private lastHeartbeatAt = new Map(); // In-memory mirror of readyPush KV fireAt values — scheduling only; KV is source of truth @@ -217,6 +263,15 @@ export class UserConnectionDO extends DurableObject { // One-shot post-eviction rebuild gate; reset whenever ensureState actually reconstructs private readyPushRebuilt = false; + // Synchronous reservation set for mutationId dedupe (prevents concurrent same-ID dispatch + // across asynchronous storage reads in the mutationId path). + private inflightMutations = new Set(); + + // Synchronous completion set for CLI responses already delivered live. + // Prevents the rehydration path from delivering a duplicate response while the + // durable write to 'done' is still in-flight. + private completedCorrelationIds = new Set(); + private stateReconstructed = false; private ensureState(): void { @@ -287,7 +342,12 @@ export class UserConnectionDO extends DurableObject { const reconnect = this.closeStaleSocket(connectionId); const kiloUserId = url.searchParams.get('kiloUserId') ?? undefined; - const attachment: WSAttachment = { role: 'cli', connectionId, sessions: [], kiloUserId }; + const attachment: WSAttachment = { + role: 'cli', + connectionId, + sessions: [], + kiloUserId, + }; this.ctx.acceptWebSocket(server, ['cli']); server.serializeAttachment(attachment); const now = Date.now(); @@ -310,7 +370,11 @@ export class UserConnectionDO extends DurableObject { } else { this.replaceWebSocket(connectionId); - const attachment: WSAttachment = { role: 'web', connectionId, subscribedSessions: [] }; + const attachment: WSAttachment = { + role: 'web', + connectionId, + subscribedSessions: [], + }; this.ctx.acceptWebSocket(server, ['web']); server.serializeAttachment(attachment); @@ -409,7 +473,9 @@ export class UserConnectionDO extends DurableObject { for (const ws of this.ctx.getWebSockets('cli')) { const att = ws.deserializeAttachment() as WSAttachment | null; if (att?.role === 'cli' && att.connectionId === connectionId) { - console.log('Closing stale CLI connection (heartbeat timeout)', { connectionId }); + console.log('Closing stale CLI connection (heartbeat timeout)', { + connectionId, + }); ws.close(4408, 'heartbeat timeout'); break; } @@ -419,6 +485,7 @@ export class UserConnectionDO extends DurableObject { } this.scheduleNextAlarm(now); + this.scheduleDurablePendingAlarm(); } // --------------------------------------------------------------------------- @@ -438,7 +505,10 @@ export class UserConnectionDO extends DurableObject { role: 'cli', connectionId: attachment.connectionId, byteCount: binaryByteCount ?? new TextEncoder().encode(raw).byteLength, - issues: result.error.issues.map(issue => ({ path: issue.path, code: issue.code })), + issues: result.error.issues.map(issue => ({ + path: issue.path, + code: issue.code, + })), }); return; } @@ -459,7 +529,19 @@ export class UserConnectionDO extends DurableObject { this.handleCliEvent(msg.sessionId, msg.parentSessionId, msg.event, msg.data); break; case 'response': - this.handleCliResponse(ws, msg.id, msg.result, msg.error); + // Extend the DO lifetime for the async durable write (fix: reply + // must persist before hibernation, but the surrounding handler stays + // synchronous for existing callers). + // Catches storage read/write failures so the waitUntil task is + // never left unhandled (rejected). + this.ctx.waitUntil( + this.handleCliResponse(ws, msg.id, msg.result, msg.error).catch((error: unknown) => { + console.error('Failed to handle CLI response (non-fatal)', { + correlationId: msg.id, + error: error instanceof Error ? error.message : String(error), + }); + }) + ); break; } } @@ -485,7 +567,14 @@ export class UserConnectionDO extends DurableObject { for (const prev of previousSessions) { if (!currentIds.has(prev.id) && this.sessionOwners.get(prev.id) === connectionId) { this.sessionOwners.delete(prev.id); - this.failPendingCommandsForOwnerChange(prev.id, undefined); + this.ctx.waitUntil( + this.failPendingCommandsForOwnerChange(prev.id, undefined).catch((error: unknown) => { + console.error('Failed to persist terminal commands for dropped session', { + sessionId: prev.id, + error: error instanceof Error ? error.message : String(error), + }); + }) + ); } } @@ -494,7 +583,16 @@ export class UserConnectionDO extends DurableObject { for (const session of sessions) { const previousOwner = this.sessionOwners.get(session.id); if (previousOwner && previousOwner !== connectionId) { - this.failPendingCommandsForOwnerChange(session.id, connectionId); + this.ctx.waitUntil( + this.failPendingCommandsForOwnerChange(session.id, connectionId).catch( + (error: unknown) => { + console.error('Failed to persist terminal commands for owner change', { + sessionId: session.id, + error: error instanceof Error ? error.message : String(error), + }); + } + ) + ); } // First sight of a main session on this DO: schedule a delayed session-ready // push (Decision 9). The durable claim in SessionIngestDO makes reconnect @@ -610,7 +708,9 @@ export class UserConnectionDO extends DurableObject { private async fireReadyPushes(now: number): Promise { let pending: Map; try { - pending = await this.ctx.storage.list({ prefix: READY_PUSH_KEY_PREFIX }); + pending = await this.ctx.storage.list({ + prefix: READY_PUSH_KEY_PREFIX, + }); } catch (error: unknown) { console.error('Failed to list readyPush entries', { error: error instanceof Error ? error.message : String(error), @@ -745,63 +845,343 @@ export class UserConnectionDO extends DurableObject { } } - private handleCliResponse( + private async handleCliResponse( respondingWs: WebSocket, id: string, result: unknown, error: unknown - ): void { - const entry = this.pendingCommands.get(id); - if (!entry || entry.targetCliWs !== respondingWs) return; - this.pendingCommands.delete(id); + ): Promise { + let entry = this.pendingCommands.get(id); + let rehydrated = false; + // Captured at the first durable read so the catalog-too-large branch + // does not need a second read (fix: second read can return undefined + // and abandon terminal processing). + let rehydratedDurable: PendingCommandEntry | undefined; + let ownsReservation = false; + + // Step 24: on in-memory miss, try to load the durable entry (D8 case 1). + if (!entry) { + // Synchronously reserve the correlationId before the async durable + // read so a concurrent duplicate CLI reply cannot also process the + // same pending durable state after a wake. + if (this.completedCorrelationIds.has(id)) return; + this.completedCorrelationIds.add(id); + ownsReservation = true; - if (CATALOG_DEDUPE_COMMANDS.has(entry.command) && result !== undefined) { - const serializedResult = JSON.stringify(result); - const resultBytes = new TextEncoder().encode(serializedResult).byteLength; - if (resultBytes > MAX_CATALOG_RESULT_BYTES) { - this.sendToWeb(entry.ws, { - type: 'response', - id: entry.originalId, - error: CATALOG_TOO_LARGE_ERROR, - }); - return; + try { + const durable = await this.getDurablePendingCommand(id); + if (!durable || durable.state !== 'pending') { + this.completedCorrelationIds.delete(id); + ownsReservation = false; + return; + } + rehydratedDurable = durable; + + // Guard: a rehydrated entry that has already expired must not deliver + // or persist a late success outcome. Write the terminal expiry error + // durably and return without delivering. + if (durable.expiresAt <= Date.now()) { + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { + ...durable, + state: 'done', + error: COMMAND_EXPIRED_ERROR, + }); + // Reservation held through the terminal write; clear after success. + this.completedCorrelationIds.delete(id); + ownsReservation = false; + return; + } + + // Validate the responding CLI by attachment connectionId. + const respondingAttachment = respondingWs.deserializeAttachment() as WSAttachment | null; + if ( + respondingAttachment?.role !== 'cli' || + respondingAttachment.connectionId !== durable.targetConnectionId + ) { + this.completedCorrelationIds.delete(id); + ownsReservation = false; + return; + } + + // Resolve the originating web socket by persisted webConnectionId. + const webWs = this.findWebByConnectionId(durable.webConnectionId); + if (!webWs) { + // D8 case 2: web socket is gone. Shape the terminal outcome + // (catalog guard, error normalization) before persisting done + // so a mutationId retry receives the same shaped outcome as + // live delivery. + const shaped = this.shapeCommandError(durable.command, result, error); + let shapedResult = this.boundDurableResult(shaped.result); + let shapedError = shaped.error; + // When the durable result is truncated, a mutationId retry must + // receive an error — not a false-success truncated marker. + if (this.isTruncatedResult(shapedResult)) { + shapedResult = undefined; + shapedError = shapedError ?? DURABLE_RESULT_TOO_LARGE_ERROR; + } + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { + ...durable, + state: 'done', + ...(shapedResult !== undefined ? { result: shapedResult } : {}), + ...(shapedError !== undefined ? { error: shapedError } : {}), + }); + // Reservation held through the terminal write; clear after success. + this.completedCorrelationIds.delete(id); + ownsReservation = false; + return; + } + + // Build an in-memory entry so the rest of handleCliResponse works. + entry = { + ws: webWs, + sessionId: durable.sessionId, + originalId: durable.originalId, + command: durable.command, + expectedOwnerConnectionId: durable.expectedOwnerConnectionId, + targetConnectionId: durable.targetConnectionId, + expiresAt: durable.expiresAt, + targetCliWs: respondingWs, + }; + rehydrated = true; + // Reservation held through the rest of handleCliResponse — cleared + // after terminal storage.put succeeds (catalog guard and normal + // path) or on any throw. + } catch (error: unknown) { + // Clean reservation on unexpected error during rehydration. + this.completedCorrelationIds.delete(id); + ownsReservation = false; + throw error; } + } else if (entry.targetCliWs !== respondingWs) { + // Strict socket-identity check for an entry still in memory. + return; } - let structuredError: { source: string; code: string; message: string } | null = null; - let stringError: string | null = null; - let sanitizeAsFailed = false; + // All code after acquisition is protected: catalog serialization, + // boundDurableResult, attachment reads, and storage.put are inside + // this try block. The finally clears any reservation that survives + // a throw. + try { + // Validate catalog result size for rehydrated entries too. + if (CATALOG_DEDUPE_COMMANDS.has(entry.command) && result !== undefined) { + const serializedResult = JSON.stringify(result); + const resultBytes = new TextEncoder().encode(serializedResult).byteLength; + if (resultBytes > MAX_CATALOG_RESULT_BYTES) { + // For non-rehydrated entries, reserve the correlationId before the + // durable write so a concurrent CLI reply is fenced. + if (!rehydrated) { + // Guard: do not override an already-expired terminal outcome. + const now = Date.now(); + if (entry.expiresAt <= now) { + this.pendingCommands.delete(id); + return; + } + this.pendingCommands.delete(id); + // Reserve the correlationId synchronously so a concurrent CLI + // reply cannot re-process the same command during the durable + // write window. + this.completedCorrelationIds.add(id); + ownsReservation = true; + } - if (typeof error === 'string' && CLI_UPGRADE_REQUIRED_COMMANDS.has(entry.command)) { - if (error === `unknown command: ${entry.command}`) { - structuredError = - entry.command === 'create_session' - ? CLI_UPGRADE_REQUIRED_CREATE_SESSION_ERROR - : CLI_UPGRADE_REQUIRED_SLASH_ERROR; - } else { + // Persist the terminal outcome durably BEFORE sending any live + // response. If storage fails, no live send occurs and reservations + // are cleaned — the error propagates to waitUntil. + let catWebConnectionId: string; + if (rehydrated) { + // Use the durable entry captured at the top of handleCliResponse. + // A second getDurablePendingCommand can miss (race), which would + // abandon terminal processing. + catWebConnectionId = rehydratedDurable?.webConnectionId ?? 'unknown'; + } else { + catWebConnectionId = + (entry.ws.deserializeAttachment() as WSAttachment | null)?.role === 'web' + ? ( + entry.ws.deserializeAttachment() as WSAttachment & { + role: 'web'; + } + ).connectionId + : 'unknown'; + } + try { + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { + sessionId: entry.sessionId, + originalId: entry.originalId, + command: entry.command, + expectedOwnerConnectionId: entry.expectedOwnerConnectionId, + targetConnectionId: entry.targetConnectionId, + expiresAt: entry.expiresAt, + webConnectionId: catWebConnectionId, + state: 'done', + error: CATALOG_TOO_LARGE_ERROR, + } satisfies PendingCommandEntry); + } catch (_error: unknown) { + this.completedCorrelationIds.delete(id); + ownsReservation = false; + // Re-throw to the surrounding waitUntil catch so the task + // failure is logged and never left unhandled. + throw _error; + } + + // Durable write succeeded — clear the marker. The durable 'done' + // state now guards against duplicate delivery. + this.completedCorrelationIds.delete(id); + ownsReservation = false; + + // Send the live response now that the durable entry is stored. + if (!rehydrated) { + this.sendToWeb(entry.ws, { + type: 'response', + id: entry.originalId, + error: CATALOG_TOO_LARGE_ERROR, + }); + } else { + const targetWeb = this.findWebByConnectionId(catWebConnectionId); + if (targetWeb) { + this.sendToWeb(targetWeb, { + type: 'response', + id: entry.originalId, + error: CATALOG_TOO_LARGE_ERROR, + }); + } + } + return; + } + } + + let structuredError: { + source: string; + code: string; + message: string; + } | null = null; + let stringError: string | null = null; + let sanitizeAsFailed = false; + + if (typeof error === 'string' && CLI_UPGRADE_REQUIRED_COMMANDS.has(entry.command)) { + if (error === `unknown command: ${entry.command}`) { + structuredError = + entry.command === 'create_session' + ? CLI_UPGRADE_REQUIRED_CREATE_SESSION_ERROR + : CLI_UPGRADE_REQUIRED_SLASH_ERROR; + } else { + stringError = error; + } + } else if (typeof error === 'string') { stringError = error; + } else if (error !== undefined) { + sanitizeAsFailed = true; } - } else if (typeof error === 'string') { - stringError = error; - } else if (error !== undefined) { - // Arbitrary non-string CLI error (e.g. a CLI-shaped structured object): - // sanitize to a generic CLI failure so the relay's own error policy - // remains the source of truth. - sanitizeAsFailed = true; - } - this.sendToWeb(entry.ws, { - type: 'response', - id: entry.originalId, - ...(result !== undefined ? { result } : {}), - ...(structuredError !== null - ? { error: structuredError } - : stringError !== null - ? { error: stringError } - : sanitizeAsFailed - ? { error: CLI_COMMAND_ERROR } - : {}), - }); + const terminalError = + structuredError ?? stringError ?? (sanitizeAsFailed ? CLI_COMMAND_ERROR : undefined); + + // Write the terminal outcome durably BEFORE any live response. + // If storage.put fails, the error propagates to the waitUntil boundary + // — reservation markers are cleaned and no live send has occurred. + // + // Guard: a live entry that has already expired must not override the + // terminal outcome set by expirePendingCommands. Drop silently. + if (!rehydrated) { + const now = Date.now(); + if (entry.expiresAt <= now) { + this.pendingCommands.delete(id); + return; + } + this.pendingCommands.delete(id); + // Reserve the correlationId synchronously so a concurrent CLI + // reply cannot re-process the same command during the durable + // write window. + this.completedCorrelationIds.add(id); + ownsReservation = true; + } + + // boundDurableResult, attachment reads, and storage.put are inside this + // try block so any throw (e.g. JSON.stringify on oversized results) + // is caught by the outer finally, which clears the reservation. + let durableResult: unknown; + let webConnectionId: string; + try { + durableResult = this.boundDurableResult(result); + // When the durable result is truncated, a mutationId retry must + // receive an error — not a false-success truncated marker. The live + // response still carries the original full result via the `result` + // variable (not `durableResult`). + let durableError = terminalError; + if (this.isTruncatedResult(durableResult)) { + durableResult = undefined; + durableError = durableError ?? DURABLE_RESULT_TOO_LARGE_ERROR; + } + webConnectionId = rehydrated + ? (rehydratedDurable?.webConnectionId ?? 'unknown') + : (entry.ws.deserializeAttachment() as WSAttachment | null)?.role === 'web' + ? (entry.ws.deserializeAttachment() as WSAttachment & { role: 'web' }).connectionId + : 'unknown'; + + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { + sessionId: entry.sessionId, + originalId: entry.originalId, + command: entry.command, + expectedOwnerConnectionId: entry.expectedOwnerConnectionId, + targetConnectionId: entry.targetConnectionId, + expiresAt: entry.expiresAt, + webConnectionId: webConnectionId ?? 'unknown', + state: 'done', + ...(durableResult !== undefined ? { result: durableResult } : {}), + ...(durableError !== undefined ? { error: durableError } : {}), + } satisfies PendingCommandEntry); + } catch (_error: unknown) { + // Clean the synchronous reservation on storage failure — the error + // propagates to waitUntil; no live send has occurred. + this.completedCorrelationIds.delete(id); + ownsReservation = false; + // Re-throw to the surrounding waitUntil catch so the task + // failure is logged and never left unhandled. + throw _error; + } + + // Durable write succeeded. Clear the in-memory marker on both paths: + // the durable 'done' state replaces it as the dedupe guard. + this.completedCorrelationIds.delete(id); + ownsReservation = false; + + if (!rehydrated) { + // Send the live response now that the durable entry is safely stored. + this.sendToWeb(entry.ws, { + type: 'response', + id: entry.originalId, + ...(result !== undefined ? { result } : {}), + ...(structuredError !== null + ? { error: structuredError } + : stringError !== null + ? { error: stringError } + : sanitizeAsFailed + ? { error: CLI_COMMAND_ERROR } + : {}), + }); + } else { + // Rehydrated entry: re-resolve the web socket by the persisted id + // and send the response now. + const targetWebWs = this.findWebByConnectionId(webConnectionId ?? 'unknown'); + if (targetWebWs) { + this.sendToWeb(targetWebWs, { + type: 'response', + id: entry.originalId, + ...(result !== undefined ? { result } : {}), + ...(structuredError !== null + ? { error: structuredError } + : stringError !== null + ? { error: stringError } + : sanitizeAsFailed + ? { error: CLI_COMMAND_ERROR } + : {}), + }); + } + } + } finally { + if (ownsReservation) { + this.completedCorrelationIds.delete(id); + } + } } // --------------------------------------------------------------------------- @@ -906,7 +1286,14 @@ export class UserConnectionDO extends DurableObject { private handleWebCommand( ws: WebSocket, - msg: { id: string; command: string; sessionId?: string; connectionId?: string; data?: unknown } + msg: { + id: string; + command: string; + sessionId?: string; + connectionId?: string; + data?: unknown; + mutationId?: string; + } ): void { const now = Date.now(); this.expirePendingCommands(now); @@ -964,7 +1351,11 @@ export class UserConnectionDO extends DurableObject { } if (!targetCli) { - this.sendToWeb(ws, { type: 'response', id: msg.id, error: 'Session owner not found' }); + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: 'Session owner not found', + }); return; } @@ -974,7 +1365,16 @@ export class UserConnectionDO extends DurableObject { msg.sessionId && msg.connectionId ? msg.connectionId : undefined; const targetConnectionId = targetAttachment.connectionId; + // Resolve the originating web socket's connectionId from its attachment. + const webAttachment = ws.deserializeAttachment() as WSAttachment | null; + const webConnectionId = + webAttachment?.role === 'web' && webAttachment.connectionId + ? webAttachment.connectionId + : 'unknown'; + + // In-memory dedupe for catalog commands from the same web socket if ( + !msg.mutationId && CATALOG_DEDUPE_COMMANDS.has(msg.command) && [...this.pendingCommands.values()].some( entry => @@ -992,6 +1392,109 @@ export class UserConnectionDO extends DurableObject { return; } + // D8 mutationId dedupe: synchronous reservation prevents concurrent + // same-ID dispatch across the async storage read below. + if (msg.mutationId) { + const mutationId = msg.mutationId; + // Reject oversized mutation IDs before they become unsafe storage keys. + if (mutationId.length > MAX_MUTATION_ID_LENGTH) { + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: MUTATION_ID_TOO_LONG_ERROR, + }); + return; + } + // Atomic reservation: check and set synchronously before any await. + if (this.inflightMutations.has(mutationId)) { + if (CATALOG_DEDUPE_COMMANDS.has(msg.command)) { + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: CATALOG_REQUEST_PENDING_ERROR, + }); + } else { + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: { + source: 'relay', + code: 'COMMAND_ALREADY_PENDING', + message: 'Command is already in flight', + }, + }); + } + return; + } + this.inflightMutations.add(mutationId); + + const capturedTargetCli = targetCli; + this.ctx.waitUntil( + (async () => { + try { + const durable = await this.getDurablePendingCommand(mutationId); + if (!durable) { + // No existing entry: check cap (counts both in-memory and durable entries). + const durableCount = await this.countDurablePendingCommands(); + const total = this.pendingCommands.size + durableCount; + if (total >= UserConnectionDO.MAX_PENDING_COMMANDS) { + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: PENDING_COMMAND_LIMIT_ERROR, + }); + return; + } + await this.dispatchDurableWebCommand( + ws, + msg, + capturedTargetCli, + targetConnectionId, + expectedOwnerConnectionId, + webConnectionId, + now + ); + return; + } + + if (durable.state === 'done') { + await this.updateDurablePendingCommandOriginalId(mutationId, msg.id); + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + ...(durable.result !== undefined ? { result: durable.result } : {}), + ...(durable.error !== undefined ? { error: durable.error } : {}), + }); + return; + } + + // Entry is 'pending': dedupe. + if (CATALOG_DEDUPE_COMMANDS.has(msg.command)) { + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: CATALOG_REQUEST_PENDING_ERROR, + }); + } else { + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: { + source: 'relay', + code: 'COMMAND_ALREADY_PENDING', + message: 'Command is already in flight', + }, + }); + } + } finally { + this.inflightMutations.delete(mutationId); + } + })() + ); + return; + } + + // In-memory cap check for the common path. if (this.pendingCommands.size >= UserConnectionDO.MAX_PENDING_COMMANDS) { this.sendToWeb(ws, { type: 'response', @@ -1001,7 +1504,37 @@ export class UserConnectionDO extends DurableObject { return; } - const correlationId = crypto.randomUUID(); + this.dispatchWebCommandSync( + ws, + msg, + targetCli, + targetConnectionId, + expectedOwnerConnectionId, + webConnectionId, + now + ); + } + + /** + * Common dispatch: persist the durable entry before forwarding the command. + */ + private async dispatchDurableWebCommand( + ws: WebSocket, + msg: { + id: string; + command: string; + sessionId?: string; + connectionId?: string; + data?: unknown; + mutationId?: string; + }, + targetCli: WebSocket, + targetConnectionId: string, + expectedOwnerConnectionId: string | undefined, + webConnectionId: string, + now: number + ): Promise { + const correlationId = msg.mutationId ?? crypto.randomUUID(); this.pendingCommands.set(correlationId, { ws, sessionId: msg.sessionId, @@ -1012,7 +1545,18 @@ export class UserConnectionDO extends DurableObject { expiresAt: now + UserConnectionDO.PENDING_COMMAND_TTL_MS, targetCliWs: targetCli, }); + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${correlationId}`, { + sessionId: msg.sessionId, + originalId: msg.id, + command: msg.command, + expectedOwnerConnectionId, + targetConnectionId, + expiresAt: now + UserConnectionDO.PENDING_COMMAND_TTL_MS, + webConnectionId, + state: 'pending' as const, + } satisfies PendingCommandEntry); this.scheduleNextAlarm(now); + this.scheduleDurablePendingAlarm(); this.sendToCli(targetCli, { type: 'command', @@ -1020,7 +1564,94 @@ export class UserConnectionDO extends DurableObject { command: msg.command, data: msg.data, ...(msg.sessionId ? { sessionId: msg.sessionId } : {}), + ...(msg.mutationId ? { mutationId: msg.mutationId } : {}), + }); + } + + private dispatchWebCommandSync( + ws: WebSocket, + msg: { + id: string; + command: string; + sessionId?: string; + connectionId?: string; + data?: unknown; + mutationId?: string; + }, + targetCli: WebSocket, + targetConnectionId: string, + expectedOwnerConnectionId: string | undefined, + webConnectionId: string, + now: number + ): void { + const correlationId = crypto.randomUUID(); + const pendingCommandKey = `${PENDING_COMMAND_KEY_PREFIX}${correlationId}`; + this.pendingCommands.set(correlationId, { + ws, + sessionId: msg.sessionId, + originalId: msg.id, + command: msg.command, + expectedOwnerConnectionId, + targetConnectionId, + expiresAt: now + UserConnectionDO.PENDING_COMMAND_TTL_MS, + targetCliWs: targetCli, }); + this.pendingInitialCommandWrites.add(correlationId); + this.ctx.waitUntil( + this.ctx.storage + .put(pendingCommandKey, { + sessionId: msg.sessionId, + originalId: msg.id, + command: msg.command, + expectedOwnerConnectionId, + targetConnectionId, + expiresAt: now + UserConnectionDO.PENDING_COMMAND_TTL_MS, + webConnectionId, + state: 'pending' as const, + } satisfies PendingCommandEntry) + .then(async () => { + const terminalEntry = this.terminalDuringInitialWrite.get(correlationId); + if (terminalEntry) { + await this.ctx.storage.put(pendingCommandKey, terminalEntry); + this.completedCorrelationIds.delete(correlationId); + this.terminalDuringInitialWrite.delete(correlationId); + return; + } + this.scheduleNextAlarm(now); + this.scheduleDurablePendingAlarm(); + this.sendToCli(targetCli, { + type: 'command', + id: correlationId, + command: msg.command, + data: msg.data, + ...(msg.sessionId ? { sessionId: msg.sessionId } : {}), + }); + }) + .catch((error: unknown) => { + // The initial durable write failed: the command was never forwarded to + // the CLI. Clean any terminal stash and its reservation marker so they + // do not leak until DO eviction. + this.terminalDuringInitialWrite.delete(correlationId); + this.completedCorrelationIds.delete(correlationId); + console.error('Failed to persist durable pending command before forwarding', { + correlationId, + error: error instanceof Error ? error.message : String(error), + }); + }) + .finally(() => { + this.pendingInitialCommandWrites.delete(correlationId); + if (!this.terminalDuringInitialWrite.has(correlationId)) { + // Only clear the reservation marker when the entry is still tracked + // in-memory. If expirePendingCommands already removed it and set + // the marker, the async sweep owns cleanup and we must not clear it + // here (otherwise a concurrent finishDurablePendingCommands could + // double-deliver). + if (this.pendingCommands.has(correlationId)) { + this.completedCorrelationIds.delete(correlationId); + } + } + }) + ); } // --------------------------------------------------------------------------- @@ -1044,11 +1675,15 @@ export class UserConnectionDO extends DurableObject { return att?.role === 'cli' && att.connectionId === connectionId; }); - // Fail pending commands that targeted this specific socket - this.failPendingCommandsForSocket(disconnectedWs, !replaced); + // Fail pending commands that targeted this specific socket. + // Await so the durable terminal entries are persisted before we proceed + // to broadcast cli.disconnected. Failed first writes retry in waitUntil. + await this.failPendingCommandsForSocket(disconnectedWs, !replaced); if (replaced) { - console.log('Stale CLI socket closed (already replaced)', { connectionId }); + console.log('Stale CLI socket closed (already replaced)', { + connectionId, + }); return; } @@ -1142,7 +1777,10 @@ export class UserConnectionDO extends DurableObject { } } - // Clean up any pending commands from this web socket + // Clean up any pending commands from this web socket. + // Step 26b: keep the durable entry — the CLI may still reply, and a + // mutationId retry needs the stored outcome. Only the in-memory entry is + // dropped; the durable entry lives until its TTL or a terminal outcome. let droppedCommands = 0; for (const [id, entry] of this.pendingCommands) { if (entry.ws === ws) { @@ -1296,7 +1934,14 @@ export class UserConnectionDO extends DurableObject { const att = ws.deserializeAttachment() as WSAttachment | null; if (att?.role === 'cli' && att.connectionId === connectionId) { console.log('Closing stale CLI socket for reconnect', { connectionId }); - this.failPendingCommandsForSocket(ws, false); + this.ctx.waitUntil( + this.failPendingCommandsForSocket(ws, false).catch((error: unknown) => { + console.error('Failed to persist terminal commands for stale socket', { + connectionId, + error: error instanceof Error ? error.message : String(error), + }); + }) + ); // Preserve session ownership — the reconnecting CLI still owns these sessions ws.close(1000, 'replaced by reconnect'); return true; @@ -1345,12 +1990,104 @@ export class UserConnectionDO extends DurableObject { return undefined; } - private failPendingCommandsForSocket(targetWs: WebSocket, cliGone: boolean): void { - for (const [id, entry] of this.pendingCommands) { - if (entry.targetCliWs !== targetWs) continue; + private findWebByConnectionId(connectionId: string): WebSocket | undefined { + for (const ws of this.ctx.getWebSockets('web')) { + const attachment = ws.deserializeAttachment() as WSAttachment | null; + if ( + attachment?.role === 'web' && + attachment.connectionId === connectionId && + !attachment.replaced + ) { + return ws; + } + } + return undefined; + } + + private async failPendingCommandsForSocket(targetWs: WebSocket, cliGone: boolean): Promise { + // Collect correlationIds handled in the in-memory sweep so the durable + // sweep in finishDurablePendingCommands does not double-deliver. + const handledIds = new Set(); + const entries = [...this.pendingCommands].filter(([, entry]) => entry.targetCliWs === targetWs); + for (const [id] of entries) { + handledIds.add(id); + this.pendingCommands.delete(id); + this.completedCorrelationIds.add(id); + } + for (const [id, entry] of entries) { // The owning CLI is really gone, so a forwarded `exit_cli` got what it // asked for: this session is no longer owned by anyone. const exited = cliGone && entry.command === 'exit_cli'; + const isOwnerFenced = !exited && Boolean(entry.expectedOwnerConnectionId); + // Live wire error: bare string for CLI-disconnect compatibility + // (the client's parseCommandError expects a string, not structured). + const liveError = exited + ? undefined + : isOwnerFenced + ? SESSION_OWNER_CHANGED_ERROR + : 'CLI disconnected'; + const durableError = liveError; + + // Step 26: pair the in-memory delete with a durable done transition. + // Persist the terminal outcome BEFORE sending the live response. + // If storage.put fails, no live send occurs. A waitUntil retry keeps + // the terminal outcome fenced while the loop continues. + const webAtt = entry.ws.deserializeAttachment() as WSAttachment | null; + const durableEntry: PendingCommandEntry = { + sessionId: entry.sessionId, + originalId: entry.originalId, + command: entry.command, + expectedOwnerConnectionId: entry.expectedOwnerConnectionId, + targetConnectionId: entry.targetConnectionId, + expiresAt: entry.expiresAt, + webConnectionId: webAtt?.role === 'web' ? webAtt.connectionId : 'unknown', + state: 'done', + ...(exited ? { result: {} } : { error: durableError }), + }; + const initialWritePending = this.pendingInitialCommandWrites.has(id); + if (initialWritePending) { + this.terminalDuringInitialWrite.set(id, durableEntry); + } + try { + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, durableEntry); + } catch (error) { + this.ctx.waitUntil( + this.ctx.storage + .put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, durableEntry) + .then(() => { + if (!initialWritePending) { + this.completedCorrelationIds.delete(id); + } + const webWs = this.findWebByConnectionId(durableEntry.webConnectionId); + if (webWs) { + this.sendToWeb( + webWs, + exited + ? { type: 'response', id: durableEntry.originalId, result: {} } + : { + type: 'response', + id: durableEntry.originalId, + error: durableError, + } + ); + } + }) + .catch(retryError => { + console.error('Failed to retry durable CLI-disconnect terminal outcome', { + correlationId: id, + error: retryError instanceof Error ? retryError.message : String(retryError), + }); + }) + ); + console.error('Failed to persist durable CLI-disconnect terminal outcome', { + correlationId: id, + error: error instanceof Error ? error.message : String(error), + }); + continue; + } + if (!initialWritePending) { + this.completedCorrelationIds.delete(id); + } this.sendToWeb( entry.ws, exited @@ -1358,46 +2095,412 @@ export class UserConnectionDO extends DurableObject { : { type: 'response', id: entry.originalId, - error: entry.expectedOwnerConnectionId - ? SESSION_OWNER_CHANGED_ERROR - : 'CLI disconnected', + error: liveError, } ); - this.pendingCommands.delete(id); + } + + const attachment = targetWs.deserializeAttachment() as WSAttachment | null; + if (cliGone && attachment?.role === 'cli') { + this.ctx.waitUntil( + this.finishDurablePendingCommands( + entry => entry.targetConnectionId === attachment.connectionId, + 'CLI disconnected', + handledIds + ) + ); } } - private failPendingCommandsForOwnerChange( + private async failPendingCommandsForOwnerChange( sessionId: string, nextOwnerConnectionId: string | undefined - ): void { - for (const [id, entry] of this.pendingCommands) { - if (entry.sessionId !== sessionId || entry.targetConnectionId === nextOwnerConnectionId) { - continue; - } + ): Promise { + const handledIds = new Set(); + const entries = [...this.pendingCommands].filter( + ([, entry]) => + entry.sessionId === sessionId && entry.targetConnectionId !== nextOwnerConnectionId + ); + for (const [id] of entries) { + handledIds.add(id); this.pendingCommands.delete(id); + this.completedCorrelationIds.add(id); + } + for (const [id, entry] of entries) { // `exit_cli` asked for exactly this: the session is no longer owned. // Ownership moving to another CLI is a genuine owner change and still fails. const exited = entry.command === 'exit_cli' && nextOwnerConnectionId === undefined; + + // Step 26: pair the in-memory delete with a durable done transition. + // Persist the terminal outcome BEFORE sending the live response. + // If storage.put fails, no live send occurs. A waitUntil retry keeps + // the terminal outcome fenced while the loop continues. + const webAtt = entry.ws.deserializeAttachment() as WSAttachment | null; + const durableEntry: PendingCommandEntry = { + sessionId: entry.sessionId, + originalId: entry.originalId, + command: entry.command, + expectedOwnerConnectionId: entry.expectedOwnerConnectionId, + targetConnectionId: entry.targetConnectionId, + expiresAt: entry.expiresAt, + webConnectionId: webAtt?.role === 'web' ? webAtt.connectionId : 'unknown', + state: 'done', + ...(exited ? { result: {} } : { error: SESSION_OWNER_CHANGED_ERROR }), + }; + const initialWritePending = this.pendingInitialCommandWrites.has(id); + if (initialWritePending) { + this.terminalDuringInitialWrite.set(id, durableEntry); + } + try { + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, durableEntry); + } catch (error) { + this.ctx.waitUntil( + this.ctx.storage + .put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, durableEntry) + .then(() => { + if (!initialWritePending) { + this.completedCorrelationIds.delete(id); + } + const webWs = this.findWebByConnectionId(durableEntry.webConnectionId); + if (webWs) { + this.sendToWeb( + webWs, + exited + ? { type: 'response', id: durableEntry.originalId, result: {} } + : { + type: 'response', + id: durableEntry.originalId, + error: SESSION_OWNER_CHANGED_ERROR, + } + ); + } + }) + .catch(retryError => { + console.error('Failed to retry durable owner-change terminal outcome', { + correlationId: id, + error: retryError instanceof Error ? retryError.message : String(retryError), + }); + }) + ); + console.error('Failed to persist durable owner-change terminal outcome', { + correlationId: id, + error: error instanceof Error ? error.message : String(error), + }); + continue; + } + if (!initialWritePending) { + this.completedCorrelationIds.delete(id); + } this.sendToWeb( entry.ws, exited ? { type: 'response', id: entry.originalId, result: {} } - : { type: 'response', id: entry.originalId, error: SESSION_OWNER_CHANGED_ERROR } + : { + type: 'response', + id: entry.originalId, + error: SESSION_OWNER_CHANGED_ERROR, + } ); } + + this.ctx.waitUntil( + this.finishDurablePendingCommands( + entry => + entry.sessionId === sessionId && entry.targetConnectionId !== nextOwnerConnectionId, + SESSION_OWNER_CHANGED_ERROR, + handledIds + ) + ); + } + + // --------------------------------------------------------------------------- + // Command error shaping + // --------------------------------------------------------------------------- + + /** + * Compute the shaped terminal outcome for a command reply. + * Applies catalog size guard and error normalization (CLI_UPGRADE_REQUIRED + * mapping, CLI_COMMAND_ERROR sanitization). Pure computation — no side effects. + */ + private shapeCommandError( + command: string, + result: unknown, + error: unknown + ): { result?: unknown; error?: unknown } { + // Catalog size guard: oversized results must not persist their raw payload. + if (CATALOG_DEDUPE_COMMANDS.has(command) && result !== undefined) { + const serializedResult = JSON.stringify(result); + const resultBytes = new TextEncoder().encode(serializedResult).byteLength; + if (resultBytes > MAX_CATALOG_RESULT_BYTES) { + return { error: CATALOG_TOO_LARGE_ERROR }; + } + } + + let structuredError: { + source: string; + code: string; + message: string; + } | null = null; + let stringError: string | null = null; + let sanitizeAsFailed = false; + + if (typeof error === 'string' && CLI_UPGRADE_REQUIRED_COMMANDS.has(command)) { + if (error === `unknown command: ${command}`) { + structuredError = + command === 'create_session' + ? CLI_UPGRADE_REQUIRED_CREATE_SESSION_ERROR + : CLI_UPGRADE_REQUIRED_SLASH_ERROR; + } else { + stringError = error; + } + } else if (typeof error === 'string') { + stringError = error; + } else if (error !== undefined) { + sanitizeAsFailed = true; + } + + const terminalError = + structuredError ?? stringError ?? (sanitizeAsFailed ? CLI_COMMAND_ERROR : undefined); + + return { + ...(result !== undefined ? { result } : {}), + ...(terminalError !== undefined ? { error: terminalError } : {}), + }; + } + + /** + * Bound a terminal result so the durable storage.put never exceeds the + * Cloudflare DO value limit. Returns the original result when under the + * threshold, or a size-exceeded marker otherwise. + */ + private boundDurableResult(result: unknown): unknown { + if (result === undefined) return undefined; + const serialized = JSON.stringify(result); + const byteCount = new TextEncoder().encode(serialized).byteLength; + if (byteCount <= MAX_DURABLE_RESULT_BYTES) return result; + return { _truncated: true, bytes: byteCount }; + } + + /** + * Detect the truncated marker returned by boundDurableResult. + * Callers must use this to route a durable-sized result to an error + * field so a mutationId retry does not return a false success. + */ + private isTruncatedResult(value: unknown): boolean { + return ( + value !== undefined && + value !== null && + typeof value === 'object' && + '_truncated' in value && + (value as Record)._truncated === true + ); + } + + // --------------------------------------------------------------------------- + // Durable pending command storage (D8) + // --------------------------------------------------------------------------- + + private async getDurablePendingCommand( + correlationId: string + ): Promise { + return this.ctx.storage.get( + `${PENDING_COMMAND_KEY_PREFIX}${correlationId}` + ); + } + + private async countDurablePendingCommands(): Promise { + const entries = await this.ctx.storage.list({ + prefix: PENDING_COMMAND_KEY_PREFIX, + }); + return entries.size; + } + + private async updateDurablePendingCommandOriginalId( + correlationId: string, + newOriginalId: string + ): Promise { + const entry = await this.getDurablePendingCommand(correlationId); + if (!entry) return; + entry.originalId = newOriginalId; + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${correlationId}`, entry); + } + + private async finishDurablePendingCommands( + matches: (entry: PendingCommandEntry) => boolean, + error: unknown, + skipIds?: ReadonlySet + ): Promise { + const entries = await this.ctx.storage.list({ + prefix: PENDING_COMMAND_KEY_PREFIX, + }); + await Promise.all( + [...entries].flatMap(([key, entry]) => { + if (entry.state !== 'pending' || !matches(entry)) return []; + const correlationId = key.slice(PENDING_COMMAND_KEY_PREFIX.length); + // Skip entries already delivered live by the in-memory sweep. + // Prevents a duplicate delivery when the in-memory storage.put and + // finishDurablePendingCommands race within separate waitUntil tasks. + if (skipIds?.has(correlationId)) return []; + // Claim the correlation BEFORE the durable put so a concurrent + // sweep with a stale list snapshot cannot pass through its + // awaited put and send after we already delivered. + if (this.completedCorrelationIds.has(correlationId)) return []; + this.completedCorrelationIds.add(correlationId); + // After the durable write succeeds, deliver to a reattached web + // socket (D8 case 3). The in-memory sweep already delivered to the + // original socket; this path covers a socket that reconnected after + // the original disconnected (different webConnectionId path). + const webWs = this.findWebByConnectionId(entry.webConnectionId); + const putPromise = this.ctx.storage + .put(key, { + ...entry, + state: 'done' as const, + error, + }) + .then(() => { + if (webWs) { + this.sendToWeb(webWs, { + type: 'response', + id: entry.originalId, + error, + }); + } + }) + .catch((_error: unknown) => { + // On storage failure: release the claim so a retry or another + // concurrent path can deliver. Log but do not rethrow — + // Promise.all already rejects on unhandled rejections; the + // caller's waitUntil catch logs upstream. + this.completedCorrelationIds.delete(correlationId); + console.error('Failed to persist finishDurablePendingCommands terminal outcome', { + correlationId, + error: _error instanceof Error ? _error.message : String(_error), + }); + }); + return [putPromise]; + }) + ); + } + + private scheduleDurablePendingAlarm(): void { + this.ctx.waitUntil( + (async () => { + const entries = await this.ctx.storage.list({ + prefix: PENDING_COMMAND_KEY_PREFIX, + }); + const now = Date.now(); + let expiresAt: number | undefined; + for (const entry of entries.values()) { + if (entry.expiresAt <= now) continue; + if (expiresAt === undefined || entry.expiresAt < expiresAt) { + expiresAt = entry.expiresAt; + } + } + if (expiresAt === undefined) return; + const currentAlarm = await this.ctx.storage.getAlarm(); + if (currentAlarm === null || expiresAt < currentAlarm) { + await this.ctx.storage.setAlarm(expiresAt); + } + })() + ); } private expirePendingCommands(now: number): void { + // Track correlationIds already delivered in-memory so the durable + // sweep does not double-deliver to the same web socket. + const deliveredInMemory = new Set(); for (const [id, entry] of this.pendingCommands) { if (entry.expiresAt > now) continue; this.pendingCommands.delete(id); + // Reserve the correlationId so a concurrent finishDurablePendingCommands + // (from failPendingCommandsForSocket/failPendingCommandsForOwnerChange) + // does not double-deliver the same terminal response. But if a + // concurrent path already reserved it, skip — that path owns delivery. + if (this.completedCorrelationIds.has(id)) { + deliveredInMemory.add(id); + continue; + } + deliveredInMemory.add(id); + this.completedCorrelationIds.add(id); this.sendToWeb(entry.ws, { type: 'response', id: entry.originalId, error: COMMAND_EXPIRED_ERROR, }); } + + // Step 26: sweep durable entries. Pending entries past their TTL are + // marked done with COMMAND_EXPIRED_ERROR; already-done entries are + // cleaned up via deletion. + this.ctx.waitUntil( + (async () => { + const entries = await this.ctx.storage.list({ + prefix: PENDING_COMMAND_KEY_PREFIX, + }); + for (const [key, durable] of entries) { + if (!durable || typeof durable.expiresAt !== 'number') continue; + if (durable.expiresAt > now) continue; + + if (durable.state === 'pending') { + const correlationId = key.slice(PENDING_COMMAND_KEY_PREFIX.length); + // Entries already delivered in the in-memory sweep above still need + // the durable 'done' mark, but should not claim or send again. + if (deliveredInMemory.has(correlationId)) { + await this.ctx.storage.put(key, { + ...durable, + state: 'done', + error: COMMAND_EXPIRED_ERROR, + }); + continue; + } + // Claim the correlation BEFORE the durable put so a concurrent + // sweep with a stale list snapshot cannot pass through its + // awaited put and send after we already delivered. + if (this.completedCorrelationIds.has(correlationId)) continue; + this.completedCorrelationIds.add(correlationId); + + try { + await this.ctx.storage.put(key, { + ...durable, + state: 'done', + error: COMMAND_EXPIRED_ERROR, + }); + } catch (_error: unknown) { + // On storage failure: release the claim so a retry or another + // concurrent path can deliver. + this.completedCorrelationIds.delete(correlationId); + console.error('Failed to persist expirePendingCommands terminal outcome', { + correlationId, + error: _error instanceof Error ? _error.message : String(_error), + }); + continue; + } + + // Deliver the terminal response to a live matching web socket + // that reconnected after the original socket disconnected (D8 case 3). + const webWs = this.findWebByConnectionId(durable.webConnectionId); + if (webWs) { + this.sendToWeb(webWs, { + type: 'response', + id: durable.originalId, + error: COMMAND_EXPIRED_ERROR, + }); + } + } else { + // Already done and past TTL: clean up. Release the completion + // marker so cross-sweep dedupe state does not accumulate through + // repeated TTL cleanup. + const correlationId = key.slice(PENDING_COMMAND_KEY_PREFIX.length); + await this.ctx.storage.delete(key); + this.completedCorrelationIds.delete(correlationId); + } + } + })().catch((error: unknown) => { + console.error('Failed to sweep durable pending commands', { + error: error instanceof Error ? error.message : String(error), + }); + }) + ); } private scheduleNextAlarm(now: number): void { diff --git a/services/session-ingest/src/types/user-connection-protocol.ts b/services/session-ingest/src/types/user-connection-protocol.ts index a7156b2e132..51aaa6a396b 100644 --- a/services/session-ingest/src/types/user-connection-protocol.ts +++ b/services/session-ingest/src/types/user-connection-protocol.ts @@ -79,6 +79,12 @@ export const CLIInboundMessageSchema = z.discriminatedUnion('type', [ command: z.string(), sessionId: z.string().optional(), data: z.unknown().optional(), + // Stable intent id forwarded from the web wire. The CLI does not + // interpret it; it is echoed back on the response for the DO's + // durable dedupe to match against the stored entry. Absent when the + // web client did not supply one. + // Bounded at 128 chars to keep storage keys safe. + mutationId: z.string().max(128).optional(), }), z.object({ type: z.literal('system'), @@ -108,6 +114,12 @@ export const WebOutboundMessageSchema = z.discriminatedUnion('type', [ connectionId: z.string().optional(), command: z.string(), data: z.unknown().optional(), + // Stable client-side intent id. When present, the DO uses it as the + // correlation id so a retry with the same mutationId runs the command + // once for the full TTL window (D8 dedupe). Absent on older clients; + // the DO falls back to crypto.randomUUID() per D5. + // Bounded at 128 chars to keep storage keys safe. + mutationId: z.string().max(128).optional(), }), z.object({ type: z.literal('ping'),