From 2d422d132681c1fff8f348bb9667150edbab2089 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Mon, 3 Aug 2026 21:22:18 +0200 Subject: [PATCH 01/24] fix(stream): close snapshot replay gap --- .../mobile-session-page-adapter.test.ts | 28 ++ .../agents/mobile-session-page-adapter.ts | 5 +- .../cloud-agent-next/cloud-agent-client.ts | 42 ++- .../routers/cli-sessions-v2-router.test.ts | 326 ++++++++++++++++-- .../web/src/routers/cli-sessions-v2-router.ts | 45 ++- .../src/cloud-agent-transport.test.ts | 200 +++++++++++ .../src/cloud-agent-transport.ts | 18 +- packages/cloud-agent-sdk/src/types.ts | 7 + .../src/persistence/CloudAgentSession.ts | 11 + services/cloud-agent-next/src/router.test.ts | 121 +++++++ .../src/router/handlers/session-management.ts | 20 ++ .../cloud-agent-next/src/router/schemas.ts | 11 + 12 files changed, 790 insertions(+), 44 deletions(-) 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/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/routers/cli-sessions-v2-router.test.ts b/apps/web/src/routers/cli-sessions-v2-router.test.ts index 7f581112b00..b531b0a37ca 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, @@ -1893,7 +2162,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]); @@ -1913,7 +2185,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 33da8d8f14d..ddb6fbf72dc 100644 --- a/apps/web/src/routers/cli-sessions-v2-router.ts +++ b/apps/web/src/routers/cli-sessions-v2-router.ts @@ -686,7 +686,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)} @@ -802,7 +805,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 { @@ -839,7 +863,7 @@ export const cliSessionsV2Router = createTRPCRouter({ }); } - return result; + return { ...result, watermarkEventId }; }), /** @@ -1047,7 +1071,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; @@ -1190,7 +1217,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 @@ -1342,7 +1372,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/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts index b0983f8f5f2..1afe555e735 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts @@ -919,4 +919,204 @@ describe('CloudAgentTransport page-seam', () => { jest.useRealTimers(); } }); + + describe('watermark cursor', () => { + it('uses fromId on first connect when the page carries a watermark', 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]); + expect(wsUrl).toContain('fromId=42'); + 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 the watermark + const firstUrl = String(webSocketConstructor.mock.calls[0]?.[0]); + expect(firstUrl).toContain('fromId=10'); + + // 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 the original watermark and skips the page refetch', 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 the watermark as fromId. + const firstUrl = String(webSocketConstructor.mock.calls[0]?.[0]); + expect(firstUrl).toContain('fromId=42'); + expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); + + // Send only sentinel events (eventId: 0) — they do NOT advance the + // live cursor, so lastEventId stays at the watermark value. + 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 must still use the watermark (42) because sentinel events + // never advance the cursor past the watermark. + expect(reconnectUrl).toContain('fromId=42'); + expect(reconnectUrl).not.toContain('replay=false'); + + // No second page fetch — lastEventId (watermark 42) is not null, so + // onReconnected skips the snapshot fallback path entirely. + expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); + + 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..95e17a55ae8 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,12 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport return null; } if (page.kind === 'success') { + // Seed lastEventId from the page's event-log watermark so the + // first WebSocket connect uses `fromId`. The DO replays any + // events persisted between the page snapshot and the socket + // open, closing the gap. On reconnect, wire events advance + // lastEventId and the live cursor takes over. + lastEventId = page.watermarkEventId ?? null; config.onInitialPageLoaded?.(page); replayPage(page); return { ticket }; 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/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; From c5542610f18a14a9213ce590fb51cc4d2d533d8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Mon, 3 Aug 2026 21:42:37 +0200 Subject: [PATCH 02/24] fix(stream): dedupe replay events --- .../src/cloud-agent-transport.test.ts | 1010 ++++++++++++----- .../src/cloud-agent-transport.ts | 7 + 2 files changed, 732 insertions(+), 285 deletions(-) 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 1afe555e735..644eba045c8 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts @@ -2,12 +2,12 @@ * Tests for CloudAgentTransport — verifies event normalization, routing to * chat/service sinks, and lifecycle generation tracking. */ -import type { CloudAgentEvent } from './event-types'; -import { createEventHelpers } from './__fixtures__/helpers'; -import type { ChatEvent, ServiceEvent } from './normalizer'; -import { createCloudAgentTransport } from './cloud-agent-transport'; -import type { SessionSnapshotPageOutcome } from './types'; -import { kiloId, cloudAgentId, makeSnapshot } from './test-helpers'; +import type { CloudAgentEvent } from "./event-types"; +import { createEventHelpers } from "./__fixtures__/helpers"; +import type { ChatEvent, ServiceEvent } from "./normalizer"; +import { createCloudAgentTransport } from "./cloud-agent-transport"; +import type { SessionSnapshotPageOutcome } from "./types"; +import { kiloId, cloudAgentId, makeSnapshot } from "./test-helpers"; // --------------------------------------------------------------------------- // WebSocket mock @@ -52,10 +52,10 @@ afterEach(() => { /** Flush microtask queue so Promise.all + .then in connect() settles. */ async function flushPromises(): Promise { - await new Promise(r => setTimeout(r, 0)); + await new Promise((r) => setTimeout(r, 0)); } -const emptySnapshot = makeSnapshot({ id: 'ses-1' }); +const emptySnapshot = makeSnapshot({ id: "ses-1" }); function sendRaw(event: CloudAgentEvent): void { mockWs.onmessage?.({ data: JSON.stringify(event) } as MessageEvent); @@ -63,35 +63,36 @@ function sendRaw(event: CloudAgentEvent): void { function createMockApi() { return { - send: jest.fn(() => Promise.resolve('sent')), - interrupt: jest.fn(() => Promise.resolve('interrupted')), - answer: jest.fn(() => Promise.resolve('answered')), - reject: jest.fn(() => Promise.resolve('rejected')), - respondToPermission: jest.fn(() => Promise.resolve('responded')), + send: jest.fn(() => Promise.resolve("sent")), + interrupt: jest.fn(() => Promise.resolve("interrupted")), + answer: jest.fn(() => Promise.resolve("answered")), + reject: jest.fn(() => Promise.resolve("rejected")), + respondToPermission: jest.fn(() => Promise.resolve("responded")), }; } function createTransportWithSinks( - getTicket: (sessionId: string) => string | Promise = () => 'test-ticket', + getTicket: (sessionId: string) => string | Promise = () => + "test-ticket", onError?: (message: string) => void, - api = createMockApi() + api = createMockApi(), ) { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; const factory = createCloudAgentTransport({ - sessionId: cloudAgentId('ses-1'), - kiloSessionId: kiloId('ses-1'), + sessionId: cloudAgentId("ses-1"), + kiloSessionId: kiloId("ses-1"), api, getTicket, fetchSnapshot: () => Promise.resolve(emptySnapshot), - websocketBaseUrl: 'ws://localhost:9999', + websocketBaseUrl: "ws://localhost:9999", onError, }); const transport = factory({ - onChatEvent: event => chatEvents.push(event), - onServiceEvent: event => serviceEvents.push(event), + onChatEvent: (event) => chatEvents.push(event), + onServiceEvent: (event) => serviceEvents.push(event), }); return { transport, chatEvents, serviceEvents, api }; @@ -107,51 +108,69 @@ beforeEach(() => { // Tests // --------------------------------------------------------------------------- -describe('CloudAgentTransport event routing', () => { - it('routes chat events to onChatEvent', async () => { +describe("CloudAgentTransport event routing", () => { + it("routes chat events to onChatEvent", async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); await flushPromises(); sendRaw( - kilocode('message.updated', { - info: { id: 'msg-1', sessionID: 'ses-1', role: 'assistant', time: { created: 1 } }, - }) + kilocode("message.updated", { + info: { + id: "msg-1", + sessionID: "ses-1", + role: "assistant", + time: { created: 1 }, + }, + }), ); expect(chatEvents).toHaveLength(1); - expect(chatEvents[0]).toEqual(expect.objectContaining({ type: 'message.updated' })); + expect(chatEvents[0]).toEqual( + expect.objectContaining({ type: "message.updated" }), + ); // session.created from snapshot replay expect(serviceEvents).toHaveLength(1); - expect(serviceEvents[0]).toEqual(expect.objectContaining({ type: 'session.created' })); + expect(serviceEvents[0]).toEqual( + expect.objectContaining({ type: "session.created" }), + ); transport.destroy(); }); - it('routes service events to onServiceEvent', async () => { + it("routes service events to onServiceEvent", async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); 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); - expect(serviceEvents[0]).toEqual(expect.objectContaining({ type: 'session.created' })); - expect(serviceEvents[1]).toEqual(expect.objectContaining({ type: 'session.status' })); + expect(serviceEvents[0]).toEqual( + expect.objectContaining({ type: "session.created" }), + ); + expect(serviceEvents[1]).toEqual( + expect.objectContaining({ type: "session.status" }), + ); expect(chatEvents).toHaveLength(0); transport.destroy(); }); - it('routes cached command catalogs emitted without an execution ID', async () => { + it("routes cached command catalogs emitted without an execution ID", async () => { const { transport, serviceEvents } = createTransportWithSinks(); const commands = [ { - name: 'deploy-prod', - description: 'Deploy production', - hints: ['$ARGUMENTS'], - source: 'command', + name: "deploy-prod", + description: "Deploy production", + hints: ["$ARGUMENTS"], + source: "command", }, ]; @@ -160,18 +179,21 @@ describe('CloudAgentTransport event routing', () => { sendRaw({ eventId: 0, executionId: null, - sessionId: 'ses-1', - streamEventType: 'commands.available', + sessionId: "ses-1", + streamEventType: "commands.available", timestamp: new Date().toISOString(), data: { commands }, }); - expect(serviceEvents).toContainEqual({ type: 'commands.available', commands }); + expect(serviceEvents).toContainEqual({ + type: "commands.available", + commands, + }); transport.destroy(); }); - it('routes mixed events to correct sinks', async () => { + it("routes mixed events to correct sinks", async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); @@ -179,42 +201,60 @@ describe('CloudAgentTransport event routing', () => { // Chat event sendRaw( - kilocode('message.updated', { - info: { id: 'msg-1', sessionID: 'ses-1', role: 'assistant', time: { created: 1 } }, - }) + kilocode("message.updated", { + 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( - kilocode('message.part.delta', { - sessionID: 'ses-1', - messageID: 'msg-1', - partID: 'part-1', - field: 'text', - delta: 'hello', - }) + kilocode("message.part.delta", { + sessionID: "ses-1", + messageID: "msg-1", + partID: "part-1", + field: "text", + delta: "hello", + }), ); expect(chatEvents).toHaveLength(2); - expect(chatEvents[0]).toEqual(expect.objectContaining({ type: 'message.updated' })); - expect(chatEvents[1]).toEqual(expect.objectContaining({ type: 'message.part.delta' })); + expect(chatEvents[0]).toEqual( + expect.objectContaining({ type: "message.updated" }), + ); + expect(chatEvents[1]).toEqual( + expect.objectContaining({ type: "message.part.delta" }), + ); // session.created from snapshot replay + session.status from sendRaw expect(serviceEvents).toHaveLength(2); - expect(serviceEvents[0]).toEqual(expect.objectContaining({ type: 'session.created' })); - expect(serviceEvents[1]).toEqual(expect.objectContaining({ type: 'session.status' })); + expect(serviceEvents[0]).toEqual( + expect.objectContaining({ type: "session.created" }), + ); + expect(serviceEvents[1]).toEqual( + expect.objectContaining({ type: "session.status" }), + ); transport.destroy(); }); - it('ignores invalid events', () => { + it("ignores invalid events", () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); - mockWs.onmessage?.({ data: 'not json at all' } as MessageEvent); + mockWs.onmessage?.({ data: "not json at all" } as MessageEvent); expect(chatEvents).toHaveLength(0); expect(serviceEvents).toHaveLength(0); @@ -222,7 +262,7 @@ describe('CloudAgentTransport event routing', () => { transport.destroy(); }); - it('drops suggestion events since cloud-agent has no accept/dismiss command path', async () => { + it("drops suggestion events since cloud-agent has no accept/dismiss command path", async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); @@ -231,14 +271,14 @@ describe('CloudAgentTransport event routing', () => { const serviceCountBefore = serviceEvents.length; sendRaw( - kilocode('suggestion.shown', { - id: 'sug-1', - text: 'review your changes', - actions: [{ label: 'review', prompt: '/review' }], - }) + kilocode("suggestion.shown", { + id: "sug-1", + text: "review your changes", + actions: [{ label: "review", prompt: "/review" }], + }), ); - sendRaw(kilocode('suggestion.accepted', { requestID: 'sug-1', index: 0 })); - sendRaw(kilocode('suggestion.dismissed', { requestID: 'sug-1' })); + sendRaw(kilocode("suggestion.accepted", { requestID: "sug-1", index: 0 })); + sendRaw(kilocode("suggestion.dismissed", { requestID: "sug-1" })); expect(serviceEvents).toHaveLength(serviceCountBefore); expect(chatEvents).toHaveLength(0); @@ -247,59 +287,86 @@ describe('CloudAgentTransport event routing', () => { }); }); -describe('CloudAgentTransport unexpected disconnect', () => { - it('synthesizes stopped event on unexpected disconnect', async () => { +describe("CloudAgentTransport unexpected disconnect", () => { + it("synthesizes stopped event on unexpected disconnect", async () => { const { transport, serviceEvents } = createTransportWithSinks(); 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'); + 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(); }); - it('suppresses synthetic stopped if already received via event pipeline', () => { + it("suppresses synthetic stopped if already received via event pipeline", () => { 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' })); + sendRaw(createEvent("complete", { currentBranch: "main" })); - const stoppedBefore = serviceEvents.filter(e => e.type === 'stopped').length; + 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); - - const stoppedAfter = serviceEvents.filter(e => e.type === 'stopped').length; + mockWs.onclose?.({ + code: 1011, + reason: "network dropped", + wasClean: false, + } as CloseEvent); + + const stoppedAfter = serviceEvents.filter( + (e) => e.type === "stopped", + ).length; expect(stoppedAfter).toBe(stoppedBefore); transport.destroy(); }); }); -describe('CloudAgentTransport ticket handling', () => { - it('calls getTicket with sessionId', () => { - const getTicket = jest.fn((_sessionId: string) => 'ticket-abc'); +describe("CloudAgentTransport ticket handling", () => { + it("calls getTicket with sessionId", () => { + const getTicket = jest.fn((_sessionId: string) => "ticket-abc"); const { transport } = createTransportWithSinks(getTicket); transport.connect(); - expect(getTicket).toHaveBeenCalledWith('ses-1'); + expect(getTicket).toHaveBeenCalledWith("ses-1"); transport.destroy(); }); - it('handles async getTicket', async () => { - const getTicket = jest.fn((_sessionId: string) => Promise.resolve('async-ticket')); + it("handles async getTicket", async () => { + const getTicket = jest.fn((_sessionId: string) => + Promise.resolve("async-ticket"), + ); const { transport, serviceEvents } = createTransportWithSinks(getTicket); transport.connect(); @@ -309,22 +376,27 @@ 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(); }); - it('refreshes an expiring ticket before opening the websocket', async () => { + it("refreshes an expiring ticket before opening the websocket", async () => { const nowSeconds = Math.floor(Date.now() / 1000); const getTicket = jest .fn() .mockResolvedValueOnce({ - ticket: 'expiring-ticket', + ticket: "expiring-ticket", expiresAt: nowSeconds + 5, }) .mockResolvedValueOnce({ - ticket: 'fresh-ticket', + ticket: "fresh-ticket", expiresAt: nowSeconds + 60, }); const { transport } = createTransportWithSinks(getTicket); @@ -336,14 +408,16 @@ describe('CloudAgentTransport ticket handling', () => { expect(getTicket).toHaveBeenCalledTimes(2); expect(webSocketConstructor).toHaveBeenCalledTimes(1); - expect(webSocketConstructor.mock.calls[0]?.[0]).toContain('ticket=fresh-ticket'); + expect(webSocketConstructor.mock.calls[0]?.[0]).toContain( + "ticket=fresh-ticket", + ); transport.destroy(); }); }); -describe('CloudAgentTransport lifecycle', () => { - it('disconnect() closes connection', async () => { +describe("CloudAgentTransport lifecycle", () => { + it("disconnect() closes connection", async () => { const { transport } = createTransportWithSinks(); transport.connect(); @@ -353,7 +427,7 @@ describe('CloudAgentTransport lifecycle', () => { expect(mockWs.close).toHaveBeenCalled(); }); - it('destroy() closes connection', async () => { + it("destroy() closes connection", async () => { const { transport } = createTransportWithSinks(); transport.connect(); @@ -363,13 +437,13 @@ describe('CloudAgentTransport lifecycle', () => { expect(mockWs.close).toHaveBeenCalled(); }); - it('stale generation after disconnect prevents connection creation', async () => { + it("stale generation after disconnect prevents connection creation", async () => { const resolveTicket: { resolve?: (value: string) => void } = {}; const getTicket = jest.fn( () => - new Promise(resolve => { + new Promise((resolve) => { resolveTicket.resolve = resolve; - }) + }), ); const { transport } = createTransportWithSinks(getTicket); @@ -379,12 +453,13 @@ describe('CloudAgentTransport lifecycle', () => { transport.disconnect(); // Now resolve the ticket — should be stale - resolveTicket.resolve?.('late-ticket'); + resolveTicket.resolve?.("late-ticket"); await flushPromises(); // Only the first WebSocket (from disconnect closing) should exist; // no new WebSocket created from the stale ticket resolution - const constructorCallsAfterDisconnect = webSocketConstructor.mock.calls.length; + const constructorCallsAfterDisconnect = + webSocketConstructor.mock.calls.length; // The initial connect() didn't create a WS (ticket was async and unresolved), // so no WS should have been constructed at all. @@ -392,124 +467,137 @@ describe('CloudAgentTransport lifecycle', () => { }); }); -describe('CloudAgentTransport command delegation', () => { - it('converts a Kilo model ref before delegating to api.send', async () => { +describe("CloudAgentTransport command delegation", () => { + it("converts a Kilo model ref before delegating to api.send", async () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); await transport.send!({ payload: { - type: 'prompt', - prompt: 'hello', - mode: 'code', - model: { providerID: 'kilo', modelID: 'gpt-4' }, + type: "prompt", + prompt: "hello", + mode: "code", + model: { providerID: "kilo", modelID: "gpt-4" }, }, }); expect(api.send).toHaveBeenCalledWith({ - sessionId: 'ses-1', - payload: { type: 'prompt', prompt: 'hello', mode: 'code', model: 'gpt-4' }, + sessionId: "ses-1", + payload: { + type: "prompt", + prompt: "hello", + mode: "code", + model: "gpt-4", + }, }); transport.destroy(); }); - it('rejects a non-Kilo model ref before calling the Cloud Agent API', async () => { + it("rejects a non-Kilo model ref before calling the Cloud Agent API", async () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); await expect( transport.send!({ payload: { - type: 'prompt', - prompt: 'hello', - mode: 'code', - model: { providerID: 'anthropic', modelID: 'claude-sonnet-4' }, + type: "prompt", + prompt: "hello", + mode: "code", + model: { providerID: "anthropic", modelID: "claude-sonnet-4" }, }, - }) - ).rejects.toThrow('Cloud Agent only supports Kilo models'); + }), + ).rejects.toThrow("Cloud Agent only supports Kilo models"); expect(api.send).not.toHaveBeenCalled(); transport.destroy(); }); - it('converts Kilo model refs while preserving canonical document attachments', async () => { + it("converts Kilo model refs while preserving canonical document attachments", async () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); const attachments = { - path: '12345678-1234-4234-9234-123456789abc', - files: ['87654321-4321-4321-8321-cba987654321.pdf'], + path: "12345678-1234-4234-9234-123456789abc", + files: ["87654321-4321-4321-8321-cba987654321.pdf"], }; await transport.send!({ payload: { - type: 'prompt', - prompt: 'read it', - mode: 'code', - model: { providerID: 'kilo', modelID: 'gpt-4' }, + type: "prompt", + prompt: "read it", + mode: "code", + model: { providerID: "kilo", modelID: "gpt-4" }, }, attachments, }); expect(api.send).toHaveBeenCalledWith({ - sessionId: 'ses-1', - payload: { type: 'prompt', prompt: 'read it', mode: 'code', model: 'gpt-4' }, + sessionId: "ses-1", + payload: { + type: "prompt", + prompt: "read it", + mode: "code", + model: "gpt-4", + }, attachments, }); transport.destroy(); }); - it('interrupt() delegates to api.interrupt with bound sessionId', () => { + it("interrupt() delegates to api.interrupt with bound sessionId", () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); void transport.interrupt!(); - expect(api.interrupt).toHaveBeenCalledWith({ sessionId: 'ses-1' }); + expect(api.interrupt).toHaveBeenCalledWith({ sessionId: "ses-1" }); transport.destroy(); }); - it('answer() delegates to api.answer with bound sessionId', () => { + it("answer() delegates to api.answer with bound sessionId", () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); - void transport.answer!({ requestId: 'req-1', answers: [['yes']] }); + void transport.answer!({ requestId: "req-1", answers: [["yes"]] }); expect(api.answer).toHaveBeenCalledWith({ - sessionId: 'ses-1', - requestId: 'req-1', - answers: [['yes']], + sessionId: "ses-1", + requestId: "req-1", + answers: [["yes"]], }); transport.destroy(); }); - it('reject() delegates to api.reject with bound sessionId', () => { + it("reject() delegates to api.reject with bound sessionId", () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); - void transport.reject!({ requestId: 'req-2' }); + void transport.reject!({ requestId: "req-2" }); expect(api.reject).toHaveBeenCalledWith({ - sessionId: 'ses-1', - requestId: 'req-2', + sessionId: "ses-1", + requestId: "req-2", }); transport.destroy(); }); - it('respondToPermission() delegates to api.respondToPermission with bound sessionId', () => { + it("respondToPermission() delegates to api.respondToPermission with bound sessionId", () => { 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', - requestId: 'req-3', - response: 'once', + sessionId: "ses-1", + requestId: "req-3", + response: "once", }); transport.destroy(); @@ -520,7 +608,7 @@ describe('CloudAgentTransport command delegation', () => { // Snapshot refetch on reconnect // --------------------------------------------------------------------------- -describe('CloudAgentTransport snapshot refetch on reconnect', () => { +describe("CloudAgentTransport snapshot refetch on reconnect", () => { // Microtask-based flush that works under jest.useFakeTimers() // (unlike flushPromises which uses setTimeout and hangs with fake timers) async function flushMicrotasks(): Promise { @@ -530,7 +618,7 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { } function createTransportWithControllableSnapshot( - snapshotOverride?: ReturnType + snapshotOverride?: ReturnType, ) { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; @@ -538,17 +626,17 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { const fetchSnapshot = jest.fn(() => Promise.resolve(snapshot)); const factory = createCloudAgentTransport({ - sessionId: cloudAgentId('ses-1'), - kiloSessionId: kiloId('ses-1'), + sessionId: cloudAgentId("ses-1"), + kiloSessionId: kiloId("ses-1"), api: createMockApi(), - getTicket: () => 'test-ticket', + getTicket: () => "test-ticket", fetchSnapshot, - websocketBaseUrl: 'ws://localhost:9999', + websocketBaseUrl: "ws://localhost:9999", }); const transport = factory({ - onChatEvent: event => chatEvents.push(event), - onServiceEvent: event => serviceEvents.push(event), + onChatEvent: (event) => chatEvents.push(event), + onServiceEvent: (event) => serviceEvents.push(event), }); return { transport, chatEvents, serviceEvents, fetchSnapshot }; @@ -563,38 +651,46 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { return { eventId: 0, executionId: null, - sessionId: 'ses-1', - streamEventType: 'kilocode', + sessionId: "ses-1", + streamEventType: "kilocode", timestamp: new Date().toISOString(), data: { - type: 'session.status', - properties: { sessionID: 'ses-1', status: { type: 'busy' } }, + type: "session.status", + properties: { sessionID: "ses-1", status: { type: "busy" } }, }, }; } /** Establish connection, simulate close + reconnect, return the new WS mock. */ - async function simulateReconnect(establishEvent?: CloudAgentEvent): Promise { - mockWs.onclose?.({ code: 1006, reason: '', wasClean: false } as CloseEvent); + async function simulateReconnect( + establishEvent?: CloudAgentEvent, + ): Promise { + mockWs.onclose?.({ code: 1006, reason: "", wasClean: false } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); - const newMockWs = webSocketConstructor.mock.results.at(-1)?.value as MockWebSocket; + const newMockWs = webSocketConstructor.mock.results.at(-1) + ?.value as MockWebSocket; - newMockWs.onopen?.(new Event('open')); + 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; } - it('resumes from the replay cursor on reconnect instead of refetching the snapshot', async () => { + it("resumes from the replay cursor on reconnect instead of refetching the snapshot", async () => { jest.useFakeTimers(); try { - const { transport, serviceEvents, fetchSnapshot } = createTransportWithControllableSnapshot(); + const { transport, serviceEvents, fetchSnapshot } = + createTransportWithControllableSnapshot(); transport.connect(); await flushMicrotasks(); @@ -602,9 +698,9 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { expect(fetchSnapshot).toHaveBeenCalledTimes(1); // Establish connection with a persisted event — its id becomes the cursor. - const establish = kilocode('session.status', { - sessionID: 'ses-1', - status: { type: 'busy' }, + const establish = kilocode("session.status", { + sessionID: "ses-1", + status: { type: "busy" }, }); sendRaw(establish); @@ -618,40 +714,44 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { expect(fetchSnapshot).toHaveBeenCalledTimes(1); const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); expect(reconnectUrl).toContain(`fromId=${establish.eventId}`); - expect(reconnectUrl).not.toContain('replay=false'); + expect(reconnectUrl).not.toContain("replay=false"); const replayedCreated = serviceEvents .slice(serviceCountBefore) - .filter(e => e.type === 'session.created'); + .filter((e) => e.type === "session.created"); 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(); } }); - it('refetches the snapshot on reconnect when no replay cursor exists and upserts it', async () => { + it("refetches the snapshot on reconnect when no replay cursor exists and upserts it", async () => { jest.useFakeTimers(); try { - const snapshotWithMessages = makeSnapshot({ id: 'ses-1' }, [ + const snapshotWithMessages = makeSnapshot({ id: "ses-1" }, [ { info: { - id: 'msg-1', - sessionID: 'ses-1', - role: 'user', + id: "msg-1", + sessionID: "ses-1", + role: "user", time: { created: 1 }, - agent: 'build', - model: { providerID: 'a', modelID: 'b' }, + agent: "build", + model: { providerID: "a", modelID: "b" }, }, parts: [ { - id: 'part-1', - sessionID: 'ses-1', - messageID: 'msg-1', - type: 'text', - text: 'hello', + id: "part-1", + sessionID: "ses-1", + messageID: "msg-1", + type: "text", + text: "hello", }, ], }, @@ -663,9 +763,15 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { transport.connect(); await flushMicrotasks(); - expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(1); - expect(chatEvents.filter(e => e.type === 'message.updated')).toHaveLength(1); - expect(chatEvents.filter(e => e.type === 'message.part.updated')).toHaveLength(1); + expect( + serviceEvents.filter((e) => e.type === "session.created"), + ).toHaveLength(1); + expect( + chatEvents.filter((e) => e.type === "message.updated"), + ).toHaveLength(1); + expect( + chatEvents.filter((e) => e.type === "message.part.updated"), + ).toHaveLength(1); // Establish connection with sentinel events only — no replay cursor, so // reconnect falls back to the snapshot refetch. @@ -675,36 +781,49 @@ describe('CloudAgentTransport snapshot refetch on reconnect', () => { await flushMicrotasks(); expect(fetchSnapshot).toHaveBeenCalledTimes(2); - expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(2); - expect(chatEvents.filter(e => e.type === 'message.updated')).toHaveLength(2); - expect(chatEvents.filter(e => e.type === 'message.part.updated')).toHaveLength(2); + expect( + serviceEvents.filter((e) => e.type === "session.created"), + ).toHaveLength(2); + expect( + chatEvents.filter((e) => e.type === "message.updated"), + ).toHaveLength(2); + 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(); } }); - it('initial connect fetches snapshot once and opens WebSocket', async () => { - const { transport, serviceEvents, fetchSnapshot } = createTransportWithControllableSnapshot(); + it("initial connect fetches snapshot once and opens WebSocket", async () => { + const { transport, serviceEvents, fetchSnapshot } = + createTransportWithControllableSnapshot(); transport.connect(); await flushPromises(); expect(fetchSnapshot).toHaveBeenCalledTimes(1); - expect(fetchSnapshot).toHaveBeenCalledWith('ses-1'); + expect(fetchSnapshot).toHaveBeenCalledWith("ses-1"); expect(webSocketConstructor).toHaveBeenCalledTimes(1); - expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(1); + expect( + serviceEvents.filter((e) => e.type === "session.created"), + ).toHaveLength(1); transport.destroy(); }); }); -describe('CloudAgentTransport page-seam', () => { +describe("CloudAgentTransport page-seam", () => { function createTransportWithPageFetch( page: SessionSnapshotPageOutcome | null, - onError?: (message: string) => void + onError?: (message: string) => void, ) { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; @@ -712,31 +831,38 @@ describe('CloudAgentTransport page-seam', () => { const onInitialPageLoaded = jest.fn(); const factory = createCloudAgentTransport({ - sessionId: cloudAgentId('ses-1'), - kiloSessionId: kiloId('ses-1'), + sessionId: cloudAgentId("ses-1"), + kiloSessionId: kiloId("ses-1"), api: createMockApi(), - getTicket: () => 'test-ticket', - fetchSnapshot: () => Promise.reject(new Error('legacy fetchSnapshot should not be called')), + getTicket: () => "test-ticket", + fetchSnapshot: () => + Promise.reject(new Error("legacy fetchSnapshot should not be called")), fetchSnapshotPage, onInitialPageLoaded, - websocketBaseUrl: 'ws://localhost:9999', + websocketBaseUrl: "ws://localhost:9999", onError, }); const transport = factory({ - onChatEvent: event => chatEvents.push(event), - onServiceEvent: event => serviceEvents.push(event), + onChatEvent: (event) => chatEvents.push(event), + 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 () => { + it("uses fetchSnapshotPage for the initial bounded read and reports the page via onInitialPageLoaded", async () => { const page = { - kind: 'success' as const, - info: { id: 'ses-1' }, + kind: "success" as const, + info: { id: "ses-1" }, messages: [], - nextCursor: 'cursor-A', + nextCursor: "cursor-A", omittedItemCount: 2, }; const { transport, fetchSnapshotPage, onInitialPageLoaded } = @@ -746,21 +872,23 @@ describe('CloudAgentTransport page-seam', () => { await flushPromises(); expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); - expect(fetchSnapshotPage).toHaveBeenCalledWith('ses-1', {}); + expect(fetchSnapshotPage).toHaveBeenCalledWith("ses-1", {}); expect(onInitialPageLoaded).toHaveBeenCalledWith(page); transport.destroy(); }); - it('surfaces typed failures on the initial read via onError', async () => { + it("surfaces typed failures on the initial read via onError", async () => { const { transport, serviceEvents } = createTransportWithPageFetch({ - kind: 'retryable_failure', + kind: "retryable_failure", }); transport.connect(); await flushPromises(); - expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(0); + expect( + serviceEvents.filter((e) => e.type === "session.created"), + ).toHaveLength(0); // The transport still connects the websocket so the user can recover // via live events. expect(webSocketConstructor).toHaveBeenCalledTimes(1); @@ -768,17 +896,21 @@ describe('CloudAgentTransport page-seam', () => { transport.destroy(); }); - it('surfaces invalid_data on the initial read via onError and still connects the websocket', async () => { + it("surfaces invalid_data on the initial read via onError and still connects the websocket", async () => { const errors: string[] = []; const { transport, chatEvents, serviceEvents, onInitialPageLoaded } = - createTransportWithPageFetch({ kind: 'invalid_data' }, message => errors.push(message)); + createTransportWithPageFetch({ kind: "invalid_data" }, (message) => + errors.push(message), + ); transport.connect(); await flushPromises(); - expect(errors).toEqual(['Session history is unavailable']); + expect(errors).toEqual(["Session history is unavailable"]); expect(onInitialPageLoaded).not.toHaveBeenCalled(); - expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(0); + expect( + serviceEvents.filter((e) => e.type === "session.created"), + ).toHaveLength(0); expect(chatEvents).toHaveLength(0); // Terminal typed failures still leave the websocket available so the user // can recover via live events; this is the intended transport behavior. @@ -787,17 +919,21 @@ describe('CloudAgentTransport page-seam', () => { transport.destroy(); }); - it('surfaces too_large on the initial read via onError and still connects the websocket', async () => { + it("surfaces too_large on the initial read via onError and still connects the websocket", async () => { const errors: string[] = []; const { transport, chatEvents, serviceEvents, onInitialPageLoaded } = - createTransportWithPageFetch({ kind: 'too_large' }, message => errors.push(message)); + createTransportWithPageFetch({ kind: "too_large" }, (message) => + errors.push(message), + ); transport.connect(); await flushPromises(); - expect(errors).toEqual(['Session history too large to load']); + expect(errors).toEqual(["Session history too large to load"]); expect(onInitialPageLoaded).not.toHaveBeenCalled(); - expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(0); + expect( + serviceEvents.filter((e) => e.type === "session.created"), + ).toHaveLength(0); expect(chatEvents).toHaveLength(0); // Terminal typed failures still leave the websocket available so the user // can recover via live events; this is the intended transport behavior. @@ -806,35 +942,37 @@ describe('CloudAgentTransport page-seam', () => { transport.destroy(); }); - it('falls back to fetchSnapshot when fetchSnapshotPage is not provided', async () => { + it("falls back to fetchSnapshot when fetchSnapshotPage is not provided", async () => { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; const fetchSnapshot = jest.fn(() => Promise.resolve(emptySnapshot)); const factory = createCloudAgentTransport({ - sessionId: cloudAgentId('ses-1'), - kiloSessionId: kiloId('ses-1'), + sessionId: cloudAgentId("ses-1"), + kiloSessionId: kiloId("ses-1"), api: createMockApi(), - getTicket: () => 'test-ticket', + getTicket: () => "test-ticket", fetchSnapshot, - websocketBaseUrl: 'ws://localhost:9999', + websocketBaseUrl: "ws://localhost:9999", }); const transport = factory({ - onChatEvent: event => chatEvents.push(event), - onServiceEvent: event => serviceEvents.push(event), + onChatEvent: (event) => chatEvents.push(event), + onServiceEvent: (event) => serviceEvents.push(event), }); transport.connect(); await flushPromises(); expect(fetchSnapshot).toHaveBeenCalledTimes(1); - expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(1); + expect( + serviceEvents.filter((e) => e.type === "session.created"), + ).toHaveLength(1); transport.destroy(); }); - it('reconnect (no eventId) uses fetchSnapshotPage and does NOT fire onInitialPageLoaded', async () => { + it("reconnect (no eventId) uses fetchSnapshotPage and does NOT fire onInitialPageLoaded", async () => { jest.useFakeTimers(); try { // Microtask-based flush that works under jest.useFakeTimers() @@ -846,10 +984,10 @@ describe('CloudAgentTransport page-seam', () => { } const page = { - kind: 'success' as const, - info: { id: 'ses-1' }, + kind: "success" as const, + info: { id: "ses-1" }, messages: [], - nextCursor: 'cursor-A', + nextCursor: "cursor-A", omittedItemCount: 0, }; const fetchSnapshotPage = jest.fn().mockResolvedValue(page); @@ -858,19 +996,19 @@ describe('CloudAgentTransport page-seam', () => { const serviceEvents: ServiceEvent[] = []; const factory = createCloudAgentTransport({ - sessionId: cloudAgentId('ses-1'), - kiloSessionId: kiloId('ses-1'), + sessionId: cloudAgentId("ses-1"), + kiloSessionId: kiloId("ses-1"), api: createMockApi(), - getTicket: () => 'test-ticket', - fetchSnapshot: () => Promise.reject(new Error('should not be called')), + getTicket: () => "test-ticket", + fetchSnapshot: () => Promise.reject(new Error("should not be called")), fetchSnapshotPage, onInitialPageLoaded, - websocketBaseUrl: 'ws://localhost:9999', + websocketBaseUrl: "ws://localhost:9999", }); const transport = factory({ - onChatEvent: event => chatEvents.push(event), - onServiceEvent: event => serviceEvents.push(event), + onChatEvent: (event) => chatEvents.push(event), + onServiceEvent: (event) => serviceEvents.push(event), }); transport.connect(); @@ -884,23 +1022,28 @@ describe('CloudAgentTransport page-seam', () => { const sentinel: CloudAgentEvent = { eventId: 0, executionId: null, - sessionId: 'ses-1', - streamEventType: 'kilocode', + sessionId: "ses-1", + streamEventType: "kilocode", timestamp: new Date().toISOString(), data: { - type: 'session.status', - properties: { sessionID: 'ses-1', status: { type: 'busy' } }, + type: "session.status", + properties: { sessionID: "ses-1", status: { type: "busy" } }, }, }; 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(); - const newMockWs = webSocketConstructor.mock.results.at(-1)?.value as MockWebSocket; - newMockWs.onopen?.(new Event('open')); + const newMockWs = webSocketConstructor.mock.results.at(-1) + ?.value as MockWebSocket; + newMockWs.onopen?.(new Event("open")); // Send a sentinel on the new socket: the connection marks itself // reconnected, then `onReconnected` falls back to `fetchSnapshotPage` // because `lastEventId` is still null. @@ -920,11 +1063,11 @@ describe('CloudAgentTransport page-seam', () => { } }); - describe('watermark cursor', () => { - it('uses fromId on first connect when the page carries a watermark', async () => { + describe("watermark cursor", () => { + it("uses fromId on first connect when the page carries a watermark", async () => { const page = { - kind: 'success' as const, - info: { id: 'ses-1' } as const, + kind: "success" as const, + info: { id: "ses-1" } as const, messages: [], nextCursor: null, omittedItemCount: 0, @@ -937,16 +1080,16 @@ describe('CloudAgentTransport page-seam', () => { expect(webSocketConstructor).toHaveBeenCalledTimes(1); const wsUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(wsUrl).toContain('fromId=42'); - expect(wsUrl).not.toContain('replay=false'); + expect(wsUrl).toContain("fromId=42"); + expect(wsUrl).not.toContain("replay=false"); transport.destroy(); }); - it('uses replay=false on first connect when the page has no watermark', async () => { + 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, + kind: "success" as const, + info: { id: "ses-1" } as const, messages: [], nextCursor: null, omittedItemCount: 0, @@ -958,13 +1101,13 @@ describe('CloudAgentTransport page-seam', () => { expect(webSocketConstructor).toHaveBeenCalledTimes(1); const wsUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(wsUrl).toContain('replay=false'); - expect(wsUrl).not.toContain('fromId'); + 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 () => { + it("uses fromId with watermark null on reconnect when wire events have set lastEventId", async () => { jest.useFakeTimers(); try { async function flushMicrotasks(): Promise { @@ -974,8 +1117,8 @@ describe('CloudAgentTransport page-seam', () => { } const page = { - kind: 'success' as const, - info: { id: 'ses-1' } as const, + kind: "success" as const, + info: { id: "ses-1" } as const, messages: [], nextCursor: null, omittedItemCount: 0, @@ -986,18 +1129,19 @@ describe('CloudAgentTransport page-seam', () => { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; const factory = createCloudAgentTransport({ - sessionId: cloudAgentId('ses-1'), - kiloSessionId: kiloId('ses-1'), + sessionId: cloudAgentId("ses-1"), + kiloSessionId: kiloId("ses-1"), api: createMockApi(), - getTicket: () => 'test-ticket', - fetchSnapshot: () => Promise.reject(new Error('should not be called')), + getTicket: () => "test-ticket", + fetchSnapshot: () => + Promise.reject(new Error("should not be called")), fetchSnapshotPage, - websocketBaseUrl: 'ws://localhost:9999', + websocketBaseUrl: "ws://localhost:9999", }); const transport = factory({ - onChatEvent: event => chatEvents.push(event), - onServiceEvent: event => serviceEvents.push(event), + onChatEvent: (event) => chatEvents.push(event), + onServiceEvent: (event) => serviceEvents.push(event), }); transport.connect(); @@ -1005,30 +1149,36 @@ describe('CloudAgentTransport page-seam', () => { // Verify first connect used the watermark const firstUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(firstUrl).toContain('fromId=10'); + expect(firstUrl).toContain("fromId=10"); // Send a wire event with eventId > 0 to advance the live cursor const establish = { eventId: 55, executionId: null, - sessionId: 'ses-1', - streamEventType: 'kilocode', + sessionId: "ses-1", + streamEventType: "kilocode", timestamp: new Date().toISOString(), data: { - type: 'session.status', - properties: { sessionID: 'ses-1', status: { type: 'busy' } }, + type: "session.status", + properties: { sessionID: "ses-1", status: { type: "busy" } }, }, }; sendRaw(establish); // Disconnect and reconnect - mockWs.onclose?.({ code: 1006, reason: '', wasClean: false } as CloseEvent); + mockWs.onclose?.({ + code: 1006, + reason: "", + wasClean: false, + } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); - const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); + 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'); + 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); @@ -1039,7 +1189,7 @@ describe('CloudAgentTransport page-seam', () => { } }); - it('reconnect on a watermarked page with only sentinel events uses the original watermark and skips the page refetch', async () => { + it("reconnect on a watermarked page with only sentinel events uses the original watermark and skips the page refetch", async () => { jest.useFakeTimers(); try { async function flushMicrotasks(): Promise { @@ -1049,8 +1199,8 @@ describe('CloudAgentTransport page-seam', () => { } const page = { - kind: 'success' as const, - info: { id: 'ses-1' } as const, + kind: "success" as const, + info: { id: "ses-1" } as const, messages: [], nextCursor: null, omittedItemCount: 0, @@ -1061,18 +1211,19 @@ describe('CloudAgentTransport page-seam', () => { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; const factory = createCloudAgentTransport({ - sessionId: cloudAgentId('ses-1'), - kiloSessionId: kiloId('ses-1'), + sessionId: cloudAgentId("ses-1"), + kiloSessionId: kiloId("ses-1"), api: createMockApi(), - getTicket: () => 'test-ticket', - fetchSnapshot: () => Promise.reject(new Error('should not be called')), + getTicket: () => "test-ticket", + fetchSnapshot: () => + Promise.reject(new Error("should not be called")), fetchSnapshotPage, - websocketBaseUrl: 'ws://localhost:9999', + websocketBaseUrl: "ws://localhost:9999", }); const transport = factory({ - onChatEvent: event => chatEvents.push(event), - onServiceEvent: event => serviceEvents.push(event), + onChatEvent: (event) => chatEvents.push(event), + onServiceEvent: (event) => serviceEvents.push(event), }); transport.connect(); @@ -1080,7 +1231,7 @@ describe('CloudAgentTransport page-seam', () => { // Verify first connect used the watermark as fromId. const firstUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(firstUrl).toContain('fromId=42'); + expect(firstUrl).toContain("fromId=42"); expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); // Send only sentinel events (eventId: 0) — they do NOT advance the @@ -1088,26 +1239,32 @@ describe('CloudAgentTransport page-seam', () => { const sentinel: CloudAgentEvent = { eventId: 0, executionId: null, - sessionId: 'ses-1', - streamEventType: 'kilocode', + sessionId: "ses-1", + streamEventType: "kilocode", timestamp: new Date().toISOString(), data: { - type: 'session.status', - properties: { sessionID: 'ses-1', status: { type: 'busy' } }, + type: "session.status", + properties: { sessionID: "ses-1", status: { type: "busy" } }, }, }; sendRaw(sentinel); // Disconnect and reconnect. - mockWs.onclose?.({ code: 1006, reason: '', wasClean: false } as CloseEvent); + mockWs.onclose?.({ + code: 1006, + reason: "", + wasClean: false, + } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); - const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); + const reconnectUrl = String( + webSocketConstructor.mock.calls.at(-1)?.[0], + ); // Reconnect must still use the watermark (42) because sentinel events // never advance the cursor past the watermark. - expect(reconnectUrl).toContain('fromId=42'); - expect(reconnectUrl).not.toContain('replay=false'); + expect(reconnectUrl).toContain("fromId=42"); + expect(reconnectUrl).not.toContain("replay=false"); // No second page fetch — lastEventId (watermark 42) is not null, so // onReconnected skips the snapshot fallback path entirely. @@ -1120,3 +1277,286 @@ describe('CloudAgentTransport page-seam', () => { }); }); }); + +// --------------------------------------------------------------------------- +// 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 { + type: "session.status"; + sessionId: string; + status: { + type: "retry"; + attempt: number; + message: string; + next: number; + }; + } => e.type === "session.status" && 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 95e17a55ae8..5f762e3ed59 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts @@ -199,6 +199,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; } From a0f06e769f51dc200563776e7661d76f56045bc7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Mon, 3 Aug 2026 22:47:33 +0200 Subject: [PATCH 03/24] fix(cli): persist remote command outcomes --- .../src/cli-live-transport.test.ts | 42 +- .../cloud-agent-sdk/src/cli-live-transport.ts | 468 +- packages/cloud-agent-sdk/src/session.test.ts | 6 +- packages/cloud-agent-sdk/src/transport.ts | 2 + .../src/user-web-connection.ts | 16 +- .../src/dos/UserConnectionDO.test.ts | 4503 ++++++++++------- .../src/dos/UserConnectionDO.ts | 1207 ++++- .../src/types/user-connection-protocol.ts | 10 + 8 files changed, 4075 insertions(+), 2179 deletions(-) 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 195afff3da3..81438b9f5d5 100644 --- a/packages/cloud-agent-sdk/src/cli-live-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cli-live-transport.test.ts @@ -2347,7 +2347,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(); @@ -2474,7 +2475,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(); @@ -2533,7 +2535,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(); }); @@ -2640,6 +2643,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 f585e5aa18d..ea8dba1f947 100644 --- a/packages/cloud-agent-sdk/src/cli-live-transport.ts +++ b/packages/cloud-agent-sdk/src/cli-live-transport.ts @@ -2,39 +2,42 @@ * CLI live transport - consumes a shared user web connection and translates * one remote CLI session into normalized transport events and commands. */ -import { normalizeCliEvent, isChatEvent } from './normalizer'; -import { parseRemoteCommandCatalog, type RemoteCommandState } from './remote-command-catalog'; -import { parseCreateSessionResponse } from './create-session'; -import { parseExitSessionResponse } from './exit-session'; +import { normalizeCliEvent, isChatEvent } from "./normalizer"; +import { + parseRemoteCommandCatalog, + type RemoteCommandState, +} from "./remote-command-catalog"; +import { parseCreateSessionResponse } from "./create-session"; +import { parseExitSessionResponse } from "./exit-session"; import { cliConnectionDataSchema, heartbeatDataSchema, remoteModelCatalogV1Schema, sessionsListDataSchema, -} from './schemas'; -import type { SlashCommandInfo } from './schemas'; -import type { RemoteModelState } from './remote-model-catalog'; +} from "./schemas"; +import type { SlashCommandInfo } from "./schemas"; +import type { RemoteModelState } from "./remote-model-catalog"; import type { CreateRemoteSessionInput, TransportFactory, TransportSendInput, TransportSink, -} from './transport'; +} from "./transport"; import type { KiloSessionId, SessionSnapshot, SessionSnapshotPage, SessionSnapshotPageOutcome, -} from './types'; +} from "./types"; import { CommandDeliveredError, UserWebCommandError, type UserWebCliEvent, type UserWebConnection, -} from './user-web-connection'; +} from "./user-web-connection"; /** Delivered error string for old-CLI strict-parse rejection of extended create_session. */ -const INVALID_CREATE_SESSION_COMMAND = 'invalid create_session command'; +const INVALID_CREATE_SESSION_COMMAND = "invalid create_session command"; /** * Build `create_session` wire data: always `protocolVersion: 1`, plus any @@ -57,7 +60,8 @@ function buildCreateSessionWireData(input?: CreateRemoteSessionInput): { type CliLiveTransportConfig = { kiloSessionId: KiloSessionId; userWebConnection: UserWebConnection; - fetchSnapshot?: ((kiloSessionId: KiloSessionId) => Promise) | undefined; + fetchSnapshot?: + ((kiloSessionId: KiloSessionId) => Promise) | undefined; /** * Page-aware root snapshot fetch. When provided, every `replayCurrentSnapshot` * (initial bounded read AND reconnect immediate + delayed resync) uses it @@ -67,7 +71,7 @@ type CliLiveTransportConfig = { */ fetchSnapshotPage?: ( kiloSessionId: KiloSessionId, - options: { cursor?: string } + options: { cursor?: string }, ) => Promise; /** * Called after a successful initial bounded page read. Reconnect and @@ -77,7 +81,8 @@ type CliLiveTransportConfig = { onInitialPageLoaded?: ((page: SessionSnapshotPage) => void) | undefined; onError?: ((message: string) => void) | undefined; onRemoteModelStateChange?: ((state: RemoteModelState) => void) | undefined; - onRemoteCommandStateChange?: ((state: RemoteCommandState) => void) | undefined; + onRemoteCommandStateChange?: + ((state: RemoteCommandState) => void) | undefined; onCapabilityChange?: (() => void) | undefined; /** * Fired whenever the per-session capabilities advertised by the owning @@ -88,7 +93,9 @@ type CliLiveTransportConfig = { * to recompute the `supportsAttachments` gate. */ onCapabilitiesChange?: - | ((capabilities: { attachments?: boolean | undefined } | undefined) => void) + | (( + capabilities: { attachments?: boolean | undefined } | undefined, + ) => void) | undefined; }; @@ -97,7 +104,7 @@ type CliLiveTransportConfig = { // still appear after long backgrounding. const RECONNECT_RESYNC_DELAY_MS = 5000; const REMOTE_SESSION_EXIT_UNAVAILABLE = - 'Remote session exit is unavailable for the current session'; + "Remote session exit is unavailable for the current session"; /** * Deep-copy a list of validated remote slash commands. @@ -108,14 +115,18 @@ const REMOTE_SESSION_EXIT_UNAVAILABLE = * schema has already validated them. Avoids `JSON.parse(JSON.stringify(...))` * because the shape is known and bounded. */ -export function deepCopyCommands(commands: SlashCommandInfo[]): SlashCommandInfo[] { - return commands.map(command => ({ +export function deepCopyCommands( + commands: SlashCommandInfo[], +): SlashCommandInfo[] { + return commands.map((command) => ({ ...command, hints: command.hints.slice(), })); } -function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactory { +function createCliLiveTransport( + config: CliLiveTransportConfig, +): TransportFactory { return (sink: TransportSink) => { let generation = 0; let cleanup: (() => void) | null = null; @@ -129,8 +140,11 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor * payload that omitted this session). Compared structurally on every * new observation; only emitted on a real change. */ - let currentCapabilities: { attachments?: boolean | undefined } | undefined = undefined; - function publishCapabilities(next: { attachments?: boolean | undefined } | undefined): void { + let currentCapabilities: { attachments?: boolean | undefined } | undefined = + undefined; + function publishCapabilities( + next: { attachments?: boolean | undefined } | undefined, + ): void { const previousAttachments = currentCapabilities?.attachments; const nextAttachments = next?.attachments; if (previousAttachments === nextAttachments) return; @@ -138,7 +152,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. @@ -149,12 +166,12 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor } | null = null; let remoteModelState: RemoteModelState = { ownerConnectionId: null, - protocol: 'unknown', - refresh: 'idle', + protocol: "unknown", + refresh: "idle", }; let remoteCommandState: RemoteCommandState = { ownerConnectionId: null, - refresh: 'idle', + refresh: "idle", commands: [], canExitSession: undefined, }; @@ -168,13 +185,17 @@ 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); } function canExitSession(): boolean { return ( - remoteCommandState.canExitSession === true && remoteCommandState.ownerConnectionId !== null + remoteCommandState.canExitSession === true && + remoteCommandState.ownerConnectionId !== null ); } @@ -184,7 +205,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[] { @@ -202,8 +226,8 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor commandCatalogRequestInFlight = null; publishRemoteModelState({ ownerConnectionId: nextOwnerConnectionId, - protocol: 'unknown', - refresh: nextOwnerConnectionId ? 'loading' : 'idle', + protocol: "unknown", + refresh: nextOwnerConnectionId ? "loading" : "idle", }); // Owner replacement clears the command cache; emit an empty catalog // event and the matching idle state so consumers can re-render. @@ -211,14 +235,14 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor publishCommands([]); publishRemoteCommandState({ ownerConnectionId: nextOwnerConnectionId, - refresh: 'idle', + refresh: "idle", commands: snapshotCommands(), canExitSession: undefined, }); } else { publishRemoteCommandState({ ownerConnectionId: nextOwnerConnectionId, - refresh: 'idle', + refresh: "idle", commands: snapshotCommands(), canExitSession: undefined, }); @@ -240,7 +264,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor error: unknown, expectedOwnerConnectionId: string, expectedGeneration: number, - expectedRequestGeneration: number + expectedRequestGeneration: number, ): void { if ( expectedGeneration !== generation || @@ -250,16 +274,19 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor return; } - if (error instanceof UserWebCommandError && error.code === 'SESSION_OWNER_CHANGED') { + if ( + error instanceof UserWebCommandError && + error.code === "SESSION_OWNER_CHANGED" + ) { setOwnerConnectionId(null); return; } - if (error instanceof Error && error.message.includes('unknown command')) { + if (error instanceof Error && error.message.includes("unknown command")) { publishRemoteModelState({ ownerConnectionId: expectedOwnerConnectionId, - protocol: 'legacy', - refresh: 'idle', + protocol: "legacy", + refresh: "idle", }); return; } @@ -267,14 +294,22 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor publishRemoteModelState({ ownerConnectionId: expectedOwnerConnectionId, protocol: remoteModelState.protocol, - ...(remoteModelState.catalog ? { catalog: remoteModelState.catalog } : {}), - refresh: 'error', - error: error instanceof Error ? error.message : 'Failed to discover remote models', + ...(remoteModelState.catalog + ? { catalog: remoteModelState.catalog } + : {}), + refresh: "error", + error: + error instanceof Error + ? error.message + : "Failed to discover remote models", }); } function discoverModels(expectedOwnerConnectionId: string): void { - if (catalogRequestInFlight?.ownerConnectionId === expectedOwnerConnectionId) return; + if ( + catalogRequestInFlight?.ownerConnectionId === expectedOwnerConnectionId + ) + return; catalogRequestGeneration += 1; const expectedRequestGeneration = catalogRequestGeneration; @@ -286,19 +321,21 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor publishRemoteModelState({ ownerConnectionId: expectedOwnerConnectionId, protocol: remoteModelState.protocol, - ...(remoteModelState.catalog ? { catalog: remoteModelState.catalog } : {}), - refresh: 'loading', + ...(remoteModelState.catalog + ? { catalog: remoteModelState.catalog } + : {}), + refresh: "loading", }); void config.userWebConnection .sendCommand( config.kiloSessionId, - 'list_models', + "list_models", { protocolVersion: 1 }, - expectedOwnerConnectionId + expectedOwnerConnectionId, ) .then( - result => { + (result) => { if ( expectedGeneration !== generation || expectedRequestGeneration !== catalogRequestGeneration || @@ -310,31 +347,33 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor const parsed = remoteModelCatalogV1Schema.safeParse(result); if (!parsed.success) { handleCatalogFailure( - new Error('Invalid remote model catalog'), + new Error("Invalid remote model catalog"), expectedOwnerConnectionId, expectedGeneration, - expectedRequestGeneration + expectedRequestGeneration, ); return; } publishRemoteModelState({ ownerConnectionId: expectedOwnerConnectionId, - protocol: 'v1', + protocol: "v1", catalog: parsed.data, - refresh: 'idle', + refresh: "idle", }); }, - error => + (error) => handleCatalogFailure( error, expectedOwnerConnectionId, expectedGeneration, - expectedRequestGeneration - ) + expectedRequestGeneration, + ), ) .finally(() => { - if (catalogRequestInFlight?.generation === expectedRequestGeneration) { + if ( + catalogRequestInFlight?.generation === expectedRequestGeneration + ) { catalogRequestInFlight = null; } }); @@ -345,7 +384,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor expectedOwnerConnectionId: string, expectedGeneration: number, expectedRequestGeneration: number, - clearCatalog: boolean + clearCatalog: boolean, ): void { if ( expectedGeneration !== generation || @@ -355,7 +394,10 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor return; } - if (error instanceof UserWebCommandError && error.code === 'SESSION_OWNER_CHANGED') { + if ( + error instanceof UserWebCommandError && + error.code === "SESSION_OWNER_CHANGED" + ) { setOwnerConnectionId(null); return; } @@ -364,11 +406,14 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // and surface actionable copy through the command state so the chat // composer can prompt the user to upgrade. Never populate the global // session error atom for this. - if (error instanceof UserWebCommandError && error.code === 'CLI_UPGRADE_REQUIRED') { + if ( + error instanceof UserWebCommandError && + error.code === "CLI_UPGRADE_REQUIRED" + ) { publishCommands([]); publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: 'upgrade-required', + refresh: "upgrade-required", commands: snapshotCommands(), canExitSession: undefined, message: error.message, @@ -378,23 +423,26 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // Legacy CLI without `list_commands` support: treat as "no commands // available" rather than a failure so the composer can hide suggestions. - if (error instanceof Error && error.message.includes('unknown command')) { + if (error instanceof Error && error.message.includes("unknown command")) { publishCommands([]); publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: 'idle', + refresh: "idle", commands: snapshotCommands(), canExitSession: undefined, }); return; } - const message = error instanceof Error ? error.message : 'Failed to discover remote commands'; + const message = + error instanceof Error + ? error.message + : "Failed to discover remote commands"; if (clearCatalog) { publishCommands([]); publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: 'error', + refresh: "error", commands: snapshotCommands(), canExitSession: undefined, message, @@ -406,7 +454,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // populating the fatal session error atom. publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: 'error', + refresh: "error", commands: snapshotCommands(), canExitSession: undefined, message, @@ -414,7 +462,11 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor } function discoverCommands(expectedOwnerConnectionId: string): void { - if (commandCatalogRequestInFlight?.ownerConnectionId === expectedOwnerConnectionId) return; + if ( + commandCatalogRequestInFlight?.ownerConnectionId === + expectedOwnerConnectionId + ) + return; commandCatalogRequestGeneration += 1; const expectedRequestGeneration = commandCatalogRequestGeneration; @@ -425,7 +477,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor }; publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: 'loading', + refresh: "loading", commands: snapshotCommands(), canExitSession: undefined, }); @@ -433,12 +485,12 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor void config.userWebConnection .sendCommand( config.kiloSessionId, - 'list_commands', + "list_commands", { protocolVersion: 1 }, - expectedOwnerConnectionId + expectedOwnerConnectionId, ) .then( - result => { + (result) => { if ( expectedGeneration !== generation || expectedRequestGeneration !== commandCatalogRequestGeneration || @@ -450,11 +502,11 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor const parsed = parseRemoteCommandCatalog(result); if (!parsed.ok) { handleCommandCatalogFailure( - new Error('Invalid remote command catalog'), + new Error("Invalid remote command catalog"), expectedOwnerConnectionId, expectedGeneration, expectedRequestGeneration, - true + true, ); return; } @@ -462,12 +514,12 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor publishCommands(parsed.commands); publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: 'idle', + refresh: "idle", commands: snapshotCommands(), canExitSession: parsed.canExitSession, }); }, - error => + (error) => handleCommandCatalogFailure( error, expectedOwnerConnectionId, @@ -476,24 +528,28 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // A relay `CATALOG_TOO_LARGE` reports the response is over // the 512 KiB size cap; clear the prior catalog so the // composer can present a clean state. - error instanceof UserWebCommandError && error.code === 'CATALOG_TOO_LARGE' - ) + error instanceof UserWebCommandError && + error.code === "CATALOG_TOO_LARGE", + ), ) .finally(() => { - if (commandCatalogRequestInFlight?.generation === expectedRequestGeneration) { + if ( + commandCatalogRequestInFlight?.generation === + expectedRequestGeneration + ) { commandCatalogRequestInFlight = null; } }); } function replayPage(page: SessionSnapshotPage): void { - sink.onServiceEvent({ type: 'session.created', info: page.info }); + sink.onServiceEvent({ type: "session.created", info: page.info }); for (const msg of page.messages) { - sink.onChatEvent({ type: 'message.updated', info: msg.info }); + sink.onChatEvent({ type: "message.updated", info: msg.info }); for (const part of msg.parts) { - sink.onChatEvent({ type: 'message.part.updated', part }); + sink.onChatEvent({ type: "message.part.updated", part }); } } } @@ -502,9 +558,13 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor sessionId: string, parentSessionId: string | undefined, event: string, - data: unknown + data: unknown, ): void { - if (sessionId !== config.kiloSessionId && parentSessionId !== config.kiloSessionId) return; + if ( + sessionId !== config.kiloSessionId && + parentSessionId !== config.kiloSessionId + ) + return; const normalized = normalizeCliEvent(event, data); if (!normalized) return; @@ -518,7 +578,7 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor function stopForDisconnectedSession(): void { if (sessionStopped) return; - sink.onServiceEvent({ type: 'stopped', reason: 'disconnected' }); + sink.onServiceEvent({ type: "stopped", reason: "disconnected" }); sessionStopped = true; // The disconnected state is only cleared by a session.status event, so // the next post-reconnect heartbeat must always forward one — even when @@ -531,18 +591,18 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // fired while the socket was dead is never replayed, which otherwise // leaves the UI stuck on a busy indicator forever. function forwardHeartbeatStatus(status: string): void { - if (status !== 'idle' && status !== 'busy') return; + if (status !== "idle" && status !== "busy") return; if (status === lastForwardedHeartbeatStatus) return; lastForwardedHeartbeatStatus = status; sink.onServiceEvent({ - type: 'session.status', + type: "session.status", sessionId: config.kiloSessionId, status: { type: status }, }); } function handleSystemMessage(event: string, data: unknown): void { - if (event === 'cli.disconnected') { + if (event === "cli.disconnected") { const parsed = cliConnectionDataSchema.safeParse(data); if (parsed.success && ownerConnectionId === parsed.data.connectionId) { setOwnerConnectionId(null); @@ -551,11 +611,13 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor return; } - if (event === 'sessions.list') { + if (event === "sessions.list") { const parsed = sessionsListDataSchema.safeParse(data); if (!parsed.success) return; - const session = parsed.data.sessions.find(item => item.id === config.kiloSessionId); + const session = parsed.data.sessions.find( + (item) => item.id === config.kiloSessionId, + ); if (session) { setOwnerConnectionId(session.connectionId); sessionStopped = false; @@ -569,11 +631,13 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor return; } - if (event === 'sessions.heartbeat') { + if (event === "sessions.heartbeat") { const parsed = heartbeatDataSchema.safeParse(data); if (!parsed.success) return; - const session = parsed.data.sessions.find(item => item.id === config.kiloSessionId); + const session = parsed.data.sessions.find( + (item) => item.id === config.kiloSessionId, + ); if (session) { setOwnerConnectionId(parsed.data.connectionId); sessionStopped = false; @@ -589,19 +653,28 @@ 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'); + if (!expectedOwnerConnectionId) + throw new Error("Remote session has no connected owner"); try { return await config.userWebConnection.sendCommand( config.kiloSessionId, command, data, - expectedOwnerConnectionId + expectedOwnerConnectionId, + ...(mutationId !== undefined ? [mutationId] : []), ); } catch (error) { - if (error instanceof UserWebCommandError && error.code === 'SESSION_OWNER_CHANGED') { + if ( + error instanceof UserWebCommandError && + error.code === "SESSION_OWNER_CHANGED" + ) { setOwnerConnectionId(null); } throw error; @@ -609,52 +682,61 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor } function getRemoteModelFields(input: TransportSendInput): - | { kind: 'none' } + | { kind: "none" } | { - kind: 'structured'; + kind: "structured"; model: { providerID: string; modelID: string }; variant?: string; } - | { kind: 'legacy'; model: string; variant?: string } { + | { kind: "legacy"; model: string; variant?: string } { const override = input.remoteModelOverride; - if (!override) return { kind: 'none' }; + if (!override) return { kind: "none" }; - if (remoteModelState.protocol === 'v1' && override.source === 'cli-catalog') { + if ( + remoteModelState.protocol === "v1" && + override.source === "cli-catalog" + ) { const provider = remoteModelState.catalog?.providers.find( - item => item.id === override.selection.model.providerID + (item) => item.id === override.selection.model.providerID, + ); + const model = provider?.models.find( + (item) => item.id === override.selection.model.modelID, ); - const model = provider?.models.find(item => item.id === override.selection.model.modelID); if (!model) { - throw new Error('Selected remote model is not available in the current CLI catalog'); + throw new Error( + "Selected remote model is not available in the current CLI catalog", + ); } const variant = override.selection.variant; if (variant && !model.variants.includes(variant)) { throw new Error( - 'Selected remote model variant is not available in the current CLI catalog' + "Selected remote model variant is not available in the current CLI catalog", ); } return { - kind: 'structured', + kind: "structured", model: override.selection.model, ...(variant ? { variant } : {}), }; } if ( - remoteModelState.protocol === 'legacy' && - override.source === 'legacy-gateway' && - override.selection.model.providerID === 'kilo' + remoteModelState.protocol === "legacy" && + override.source === "legacy-gateway" && + override.selection.model.providerID === "kilo" ) { return { - kind: 'legacy', + kind: "legacy", model: override.selection.model.modelID, - ...(override.selection.variant ? { variant: override.selection.variant } : {}), + ...(override.selection.variant + ? { variant: override.selection.variant } + : {}), }; } throw new Error( - 'Selected remote model override is incompatible with the connected CLI model protocol' + "Selected remote model override is incompatible with the connected CLI model protocol", ); } @@ -678,12 +760,12 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor publishCapabilities(undefined); publishRemoteModelState({ ownerConnectionId: null, - protocol: 'unknown', - refresh: 'idle', + protocol: "unknown", + refresh: "idle", }); publishRemoteCommandState({ ownerConnectionId: null, - refresh: 'idle', + refresh: "idle", commands: snapshotCommands(), canExitSession: undefined, }); @@ -698,7 +780,12 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor const events = bufferedCliEvents; bufferedCliEvents = null; for (const msg of events ?? []) { - handleEventMessage(msg.sessionId, msg.parentSessionId, msg.event, msg.data); + handleEventMessage( + msg.sessionId, + msg.parentSessionId, + msg.event, + msg.data, + ); } sink.onReplayComplete?.(); }; @@ -748,7 +835,10 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor return; } if (reportError) { - const message = error instanceof Error ? error.message : 'Failed to fetch snapshot'; + const message = + error instanceof Error + ? error.message + : "Failed to fetch snapshot"; config.onError?.(message); } bufferedCliEvents = [ @@ -760,28 +850,30 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor }; if (config.fetchSnapshotPage) { - void config.fetchSnapshotPage(config.kiloSessionId, {}).then(page => { - if (page && page.kind === 'success') { - onPage(page); - } else { - onError( - new Error( - page === null - ? 'Session not found' - : page.kind === 'retryable_failure' - ? 'Session history temporarily unavailable' - : page.kind === 'too_large' - ? 'Session history too large to load' - : 'Session history is unavailable' - ) - ); - } - }, onError); + void config + .fetchSnapshotPage(config.kiloSessionId, {}) + .then((page) => { + if (page && page.kind === "success") { + onPage(page); + } else { + onError( + new Error( + page === null + ? "Session not found" + : page.kind === "retryable_failure" + ? "Session history temporarily unavailable" + : page.kind === "too_large" + ? "Session history too large to load" + : "Session history is unavailable", + ), + ); + } + }, onError); return; } if (config.fetchSnapshot) { - void config.fetchSnapshot(config.kiloSessionId).then(snapshot => { + void config.fetchSnapshot(config.kiloSessionId).then((snapshot) => { onPage({ info: snapshot.info, messages: snapshot.messages, @@ -793,20 +885,28 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor }; replayCurrentSnapshot(true); - const offCli = config.userWebConnection.onCliEvent(config.kiloSessionId, msg => { - const normalized = normalizeCliEvent(msg.event, msg.data); - const shouldBufferForSnapshot = - normalized && - (isChatEvent(normalized) || - normalized.type === 'session.created' || - normalized.type === 'session.updated'); - if (shouldBufferForSnapshot && bufferedCliEvents !== null) { - bufferedCliEvents.push(msg); - return; - } - handleEventMessage(msg.sessionId, msg.parentSessionId, msg.event, msg.data); - }); - const offSystem = config.userWebConnection.onSystemEvent(msg => { + const offCli = config.userWebConnection.onCliEvent( + config.kiloSessionId, + (msg) => { + const normalized = normalizeCliEvent(msg.event, msg.data); + const shouldBufferForSnapshot = + normalized && + (isChatEvent(normalized) || + normalized.type === "session.created" || + normalized.type === "session.updated"); + if (shouldBufferForSnapshot && bufferedCliEvents !== null) { + bufferedCliEvents.push(msg); + return; + } + handleEventMessage( + msg.sessionId, + msg.parentSessionId, + msg.event, + msg.data, + ); + }, + ); + const offSystem = config.userWebConnection.onSystemEvent((msg) => { handleSystemMessage(msg.event, msg.data); }); const offReconnect = config.userWebConnection.onReconnect(() => { @@ -829,9 +929,8 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor discoverCommands(ownerConnectionId); } }); - const releaseSubscription = config.userWebConnection.subscribeToCliSession( - config.kiloSessionId - ); + const releaseSubscription = + config.userWebConnection.subscribeToCliSession(config.kiloSessionId); let released = false; cleanup = () => { if (released) return; @@ -861,6 +960,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 +968,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,22 +977,26 @@ 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; } } const parsed = parseCreateSessionResponse(result); if (!parsed.ok) { - throw new Error('Invalid create_session response'); + throw new Error("Invalid create_session response"); } return parsed.kiloSessionId; }, exitSession: async () => { - if (remoteCommandState.refresh === 'upgrade-required') { + if (remoteCommandState.refresh === "upgrade-required") { throw new Error( remoteCommandState.message ?? - 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.' + "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", ); } if (!canExitSession()) { @@ -901,20 +1005,20 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // `exit_cli` is the compatibility wire command name for // session-detach — the product-facing name is "session exit". - const result = await sendCommand('exit_cli', { protocolVersion: 1 }); + const result = await sendCommand("exit_cli", { protocolVersion: 1 }); if (!parseExitSessionResponse(result).ok) { - throw new Error('Invalid exit_cli response'); + throw new Error("Invalid exit_cli response"); } }, send: async (input: TransportSendInput) => { - if (input.payload.type === 'command') { + if (input.payload.type === "command") { const remoteModel = getRemoteModelFields(input); - return sendCommand('send_command', { + return sendCommand("send_command", { protocolVersion: 1, command: input.payload.command, arguments: input.payload.arguments, ...(input.messageId ? { messageID: input.messageId } : {}), - ...(remoteModel.kind === 'none' + ...(remoteModel.kind === "none" ? {} : { // `send_command` requires a structured model: never emit a @@ -924,71 +1028,75 @@ function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactor // `dispatchedKilocodeModelId` would emit for the // equivalent prompt wire. model: - remoteModel.kind === 'structured' + remoteModel.kind === "structured" ? remoteModel.model : { - providerID: 'kilo', - modelID: remoteModel.model.replace(/^kilo\//, ''), + providerID: "kilo", + modelID: remoteModel.model.replace(/^kilo\//, ""), }, - ...(remoteModel.variant ? { variant: remoteModel.variant } : {}), + ...(remoteModel.variant + ? { variant: remoteModel.variant } + : {}), }), }); } const payload = input.payload; const remoteModel = getRemoteModelFields(input); const parts: Array< - | { type: 'text'; text: string } + | { type: "text"; text: string } | { - type: 'file'; + type: "file"; mime: string; filename: string; url: string; } - > = [{ type: 'text', text: payload.prompt }]; + > = [{ type: "text", text: payload.prompt }]; if (input.attachmentParts && input.attachmentParts.length > 0) { for (const part of input.attachmentParts) { parts.push({ - type: 'file', + type: "file", mime: part.mime, filename: part.filename, url: part.url, }); } } - return sendCommand('send_message', { + return sendCommand("send_message", { sessionID: config.kiloSessionId, parts, ...(payload.mode ? { agent: payload.mode } : {}), - ...(remoteModel.kind === 'none' + ...(remoteModel.kind === "none" ? {} : { model: remoteModel.model, - ...(remoteModel.variant ? { variant: remoteModel.variant } : {}), + ...(remoteModel.variant + ? { variant: remoteModel.variant } + : {}), }), }); }, - interrupt: () => sendCommand('interrupt', {}), - answer: payload => - sendCommand('question_reply', { + interrupt: () => sendCommand("interrupt", {}), + answer: (payload) => + sendCommand("question_reply", { requestID: payload.requestId, answers: payload.answers, }), - reject: payload => - sendCommand('question_reject', { + reject: (payload) => + sendCommand("question_reject", { requestID: payload.requestId, }), - respondToPermission: payload => - sendCommand('permission_respond', { + respondToPermission: (payload) => + sendCommand("permission_respond", { requestID: payload.requestId, reply: payload.response, }), - acceptSuggestion: payload => - sendCommand('suggestion_accept', { + acceptSuggestion: (payload) => + sendCommand("suggestion_accept", { requestID: payload.requestId, index: payload.index, }), - dismissSuggestion: payload => - sendCommand('suggestion_dismiss', { + dismissSuggestion: (payload) => + sendCommand("suggestion_dismiss", { requestID: payload.requestId, }), diff --git a/packages/cloud-agent-sdk/src/session.test.ts b/packages/cloud-agent-sdk/src/session.test.ts index 97fb0d4dd1a..3bceaf78298 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/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/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 8bfd31904b5..1e10186d009 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -1,7 +1,7 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; // Mock cloudflare:workers before importing UserConnectionDO -vi.mock('cloudflare:workers', () => ({ +vi.mock("cloudflare:workers", () => ({ DurableObject: class { ctx: unknown; env: unknown; @@ -18,11 +18,11 @@ const sessionIngestMocks = vi.hoisted(() => ({ getSessionIngestDO: vi.fn(), })); -vi.mock('./SessionIngestDO', () => ({ +vi.mock("./SessionIngestDO", () => ({ getSessionIngestDO: sessionIngestMocks.getSessionIngestDO, })); -import { MAX_CATALOG_RESULT_BYTES, UserConnectionDO } from './UserConnectionDO'; +import { MAX_CATALOG_RESULT_BYTES, UserConnectionDO } from "./UserConnectionDO"; // --------------------------------------------------------------------------- // Mock WebSocket @@ -73,16 +73,20 @@ function makeStorageFake() { delete: vi.fn(async (key: string) => { store.delete(key); }), - list: vi.fn(async (opts?: { prefix?: string }): Promise> => { - const result = new Map(); - const prefix = opts?.prefix ?? ''; - for (const [key, value] of store) { - if (key.startsWith(prefix)) { - result.set(key, value as T); + list: vi.fn( + async (opts?: { + prefix?: string; + }): Promise> => { + const result = new Map(); + const prefix = opts?.prefix ?? ""; + for (const [key, value] of store) { + if (key.startsWith(prefix)) { + result.set(key, value as T); + } } - } - return result; - }), + return result; + }, + ), setAlarm: vi.fn(), }; } @@ -105,7 +109,7 @@ function createMockCtx() { return { getWebSockets(tag?: string): MockWS[] { if (!tag) return [...sockets]; - return sockets.filter(ws => ws._tags.includes(tag)); + return sockets.filter((ws) => ws._tags.includes(tag)); }, acceptWebSocket(ws: MockWS, tags: string[]) { ws._tags = tags; @@ -126,10 +130,10 @@ function createMockCtx() { /** Drain microtasks so waitUntil-scheduled async IIFEs settle. */ async function flushAsync(): Promise { - await new Promise(resolve => { + await new Promise((resolve) => { setTimeout(resolve, 0); }); - await new Promise(resolve => { + await new Promise((resolve) => { setTimeout(resolve, 0); }); } @@ -140,12 +144,14 @@ async function flushAsync(): Promise { function makeSession( id: string, - status = 'busy', - title = 'Test', + status = "busy", + title = "Test", parentSessionId?: string, - platform?: string + platform?: string, ) { - const base = platform ? { id, status, title, platform } : { id, status, title }; + const base = platform + ? { id, status, title, platform } + : { id, status, title }; return parentSessionId ? { ...base, parentSessionId } : base; } @@ -156,11 +162,11 @@ function parseSent(ws: MockWS, callIndex = 0): unknown { } function isRecord(v: unknown): v is Record { - return typeof v === 'object' && v !== null && !Array.isArray(v); + return typeof v === "object" && v !== null && !Array.isArray(v); } function allSent(ws: MockWS): Record[] { - return ws.send.mock.calls.map(c => { + return ws.send.mock.calls.map((c) => { const parsed: unknown = JSON.parse(String(c[0])); if (!isRecord(parsed)) { throw new Error(`Expected JSON object but got: ${String(c[0])}`); @@ -172,7 +178,7 @@ function allSent(ws: MockWS): Record[] { /** Extract the correlationId that was sent to CLI for a given command. */ function getCorrelationId(cliWs: MockWS, callIndex = 0): string { const msgs = allSent(cliWs); - const cmdMsgs = msgs.filter(m => m.type === 'command'); + const cmdMsgs = msgs.filter((m) => m.type === "command"); const msg = cmdMsgs[callIndex]; if (!msg) throw new Error(`No command call at index ${callIndex}`); return msg.id as string; @@ -186,52 +192,58 @@ function setup() { return { doInstance, ctx, mockCtx }; } -function connectWebSocket(doInstance: UserConnectionDO, connectionId: string): MockWS { +function connectWebSocket( + doInstance: UserConnectionDO, + connectionId: string, +): MockWS { const client = createMockWs(); const server = createMockWs(); vi.stubGlobal( - 'WebSocketPair', + "WebSocketPair", class { 0 = client; 1 = server; - } + }, ); vi.stubGlobal( - 'Response', + "Response", class { constructor(_body?: BodyInit | null, _init?: ResponseInit) {} - } + }, ); doInstance.fetch( new Request(`http://local/web?connectionId=${connectionId}`, { - headers: { Upgrade: 'websocket' }, - }) + headers: { Upgrade: "websocket" }, + }), ); return server; } -function connectCliSocket(doInstance: UserConnectionDO, connectionId: string): MockWS { +function connectCliSocket( + doInstance: UserConnectionDO, + connectionId: string, +): MockWS { const client = createMockWs(); const server = createMockWs(); vi.stubGlobal( - 'WebSocketPair', + "WebSocketPair", class { 0 = client; 1 = server; - } + }, ); vi.stubGlobal( - 'Response', + "Response", class { constructor(_body?: BodyInit | null, _init?: ResponseInit) {} - } + }, ); doInstance.fetch( new Request(`http://local/cli?connectionId=${connectionId}`, { - headers: { Upgrade: 'websocket' }, - }) + headers: { Upgrade: "websocket" }, + }), ); return server; } @@ -247,18 +259,18 @@ function addCliSocket( platform?: string; }> = [], instance?: { name: string; projectName: string; version?: string }, - kiloUserId?: string + kiloUserId?: string, ): MockWS { const attachment: { - role: 'cli'; + role: "cli"; connectionId: string; sessions: typeof sessions; instance?: typeof instance; kiloUserId?: string; - } = { role: 'cli', connectionId, sessions }; + } = { role: "cli", connectionId, sessions }; if (instance) attachment.instance = instance; if (kiloUserId) attachment.kiloUserId = kiloUserId; - const ws = createMockWs(['cli'], attachment); + const ws = createMockWs(["cli"], attachment); mockCtx.addSocket(ws); return ws; } @@ -266,11 +278,11 @@ function addCliSocket( /** Create a web WebSocket and add it to the context. */ function addWebSocket( mockCtx: ReturnType, - connectionId = 'web-1', - subscribedSessions: string[] = [] + connectionId = "web-1", + subscribedSessions: string[] = [], ): MockWS { - const attachment = { role: 'web' as const, connectionId, subscribedSessions }; - const ws = createMockWs(['web'], attachment); + const attachment = { role: "web" as const, connectionId, subscribedSessions }; + const ws = createMockWs(["web"], attachment); mockCtx.addSocket(ws); return ws; } @@ -289,12 +301,14 @@ function sendHeartbeat( protocolVersion?: string; capabilities?: { attachments?: boolean }; instance?: { name: string; projectName: string; version?: string }; - } = {} + } = {}, ) { const msg = JSON.stringify({ - type: 'heartbeat', + type: "heartbeat", sessions, - ...(options.protocolVersion ? { protocolVersion: options.protocolVersion } : {}), + ...(options.protocolVersion + ? { protocolVersion: options.protocolVersion } + : {}), ...(options.capabilities ? { capabilities: options.capabilities } : {}), ...(options.instance ? { instance: options.instance } : {}), }); @@ -302,20 +316,28 @@ function sendHeartbeat( } /** Send a subscribe from a web ws */ -function sendSubscribe(doInstance: UserConnectionDO, webWs: MockWS, sessionId: string) { - const msg = JSON.stringify({ type: 'subscribe', sessionId }); +function sendSubscribe( + doInstance: UserConnectionDO, + webWs: MockWS, + sessionId: string, +) { + const msg = JSON.stringify({ type: "subscribe", sessionId }); doInstance.webSocketMessage(webWs as never, msg); } /** Send an unsubscribe from a web ws */ -function sendUnsubscribe(doInstance: UserConnectionDO, webWs: MockWS, sessionId: string) { - const msg = JSON.stringify({ type: 'unsubscribe', sessionId }); +function sendUnsubscribe( + doInstance: UserConnectionDO, + webWs: MockWS, + sessionId: string, +) { + const msg = JSON.stringify({ type: "unsubscribe", sessionId }); doInstance.webSocketMessage(webWs as never, msg); } /** Send a viewer ping from a web ws */ function sendPing(doInstance: UserConnectionDO, webWs: MockWS, nonce: string) { - const msg = JSON.stringify({ type: 'ping', nonce }); + const msg = JSON.stringify({ type: "ping", nonce }); doInstance.webSocketMessage(webWs as never, msg); } @@ -323,68 +345,91 @@ function sendPing(doInstance: UserConnectionDO, webWs: MockWS, nonce: string) { 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; + }, ) { - const msg = JSON.stringify({ type: 'command', ...opts }); - doInstance.webSocketMessage(webWs as never, msg); + const msg = JSON.stringify({ type: "command", ...opts }); + return doInstance.webSocketMessage(webWs as never, msg); } /** Send a response from a CLI ws */ function sendCliResponse( doInstance: UserConnectionDO, cliWs: MockWS, - opts: { id: string; result?: unknown; error?: unknown } + opts: { id: string; result?: unknown; error?: unknown }, ) { - const msg = JSON.stringify({ type: 'response', ...opts }); - doInstance.webSocketMessage(cliWs as never, msg); + const msg = JSON.stringify({ type: "response", ...opts }); + return doInstance.webSocketMessage(cliWs as never, msg); } -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) { - throw new Error(`Result fixture does not serialize to ${targetBytes} bytes`); +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 + ) { + throw new Error( + `Result fixture does not serialize to ${targetBytes} bytes`, + ); } return result; } function createUtf8OversizedResult(): { padding: string } { - const framingBytes = JSON.stringify({ padding: '' }).length; + const framingBytes = JSON.stringify({ padding: "" }).length; const result = { - padding: 'é'.repeat(Math.floor((MAX_CATALOG_RESULT_BYTES - framingBytes) / 2) + 1), + padding: "é".repeat( + Math.floor((MAX_CATALOG_RESULT_BYTES - framingBytes) / 2) + 1, + ), }; if ( JSON.stringify(result).length >= MAX_CATALOG_RESULT_BYTES || - new TextEncoder().encode(JSON.stringify(result)).byteLength <= MAX_CATALOG_RESULT_BYTES + new TextEncoder().encode(JSON.stringify(result)).byteLength <= + MAX_CATALOG_RESULT_BYTES ) { - throw new Error('UTF-8 catalog fixture does not cross the byte-only boundary'); + throw new Error( + "UTF-8 catalog fixture does not cross the byte-only boundary", + ); } return result; } /** Trigger CLI disconnect (awaits attention reset before broadcast). */ async function disconnectCli(doInstance: UserConnectionDO, cliWs: MockWS) { - await doInstance.webSocketClose(cliWs as never, 0, '', false); + await doInstance.webSocketClose(cliWs as never, 0, "", false); } /** Trigger web disconnect */ function disconnectWeb(doInstance: UserConnectionDO, webWs: MockWS) { - void doInstance.webSocketClose(webWs as never, 0, '', false); + void doInstance.webSocketClose(webWs as never, 0, "", false); } // =========================================================================== // Tests // =========================================================================== -describe('UserConnectionDO', () => { +describe("UserConnectionDO", () => { beforeEach(() => { vi.restoreAllMocks(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockReset(); - sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockResolvedValue(undefined); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockResolvedValue( + undefined, + ); sessionIngestMocks.claimSessionReadyPush.mockReset(); sessionIngestMocks.getSessionIngestDO.mockReset(); sessionIngestMocks.getSessionIngestDO.mockReturnValue({ - resetAttentionStatusOnCliDisconnect: sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + resetAttentionStatusOnCliDisconnect: + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, claimSessionReadyPush: sessionIngestMocks.claimSessionReadyPush, }); }); @@ -393,73 +438,76 @@ describe('UserConnectionDO', () => { vi.unstubAllGlobals(); }); - describe('notifySessionEvent', () => { - it('broadcasts semantic session events to web sockets only', async () => { + describe("notifySessionEvent", () => { + it("broadcasts semantic session events to web sockets only", async () => { const { doInstance, mockCtx } = setup(); const webWs = addWebSocket(mockCtx); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); const session = { - source: 'v2' as const, - sessionId: 'ses_12345678901234567890123456', - createdAt: '2026-01-01T00:00:00.000Z', - updatedAt: '2026-01-01T00:00:01.000Z', - title: 'Test', - createdOnPlatform: 'web', + source: "v2" as const, + sessionId: "ses_12345678901234567890123456", + createdAt: "2026-01-01T00:00:00.000Z", + updatedAt: "2026-01-01T00:00:01.000Z", + title: "Test", + createdOnPlatform: "web", organizationId: null, gitUrl: null, gitBranch: null, parentSessionId: null, - status: 'idle' as const, + status: "idle" as const, statusUpdatedAt: null, }; const result = await doInstance.notifySessionEvent({ - type: 'session.created', - data: { source: 'v2', session, changedAt: session.updatedAt }, + type: "session.created", + data: { source: "v2", session, changedAt: session.updatedAt }, }); expect(result).toEqual({ delivered: 1 }); expect(parseSent(webWs)).toEqual({ - type: 'system', - event: 'session.created', - data: { source: 'v2', session, changedAt: session.updatedAt }, + type: "system", + event: "session.created", + data: { source: "v2", session, changedAt: session.updatedAt }, }); expect(cliWs.send).not.toHaveBeenCalled(); }); - it('rejects invalid session event payloads without broadcasting', async () => { + it("rejects invalid session event payloads without broadcasting", async () => { const { doInstance, mockCtx } = setup(); 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(); }); }); - describe('hasActiveCliSession', () => { - it('tracks whether a connected CLI heartbeat currently owns the session', async () => { + describe("hasActiveCliSession", () => { + it("tracks whether a connected CLI heartbeat currently owns the session", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - expect(doInstance.hasActiveCliSession('ses_1')).toBe(false); + expect(doInstance.hasActiveCliSession("ses_1")).toBe(false); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_1")]); - expect(doInstance.hasActiveCliSession('ses_1')).toBe(true); + expect(doInstance.hasActiveCliSession("ses_1")).toBe(true); mockCtx.removeSocket(cliWs); await disconnectCli(doInstance, cliWs); - expect(doInstance.hasActiveCliSession('ses_1')).toBe(false); + 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", () => { const { doInstance, mockCtx } = setup(); - addCliSocket(mockCtx, 'cli-1', [makeSession('ses_1')]); + addCliSocket(mockCtx, "cli-1", [makeSession("ses_1")]); - expect(doInstance.hasActiveCliSession('ses_1')).toBe(true); + expect(doInstance.hasActiveCliSession("ses_1")).toBe(true); }); }); @@ -467,13 +515,13 @@ describe('UserConnectionDO', () => { // Heartbeat processing // ------------------------------------------------------------------------- - describe('heartbeat processing', () => { - it('updates session ownership and persists attachment', () => { + describe("heartbeat processing", () => { + it("updates session ownership and persists attachment", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + addWebSocket(mockCtx, "web-1"); - const sessions = [makeSession('s1'), makeSession('s2')]; + const sessions = [makeSession("s1"), makeSession("s2")]; sendHeartbeat(doInstance, cliWs, sessions); // CLI attachment updated with sessions @@ -481,62 +529,71 @@ 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + addWebSocket(mockCtx, "web-1"); // First heartbeat: owns s1 and s2 - sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")]); // Second heartbeat: only s1 - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); // 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' }); + const webWs2 = addWebSocket(mockCtx, "web-2"); + sendCommand(doInstance, webWs2, { + id: "cmd-1", + command: "send_message", + sessionId: "s2", + }); const resp = parseSent(webWs2); expect(resp).toMatchObject({ - type: 'response', - id: 'cmd-1', - error: 'Session owner not found', + type: "response", + id: "cmd-1", + error: "Session owner not found", }); }); - it('fails an in-flight command when the session owner changes', () => { + it("fails an in-flight command when the session owner changes", () => { const { doInstance, mockCtx } = setup(); - const firstOwner = addCliSocket(mockCtx, 'cli-1'); - const nextOwner = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + 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, firstOwner, [makeSession("s1")]); sendHeartbeat(doInstance, nextOwner, []); firstOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(firstOwner); webWs.send.mockClear(); - sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, nextOwner, [makeSession("s1")]); // The owner-change heartbeat broadcasts sessions.heartbeat and also fires // the SESSION_OWNER_CHANGED error response for the in-flight command. The // test cares about the latter; find it by type+id. - expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ - type: 'response', - id: 'cmd-1', + expect( + allSent(webWs).find((m) => m.type === "response" && m.id === "cmd-1"), + ).toEqual({ + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); - sendCliResponse(doInstance, firstOwner, { id: correlationId, result: 'late' }); + 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,78 +601,78 @@ 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", () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1'); - const cli2 = addCliSocket(mockCtx, 'cli-2'); - addWebSocket(mockCtx, 'web-1'); + const cli1 = addCliSocket(mockCtx, "cli-1"); + const cli2 = addCliSocket(mockCtx, "cli-2"); + addWebSocket(mockCtx, "web-1"); // cli1 owns s1 - sendHeartbeat(doInstance, cli1, [makeSession('s1')]); + sendHeartbeat(doInstance, cli1, [makeSession("s1")]); // web subscribes to s1 — subscribe sent to cli1 (the current owner) - const webWs = mockCtx.sockets.find(s => s._tags.includes('web'))!; - sendSubscribe(doInstance, webWs, 's1'); + const webWs = mockCtx.sockets.find((s) => s._tags.includes("web"))!; + sendSubscribe(doInstance, webWs, "s1"); cli1.send.mockClear(); cli2.send.mockClear(); // cli2 now reports s1 — becomes new owner - sendHeartbeat(doInstance, cli2, [makeSession('s1')]); + sendHeartbeat(doInstance, cli2, [makeSession("s1")]); // cli2 should have received the replayed subscribe for s1 const cli2Msgs = allSent(cli2); - expect(cli2Msgs).toContainEqual({ type: 'subscribe', sessionId: 's1' }); + 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", () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1'); - const cli2 = addCliSocket(mockCtx, 'cli-2'); - const subWeb = addWebSocket(mockCtx, 'web-sub'); - const otherWeb = addWebSocket(mockCtx, 'web-other'); + const cli1 = addCliSocket(mockCtx, "cli-1"); + const cli2 = addCliSocket(mockCtx, "cli-2"); + const subWeb = addWebSocket(mockCtx, "web-sub"); + const otherWeb = addWebSocket(mockCtx, "web-other"); // cli1 owns s1, cli2 owns s2 - sendHeartbeat(doInstance, cli1, [makeSession('s1')]); - sendHeartbeat(doInstance, cli2, [makeSession('s2')]); + sendHeartbeat(doInstance, cli1, [makeSession("s1")]); + sendHeartbeat(doInstance, cli2, [makeSession("s2")]); // subWeb subscribes to s1, otherWeb subscribes to s2 (subscriptions are // irrelevant to broadcast delivery — both should still receive cli1's heartbeat) - sendSubscribe(doInstance, subWeb, 's1'); - sendSubscribe(doInstance, otherWeb, 's2'); + sendSubscribe(doInstance, subWeb, "s1"); + sendSubscribe(doInstance, otherWeb, "s2"); subWeb.send.mockClear(); otherWeb.send.mockClear(); // cli1 sends heartbeat — both viewers must receive it - sendHeartbeat(doInstance, cli1, [makeSession('s1')]); + sendHeartbeat(doInstance, cli1, [makeSession("s1")]); expect(subWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(subWeb)).toMatchObject({ - type: 'system', - event: 'sessions.heartbeat', - data: { connectionId: 'cli-1' }, + type: "system", + event: "sessions.heartbeat", + data: { connectionId: "cli-1" }, }); expect(otherWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(otherWeb)).toMatchObject({ - type: 'system', - event: 'sessions.heartbeat', - data: { connectionId: 'cli-1' }, + type: "system", + event: "sessions.heartbeat", + data: { connectionId: "cli-1" }, }); }); - it('delivers one heartbeat per web socket (delivery count equals ws count)', () => { + it("delivers one heartbeat per web socket (delivery count equals ws count)", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const web1 = addWebSocket(mockCtx, 'web-1'); - const web2 = addWebSocket(mockCtx, 'web-2'); - const web3 = addWebSocket(mockCtx, 'web-3'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const web1 = addWebSocket(mockCtx, "web-1"); + const web2 = addWebSocket(mockCtx, "web-2"); + const web3 = addWebSocket(mockCtx, "web-3"); // No subscriptions — the broadcast must still hit every web socket. web1.send.mockClear(); web2.send.mockClear(); web3.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); // One heartbeat frame per active web socket. expect(web1.send).toHaveBeenCalledTimes(1); @@ -624,58 +681,64 @@ 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', - data: { connectionId: 'cli-1', sessions: [{ id: 's1' }] }, + type: "system", + event: "sessions.heartbeat", + data: { connectionId: "cli-1", sessions: [{ id: "s1" }] }, }); } }); - it('forwards the CLI-reported protocolVersion to every web socket', () => { + it("forwards the CLI-reported protocolVersion to every web socket", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { protocolVersion: '1' }); - sendSubscribe(doInstance, webWs, 's1'); + 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', - event: 'sessions.heartbeat', - data: { connectionId: 'cli-1', protocolVersion: '1' }, + type: "system", + event: "sessions.heartbeat", + data: { connectionId: "cli-1", protocolVersion: "1" }, }); }); - it('omits protocolVersion for a legacy CLI that never reports one', () => { + it("omits protocolVersion for a legacy CLI that never reports one", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, webWs, 's1'); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendSubscribe(doInstance, webWs, "s1"); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); const sent = parseSent(webWs) as { data: Record }; - expect(sent.data).not.toHaveProperty('protocolVersion'); + 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)", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); // subWeb subscribed to s1, otherWeb is unrelated — both must learn s1 is gone. - const subWeb = addWebSocket(mockCtx, 'web-sub'); - const otherWeb = addWebSocket(mockCtx, 'web-other'); + const subWeb = addWebSocket(mockCtx, "web-sub"); + const otherWeb = addWebSocket(mockCtx, "web-other"); // cli1 owns s1 - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, subWeb, 's1'); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendSubscribe(doInstance, subWeb, "s1"); subWeb.send.mockClear(); otherWeb.send.mockClear(); @@ -684,50 +747,50 @@ describe('UserConnectionDO', () => { expect(subWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(subWeb)).toMatchObject({ - type: 'system', - event: 'sessions.heartbeat', - data: { connectionId: 'cli-1', sessions: [] }, + type: "system", + event: "sessions.heartbeat", + data: { connectionId: "cli-1", sessions: [] }, }); expect(otherWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(otherWeb)).toMatchObject({ - type: 'system', - event: 'sessions.heartbeat', - data: { connectionId: 'cli-1', sessions: [] }, + type: "system", + event: "sessions.heartbeat", + data: { connectionId: "cli-1", sessions: [] }, }); }); - it('delivers heartbeat to web sockets that are not subscribed to anything', () => { + it("delivers heartbeat to web sockets that are not subscribed to anything", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); // webWs has no subscriptions - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toMatchObject({ - type: 'system', - event: 'sessions.heartbeat', - data: { connectionId: 'cli-1', sessions: [{ id: 's1' }] }, + type: "system", + event: "sessions.heartbeat", + data: { connectionId: "cli-1", sessions: [{ id: "s1" }] }, }); }); - it('schedules stale alarm on heartbeat', () => { + it("schedules stale alarm on heartbeat", () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); expect(ctx.storage.setAlarm).toHaveBeenCalled(); }); - it('sends heartbeat_ack to CLI socket', () => { + it("sends heartbeat_ack to CLI socket", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); const msgs = allSent(cliWs); - expect(msgs).toContainEqual({ type: 'heartbeat_ack' }); + expect(msgs).toContainEqual({ type: "heartbeat_ack" }); }); // ----------------------------------------------------------------------- @@ -738,11 +801,11 @@ 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { + sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")], { capabilities: { attachments: true }, }); @@ -753,98 +816,103 @@ describe('UserConnectionDO', () => { } }); - it('projects the latest connection capabilities onto every sessions.heartbeat row', () => { + it("projects the latest connection capabilities onto every sessions.heartbeat row", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); // Establish ownership and subscription first so the heartbeat broadcast // is targeted at this web socket. - sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); - sendSubscribe(doInstance, webWs, 's1'); + sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")]); + sendSubscribe(doInstance, webWs, "s1"); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { + sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")], { capabilities: { attachments: true }, }); expect(parseSent(webWs)).toMatchObject({ data: { sessions: [ - { id: 's1', capabilities: { attachments: true } }, - { id: 's2', capabilities: { attachments: true } }, + { id: "s1", capabilities: { attachments: true } }, + { id: "s2", capabilities: { attachments: true } }, ], }, }); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")]); expect(parseSent(webWs)).toMatchObject({ - data: { sessions: [{ id: 's1' }, { id: 's2' }] }, + data: { sessions: [{ id: "s1" }, { id: "s2" }] }, }); - const legacyRows = (parseSent(webWs) as { data: { sessions: Record[] } }) - .data.sessions; - expect(legacyRows.every(row => !Object.hasOwn(row, 'capabilities'))).toBe(true); + const legacyRows = ( + parseSent(webWs) as { data: { sessions: Record[] } } + ).data.sessions; + expect( + legacyRows.every((row) => !Object.hasOwn(row, "capabilities")), + ).toBe(true); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { + sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")], { capabilities: { attachments: false }, }); expect(parseSent(webWs)).toMatchObject({ data: { sessions: [ - { id: 's1', capabilities: { attachments: false } }, - { id: 's2', capabilities: { attachments: false } }, + { id: "s1", capabilities: { attachments: false } }, + { id: "s2", capabilities: { attachments: false } }, ], }, }); }); - 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)", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); // First heartbeat with capabilities - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { capabilities: { attachments: true }, }); // Second heartbeat omits capabilities (CLI rollback / legacy) - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); const rows = doInstance.getActiveSessions(); expect(rows).toHaveLength(1); - expect(rows[0]).not.toHaveProperty('capabilities'); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { capabilities: { attachments: true }, }); - sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, "s1"); webWs.send.mockClear(); // Latest heartbeat omits capabilities. - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); const sent = parseSent(webWs) as { data: Record }; - expect(sent.data).not.toHaveProperty('capabilities'); + 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)", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + 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')], { + sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { capabilities: { attachments: false }, }); @@ -853,17 +921,17 @@ 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); // Legacy heartbeat — no capabilities field - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); const legacy = doInstance.getActiveSessions(); - expect(legacy[0]).not.toHaveProperty('capabilities'); + expect(legacy[0]).not.toHaveProperty("capabilities"); // Upgraded CLI starts advertising attachments=true - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { capabilities: { attachments: true }, }); @@ -871,27 +939,31 @@ 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { + 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')], { + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { + sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")], { capabilities: { attachments: false }, }); @@ -901,14 +973,14 @@ 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", () => { const { doInstance, mockCtx } = setup(); // Pre-existing attachment with capabilities — simulates a socket that // was accepted before the DO was evicted. - const cliWs = createMockWs(['cli'], { - role: 'cli', - connectionId: 'cli-hiber', - sessions: [makeSession('s1')], + const cliWs = createMockWs(["cli"], { + role: "cli", + connectionId: "cli-hiber", + sessions: [makeSession("s1")], capabilities: { attachments: true }, }); mockCtx.addSocket(cliWs); @@ -923,32 +995,32 @@ describe('UserConnectionDO', () => { // Stale connection eviction // ------------------------------------------------------------------------- - describe('stale connection eviction', () => { - it('closes stale connection after timeout', async () => { + describe("stale connection eviction", () => { + it("closes stale connection after timeout", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); // Send heartbeat to register the connection and set lastHeartbeatAt - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); // Fast-forward time so the connection appears stale - vi.spyOn(Date, 'now') + vi.spyOn(Date, "now") .mockReturnValueOnce(Date.now() + 31_000) // for ensureState check .mockReturnValue(Date.now() + 31_000); // for alarm's Date.now() await doInstance.alarm(); - expect(cliWs.close).toHaveBeenCalledWith(4408, 'heartbeat timeout'); + expect(cliWs.close).toHaveBeenCalledWith(4408, "heartbeat timeout"); }); - it('reschedules alarm if other live connections remain', async () => { + it("reschedules alarm if other live connections remain", async () => { const { doInstance, mockCtx, ctx } = setup(); - const staleCli = addCliSocket(mockCtx, 'stale-1'); - const freshCli = addCliSocket(mockCtx, 'fresh-1'); + const staleCli = addCliSocket(mockCtx, "stale-1"); + const freshCli = addCliSocket(mockCtx, "fresh-1"); // Both send heartbeats - sendHeartbeat(doInstance, staleCli, [makeSession('s1')]); - sendHeartbeat(doInstance, freshCli, [makeSession('s2')]); + sendHeartbeat(doInstance, staleCli, [makeSession("s1")]); + sendHeartbeat(doInstance, freshCli, [makeSession("s2")]); // Reset setAlarm call count ctx.storage.setAlarm.mockClear(); @@ -956,31 +1028,31 @@ describe('UserConnectionDO', () => { // Make stale-1 appear stale but fresh-1 stays fresh const now = Date.now(); const staleTime = now + 31_000; - vi.spyOn(Date, 'now').mockReturnValue(staleTime); + vi.spyOn(Date, "now").mockReturnValue(staleTime); // Manually set lastHeartbeatAt for fresh-1 to "just now" (staleTime) // by sending another heartbeat from fresh-1 - sendHeartbeat(doInstance, freshCli, [makeSession('s2')]); + sendHeartbeat(doInstance, freshCli, [makeSession("s2")]); ctx.storage.setAlarm.mockClear(); await doInstance.alarm(); // Stale one closed - expect(staleCli.close).toHaveBeenCalledWith(4408, 'heartbeat timeout'); + expect(staleCli.close).toHaveBeenCalledWith(4408, "heartbeat timeout"); // Fresh one alive expect(freshCli.close).not.toHaveBeenCalled(); // Alarm rescheduled because fresh-1 remains expect(ctx.storage.setAlarm).toHaveBeenCalled(); }); - it('does not evict connection with recent heartbeat', async () => { + it("does not evict connection with recent heartbeat", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); // Time is within timeout window - vi.spyOn(Date, 'now').mockReturnValue(Date.now() + 10_000); + vi.spyOn(Date, "now").mockReturnValue(Date.now() + 10_000); await doInstance.alarm(); @@ -992,106 +1064,123 @@ describe('UserConnectionDO', () => { // Subscribe / Unsubscribe // ------------------------------------------------------------------------- - describe('subscribe/unsubscribe', () => { - it('sends subscribe to owning CLI when web subscribes', () => { + describe("subscribe/unsubscribe", () => { + it("sends subscribe to owning CLI when web subscribes", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); // CLI owns s1 - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, "s1"); // CLI should receive subscribe expect(cliWs.send).toHaveBeenCalledTimes(1); - expect(parseSent(cliWs)).toEqual({ type: 'subscribe', sessionId: 's1' }); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Fix bug')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1", "busy", "Fix bug")]); webWs.send.mockClear(); - sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, "s1"); expect(parseSent(webWs)).toEqual({ - type: 'system', - event: 'sessions.list', - data: { sessions: [{ id: 's1', status: 'busy', title: 'Fix bug', connectionId: 'cli-1' }] }, + type: "system", + event: "sessions.list", + 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", () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1'); - const cli2 = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cli1 = addCliSocket(mockCtx, "cli-1"); + const cli2 = addCliSocket(mockCtx, "cli-2"); + const webWs = addWebSocket(mockCtx, "web-1"); // No heartbeat sent, so no owner for 's1' // Trigger ensureState via a harmless message first - sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, "s1"); // Both CLIs should receive subscribe expect(cli1.send).toHaveBeenCalled(); expect(cli2.send).toHaveBeenCalled(); - expect(parseSent(cli1)).toEqual({ type: 'subscribe', sessionId: 's1' }); - expect(parseSent(cli2)).toEqual({ type: 'subscribe', sessionId: 's1' }); + expect(parseSent(cli1)).toEqual({ type: "subscribe", sessionId: "s1" }); + expect(parseSent(cli2)).toEqual({ type: "subscribe", sessionId: "s1" }); }); - it('duplicate subscribe is idempotent for attachment', () => { + it("duplicate subscribe is idempotent for attachment", () => { const { doInstance, mockCtx } = setup(); - addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendSubscribe(doInstance, webWs, 's1'); - sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, "s1"); - const att = webWs.deserializeAttachment() as { subscribedSessions: string[] }; - expect(att.subscribedSessions).toEqual(['s1']); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, "s1"); cliWs.send.mockClear(); - sendUnsubscribe(doInstance, webWs, 's1'); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const web1 = addWebSocket(mockCtx, 'web-1'); - const web2 = addWebSocket(mockCtx, 'web-2'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const web1 = addWebSocket(mockCtx, "web-1"); + const web2 = addWebSocket(mockCtx, "web-2"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendSubscribe(doInstance, web1, 's1'); - sendSubscribe(doInstance, web2, 's1'); + sendSubscribe(doInstance, web1, "s1"); + sendSubscribe(doInstance, web2, "s1"); cliWs.send.mockClear(); // Unsubscribe first — CLI should NOT get unsubscribe - sendUnsubscribe(doInstance, web1, 's1'); + sendUnsubscribe(doInstance, web1, "s1"); expect(cliWs.send).not.toHaveBeenCalled(); // Unsubscribe second — CLI SHOULD get unsubscribe - sendUnsubscribe(doInstance, web2, 's1'); + 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", + }); }); }); @@ -1099,114 +1188,119 @@ describe('UserConnectionDO', () => { // Viewer liveness // ------------------------------------------------------------------------- - describe('viewer liveness', () => { - it('replies to a viewer ping with the matching nonce only', () => { + describe("viewer liveness", () => { + it("replies to a viewer ping with the matching nonce only", () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'viewer-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "viewer-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); ctx.storage.setAlarm.mockClear(); cliWs.send.mockClear(); webWs.send.mockClear(); - sendPing(doInstance, webWs, 'nonce-1'); + sendPing(doInstance, webWs, "nonce-1"); expect(webWs.send).toHaveBeenCalledTimes(1); - expect(parseSent(webWs)).toEqual({ type: 'pong', nonce: 'nonce-1' }); + expect(parseSent(webWs)).toEqual({ type: "pong", nonce: "nonce-1" }); expect(doInstance.getActiveSessions()).toEqual([ - { id: 's1', status: 'busy', title: 'Test', connectionId: 'cli-1' }, + { id: "s1", status: "busy", title: "Test", connectionId: "cli-1" }, ]); expect(cliWs.send).not.toHaveBeenCalled(); expect(ctx.storage.setAlarm).not.toHaveBeenCalled(); expect(webWs.deserializeAttachment()).toEqual({ - role: 'web', - connectionId: 'viewer-1', + role: "web", + connectionId: "viewer-1", subscribedSessions: [], }); }); }); - describe('viewer connection identity', () => { - it('replaces an older web viewer with the same connectionId and broadcasts only to its replacement', async () => { + describe("viewer connection identity", () => { + it("replaces an older web viewer with the same connectionId and broadcasts only to its replacement", async () => { const { doInstance, mockCtx } = setup(); - const oldWeb = connectWebSocket(doInstance, 'viewer-1'); + const oldWeb = connectWebSocket(doInstance, "viewer-1"); oldWeb.send.mockClear(); - const newWeb = connectWebSocket(doInstance, 'viewer-1'); + const newWeb = connectWebSocket(doInstance, "viewer-1"); newWeb.send.mockClear(); - expect(oldWeb.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); + expect(oldWeb.close).toHaveBeenCalledWith(1000, "replaced by reconnect"); await doInstance.notifySessionEvent({ - type: 'session.deleted', + type: "session.deleted", data: { - source: 'v2', - sessionId: 's1', + source: "v2", + sessionId: "s1", parentSessionId: null, organizationId: null, gitUrl: null, gitBranch: null, - createdOnPlatform: 'web', - deletedAt: '2026-01-01T00:00:02.000Z', + createdOnPlatform: "web", + deletedAt: "2026-01-01T00:00:02.000Z", }, }); expect(oldWeb.send).not.toHaveBeenCalled(); expect(newWeb.send).toHaveBeenCalledTimes(1); - expect(mockCtx.sockets.filter(socket => socket._tags.includes('web'))).toHaveLength(2); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + const cliWs = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - const oldWeb = connectWebSocket(doInstance, 'viewer-1'); - sendSubscribe(doInstance, oldWeb, 's1'); + const oldWeb = connectWebSocket(doInstance, "viewer-1"); + sendSubscribe(doInstance, oldWeb, "s1"); cliWs.send.mockClear(); - const newWeb = connectWebSocket(doInstance, 'viewer-1'); + 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', + role: "web", + connectionId: "viewer-1", subscribedSessions: [], }); }); - it('ignores messages from a viewer that has been replaced', () => { + it("ignores messages from a viewer that has been replaced", () => { const { doInstance } = setup(); - const oldWeb = connectWebSocket(doInstance, 'viewer-1'); - connectWebSocket(doInstance, 'viewer-1'); + const oldWeb = connectWebSocket(doInstance, "viewer-1"); + connectWebSocket(doInstance, "viewer-1"); oldWeb.send.mockClear(); - sendPing(doInstance, oldWeb, 'stale-ping'); + sendPing(doInstance, oldWeb, "stale-ping"); expect(oldWeb.send).not.toHaveBeenCalled(); }); - it('keeps distinct viewer identities connected for independent broadcasts', async () => { + it("keeps distinct viewer identities connected for independent broadcasts", async () => { const { doInstance } = setup(); - const firstWeb = connectWebSocket(doInstance, 'viewer-1'); - const secondWeb = connectWebSocket(doInstance, 'viewer-2'); + const firstWeb = connectWebSocket(doInstance, "viewer-1"); + const secondWeb = connectWebSocket(doInstance, "viewer-2"); firstWeb.send.mockClear(); secondWeb.send.mockClear(); await doInstance.notifySessionEvent({ - type: 'session.deleted', + type: "session.deleted", data: { - source: 'v2', - sessionId: 's1', + source: "v2", + sessionId: "s1", parentSessionId: null, organizationId: null, gitUrl: null, gitBranch: null, - createdOnPlatform: 'web', - deletedAt: '2026-01-01T00:00:02.000Z', + createdOnPlatform: "web", + deletedAt: "2026-01-01T00:00:02.000Z", }, }); @@ -1215,11 +1309,11 @@ 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'shared-id'); + const cliWs = addCliSocket(mockCtx, "shared-id"); - connectWebSocket(doInstance, 'shared-id'); + connectWebSocket(doInstance, "shared-id"); expect(cliWs.close).not.toHaveBeenCalled(); }); @@ -1229,13 +1323,13 @@ describe('UserConnectionDO', () => { // CLI disconnect // ------------------------------------------------------------------------- - describe('CLI disconnect', () => { - it('cleans up session ownership and broadcasts cli.disconnected', async () => { + describe("CLI disconnect", () => { + it("cleans up session ownership and broadcasts cli.disconnected", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); webWs.send.mockClear(); // Remove from sockets before disconnect (simulates runtime closing) @@ -1246,29 +1340,41 @@ describe('UserConnectionDO', () => { expect(webWs.send).toHaveBeenCalled(); const msgs = allSent(webWs); const disconnectMsg = msgs.find( - (m: Record) => m.type === 'system' && m.event === 'cli.disconnected' + (m: Record) => + m.type === "system" && m.event === "cli.disconnected", ); expect(disconnectMsg).toEqual({ - type: 'system', - event: 'cli.disconnected', - data: { connectionId: 'cli-1' }, + type: "system", + event: "cli.disconnected", + data: { connectionId: "cli-1" }, }); // 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' }); + 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", + }); }); - it('sends error responses for pending commands on disconnect', async () => { + it("sends error responses for pending commands on disconnect", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); // Send command from web - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + sendCommand(doInstance, webWs, { + id: "cmd-1", + command: "send_message", + sessionId: "s1", + }); webWs.send.mockClear(); // CLI disconnects @@ -1278,22 +1384,27 @@ describe('UserConnectionDO', () => { // Web receives error response with original id const msgs = allSent(webWs); const errorResp = msgs.find( - (m: Record) => m.type === 'response' && m.id === 'cmd-1' + (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 () => { + it("reports owner change when an owner-fenced command target disconnects", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); webWs.send.mockClear(); @@ -1301,57 +1412,60 @@ describe('UserConnectionDO', () => { await disconnectCli(doInstance, cliWs); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); }); - it('fails pending commands as soon as their target socket is replaced', () => { + it("fails pending commands as soon as their target socket is replaced", () => { const { doInstance, mockCtx } = setup(); - const firstCli = connectCliSocket(doInstance, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const firstCli = connectCliSocket(doInstance, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, firstCli, [makeSession('s1')]); + sendHeartbeat(doInstance, firstCli, [makeSession("s1")]); firstCli.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); webWs.send.mockClear(); - connectCliSocket(doInstance, 'cli-1'); + connectCliSocket(doInstance, "cli-1"); - expect(firstCli.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); + expect(firstCli.close).toHaveBeenCalledWith( + 1000, + "replaced by reconnect", + ); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); }); - it('sends error for connection-routed pending commands on CLI disconnect', async () => { + it("sends error for connection-routed pending commands on CLI disconnect", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); sendHeartbeat(doInstance, cliWs, []); // Send command routed by connectionId (no sessionId) sendCommand(doInstance, webWs, { - id: 'cmd-conn', - command: 'send_message', - connectionId: 'cli-1', + id: "cmd-conn", + command: "send_message", + connectionId: "cli-1", }); webWs.send.mockClear(); @@ -1361,24 +1475,28 @@ describe('UserConnectionDO', () => { const msgs = allSent(webWs); const errorResp = msgs.find( - (m: Record) => m.type === 'response' && m.id === 'cmd-conn' + (m: Record) => + m.type === "response" && m.id === "cmd-conn", ); expect(errorResp).toMatchObject({ - type: 'response', - id: 'cmd-conn', - error: 'CLI disconnected', + type: "response", + id: "cmd-conn", + error: "CLI disconnected", }); }); - it('sends error for fallback-routed pending commands on CLI disconnect', async () => { + it("sends error for fallback-routed pending commands on CLI disconnect", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); sendHeartbeat(doInstance, cliWs, []); // Send command with no sessionId or connectionId (fallback routing) - sendCommand(doInstance, webWs, { id: 'cmd-fallback', command: 'send_message' }); + sendCommand(doInstance, webWs, { + id: "cmd-fallback", + command: "send_message", + }); webWs.send.mockClear(); mockCtx.removeSocket(cliWs); @@ -1386,25 +1504,26 @@ describe('UserConnectionDO', () => { const msgs = allSent(webWs); const errorResp = msgs.find( - (m: Record) => m.type === 'response' && m.id === 'cmd-fallback' + (m: Record) => + m.type === "response" && m.id === "cmd-fallback", ); expect(errorResp).toMatchObject({ - type: 'response', - id: 'cmd-fallback', - error: 'CLI disconnected', + type: "response", + id: "cmd-fallback", + error: "CLI disconnected", }); }); - it('reconnecting CLI — old socket close does not destroy state', async () => { + it("reconnecting CLI — old socket close does not destroy state", async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cli1 = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cli1, [makeSession('s1')]); + sendHeartbeat(doInstance, cli1, [makeSession("s1")]); // CLI2 connects with same connectionId (simulates reconnect) - const cli2 = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cli2, [makeSession('s1')]); + const cli2 = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cli2, [makeSession("s1")]); // CLI1's close event fires (stale socket), but cli2 still holds the connectionId // DON'T remove cli2 from sockets — cli2 is the replacement @@ -1415,30 +1534,34 @@ describe('UserConnectionDO', () => { // State should NOT be cleaned up — cli2 is live // Verify by routing a command to s1 — should reach cli2 cli2.send.mockClear(); - sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, "s1"); expect(cli2.send).toHaveBeenCalled(); - expect(parseSent(cli2)).toEqual({ type: 'subscribe', sessionId: 's1' }); + expect(parseSent(cli2)).toEqual({ type: "subscribe", sessionId: "s1" }); }); - it('reconnecting CLI — commands sent to replacement socket are not spuriously failed', async () => { + it("reconnecting CLI — commands sent to replacement socket are not spuriously failed", async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cli1 = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cli1, [makeSession('s1')]); + sendHeartbeat(doInstance, cli1, [makeSession("s1")]); // cli2 connects with the same connectionId (reconnect). // In production, closeStaleSocket removes cli1 before cli2 is accepted. // Simulate that by removing cli1 from the socket list first. mockCtx.removeSocket(cli1); - const cli2 = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cli2, [makeSession('s1')]); + const cli2 = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cli2, [makeSession("s1")]); cli2.send.mockClear(); 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' }); + sendCommand(doInstance, webWs, { + id: "cmd-new", + command: "send_message", + sessionId: "s1", + }); expect(cli2.send).toHaveBeenCalled(); const correlationId = getCorrelationId(cli2); @@ -1449,36 +1572,40 @@ describe('UserConnectionDO', () => { // Web should NOT have received an error for cmd-new — it was sent to cli2, not cli1 const errorMsgs = allSent(webWs).filter( - m => m.type === 'response' && m.id === 'cmd-new' && m.error + (m) => m.type === "response" && m.id === "cmd-new" && m.error, ); expect(errorMsgs).toHaveLength(0); // cli2 responds successfully webWs.send.mockClear(); - sendCliResponse(doInstance, cli2, { id: correlationId, result: 'ok' }); + sendCliResponse(doInstance, cli2, { id: correlationId, result: "ok" }); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-new', - result: 'ok', + type: "response", + id: "cmd-new", + result: "ok", }); }); - it('reconnecting CLI — pending commands from old socket get error responses', async () => { + it("reconnecting CLI — pending commands from old socket get error responses", async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cli1 = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cli1, [makeSession('s1')]); + 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' }); + sendCommand(doInstance, webWs, { + id: "cmd-1", + command: "send_message", + sessionId: "s1", + }); webWs.send.mockClear(); // CLI2 connects with the same connectionId (reconnect) - const cli2 = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cli2, [makeSession('s1')]); + const cli2 = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cli2, [makeSession("s1")]); // cli1's close event fires — cmd-1 was sent on cli1's wire, cli2 never saw it mockCtx.removeSocket(cli1); @@ -1487,75 +1614,88 @@ describe('UserConnectionDO', () => { // Web should receive an error for the stranded command const msgs = allSent(webWs); const errorResp = msgs.find( - (m: Record) => m.type === 'response' && m.id === 'cmd-1' + (m: Record) => + m.type === "response" && m.id === "cmd-1", ); expect(errorResp).toMatchObject({ - type: 'response', - id: 'cmd-1', - error: 'CLI disconnected', + type: "response", + id: "cmd-1", + error: "CLI disconnected", }); }); - it('resets attention status for owned sessions before broadcasting cli.disconnected', async () => { + it("resets attention status for owned sessions before broadcasting cli.disconnected", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); + const webWs = addWebSocket(mockCtx, "web-1"); sendHeartbeat(doInstance, cliWs, [ - makeSession('s-question', 'question'), - makeSession('s-busy', 'busy'), + makeSession("s-question", "question"), + makeSession("s-busy", "busy"), ]); webWs.send.mockClear(); const callOrder: string[] = []; - sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockImplementation(async () => { - callOrder.push('reset'); - // Disconnect must not have been broadcast yet (ordering guarantee). - expect( - allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected') - ).toBe(false); - }); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockImplementation( + async () => { + callOrder.push("reset"); + // Disconnect must not have been broadcast yet (ordering guarantee). + expect( + allSent(webWs).some( + (m) => m.type === "system" && m.event === "cli.disconnected", + ), + ).toBe(false); + }, + ); // Leave the socket in getWebSockets() — matches workerd during webSocketClose. await disconnectCli(doInstance, cliWs); - callOrder.push('disconnect'); + callOrder.push("disconnect"); - expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith(expect.anything(), { - kiloUserId: 'usr_1', - sessionId: 's-question', - }); - expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith(expect.anything(), { - kiloUserId: 'usr_1', - sessionId: 's-busy', - }); - // Both owned sessions are delegated; attention-only filtering is on the metadata side. - expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledTimes(2); - expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( - 'usr_1', - 's-question' - ); - expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( - 'usr_1', - 's-busy' + expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith( + expect.anything(), + { + kiloUserId: "usr_1", + sessionId: "s-question", + }, ); - expect(callOrder.filter(step => step === 'reset')).toHaveLength(2); - expect(callOrder.at(-1)).toBe('disconnect'); - expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( - true + expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith( + expect.anything(), + { + kiloUserId: "usr_1", + sessionId: "s-busy", + }, ); + // Both owned sessions are delegated; attention-only filtering is on the metadata side. + expect( + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + ).toHaveBeenCalledTimes(2); + expect( + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + ).toHaveBeenCalledWith("usr_1", "s-question"); + expect( + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + ).toHaveBeenCalledWith("usr_1", "s-busy"); + expect(callOrder.filter((step) => step === "reset")).toHaveLength(2); + expect(callOrder.at(-1)).toBe("disconnect"); + expect( + allSent(webWs).some( + (m) => m.type === "system" && m.event === "cli.disconnected", + ), + ).toBe(true); }); - it('resets attention when the closing socket is still listed in getWebSockets (workerd)', async () => { + it("resets attention when the closing socket is still listed in getWebSockets (workerd)", async () => { // Production wrangler/workerd keeps the closing WebSocket in getWebSockets() // during webSocketClose. Matching connectionId without excluding self would // treat every disconnect as a stale reconnect and skip the attention reset. // Prior unit tests always called removeSocket first, so they never caught this. const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s-question', 'question')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s-question", "question")]); webWs.send.mockClear(); sessionIngestMocks.getSessionIngestDO.mockClear(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); @@ -1564,48 +1704,57 @@ describe('UserConnectionDO', () => { expect(mockCtx.sockets).toContain(cliWs); await disconnectCli(doInstance, cliWs); - expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledTimes(1); - expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( - 'usr_1', - 's-question' - ); - expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( - true - ); + expect( + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + ).toHaveBeenCalledTimes(1); + expect( + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + ).toHaveBeenCalledWith("usr_1", "s-question"); + expect( + allSent(webWs).some( + (m) => m.type === "system" && m.event === "cli.disconnected", + ), + ).toBe(true); }); - it('does not reset attention when kiloUserId is missing on the attachment', async () => { + it("does not reset attention when kiloUserId is missing on the attachment", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); - const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); + const warn = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); - sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'question')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1", "question")]); webWs.send.mockClear(); mockCtx.removeSocket(cliWs); await disconnectCli(doInstance, cliWs); expect(sessionIngestMocks.getSessionIngestDO).not.toHaveBeenCalled(); - expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); + expect( + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + ).not.toHaveBeenCalled(); expect(warn).toHaveBeenCalledWith( - 'Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment', - { ownedSessionCount: 1 } - ); - expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( - true + "Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment", + { ownedSessionCount: 1 }, ); + expect( + allSent(webWs).some( + (m) => m.type === "system" && m.event === "cli.disconnected", + ), + ).toBe(true); }); - it('does not reset attention for sessions owned by another live connection', async () => { + it("does not reset attention for sessions owned by another live connection", async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); - const cli2 = addCliSocket(mockCtx, 'cli-2', [], undefined, 'usr_1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cli1 = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); + const cli2 = addCliSocket(mockCtx, "cli-2", [], undefined, "usr_1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cli1, [makeSession('s1', 'question')]); + sendHeartbeat(doInstance, cli1, [makeSession("s1", "question")]); // cli2 takes ownership of s1 - sendHeartbeat(doInstance, cli2, [makeSession('s1', 'question')]); + sendHeartbeat(doInstance, cli2, [makeSession("s1", "question")]); webWs.send.mockClear(); sessionIngestMocks.getSessionIngestDO.mockClear(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); @@ -1614,19 +1763,23 @@ describe('UserConnectionDO', () => { await disconnectCli(doInstance, cli1); // cli1 no longer owns s1, so no reset for that session - expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); - expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( - true - ); + expect( + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + ).not.toHaveBeenCalled(); + expect( + allSent(webWs).some( + (m) => m.type === "system" && m.event === "cli.disconnected", + ), + ).toBe(true); }); - it('stale reconnect close does not reset attention status', async () => { + it("stale reconnect close does not reset attention status", async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); - const cli2 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const cli1 = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); + const cli2 = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); - sendHeartbeat(doInstance, cli1, [makeSession('s1', 'question')]); - sendHeartbeat(doInstance, cli2, [makeSession('s1', 'question')]); + sendHeartbeat(doInstance, cli1, [makeSession("s1", "question")]); + sendHeartbeat(doInstance, cli2, [makeSession("s1", "question")]); sessionIngestMocks.getSessionIngestDO.mockClear(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); @@ -1634,30 +1787,36 @@ describe('UserConnectionDO', () => { await disconnectCli(doInstance, cli1); expect(sessionIngestMocks.getSessionIngestDO).not.toHaveBeenCalled(); - expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); + expect( + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + ).not.toHaveBeenCalled(); }); - it('still broadcasts cli.disconnected when attention reset fails', async () => { + it("still broadcasts cli.disconnected when attention reset fails", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); - const webWs = addWebSocket(mockCtx, 'web-1'); - const error = vi.spyOn(console, 'error').mockImplementation(() => undefined); + const cliWs = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); + const webWs = addWebSocket(mockCtx, "web-1"); + const error = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); - sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'question')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1", "question")]); webWs.send.mockClear(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockRejectedValueOnce( - new Error('db down') + new Error("db down"), ); // Leave socket listed (workerd close semantics). await disconnectCli(doInstance, cliWs); - expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( - true - ); + expect( + allSent(webWs).some( + (m) => m.type === "system" && m.event === "cli.disconnected", + ), + ).toBe(true); expect(error).toHaveBeenCalledWith( - 'Failed to reset attention status on CLI disconnect', - expect.objectContaining({ error: 'db down' }) + "Failed to reset attention status on CLI disconnect", + expect.objectContaining({ error: "db down" }), ); }); }); @@ -1666,37 +1825,37 @@ describe('UserConnectionDO', () => { // Web disconnect // ------------------------------------------------------------------------- - describe('web disconnect', () => { - it('removes from all subscription sets', () => { + describe("web disconnect", () => { + it("removes from all subscription sets", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); - sendSubscribe(doInstance, webWs, 's1'); - sendSubscribe(doInstance, webWs, 's2'); + sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")]); + sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, "s2"); mockCtx.removeSocket(webWs); disconnectWeb(doInstance, webWs); // Verify: CLI events for s1 and s2 go nowhere (no crash) const cliEventMsg = JSON.stringify({ - type: 'event', - sessionId: 's1', - event: 'message.updated', + type: "event", + sessionId: "s1", + event: "message.updated", data: {}, }); doInstance.webSocketMessage(cliWs as never, cliEventMsg); // 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, webWs, 's1'); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendSubscribe(doInstance, webWs, "s1"); cliWs.send.mockClear(); mockCtx.removeSocket(webWs); @@ -1704,25 +1863,31 @@ describe('UserConnectionDO', () => { // CLI should get unsubscribe for s1 const msgs = allSent(cliWs); - const unsub = msgs.find((m: Record) => m.type === 'unsubscribe'); - expect(unsub).toEqual({ type: 'unsubscribe', sessionId: 's1' }); + const unsub = msgs.find( + (m: Record) => m.type === "unsubscribe", + ); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + 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' }); + sendCliResponse(doInstance, cliWs, { id: correlationId, result: "ok" }); }); }); @@ -1730,294 +1895,337 @@ describe('UserConnectionDO', () => { // Command routing // ------------------------------------------------------------------------- - describe('command routing', () => { - it('routes web command to correct CLI by sessionId', () => { + describe("command routing", () => { + it("routes web command to correct CLI by sessionId", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_message', - sessionId: 's1', - data: { text: 'hello' }, + id: "cmd-1", + command: "send_message", + sessionId: "s1", + data: { text: "hello" }, }); expect(cliWs.send).toHaveBeenCalledTimes(1); const sent = parseSent(cliWs) as Record; expect(sent).toMatchObject({ - type: 'command', - command: 'send_message', - sessionId: 's1', - data: { text: 'hello' }, + type: "command", + command: "send_message", + sessionId: "s1", + data: { text: "hello" }, }); - expect(typeof sent.id).toBe('string'); - expect(sent.id).not.toBe('cmd-1'); + expect(typeof sent.id).toBe("string"); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + 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 } }); + sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { success: true }, + }); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", result: { success: true }, }); }); - it('sanitizes a relay-shaped CLI error before forwarding it to web', () => { + it("sanitizes a relay-shaped CLI error before forwarding it to web", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + sendCommand(doInstance, webWs, { + id: "cmd-1", + command: "send_message", + sessionId: "s1", + }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'cli', - message: 'Command failed', + source: "cli", + message: "Command failed", }, }); }); - it('preserves CLI string errors for old-CLI compatibility', () => { + it("preserves CLI string errors for old-CLI compatibility", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1' }); + sendCommand(doInstance, webWs, { + id: "cmd-1", + command: "list_models", + sessionId: "s1", + }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'unknown command: list_models', + error: "unknown command: list_models", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', - error: 'unknown command: list_models', + type: "response", + id: "cmd-1", + error: "unknown command: list_models", }); }); - it('accepts a pending response only from the targeted CLI socket', () => { + it("accepts a pending response only from the targeted CLI socket", () => { const { doInstance, mockCtx } = setup(); - const targetCli = addCliSocket(mockCtx, 'cli-1'); - const otherCli = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const targetCli = addCliSocket(mockCtx, "cli-1"); + const otherCli = addCliSocket(mockCtx, "cli-2"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, targetCli, [makeSession('s1')]); + sendHeartbeat(doInstance, targetCli, [makeSession("s1")]); sendHeartbeat(doInstance, otherCli, []); targetCli.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + 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' }); + 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' }); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); sendCommand(doInstance, webWs, { - id: 'cmd-2', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-2", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-2', + type: "response", + id: "cmd-2", error: { - source: 'relay', - code: 'CATALOG_REQUEST_PENDING', - message: 'Model catalog request already pending', + source: "relay", + code: "CATALOG_REQUEST_PENDING", + message: "Model catalog request already pending", }, }); - expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(1); + 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", () => { const now = 1_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); + vi.spyOn(Date, "now").mockReturnValue(now); const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); vi.mocked(Date.now).mockReturnValue(now + 35_001); sendCommand(doInstance, webWs, { - id: 'cmd-2', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-2", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'COMMAND_EXPIRED', - message: 'Command expired', + source: "relay", + code: "COMMAND_EXPIRED", + message: "Command expired", }, }); - expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(2); + 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", () => { const now = 1_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); + vi.spyOn(Date, "now").mockReturnValue(now); const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendCommand(doInstance, webWs, { + id: "cmd-1", + command: "send_message", + sessionId: "s1", + }); ctx.storage.setAlarm.mockClear(); vi.mocked(Date.now).mockReturnValue(now + 20_000); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); expect(ctx.storage.setAlarm).toHaveBeenCalledWith(now + 35_000); }); - it('expires pending commands during alarm processing', async () => { + it("expires pending commands during alarm processing", async () => { const now = 1_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); + vi.spyOn(Date, "now").mockReturnValue(now); const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1' }); + sendCommand(doInstance, webWs, { + id: "cmd-1", + command: "send_message", + sessionId: "s1", + }); const correlationId = getCorrelationId(cliWs); vi.mocked(Date.now).mockReturnValue(now + 34_000); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); webWs.send.mockClear(); vi.mocked(Date.now).mockReturnValue(now + 35_001); await doInstance.alarm(); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'COMMAND_EXPIRED', - message: 'Command expired', + source: "relay", + code: "COMMAND_EXPIRED", + message: "Command expired", }, }); - sendCliResponse(doInstance, cliWs, { id: correlationId, result: 'late' }); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); for (let index = 0; index < 128; index++) { sendCommand(doInstance, webWs, { id: `cmd-${index}`, - command: 'send_message', - sessionId: 's1', + command: "send_message", + sessionId: "s1", }); } sendCommand(doInstance, webWs, { - id: 'cmd-over-cap', - command: 'send_message', - sessionId: 's1', + id: "cmd-over-cap", + command: "send_message", + sessionId: "s1", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-over-cap', + type: "response", + id: "cmd-over-cap", error: { - source: 'relay', - code: 'PENDING_COMMAND_LIMIT', - message: 'Too many pending commands', + source: "relay", + code: "PENDING_COMMAND_LIMIT", + message: "Too many pending commands", }, }); - expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(128); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2025,21 +2233,25 @@ describe('UserConnectionDO', () => { 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2050,28 +2262,28 @@ describe('UserConnectionDO', () => { }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'CATALOG_TOO_LARGE', - message: 'Model catalog response is too large', + source: "relay", + code: "CATALOG_TOO_LARGE", + message: "Model catalog response is too large", }, }); }); - it('rejects a multibyte list_models result over 512 KiB', () => { + it("rejects a multibyte list_models result over 512 KiB", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2082,71 +2294,71 @@ describe('UserConnectionDO', () => { }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'CATALOG_TOO_LARGE', - message: 'Model catalog response is too large', + source: "relay", + code: "CATALOG_TOO_LARGE", + message: "Model catalog response is too large", }, }); }); - it('returns error when CLI not found for session', () => { + it("returns error when CLI not found for session", () => { const { doInstance, mockCtx } = setup(); - const webWs = addWebSocket(mockCtx, 'web-1'); + const webWs = addWebSocket(mockCtx, "web-1"); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_message', - sessionId: 'unknown-session', + id: "cmd-1", + command: "send_message", + sessionId: "unknown-session", }); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', - error: 'Session owner not found', + type: "response", + id: "cmd-1", + error: "Session owner not found", }); }); - it('rejects a stale expected session owner without forwarding', () => { + it("rejects a stale expected session owner without forwarding", () => { const { doInstance, mockCtx } = setup(); - const currentOwner = addCliSocket(mockCtx, 'cli-1'); - const staleOwner = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const currentOwner = addCliSocket(mockCtx, "cli-1"); + const staleOwner = addCliSocket(mockCtx, "cli-2"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, currentOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, currentOwner, [makeSession("s1")]); sendHeartbeat(doInstance, staleOwner, []); currentOwner.send.mockClear(); staleOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_message', - sessionId: 's1', - connectionId: 'cli-2', + id: "cmd-1", + command: "send_message", + sessionId: "s1", + connectionId: "cli-2", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); expect(currentOwner.send).not.toHaveBeenCalled(); expect(staleOwner.send).not.toHaveBeenCalled(); }); - it('routes command by connectionId to specific CLI', () => { + it("routes command by connectionId to specific CLI", () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, 'cli-1'); - const cli2 = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cli1 = addCliSocket(mockCtx, "cli-1"); + const cli2 = addCliSocket(mockCtx, "cli-2"); + const webWs = addWebSocket(mockCtx, "web-1"); // Trigger ensureState sendHeartbeat(doInstance, cli1, []); @@ -2155,29 +2367,37 @@ describe('UserConnectionDO', () => { cli2.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_message', - connectionId: 'cli-2', + id: "cmd-1", + command: "send_message", + connectionId: "cli-2", }); expect(cli1.send).not.toHaveBeenCalled(); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const web1 = addWebSocket(mockCtx, 'web-1'); - const web2 = addWebSocket(mockCtx, 'web-2'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const web1 = addWebSocket(mockCtx, "web-1"); + const web2 = addWebSocket(mockCtx, "web-2"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); // Both web sockets send commands with the same id - sendCommand(doInstance, web1, { id: 'dup-id', command: 'send_message', sessionId: 's1' }); + 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' }); + sendCommand(doInstance, web2, { + id: "dup-id", + command: "send_message", + sessionId: "s1", + }); const corr2 = getCorrelationId(cliWs, 1); expect(corr1).not.toBe(corr2); @@ -2185,22 +2405,30 @@ describe('UserConnectionDO', () => { web1.send.mockClear(); web2.send.mockClear(); - sendCliResponse(doInstance, cliWs, { id: corr1, result: 'result-1' }); - sendCliResponse(doInstance, cliWs, { id: corr2, result: 'result-2' }); + sendCliResponse(doInstance, cliWs, { id: corr1, result: "result-1" }); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message' }); + sendCommand(doInstance, webWs, { id: "cmd-1", command: "send_message" }); expect(cliWs.send).toHaveBeenCalledTimes(1); }); }); @@ -2209,168 +2437,183 @@ describe('UserConnectionDO', () => { // Command allowlist // ------------------------------------------------------------------------- - describe('command allowlist', () => { + describe("command allowlist", () => { const ALLOWED = [ - 'send_message', - 'interrupt', - 'question_reply', - 'question_reject', - 'permission_respond', - 'suggestion_accept', - 'suggestion_dismiss', - 'list_models', - 'list_commands', - 'send_command', - 'create_session', - 'exit_cli', + "send_message", + "interrupt", + "question_reply", + "question_reject", + "permission_respond", + "suggestion_accept", + "suggestion_dismiss", + "list_models", + "list_commands", + "send_command", + "create_session", + "exit_cli", ]; - it('forwards every allowed viewer command to the owning CLI', () => { + it("forwards every allowed viewer command to the owning CLI", () => { for (const command of ALLOWED) { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', + id: "cmd-1", command, - sessionId: 's1', - data: command === 'exit_cli' ? { protocolVersion: 1 } : { hello: 'world' }, + sessionId: "s1", + data: + command === "exit_cli" + ? { protocolVersion: 1 } + : { hello: "world" }, }); expect(cliWs.send).toHaveBeenCalledTimes(1); const sent = parseSent(cliWs) as Record; - expect(sent).toMatchObject({ type: 'command', command, sessionId: 's1' }); - expect(typeof sent.id).toBe('string'); - expect(sent.id).not.toBe('cmd-1'); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'eval', sessionId: 's1' }); + sendCommand(doInstance, webWs, { + id: "cmd-1", + command: "eval", + sessionId: "s1", + }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'COMMAND_NOT_ALLOWED', - message: 'Command is not allowed', + source: "relay", + code: "COMMAND_NOT_ALLOWED", + message: "Command is not allowed", }, }); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'shell', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "shell", + sessionId: "s1", + connectionId: "cli-1", }); // No owner-fencing error — allowlist runs first. expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'COMMAND_NOT_ALLOWED', - message: 'Command is not allowed', + source: "relay", + code: "COMMAND_NOT_ALLOWED", + message: "Command is not allowed", }, }); 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", () => { const { doInstance, mockCtx } = setup(); - const webWs = addWebSocket(mockCtx, 'web-1'); + const webWs = addWebSocket(mockCtx, "web-1"); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'eval', - sessionId: 'unknown-session', + id: "cmd-1", + command: "eval", + sessionId: "unknown-session", }); // COMMAND_NOT_ALLOWED wins over "Session owner not found". expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'COMMAND_NOT_ALLOWED', - message: 'Command is not allowed', + source: "relay", + code: "COMMAND_NOT_ALLOWED", + message: "Command is not allowed", }, }); }); - it('does not allocate a pending entry or forward a disallowed command', () => { + it("does not allocate a pending entry or forward a disallowed command", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'eval', sessionId: 's1' }); + 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'); + expect(sent).toMatchObject({ type: "response", id: "cmd-1" }); + expect(sent.id).toBe("cmd-1"); // The CLI must not have received any command envelope, and the DO must // 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' }); + 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", () => { const { doInstance, mockCtx } = setup(); - const currentOwner = addCliSocket(mockCtx, 'cli-1'); - const staleOwner = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const currentOwner = addCliSocket(mockCtx, "cli-1"); + const staleOwner = addCliSocket(mockCtx, "cli-2"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, currentOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, currentOwner, [makeSession("s1")]); sendHeartbeat(doInstance, staleOwner, []); currentOwner.send.mockClear(); staleOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_message', - sessionId: 's1', - connectionId: 'cli-2', + id: "cmd-1", + command: "send_message", + sessionId: "s1", + connectionId: "cli-2", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); expect(currentOwner.send).not.toHaveBeenCalled(); @@ -2382,75 +2625,79 @@ describe('UserConnectionDO', () => { // list_commands dedupe and size cap // ------------------------------------------------------------------------- - describe('list_commands dedupe and size cap', () => { - it('rejects a duplicate in-flight list_commands request for the same viewer session and owner', () => { + describe("list_commands dedupe and size cap", () => { + it("rejects a duplicate in-flight list_commands request for the same viewer session and owner", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_commands', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_commands", + sessionId: "s1", + connectionId: "cli-1", }); sendCommand(doInstance, webWs, { - id: 'cmd-2', - command: 'list_commands', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-2", + command: "list_commands", + sessionId: "s1", + connectionId: "cli-1", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-2', + type: "response", + id: "cmd-2", error: { - source: 'relay', - code: 'CATALOG_REQUEST_PENDING', - message: 'Model catalog request already pending', + source: "relay", + code: "CATALOG_REQUEST_PENDING", + message: "Model catalog request already pending", }, }); - expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(1); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_models', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_models", + sessionId: "s1", + connectionId: "cli-1", }); sendCommand(doInstance, webWs, { - id: 'cmd-2', - command: 'list_commands', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-2", + command: "list_commands", + sessionId: "s1", + connectionId: "cli-1", }); - expect(allSent(cliWs).filter(message => message.type === 'command')).toHaveLength(2); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_commands', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_commands", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2458,21 +2705,25 @@ describe('UserConnectionDO', () => { 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_commands', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_commands", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2483,28 +2734,28 @@ describe('UserConnectionDO', () => { }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'CATALOG_TOO_LARGE', - message: 'Model catalog response is too large', + source: "relay", + code: "CATALOG_TOO_LARGE", + message: "Model catalog response is too large", }, }); }); - it('rejects a multibyte list_commands result over 512 KiB', () => { + it("rejects a multibyte list_commands result over 512 KiB", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_commands', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_commands", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2515,12 +2766,12 @@ describe('UserConnectionDO', () => { }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'CATALOG_TOO_LARGE', - message: 'Model catalog response is too large', + source: "relay", + code: "CATALOG_TOO_LARGE", + message: "Model catalog response is too large", }, }); }); @@ -2530,84 +2781,86 @@ describe('UserConnectionDO', () => { // Old CLI upgrade-required mapping // ------------------------------------------------------------------------- - describe('old CLI upgrade-required mapping', () => { + describe("old CLI upgrade-required mapping", () => { it('maps "unknown command: list_commands" to CLI_UPGRADE_REQUIRED with slash message', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_commands', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_commands", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'unknown command: list_commands', + error: "unknown command: list_commands", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'CLI_UPGRADE_REQUIRED', - message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', + source: "relay", + code: "CLI_UPGRADE_REQUIRED", + message: + "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", }, }); }); it('maps "unknown command: send_command" to CLI_UPGRADE_REQUIRED with slash message', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_command', - sessionId: 's1', - connectionId: 'cli-1', - data: { command: 'init' }, + id: "cmd-1", + command: "send_command", + sessionId: "s1", + connectionId: "cli-1", + data: { command: "init" }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'unknown command: send_command', + error: "unknown command: send_command", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'CLI_UPGRADE_REQUIRED', - message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', + source: "relay", + code: "CLI_UPGRADE_REQUIRED", + message: + "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", }, }); }); it('maps "unknown command: exit_cli" to CLI_UPGRADE_REQUIRED with slash message', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); const correlationId = getCorrelationId(cliWs); @@ -2615,145 +2868,150 @@ describe('UserConnectionDO', () => { sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'unknown command: exit_cli', + error: "unknown command: exit_cli", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'CLI_UPGRADE_REQUIRED', - message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', + source: "relay", + code: "CLI_UPGRADE_REQUIRED", + message: + "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", }, }); }); it('maps "unknown command: create_session" to CLI_UPGRADE_REQUIRED with create_session message', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'create_session', - connectionId: 'cli-1', - data: { title: 'New session' }, + id: "cmd-1", + command: "create_session", + connectionId: "cli-1", + data: { title: "New session" }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'unknown command: create_session', + error: "unknown command: create_session", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'CLI_UPGRADE_REQUIRED', + source: "relay", + code: "CLI_UPGRADE_REQUIRED", message: - 'Creating remote sessions from mobile requires a newer Kilo CLI. Update Kilo CLI and reconnect.', + "Creating remote sessions from mobile requires a newer Kilo CLI. Update Kilo CLI and reconnect.", }, }); }); it('preserves "unknown command: list_models" because list_models is not in the upgrade-required set', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1' }); + sendCommand(doInstance, webWs, { + id: "cmd-1", + command: "list_models", + sessionId: "s1", + }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'unknown command: list_models', + error: "unknown command: list_models", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', - error: 'unknown command: list_models', + type: "response", + id: "cmd-1", + error: "unknown command: list_models", }); }); - it('preserves an unrelated CLI string error for send_command', () => { + it("preserves an unrelated CLI string error for send_command", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_command', - sessionId: 's1', - data: { command: 'init' }, + id: "cmd-1", + command: "send_command", + sessionId: "s1", + data: { command: "init" }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'session not ready', + error: "session not ready", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', - error: 'session not ready', + type: "response", + id: "cmd-1", + error: "session not ready", }); }); it('does not match a longer error that merely starts with "unknown command: list_commands"', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'list_commands', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "list_commands", + sessionId: "s1", + connectionId: "cli-1", }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'unknown command: list_commands: try again', + error: "unknown command: list_commands: try again", }); // Exact-match only — do not misclassify longer error strings. expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', - error: 'unknown command: list_commands: try again', + type: "response", + id: "cmd-1", + error: "unknown command: list_commands: try again", }); }); - it('preserves a longer exit_cli unknown-command error', () => { + it("preserves a longer exit_cli unknown-command error", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); const correlationId = getCorrelationId(cliWs); @@ -2761,13 +3019,13 @@ describe('UserConnectionDO', () => { sendCliResponse(doInstance, cliWs, { id: correlationId, - error: 'unknown command: exit_cli: session not ready', + error: "unknown command: exit_cli: session not ready", }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', - error: 'unknown command: exit_cli: session not ready', + type: "response", + id: "cmd-1", + error: "unknown command: exit_cli: session not ready", }); }); }); @@ -2778,134 +3036,142 @@ describe('UserConnectionDO', () => { // and must NOT be subject to the 512 KiB catalog response cap.) // ------------------------------------------------------------------------- - describe('send_command / create_session negative coverage', () => { - it('forwards two in-flight same-owner/same-session send_command requests without deduping', () => { + describe("send_command / create_session negative coverage", () => { + it("forwards two in-flight same-owner/same-session send_command requests without deduping", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_command', - sessionId: 's1', - connectionId: 'cli-1', - data: { command: 'init' }, + id: "cmd-1", + command: "send_command", + sessionId: "s1", + connectionId: "cli-1", + data: { command: "init" }, }); sendCommand(doInstance, webWs, { - id: 'cmd-2', - command: 'send_command', - sessionId: 's1', - connectionId: 'cli-1', - data: { command: 'plan' }, + id: "cmd-2", + command: "send_command", + sessionId: "s1", + connectionId: "cli-1", + data: { command: "plan" }, }); - const cliCommands = allSent(cliWs).filter(message => message.type === 'command'); + const cliCommands = allSent(cliWs).filter( + (message) => message.type === "command", + ); expect(cliCommands).toHaveLength(2); expect(cliCommands[0]).toMatchObject({ - type: 'command', - command: 'send_command', - sessionId: 's1', - data: { command: 'init' }, + type: "command", + command: "send_command", + sessionId: "s1", + data: { command: "init" }, }); expect(cliCommands[1]).toMatchObject({ - type: 'command', - command: 'send_command', - sessionId: 's1', - data: { command: 'plan' }, + type: "command", + command: "send_command", + sessionId: "s1", + data: { command: "plan" }, }); expect(cliCommands[0].id).not.toBe(cliCommands[1].id); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'create_session', - connectionId: 'cli-1', - data: { title: 'First session' }, + id: "cmd-1", + command: "create_session", + connectionId: "cli-1", + data: { title: "First session" }, }); sendCommand(doInstance, webWs, { - id: 'cmd-2', - command: 'create_session', - connectionId: 'cli-1', - data: { title: 'Second session' }, + id: "cmd-2", + command: "create_session", + connectionId: "cli-1", + data: { title: "Second session" }, }); - const cliCommands = allSent(cliWs).filter(message => message.type === 'command'); + const cliCommands = allSent(cliWs).filter( + (message) => message.type === "command", + ); expect(cliCommands).toHaveLength(2); expect(cliCommands[0]).toMatchObject({ - type: 'command', - command: 'create_session', - data: { title: 'First session' }, + type: "command", + command: "create_session", + data: { title: "First session" }, }); expect(cliCommands[1]).toMatchObject({ - type: 'command', - command: 'create_session', - data: { title: 'Second session' }, + type: "command", + command: "create_session", + data: { title: "Second session" }, }); expect(cliCommands[0].id).not.toBe(cliCommands[1].id); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'send_command', - sessionId: 's1', - connectionId: 'cli-1', - data: { command: 'init' }, + id: "cmd-1", + command: "send_command", + sessionId: "s1", + connectionId: "cli-1", + data: { command: "init" }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); + const result = createResultWithSerializedBytes( + MAX_CATALOG_RESULT_BYTES + 1, + ); sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", result, }); }); - it('relays a create_session result over 512 KiB unchanged', () => { + it("relays a create_session result over 512 KiB unchanged", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'create_session', - connectionId: 'cli-1', - data: { title: 'Big session' }, + id: "cmd-1", + command: "create_session", + connectionId: "cli-1", + data: { title: "Big session" }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); + const result = createResultWithSerializedBytes( + MAX_CATALOG_RESULT_BYTES + 1, + ); sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", result, }); }); @@ -2915,191 +3181,210 @@ describe('UserConnectionDO', () => { // exit_cli routing and relay policy // ------------------------------------------------------------------------- - describe('exit_cli routing and relay policy', () => { + describe("exit_cli routing and relay policy", () => { 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: "missing sessionId", input: { data: { protocolVersion: 1 } } }, + { label: "missing data", input: { sessionId: "s1" } }, + { + 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: 'extra data field', - input: { sessionId: 's1', data: { protocolVersion: 1, extra: true } }, + label: "primitive data", + input: { sessionId: "s1", data: "protocolVersion=1" }, }, - { 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 }) => { + ])("rejects $label before routing or pending allocation", ({ input }) => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', + id: "cmd-1", + command: "exit_cli", ...input, }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'INVALID_COMMAND', - message: 'Invalid command', + source: "relay", + code: "INVALID_COMMAND", + message: "Invalid command", }, }); expect(cliWs.send).not.toHaveBeenCalled(); - expect(Reflect.get(doInstance, 'pendingCommands')).toEqual(new Map()); + 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", () => { const { doInstance, mockCtx } = setup(); - const selectedOwner = addCliSocket(mockCtx, 'cli-1'); - const otherCli = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const selectedOwner = addCliSocket(mockCtx, "cli-1"); + const otherCli = addCliSocket(mockCtx, "cli-2"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, selectedOwner, [makeSession('s1')]); - sendHeartbeat(doInstance, otherCli, [makeSession('s2')]); + sendHeartbeat(doInstance, selectedOwner, [makeSession("s1")]); + sendHeartbeat(doInstance, otherCli, [makeSession("s2")]); selectedOwner.send.mockClear(); otherCli.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); - expect(allSent(selectedOwner).filter(message => message.type === 'command')).toEqual([ + expect( + allSent(selectedOwner).filter((message) => message.type === "command"), + ).toEqual([ expect.objectContaining({ - type: 'command', - command: 'exit_cli', - sessionId: 's1', + type: "command", + command: "exit_cli", + sessionId: "s1", data: { protocolVersion: 1 }, }), ]); 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", () => { const { doInstance, mockCtx } = setup(); - const currentOwner = addCliSocket(mockCtx, 'cli-1'); - const staleOwner = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const currentOwner = addCliSocket(mockCtx, "cli-1"); + const staleOwner = addCliSocket(mockCtx, "cli-2"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, currentOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, currentOwner, [makeSession("s1")]); sendHeartbeat(doInstance, staleOwner, []); currentOwner.send.mockClear(); staleOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-2', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-2", data: { protocolVersion: 1 }, }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); expect(currentOwner.send).not.toHaveBeenCalled(); 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", () => { const { doInstance, mockCtx } = setup(); - const webWs = addWebSocket(mockCtx, 'web-1'); + const webWs = addWebSocket(mockCtx, "web-1"); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", data: { protocolVersion: 1 }, }); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', - error: 'Session owner not found', + type: "response", + id: "cmd-1", + error: "Session owner not found", }); }); - it('does not dedupe concurrent exit_cli requests', () => { + it("does not dedupe concurrent exit_cli requests", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); sendCommand(doInstance, webWs, { - id: 'cmd-2', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-2", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); - const commands = allSent(cliWs).filter(message => message.type === 'command'); + const commands = allSent(cliWs).filter( + (message) => message.type === "command", + ); expect(commands).toHaveLength(2); expect(commands[0].id).not.toBe(commands[1].id); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); const correlationId = getCorrelationId(cliWs); - const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); + const result = createResultWithSerializedBytes( + MAX_CATALOG_RESULT_BYTES + 1, + ); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); const correlationId = getCorrelationId(cliWs); @@ -3108,42 +3393,44 @@ describe('UserConnectionDO', () => { // Exit's own effect: CLI heartbeat no longer lists the session. sendHeartbeat(doInstance, cliWs, []); - expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ - type: 'response', - id: 'cmd-1', + expect( + allSent(webWs).find((m) => m.type === "response" && m.id === "cmd-1"), + ).toEqual({ + type: "response", + id: "cmd-1", result: {}, }); expect( allSent(webWs).some( - m => - m.type === 'response' && - m.id === 'cmd-1' && + (m) => + m.type === "response" && + m.id === "cmd-1" && isRecord(m.error) && - m.error.code === 'SESSION_OWNER_CHANGED' - ) + m.error.code === "SESSION_OWNER_CHANGED", + ), ).toBe(false); // CLI's late ACK must not produce a second cmd-1 response. const responsesBeforeLateAck = allSent(webWs).filter( - m => m.type === 'response' && m.id === 'cmd-1' + (m) => m.type === "response" && m.id === "cmd-1", ).length; sendCliResponse(doInstance, cliWs, { id: correlationId, result: {} }); - expect(allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-1')).toHaveLength( - responsesBeforeLateAck - ); + expect( + allSent(webWs).filter((m) => m.type === "response" && m.id === "cmd-1"), + ).toHaveLength(responsesBeforeLateAck); }); - it('resolves exit_cli successfully when the owning socket closes', async () => { + it("resolves exit_cli successfully when the owning socket closes", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); webWs.send.mockClear(); @@ -3152,102 +3439,109 @@ describe('UserConnectionDO', () => { await disconnectCli(doInstance, cliWs); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", result: {}, }); }); - it.each(['list_models', 'send_message'] as const)( - 'still fails %s with SESSION_OWNER_CHANGED when heartbeat drops the session', - command => { + it.each(["list_models", "send_message"] as const)( + "still fails %s with SESSION_OWNER_CHANGED when heartbeat drops the session", + (command) => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', + id: "cmd-1", command, - sessionId: 's1', - connectionId: 'cli-1', + sessionId: "s1", + connectionId: "cli-1", }); webWs.send.mockClear(); sendHeartbeat(doInstance, cliWs, []); - expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ - type: 'response', - id: 'cmd-1', + expect( + allSent(webWs).find((m) => m.type === "response" && m.id === "cmd-1"), + ).toEqual({ + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); - } + }, ); - it('still fails exit_cli with SESSION_OWNER_CHANGED on genuine takeover', () => { + it("still fails exit_cli with SESSION_OWNER_CHANGED on genuine takeover", () => { const { doInstance, mockCtx } = setup(); - const firstOwner = addCliSocket(mockCtx, 'cli-1'); - const nextOwner = addCliSocket(mockCtx, 'cli-2'); - const webWs = addWebSocket(mockCtx, 'web-1'); + 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, firstOwner, [makeSession("s1")]); sendHeartbeat(doInstance, nextOwner, []); firstOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); webWs.send.mockClear(); - sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, nextOwner, [makeSession("s1")]); - expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ - type: 'response', - id: 'cmd-1', + expect( + allSent(webWs).find((m) => m.type === "response" && m.id === "cmd-1"), + ).toEqual({ + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); }); - 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", () => { const { doInstance, mockCtx } = setup(); - const firstCli = connectCliSocket(doInstance, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const firstCli = connectCliSocket(doInstance, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, firstCli, [makeSession('s1')]); + sendHeartbeat(doInstance, firstCli, [makeSession("s1")]); firstCli.send.mockClear(); sendCommand(doInstance, webWs, { - id: 'cmd-1', - command: 'exit_cli', - sessionId: 's1', - connectionId: 'cli-1', + id: "cmd-1", + command: "exit_cli", + sessionId: "s1", + connectionId: "cli-1", data: { protocolVersion: 1 }, }); webWs.send.mockClear(); - connectCliSocket(doInstance, 'cli-1'); + connectCliSocket(doInstance, "cli-1"); - expect(firstCli.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); + expect(firstCli.close).toHaveBeenCalledWith( + 1000, + "replaced by reconnect", + ); expect(parseSent(webWs)).toEqual({ - type: 'response', - id: 'cmd-1', + type: "response", + id: "cmd-1", error: { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }, }); }); @@ -3257,87 +3551,87 @@ describe('UserConnectionDO', () => { // CLI event forwarding // ------------------------------------------------------------------------- - describe('CLI event forwarding', () => { - it('forwards events to subscribed web sockets only', () => { + describe("CLI event forwarding", () => { + it("forwards events to subscribed web sockets only", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const subWeb = addWebSocket(mockCtx, 'web-sub'); - const otherWeb = addWebSocket(mockCtx, 'web-other'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const subWeb = addWebSocket(mockCtx, "web-sub"); + const otherWeb = addWebSocket(mockCtx, "web-other"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, subWeb, 's1'); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendSubscribe(doInstance, subWeb, "s1"); subWeb.send.mockClear(); otherWeb.send.mockClear(); // CLI sends event for s1 const eventMsg = JSON.stringify({ - type: 'event', - sessionId: 's1', - event: 'message.updated', - data: { id: 'msg-1' }, + type: "event", + sessionId: "s1", + event: "message.updated", + data: { id: "msg-1" }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(subWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(subWeb)).toEqual({ - type: 'event', - sessionId: 's1', - event: 'message.updated', - data: { id: 'msg-1' }, + type: "event", + sessionId: "s1", + event: "message.updated", + data: { id: "msg-1" }, }); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const parentWeb = addWebSocket(mockCtx, 'web-parent'); - const childWeb = addWebSocket(mockCtx, 'web-child'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const parentWeb = addWebSocket(mockCtx, "web-parent"); + const childWeb = addWebSocket(mockCtx, "web-child"); - sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); - sendSubscribe(doInstance, parentWeb, 'parent-session'); - sendSubscribe(doInstance, childWeb, 'child-session-1'); + sendHeartbeat(doInstance, cliWs, [makeSession("parent-session")]); + sendSubscribe(doInstance, parentWeb, "parent-session"); + sendSubscribe(doInstance, childWeb, "child-session-1"); parentWeb.send.mockClear(); childWeb.send.mockClear(); const eventMsg = JSON.stringify({ - type: 'event', - sessionId: 'child-session-1', - parentSessionId: 'parent-session', - event: 'message.updated', - data: { id: 'msg-1' }, + type: "event", + sessionId: "child-session-1", + parentSessionId: "parent-session", + event: "message.updated", + data: { id: "msg-1" }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(parentWeb.send).toHaveBeenCalledTimes(1); expect(childWeb.send).toHaveBeenCalledTimes(1); const expected = { - type: 'event', - sessionId: 'child-session-1', - parentSessionId: 'parent-session', - event: 'message.updated', - data: { id: 'msg-1' }, + type: "event", + sessionId: "child-session-1", + parentSessionId: "parent-session", + event: "message.updated", + data: { id: "msg-1" }, }; expect(parseSent(parentWeb)).toEqual(expected); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); - sendSubscribe(doInstance, webWs, 'parent-session'); - sendSubscribe(doInstance, webWs, 'child-session-1'); + sendHeartbeat(doInstance, cliWs, [makeSession("parent-session")]); + sendSubscribe(doInstance, webWs, "parent-session"); + sendSubscribe(doInstance, webWs, "child-session-1"); webWs.send.mockClear(); const eventMsg = JSON.stringify({ - type: 'event', - sessionId: 'child-session-1', - parentSessionId: 'parent-session', - event: 'message.updated', - data: { id: 'msg-1' }, + type: "event", + sessionId: "child-session-1", + parentSessionId: "parent-session", + event: "message.updated", + data: { id: "msg-1" }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); @@ -3345,102 +3639,102 @@ 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); - sendSubscribe(doInstance, webWs, 'parent-session'); + sendHeartbeat(doInstance, cliWs, [makeSession("parent-session")]); + sendSubscribe(doInstance, webWs, "parent-session"); webWs.send.mockClear(); // CLI sends event for a child session with parentSessionId const eventMsg = JSON.stringify({ - type: 'event', - sessionId: 'child-session-1', - parentSessionId: 'parent-session', - event: 'message.updated', - data: { id: 'msg-child-1' }, + type: "event", + sessionId: "child-session-1", + parentSessionId: "parent-session", + event: "message.updated", + data: { id: "msg-child-1" }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: 'event', - sessionId: 'child-session-1', - parentSessionId: 'parent-session', - event: 'message.updated', - data: { id: 'msg-child-1' }, + type: "event", + sessionId: "child-session-1", + parentSessionId: "parent-session", + event: "message.updated", + data: { id: "msg-child-1" }, }); }); - it('drops child event when neither sessionId nor parentSessionId has subscribers', () => { + it("drops child event when neither sessionId nor parentSessionId has subscribers", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('other-session')]); - sendSubscribe(doInstance, webWs, 'other-session'); + sendHeartbeat(doInstance, cliWs, [makeSession("other-session")]); + sendSubscribe(doInstance, webWs, "other-session"); webWs.send.mockClear(); // Child event with parent that nobody subscribes to const eventMsg = JSON.stringify({ - type: 'event', - sessionId: 'child-session-1', - parentSessionId: 'unknown-parent', - event: 'message.updated', - data: { id: 'msg-child-1' }, + type: "event", + sessionId: "child-session-1", + parentSessionId: "unknown-parent", + event: "message.updated", + data: { id: "msg-child-1" }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).not.toHaveBeenCalled(); }); - it('events without parentSessionId still route normally (backward compat)', () => { + it("events without parentSessionId still route normally (backward compat)", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, webWs, 's1'); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendSubscribe(doInstance, webWs, "s1"); webWs.send.mockClear(); const eventMsg = JSON.stringify({ - type: 'event', - sessionId: 's1', - event: 'message.updated', - data: { id: 'msg-1' }, + type: "event", + sessionId: "s1", + event: "message.updated", + data: { id: "msg-1" }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: 'event', - sessionId: 's1', - event: 'message.updated', - data: { id: 'msg-1' }, + type: "event", + sessionId: "s1", + event: "message.updated", + data: { id: "msg-1" }, }); }); - it('child event does not include parentSessionId when not set', () => { + it("child event does not include parentSessionId when not set", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, webWs, 's1'); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendSubscribe(doInstance, webWs, "s1"); webWs.send.mockClear(); const eventMsg = JSON.stringify({ - type: 'event', - sessionId: 's1', - event: 'session.status', + type: "event", + sessionId: "s1", + event: "session.status", data: {}, }); doInstance.webSocketMessage(cliWs as never, eventMsg); const sent = parseSent(webWs); - expect(sent).not.toHaveProperty('parentSessionId'); + expect(sent).not.toHaveProperty("parentSessionId"); }); }); @@ -3448,12 +3742,12 @@ describe('UserConnectionDO', () => { // Broadcast resilience // ------------------------------------------------------------------------- - describe('broadcast resilience', () => { - it('one closed socket does not abort send to other web sockets', () => { + describe("broadcast resilience", () => { + it("one closed socket does not abort send to other web sockets", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const failWeb = addWebSocket(mockCtx, 'web-fail'); - const okWeb = addWebSocket(mockCtx, 'web-ok'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const failWeb = addWebSocket(mockCtx, "web-fail"); + const okWeb = addWebSocket(mockCtx, "web-ok"); // Both web sockets receive heartbeats via broadcast (no subscription needed). failWeb.send.mockClear(); @@ -3461,16 +3755,16 @@ describe('UserConnectionDO', () => { // Make failWeb throw on send failWeb.send.mockImplementation(() => { - throw new Error('socket closed'); + throw new Error("socket closed"); }); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); // okWeb should still receive the message expect(okWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(okWeb)).toMatchObject({ - type: 'system', - event: 'sessions.heartbeat', + type: "system", + event: "sessions.heartbeat", }); }); }); @@ -3479,42 +3773,51 @@ describe('UserConnectionDO', () => { // Hibernation recovery (ensureState) // ------------------------------------------------------------------------- - describe('ensureState (hibernation recovery)', () => { - it('reconstructs sessionOwners and connectionSessions from CLI attachments', () => { + describe("ensureState (hibernation recovery)", () => { + it("reconstructs sessionOwners and connectionSessions from CLI attachments", () => { const { doInstance, mockCtx } = setup(); // Simulate hibernation: sockets exist with pre-set attachments - const sessions = [makeSession('s1'), makeSession('s2')]; - addCliSocket(mockCtx, 'cli-1', sessions); - const webWs = addWebSocket(mockCtx, 'web-1'); + const sessions = [makeSession("s1"), makeSession("s2")]; + addCliSocket(mockCtx, "cli-1", sessions); + const webWs = addWebSocket(mockCtx, "web-1"); // Trigger ensureState by calling any method (e.g., webSocketMessage with subscribe) - sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, "s1"); // 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' }); + const web2 = addWebSocket(mockCtx, "web-2"); + 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')); + 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' }); + const cmdMsg = cliMsgs.find( + (m: Record) => m.type === "command", + ); + expect(cmdMsg).toMatchObject({ + type: "command", + command: "send_message", + }); }); - it('reconstructs webSubscriptions from web attachments', () => { + it("reconstructs webSubscriptions from web attachments", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1', [makeSession('s1')]); + const cliWs = addCliSocket(mockCtx, "cli-1", [makeSession("s1")]); // Web socket with pre-existing subscription (from hibernation) - const webWs = addWebSocket(mockCtx, 'web-1', ['s1']); + const webWs = addWebSocket(mockCtx, "web-1", ["s1"]); // Trigger ensureState by calling any method const triggerMsg = JSON.stringify({ - type: 'event', - sessionId: 's1', - event: 'test', + type: "event", + sessionId: "s1", + event: "test", data: {}, }); doInstance.webSocketMessage(cliWs as never, triggerMsg); @@ -3522,25 +3825,30 @@ describe('UserConnectionDO', () => { // webWs should have received the event because it was subscribed via attachment expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toMatchObject({ - type: 'event', - sessionId: 's1', + type: "event", + sessionId: "s1", }); }); - it('does not restore subscriptions from a viewer already replaced before hibernation', () => { + it("does not restore subscriptions from a viewer already replaced before hibernation", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1', [makeSession('s1')]); - const replacedWeb = addWebSocket(mockCtx, 'web-old', ['s1']); + const cliWs = addCliSocket(mockCtx, "cli-1", [makeSession("s1")]); + const replacedWeb = addWebSocket(mockCtx, "web-old", ["s1"]); replacedWeb.serializeAttachment({ - role: 'web', - connectionId: 'web-old', - subscribedSessions: ['s1'], + role: "web", + connectionId: "web-old", + subscribedSessions: ["s1"], replaced: true, }); 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(); @@ -3551,41 +3859,47 @@ describe('UserConnectionDO', () => { // getActiveSessions RPC // ------------------------------------------------------------------------- - describe('getActiveSessions', () => { - it('returns sessions from live CLI connections', () => { + describe("getActiveSessions", () => { + it("returns sessions from live CLI connections", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); sendHeartbeat(doInstance, cliWs, [ - makeSession('s1', 'busy', 'Fix bug'), - makeSession('s2', 'idle', 'Review PR'), + makeSession("s1", "busy", "Fix bug"), + makeSession("s2", "idle", "Review PR"), ]); const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: 's1', status: 'busy', title: 'Fix bug', connectionId: 'cli-1' }, - { id: 's2', status: 'idle', title: 'Review PR', connectionId: 'cli-1' }, + { id: "s1", status: "busy", title: "Fix bug", connectionId: "cli-1" }, + { id: "s2", status: "idle", title: "Review PR", connectionId: "cli-1" }, ]); }); - it('includes the CLI-reported protocolVersion on each session row', () => { + it("includes the CLI-reported protocolVersion on each session row", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Fix bug')], { - protocolVersion: '1', + sendHeartbeat(doInstance, cliWs, [makeSession("s1", "busy", "Fix bug")], { + protocolVersion: "1", }); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + const cliWs = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); // Remove from sockets (simulates close) mockCtx.removeSocket(cliWs); @@ -3594,68 +3908,86 @@ describe('UserConnectionDO', () => { expect(result).toEqual([]); }); - it('excludes child sessions reported with parentSessionId in heartbeat', () => { + it("excludes child sessions reported with parentSessionId in heartbeat", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); sendHeartbeat(doInstance, cliWs, [ - makeSession('root-1', 'busy', 'Root session'), - makeSession('child-1', 'busy', 'Child session', 'root-1'), + makeSession("root-1", "busy", "Root session"), + makeSession("child-1", "busy", "Child session", "root-1"), ]); 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); // First heartbeat: root + child sendHeartbeat(doInstance, cliWs, [ - makeSession('root-1', 'busy', 'Root session'), - makeSession('child-1', 'busy', 'Child session', 'root-1'), + makeSession("root-1", "busy", "Root session"), + makeSession("child-1", "busy", "Child session", "root-1"), ]); // Second heartbeat: only root (child finished) - sendHeartbeat(doInstance, cliWs, [makeSession('root-1', 'idle', 'Root session')]); + sendHeartbeat(doInstance, cliWs, [ + makeSession("root-1", "idle", "Root session"), + ]); 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)", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); sendHeartbeat(doInstance, cliWs, [ - makeSession('s1', 'busy', 'On a Mac', undefined, 'darwin'), - makeSession('s2', 'idle', 'Other'), + makeSession("s1", "busy", "On a Mac", undefined, "darwin"), + makeSession("s2", "idle", "Other"), ]); const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: 's1', status: 'busy', title: 'On a Mac', connectionId: 'cli-1', platform: 'darwin' }, - { id: 's2', status: 'idle', title: 'Other', connectionId: 'cli-1' }, + { + 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)", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); // Legacy CLI heartbeat without platform. - sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Legacy')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1", "busy", "Legacy")]); const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: 's1', status: 'busy', title: 'Legacy', connectionId: 'cli-1' }, + { id: "s1", status: "busy", title: "Legacy", connectionId: "cli-1" }, ]); - expect(result[0]).not.toHaveProperty('platform'); + expect(result[0]).not.toHaveProperty("platform"); }); }); @@ -3663,56 +3995,69 @@ describe('UserConnectionDO', () => { // getConnectedInstances RPC (W3) // ------------------------------------------------------------------------- - describe('getConnectedInstances', () => { - it('returns one row per CLI socket that has an `instance` attachment', () => { + describe("getConnectedInstances", () => { + it("returns one row per CLI socket that has an `instance` attachment", () => { const { doInstance, mockCtx } = setup(); // Use the hibernated-attachment pattern (no heartbeat) — the live // scan reads the `instance` directly from the attachment, which is // what the spec requires: a fresh value with no in-memory map. - addCliSocket(mockCtx, 'cli-A', [], { - name: 'laptop-A', - projectName: 'kilo', - version: '0.1.2', + addCliSocket(mockCtx, "cli-A", [], { + name: "laptop-A", + 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-B', name: 'laptop-B', projectName: 'kilo' }, - ]) + { + 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", () => { 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'); + 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`", () => { const { doInstance, mockCtx } = setup(); // Legacy CLI: pre-spawner heartbeat has no `instance`. - const cliWs = addCliSocket(mockCtx, 'legacy-1'); + const cliWs = addCliSocket(mockCtx, "legacy-1"); sendHeartbeat(doInstance, cliWs, []); const { instances } = doInstance.getConnectedInstances(); expect(instances).toEqual([]); }); - it('excludes web sockets', () => { + it("excludes web sockets", () => { const { doInstance, mockCtx } = setup(); addWebSocket(mockCtx); // A web socket with an `instance`-shaped attachment must still be skipped. - const webWithInstance = createMockWs(['web'], { - role: 'web', - connectionId: 'web-1', + const webWithInstance = createMockWs(["web"], { + role: "web", + connectionId: "web-1", subscribedSessions: [], } as never); mockCtx.addSocket(webWithInstance); @@ -3721,31 +4066,36 @@ 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)", () => { 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. - addCliSocket(mockCtx, 'cli-h', [], { - name: 'laptop-h', - projectName: 'kilo', - version: '1.0.0', + addCliSocket(mockCtx, "cli-h", [], { + name: "laptop-h", + projectName: "kilo", + version: "1.0.0", }); 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", () => { const { doInstance, mockCtx } = setup(); // Hibernated attachment carries capabilities — same source // getConnectedInstances already uses for instance/version. - const cliWs = createMockWs(['cli'], { - role: 'cli', - connectionId: 'cli-cap', + const cliWs = createMockWs(["cli"], { + role: "cli", + connectionId: "cli-cap", sessions: [], - instance: { name: 'laptop-cap', projectName: 'kilo' }, + instance: { name: "laptop-cap", projectName: "kilo" }, capabilities: { attachments: true }, }); mockCtx.addSocket(cliWs); @@ -3753,46 +4103,56 @@ describe('UserConnectionDO', () => { const { instances } = doInstance.getConnectedInstances(); expect(instances).toEqual([ { - connectionId: 'cli-cap', - name: 'laptop-cap', - projectName: 'kilo', + connectionId: "cli-cap", + name: "laptop-cap", + projectName: "kilo", capabilities: { attachments: true }, }, ]); }); - it('omits capabilities when the CLI attachment has none (legacy CLI)', () => { + it("omits capabilities when the CLI attachment has none (legacy CLI)", () => { const { doInstance, mockCtx } = setup(); - addCliSocket(mockCtx, 'cli-legacy-cap', [], { - name: 'laptop-legacy', - projectName: 'kilo', + addCliSocket(mockCtx, "cli-legacy-cap", [], { + name: "laptop-legacy", + projectName: "kilo", }); 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'); + expect(instances[0]).not.toHaveProperty("capabilities"); }); - it('persists `instance` in the WS attachment across heartbeats', () => { + it("persists `instance` in the WS attachment across heartbeats", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); sendHeartbeat(doInstance, cliWs, [], { - protocolVersion: '1', - instance: { name: 'laptop-1', projectName: 'kilo', version: '0.1.0' }, + protocolVersion: "1", + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); // First heartbeat: with instance. sendHeartbeat(doInstance, cliWs, [], { - instance: { name: 'laptop-1', projectName: 'kilo' }, + instance: { name: "laptop-1", projectName: "kilo" }, }); // Second heartbeat: instance removed (legacy fallback). The DO must not // keep a stale `instance` value in the attachment. @@ -3807,7 +4167,7 @@ describe('UserConnectionDO', () => { // WS attachment size guardrail (W3) // ------------------------------------------------------------------------- - describe('WS attachment size', () => { + describe("WS attachment size", () => { // The Cloudflare `serializeAttachment` budget is ~2 KiB. A bounded // instance object (name+projectName+version, all max length) adds well // under 200 bytes; this test pins that contract so a future schema @@ -3819,38 +4179,42 @@ 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", () => { const worstCaseInstance = { - name: 'x'.repeat(64), - projectName: 'x'.repeat(64), - version: 'x'.repeat(32), + name: "x".repeat(64), + projectName: "x".repeat(64), + version: "x".repeat(32), }; // 4 sessions with realistic-but-large titles, git URLs, and branches. // (4 is a generous upper bound for a single CLI owning a live session // fleet; the actual HeartbeatSession shape imposes tighter per-field // limits at the protocol layer.) const sessions = Array.from({ length: 4 }, (_, i) => ({ - id: `ses_${String(i).padStart(26, '0')}`, - status: 'busy', - title: 'T'.repeat(120), - gitUrl: 'https://github.com/org/' + 'x'.repeat(60) + '.git', - gitBranch: 'b'.repeat(40), + id: `ses_${String(i).padStart(26, "0")}`, + status: "busy", + title: "T".repeat(120), + gitUrl: "https://github.com/org/" + "x".repeat(60) + ".git", + gitBranch: "b".repeat(40), })); const attachment = { - role: 'cli' as const, - connectionId: 'cli-1', + role: "cli" as const, + connectionId: "cli-1", sessions, - protocolVersion: '255.255.65535', - kiloUserId: 'usr_' + 'x'.repeat(28), + protocolVersion: "255.255.65535", + kiloUserId: "usr_" + "x".repeat(28), instance: worstCaseInstance, }; - const serialized = new TextEncoder().encode(JSON.stringify(attachment)).byteLength; + const serialized = new TextEncoder().encode( + JSON.stringify(attachment), + ).byteLength; expect(serialized).toBeLessThan(SERIALIZE_ATTACHMENT_BUDGET); // Sanity: the bounded instance alone is far below the headroom. - const instanceBytes = new TextEncoder().encode(JSON.stringify(worstCaseInstance)).byteLength; + const instanceBytes = new TextEncoder().encode( + JSON.stringify(worstCaseInstance), + ).byteLength; expect(instanceBytes).toBeLessThan(INSTANCE_HEADROOM); }); }); @@ -3859,23 +4223,29 @@ describe('UserConnectionDO', () => { // Owner-unique active sessions (W-followup) // ------------------------------------------------------------------------- - describe('owner-unique active sessions', () => { - it('emits owner-unique rows: ownership transfer with both CLIs live yields exactly one row under the new owner', () => { + describe("owner-unique active sessions", () => { + it("emits owner-unique rows: ownership transfer with both CLIs live yields exactly one row under the new owner", () => { const { doInstance, ctx, mockCtx } = setup(); - const oldOwner = addCliSocket(mockCtx, 'cli-old'); - const newOwner = addCliSocket(mockCtx, 'cli-new'); + const oldOwner = addCliSocket(mockCtx, "cli-old"); + const newOwner = addCliSocket(mockCtx, "cli-new"); // cli-old claims the session - sendHeartbeat(doInstance, oldOwner, [makeSession('ses_transfer', 'busy', 'Transfer me')]); + sendHeartbeat(doInstance, oldOwner, [ + makeSession("ses_transfer", "busy", "Transfer me"), + ]); // cli-new also claims the same session id while cli-old is still connected. // The DO routes the session to the new owner (sessionOwners.get === 'cli-new'). - sendHeartbeat(doInstance, newOwner, [makeSession('ses_transfer', 'busy', 'Transfer me')]); + sendHeartbeat(doInstance, newOwner, [ + makeSession("ses_transfer", "busy", "Transfer me"), + ]); // Both CLIs are still live sockets — the snapshot should see them both. - expect(ctx.getWebSockets('cli').map(ws => ws.deserializeAttachment())).toEqual([ - expect.objectContaining({ role: 'cli', connectionId: 'cli-old' }), - expect.objectContaining({ role: 'cli', connectionId: 'cli-new' }), + expect( + ctx.getWebSockets("cli").map((ws) => ws.deserializeAttachment()), + ).toEqual([ + expect.objectContaining({ role: "cli", connectionId: "cli-old" }), + expect.objectContaining({ role: "cli", connectionId: "cli-new" }), ]); const result = doInstance.getActiveSessions(); @@ -3883,10 +4253,10 @@ describe('UserConnectionDO', () => { // Exactly one row for the transferred session id, under the new owner. expect(result).toEqual([ { - id: 'ses_transfer', - status: 'busy', - title: 'Transfer me', - connectionId: 'cli-new', + id: "ses_transfer", + status: "busy", + title: "Transfer me", + connectionId: "cli-new", }, ]); }); @@ -3896,85 +4266,97 @@ describe('UserConnectionDO', () => { // Edge cases // ------------------------------------------------------------------------- - describe('edge cases', () => { - it('ignores non-JSON messages', () => { + describe("edge cases", () => { + it("ignores non-JSON messages", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); // Should not throw - doInstance.webSocketMessage(cliWs as never, 'not-json'); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const warn = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); const malformed = '{"secret":"raw-secret-must-not-be-logged"'; doInstance.webSocketMessage(cliWs as never, malformed); - expect(warn).toHaveBeenCalledWith('Failed to parse WebSocket message as JSON', { - role: 'cli', - connectionId: 'cli-1', - byteCount: new TextEncoder().encode(malformed).byteLength, - }); - expect(JSON.stringify(warn.mock.calls)).not.toContain('raw-secret-must-not-be-logged'); + expect(warn).toHaveBeenCalledWith( + "Failed to parse WebSocket message as JSON", + { + role: "cli", + connectionId: "cli-1", + byteCount: new TextEncoder().encode(malformed).byteLength, + }, + ); + 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", () => { const { doInstance, mockCtx } = setup(); - const ws = createMockWs(['cli'], null); + const ws = createMockWs(["cli"], null); mockCtx.addSocket(ws); // Trigger ensureState first - doInstance.webSocketMessage(ws as never, JSON.stringify({ type: 'heartbeat', sessions: [] })); + doInstance.webSocketMessage( + ws as never, + JSON.stringify({ type: "heartbeat", sessions: [] }), + ); // Should not throw }); - it('ignores messages that fail Zod validation', () => { + it("ignores messages that fail Zod validation", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); sendHeartbeat(doInstance, cliWs, []); // trigger ensureState // Invalid CLI message - const badMsg = JSON.stringify({ type: 'invalid_type' }); + const badMsg = JSON.stringify({ type: "invalid_type" }); doInstance.webSocketMessage(cliWs as never, badMsg); // Should not throw // Invalid web message - const webWs = addWebSocket(mockCtx, 'web-1'); + const webWs = addWebSocket(mockCtx, "web-1"); doInstance.webSocketMessage(webWs as never, badMsg); // Should not throw }); - it('logs malformed CLI message metadata without raw payload content', () => { + it("logs malformed CLI message metadata without raw payload content", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); - const secret = 'raw-secret-must-not-be-logged'; + const cliWs = addCliSocket(mockCtx, "cli-1"); + const warn = vi + .spyOn(console, "warn") + .mockImplementation(() => undefined); + const secret = "raw-secret-must-not-be-logged"; const malformed = JSON.stringify({ - type: 'response', + type: "response", id: 123, result: { secret }, }); doInstance.webSocketMessage(cliWs as never, malformed); - expect(warn).toHaveBeenCalledWith('CLI message parse failed', { - role: 'cli', - connectionId: 'cli-1', + expect(warn).toHaveBeenCalledWith("CLI message parse failed", { + role: "cli", + connectionId: "cli-1", byteCount: new TextEncoder().encode(malformed).byteLength, - issues: [{ path: ['id'], code: 'invalid_type' }], + issues: [{ path: ["id"], code: "invalid_type" }], }); expect(JSON.stringify(warn.mock.calls)).not.toContain(secret); }); - it('webSocketError triggers webSocketClose', async () => { + it("webSocketError triggers webSocketClose", async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - const webWs = addWebSocket(mockCtx, 'web-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + const webWs = addWebSocket(mockCtx, "web-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); webWs.send.mockClear(); // Remove CLI so disconnect can clean up @@ -3983,16 +4365,20 @@ describe('UserConnectionDO', () => { // Should broadcast cli.disconnected const msgs = allSent(webWs); - expect(msgs.some((m: Record) => m.event === 'cli.disconnected')).toBe(true); + 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", () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + const cliWs = addCliSocket(mockCtx, "cli-1"); sendHeartbeat(doInstance, cliWs, []); // Should not throw - sendCliResponse(doInstance, cliWs, { id: 'nonexistent', result: 'ok' }); + sendCliResponse(doInstance, cliWs, { id: "nonexistent", result: "ok" }); }); }); @@ -4000,14 +4386,15 @@ describe('UserConnectionDO', () => { // Session-ready push (delayed, Decision 9) // ------------------------------------------------------------------------- - describe('session-ready delayed push', () => { + describe("session-ready delayed push", () => { function setupWithIngestDO() { const mockCtx = createMockCtx(); const ctx = mockCtx.build(); const claimSessionReadyPush = vi.fn(async () => {}); sessionIngestMocks.getSessionIngestDO.mockReturnValue({ claimSessionReadyPush, - resetAttentionStatusOnCliDisconnect: sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + resetAttentionStatusOnCliDisconnect: + sessionIngestMocks.resetAttentionStatusOnCliDisconnect, }); const doInstance = new UserConnectionDO(ctx as never, {} as never); return { doInstance, mockCtx, ctx, claimSessionReadyPush }; @@ -4016,76 +4403,86 @@ describe('UserConnectionDO', () => { function addCliSocketForUser( mockCtx: ReturnType, connectionId: string, - kiloUserId: string + kiloUserId: string, ): MockWS { - const attachment = { role: 'cli' as const, connectionId, sessions: [], kiloUserId }; - const ws = createMockWs(['cli'], attachment); + const attachment = { + role: "cli" as const, + connectionId, + sessions: [], + kiloUserId, + }; + const ws = createMockWs(["cli"], attachment); mockCtx.addSocket(ws); return ws; } - it('writes a pending readyPush entry on first sight and does not claim immediately', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); - const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + it("writes a pending readyPush entry on first sight and does not claim immediately", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); + vi.spyOn(Date, "now").mockReturnValue(now); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_main')]); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_main")]); await flushAsync(); expect(claimSessionReadyPush).not.toHaveBeenCalled(); - const entry = ctx.storage.store.get('readyPush:ses_main') as { + const entry = ctx.storage.store.get("readyPush:ses_main") as { kiloUserId: string; title: string; fireAt: number; attempts: number; }; expect(entry).toMatchObject({ - kiloUserId: 'usr_1', - title: 'Test', + kiloUserId: "usr_1", + title: "Test", fireAt: now + 5_000, attempts: 0, }); // Subsequent heartbeats for the same session must not re-schedule. - sendHeartbeat(doInstance, cliWs, [makeSession('ses_main')]); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_main")]); await flushAsync(); expect(ctx.storage.put).toHaveBeenCalledTimes(1); }); - it('reconnect first-sight does not reset a pending readyPush fireAt or attempts', async () => { + it("reconnect first-sight does not reset a pending readyPush fireAt or attempts", async () => { const { doInstance, mockCtx, ctx } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); + vi.spyOn(Date, "now").mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_reconnect')]); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_reconnect")]); await flushAsync(); const originalFireAt = now + 5_000; // Simulate a prior retry so attempts is non-zero. - await ctx.storage.put('readyPush:ses_reconnect', { - kiloUserId: 'usr_1', - title: 'Test', + await ctx.storage.put("readyPush:ses_reconnect", { + kiloUserId: "usr_1", + title: "Test", fireAt: originalFireAt, attempts: 2, }); - const internal = doInstance as unknown as { readyPushFireAt: Map }; - internal.readyPushFireAt.set('ses_reconnect', originalFireAt); + const internal = doInstance as unknown as { + readyPushFireAt: Map; + }; + internal.readyPushFireAt.set("ses_reconnect", originalFireAt); // Disconnect clears sessionOwners; reconnect is a new "first sight". await disconnectCli(doInstance, cliWs); - expect(internal.readyPushFireAt.get('ses_reconnect')).toBe(originalFireAt); + expect(internal.readyPushFireAt.get("ses_reconnect")).toBe( + originalFireAt, + ); const later = now + 2_000; - vi.spyOn(Date, 'now').mockReturnValue(later); + vi.spyOn(Date, "now").mockReturnValue(later); ctx.storage.put.mockClear(); - const cliWs2 = addCliSocketForUser(mockCtx, 'cli-2', 'usr_1'); - sendHeartbeat(doInstance, cliWs2, [makeSession('ses_reconnect')]); + const cliWs2 = addCliSocketForUser(mockCtx, "cli-2", "usr_1"); + sendHeartbeat(doInstance, cliWs2, [makeSession("ses_reconnect")]); await flushAsync(); - const entry = ctx.storage.store.get('readyPush:ses_reconnect') as { + const entry = ctx.storage.store.get("readyPush:ses_reconnect") as { fireAt: number; attempts: number; }; @@ -4095,39 +4492,41 @@ describe('UserConnectionDO', () => { expect(ctx.storage.put).not.toHaveBeenCalled(); }); - it('never schedules for subagent sessions', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); - const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + it("never schedules for subagent sessions", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); sendHeartbeat(doInstance, cliWs, [ - makeSession('ses_main'), - makeSession('ses_sub', 'busy', 'Sub', 'ses_main'), + makeSession("ses_main"), + makeSession("ses_sub", "busy", "Sub", "ses_main"), ]); await flushAsync(); - expect(ctx.storage.store.has('readyPush:ses_main')).toBe(true); - expect(ctx.storage.store.has('readyPush:ses_sub')).toBe(false); + expect(ctx.storage.store.has("readyPush:ses_main")).toBe(true); + expect(ctx.storage.store.has("readyPush:ses_sub")).toBe(false); expect(claimSessionReadyPush).not.toHaveBeenCalled(); }); - it('does not schedule on sockets without a kiloUserId (legacy attachment)', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); + it("does not schedule on sockets without a kiloUserId (legacy attachment)", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); + const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_main')]); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_main")]); await flushAsync(); expect(ctx.storage.store.size).toBe(0); expect(claimSessionReadyPush).not.toHaveBeenCalled(); }); - it('arms the alarm for a readyPush even with no heartbeat candidates', async () => { + it("arms the alarm for a readyPush even with no heartbeat candidates", async () => { const { doInstance, mockCtx, ctx } = setupWithIngestDO(); // No CLI sockets — only a pending readyPush in KV, rebuilt via one-shot. const fireAt = Date.now() + 5_000; - await ctx.storage.put('readyPush:ses_orphan', { - kiloUserId: 'usr_1', - title: 'Orphan', + await ctx.storage.put("readyPush:ses_orphan", { + kiloUserId: "usr_1", + title: "Orphan", fireAt, attempts: 0, }); @@ -4143,10 +4542,10 @@ describe('UserConnectionDO', () => { ctx.storage.setAlarm.mockClear(); // Force a schedule by going through a heartbeat that schedules readyPush // then clear lastHeartbeat by disconnecting — still has readyPush mirror. - const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_arm')]); + vi.spyOn(Date, "now").mockReturnValue(now); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_arm")]); await flushAsync(); // Disconnect removes lastHeartbeatAt for this connection await disconnectCli(doInstance, cliWs); @@ -4162,24 +4561,27 @@ describe('UserConnectionDO', () => { expect(armedAt).toBe(now + 5_000); }); - it('arms immediately for overdue readyPush entries (never filtered by fireAt > now)', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + it("arms immediately for overdue readyPush entries (never filtered by fireAt > now)", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); - await ctx.storage.put('readyPush:ses_overdue', { - kiloUserId: 'usr_1', - title: 'Old', + vi.spyOn(Date, "now").mockReturnValue(now); + await ctx.storage.put("readyPush:ses_overdue", { + kiloUserId: "usr_1", + title: "Old", fireAt: now - 1_000, attempts: 0, }); // Seed mirror via first-sight path, then replace with overdue - const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_overdue', 'idle', 'Old')]); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_overdue", "idle", "Old"), + ]); await flushAsync(); // Overwrite with overdue fireAt and mirror - await ctx.storage.put('readyPush:ses_overdue', { - kiloUserId: 'usr_1', - title: 'Old', + await ctx.storage.put("readyPush:ses_overdue", { + kiloUserId: "usr_1", + title: "Old", fireAt: now - 1_000, attempts: 0, }); @@ -4187,7 +4589,7 @@ describe('UserConnectionDO', () => { const internal = doInstance as unknown as { readyPushFireAt: Map; }; - internal.readyPushFireAt.set('ses_overdue', now - 1_000); + internal.readyPushFireAt.set("ses_overdue", now - 1_000); ctx.storage.setAlarm.mockClear(); claimSessionReadyPush.mockClear(); @@ -4197,13 +4599,13 @@ describe('UserConnectionDO', () => { // After fire, entry deleted; remaining schedule may or may not arm }); - it('one-shot mirror rebuild sets flag only after successful refresh', async () => { + it("one-shot mirror rebuild sets flag only after successful refresh", async () => { const { doInstance, mockCtx, ctx } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); - await ctx.storage.put('readyPush:ses_rebuild', { - kiloUserId: 'usr_1', - title: 'Rebuild', + vi.spyOn(Date, "now").mockReturnValue(now); + await ctx.storage.put("readyPush:ses_rebuild", { + kiloUserId: "usr_1", + title: "Rebuild", fireAt: now + 2_000, attempts: 0, }); @@ -4217,219 +4619,250 @@ describe('UserConnectionDO', () => { expect(internal.readyPushRebuilt).toBe(false); // Kick schedule with empty mirror → async rebuild - const cliWs = addCliSocket(mockCtx, 'cli-legacy'); // no kiloUserId → no new readyPush + const cliWs = addCliSocket(mockCtx, "cli-legacy"); // no kiloUserId → no new readyPush sendHeartbeat(doInstance, cliWs, []); await flushAsync(); expect(internal.readyPushRebuilt).toBe(true); - expect(internal.readyPushFireAt.get('ses_rebuild')).toBe(now + 2_000); + expect(internal.readyPushFireAt.get("ses_rebuild")).toBe(now + 2_000); expect(ctx.storage.setAlarm).toHaveBeenCalled(); // Failed list leaves flag false internal.readyPushFireAt.clear(); internal.readyPushRebuilt = false; - ctx.storage.list.mockRejectedValueOnce(new Error('kv down')); + ctx.storage.list.mockRejectedValueOnce(new Error("kv down")); sendHeartbeat(doInstance, cliWs, []); await flushAsync(); expect(internal.readyPushRebuilt).toBe(false); }); - it('alarm passes changed heartbeat title when connectionSessions diverged, else undefined', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + it("alarm passes changed heartbeat title when connectionSessions diverged, else undefined", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); 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 }; + vi.spyOn(Date, "now").mockReturnValue(now); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + const internal = doInstance as unknown as { + readyPushFireAt: Map; + }; // Establish ownership for both sessions first (first-sight schedules entries). sendHeartbeat(doInstance, cliWs, [ - makeSession('ses_t', 'idle', 'Stored'), - makeSession('ses_same', 'idle', 'Same'), + makeSession("ses_t", "idle", "Stored"), + makeSession("ses_same", "idle", "Same"), ]); await flushAsync(); // Diverged title for ses_t; matching title for ses_same. Force both overdue. sendHeartbeat(doInstance, cliWs, [ - makeSession('ses_t', 'idle', 'Generated'), - makeSession('ses_same', 'idle', 'Same'), + makeSession("ses_t", "idle", "Generated"), + makeSession("ses_same", "idle", "Same"), ]); - await ctx.storage.put('readyPush:ses_t', { - kiloUserId: 'usr_1', - title: 'Stored', + await ctx.storage.put("readyPush:ses_t", { + kiloUserId: "usr_1", + title: "Stored", fireAt: now - 1, attempts: 0, }); - await ctx.storage.put('readyPush:ses_same', { - kiloUserId: 'usr_1', - title: 'Same', + await ctx.storage.put("readyPush:ses_same", { + kiloUserId: "usr_1", + title: "Same", fireAt: now - 1, attempts: 0, }); - internal.readyPushFireAt.set('ses_t', now - 1); - internal.readyPushFireAt.set('ses_same', now - 1); + internal.readyPushFireAt.set("ses_t", now - 1); + internal.readyPushFireAt.set("ses_same", now - 1); claimSessionReadyPush.mockClear(); await doInstance.alarm(); - expect(claimSessionReadyPush).toHaveBeenCalledWith('usr_1', 'ses_t', 'Generated'); - expect(claimSessionReadyPush).toHaveBeenCalledWith('usr_1', 'ses_same', undefined); + expect(claimSessionReadyPush).toHaveBeenCalledWith( + "usr_1", + "ses_t", + "Generated", + ); + expect(claimSessionReadyPush).toHaveBeenCalledWith( + "usr_1", + "ses_same", + undefined, + ); }); - it('claim-then-delete: RPC rejection keeps the entry; drops after 3 attempts', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + it("claim-then-delete: RPC rejection keeps the entry; drops after 3 attempts", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_retry')]); + vi.spyOn(Date, "now").mockReturnValue(now); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_retry")]); await flushAsync(); - 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', - title: 'Test', + 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", + title: "Test", fireAt: now - 1, attempts: 0, }); - claimSessionReadyPush.mockRejectedValueOnce(new Error('DO down')); + claimSessionReadyPush.mockRejectedValueOnce(new Error("DO down")); await doInstance.alarm(); - const after1 = ctx.storage.store.get('readyPush:ses_retry') as { + const after1 = ctx.storage.store.get("readyPush:ses_retry") as { attempts: number; fireAt: number; }; expect(after1.attempts).toBe(1); // Retry must back off so attempts span real time (not fireAt already ≤ now). expect(after1.fireAt).toBe(now + 5_000); - expect(internal.readyPushFireAt.get('ses_retry')).toBe(now + 5_000); + expect(internal.readyPushFireAt.get("ses_retry")).toBe(now + 5_000); // Advance past backoff so attempt 2 is due. - vi.spyOn(Date, 'now').mockReturnValue(now + 5_000); - claimSessionReadyPush.mockRejectedValueOnce(new Error('DO down')); + vi.spyOn(Date, "now").mockReturnValue(now + 5_000); + claimSessionReadyPush.mockRejectedValueOnce(new Error("DO down")); await doInstance.alarm(); - const after2 = ctx.storage.store.get('readyPush:ses_retry') as { + const after2 = ctx.storage.store.get("readyPush:ses_retry") as { attempts: number; fireAt: number; }; expect(after2.attempts).toBe(2); expect(after2.fireAt).toBe(now + 10_000); - vi.spyOn(Date, 'now').mockReturnValue(now + 10_000); - claimSessionReadyPush.mockRejectedValueOnce(new Error('DO down')); + vi.spyOn(Date, "now").mockReturnValue(now + 10_000); + claimSessionReadyPush.mockRejectedValueOnce(new Error("DO down")); await doInstance.alarm(); - expect(ctx.storage.store.has('readyPush:ses_retry')).toBe(false); - expect(internal.readyPushFireAt.has('ses_retry')).toBe(false); + expect(ctx.storage.store.has("readyPush:ses_retry")).toBe(false); + expect(internal.readyPushFireAt.has("ses_retry")).toBe(false); }); - it('claim rejection re-arms at now+backoff, not now', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + it("claim rejection re-arms at now+backoff, not now", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_backoff')]); + vi.spyOn(Date, "now").mockReturnValue(now); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_backoff")]); await flushAsync(); - 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', - title: 'Test', + 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", + title: "Test", fireAt: now - 1, attempts: 0, }); - claimSessionReadyPush.mockRejectedValueOnce(new Error('transport blip')); + claimSessionReadyPush.mockRejectedValueOnce(new Error("transport blip")); await doInstance.alarm(); - const entry = ctx.storage.store.get('readyPush:ses_backoff') as { + const entry = ctx.storage.store.get("readyPush:ses_backoff") as { attempts: number; fireAt: number; }; expect(entry.attempts).toBe(1); expect(entry.fireAt).toBe(now + 5_000); - expect(internal.readyPushFireAt.get('ses_backoff')).toBe(now + 5_000); + expect(internal.readyPushFireAt.get("ses_backoff")).toBe(now + 5_000); // Same instant must not exhaust further attempts. claimSessionReadyPush.mockClear(); - claimSessionReadyPush.mockRejectedValue(new Error('still down')); + claimSessionReadyPush.mockRejectedValue(new Error("still down")); await doInstance.alarm(); expect(claimSessionReadyPush).not.toHaveBeenCalled(); expect( - (ctx.storage.store.get('readyPush:ses_backoff') as { attempts: number }).attempts + (ctx.storage.store.get("readyPush:ses_backoff") as { attempts: number }) + .attempts, ).toBe(1); }); - it('does not refire after a successful claim', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + it("does not refire after a successful claim", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_ok')]); + vi.spyOn(Date, "now").mockReturnValue(now); + const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_ok")]); await flushAsync(); - 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', - title: 'Test', + 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", + title: "Test", fireAt: now - 1, attempts: 0, }); await doInstance.alarm(); expect(claimSessionReadyPush).toHaveBeenCalledTimes(1); - expect(ctx.storage.store.has('readyPush:ses_ok')).toBe(false); + expect(ctx.storage.store.has("readyPush:ses_ok")).toBe(false); claimSessionReadyPush.mockClear(); await doInstance.alarm(); expect(claimSessionReadyPush).not.toHaveBeenCalled(); }); - it('KV is source of truth for alarm even without the mirror', async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + it("KV is source of truth for alarm even without the mirror", async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = + setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); + vi.spyOn(Date, "now").mockReturnValue(now); // Ensure ensureState runs so DO is live - addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); - await ctx.storage.put('readyPush:ses_kv_only', { - kiloUserId: 'usr_1', - title: 'KV', + addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + await ctx.storage.put("readyPush:ses_kv_only", { + kiloUserId: "usr_1", + title: "KV", fireAt: now - 1, 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(); - expect(claimSessionReadyPush).toHaveBeenCalledWith('usr_1', 'ses_kv_only', undefined); - expect(ctx.storage.store.has('readyPush:ses_kv_only')).toBe(false); + expect(claimSessionReadyPush).toHaveBeenCalledWith( + "usr_1", + "ses_kv_only", + undefined, + ); + 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", () => { const { doInstance } = setupWithIngestDO(); const client = createMockWs(); const server = createMockWs(); vi.stubGlobal( - 'WebSocketPair', + "WebSocketPair", class { 0 = client; 1 = server; - } + }, ); vi.stubGlobal( - 'Response', + "Response", class { constructor(_body?: BodyInit | null, _init?: ResponseInit) {} - } + }, ); doInstance.fetch( - new Request('http://local/cli?connectionId=cli-1&kiloUserId=usr_1', { - headers: { Upgrade: 'websocket' }, - }) + new Request("http://local/cli?connectionId=cli-1&kiloUserId=usr_1", { + headers: { Upgrade: "websocket" }, + }), ); - expect(server.deserializeAttachment()).toMatchObject({ role: 'cli', kiloUserId: 'usr_1' }); + expect(server.deserializeAttachment()).toMatchObject({ + role: "cli", + kiloUserId: "usr_1", + }); }); }); @@ -4437,135 +4870,763 @@ describe('UserConnectionDO', () => { // notifySessionRenamed + rename catch-up (lazy KV per heartbeat) // ------------------------------------------------------------------------- - describe('notifySessionRenamed', () => { - it('delivers session.renamed to the owning CLI and always persists KV', async () => { + describe("notifySessionRenamed", () => { + it("delivers session.renamed to the owning CLI and always persists KV", async () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1', [makeSession('ses_r')]); + const cliWs = addCliSocket(mockCtx, "cli-1", [makeSession("ses_r")]); // Establish ownership via heartbeat - sendHeartbeat(doInstance, cliWs, [makeSession('ses_r')]); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_r")]); - const result = await doInstance.notifySessionRenamed('ses_r', 'Renamed Title'); + 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' + (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 () => { + it("returns delivered:false when no owner but still persists KV", async () => { const { doInstance, ctx } = setup(); - const result = await doInstance.notifySessionRenamed('ses_missing', 'Offline Rename'); + const result = await doInstance.notifySessionRenamed( + "ses_missing", + "Offline Rename", + ); expect(result).toEqual({ delivered: false }); - expect(ctx.storage.store.get('rename:ses_missing')).toMatchObject({ - title: 'Offline Rename', + expect(ctx.storage.store.get("rename:ses_missing")).toMatchObject({ + title: "Offline Rename", }); }); - it('re-emits session.renamed on heartbeat title mismatch; deletes on match', async () => { + it("re-emits session.renamed on heartbeat title mismatch; deletes on match", async () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_catch', 'idle', 'Old Title')]); + const cliWs = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_catch", "idle", "Old Title"), + ]); - await doInstance.notifySessionRenamed('ses_catch', 'New Title'); + await doInstance.notifySessionRenamed("ses_catch", "New Title"); cliWs.send.mockClear(); // Mismatch → re-emit - sendHeartbeat(doInstance, cliWs, [makeSession('ses_catch', 'idle', 'Old Title')]); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_catch", "idle", "Old Title"), + ]); await flushAsync(); const reEmits = allSent(cliWs).filter( - m => m.type === 'system' && m.event === 'session.renamed' + (m) => m.type === "system" && m.event === "session.renamed", ); expect(reEmits).toHaveLength(1); - expect(reEmits[0].data).toEqual({ sessionId: 'ses_catch', title: 'New Title' }); - expect(ctx.storage.store.has('rename:ses_catch')).toBe(true); + expect(reEmits[0].data).toEqual({ + sessionId: "ses_catch", + title: "New Title", + }); + expect(ctx.storage.store.has("rename:ses_catch")).toBe(true); // Match → delete entry cliWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_catch', 'idle', 'New Title')]); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_catch", "idle", "New Title"), + ]); await flushAsync(); const afterMatch = allSent(cliWs).filter( - m => m.type === 'system' && m.event === 'session.renamed' + (m) => m.type === "system" && m.event === "session.renamed", ); expect(afterMatch).toHaveLength(0); - expect(ctx.storage.store.has('rename:ses_catch')).toBe(false); + expect(ctx.storage.store.has("rename:ses_catch")).toBe(false); }); - it('registers rename catch-up with waitUntil on heartbeat', async () => { + it("registers rename catch-up with waitUntil on heartbeat", async () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_wu')]); - await doInstance.notifySessionRenamed('ses_wu', 'Catch Up Title'); + const cliWs = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cliWs, [makeSession("ses_wu")]); + await doInstance.notifySessionRenamed("ses_wu", "Catch Up Title"); const waitUntil = ctx.waitUntil as ReturnType; waitUntil.mockClear(); cliWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_wu', 'idle', 'Stale')]); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_wu", "idle", "Stale"), + ]); expect(waitUntil).toHaveBeenCalled(); - const registered = waitUntil.mock.calls.map(c => c[0]); + const registered = waitUntil.mock.calls.map((c) => c[0]); expect( registered.some( - p => p instanceof Promise || typeof (p as PromiseLike)?.then === 'function' - ) + (p) => + p instanceof Promise || + typeof (p as PromiseLike)?.then === "function", + ), ).toBe(true); await flushAsync(); const reEmits = allSent(cliWs).filter( - m => m.type === 'system' && m.event === 'session.renamed' + (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 () => { + it("prunes rename entries older than TTL and does not re-emit them", async () => { const { doInstance, mockCtx, ctx } = setup(); const now = 1_700_000_000_000; const RENAME_ENTRY_TTL_MS = 7 * 24 * 60 * 60 * 1_000; - vi.spyOn(Date, 'now').mockReturnValue(now); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_stale', 'idle', 'Old')]); + vi.spyOn(Date, "now").mockReturnValue(now); + const cliWs = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_stale", "idle", "Old"), + ]); - await ctx.storage.put('rename:ses_stale', { - title: 'Never Applied', + await ctx.storage.put("rename:ses_stale", { + title: "Never Applied", at: now - RENAME_ENTRY_TTL_MS - 1, }); cliWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_stale', 'idle', 'Old')]); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_stale", "idle", "Old"), + ]); await flushAsync(); const reEmits = allSent(cliWs).filter( - m => m.type === 'system' && m.event === 'session.renamed' + (m) => m.type === "system" && m.event === "session.renamed", ); expect(reEmits).toHaveLength(0); - expect(ctx.storage.store.has('rename:ses_stale')).toBe(false); + expect(ctx.storage.store.has("rename:ses_stale")).toBe(false); }); - it('still re-emits a fresh rename entry within TTL', async () => { + it("still re-emits a fresh rename entry within TTL", async () => { const { doInstance, mockCtx, ctx } = setup(); const now = 1_700_000_000_000; - vi.spyOn(Date, 'now').mockReturnValue(now); - const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_fresh', 'idle', 'Old')]); + vi.spyOn(Date, "now").mockReturnValue(now); + const cliWs = addCliSocket(mockCtx, "cli-1"); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_fresh", "idle", "Old"), + ]); - await ctx.storage.put('rename:ses_fresh', { - title: 'Within TTL', + await ctx.storage.put("rename:ses_fresh", { + title: "Within TTL", at: now - 60_000, }); cliWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession('ses_fresh', 'idle', 'Old')]); + sendHeartbeat(doInstance, cliWs, [ + makeSession("ses_fresh", "idle", "Old"), + ]); await flushAsync(); const reEmits = allSent(cliWs).filter( - m => m.type === 'system' && m.event === 'session.renamed' + (m) => m.type === "system" && m.event === "session.renamed", ); expect(reEmits).toHaveLength(1); - expect(reEmits[0].data).toEqual({ sessionId: 'ses_fresh', title: 'Within TTL' }); - expect(ctx.storage.store.has('rename:ses_fresh')).toBe(true); + 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 }, + }); + await flushAsync(); + + // 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 }, + }); + await flushAsync(); + + // 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("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 }, + }); + await flushAsync(); + + 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. + 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")]); + 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", + }); + await flushAsync(); + + // 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")]); + 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", + }); + await flushAsync(); + + // 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")]); + 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. + sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { durable: true }, + }); + await flushAsync(); + + // 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, + }); + await flushAsync(); + + // 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")]); + 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")]); + 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"); }); }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 1c092be5e1a..9abb803b913 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -1,7 +1,7 @@ -import { DurableObject } from 'cloudflare:workers'; +import { DurableObject } from "cloudflare:workers"; -import type { Env } from '../env'; -import { getSessionIngestDO } from './SessionIngestDO'; +import type { Env } from "../env"; +import { getSessionIngestDO } from "./SessionIngestDO"; import { CLIOutboundMessageSchema, type CLIInboundMessage, @@ -10,7 +10,7 @@ import { SessionEventPayloadSchema, type WebInboundMessage, WebOutboundMessageSchema, -} from '../types/user-connection-protocol'; +} from "../types/user-connection-protocol"; type HeartbeatSession = { id: string; @@ -31,7 +31,7 @@ type ConnectionCapabilities = { attachments?: boolean }; type WSAttachment = | { - role: 'cli'; + role: "cli"; connectionId: string; sessions: HeartbeatSession[]; // Undefined means no protocolVersion has been reported yet — either the @@ -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. @@ -73,91 +78,95 @@ export const MAX_CATALOG_RESULT_BYTES = 512 * 1024; // 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([ - 'send_message', - 'interrupt', - 'question_reply', - 'question_reject', - 'permission_respond', - 'suggestion_accept', - 'suggestion_dismiss', - 'list_models', - 'list_commands', - 'send_command', - 'create_session', - 'exit_cli', + "send_message", + "interrupt", + "question_reply", + "question_reject", + "permission_respond", + "suggestion_accept", + "suggestion_dismiss", + "list_models", + "list_commands", + "send_command", + "create_session", + "exit_cli", ]); // In-flight dedupe and 512 KiB response cap apply to these catalog-style reads. -const CATALOG_DEDUPE_COMMANDS: ReadonlySet = new Set(['list_models', 'list_commands']); +const CATALOG_DEDUPE_COMMANDS: ReadonlySet = new Set([ + "list_models", + "list_commands", +]); // Operations that older CLIs reject with a precise "unknown command: " // string. Only these commands get mapped to a structured CLI_UPGRADE_REQUIRED // response; any other CLI error is preserved verbatim. const CLI_UPGRADE_REQUIRED_COMMANDS: ReadonlySet = new Set([ - 'list_commands', - 'send_command', - 'create_session', - 'exit_cli', + "list_commands", + "send_command", + "create_session", + "exit_cli", ]); const SESSION_OWNER_CHANGED_ERROR = { - source: 'relay', - code: 'SESSION_OWNER_CHANGED', - message: 'Session owner changed', + source: "relay", + code: "SESSION_OWNER_CHANGED", + message: "Session owner changed", }; const CATALOG_TOO_LARGE_ERROR = { - source: 'relay', - code: 'CATALOG_TOO_LARGE', - message: 'Model catalog response is too large', + source: "relay", + code: "CATALOG_TOO_LARGE", + message: "Model catalog response is too large", }; const CATALOG_REQUEST_PENDING_ERROR = { - source: 'relay', - code: 'CATALOG_REQUEST_PENDING', - message: 'Model catalog request already pending', + source: "relay", + code: "CATALOG_REQUEST_PENDING", + message: "Model catalog request already pending", }; const PENDING_COMMAND_LIMIT_ERROR = { - source: 'relay', - code: 'PENDING_COMMAND_LIMIT', - message: 'Too many pending commands', + source: "relay", + code: "PENDING_COMMAND_LIMIT", + message: "Too many pending commands", }; const COMMAND_EXPIRED_ERROR = { - source: 'relay', - code: 'COMMAND_EXPIRED', - message: 'Command expired', + source: "relay", + code: "COMMAND_EXPIRED", + message: "Command expired", }; const COMMAND_NOT_ALLOWED_ERROR = { - source: 'relay', - code: 'COMMAND_NOT_ALLOWED', - message: 'Command is not allowed', + source: "relay", + code: "COMMAND_NOT_ALLOWED", + message: "Command is not allowed", }; const INVALID_COMMAND_ERROR = { - source: 'relay', - code: 'INVALID_COMMAND', - message: 'Invalid command', + source: "relay", + code: "INVALID_COMMAND", + message: "Invalid command", }; const CLI_UPGRADE_REQUIRED_SLASH_ERROR = { - source: 'relay', - code: 'CLI_UPGRADE_REQUIRED', - message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', + source: "relay", + code: "CLI_UPGRADE_REQUIRED", + message: + "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", }; const CLI_UPGRADE_REQUIRED_CREATE_SESSION_ERROR = { - source: 'relay', - code: 'CLI_UPGRADE_REQUIRED', + source: "relay", + code: "CLI_UPGRADE_REQUIRED", message: - 'Creating remote sessions from mobile requires a newer Kilo CLI. Update Kilo CLI and reconnect.', + "Creating remote sessions from mobile requires a newer Kilo CLI. Update Kilo CLI and reconnect.", }; const CLI_COMMAND_ERROR = { - source: 'cli', - message: 'Command failed', + source: "cli", + message: "Command failed", }; type ReadyPushEntry = { @@ -172,8 +181,25 @@ type RenameEntry = { at: number; }; -const READY_PUSH_KEY_PREFIX = 'readyPush:'; -const RENAME_KEY_PREFIX = 'rename:'; +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; @@ -195,7 +221,10 @@ export class UserConnectionDO extends DurableObject { // Protocol version per CLI connection (from heartbeat); absent = legacy CLI private connectionProtocolVersion = new Map(); // Capabilities per CLI connection (from heartbeat); absent = legacy CLI - private connectionCapabilities = new Map(); + private connectionCapabilities = new Map< + string, + ConnectionCapabilities | undefined + >(); // Pending command responses: correlationId → originating web socket private pendingCommands = new Map< string, @@ -217,6 +246,10 @@ 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(); + private stateReconstructed = false; private ensureState(): void { @@ -230,9 +263,10 @@ export class UserConnectionDO extends DurableObject { const attachment = ws.deserializeAttachment() as WSAttachment | null; if (!attachment) continue; - if (attachment.role === 'cli') { + if (attachment.role === "cli") { cliCount++; - const { connectionId, sessions, protocolVersion, capabilities } = attachment; + const { connectionId, sessions, protocolVersion, capabilities } = + attachment; this.connectionSessions.set(connectionId, sessions); this.connectionProtocolVersion.set(connectionId, protocolVersion); this.connectionCapabilities.set(connectionId, capabilities); @@ -255,7 +289,7 @@ export class UserConnectionDO extends DurableObject { } } - console.log('State reconstructed after hibernation', { + console.log("State reconstructed after hibernation", { cliSockets: cliCount, webSockets: webCount, sessions: sessionCount, @@ -268,63 +302,73 @@ export class UserConnectionDO extends DurableObject { } fetch(request: Request): Response { - const upgradeHeader = request.headers.get('Upgrade'); - if (upgradeHeader !== 'websocket') { - return new Response('Expected WebSocket upgrade', { status: 426 }); + const upgradeHeader = request.headers.get("Upgrade"); + if (upgradeHeader !== "websocket") { + return new Response("Expected WebSocket upgrade", { status: 426 }); } this.ensureState(); const url = new URL(request.url); - const role = url.pathname.endsWith('/cli') ? 'cli' : 'web'; + const role = url.pathname.endsWith("/cli") ? "cli" : "web"; const pair = new WebSocketPair(); const [client, server] = [pair[0], pair[1]]; - const connectionId = url.searchParams.get('connectionId') ?? crypto.randomUUID(); + const connectionId = + url.searchParams.get("connectionId") ?? crypto.randomUUID(); - if (role === 'cli') { + if (role === "cli") { // Close any stale socket from a previous connection with the same ID (CLI reconnect) const reconnect = this.closeStaleSocket(connectionId); - const kiloUserId = url.searchParams.get('kiloUserId') ?? undefined; - const attachment: WSAttachment = { role: 'cli', connectionId, sessions: [], kiloUserId }; - this.ctx.acceptWebSocket(server, ['cli']); + const kiloUserId = url.searchParams.get("kiloUserId") ?? undefined; + const attachment: WSAttachment = { + role: "cli", + connectionId, + sessions: [], + kiloUserId, + }; + this.ctx.acceptWebSocket(server, ["cli"]); server.serializeAttachment(attachment); const now = Date.now(); this.lastHeartbeatAt.set(connectionId, now); this.scheduleNextAlarm(now); - console.log('CLI socket connected', { + console.log("CLI socket connected", { connectionId, reconnect, - totalCliSockets: this.ctx.getWebSockets('cli').length, + totalCliSockets: this.ctx.getWebSockets("cli").length, }); if (!reconnect) { this.broadcastToWeb({ - type: 'system', - event: 'cli.connected', + type: "system", + event: "cli.connected", data: { connectionId }, }); } } else { this.replaceWebSocket(connectionId); - const attachment: WSAttachment = { role: 'web', connectionId, subscribedSessions: [] }; - this.ctx.acceptWebSocket(server, ['web']); + const attachment: WSAttachment = { + role: "web", + connectionId, + subscribedSessions: [], + }; + this.ctx.acceptWebSocket(server, ["web"]); server.serializeAttachment(attachment); const sessions = this.aggregateSessions(); - console.log('Web socket connected', { + console.log("Web socket connected", { connectionId, - totalWebSockets: this.ctx.getWebSockets('web').length, + totalWebSockets: this.ctx.getWebSockets("web").length, activeSessions: sessions.length, }); this.sendToWeb(server, { - type: 'system', - event: 'sessions.list', + type: "system", + event: "sessions.list", data: { sessions }, }); } @@ -332,22 +376,27 @@ export class UserConnectionDO extends DurableObject { return new Response(null, { status: 101, webSocket: client }); } - webSocketMessage(ws: WebSocket, message: string | ArrayBuffer): void { + webSocketMessage( + ws: WebSocket, + message: string | ArrayBuffer, + ): void { this.ensureState(); const attachment = ws.deserializeAttachment() as WSAttachment | null; if (!attachment) { - console.warn('WebSocket message from socket with no attachment'); + console.warn("WebSocket message from socket with no attachment"); return; } - const raw = typeof message === 'string' ? message : new TextDecoder().decode(message); - const binaryByteCount = typeof message === 'string' ? undefined : message.byteLength; + const raw = + typeof message === "string" ? message : new TextDecoder().decode(message); + const binaryByteCount = + typeof message === "string" ? undefined : message.byteLength; let parsed: unknown; try { parsed = JSON.parse(raw); } catch { - console.warn('Failed to parse WebSocket message as JSON', { + console.warn("Failed to parse WebSocket message as JSON", { role: attachment.role, connectionId: attachment.connectionId, byteCount: binaryByteCount ?? new TextEncoder().encode(raw).byteLength, @@ -355,8 +404,14 @@ export class UserConnectionDO extends DurableObject { return; } - if (attachment.role === 'cli') { - this.handleCliMessage(ws, attachment, parsed, raw, binaryByteCount); + if (attachment.role === "cli") { + this.handleCliMessage( + ws, + attachment, + parsed, + raw, + binaryByteCount, + ); } else if (!attachment.replaced) { this.handleWebMessage(ws, attachment, parsed); } @@ -366,14 +421,14 @@ export class UserConnectionDO extends DurableObject { ws: WebSocket, _code: number, _reason: string, - _wasClean: boolean + _wasClean: boolean, ): void | Promise { this.ensureState(); const attachment = ws.deserializeAttachment() as WSAttachment | null; if (!attachment) return; - if (attachment.role === 'cli') { + if (attachment.role === "cli") { // Await attention resets so cli.disconnected is not broadcast until the // stored status write has committed (mobile history refetch races otherwise). return this.handleCliDisconnect(ws, attachment); @@ -383,11 +438,11 @@ export class UserConnectionDO extends DurableObject { webSocketError(ws: WebSocket): void | Promise { const attachment = ws.deserializeAttachment() as WSAttachment | null; - console.error('WebSocket error', { - role: attachment?.role ?? 'unknown', - connectionId: attachment?.connectionId ?? 'unknown', + console.error("WebSocket error", { + role: attachment?.role ?? "unknown", + connectionId: attachment?.connectionId ?? "unknown", }); - return this.webSocketClose(ws, 0, '', false); + return this.webSocketClose(ws, 0, "", false); } async alarm(): Promise { @@ -406,11 +461,13 @@ export class UserConnectionDO extends DurableObject { for (const connectionId of staleConnectionIds) { // Find and close the stale CLI WebSocket - for (const ws of this.ctx.getWebSockets('cli')) { + 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 }); - ws.close(4408, 'heartbeat timeout'); + if (att?.role === "cli" && att.connectionId === connectionId) { + console.log("Closing stale CLI connection (heartbeat timeout)", { + connectionId, + }); + ws.close(4408, "heartbeat timeout"); break; } } @@ -419,6 +476,7 @@ export class UserConnectionDO extends DurableObject { } this.scheduleNextAlarm(now); + this.scheduleDurablePendingAlarm(); } // --------------------------------------------------------------------------- @@ -427,50 +485,63 @@ export class UserConnectionDO extends DurableObject { private handleCliMessage( ws: WebSocket, - attachment: WSAttachment & { role: 'cli' }, + attachment: WSAttachment & { role: "cli" }, parsed: unknown, raw: string, - binaryByteCount: number | undefined + binaryByteCount: number | undefined, ): void { const result = CLIOutboundMessageSchema.safeParse(parsed); if (!result.success) { - console.warn('CLI message parse failed', { - role: 'cli', + console.warn("CLI message parse failed", { + 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; } const msg = result.data; switch (msg.type) { - case 'heartbeat': + case "heartbeat": this.handleHeartbeat( ws, attachment, msg.sessions, msg.protocolVersion, msg.capabilities, - msg.instance + msg.instance, ); break; - case 'event': - this.handleCliEvent(msg.sessionId, msg.parentSessionId, msg.event, msg.data); + case "event": + this.handleCliEvent( + msg.sessionId, + msg.parentSessionId, + msg.event, + msg.data, + ); break; - case 'response': - this.handleCliResponse(ws, msg.id, msg.result, msg.error); + case "response": + // Extend the DO lifetime for the async durable write (fix: reply + // must persist before hibernation, but the surrounding handler stays + // synchronous for existing callers). + this.ctx.waitUntil( + this.handleCliResponse(ws, msg.id, msg.result, msg.error), + ); break; } } private handleHeartbeat( ws: WebSocket, - attachment: WSAttachment & { role: 'cli' }, + attachment: WSAttachment & { role: "cli" }, sessions: HeartbeatSession[], protocolVersion: string | undefined, capabilities: ConnectionCapabilities | undefined, - instance: Instance | undefined + instance: Instance | undefined, ): void { const { connectionId } = attachment; const now = Date.now(); @@ -481,9 +552,12 @@ export class UserConnectionDO extends DurableObject { // Remove sessions this connection previously owned but no longer reports const previousSessions = this.connectionSessions.get(connectionId) ?? []; - const currentIds = new Set(sessions.map(s => s.id)); + const currentIds = new Set(sessions.map((s) => s.id)); for (const prev of previousSessions) { - if (!currentIds.has(prev.id) && this.sessionOwners.get(prev.id) === connectionId) { + if ( + !currentIds.has(prev.id) && + this.sessionOwners.get(prev.id) === connectionId + ) { this.sessionOwners.delete(prev.id); this.failPendingCommandsForOwnerChange(prev.id, undefined); } @@ -500,7 +574,11 @@ export class UserConnectionDO extends DurableObject { // push (Decision 9). The durable claim in SessionIngestDO makes reconnect // re-sights no-ops once the delayed claim lands. if (!previousOwner && !session.parentSessionId && attachment.kiloUserId) { - this.scheduleSessionReadyPush(attachment.kiloUserId, session.id, session.title); + this.scheduleSessionReadyPush( + attachment.kiloUserId, + session.id, + session.title, + ); } this.sessionOwners.set(session.id, connectionId); } @@ -511,23 +589,26 @@ export class UserConnectionDO extends DurableObject { // waitUntil so the DO does not hibernate before catch-up finishes. this.ctx.waitUntil( this.catchUpPendingRenames(ws, sessions).catch((error: unknown) => { - console.error('Failed to catch up pending renames (non-fatal)', { + console.error("Failed to catch up pending renames (non-fatal)", { error: error instanceof Error ? error.message : String(error), }); - }) + }), ); // Replay existing subscriptions for sessions newly owned by this CLI - const previousIds = new Set(previousSessions.map(s => s.id)); + const previousIds = new Set(previousSessions.map((s) => s.id)); for (const session of sessions) { - if (!previousIds.has(session.id) && this.webSubscriptions.has(session.id)) { - this.sendToCli(ws, { type: 'subscribe', sessionId: session.id }); + if ( + !previousIds.has(session.id) && + this.webSubscriptions.has(session.id) + ) { + this.sendToCli(ws, { type: "subscribe", sessionId: session.id }); } } // Persist to attachment for hibernation recovery const updatedAttachment: WSAttachment = { - role: 'cli', + role: "cli", connectionId, sessions, protocolVersion, @@ -541,20 +622,20 @@ export class UserConnectionDO extends DurableObject { // and non-subscribers both receive it: a removed session id is detectable // from its absence in the payload, so no subscriber special-case is needed. this.broadcastToWeb({ - type: 'system', - event: 'sessions.heartbeat', + type: "system", + event: "sessions.heartbeat", data: { connectionId, protocolVersion, capabilities, - sessions: sessions.map(session => ({ + sessions: sessions.map((session) => ({ ...session, ...(capabilities ? { capabilities } : {}), })), }, }); - this.sendToCli(ws, { type: 'heartbeat_ack' }); + this.sendToCli(ws, { type: "heartbeat_ack" }); } /** @@ -563,7 +644,11 @@ export class UserConnectionDO extends DurableObject { * No-op when a pending entry already exists (mirror or KV) so reconnect * first-sights do not slide fireAt or reset attempts. */ - private scheduleSessionReadyPush(kiloUserId: string, sessionId: string, title: string): void { + private scheduleSessionReadyPush( + kiloUserId: string, + sessionId: string, + title: string, + ): void { if (this.readyPushFireAt.has(sessionId)) return; const key = `${READY_PUSH_KEY_PREFIX}${sessionId}`; @@ -585,11 +670,11 @@ export class UserConnectionDO extends DurableObject { this.readyPushFireAt.set(sessionId, fireAt); this.scheduleNextAlarm(Date.now()); })().catch((error: unknown) => { - console.error('Failed to schedule session-ready push (non-fatal)', { + console.error("Failed to schedule session-ready push (non-fatal)", { sessionId, error: error instanceof Error ? error.message : String(error), }); - }) + }), ); } @@ -601,7 +686,7 @@ export class UserConnectionDO extends DurableObject { private async claimSessionReadyPush( kiloUserId: string, sessionId: string, - title?: string + title?: string, ): Promise { const stub = getSessionIngestDO(this.env, { kiloUserId, sessionId }); await stub.claimSessionReadyPush(kiloUserId, sessionId, title); @@ -610,16 +695,18 @@ 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', { + console.error("Failed to list readyPush entries", { error: error instanceof Error ? error.message : String(error), }); return; } for (const [key, entry] of pending) { - if (!entry || typeof entry.fireAt !== 'number') continue; + if (!entry || typeof entry.fireAt !== "number") continue; if (entry.fireAt > now) continue; const sessionId = key.slice(READY_PUSH_KEY_PREFIX.length); @@ -629,7 +716,9 @@ export class UserConnectionDO extends DurableObject { // else undefined so notifications falls back to the live DB title. const heartbeatTitle = this.findHeartbeatTitle(sessionId); const freshest = - heartbeatTitle !== undefined && heartbeatTitle !== entry.title ? heartbeatTitle : undefined; + heartbeatTitle !== undefined && heartbeatTitle !== entry.title + ? heartbeatTitle + : undefined; try { await this.claimSessionReadyPush(entry.kiloUserId, sessionId, freshest); @@ -640,7 +729,7 @@ export class UserConnectionDO extends DurableObject { if (attempts >= READY_PUSH_MAX_ATTEMPTS) { await this.ctx.storage.delete(key); this.readyPushFireAt.delete(sessionId); - console.error('Dropping session-ready push after max attempts', { + console.error("Dropping session-ready push after max attempts", { sessionId, attempts, error: error instanceof Error ? error.message : String(error), @@ -649,7 +738,7 @@ export class UserConnectionDO extends DurableObject { const fireAt = now + READY_PUSH_RETRY_BACKOFF_MS; await this.ctx.storage.put(key, { ...entry, attempts, fireAt }); this.readyPushFireAt.set(sessionId, fireAt); - console.error('Session-ready push claim failed; will retry', { + console.error("Session-ready push claim failed; will retry", { sessionId, attempts, fireAt, @@ -669,7 +758,10 @@ export class UserConnectionDO extends DurableObject { return undefined; } - private async catchUpPendingRenames(ws: WebSocket, sessions: HeartbeatSession[]): Promise { + private async catchUpPendingRenames( + ws: WebSocket, + sessions: HeartbeatSession[], + ): Promise { const now = Date.now(); for (const session of sessions) { const key = `${RENAME_KEY_PREFIX}${session.id}`; @@ -677,7 +769,7 @@ export class UserConnectionDO extends DurableObject { try { entry = await this.ctx.storage.get(key); } catch (error: unknown) { - console.error('Failed to read rename catch-up entry', { + console.error("Failed to read rename catch-up entry", { sessionId: session.id, error: error instanceof Error ? error.message : String(error), }); @@ -686,11 +778,14 @@ export class UserConnectionDO extends DurableObject { if (!entry) continue; // Prune entries past TTL so finished/ignored renames do not re-emit forever. - if (typeof entry.at === 'number' && now - entry.at > RENAME_ENTRY_TTL_MS) { + if ( + typeof entry.at === "number" && + now - entry.at > RENAME_ENTRY_TTL_MS + ) { try { await this.ctx.storage.delete(key); } catch (error: unknown) { - console.error('Failed to delete expired rename entry', { + console.error("Failed to delete expired rename entry", { sessionId: session.id, error: error instanceof Error ? error.message : String(error), }); @@ -702,7 +797,7 @@ export class UserConnectionDO extends DurableObject { try { await this.ctx.storage.delete(key); } catch (error: unknown) { - console.error('Failed to delete matched rename entry', { + console.error("Failed to delete matched rename entry", { sessionId: session.id, error: error instanceof Error ? error.message : String(error), }); @@ -711,8 +806,8 @@ export class UserConnectionDO extends DurableObject { } this.sendToCli(ws, { - type: 'system', - event: 'session.renamed', + type: "system", + event: "session.renamed", data: { sessionId: session.id, title: entry.title }, }); } @@ -722,10 +817,12 @@ export class UserConnectionDO extends DurableObject { sessionId: string, parentSessionId: string | undefined, event: string, - data: unknown + data: unknown, ): void { const childSubs = this.webSubscriptions.get(sessionId); - const parentSubs = parentSessionId ? this.webSubscriptions.get(parentSessionId) : undefined; + const parentSubs = parentSessionId + ? this.webSubscriptions.get(parentSessionId) + : undefined; if (!childSubs && !parentSubs) return; const merged = new Set(); @@ -734,7 +831,7 @@ export class UserConnectionDO extends DurableObject { if (merged.size === 0) return; const msg: WebInboundMessage = { - type: 'event', + type: "event", sessionId, ...(parentSessionId ? { parentSessionId } : {}), event, @@ -745,63 +842,201 @@ 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); + error: unknown, + ): Promise { + let entry = this.pendingCommands.get(id); + let rehydrated = false; + + // Step 24: on in-memory miss, try to load the durable entry (D8 case 1). + if (!entry) { + const durable = await this.getDurablePendingCommand(id); + if (!durable || durable.state !== "pending") return; + + // Validate the responding CLI by attachment connectionId. + const respondingAttachment = + respondingWs.deserializeAttachment() as WSAttachment | null; + if ( + respondingAttachment?.role !== "cli" || + respondingAttachment.connectionId !== durable.targetConnectionId + ) { + return; + } + + // Resolve the originating web socket by persisted webConnectionId. + const webWs = this.findWebByConnectionId(durable.webConnectionId); + + // 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; + } else if (entry.targetCliWs !== respondingWs) { + // Strict socket-identity check for an entry still in memory. + return; + } + // 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) { - this.sendToWeb(entry.ws, { - type: 'response', - id: entry.originalId, + // For non-rehydrated entries, send the oversized error synchronously + // before the durable write so the mock harness sees the send. + if (!rehydrated) { + this.pendingCommands.delete(id); + this.sendToWeb(entry.ws, { + type: "response", + id: entry.originalId, + error: CATALOG_TOO_LARGE_ERROR, + }); + } + + // Persist the terminal outcome durably. + const catWebConnectionId = rehydrated + ? (await this.getDurablePendingCommand(id))!.webConnectionId + : (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: catWebConnectionId, + state: "done", error: CATALOG_TOO_LARGE_ERROR, - }); + } satisfies PendingCommandEntry); + + // For a rehydrated entry, re-resolve and send now. + if (rehydrated) { + 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 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 ( + typeof error === "string" && + CLI_UPGRADE_REQUIRED_COMMANDS.has(entry.command) + ) { if (error === `unknown command: ${entry.command}`) { structuredError = - entry.command === 'create_session' + entry.command === "create_session" ? CLI_UPGRADE_REQUIRED_CREATE_SESSION_ERROR : CLI_UPGRADE_REQUIRED_SLASH_ERROR; } else { stringError = error; } - } else if (typeof error === 'string') { + } 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, + const terminalError = + structuredError ?? + stringError ?? + (sanitizeAsFailed ? CLI_COMMAND_ERROR : undefined); + + // Step 26a: write terminal outcome durably, always. + // For non-rehydrated entries, send the live response synchronously + // before the durable write so the response reaches the client without + // yielding to an async storage put. + if (!rehydrated) { + this.pendingCommands.delete(id); + 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 } + : {}), + }); + } + + // Persist the durable entry. Storage write is extended by the + // surrounding ctx.waitUntil; the live response has already been + // delivered for the non-rehydrated case above. + const webConnectionId = rehydrated + ? (await this.getDurablePendingCommand(id))?.webConnectionId + : (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", ...(result !== undefined ? { result } : {}), - ...(structuredError !== null - ? { error: structuredError } - : stringError !== null - ? { error: stringError } - : sanitizeAsFailed - ? { error: CLI_COMMAND_ERROR } - : {}), - }); + ...(terminalError !== undefined ? { error: terminalError } : {}), + } satisfies PendingCommandEntry); + + // For a rehydrated entry, re-resolve the web socket by the persisted + // id and send the response now (the non-rehydrated path already sent + // it synchronously above). + if (rehydrated) { + 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 } + : {}), + }); + } + } } // --------------------------------------------------------------------------- @@ -810,39 +1045,39 @@ export class UserConnectionDO extends DurableObject { private handleWebMessage( ws: WebSocket, - attachment: WSAttachment & { role: 'web' }, - parsed: unknown + attachment: WSAttachment & { role: "web" }, + parsed: unknown, ): void { const result = WebOutboundMessageSchema.safeParse(parsed); if (!result.success) { - console.warn('Invalid web message', { + console.warn("Invalid web message", { connectionId: attachment.connectionId, - errors: result.error.issues.map(i => i.message), + errors: result.error.issues.map((i) => i.message), }); return; } const msg = result.data; switch (msg.type) { - case 'subscribe': + case "subscribe": this.handleWebSubscribe(ws, attachment, msg.sessionId); break; - case 'unsubscribe': + case "unsubscribe": this.handleWebUnsubscribe(ws, attachment, msg.sessionId); break; - case 'command': + case "command": this.handleWebCommand(ws, msg); break; - case 'ping': - this.sendToWeb(ws, { type: 'pong', nonce: msg.nonce }); + case "ping": + this.sendToWeb(ws, { type: "pong", nonce: msg.nonce }); break; } } private handleWebSubscribe( ws: WebSocket, - attachment: WSAttachment & { role: 'web' }, - sessionId: string + attachment: WSAttachment & { role: "web" }, + sessionId: string, ): void { let subs = this.webSubscriptions.get(sessionId); if (!subs) { @@ -858,8 +1093,8 @@ export class UserConnectionDO extends DurableObject { } this.sendToWeb(ws, { - type: 'system', - event: 'sessions.list', + type: "system", + event: "sessions.list", data: { sessions: this.aggregateSessions() }, }); @@ -869,18 +1104,18 @@ export class UserConnectionDO extends DurableObject { // so it wasn't reported in the most recent heartbeat. const cliWs = this.findCliForSession(sessionId); if (cliWs) { - this.sendToCli(cliWs, { type: 'subscribe', sessionId }); + this.sendToCli(cliWs, { type: "subscribe", sessionId }); } else { - for (const ws of this.ctx.getWebSockets('cli')) { - this.sendToCli(ws, { type: 'subscribe', sessionId }); + for (const ws of this.ctx.getWebSockets("cli")) { + this.sendToCli(ws, { type: "subscribe", sessionId }); } } } private handleWebUnsubscribe( ws: WebSocket, - attachment: WSAttachment & { role: 'web' }, - sessionId: string + attachment: WSAttachment & { role: "web" }, + sessionId: string, ): void { const subs = this.webSubscriptions.get(sessionId); if (subs) { @@ -891,7 +1126,7 @@ export class UserConnectionDO extends DurableObject { this.webSubscriptions.delete(sessionId); const cliWs = this.findCliForSession(sessionId); if (cliWs) { - this.sendToCli(cliWs, { type: 'unsubscribe', sessionId }); + this.sendToCli(cliWs, { type: "unsubscribe", sessionId }); } } } @@ -906,7 +1141,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); @@ -915,7 +1157,7 @@ export class UserConnectionDO extends DurableObject { // ownership, allocate a pending slot, or forward to the CLI. if (!ALLOWED_VIEWER_COMMANDS.has(msg.command)) { this.sendToWeb(ws, { - type: 'response', + type: "response", id: msg.id, error: COMMAND_NOT_ALLOWED_ERROR, }); @@ -923,17 +1165,17 @@ export class UserConnectionDO extends DurableObject { } if ( - msg.command === 'exit_cli' && + msg.command === "exit_cli" && (!msg.sessionId || - typeof msg.data !== 'object' || + typeof msg.data !== "object" || msg.data === null || Array.isArray(msg.data) || Object.keys(msg.data).length !== 1 || - !Object.hasOwn(msg.data, 'protocolVersion') || - Reflect.get(msg.data, 'protocolVersion') !== 1) + !Object.hasOwn(msg.data, "protocolVersion") || + Reflect.get(msg.data, "protocolVersion") !== 1) ) { this.sendToWeb(ws, { - type: 'response', + type: "response", id: msg.id, error: INVALID_COMMAND_ERROR, }); @@ -945,9 +1187,12 @@ export class UserConnectionDO extends DurableObject { if (msg.sessionId && msg.connectionId) { targetCli = this.findCliByConnectionId(msg.connectionId); - if (this.sessionOwners.get(msg.sessionId) !== msg.connectionId || !targetCli) { + if ( + this.sessionOwners.get(msg.sessionId) !== msg.connectionId || + !targetCli + ) { this.sendToWeb(ws, { - type: 'response', + type: "response", id: msg.id, error: SESSION_OWNER_CHANGED_ERROR, }); @@ -959,49 +1204,194 @@ export class UserConnectionDO extends DurableObject { targetCli = this.findCliForSession(msg.sessionId); } else { // Fall back to first available CLI - const cliSockets = this.ctx.getWebSockets('cli'); + const cliSockets = this.ctx.getWebSockets("cli"); targetCli = cliSockets[0]; } 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; } - const targetAttachment = targetCli.deserializeAttachment() as WSAttachment | null; - if (targetAttachment?.role !== 'cli') return; + const targetAttachment = + targetCli.deserializeAttachment() as WSAttachment | null; + if (targetAttachment?.role !== "cli") return; const expectedOwnerConnectionId = 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 => + (entry) => entry.ws === ws && entry.command === msg.command && entry.sessionId === msg.sessionId && - entry.targetConnectionId === targetConnectionId + entry.targetConnectionId === targetConnectionId, ) ) { this.sendToWeb(ws, { - type: 'response', + type: "response", id: msg.id, error: CATALOG_REQUEST_PENDING_ERROR, }); return; } + // D8 mutationId dedupe: synchronous reservation prevents concurrent + // same-ID dispatch across the async storage read below. + if (msg.mutationId) { + const mutationId = msg.mutationId; + // 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', + type: "response", id: msg.id, error: PENDING_COMMAND_LIMIT_ERROR, }); 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,10 +1402,72 @@ 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", + id: correlationId, + 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(); + 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.ctx.waitUntil( + 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', + type: "command", id: correlationId, command: msg.command, data: msg.data, @@ -1029,7 +1481,7 @@ export class UserConnectionDO extends DurableObject { private async handleCliDisconnect( disconnectedWs: WebSocket, - attachment: WSAttachment & { role: 'cli' } + attachment: WSAttachment & { role: "cli" }, ): Promise { const { connectionId } = attachment; @@ -1038,17 +1490,19 @@ export class UserConnectionDO extends DurableObject { // Exclude the closing socket: under wrangler/workerd, getWebSockets() still // includes it during webSocketClose, so matching self would always look "replaced" // and skip ownership cleanup + attention reset (DEF-5 E2E failure). - const replaced = this.ctx.getWebSockets('cli').some(ws => { + const replaced = this.ctx.getWebSockets("cli").some((ws) => { if (ws === disconnectedWs) return false; const att = ws.deserializeAttachment() as WSAttachment | null; - return att?.role === 'cli' && att.connectionId === connectionId; + return att?.role === "cli" && att.connectionId === connectionId; }); // Fail pending commands that targeted this specific socket 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; } @@ -1066,10 +1520,10 @@ export class UserConnectionDO extends DurableObject { this.connectionCapabilities.delete(connectionId); this.lastHeartbeatAt.delete(connectionId); - console.log('CLI socket disconnected', { + console.log("CLI socket disconnected", { connectionId, droppedSessions: ownedSessions.size, - remainingCliSockets: this.ctx.getWebSockets('cli').length, + remainingCliSockets: this.ctx.getWebSockets("cli").length, }); // Leave webSubscriptions intact — a reconnecting CLI can resume @@ -1078,11 +1532,14 @@ export class UserConnectionDO extends DurableObject { // departure refetch observes `retry` rather than a stuck `question`. // kiloUserId comes from the CLI attachment (authenticated /user/cli route); // without it we cannot safely target rows and must no-op. - await this.resetOwnedSessionAttentionOnDisconnect(attachment.kiloUserId, ownedSessions); + await this.resetOwnedSessionAttentionOnDisconnect( + attachment.kiloUserId, + ownedSessions, + ); this.broadcastToWeb({ - type: 'system', - event: 'cli.disconnected', + type: "system", + event: "cli.disconnected", data: { connectionId }, }); } @@ -1093,29 +1550,32 @@ export class UserConnectionDO extends DurableObject { */ private async resetOwnedSessionAttentionOnDisconnect( kiloUserId: string | undefined, - ownedSessions: ReadonlySet + ownedSessions: ReadonlySet, ): Promise { if (ownedSessions.size === 0) return; if (!kiloUserId) { console.warn( - 'Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment', - { ownedSessionCount: ownedSessions.size } + "Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment", + { ownedSessionCount: ownedSessions.size }, ); return; } const results = await Promise.allSettled( - [...ownedSessions].map(async sessionId => { + [...ownedSessions].map(async (sessionId) => { const stub = getSessionIngestDO(this.env, { kiloUserId, sessionId }); await stub.resetAttentionStatusOnCliDisconnect(kiloUserId, sessionId); - }) + }), ); for (const result of results) { - if (result.status === 'rejected') { - console.error('Failed to reset attention status on CLI disconnect', { - error: result.reason instanceof Error ? result.reason.message : String(result.reason), + if (result.status === "rejected") { + console.error("Failed to reset attention status on CLI disconnect", { + error: + result.reason instanceof Error + ? result.reason.message + : String(result.reason), }); } } @@ -1123,7 +1583,8 @@ export class UserConnectionDO extends DurableObject { private handleWebDisconnect(ws: WebSocket): void { const attachment = ws.deserializeAttachment() as WSAttachment | null; - const connectionId = attachment?.role === 'web' ? attachment.connectionId : 'unknown'; + const connectionId = + attachment?.role === "web" ? attachment.connectionId : "unknown"; // Remove from all subscription sets let droppedSubscriptions = 0; @@ -1137,12 +1598,15 @@ export class UserConnectionDO extends DurableObject { // Tell owning CLI to stop forwarding const cliWs = this.findCliForSession(sessionId); if (cliWs) { - this.sendToCli(cliWs, { type: 'unsubscribe', sessionId }); + this.sendToCli(cliWs, { type: "unsubscribe", sessionId }); } } } - // 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) { @@ -1151,11 +1615,11 @@ export class UserConnectionDO extends DurableObject { } } - console.log('Web socket disconnected', { + console.log("Web socket disconnected", { connectionId, droppedSubscriptions, droppedCommands, - remainingWebSockets: this.ctx.getWebSockets('web').length, + remainingWebSockets: this.ctx.getWebSockets("web").length, }); } @@ -1188,9 +1652,9 @@ export class UserConnectionDO extends DurableObject { getConnectedInstances(): { instances: ConnectedInstanceRow[] } { this.ensureState(); const instances: ConnectedInstanceRow[] = []; - for (const ws of this.ctx.getWebSockets('cli')) { + for (const ws of this.ctx.getWebSockets("cli")) { const att = ws.deserializeAttachment() as WSAttachment | null; - if (att?.role !== 'cli' || !att.instance) continue; + if (att?.role !== "cli" || !att.instance) continue; instances.push({ connectionId: att.connectionId, name: att.instance.name, @@ -1202,11 +1666,13 @@ export class UserConnectionDO extends DurableObject { return { instances }; } - async notifySessionEvent(event: SessionEventPayload): Promise<{ delivered: number }> { + async notifySessionEvent( + event: SessionEventPayload, + ): Promise<{ delivered: number }> { this.ensureState(); const parsed = SessionEventPayloadSchema.parse(event); const msg: WebInboundMessage = { - type: 'system', + type: "system", event: parsed.type, data: parsed.data, }; @@ -1218,7 +1684,7 @@ export class UserConnectionDO extends DurableObject { ws.send(json); delivered++; } catch (err) { - console.warn('notifySessionEvent: skipping failed web socket:', err); + console.warn("notifySessionEvent: skipping failed web socket:", err); } } return { delivered }; @@ -1228,7 +1694,10 @@ export class UserConnectionDO extends DurableObject { * Persist a web→CLI rename under KV (offline catch-up) and deliver * `session.renamed` to the owning CLI socket when one is connected. */ - async notifySessionRenamed(sessionId: string, title: string): Promise<{ delivered: boolean }> { + async notifySessionRenamed( + sessionId: string, + title: string, + ): Promise<{ delivered: boolean }> { this.ensureState(); await this.ctx.storage.put(`${RENAME_KEY_PREFIX}${sessionId}`, { title, @@ -1245,8 +1714,8 @@ export class UserConnectionDO extends DurableObject { } this.sendToCli(cliWs, { - type: 'system', - event: 'session.renamed', + type: "system", + event: "session.renamed", data: { sessionId, title }, }); return { delivered: true }; @@ -1265,7 +1734,7 @@ export class UserConnectionDO extends DurableObject { try { ws.send(JSON.stringify(msg)); } catch (err) { - console.warn('sendToCli failed:', err); + console.warn("sendToCli failed:", err); } } @@ -1273,7 +1742,7 @@ export class UserConnectionDO extends DurableObject { try { ws.send(JSON.stringify(msg)); } catch (err) { - console.warn('sendToWeb failed:', err); + console.warn("sendToWeb failed:", err); } } @@ -1284,7 +1753,7 @@ export class UserConnectionDO extends DurableObject { try { ws.send(json); } catch (err) { - console.warn('broadcastToWeb: skipping failed socket:', err); + console.warn("broadcastToWeb: skipping failed socket:", err); } } } @@ -1292,13 +1761,13 @@ export class UserConnectionDO extends DurableObject { /** Close a stale CLI socket that has the same connectionId (from a previous connection). Returns true if one was found. */ private closeStaleSocket(connectionId: string): boolean { - for (const ws of this.ctx.getWebSockets('cli')) { + 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 socket for reconnect', { connectionId }); + if (att?.role === "cli" && att.connectionId === connectionId) { + console.log("Closing stale CLI socket for reconnect", { connectionId }); this.failPendingCommandsForSocket(ws, false); // Preserve session ownership — the reconnecting CLI still owns these sessions - ws.close(1000, 'replaced by reconnect'); + ws.close(1000, "replaced by reconnect"); return true; } } @@ -1306,10 +1775,10 @@ export class UserConnectionDO extends DurableObject { } private replaceWebSocket(connectionId: string): void { - for (const ws of this.ctx.getWebSockets('web')) { + for (const ws of this.ctx.getWebSockets("web")) { const attachment = ws.deserializeAttachment() as WSAttachment | null; if ( - attachment?.role !== 'web' || + attachment?.role !== "web" || attachment.connectionId !== connectionId || attachment.replaced ) { @@ -1318,14 +1787,14 @@ export class UserConnectionDO extends DurableObject { ws.serializeAttachment({ ...attachment, replaced: true }); this.handleWebDisconnect(ws); - ws.close(1000, 'replaced by reconnect'); + ws.close(1000, "replaced by reconnect"); } } private activeWebSockets(): WebSocket[] { - return this.ctx.getWebSockets('web').filter(ws => { + return this.ctx.getWebSockets("web").filter((ws) => { const attachment = ws.deserializeAttachment() as WSAttachment | null; - return attachment?.role === 'web' && !attachment.replaced; + return attachment?.role === "web" && !attachment.replaced; }); } @@ -1336,56 +1805,215 @@ export class UserConnectionDO extends DurableObject { } private findCliByConnectionId(connectionId: string): WebSocket | undefined { - for (const ws of this.ctx.getWebSockets('cli')) { + for (const ws of this.ctx.getWebSockets("cli")) { + const attachment = ws.deserializeAttachment() as WSAttachment | null; + if ( + attachment?.role === "cli" && + attachment.connectionId === connectionId + ) { + return ws; + } + } + return undefined; + } + + private findWebByConnectionId(connectionId: string): WebSocket | undefined { + for (const ws of this.ctx.getWebSockets("web")) { const attachment = ws.deserializeAttachment() as WSAttachment | null; - if (attachment?.role === 'cli' && attachment.connectionId === connectionId) { + if ( + attachment?.role === "web" && + attachment.connectionId === connectionId && + !attachment.replaced + ) { return ws; } } return undefined; } - private failPendingCommandsForSocket(targetWs: WebSocket, cliGone: boolean): void { + private failPendingCommandsForSocket( + targetWs: WebSocket, + cliGone: boolean, + ): void { for (const [id, entry] of this.pendingCommands) { if (entry.targetCliWs !== targetWs) continue; // 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 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; this.sendToWeb( entry.ws, exited - ? { type: 'response', id: entry.originalId, result: {} } + ? { type: "response", id: entry.originalId, result: {} } : { - type: 'response', + type: "response", id: entry.originalId, - error: entry.expectedOwnerConnectionId - ? SESSION_OWNER_CHANGED_ERROR - : 'CLI disconnected', - } + error: liveError, + }, ); this.pendingCommands.delete(id); + + // Step 26: pair the in-memory delete with a durable done transition. + const webAtt = entry.ws.deserializeAttachment() as WSAttachment | null; + this.ctx.waitUntil( + 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: + webAtt?.role === "web" ? webAtt.connectionId : "unknown", + state: "done", + ...(exited ? { result: {} } : { error: durableError }), + } satisfies PendingCommandEntry), + ); + } + + const attachment = targetWs.deserializeAttachment() as WSAttachment | null; + if (cliGone && attachment?.role === "cli") { + this.ctx.waitUntil( + this.finishDurablePendingCommands( + entry => entry.targetConnectionId === attachment.connectionId, + "CLI disconnected", + ), + ); } } private failPendingCommandsForOwnerChange( sessionId: string, - nextOwnerConnectionId: string | undefined + nextOwnerConnectionId: string | undefined, ): void { for (const [id, entry] of this.pendingCommands) { - if (entry.sessionId !== sessionId || entry.targetConnectionId === nextOwnerConnectionId) { + if ( + entry.sessionId !== sessionId || + entry.targetConnectionId === nextOwnerConnectionId + ) { continue; } this.pendingCommands.delete(id); // `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; + const exited = + entry.command === "exit_cli" && nextOwnerConnectionId === undefined; 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, result: {} } + : { + type: "response", + id: entry.originalId, + error: SESSION_OWNER_CHANGED_ERROR, + }, + ); + + // Step 26: pair the in-memory delete with a durable done transition. + const webAtt = entry.ws.deserializeAttachment() as WSAttachment | null; + this.ctx.waitUntil( + 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: + webAtt?.role === "web" ? webAtt.connectionId : "unknown", + state: "done", + ...(exited ? { result: {} } : { error: SESSION_OWNER_CHANGED_ERROR }), + } satisfies PendingCommandEntry), ); } + + this.ctx.waitUntil( + this.finishDurablePendingCommands( + entry => + entry.sessionId === sessionId && + entry.targetConnectionId !== nextOwnerConnectionId, + SESSION_OWNER_CHANGED_ERROR, + ), + ); + } + + // --------------------------------------------------------------------------- + // 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, + ): Promise { + const entries = await this.ctx.storage.list({ + prefix: PENDING_COMMAND_KEY_PREFIX, + }); + await Promise.all( + [...entries].flatMap(([key, entry]) => + entry.state === "pending" && matches(entry) + ? [this.ctx.storage.put(key, { ...entry, state: "done", error })] + : [], + ), + ); + } + + 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 { @@ -1393,11 +2021,42 @@ export class UserConnectionDO extends DurableObject { if (entry.expiresAt > now) continue; this.pendingCommands.delete(id); this.sendToWeb(entry.ws, { - type: 'response', + 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") { + // Write terminal outcome durably (step 26a). + await this.ctx.storage.put(key, { + ...durable, + state: "done", + error: COMMAND_EXPIRED_ERROR, + }); + } else { + // Already done and past TTL: clean up. + await this.ctx.storage.delete(key); + } + } + })().catch((error: unknown) => { + console.error("Failed to sweep durable pending commands", { + error: error instanceof Error ? error.message : String(error), + }); + }), + ); } private scheduleNextAlarm(now: number): void { @@ -1411,7 +2070,7 @@ export class UserConnectionDO extends DurableObject { prefix: READY_PUSH_KEY_PREFIX, }); for (const [key, entry] of pending) { - if (!entry || typeof entry.fireAt !== 'number') continue; + if (!entry || typeof entry.fireAt !== "number") continue; const sessionId = key.slice(READY_PUSH_KEY_PREFIX.length); if (sessionId) this.readyPushFireAt.set(sessionId, entry.fireAt); } @@ -1419,11 +2078,11 @@ export class UserConnectionDO extends DurableObject { this.scheduleNextAlarm(Date.now()); } catch (error: unknown) { // Leave readyPushRebuilt false so the next schedule retries. - console.error('Failed to rebuild readyPush mirror', { + console.error("Failed to rebuild readyPush mirror", { error: error instanceof Error ? error.message : String(error), }); } - })() + })(), ); } @@ -1431,13 +2090,19 @@ export class UserConnectionDO extends DurableObject { for (const lastSeen of this.lastHeartbeatAt.values()) { const staleAt = lastSeen + UserConnectionDO.HEARTBEAT_TIMEOUT_MS; - if (staleAt > now && (nextAlarmAt === undefined || staleAt < nextAlarmAt)) { + if ( + staleAt > now && + (nextAlarmAt === undefined || staleAt < nextAlarmAt) + ) { nextAlarmAt = staleAt; } } for (const entry of this.pendingCommands.values()) { - if (entry.expiresAt > now && (nextAlarmAt === undefined || entry.expiresAt < nextAlarmAt)) { + if ( + entry.expiresAt > now && + (nextAlarmAt === undefined || entry.expiresAt < nextAlarmAt) + ) { nextAlarmAt = entry.expiresAt; } } @@ -1471,9 +2136,9 @@ export class UserConnectionDO extends DurableObject { // Build set of connectionIds that still have a live CLI WebSocket. // This guards against stale entries that persist if a close event is delayed. const liveConnectionIds = new Set(); - for (const ws of this.ctx.getWebSockets('cli')) { + for (const ws of this.ctx.getWebSockets("cli")) { const att = ws.deserializeAttachment() as WSAttachment | null; - if (att?.role === 'cli') liveConnectionIds.add(att.connectionId); + if (att?.role === "cli") liveConnectionIds.add(att.connectionId); } const result: Array< diff --git a/services/session-ingest/src/types/user-connection-protocol.ts b/services/session-ingest/src/types/user-connection-protocol.ts index a7156b2e132..8b612788bff 100644 --- a/services/session-ingest/src/types/user-connection-protocol.ts +++ b/services/session-ingest/src/types/user-connection-protocol.ts @@ -79,6 +79,11 @@ 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. + mutationId: z.string().optional(), }), z.object({ type: z.literal('system'), @@ -108,6 +113,11 @@ 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. + mutationId: z.string().optional(), }), z.object({ type: z.literal('ping'), From 329009169339620e5d010996579174b407220e10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Mon, 3 Aug 2026 23:13:11 +0200 Subject: [PATCH 04/24] fix(events): fence stale connection attempts --- .../src/__tests__/client.test.ts | 59 +++++++++++++++++++ packages/event-service/src/client.ts | 12 +++- 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/packages/event-service/src/__tests__/client.test.ts b/packages/event-service/src/__tests__/client.test.ts index f5383219ca1..ce725aeb87b 100644 --- a/packages/event-service/src/__tests__/client.test.ts +++ b/packages/event-service/src/__tests__/client.test.ts @@ -572,6 +572,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(); diff --git a/packages/event-service/src/client.ts b/packages/event-service/src/client.ts index 14230344640..7dc1b973145 100644 --- a/packages/event-service/src/client.ts +++ b/packages/event-service/src/client.ts @@ -75,6 +75,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 @@ -99,6 +100,7 @@ export class EventServiceClient { } async connect(): Promise { + this.generation++; this.destroyed = false; this.reconnectAttempts = 0; this.authRecoveryAttempts = 0; @@ -146,6 +148,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 +158,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,6 +201,7 @@ export class EventServiceClient { }, HANDSHAKE_TIMEOUT_MS); ws.addEventListener('open', () => { + if (this.ws !== ws) return; const isReconnect = this.hasConnectedBefore; this.connected = true; this.hasConnectedBefore = true; @@ -216,6 +221,7 @@ export class EventServiceClient { }); ws.addEventListener('message', (event: MessageEvent) => { + if (this.ws !== ws) return; this.handleMessage(event.data as string); }); @@ -250,6 +256,7 @@ export class EventServiceClient { } disconnect(): void { + this.generation++; this.destroyed = true; if (this.reconnectTimer !== null) { clearTimeout(this.reconnectTimer); @@ -422,6 +429,7 @@ export class EventServiceClient { this.reconnectAttempts++; this.reconnectTimer = setTimeout(() => { this.reconnectTimer = null; + this.generation++; this.connectOnce().catch(err => { void this.handleReconnectFailure(err); }); From 70ed9b6576d07743bf8cd2a5dbc1b7eba44e550d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Mon, 3 Aug 2026 23:28:27 +0200 Subject: [PATCH 05/24] chore(format): restore oxfmt style Groups 2 and 3 were formatted with prettier, which this repo does not declare. Reformat the affected files with oxfmt. No code changes. --- .../src/cli-live-transport.test.ts | 6 +- .../cloud-agent-sdk/src/cli-live-transport.ts | 447 +- .../src/cloud-agent-transport.test.ts | 752 ++- packages/cloud-agent-sdk/src/session.test.ts | 4 +- .../src/dos/UserConnectionDO.test.ts | 4236 ++++++++--------- .../src/dos/UserConnectionDO.ts | 750 ++- 6 files changed, 2844 insertions(+), 3351 deletions(-) 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 81438b9f5d5..7df839b072f 100644 --- a/packages/cloud-agent-sdk/src/cli-live-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cli-live-transport.test.ts @@ -2348,7 +2348,7 @@ describe('CliLiveTransport createSession', () => { 'create_session', { protocolVersion: 1 }, 'owner', - expect.any(String) as string, + expect.any(String) as string ); expect(userWebConnection.sendCommandToConnection).not.toHaveBeenCalled(); transport.destroy(); @@ -2476,7 +2476,7 @@ describe('CliLiveTransport createSession', () => { 'create_session', { protocolVersion: 1 }, 'owner-a', - expect.any(String) as string, + expect.any(String) as string ); expect(userWebConnection.sendCommandToConnection).not.toHaveBeenCalled(); transport.destroy(); @@ -2536,7 +2536,7 @@ describe('CliLiveTransport createSession', () => { orgId: 'aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee', }, 'owner', - expect.any(String) as string, + expect.any(String) as string ); transport.destroy(); }); diff --git a/packages/cloud-agent-sdk/src/cli-live-transport.ts b/packages/cloud-agent-sdk/src/cli-live-transport.ts index ea8dba1f947..26d930413b7 100644 --- a/packages/cloud-agent-sdk/src/cli-live-transport.ts +++ b/packages/cloud-agent-sdk/src/cli-live-transport.ts @@ -2,42 +2,39 @@ * CLI live transport - consumes a shared user web connection and translates * one remote CLI session into normalized transport events and commands. */ -import { normalizeCliEvent, isChatEvent } from "./normalizer"; -import { - parseRemoteCommandCatalog, - type RemoteCommandState, -} from "./remote-command-catalog"; -import { parseCreateSessionResponse } from "./create-session"; -import { parseExitSessionResponse } from "./exit-session"; +import { normalizeCliEvent, isChatEvent } from './normalizer'; +import { parseRemoteCommandCatalog, type RemoteCommandState } from './remote-command-catalog'; +import { parseCreateSessionResponse } from './create-session'; +import { parseExitSessionResponse } from './exit-session'; import { cliConnectionDataSchema, heartbeatDataSchema, remoteModelCatalogV1Schema, sessionsListDataSchema, -} from "./schemas"; -import type { SlashCommandInfo } from "./schemas"; -import type { RemoteModelState } from "./remote-model-catalog"; +} from './schemas'; +import type { SlashCommandInfo } from './schemas'; +import type { RemoteModelState } from './remote-model-catalog'; import type { CreateRemoteSessionInput, TransportFactory, TransportSendInput, TransportSink, -} from "./transport"; +} from './transport'; import type { KiloSessionId, SessionSnapshot, SessionSnapshotPage, SessionSnapshotPageOutcome, -} from "./types"; +} from './types'; import { CommandDeliveredError, UserWebCommandError, type UserWebCliEvent, type UserWebConnection, -} from "./user-web-connection"; +} from './user-web-connection'; /** Delivered error string for old-CLI strict-parse rejection of extended create_session. */ -const INVALID_CREATE_SESSION_COMMAND = "invalid create_session command"; +const INVALID_CREATE_SESSION_COMMAND = 'invalid create_session command'; /** * Build `create_session` wire data: always `protocolVersion: 1`, plus any @@ -60,8 +57,7 @@ function buildCreateSessionWireData(input?: CreateRemoteSessionInput): { type CliLiveTransportConfig = { kiloSessionId: KiloSessionId; userWebConnection: UserWebConnection; - fetchSnapshot?: - ((kiloSessionId: KiloSessionId) => Promise) | undefined; + fetchSnapshot?: ((kiloSessionId: KiloSessionId) => Promise) | undefined; /** * Page-aware root snapshot fetch. When provided, every `replayCurrentSnapshot` * (initial bounded read AND reconnect immediate + delayed resync) uses it @@ -71,7 +67,7 @@ type CliLiveTransportConfig = { */ fetchSnapshotPage?: ( kiloSessionId: KiloSessionId, - options: { cursor?: string }, + options: { cursor?: string } ) => Promise; /** * Called after a successful initial bounded page read. Reconnect and @@ -81,8 +77,7 @@ type CliLiveTransportConfig = { onInitialPageLoaded?: ((page: SessionSnapshotPage) => void) | undefined; onError?: ((message: string) => void) | undefined; onRemoteModelStateChange?: ((state: RemoteModelState) => void) | undefined; - onRemoteCommandStateChange?: - ((state: RemoteCommandState) => void) | undefined; + onRemoteCommandStateChange?: ((state: RemoteCommandState) => void) | undefined; onCapabilityChange?: (() => void) | undefined; /** * Fired whenever the per-session capabilities advertised by the owning @@ -93,9 +88,7 @@ type CliLiveTransportConfig = { * to recompute the `supportsAttachments` gate. */ onCapabilitiesChange?: - | (( - capabilities: { attachments?: boolean | undefined } | undefined, - ) => void) + | ((capabilities: { attachments?: boolean | undefined } | undefined) => void) | undefined; }; @@ -104,7 +97,7 @@ type CliLiveTransportConfig = { // still appear after long backgrounding. const RECONNECT_RESYNC_DELAY_MS = 5000; const REMOTE_SESSION_EXIT_UNAVAILABLE = - "Remote session exit is unavailable for the current session"; + 'Remote session exit is unavailable for the current session'; /** * Deep-copy a list of validated remote slash commands. @@ -115,18 +108,14 @@ const REMOTE_SESSION_EXIT_UNAVAILABLE = * schema has already validated them. Avoids `JSON.parse(JSON.stringify(...))` * because the shape is known and bounded. */ -export function deepCopyCommands( - commands: SlashCommandInfo[], -): SlashCommandInfo[] { - return commands.map((command) => ({ +export function deepCopyCommands(commands: SlashCommandInfo[]): SlashCommandInfo[] { + return commands.map(command => ({ ...command, hints: command.hints.slice(), })); } -function createCliLiveTransport( - config: CliLiveTransportConfig, -): TransportFactory { +function createCliLiveTransport(config: CliLiveTransportConfig): TransportFactory { return (sink: TransportSink) => { let generation = 0; let cleanup: (() => void) | null = null; @@ -140,11 +129,8 @@ function createCliLiveTransport( * payload that omitted this session). Compared structurally on every * new observation; only emitted on a real change. */ - let currentCapabilities: { attachments?: boolean | undefined } | undefined = - undefined; - function publishCapabilities( - next: { attachments?: boolean | undefined } | undefined, - ): void { + let currentCapabilities: { attachments?: boolean | undefined } | undefined = undefined; + function publishCapabilities(next: { attachments?: boolean | undefined } | undefined): void { const previousAttachments = currentCapabilities?.attachments; const nextAttachments = next?.attachments; if (previousAttachments === nextAttachments) return; @@ -166,12 +152,12 @@ function createCliLiveTransport( } | null = null; let remoteModelState: RemoteModelState = { ownerConnectionId: null, - protocol: "unknown", - refresh: "idle", + protocol: 'unknown', + refresh: 'idle', }; let remoteCommandState: RemoteCommandState = { ownerConnectionId: null, - refresh: "idle", + refresh: 'idle', commands: [], canExitSession: undefined, }; @@ -194,8 +180,7 @@ function createCliLiveTransport( function canExitSession(): boolean { return ( - remoteCommandState.canExitSession === true && - remoteCommandState.ownerConnectionId !== null + remoteCommandState.canExitSession === true && remoteCommandState.ownerConnectionId !== null ); } @@ -206,7 +191,7 @@ function createCliLiveTransport( const snapshotted = deepCopyCommands(commands); lastValidCommands = snapshotted; sink.onServiceEvent({ - type: "commands.available", + type: 'commands.available', commands: deepCopyCommands(snapshotted), }); } @@ -226,8 +211,8 @@ function createCliLiveTransport( commandCatalogRequestInFlight = null; publishRemoteModelState({ ownerConnectionId: nextOwnerConnectionId, - protocol: "unknown", - refresh: nextOwnerConnectionId ? "loading" : "idle", + protocol: 'unknown', + refresh: nextOwnerConnectionId ? 'loading' : 'idle', }); // Owner replacement clears the command cache; emit an empty catalog // event and the matching idle state so consumers can re-render. @@ -235,14 +220,14 @@ function createCliLiveTransport( publishCommands([]); publishRemoteCommandState({ ownerConnectionId: nextOwnerConnectionId, - refresh: "idle", + refresh: 'idle', commands: snapshotCommands(), canExitSession: undefined, }); } else { publishRemoteCommandState({ ownerConnectionId: nextOwnerConnectionId, - refresh: "idle", + refresh: 'idle', commands: snapshotCommands(), canExitSession: undefined, }); @@ -264,7 +249,7 @@ function createCliLiveTransport( error: unknown, expectedOwnerConnectionId: string, expectedGeneration: number, - expectedRequestGeneration: number, + expectedRequestGeneration: number ): void { if ( expectedGeneration !== generation || @@ -274,19 +259,16 @@ function createCliLiveTransport( return; } - if ( - error instanceof UserWebCommandError && - error.code === "SESSION_OWNER_CHANGED" - ) { + if (error instanceof UserWebCommandError && error.code === 'SESSION_OWNER_CHANGED') { setOwnerConnectionId(null); return; } - if (error instanceof Error && error.message.includes("unknown command")) { + if (error instanceof Error && error.message.includes('unknown command')) { publishRemoteModelState({ ownerConnectionId: expectedOwnerConnectionId, - protocol: "legacy", - refresh: "idle", + protocol: 'legacy', + refresh: 'idle', }); return; } @@ -294,22 +276,14 @@ function createCliLiveTransport( publishRemoteModelState({ ownerConnectionId: expectedOwnerConnectionId, protocol: remoteModelState.protocol, - ...(remoteModelState.catalog - ? { catalog: remoteModelState.catalog } - : {}), - refresh: "error", - error: - error instanceof Error - ? error.message - : "Failed to discover remote models", + ...(remoteModelState.catalog ? { catalog: remoteModelState.catalog } : {}), + refresh: 'error', + error: error instanceof Error ? error.message : 'Failed to discover remote models', }); } function discoverModels(expectedOwnerConnectionId: string): void { - if ( - catalogRequestInFlight?.ownerConnectionId === expectedOwnerConnectionId - ) - return; + if (catalogRequestInFlight?.ownerConnectionId === expectedOwnerConnectionId) return; catalogRequestGeneration += 1; const expectedRequestGeneration = catalogRequestGeneration; @@ -321,21 +295,19 @@ function createCliLiveTransport( publishRemoteModelState({ ownerConnectionId: expectedOwnerConnectionId, protocol: remoteModelState.protocol, - ...(remoteModelState.catalog - ? { catalog: remoteModelState.catalog } - : {}), - refresh: "loading", + ...(remoteModelState.catalog ? { catalog: remoteModelState.catalog } : {}), + refresh: 'loading', }); void config.userWebConnection .sendCommand( config.kiloSessionId, - "list_models", + 'list_models', { protocolVersion: 1 }, - expectedOwnerConnectionId, + expectedOwnerConnectionId ) .then( - (result) => { + result => { if ( expectedGeneration !== generation || expectedRequestGeneration !== catalogRequestGeneration || @@ -347,33 +319,31 @@ function createCliLiveTransport( const parsed = remoteModelCatalogV1Schema.safeParse(result); if (!parsed.success) { handleCatalogFailure( - new Error("Invalid remote model catalog"), + new Error('Invalid remote model catalog'), expectedOwnerConnectionId, expectedGeneration, - expectedRequestGeneration, + expectedRequestGeneration ); return; } publishRemoteModelState({ ownerConnectionId: expectedOwnerConnectionId, - protocol: "v1", + protocol: 'v1', catalog: parsed.data, - refresh: "idle", + refresh: 'idle', }); }, - (error) => + error => handleCatalogFailure( error, expectedOwnerConnectionId, expectedGeneration, - expectedRequestGeneration, - ), + expectedRequestGeneration + ) ) .finally(() => { - if ( - catalogRequestInFlight?.generation === expectedRequestGeneration - ) { + if (catalogRequestInFlight?.generation === expectedRequestGeneration) { catalogRequestInFlight = null; } }); @@ -384,7 +354,7 @@ function createCliLiveTransport( expectedOwnerConnectionId: string, expectedGeneration: number, expectedRequestGeneration: number, - clearCatalog: boolean, + clearCatalog: boolean ): void { if ( expectedGeneration !== generation || @@ -394,10 +364,7 @@ function createCliLiveTransport( return; } - if ( - error instanceof UserWebCommandError && - error.code === "SESSION_OWNER_CHANGED" - ) { + if (error instanceof UserWebCommandError && error.code === 'SESSION_OWNER_CHANGED') { setOwnerConnectionId(null); return; } @@ -406,14 +373,11 @@ function createCliLiveTransport( // and surface actionable copy through the command state so the chat // composer can prompt the user to upgrade. Never populate the global // session error atom for this. - if ( - error instanceof UserWebCommandError && - error.code === "CLI_UPGRADE_REQUIRED" - ) { + if (error instanceof UserWebCommandError && error.code === 'CLI_UPGRADE_REQUIRED') { publishCommands([]); publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: "upgrade-required", + refresh: 'upgrade-required', commands: snapshotCommands(), canExitSession: undefined, message: error.message, @@ -423,26 +387,23 @@ function createCliLiveTransport( // Legacy CLI without `list_commands` support: treat as "no commands // available" rather than a failure so the composer can hide suggestions. - if (error instanceof Error && error.message.includes("unknown command")) { + if (error instanceof Error && error.message.includes('unknown command')) { publishCommands([]); publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: "idle", + refresh: 'idle', commands: snapshotCommands(), canExitSession: undefined, }); return; } - const message = - error instanceof Error - ? error.message - : "Failed to discover remote commands"; + const message = error instanceof Error ? error.message : 'Failed to discover remote commands'; if (clearCatalog) { publishCommands([]); publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: "error", + refresh: 'error', commands: snapshotCommands(), canExitSession: undefined, message, @@ -454,7 +415,7 @@ function createCliLiveTransport( // populating the fatal session error atom. publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: "error", + refresh: 'error', commands: snapshotCommands(), canExitSession: undefined, message, @@ -462,11 +423,7 @@ function createCliLiveTransport( } function discoverCommands(expectedOwnerConnectionId: string): void { - if ( - commandCatalogRequestInFlight?.ownerConnectionId === - expectedOwnerConnectionId - ) - return; + if (commandCatalogRequestInFlight?.ownerConnectionId === expectedOwnerConnectionId) return; commandCatalogRequestGeneration += 1; const expectedRequestGeneration = commandCatalogRequestGeneration; @@ -477,7 +434,7 @@ function createCliLiveTransport( }; publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: "loading", + refresh: 'loading', commands: snapshotCommands(), canExitSession: undefined, }); @@ -485,12 +442,12 @@ function createCliLiveTransport( void config.userWebConnection .sendCommand( config.kiloSessionId, - "list_commands", + 'list_commands', { protocolVersion: 1 }, - expectedOwnerConnectionId, + expectedOwnerConnectionId ) .then( - (result) => { + result => { if ( expectedGeneration !== generation || expectedRequestGeneration !== commandCatalogRequestGeneration || @@ -502,11 +459,11 @@ function createCliLiveTransport( const parsed = parseRemoteCommandCatalog(result); if (!parsed.ok) { handleCommandCatalogFailure( - new Error("Invalid remote command catalog"), + new Error('Invalid remote command catalog'), expectedOwnerConnectionId, expectedGeneration, expectedRequestGeneration, - true, + true ); return; } @@ -514,12 +471,12 @@ function createCliLiveTransport( publishCommands(parsed.commands); publishRemoteCommandState({ ownerConnectionId: expectedOwnerConnectionId, - refresh: "idle", + refresh: 'idle', commands: snapshotCommands(), canExitSession: parsed.canExitSession, }); }, - (error) => + error => handleCommandCatalogFailure( error, expectedOwnerConnectionId, @@ -528,28 +485,24 @@ function createCliLiveTransport( // A relay `CATALOG_TOO_LARGE` reports the response is over // the 512 KiB size cap; clear the prior catalog so the // composer can present a clean state. - error instanceof UserWebCommandError && - error.code === "CATALOG_TOO_LARGE", - ), + error instanceof UserWebCommandError && error.code === 'CATALOG_TOO_LARGE' + ) ) .finally(() => { - if ( - commandCatalogRequestInFlight?.generation === - expectedRequestGeneration - ) { + if (commandCatalogRequestInFlight?.generation === expectedRequestGeneration) { commandCatalogRequestInFlight = null; } }); } function replayPage(page: SessionSnapshotPage): void { - sink.onServiceEvent({ type: "session.created", info: page.info }); + sink.onServiceEvent({ type: 'session.created', info: page.info }); for (const msg of page.messages) { - sink.onChatEvent({ type: "message.updated", info: msg.info }); + sink.onChatEvent({ type: 'message.updated', info: msg.info }); for (const part of msg.parts) { - sink.onChatEvent({ type: "message.part.updated", part }); + sink.onChatEvent({ type: 'message.part.updated', part }); } } } @@ -558,13 +511,9 @@ function createCliLiveTransport( sessionId: string, parentSessionId: string | undefined, event: string, - data: unknown, + data: unknown ): void { - if ( - sessionId !== config.kiloSessionId && - parentSessionId !== config.kiloSessionId - ) - return; + if (sessionId !== config.kiloSessionId && parentSessionId !== config.kiloSessionId) return; const normalized = normalizeCliEvent(event, data); if (!normalized) return; @@ -578,7 +527,7 @@ function createCliLiveTransport( function stopForDisconnectedSession(): void { if (sessionStopped) return; - sink.onServiceEvent({ type: "stopped", reason: "disconnected" }); + sink.onServiceEvent({ type: 'stopped', reason: 'disconnected' }); sessionStopped = true; // The disconnected state is only cleared by a session.status event, so // the next post-reconnect heartbeat must always forward one — even when @@ -591,18 +540,18 @@ function createCliLiveTransport( // fired while the socket was dead is never replayed, which otherwise // leaves the UI stuck on a busy indicator forever. function forwardHeartbeatStatus(status: string): void { - if (status !== "idle" && status !== "busy") return; + if (status !== 'idle' && status !== 'busy') return; if (status === lastForwardedHeartbeatStatus) return; lastForwardedHeartbeatStatus = status; sink.onServiceEvent({ - type: "session.status", + type: 'session.status', sessionId: config.kiloSessionId, status: { type: status }, }); } function handleSystemMessage(event: string, data: unknown): void { - if (event === "cli.disconnected") { + if (event === 'cli.disconnected') { const parsed = cliConnectionDataSchema.safeParse(data); if (parsed.success && ownerConnectionId === parsed.data.connectionId) { setOwnerConnectionId(null); @@ -611,13 +560,11 @@ function createCliLiveTransport( return; } - if (event === "sessions.list") { + if (event === 'sessions.list') { const parsed = sessionsListDataSchema.safeParse(data); if (!parsed.success) return; - const session = parsed.data.sessions.find( - (item) => item.id === config.kiloSessionId, - ); + const session = parsed.data.sessions.find(item => item.id === config.kiloSessionId); if (session) { setOwnerConnectionId(session.connectionId); sessionStopped = false; @@ -631,13 +578,11 @@ function createCliLiveTransport( return; } - if (event === "sessions.heartbeat") { + if (event === 'sessions.heartbeat') { const parsed = heartbeatDataSchema.safeParse(data); if (!parsed.success) return; - const session = parsed.data.sessions.find( - (item) => item.id === config.kiloSessionId, - ); + const session = parsed.data.sessions.find(item => item.id === config.kiloSessionId); if (session) { setOwnerConnectionId(parsed.data.connectionId); sessionStopped = false; @@ -656,11 +601,10 @@ function createCliLiveTransport( async function sendCommand( command: string, data: unknown, - mutationId?: string, + mutationId?: string ): Promise { const expectedOwnerConnectionId = ownerConnectionId; - if (!expectedOwnerConnectionId) - throw new Error("Remote session has no connected owner"); + if (!expectedOwnerConnectionId) throw new Error('Remote session has no connected owner'); try { return await config.userWebConnection.sendCommand( @@ -668,13 +612,10 @@ function createCliLiveTransport( command, data, expectedOwnerConnectionId, - ...(mutationId !== undefined ? [mutationId] : []), + ...(mutationId !== undefined ? [mutationId] : []) ); } catch (error) { - if ( - error instanceof UserWebCommandError && - error.code === "SESSION_OWNER_CHANGED" - ) { + if (error instanceof UserWebCommandError && error.code === 'SESSION_OWNER_CHANGED') { setOwnerConnectionId(null); } throw error; @@ -682,61 +623,52 @@ function createCliLiveTransport( } function getRemoteModelFields(input: TransportSendInput): - | { kind: "none" } + | { kind: 'none' } | { - kind: "structured"; + kind: 'structured'; model: { providerID: string; modelID: string }; variant?: string; } - | { kind: "legacy"; model: string; variant?: string } { + | { kind: 'legacy'; model: string; variant?: string } { const override = input.remoteModelOverride; - if (!override) return { kind: "none" }; + if (!override) return { kind: 'none' }; - if ( - remoteModelState.protocol === "v1" && - override.source === "cli-catalog" - ) { + if (remoteModelState.protocol === 'v1' && override.source === 'cli-catalog') { const provider = remoteModelState.catalog?.providers.find( - (item) => item.id === override.selection.model.providerID, - ); - const model = provider?.models.find( - (item) => item.id === override.selection.model.modelID, + item => item.id === override.selection.model.providerID ); + const model = provider?.models.find(item => item.id === override.selection.model.modelID); if (!model) { - throw new Error( - "Selected remote model is not available in the current CLI catalog", - ); + throw new Error('Selected remote model is not available in the current CLI catalog'); } const variant = override.selection.variant; if (variant && !model.variants.includes(variant)) { throw new Error( - "Selected remote model variant is not available in the current CLI catalog", + 'Selected remote model variant is not available in the current CLI catalog' ); } return { - kind: "structured", + kind: 'structured', model: override.selection.model, ...(variant ? { variant } : {}), }; } if ( - remoteModelState.protocol === "legacy" && - override.source === "legacy-gateway" && - override.selection.model.providerID === "kilo" + remoteModelState.protocol === 'legacy' && + override.source === 'legacy-gateway' && + override.selection.model.providerID === 'kilo' ) { return { - kind: "legacy", + kind: 'legacy', model: override.selection.model.modelID, - ...(override.selection.variant - ? { variant: override.selection.variant } - : {}), + ...(override.selection.variant ? { variant: override.selection.variant } : {}), }; } throw new Error( - "Selected remote model override is incompatible with the connected CLI model protocol", + 'Selected remote model override is incompatible with the connected CLI model protocol' ); } @@ -760,12 +692,12 @@ function createCliLiveTransport( publishCapabilities(undefined); publishRemoteModelState({ ownerConnectionId: null, - protocol: "unknown", - refresh: "idle", + protocol: 'unknown', + refresh: 'idle', }); publishRemoteCommandState({ ownerConnectionId: null, - refresh: "idle", + refresh: 'idle', commands: snapshotCommands(), canExitSession: undefined, }); @@ -780,12 +712,7 @@ function createCliLiveTransport( const events = bufferedCliEvents; bufferedCliEvents = null; for (const msg of events ?? []) { - handleEventMessage( - msg.sessionId, - msg.parentSessionId, - msg.event, - msg.data, - ); + handleEventMessage(msg.sessionId, msg.parentSessionId, msg.event, msg.data); } sink.onReplayComplete?.(); }; @@ -835,10 +762,7 @@ function createCliLiveTransport( return; } if (reportError) { - const message = - error instanceof Error - ? error.message - : "Failed to fetch snapshot"; + const message = error instanceof Error ? error.message : 'Failed to fetch snapshot'; config.onError?.(message); } bufferedCliEvents = [ @@ -850,30 +774,28 @@ function createCliLiveTransport( }; if (config.fetchSnapshotPage) { - void config - .fetchSnapshotPage(config.kiloSessionId, {}) - .then((page) => { - if (page && page.kind === "success") { - onPage(page); - } else { - onError( - new Error( - page === null - ? "Session not found" - : page.kind === "retryable_failure" - ? "Session history temporarily unavailable" - : page.kind === "too_large" - ? "Session history too large to load" - : "Session history is unavailable", - ), - ); - } - }, onError); + void config.fetchSnapshotPage(config.kiloSessionId, {}).then(page => { + if (page && page.kind === 'success') { + onPage(page); + } else { + onError( + new Error( + page === null + ? 'Session not found' + : page.kind === 'retryable_failure' + ? 'Session history temporarily unavailable' + : page.kind === 'too_large' + ? 'Session history too large to load' + : 'Session history is unavailable' + ) + ); + } + }, onError); return; } if (config.fetchSnapshot) { - void config.fetchSnapshot(config.kiloSessionId).then((snapshot) => { + void config.fetchSnapshot(config.kiloSessionId).then(snapshot => { onPage({ info: snapshot.info, messages: snapshot.messages, @@ -885,28 +807,20 @@ function createCliLiveTransport( }; replayCurrentSnapshot(true); - const offCli = config.userWebConnection.onCliEvent( - config.kiloSessionId, - (msg) => { - const normalized = normalizeCliEvent(msg.event, msg.data); - const shouldBufferForSnapshot = - normalized && - (isChatEvent(normalized) || - normalized.type === "session.created" || - normalized.type === "session.updated"); - if (shouldBufferForSnapshot && bufferedCliEvents !== null) { - bufferedCliEvents.push(msg); - return; - } - handleEventMessage( - msg.sessionId, - msg.parentSessionId, - msg.event, - msg.data, - ); - }, - ); - const offSystem = config.userWebConnection.onSystemEvent((msg) => { + const offCli = config.userWebConnection.onCliEvent(config.kiloSessionId, msg => { + const normalized = normalizeCliEvent(msg.event, msg.data); + const shouldBufferForSnapshot = + normalized && + (isChatEvent(normalized) || + normalized.type === 'session.created' || + normalized.type === 'session.updated'); + if (shouldBufferForSnapshot && bufferedCliEvents !== null) { + bufferedCliEvents.push(msg); + return; + } + handleEventMessage(msg.sessionId, msg.parentSessionId, msg.event, msg.data); + }); + const offSystem = config.userWebConnection.onSystemEvent(msg => { handleSystemMessage(msg.event, msg.data); }); const offReconnect = config.userWebConnection.onReconnect(() => { @@ -929,8 +843,9 @@ function createCliLiveTransport( discoverCommands(ownerConnectionId); } }); - const releaseSubscription = - config.userWebConnection.subscribeToCliSession(config.kiloSessionId); + const releaseSubscription = config.userWebConnection.subscribeToCliSession( + config.kiloSessionId + ); let released = false; cleanup = () => { if (released) return; @@ -968,7 +883,7 @@ function createCliLiveTransport( wireData.orgId !== undefined; let result: unknown; try { - result = await sendCommand("create_session", wireData, mutationId); + 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. @@ -978,9 +893,9 @@ function createCliLiveTransport( error.message === INVALID_CREATE_SESSION_COMMAND ) { result = await sendCommand( - "create_session", + 'create_session', { protocolVersion: 1 }, - crypto.randomUUID(), + crypto.randomUUID() ); } else { throw error; @@ -988,15 +903,15 @@ function createCliLiveTransport( } const parsed = parseCreateSessionResponse(result); if (!parsed.ok) { - throw new Error("Invalid create_session response"); + throw new Error('Invalid create_session response'); } return parsed.kiloSessionId; }, exitSession: async () => { - if (remoteCommandState.refresh === "upgrade-required") { + if (remoteCommandState.refresh === 'upgrade-required') { throw new Error( remoteCommandState.message ?? - "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", + 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.' ); } if (!canExitSession()) { @@ -1005,20 +920,20 @@ function createCliLiveTransport( // `exit_cli` is the compatibility wire command name for // session-detach — the product-facing name is "session exit". - const result = await sendCommand("exit_cli", { protocolVersion: 1 }); + const result = await sendCommand('exit_cli', { protocolVersion: 1 }); if (!parseExitSessionResponse(result).ok) { - throw new Error("Invalid exit_cli response"); + throw new Error('Invalid exit_cli response'); } }, send: async (input: TransportSendInput) => { - if (input.payload.type === "command") { + if (input.payload.type === 'command') { const remoteModel = getRemoteModelFields(input); - return sendCommand("send_command", { + return sendCommand('send_command', { protocolVersion: 1, command: input.payload.command, arguments: input.payload.arguments, ...(input.messageId ? { messageID: input.messageId } : {}), - ...(remoteModel.kind === "none" + ...(remoteModel.kind === 'none' ? {} : { // `send_command` requires a structured model: never emit a @@ -1028,75 +943,71 @@ function createCliLiveTransport( // `dispatchedKilocodeModelId` would emit for the // equivalent prompt wire. model: - remoteModel.kind === "structured" + remoteModel.kind === 'structured' ? remoteModel.model : { - providerID: "kilo", - modelID: remoteModel.model.replace(/^kilo\//, ""), + providerID: 'kilo', + modelID: remoteModel.model.replace(/^kilo\//, ''), }, - ...(remoteModel.variant - ? { variant: remoteModel.variant } - : {}), + ...(remoteModel.variant ? { variant: remoteModel.variant } : {}), }), }); } const payload = input.payload; const remoteModel = getRemoteModelFields(input); const parts: Array< - | { type: "text"; text: string } + | { type: 'text'; text: string } | { - type: "file"; + type: 'file'; mime: string; filename: string; url: string; } - > = [{ type: "text", text: payload.prompt }]; + > = [{ type: 'text', text: payload.prompt }]; if (input.attachmentParts && input.attachmentParts.length > 0) { for (const part of input.attachmentParts) { parts.push({ - type: "file", + type: 'file', mime: part.mime, filename: part.filename, url: part.url, }); } } - return sendCommand("send_message", { + return sendCommand('send_message', { sessionID: config.kiloSessionId, parts, ...(payload.mode ? { agent: payload.mode } : {}), - ...(remoteModel.kind === "none" + ...(remoteModel.kind === 'none' ? {} : { model: remoteModel.model, - ...(remoteModel.variant - ? { variant: remoteModel.variant } - : {}), + ...(remoteModel.variant ? { variant: remoteModel.variant } : {}), }), }); }, - interrupt: () => sendCommand("interrupt", {}), - answer: (payload) => - sendCommand("question_reply", { + interrupt: () => sendCommand('interrupt', {}), + answer: payload => + sendCommand('question_reply', { requestID: payload.requestId, answers: payload.answers, }), - reject: (payload) => - sendCommand("question_reject", { + reject: payload => + sendCommand('question_reject', { requestID: payload.requestId, }), - respondToPermission: (payload) => - sendCommand("permission_respond", { + respondToPermission: payload => + sendCommand('permission_respond', { requestID: payload.requestId, reply: payload.response, }), - acceptSuggestion: (payload) => - sendCommand("suggestion_accept", { + acceptSuggestion: payload => + sendCommand('suggestion_accept', { requestID: payload.requestId, index: payload.index, }), - dismissSuggestion: (payload) => - sendCommand("suggestion_dismiss", { + dismissSuggestion: payload => + sendCommand('suggestion_dismiss', { requestID: payload.requestId, }), 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 644eba045c8..77203ca4829 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts @@ -2,12 +2,12 @@ * Tests for CloudAgentTransport — verifies event normalization, routing to * chat/service sinks, and lifecycle generation tracking. */ -import type { CloudAgentEvent } from "./event-types"; -import { createEventHelpers } from "./__fixtures__/helpers"; -import type { ChatEvent, ServiceEvent } from "./normalizer"; -import { createCloudAgentTransport } from "./cloud-agent-transport"; -import type { SessionSnapshotPageOutcome } from "./types"; -import { kiloId, cloudAgentId, makeSnapshot } from "./test-helpers"; +import type { CloudAgentEvent } from './event-types'; +import { createEventHelpers } from './__fixtures__/helpers'; +import type { ChatEvent, ServiceEvent } from './normalizer'; +import { createCloudAgentTransport } from './cloud-agent-transport'; +import type { SessionSnapshotPageOutcome } from './types'; +import { kiloId, cloudAgentId, makeSnapshot } from './test-helpers'; // --------------------------------------------------------------------------- // WebSocket mock @@ -52,10 +52,10 @@ afterEach(() => { /** Flush microtask queue so Promise.all + .then in connect() settles. */ async function flushPromises(): Promise { - await new Promise((r) => setTimeout(r, 0)); + await new Promise(r => setTimeout(r, 0)); } -const emptySnapshot = makeSnapshot({ id: "ses-1" }); +const emptySnapshot = makeSnapshot({ id: 'ses-1' }); function sendRaw(event: CloudAgentEvent): void { mockWs.onmessage?.({ data: JSON.stringify(event) } as MessageEvent); @@ -63,36 +63,35 @@ function sendRaw(event: CloudAgentEvent): void { function createMockApi() { return { - send: jest.fn(() => Promise.resolve("sent")), - interrupt: jest.fn(() => Promise.resolve("interrupted")), - answer: jest.fn(() => Promise.resolve("answered")), - reject: jest.fn(() => Promise.resolve("rejected")), - respondToPermission: jest.fn(() => Promise.resolve("responded")), + send: jest.fn(() => Promise.resolve('sent')), + interrupt: jest.fn(() => Promise.resolve('interrupted')), + answer: jest.fn(() => Promise.resolve('answered')), + reject: jest.fn(() => Promise.resolve('rejected')), + respondToPermission: jest.fn(() => Promise.resolve('responded')), }; } function createTransportWithSinks( - getTicket: (sessionId: string) => string | Promise = () => - "test-ticket", + getTicket: (sessionId: string) => string | Promise = () => 'test-ticket', onError?: (message: string) => void, - api = createMockApi(), + api = createMockApi() ) { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; const factory = createCloudAgentTransport({ - sessionId: cloudAgentId("ses-1"), - kiloSessionId: kiloId("ses-1"), + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), api, getTicket, fetchSnapshot: () => Promise.resolve(emptySnapshot), - websocketBaseUrl: "ws://localhost:9999", + websocketBaseUrl: 'ws://localhost:9999', onError, }); const transport = factory({ - onChatEvent: (event) => chatEvents.push(event), - onServiceEvent: (event) => serviceEvents.push(event), + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), }); return { transport, chatEvents, serviceEvents, api }; @@ -108,69 +107,61 @@ beforeEach(() => { // Tests // --------------------------------------------------------------------------- -describe("CloudAgentTransport event routing", () => { - it("routes chat events to onChatEvent", async () => { +describe('CloudAgentTransport event routing', () => { + it('routes chat events to onChatEvent', async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); await flushPromises(); sendRaw( - kilocode("message.updated", { + kilocode('message.updated', { info: { - id: "msg-1", - sessionID: "ses-1", - role: "assistant", + id: 'msg-1', + sessionID: 'ses-1', + role: 'assistant', time: { created: 1 }, }, - }), + }) ); expect(chatEvents).toHaveLength(1); - expect(chatEvents[0]).toEqual( - expect.objectContaining({ type: "message.updated" }), - ); + expect(chatEvents[0]).toEqual(expect.objectContaining({ type: 'message.updated' })); // session.created from snapshot replay expect(serviceEvents).toHaveLength(1); - expect(serviceEvents[0]).toEqual( - expect.objectContaining({ type: "session.created" }), - ); + expect(serviceEvents[0]).toEqual(expect.objectContaining({ type: 'session.created' })); transport.destroy(); }); - it("routes service events to onServiceEvent", async () => { + it('routes service events to onServiceEvent', async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); await flushPromises(); sendRaw( - kilocode("session.status", { - sessionID: "ses-1", - status: { type: "busy" }, - }), + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) ); // session.created from snapshot replay + session.status from sendRaw expect(serviceEvents).toHaveLength(2); - expect(serviceEvents[0]).toEqual( - expect.objectContaining({ type: "session.created" }), - ); - expect(serviceEvents[1]).toEqual( - expect.objectContaining({ type: "session.status" }), - ); + expect(serviceEvents[0]).toEqual(expect.objectContaining({ type: 'session.created' })); + expect(serviceEvents[1]).toEqual(expect.objectContaining({ type: 'session.status' })); expect(chatEvents).toHaveLength(0); transport.destroy(); }); - it("routes cached command catalogs emitted without an execution ID", async () => { + it('routes cached command catalogs emitted without an execution ID', async () => { const { transport, serviceEvents } = createTransportWithSinks(); const commands = [ { - name: "deploy-prod", - description: "Deploy production", - hints: ["$ARGUMENTS"], - source: "command", + name: 'deploy-prod', + description: 'Deploy production', + hints: ['$ARGUMENTS'], + source: 'command', }, ]; @@ -179,21 +170,21 @@ describe("CloudAgentTransport event routing", () => { sendRaw({ eventId: 0, executionId: null, - sessionId: "ses-1", - streamEventType: "commands.available", + sessionId: 'ses-1', + streamEventType: 'commands.available', timestamp: new Date().toISOString(), data: { commands }, }); expect(serviceEvents).toContainEqual({ - type: "commands.available", + type: 'commands.available', commands, }); transport.destroy(); }); - it("routes mixed events to correct sinks", async () => { + it('routes mixed events to correct sinks', async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); @@ -201,60 +192,52 @@ describe("CloudAgentTransport event routing", () => { // Chat event sendRaw( - kilocode("message.updated", { + kilocode('message.updated', { info: { - id: "msg-1", - sessionID: "ses-1", - role: "assistant", + id: 'msg-1', + sessionID: 'ses-1', + role: 'assistant', time: { created: 1 }, }, - }), + }) ); // Service event sendRaw( - kilocode("session.status", { - sessionID: "ses-1", - status: { type: "busy" }, - }), + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) ); // Chat event (delta) sendRaw( - kilocode("message.part.delta", { - sessionID: "ses-1", - messageID: "msg-1", - partID: "part-1", - field: "text", - delta: "hello", - }), + kilocode('message.part.delta', { + sessionID: 'ses-1', + messageID: 'msg-1', + partID: 'part-1', + field: 'text', + delta: 'hello', + }) ); expect(chatEvents).toHaveLength(2); - expect(chatEvents[0]).toEqual( - expect.objectContaining({ type: "message.updated" }), - ); - expect(chatEvents[1]).toEqual( - expect.objectContaining({ type: "message.part.delta" }), - ); + expect(chatEvents[0]).toEqual(expect.objectContaining({ type: 'message.updated' })); + expect(chatEvents[1]).toEqual(expect.objectContaining({ type: 'message.part.delta' })); // session.created from snapshot replay + session.status from sendRaw expect(serviceEvents).toHaveLength(2); - expect(serviceEvents[0]).toEqual( - expect.objectContaining({ type: "session.created" }), - ); - expect(serviceEvents[1]).toEqual( - expect.objectContaining({ type: "session.status" }), - ); + expect(serviceEvents[0]).toEqual(expect.objectContaining({ type: 'session.created' })); + expect(serviceEvents[1]).toEqual(expect.objectContaining({ type: 'session.status' })); transport.destroy(); }); - it("ignores invalid events", () => { + it('ignores invalid events', () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); - mockWs.onmessage?.({ data: "not json at all" } as MessageEvent); + mockWs.onmessage?.({ data: 'not json at all' } as MessageEvent); expect(chatEvents).toHaveLength(0); expect(serviceEvents).toHaveLength(0); @@ -262,7 +245,7 @@ describe("CloudAgentTransport event routing", () => { transport.destroy(); }); - it("drops suggestion events since cloud-agent has no accept/dismiss command path", async () => { + it('drops suggestion events since cloud-agent has no accept/dismiss command path', async () => { const { transport, chatEvents, serviceEvents } = createTransportWithSinks(); transport.connect(); @@ -271,14 +254,14 @@ describe("CloudAgentTransport event routing", () => { const serviceCountBefore = serviceEvents.length; sendRaw( - kilocode("suggestion.shown", { - id: "sug-1", - text: "review your changes", - actions: [{ label: "review", prompt: "/review" }], - }), + kilocode('suggestion.shown', { + id: 'sug-1', + text: 'review your changes', + actions: [{ label: 'review', prompt: '/review' }], + }) ); - sendRaw(kilocode("suggestion.accepted", { requestID: "sug-1", index: 0 })); - sendRaw(kilocode("suggestion.dismissed", { requestID: "sug-1" })); + sendRaw(kilocode('suggestion.accepted', { requestID: 'sug-1', index: 0 })); + sendRaw(kilocode('suggestion.dismissed', { requestID: 'sug-1' })); expect(serviceEvents).toHaveLength(serviceCountBefore); expect(chatEvents).toHaveLength(0); @@ -287,86 +270,80 @@ describe("CloudAgentTransport event routing", () => { }); }); -describe("CloudAgentTransport unexpected disconnect", () => { - it("synthesizes stopped event on unexpected disconnect", async () => { +describe('CloudAgentTransport unexpected disconnect', () => { + it('synthesizes stopped event on unexpected disconnect', async () => { const { transport, serviceEvents } = createTransportWithSinks(); transport.connect(); await flushPromises(); sendRaw( - kilocode("session.status", { - sessionID: "ses-1", - status: { type: "busy" }, - }), + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) ); // Non-auth close triggers onUnexpectedDisconnect in connection.ts mockWs.onclose?.({ code: 1011, - reason: "network dropped", + reason: 'network dropped', wasClean: false, } as CloseEvent); - const stoppedEvents = serviceEvents.filter((e) => e.type === "stopped"); + const stoppedEvents = serviceEvents.filter(e => e.type === 'stopped'); expect(stoppedEvents).toHaveLength(1); expect(stoppedEvents[0]).toEqual({ - type: "stopped", - reason: "transport-disconnected", + type: 'stopped', + reason: 'transport-disconnected', }); transport.destroy(); }); - it("suppresses synthetic stopped if already received via event pipeline", () => { + it('suppresses synthetic stopped if already received via event pipeline', () => { const { transport, serviceEvents } = createTransportWithSinks(); transport.connect(); sendRaw( - kilocode("session.status", { - sessionID: "ses-1", - status: { type: "busy" }, - }), + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) ); // complete → stopped(complete) through normal pipeline - sendRaw(createEvent("complete", { currentBranch: "main" })); + sendRaw(createEvent('complete', { currentBranch: 'main' })); - const stoppedBefore = serviceEvents.filter( - (e) => e.type === "stopped", - ).length; + const stoppedBefore = serviceEvents.filter(e => e.type === 'stopped').length; // Now close unexpectedly — should NOT generate another stopped mockWs.onclose?.({ code: 1011, - reason: "network dropped", + reason: 'network dropped', wasClean: false, } as CloseEvent); - const stoppedAfter = serviceEvents.filter( - (e) => e.type === "stopped", - ).length; + const stoppedAfter = serviceEvents.filter(e => e.type === 'stopped').length; expect(stoppedAfter).toBe(stoppedBefore); transport.destroy(); }); }); -describe("CloudAgentTransport ticket handling", () => { - it("calls getTicket with sessionId", () => { - const getTicket = jest.fn((_sessionId: string) => "ticket-abc"); +describe('CloudAgentTransport ticket handling', () => { + it('calls getTicket with sessionId', () => { + const getTicket = jest.fn((_sessionId: string) => 'ticket-abc'); const { transport } = createTransportWithSinks(getTicket); transport.connect(); - expect(getTicket).toHaveBeenCalledWith("ses-1"); + expect(getTicket).toHaveBeenCalledWith('ses-1'); transport.destroy(); }); - it("handles async getTicket", async () => { - const getTicket = jest.fn((_sessionId: string) => - Promise.resolve("async-ticket"), - ); + it('handles async getTicket', async () => { + const getTicket = jest.fn((_sessionId: string) => Promise.resolve('async-ticket')); const { transport, serviceEvents } = createTransportWithSinks(getTicket); transport.connect(); @@ -377,26 +354,26 @@ describe("CloudAgentTransport ticket handling", () => { // Events still route correctly (session.created from replay + session.status from sendRaw) sendRaw( - kilocode("session.status", { - sessionID: "ses-1", - status: { type: "busy" }, - }), + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) ); expect(serviceEvents).toHaveLength(2); transport.destroy(); }); - it("refreshes an expiring ticket before opening the websocket", async () => { + it('refreshes an expiring ticket before opening the websocket', async () => { const nowSeconds = Math.floor(Date.now() / 1000); const getTicket = jest .fn() .mockResolvedValueOnce({ - ticket: "expiring-ticket", + ticket: 'expiring-ticket', expiresAt: nowSeconds + 5, }) .mockResolvedValueOnce({ - ticket: "fresh-ticket", + ticket: 'fresh-ticket', expiresAt: nowSeconds + 60, }); const { transport } = createTransportWithSinks(getTicket); @@ -408,16 +385,14 @@ describe("CloudAgentTransport ticket handling", () => { expect(getTicket).toHaveBeenCalledTimes(2); expect(webSocketConstructor).toHaveBeenCalledTimes(1); - expect(webSocketConstructor.mock.calls[0]?.[0]).toContain( - "ticket=fresh-ticket", - ); + expect(webSocketConstructor.mock.calls[0]?.[0]).toContain('ticket=fresh-ticket'); transport.destroy(); }); }); -describe("CloudAgentTransport lifecycle", () => { - it("disconnect() closes connection", async () => { +describe('CloudAgentTransport lifecycle', () => { + it('disconnect() closes connection', async () => { const { transport } = createTransportWithSinks(); transport.connect(); @@ -427,7 +402,7 @@ describe("CloudAgentTransport lifecycle", () => { expect(mockWs.close).toHaveBeenCalled(); }); - it("destroy() closes connection", async () => { + it('destroy() closes connection', async () => { const { transport } = createTransportWithSinks(); transport.connect(); @@ -437,13 +412,13 @@ describe("CloudAgentTransport lifecycle", () => { expect(mockWs.close).toHaveBeenCalled(); }); - it("stale generation after disconnect prevents connection creation", async () => { + it('stale generation after disconnect prevents connection creation', async () => { const resolveTicket: { resolve?: (value: string) => void } = {}; const getTicket = jest.fn( () => - new Promise((resolve) => { + new Promise(resolve => { resolveTicket.resolve = resolve; - }), + }) ); const { transport } = createTransportWithSinks(getTicket); @@ -453,13 +428,12 @@ describe("CloudAgentTransport lifecycle", () => { transport.disconnect(); // Now resolve the ticket — should be stale - resolveTicket.resolve?.("late-ticket"); + resolveTicket.resolve?.('late-ticket'); await flushPromises(); // Only the first WebSocket (from disconnect closing) should exist; // no new WebSocket created from the stale ticket resolution - const constructorCallsAfterDisconnect = - webSocketConstructor.mock.calls.length; + const constructorCallsAfterDisconnect = webSocketConstructor.mock.calls.length; // The initial connect() didn't create a WS (ticket was async and unresolved), // so no WS should have been constructed at all. @@ -467,77 +441,77 @@ describe("CloudAgentTransport lifecycle", () => { }); }); -describe("CloudAgentTransport command delegation", () => { - it("converts a Kilo model ref before delegating to api.send", async () => { +describe('CloudAgentTransport command delegation', () => { + it('converts a Kilo model ref before delegating to api.send', async () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); await transport.send!({ payload: { - type: "prompt", - prompt: "hello", - mode: "code", - model: { providerID: "kilo", modelID: "gpt-4" }, + type: 'prompt', + prompt: 'hello', + mode: 'code', + model: { providerID: 'kilo', modelID: 'gpt-4' }, }, }); expect(api.send).toHaveBeenCalledWith({ - sessionId: "ses-1", + sessionId: 'ses-1', payload: { - type: "prompt", - prompt: "hello", - mode: "code", - model: "gpt-4", + type: 'prompt', + prompt: 'hello', + mode: 'code', + model: 'gpt-4', }, }); transport.destroy(); }); - it("rejects a non-Kilo model ref before calling the Cloud Agent API", async () => { + it('rejects a non-Kilo model ref before calling the Cloud Agent API', async () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); await expect( transport.send!({ payload: { - type: "prompt", - prompt: "hello", - mode: "code", - model: { providerID: "anthropic", modelID: "claude-sonnet-4" }, + type: 'prompt', + prompt: 'hello', + mode: 'code', + model: { providerID: 'anthropic', modelID: 'claude-sonnet-4' }, }, - }), - ).rejects.toThrow("Cloud Agent only supports Kilo models"); + }) + ).rejects.toThrow('Cloud Agent only supports Kilo models'); expect(api.send).not.toHaveBeenCalled(); transport.destroy(); }); - it("converts Kilo model refs while preserving canonical document attachments", async () => { + it('converts Kilo model refs while preserving canonical document attachments', async () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); const attachments = { - path: "12345678-1234-4234-9234-123456789abc", - files: ["87654321-4321-4321-8321-cba987654321.pdf"], + path: '12345678-1234-4234-9234-123456789abc', + files: ['87654321-4321-4321-8321-cba987654321.pdf'], }; await transport.send!({ payload: { - type: "prompt", - prompt: "read it", - mode: "code", - model: { providerID: "kilo", modelID: "gpt-4" }, + type: 'prompt', + prompt: 'read it', + mode: 'code', + model: { providerID: 'kilo', modelID: 'gpt-4' }, }, attachments, }); expect(api.send).toHaveBeenCalledWith({ - sessionId: "ses-1", + sessionId: 'ses-1', payload: { - type: "prompt", - prompt: "read it", - mode: "code", - model: "gpt-4", + type: 'prompt', + prompt: 'read it', + mode: 'code', + model: 'gpt-4', }, attachments, }); @@ -545,59 +519,59 @@ describe("CloudAgentTransport command delegation", () => { transport.destroy(); }); - it("interrupt() delegates to api.interrupt with bound sessionId", () => { + it('interrupt() delegates to api.interrupt with bound sessionId', () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); void transport.interrupt!(); - expect(api.interrupt).toHaveBeenCalledWith({ sessionId: "ses-1" }); + expect(api.interrupt).toHaveBeenCalledWith({ sessionId: 'ses-1' }); transport.destroy(); }); - it("answer() delegates to api.answer with bound sessionId", () => { + it('answer() delegates to api.answer with bound sessionId', () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); - void transport.answer!({ requestId: "req-1", answers: [["yes"]] }); + void transport.answer!({ requestId: 'req-1', answers: [['yes']] }); expect(api.answer).toHaveBeenCalledWith({ - sessionId: "ses-1", - requestId: "req-1", - answers: [["yes"]], + sessionId: 'ses-1', + requestId: 'req-1', + answers: [['yes']], }); transport.destroy(); }); - it("reject() delegates to api.reject with bound sessionId", () => { + it('reject() delegates to api.reject with bound sessionId', () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); - void transport.reject!({ requestId: "req-2" }); + void transport.reject!({ requestId: 'req-2' }); expect(api.reject).toHaveBeenCalledWith({ - sessionId: "ses-1", - requestId: "req-2", + sessionId: 'ses-1', + requestId: 'req-2', }); transport.destroy(); }); - it("respondToPermission() delegates to api.respondToPermission with bound sessionId", () => { + it('respondToPermission() delegates to api.respondToPermission with bound sessionId', () => { const api = createMockApi(); const { transport } = createTransportWithSinks(undefined, undefined, api); void transport.respondToPermission!({ - requestId: "req-3", - response: "once", + requestId: 'req-3', + response: 'once', }); expect(api.respondToPermission).toHaveBeenCalledWith({ - sessionId: "ses-1", - requestId: "req-3", - response: "once", + sessionId: 'ses-1', + requestId: 'req-3', + response: 'once', }); transport.destroy(); @@ -608,7 +582,7 @@ describe("CloudAgentTransport command delegation", () => { // Snapshot refetch on reconnect // --------------------------------------------------------------------------- -describe("CloudAgentTransport snapshot refetch on reconnect", () => { +describe('CloudAgentTransport snapshot refetch on reconnect', () => { // Microtask-based flush that works under jest.useFakeTimers() // (unlike flushPromises which uses setTimeout and hangs with fake timers) async function flushMicrotasks(): Promise { @@ -618,7 +592,7 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { } function createTransportWithControllableSnapshot( - snapshotOverride?: ReturnType, + snapshotOverride?: ReturnType ) { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; @@ -626,17 +600,17 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { const fetchSnapshot = jest.fn(() => Promise.resolve(snapshot)); const factory = createCloudAgentTransport({ - sessionId: cloudAgentId("ses-1"), - kiloSessionId: kiloId("ses-1"), + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), api: createMockApi(), - getTicket: () => "test-ticket", + getTicket: () => 'test-ticket', fetchSnapshot, - websocketBaseUrl: "ws://localhost:9999", + websocketBaseUrl: 'ws://localhost:9999', }); const transport = factory({ - onChatEvent: (event) => chatEvents.push(event), - onServiceEvent: (event) => serviceEvents.push(event), + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), }); return { transport, chatEvents, serviceEvents, fetchSnapshot }; @@ -651,46 +625,42 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { return { eventId: 0, executionId: null, - sessionId: "ses-1", - streamEventType: "kilocode", + sessionId: 'ses-1', + streamEventType: 'kilocode', timestamp: new Date().toISOString(), data: { - type: "session.status", - properties: { sessionID: "ses-1", status: { type: "busy" } }, + type: 'session.status', + properties: { sessionID: 'ses-1', status: { type: 'busy' } }, }, }; } /** Establish connection, simulate close + reconnect, return the new WS mock. */ - async function simulateReconnect( - establishEvent?: CloudAgentEvent, - ): Promise { - mockWs.onclose?.({ code: 1006, reason: "", wasClean: false } as CloseEvent); + async function simulateReconnect(establishEvent?: CloudAgentEvent): Promise { + mockWs.onclose?.({ code: 1006, reason: '', wasClean: false } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); - const newMockWs = webSocketConstructor.mock.results.at(-1) - ?.value as MockWebSocket; + const newMockWs = webSocketConstructor.mock.results.at(-1)?.value as MockWebSocket; - newMockWs.onopen?.(new Event("open")); + newMockWs.onopen?.(new Event('open')); sendRawOn( newMockWs, establishEvent ?? - kilocode("session.status", { - sessionID: "ses-1", - status: { type: "busy" }, - }), + kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, + }) ); return newMockWs; } - it("resumes from the replay cursor on reconnect instead of refetching the snapshot", async () => { + it('resumes from the replay cursor on reconnect instead of refetching the snapshot', async () => { jest.useFakeTimers(); try { - const { transport, serviceEvents, fetchSnapshot } = - createTransportWithControllableSnapshot(); + const { transport, serviceEvents, fetchSnapshot } = createTransportWithControllableSnapshot(); transport.connect(); await flushMicrotasks(); @@ -698,9 +668,9 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { expect(fetchSnapshot).toHaveBeenCalledTimes(1); // Establish connection with a persisted event — its id becomes the cursor. - const establish = kilocode("session.status", { - sessionID: "ses-1", - status: { type: "busy" }, + const establish = kilocode('session.status', { + sessionID: 'ses-1', + status: { type: 'busy' }, }); sendRaw(establish); @@ -714,17 +684,17 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { expect(fetchSnapshot).toHaveBeenCalledTimes(1); const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); expect(reconnectUrl).toContain(`fromId=${establish.eventId}`); - expect(reconnectUrl).not.toContain("replay=false"); + expect(reconnectUrl).not.toContain('replay=false'); const replayedCreated = serviceEvents .slice(serviceCountBefore) - .filter((e) => e.type === "session.created"); + .filter(e => e.type === 'session.created'); expect(replayedCreated).toHaveLength(0); transport.destroy(); newMockWs.onclose?.({ code: 1000, - reason: "", + reason: '', wasClean: true, } as CloseEvent); } finally { @@ -732,26 +702,26 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { } }); - it("refetches the snapshot on reconnect when no replay cursor exists and upserts it", async () => { + it('refetches the snapshot on reconnect when no replay cursor exists and upserts it', async () => { jest.useFakeTimers(); try { - const snapshotWithMessages = makeSnapshot({ id: "ses-1" }, [ + const snapshotWithMessages = makeSnapshot({ id: 'ses-1' }, [ { info: { - id: "msg-1", - sessionID: "ses-1", - role: "user", + id: 'msg-1', + sessionID: 'ses-1', + role: 'user', time: { created: 1 }, - agent: "build", - model: { providerID: "a", modelID: "b" }, + agent: 'build', + model: { providerID: 'a', modelID: 'b' }, }, parts: [ { - id: "part-1", - sessionID: "ses-1", - messageID: "msg-1", - type: "text", - text: "hello", + id: 'part-1', + sessionID: 'ses-1', + messageID: 'msg-1', + type: 'text', + text: 'hello', }, ], }, @@ -763,15 +733,9 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { transport.connect(); await flushMicrotasks(); - expect( - serviceEvents.filter((e) => e.type === "session.created"), - ).toHaveLength(1); - expect( - chatEvents.filter((e) => e.type === "message.updated"), - ).toHaveLength(1); - expect( - chatEvents.filter((e) => e.type === "message.part.updated"), - ).toHaveLength(1); + expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(1); + expect(chatEvents.filter(e => e.type === 'message.updated')).toHaveLength(1); + expect(chatEvents.filter(e => e.type === 'message.part.updated')).toHaveLength(1); // Establish connection with sentinel events only — no replay cursor, so // reconnect falls back to the snapshot refetch. @@ -781,20 +745,14 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { await flushMicrotasks(); expect(fetchSnapshot).toHaveBeenCalledTimes(2); - expect( - serviceEvents.filter((e) => e.type === "session.created"), - ).toHaveLength(2); - expect( - chatEvents.filter((e) => e.type === "message.updated"), - ).toHaveLength(2); - expect( - chatEvents.filter((e) => e.type === "message.part.updated"), - ).toHaveLength(2); + expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(2); + expect(chatEvents.filter(e => e.type === 'message.updated')).toHaveLength(2); + expect(chatEvents.filter(e => e.type === 'message.part.updated')).toHaveLength(2); transport.destroy(); newMockWs.onclose?.({ code: 1000, - reason: "", + reason: '', wasClean: true, } as CloseEvent); } finally { @@ -802,28 +760,25 @@ describe("CloudAgentTransport snapshot refetch on reconnect", () => { } }); - it("initial connect fetches snapshot once and opens WebSocket", async () => { - const { transport, serviceEvents, fetchSnapshot } = - createTransportWithControllableSnapshot(); + it('initial connect fetches snapshot once and opens WebSocket', async () => { + const { transport, serviceEvents, fetchSnapshot } = createTransportWithControllableSnapshot(); transport.connect(); await flushPromises(); expect(fetchSnapshot).toHaveBeenCalledTimes(1); - expect(fetchSnapshot).toHaveBeenCalledWith("ses-1"); + expect(fetchSnapshot).toHaveBeenCalledWith('ses-1'); expect(webSocketConstructor).toHaveBeenCalledTimes(1); - expect( - serviceEvents.filter((e) => e.type === "session.created"), - ).toHaveLength(1); + expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(1); transport.destroy(); }); }); -describe("CloudAgentTransport page-seam", () => { +describe('CloudAgentTransport page-seam', () => { function createTransportWithPageFetch( page: SessionSnapshotPageOutcome | null, - onError?: (message: string) => void, + onError?: (message: string) => void ) { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; @@ -831,21 +786,20 @@ describe("CloudAgentTransport page-seam", () => { const onInitialPageLoaded = jest.fn(); const factory = createCloudAgentTransport({ - sessionId: cloudAgentId("ses-1"), - kiloSessionId: kiloId("ses-1"), + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), api: createMockApi(), - getTicket: () => "test-ticket", - fetchSnapshot: () => - Promise.reject(new Error("legacy fetchSnapshot should not be called")), + getTicket: () => 'test-ticket', + fetchSnapshot: () => Promise.reject(new Error('legacy fetchSnapshot should not be called')), fetchSnapshotPage, onInitialPageLoaded, - websocketBaseUrl: "ws://localhost:9999", + websocketBaseUrl: 'ws://localhost:9999', onError, }); const transport = factory({ - onChatEvent: (event) => chatEvents.push(event), - onServiceEvent: (event) => serviceEvents.push(event), + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), }); return { @@ -857,12 +811,12 @@ describe("CloudAgentTransport page-seam", () => { }; } - it("uses fetchSnapshotPage for the initial bounded read and reports the page via onInitialPageLoaded", async () => { + it('uses fetchSnapshotPage for the initial bounded read and reports the page via onInitialPageLoaded', async () => { const page = { - kind: "success" as const, - info: { id: "ses-1" }, + kind: 'success' as const, + info: { id: 'ses-1' }, messages: [], - nextCursor: "cursor-A", + nextCursor: 'cursor-A', omittedItemCount: 2, }; const { transport, fetchSnapshotPage, onInitialPageLoaded } = @@ -872,23 +826,21 @@ describe("CloudAgentTransport page-seam", () => { await flushPromises(); expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); - expect(fetchSnapshotPage).toHaveBeenCalledWith("ses-1", {}); + expect(fetchSnapshotPage).toHaveBeenCalledWith('ses-1', {}); expect(onInitialPageLoaded).toHaveBeenCalledWith(page); transport.destroy(); }); - it("surfaces typed failures on the initial read via onError", async () => { + it('surfaces typed failures on the initial read via onError', async () => { const { transport, serviceEvents } = createTransportWithPageFetch({ - kind: "retryable_failure", + kind: 'retryable_failure', }); transport.connect(); await flushPromises(); - expect( - serviceEvents.filter((e) => e.type === "session.created"), - ).toHaveLength(0); + expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(0); // The transport still connects the websocket so the user can recover // via live events. expect(webSocketConstructor).toHaveBeenCalledTimes(1); @@ -896,21 +848,17 @@ describe("CloudAgentTransport page-seam", () => { transport.destroy(); }); - it("surfaces invalid_data on the initial read via onError and still connects the websocket", async () => { + it('surfaces invalid_data on the initial read via onError and still connects the websocket', async () => { const errors: string[] = []; const { transport, chatEvents, serviceEvents, onInitialPageLoaded } = - createTransportWithPageFetch({ kind: "invalid_data" }, (message) => - errors.push(message), - ); + createTransportWithPageFetch({ kind: 'invalid_data' }, message => errors.push(message)); transport.connect(); await flushPromises(); - expect(errors).toEqual(["Session history is unavailable"]); + expect(errors).toEqual(['Session history is unavailable']); expect(onInitialPageLoaded).not.toHaveBeenCalled(); - expect( - serviceEvents.filter((e) => e.type === "session.created"), - ).toHaveLength(0); + expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(0); expect(chatEvents).toHaveLength(0); // Terminal typed failures still leave the websocket available so the user // can recover via live events; this is the intended transport behavior. @@ -919,21 +867,17 @@ describe("CloudAgentTransport page-seam", () => { transport.destroy(); }); - it("surfaces too_large on the initial read via onError and still connects the websocket", async () => { + it('surfaces too_large on the initial read via onError and still connects the websocket', async () => { const errors: string[] = []; const { transport, chatEvents, serviceEvents, onInitialPageLoaded } = - createTransportWithPageFetch({ kind: "too_large" }, (message) => - errors.push(message), - ); + createTransportWithPageFetch({ kind: 'too_large' }, message => errors.push(message)); transport.connect(); await flushPromises(); - expect(errors).toEqual(["Session history too large to load"]); + expect(errors).toEqual(['Session history too large to load']); expect(onInitialPageLoaded).not.toHaveBeenCalled(); - expect( - serviceEvents.filter((e) => e.type === "session.created"), - ).toHaveLength(0); + expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(0); expect(chatEvents).toHaveLength(0); // Terminal typed failures still leave the websocket available so the user // can recover via live events; this is the intended transport behavior. @@ -942,37 +886,35 @@ describe("CloudAgentTransport page-seam", () => { transport.destroy(); }); - it("falls back to fetchSnapshot when fetchSnapshotPage is not provided", async () => { + it('falls back to fetchSnapshot when fetchSnapshotPage is not provided', async () => { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; const fetchSnapshot = jest.fn(() => Promise.resolve(emptySnapshot)); const factory = createCloudAgentTransport({ - sessionId: cloudAgentId("ses-1"), - kiloSessionId: kiloId("ses-1"), + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), api: createMockApi(), - getTicket: () => "test-ticket", + getTicket: () => 'test-ticket', fetchSnapshot, - websocketBaseUrl: "ws://localhost:9999", + websocketBaseUrl: 'ws://localhost:9999', }); const transport = factory({ - onChatEvent: (event) => chatEvents.push(event), - onServiceEvent: (event) => serviceEvents.push(event), + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), }); transport.connect(); await flushPromises(); expect(fetchSnapshot).toHaveBeenCalledTimes(1); - expect( - serviceEvents.filter((e) => e.type === "session.created"), - ).toHaveLength(1); + expect(serviceEvents.filter(e => e.type === 'session.created')).toHaveLength(1); transport.destroy(); }); - it("reconnect (no eventId) uses fetchSnapshotPage and does NOT fire onInitialPageLoaded", async () => { + it('reconnect (no eventId) uses fetchSnapshotPage and does NOT fire onInitialPageLoaded', async () => { jest.useFakeTimers(); try { // Microtask-based flush that works under jest.useFakeTimers() @@ -984,10 +926,10 @@ describe("CloudAgentTransport page-seam", () => { } const page = { - kind: "success" as const, - info: { id: "ses-1" }, + kind: 'success' as const, + info: { id: 'ses-1' }, messages: [], - nextCursor: "cursor-A", + nextCursor: 'cursor-A', omittedItemCount: 0, }; const fetchSnapshotPage = jest.fn().mockResolvedValue(page); @@ -996,19 +938,19 @@ describe("CloudAgentTransport page-seam", () => { const serviceEvents: ServiceEvent[] = []; const factory = createCloudAgentTransport({ - sessionId: cloudAgentId("ses-1"), - kiloSessionId: kiloId("ses-1"), + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), api: createMockApi(), - getTicket: () => "test-ticket", - fetchSnapshot: () => Promise.reject(new Error("should not be called")), + getTicket: () => 'test-ticket', + fetchSnapshot: () => Promise.reject(new Error('should not be called')), fetchSnapshotPage, onInitialPageLoaded, - websocketBaseUrl: "ws://localhost:9999", + websocketBaseUrl: 'ws://localhost:9999', }); const transport = factory({ - onChatEvent: (event) => chatEvents.push(event), - onServiceEvent: (event) => serviceEvents.push(event), + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), }); transport.connect(); @@ -1022,12 +964,12 @@ describe("CloudAgentTransport page-seam", () => { const sentinel: CloudAgentEvent = { eventId: 0, executionId: null, - sessionId: "ses-1", - streamEventType: "kilocode", + sessionId: 'ses-1', + streamEventType: 'kilocode', timestamp: new Date().toISOString(), data: { - type: "session.status", - properties: { sessionID: "ses-1", status: { type: "busy" } }, + type: 'session.status', + properties: { sessionID: 'ses-1', status: { type: 'busy' } }, }, }; sendRaw(sentinel); @@ -1035,15 +977,14 @@ describe("CloudAgentTransport page-seam", () => { // Trigger an unexpected disconnect on the first socket. mockWs.onclose?.({ code: 1006, - reason: "", + reason: '', wasClean: false, } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); - const newMockWs = webSocketConstructor.mock.results.at(-1) - ?.value as MockWebSocket; - newMockWs.onopen?.(new Event("open")); + const newMockWs = webSocketConstructor.mock.results.at(-1)?.value as MockWebSocket; + newMockWs.onopen?.(new Event('open')); // Send a sentinel on the new socket: the connection marks itself // reconnected, then `onReconnected` falls back to `fetchSnapshotPage` // because `lastEventId` is still null. @@ -1063,11 +1004,11 @@ describe("CloudAgentTransport page-seam", () => { } }); - describe("watermark cursor", () => { - it("uses fromId on first connect when the page carries a watermark", async () => { + describe('watermark cursor', () => { + it('uses fromId on first connect when the page carries a watermark', async () => { const page = { - kind: "success" as const, - info: { id: "ses-1" } as const, + kind: 'success' as const, + info: { id: 'ses-1' } as const, messages: [], nextCursor: null, omittedItemCount: 0, @@ -1080,16 +1021,16 @@ describe("CloudAgentTransport page-seam", () => { expect(webSocketConstructor).toHaveBeenCalledTimes(1); const wsUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(wsUrl).toContain("fromId=42"); - expect(wsUrl).not.toContain("replay=false"); + expect(wsUrl).toContain('fromId=42'); + expect(wsUrl).not.toContain('replay=false'); transport.destroy(); }); - it("uses replay=false on first connect when the page has no watermark", async () => { + 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, + kind: 'success' as const, + info: { id: 'ses-1' } as const, messages: [], nextCursor: null, omittedItemCount: 0, @@ -1101,13 +1042,13 @@ describe("CloudAgentTransport page-seam", () => { expect(webSocketConstructor).toHaveBeenCalledTimes(1); const wsUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(wsUrl).toContain("replay=false"); - expect(wsUrl).not.toContain("fromId"); + 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 () => { + it('uses fromId with watermark null on reconnect when wire events have set lastEventId', async () => { jest.useFakeTimers(); try { async function flushMicrotasks(): Promise { @@ -1117,8 +1058,8 @@ describe("CloudAgentTransport page-seam", () => { } const page = { - kind: "success" as const, - info: { id: "ses-1" } as const, + kind: 'success' as const, + info: { id: 'ses-1' } as const, messages: [], nextCursor: null, omittedItemCount: 0, @@ -1129,19 +1070,18 @@ describe("CloudAgentTransport page-seam", () => { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; const factory = createCloudAgentTransport({ - sessionId: cloudAgentId("ses-1"), - kiloSessionId: kiloId("ses-1"), + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), api: createMockApi(), - getTicket: () => "test-ticket", - fetchSnapshot: () => - Promise.reject(new Error("should not be called")), + getTicket: () => 'test-ticket', + fetchSnapshot: () => Promise.reject(new Error('should not be called')), fetchSnapshotPage, - websocketBaseUrl: "ws://localhost:9999", + websocketBaseUrl: 'ws://localhost:9999', }); const transport = factory({ - onChatEvent: (event) => chatEvents.push(event), - onServiceEvent: (event) => serviceEvents.push(event), + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), }); transport.connect(); @@ -1149,18 +1089,18 @@ describe("CloudAgentTransport page-seam", () => { // Verify first connect used the watermark const firstUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(firstUrl).toContain("fromId=10"); + expect(firstUrl).toContain('fromId=10'); // Send a wire event with eventId > 0 to advance the live cursor const establish = { eventId: 55, executionId: null, - sessionId: "ses-1", - streamEventType: "kilocode", + sessionId: 'ses-1', + streamEventType: 'kilocode', timestamp: new Date().toISOString(), data: { - type: "session.status", - properties: { sessionID: "ses-1", status: { type: "busy" } }, + type: 'session.status', + properties: { sessionID: 'ses-1', status: { type: 'busy' } }, }, }; sendRaw(establish); @@ -1168,17 +1108,15 @@ describe("CloudAgentTransport page-seam", () => { // Disconnect and reconnect mockWs.onclose?.({ code: 1006, - reason: "", + reason: '', wasClean: false, } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); - const reconnectUrl = String( - webSocketConstructor.mock.calls.at(-1)?.[0], - ); + 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"); + 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); @@ -1189,7 +1127,7 @@ describe("CloudAgentTransport page-seam", () => { } }); - it("reconnect on a watermarked page with only sentinel events uses the original watermark and skips the page refetch", async () => { + it('reconnect on a watermarked page with only sentinel events uses the original watermark and skips the page refetch', async () => { jest.useFakeTimers(); try { async function flushMicrotasks(): Promise { @@ -1199,8 +1137,8 @@ describe("CloudAgentTransport page-seam", () => { } const page = { - kind: "success" as const, - info: { id: "ses-1" } as const, + kind: 'success' as const, + info: { id: 'ses-1' } as const, messages: [], nextCursor: null, omittedItemCount: 0, @@ -1211,19 +1149,18 @@ describe("CloudAgentTransport page-seam", () => { const chatEvents: ChatEvent[] = []; const serviceEvents: ServiceEvent[] = []; const factory = createCloudAgentTransport({ - sessionId: cloudAgentId("ses-1"), - kiloSessionId: kiloId("ses-1"), + sessionId: cloudAgentId('ses-1'), + kiloSessionId: kiloId('ses-1'), api: createMockApi(), - getTicket: () => "test-ticket", - fetchSnapshot: () => - Promise.reject(new Error("should not be called")), + getTicket: () => 'test-ticket', + fetchSnapshot: () => Promise.reject(new Error('should not be called')), fetchSnapshotPage, - websocketBaseUrl: "ws://localhost:9999", + websocketBaseUrl: 'ws://localhost:9999', }); const transport = factory({ - onChatEvent: (event) => chatEvents.push(event), - onServiceEvent: (event) => serviceEvents.push(event), + onChatEvent: event => chatEvents.push(event), + onServiceEvent: event => serviceEvents.push(event), }); transport.connect(); @@ -1231,7 +1168,7 @@ describe("CloudAgentTransport page-seam", () => { // Verify first connect used the watermark as fromId. const firstUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(firstUrl).toContain("fromId=42"); + expect(firstUrl).toContain('fromId=42'); expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); // Send only sentinel events (eventId: 0) — they do NOT advance the @@ -1239,12 +1176,12 @@ describe("CloudAgentTransport page-seam", () => { const sentinel: CloudAgentEvent = { eventId: 0, executionId: null, - sessionId: "ses-1", - streamEventType: "kilocode", + sessionId: 'ses-1', + streamEventType: 'kilocode', timestamp: new Date().toISOString(), data: { - type: "session.status", - properties: { sessionID: "ses-1", status: { type: "busy" } }, + type: 'session.status', + properties: { sessionID: 'ses-1', status: { type: 'busy' } }, }, }; sendRaw(sentinel); @@ -1252,19 +1189,17 @@ describe("CloudAgentTransport page-seam", () => { // Disconnect and reconnect. mockWs.onclose?.({ code: 1006, - reason: "", + reason: '', wasClean: false, } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); - const reconnectUrl = String( - webSocketConstructor.mock.calls.at(-1)?.[0], - ); + const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); // Reconnect must still use the watermark (42) because sentinel events // never advance the cursor past the watermark. - expect(reconnectUrl).toContain("fromId=42"); - expect(reconnectUrl).not.toContain("replay=false"); + expect(reconnectUrl).toContain('fromId=42'); + expect(reconnectUrl).not.toContain('replay=false'); // No second page fetch — lastEventId (watermark 42) is not null, so // onReconnected skips the snapshot fallback path entirely. @@ -1282,7 +1217,7 @@ describe("CloudAgentTransport page-seam", () => { // Exactly-once event replay // --------------------------------------------------------------------------- -describe("CloudAgentTransport 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. @@ -1291,17 +1226,17 @@ describe("CloudAgentTransport exactly-once event replay", () => { return { eventId, executionId: null, - sessionId: "ses-1", - streamEventType: "kilocode", + sessionId: 'ses-1', + streamEventType: 'kilocode', timestamp: new Date().toISOString(), data: { - type: "session.status", - properties: { sessionID: "ses-1", status: { type: "busy" } }, + type: 'session.status', + properties: { sessionID: 'ses-1', status: { type: 'busy' } }, }, }; } - it("drops a duplicate event ID at the cursor", async () => { + it('drops a duplicate event ID at the cursor', async () => { const { transport, serviceEvents } = createTransportWithSinks(); transport.connect(); @@ -1321,7 +1256,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { transport.destroy(); }); - it("drops a positive event ID below the cursor", async () => { + it('drops a positive event ID below the cursor', async () => { const { transport, serviceEvents } = createTransportWithSinks(); transport.connect(); @@ -1340,7 +1275,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { transport.destroy(); }); - it("delivers a higher positive event ID and advances the cursor", async () => { + it('delivers a higher positive event ID and advances the cursor', async () => { const { transport, serviceEvents } = createTransportWithSinks(); transport.connect(); @@ -1357,7 +1292,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { transport.destroy(); }); - it("always delivers event ID zero (synthetic sentinel)", async () => { + it('always delivers event ID zero (synthetic sentinel)', async () => { const { transport, serviceEvents } = createTransportWithSinks(); transport.connect(); @@ -1375,7 +1310,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { transport.destroy(); }); - it("replays 3 through 8 after 1 through 5 on resume, delivering each once", async () => { + it('replays 3 through 8 after 1 through 5 on resume, delivering each once', async () => { jest.useFakeTimers(); try { async function flushMicrotasks(): Promise { @@ -1390,15 +1325,15 @@ describe("CloudAgentTransport exactly-once event replay", () => { return { eventId, executionId: null, - sessionId: "ses-1", - streamEventType: "kilocode", + sessionId: 'ses-1', + streamEventType: 'kilocode', timestamp: new Date().toISOString(), data: { - type: "session.status", + type: 'session.status', properties: { - sessionID: "ses-1", + sessionID: 'ses-1', status: { - type: "retry", + type: 'retry', attempt: eventId, message: `exactly-once-${eventId}`, next: eventId + 1, @@ -1421,7 +1356,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { // Disconnect (triggers reconnect via connection.ts reconnection logic). mockWs.onclose?.({ code: 1006, - reason: "", + reason: '', wasClean: false, } as CloseEvent); jest.advanceTimersByTime(2000); @@ -1430,9 +1365,8 @@ describe("CloudAgentTransport exactly-once event replay", () => { // 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")); + 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)), @@ -1444,19 +1378,19 @@ describe("CloudAgentTransport exactly-once event replay", () => { const deliveredIds = serviceEvents .filter( ( - e, + e ): e is { - type: "session.status"; + type: 'session.status'; sessionId: string; status: { - type: "retry"; + type: 'retry'; attempt: number; message: string; next: number; }; - } => e.type === "session.status" && e.status.type === "retry", + } => e.type === 'session.status' && e.status.type === 'retry' ) - .map((e) => e.status.attempt) + .map(e => e.status.attempt) .sort((a, b) => a - b); // Each ID from 1 through 8 must appear exactly once. @@ -1466,7 +1400,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { transport.destroy(); newMockWs.onclose?.({ code: 1000, - reason: "", + reason: '', wasClean: true, } as CloseEvent); } finally { @@ -1474,7 +1408,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { } }); - it("cursor never moves backward after delivering a higher ID", async () => { + it('cursor never moves backward after delivering a higher ID', async () => { jest.useFakeTimers(); try { async function flushMicrotasks(): Promise { @@ -1501,7 +1435,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { // Disconnect and reconnect to inspect the replay cursor. mockWs.onclose?.({ code: 1006, - reason: "", + reason: '', wasClean: false, } as CloseEvent); jest.advanceTimersByTime(2000); @@ -1509,7 +1443,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); // Cursor must still be 10, not 5. - expect(reconnectUrl).toContain("fromId=10"); + expect(reconnectUrl).toContain('fromId=10'); transport.destroy(); } finally { @@ -1517,7 +1451,7 @@ describe("CloudAgentTransport exactly-once event replay", () => { } }); - it("passes through sentinel (ID 0) events without affecting the cursor for replay URLs", async () => { + it('passes through sentinel (ID 0) events without affecting the cursor for replay URLs', async () => { jest.useFakeTimers(); try { async function flushMicrotasks(): Promise { @@ -1545,14 +1479,14 @@ describe("CloudAgentTransport exactly-once event replay", () => { // Disconnect and reconnect — cursor must still be 7. mockWs.onclose?.({ code: 1006, - reason: "", + reason: '', wasClean: false, } as CloseEvent); jest.advanceTimersByTime(2000); await flushMicrotasks(); const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); - expect(reconnectUrl).toContain("fromId=7"); + expect(reconnectUrl).toContain('fromId=7'); transport.destroy(); } finally { diff --git a/packages/cloud-agent-sdk/src/session.test.ts b/packages/cloud-agent-sdk/src/session.test.ts index 3bceaf78298..0d7b8686f3a 100644 --- a/packages/cloud-agent-sdk/src/session.test.ts +++ b/packages/cloud-agent-sdk/src/session.test.ts @@ -511,7 +511,7 @@ describe('remote session create and retry commands', () => { 'create_session', { protocolVersion: 1 }, 'owner', - expect.any(String) as string, + expect.any(String) as string ); expect(userWebConnection.sendCommandToConnection).not.toHaveBeenCalled(); session.destroy(); @@ -546,7 +546,7 @@ describe('remote session create and retry commands', () => { orgId: 'aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee', }, 'owner', - expect.any(String) as string, + expect.any(String) as string ); session.destroy(); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 1e10186d009..5eeb84eec49 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -1,7 +1,7 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; // Mock cloudflare:workers before importing UserConnectionDO -vi.mock("cloudflare:workers", () => ({ +vi.mock('cloudflare:workers', () => ({ DurableObject: class { ctx: unknown; env: unknown; @@ -18,11 +18,11 @@ const sessionIngestMocks = vi.hoisted(() => ({ getSessionIngestDO: vi.fn(), })); -vi.mock("./SessionIngestDO", () => ({ +vi.mock('./SessionIngestDO', () => ({ getSessionIngestDO: sessionIngestMocks.getSessionIngestDO, })); -import { MAX_CATALOG_RESULT_BYTES, UserConnectionDO } from "./UserConnectionDO"; +import { MAX_CATALOG_RESULT_BYTES, UserConnectionDO } from './UserConnectionDO'; // --------------------------------------------------------------------------- // Mock WebSocket @@ -73,20 +73,16 @@ function makeStorageFake() { delete: vi.fn(async (key: string) => { store.delete(key); }), - list: vi.fn( - async (opts?: { - prefix?: string; - }): Promise> => { - const result = new Map(); - const prefix = opts?.prefix ?? ""; - for (const [key, value] of store) { - if (key.startsWith(prefix)) { - result.set(key, value as T); - } + list: vi.fn(async (opts?: { prefix?: string }): Promise> => { + const result = new Map(); + const prefix = opts?.prefix ?? ''; + for (const [key, value] of store) { + if (key.startsWith(prefix)) { + result.set(key, value as T); } - return result; - }, - ), + } + return result; + }), setAlarm: vi.fn(), }; } @@ -109,7 +105,7 @@ function createMockCtx() { return { getWebSockets(tag?: string): MockWS[] { if (!tag) return [...sockets]; - return sockets.filter((ws) => ws._tags.includes(tag)); + return sockets.filter(ws => ws._tags.includes(tag)); }, acceptWebSocket(ws: MockWS, tags: string[]) { ws._tags = tags; @@ -130,10 +126,10 @@ function createMockCtx() { /** Drain microtasks so waitUntil-scheduled async IIFEs settle. */ async function flushAsync(): Promise { - await new Promise((resolve) => { + await new Promise(resolve => { setTimeout(resolve, 0); }); - await new Promise((resolve) => { + await new Promise(resolve => { setTimeout(resolve, 0); }); } @@ -144,14 +140,12 @@ async function flushAsync(): Promise { function makeSession( id: string, - status = "busy", - title = "Test", + status = 'busy', + title = 'Test', parentSessionId?: string, - platform?: string, + platform?: string ) { - const base = platform - ? { id, status, title, platform } - : { id, status, title }; + const base = platform ? { id, status, title, platform } : { id, status, title }; return parentSessionId ? { ...base, parentSessionId } : base; } @@ -162,11 +156,11 @@ function parseSent(ws: MockWS, callIndex = 0): unknown { } function isRecord(v: unknown): v is Record { - return typeof v === "object" && v !== null && !Array.isArray(v); + return typeof v === 'object' && v !== null && !Array.isArray(v); } function allSent(ws: MockWS): Record[] { - return ws.send.mock.calls.map((c) => { + return ws.send.mock.calls.map(c => { const parsed: unknown = JSON.parse(String(c[0])); if (!isRecord(parsed)) { throw new Error(`Expected JSON object but got: ${String(c[0])}`); @@ -178,7 +172,7 @@ function allSent(ws: MockWS): Record[] { /** Extract the correlationId that was sent to CLI for a given command. */ function getCorrelationId(cliWs: MockWS, callIndex = 0): string { const msgs = allSent(cliWs); - const cmdMsgs = msgs.filter((m) => m.type === "command"); + const cmdMsgs = msgs.filter(m => m.type === 'command'); const msg = cmdMsgs[callIndex]; if (!msg) throw new Error(`No command call at index ${callIndex}`); return msg.id as string; @@ -192,58 +186,52 @@ function setup() { return { doInstance, ctx, mockCtx }; } -function connectWebSocket( - doInstance: UserConnectionDO, - connectionId: string, -): MockWS { +function connectWebSocket(doInstance: UserConnectionDO, connectionId: string): MockWS { const client = createMockWs(); const server = createMockWs(); vi.stubGlobal( - "WebSocketPair", + 'WebSocketPair', class { 0 = client; 1 = server; - }, + } ); vi.stubGlobal( - "Response", + 'Response', class { constructor(_body?: BodyInit | null, _init?: ResponseInit) {} - }, + } ); doInstance.fetch( new Request(`http://local/web?connectionId=${connectionId}`, { - headers: { Upgrade: "websocket" }, - }), + headers: { Upgrade: 'websocket' }, + }) ); return server; } -function connectCliSocket( - doInstance: UserConnectionDO, - connectionId: string, -): MockWS { +function connectCliSocket(doInstance: UserConnectionDO, connectionId: string): MockWS { const client = createMockWs(); const server = createMockWs(); vi.stubGlobal( - "WebSocketPair", + 'WebSocketPair', class { 0 = client; 1 = server; - }, + } ); vi.stubGlobal( - "Response", + 'Response', class { constructor(_body?: BodyInit | null, _init?: ResponseInit) {} - }, + } ); doInstance.fetch( new Request(`http://local/cli?connectionId=${connectionId}`, { - headers: { Upgrade: "websocket" }, - }), + headers: { Upgrade: 'websocket' }, + }) ); return server; } @@ -259,18 +247,18 @@ function addCliSocket( platform?: string; }> = [], instance?: { name: string; projectName: string; version?: string }, - kiloUserId?: string, + kiloUserId?: string ): MockWS { const attachment: { - role: "cli"; + role: 'cli'; connectionId: string; sessions: typeof sessions; instance?: typeof instance; kiloUserId?: string; - } = { role: "cli", connectionId, sessions }; + } = { role: 'cli', connectionId, sessions }; if (instance) attachment.instance = instance; if (kiloUserId) attachment.kiloUserId = kiloUserId; - const ws = createMockWs(["cli"], attachment); + const ws = createMockWs(['cli'], attachment); mockCtx.addSocket(ws); return ws; } @@ -278,11 +266,11 @@ function addCliSocket( /** Create a web WebSocket and add it to the context. */ function addWebSocket( mockCtx: ReturnType, - connectionId = "web-1", - subscribedSessions: string[] = [], + connectionId = 'web-1', + subscribedSessions: string[] = [] ): MockWS { - const attachment = { role: "web" as const, connectionId, subscribedSessions }; - const ws = createMockWs(["web"], attachment); + const attachment = { role: 'web' as const, connectionId, subscribedSessions }; + const ws = createMockWs(['web'], attachment); mockCtx.addSocket(ws); return ws; } @@ -301,14 +289,12 @@ function sendHeartbeat( protocolVersion?: string; capabilities?: { attachments?: boolean }; instance?: { name: string; projectName: string; version?: string }; - } = {}, + } = {} ) { const msg = JSON.stringify({ - type: "heartbeat", + type: 'heartbeat', sessions, - ...(options.protocolVersion - ? { protocolVersion: options.protocolVersion } - : {}), + ...(options.protocolVersion ? { protocolVersion: options.protocolVersion } : {}), ...(options.capabilities ? { capabilities: options.capabilities } : {}), ...(options.instance ? { instance: options.instance } : {}), }); @@ -316,28 +302,20 @@ function sendHeartbeat( } /** Send a subscribe from a web ws */ -function sendSubscribe( - doInstance: UserConnectionDO, - webWs: MockWS, - sessionId: string, -) { - const msg = JSON.stringify({ type: "subscribe", sessionId }); +function sendSubscribe(doInstance: UserConnectionDO, webWs: MockWS, sessionId: string) { + const msg = JSON.stringify({ type: 'subscribe', sessionId }); doInstance.webSocketMessage(webWs as never, msg); } /** Send an unsubscribe from a web ws */ -function sendUnsubscribe( - doInstance: UserConnectionDO, - webWs: MockWS, - sessionId: string, -) { - const msg = JSON.stringify({ type: "unsubscribe", sessionId }); +function sendUnsubscribe(doInstance: UserConnectionDO, webWs: MockWS, sessionId: string) { + const msg = JSON.stringify({ type: 'unsubscribe', sessionId }); doInstance.webSocketMessage(webWs as never, msg); } /** Send a viewer ping from a web ws */ function sendPing(doInstance: UserConnectionDO, webWs: MockWS, nonce: string) { - const msg = JSON.stringify({ type: "ping", nonce }); + const msg = JSON.stringify({ type: 'ping', nonce }); doInstance.webSocketMessage(webWs as never, msg); } @@ -352,9 +330,9 @@ function sendCommand( connectionId?: string; data?: unknown; mutationId?: string; - }, + } ) { - const msg = JSON.stringify({ type: "command", ...opts }); + const msg = JSON.stringify({ type: 'command', ...opts }); return doInstance.webSocketMessage(webWs as never, msg); } @@ -362,74 +340,60 @@ function sendCommand( function sendCliResponse( doInstance: UserConnectionDO, cliWs: MockWS, - opts: { id: string; result?: unknown; error?: unknown }, + opts: { id: string; result?: unknown; error?: unknown } ) { - const msg = JSON.stringify({ type: "response", ...opts }); + const msg = JSON.stringify({ type: 'response', ...opts }); return doInstance.webSocketMessage(cliWs as never, msg); } 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 - ) { - throw new Error( - `Result fixture does not serialize to ${targetBytes} bytes`, - ); + 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) { + throw new Error(`Result fixture does not serialize to ${targetBytes} bytes`); } return result; } function createUtf8OversizedResult(): { padding: string } { - const framingBytes = JSON.stringify({ padding: "" }).length; + const framingBytes = JSON.stringify({ padding: '' }).length; const result = { - padding: "é".repeat( - Math.floor((MAX_CATALOG_RESULT_BYTES - framingBytes) / 2) + 1, - ), + padding: 'é'.repeat(Math.floor((MAX_CATALOG_RESULT_BYTES - framingBytes) / 2) + 1), }; if ( JSON.stringify(result).length >= MAX_CATALOG_RESULT_BYTES || - new TextEncoder().encode(JSON.stringify(result)).byteLength <= - MAX_CATALOG_RESULT_BYTES + new TextEncoder().encode(JSON.stringify(result)).byteLength <= MAX_CATALOG_RESULT_BYTES ) { - throw new Error( - "UTF-8 catalog fixture does not cross the byte-only boundary", - ); + throw new Error('UTF-8 catalog fixture does not cross the byte-only boundary'); } return result; } /** Trigger CLI disconnect (awaits attention reset before broadcast). */ async function disconnectCli(doInstance: UserConnectionDO, cliWs: MockWS) { - await doInstance.webSocketClose(cliWs as never, 0, "", false); + await doInstance.webSocketClose(cliWs as never, 0, '', false); } /** Trigger web disconnect */ function disconnectWeb(doInstance: UserConnectionDO, webWs: MockWS) { - void doInstance.webSocketClose(webWs as never, 0, "", false); + void doInstance.webSocketClose(webWs as never, 0, '', false); } // =========================================================================== // Tests // =========================================================================== -describe("UserConnectionDO", () => { +describe('UserConnectionDO', () => { beforeEach(() => { vi.restoreAllMocks(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockReset(); - sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockResolvedValue( - undefined, - ); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockResolvedValue(undefined); sessionIngestMocks.claimSessionReadyPush.mockReset(); sessionIngestMocks.getSessionIngestDO.mockReset(); sessionIngestMocks.getSessionIngestDO.mockReturnValue({ - resetAttentionStatusOnCliDisconnect: - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + resetAttentionStatusOnCliDisconnect: sessionIngestMocks.resetAttentionStatusOnCliDisconnect, claimSessionReadyPush: sessionIngestMocks.claimSessionReadyPush, }); }); @@ -438,76 +402,76 @@ describe("UserConnectionDO", () => { vi.unstubAllGlobals(); }); - describe("notifySessionEvent", () => { - it("broadcasts semantic session events to web sockets only", async () => { + describe('notifySessionEvent', () => { + it('broadcasts semantic session events to web sockets only', async () => { const { doInstance, mockCtx } = setup(); const webWs = addWebSocket(mockCtx); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); const session = { - source: "v2" as const, - sessionId: "ses_12345678901234567890123456", - createdAt: "2026-01-01T00:00:00.000Z", - updatedAt: "2026-01-01T00:00:01.000Z", - title: "Test", - createdOnPlatform: "web", + source: 'v2' as const, + sessionId: 'ses_12345678901234567890123456', + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:01.000Z', + title: 'Test', + createdOnPlatform: 'web', organizationId: null, gitUrl: null, gitBranch: null, parentSessionId: null, - status: "idle" as const, + status: 'idle' as const, statusUpdatedAt: null, }; const result = await doInstance.notifySessionEvent({ - type: "session.created", - data: { source: "v2", session, changedAt: session.updatedAt }, + type: 'session.created', + data: { source: 'v2', session, changedAt: session.updatedAt }, }); expect(result).toEqual({ delivered: 1 }); expect(parseSent(webWs)).toEqual({ - type: "system", - event: "session.created", - data: { source: "v2", session, changedAt: session.updatedAt }, + type: 'system', + event: 'session.created', + data: { source: 'v2', session, changedAt: session.updatedAt }, }); expect(cliWs.send).not.toHaveBeenCalled(); }); - it("rejects invalid session event payloads without broadcasting", async () => { + it('rejects invalid session event payloads without broadcasting', async () => { const { doInstance, mockCtx } = setup(); const webWs = addWebSocket(mockCtx); await expect( doInstance.notifySessionEvent({ - type: "session.created", - data: { source: "v1" }, - } as never), + type: 'session.created', + data: { source: 'v1' }, + } as never) ).rejects.toThrow(); expect(webWs.send).not.toHaveBeenCalled(); }); }); - describe("hasActiveCliSession", () => { - it("tracks whether a connected CLI heartbeat currently owns the session", async () => { + describe('hasActiveCliSession', () => { + it('tracks whether a connected CLI heartbeat currently owns the session', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - expect(doInstance.hasActiveCliSession("ses_1")).toBe(false); + expect(doInstance.hasActiveCliSession('ses_1')).toBe(false); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_1')]); - expect(doInstance.hasActiveCliSession("ses_1")).toBe(true); + expect(doInstance.hasActiveCliSession('ses_1')).toBe(true); mockCtx.removeSocket(cliWs); await disconnectCli(doInstance, cliWs); - expect(doInstance.hasActiveCliSession("ses_1")).toBe(false); + 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', () => { const { doInstance, mockCtx } = setup(); - addCliSocket(mockCtx, "cli-1", [makeSession("ses_1")]); + addCliSocket(mockCtx, 'cli-1', [makeSession('ses_1')]); - expect(doInstance.hasActiveCliSession("ses_1")).toBe(true); + expect(doInstance.hasActiveCliSession('ses_1')).toBe(true); }); }); @@ -515,13 +479,13 @@ describe("UserConnectionDO", () => { // Heartbeat processing // ------------------------------------------------------------------------- - describe("heartbeat processing", () => { - it("updates session ownership and persists attachment", () => { + describe('heartbeat processing', () => { + it('updates session ownership and persists attachment', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + addWebSocket(mockCtx, 'web-1'); - const sessions = [makeSession("s1"), makeSession("s2")]; + const sessions = [makeSession('s1'), makeSession('s2')]; sendHeartbeat(doInstance, cliWs, sessions); // CLI attachment updated with sessions @@ -529,70 +493,68 @@ 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + addWebSocket(mockCtx, 'web-1'); // First heartbeat: owns s1 and s2 - sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); // Second heartbeat: only s1 - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // Verify via command routing: command to s2 should fail (no owner) - const webWs2 = addWebSocket(mockCtx, "web-2"); + const webWs2 = addWebSocket(mockCtx, 'web-2'); sendCommand(doInstance, webWs2, { - id: "cmd-1", - command: "send_message", - sessionId: "s2", + id: 'cmd-1', + command: 'send_message', + sessionId: 's2', }); const resp = parseSent(webWs2); expect(resp).toMatchObject({ - type: "response", - id: "cmd-1", - error: "Session owner not found", + type: 'response', + id: 'cmd-1', + error: 'Session owner not found', }); }); - it("fails an in-flight command when the session owner changes", () => { + it('fails an in-flight command when the session owner changes', () => { const { doInstance, mockCtx } = setup(); - const firstOwner = addCliSocket(mockCtx, "cli-1"); - const nextOwner = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + 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, firstOwner, [makeSession('s1')]); sendHeartbeat(doInstance, nextOwner, []); firstOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(firstOwner); webWs.send.mockClear(); - sendHeartbeat(doInstance, nextOwner, [makeSession("s1")]); + sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); // The owner-change heartbeat broadcasts sessions.heartbeat and also fires // the SESSION_OWNER_CHANGED error response for the in-flight command. The // test cares about the latter; find it by type+id. - expect( - allSent(webWs).find((m) => m.type === "response" && m.id === "cmd-1"), - ).toEqual({ - type: "response", - id: "cmd-1", + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); sendCliResponse(doInstance, firstOwner, { id: correlationId, - result: "late", + result: 'late', }); // The late sendCliResponse is filtered (the pending entry was already // removed by the owner-change path) so webWs sees exactly the two @@ -601,78 +563,78 @@ 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', () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1"); - const cli2 = addCliSocket(mockCtx, "cli-2"); - addWebSocket(mockCtx, "web-1"); + const cli1 = addCliSocket(mockCtx, 'cli-1'); + const cli2 = addCliSocket(mockCtx, 'cli-2'); + addWebSocket(mockCtx, 'web-1'); // cli1 owns s1 - sendHeartbeat(doInstance, cli1, [makeSession("s1")]); + sendHeartbeat(doInstance, cli1, [makeSession('s1')]); // web subscribes to s1 — subscribe sent to cli1 (the current owner) - const webWs = mockCtx.sockets.find((s) => s._tags.includes("web"))!; - sendSubscribe(doInstance, webWs, "s1"); + const webWs = mockCtx.sockets.find(s => s._tags.includes('web'))!; + sendSubscribe(doInstance, webWs, 's1'); cli1.send.mockClear(); cli2.send.mockClear(); // cli2 now reports s1 — becomes new owner - sendHeartbeat(doInstance, cli2, [makeSession("s1")]); + sendHeartbeat(doInstance, cli2, [makeSession('s1')]); // cli2 should have received the replayed subscribe for s1 const cli2Msgs = allSent(cli2); - expect(cli2Msgs).toContainEqual({ type: "subscribe", sessionId: "s1" }); + 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', () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1"); - const cli2 = addCliSocket(mockCtx, "cli-2"); - const subWeb = addWebSocket(mockCtx, "web-sub"); - const otherWeb = addWebSocket(mockCtx, "web-other"); + const cli1 = addCliSocket(mockCtx, 'cli-1'); + const cli2 = addCliSocket(mockCtx, 'cli-2'); + const subWeb = addWebSocket(mockCtx, 'web-sub'); + const otherWeb = addWebSocket(mockCtx, 'web-other'); // cli1 owns s1, cli2 owns s2 - sendHeartbeat(doInstance, cli1, [makeSession("s1")]); - sendHeartbeat(doInstance, cli2, [makeSession("s2")]); + sendHeartbeat(doInstance, cli1, [makeSession('s1')]); + sendHeartbeat(doInstance, cli2, [makeSession('s2')]); // subWeb subscribes to s1, otherWeb subscribes to s2 (subscriptions are // irrelevant to broadcast delivery — both should still receive cli1's heartbeat) - sendSubscribe(doInstance, subWeb, "s1"); - sendSubscribe(doInstance, otherWeb, "s2"); + sendSubscribe(doInstance, subWeb, 's1'); + sendSubscribe(doInstance, otherWeb, 's2'); subWeb.send.mockClear(); otherWeb.send.mockClear(); // cli1 sends heartbeat — both viewers must receive it - sendHeartbeat(doInstance, cli1, [makeSession("s1")]); + sendHeartbeat(doInstance, cli1, [makeSession('s1')]); expect(subWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(subWeb)).toMatchObject({ - type: "system", - event: "sessions.heartbeat", - data: { connectionId: "cli-1" }, + type: 'system', + event: 'sessions.heartbeat', + data: { connectionId: 'cli-1' }, }); expect(otherWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(otherWeb)).toMatchObject({ - type: "system", - event: "sessions.heartbeat", - data: { connectionId: "cli-1" }, + type: 'system', + event: 'sessions.heartbeat', + data: { connectionId: 'cli-1' }, }); }); - it("delivers one heartbeat per web socket (delivery count equals ws count)", () => { + it('delivers one heartbeat per web socket (delivery count equals ws count)', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const web1 = addWebSocket(mockCtx, "web-1"); - const web2 = addWebSocket(mockCtx, "web-2"); - const web3 = addWebSocket(mockCtx, "web-3"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const web1 = addWebSocket(mockCtx, 'web-1'); + const web2 = addWebSocket(mockCtx, 'web-2'); + const web3 = addWebSocket(mockCtx, 'web-3'); // No subscriptions — the broadcast must still hit every web socket. web1.send.mockClear(); web2.send.mockClear(); web3.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // One heartbeat frame per active web socket. expect(web1.send).toHaveBeenCalledTimes(1); @@ -685,60 +647,60 @@ describe("UserConnectionDO", () => { data: { connectionId: string; sessions: unknown[] }; }; expect(sent).toMatchObject({ - type: "system", - event: "sessions.heartbeat", - data: { connectionId: "cli-1", sessions: [{ id: "s1" }] }, + type: 'system', + event: 'sessions.heartbeat', + data: { connectionId: 'cli-1', sessions: [{ id: 's1' }] }, }); } }); - it("forwards the CLI-reported protocolVersion to every web socket", () => { + it('forwards the CLI-reported protocolVersion to every web socket', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + 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"); + 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", - event: "sessions.heartbeat", - data: { connectionId: "cli-1", protocolVersion: "1" }, + type: 'system', + event: 'sessions.heartbeat', + data: { connectionId: 'cli-1', protocolVersion: '1' }, }); }); - it("omits protocolVersion for a legacy CLI that never reports one", () => { + it('omits protocolVersion for a legacy CLI that never reports one', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); - sendSubscribe(doInstance, webWs, "s1"); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); const sent = parseSent(webWs) as { data: Record }; - expect(sent.data).not.toHaveProperty("protocolVersion"); + 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)', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); // subWeb subscribed to s1, otherWeb is unrelated — both must learn s1 is gone. - const subWeb = addWebSocket(mockCtx, "web-sub"); - const otherWeb = addWebSocket(mockCtx, "web-other"); + const subWeb = addWebSocket(mockCtx, 'web-sub'); + const otherWeb = addWebSocket(mockCtx, 'web-other'); // cli1 owns s1 - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); - sendSubscribe(doInstance, subWeb, "s1"); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendSubscribe(doInstance, subWeb, 's1'); subWeb.send.mockClear(); otherWeb.send.mockClear(); @@ -747,50 +709,50 @@ describe("UserConnectionDO", () => { expect(subWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(subWeb)).toMatchObject({ - type: "system", - event: "sessions.heartbeat", - data: { connectionId: "cli-1", sessions: [] }, + type: 'system', + event: 'sessions.heartbeat', + data: { connectionId: 'cli-1', sessions: [] }, }); expect(otherWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(otherWeb)).toMatchObject({ - type: "system", - event: "sessions.heartbeat", - data: { connectionId: "cli-1", sessions: [] }, + type: 'system', + event: 'sessions.heartbeat', + data: { connectionId: 'cli-1', sessions: [] }, }); }); - it("delivers heartbeat to web sockets that are not subscribed to anything", () => { + it('delivers heartbeat to web sockets that are not subscribed to anything', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); // webWs has no subscriptions - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toMatchObject({ - type: "system", - event: "sessions.heartbeat", - data: { connectionId: "cli-1", sessions: [{ id: "s1" }] }, + type: 'system', + event: 'sessions.heartbeat', + data: { connectionId: 'cli-1', sessions: [{ id: 's1' }] }, }); }); - it("schedules stale alarm on heartbeat", () => { + it('schedules stale alarm on heartbeat', () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); expect(ctx.storage.setAlarm).toHaveBeenCalled(); }); - it("sends heartbeat_ack to CLI socket", () => { + it('sends heartbeat_ack to CLI socket', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); const msgs = allSent(cliWs); - expect(msgs).toContainEqual({ type: "heartbeat_ack" }); + expect(msgs).toContainEqual({ type: 'heartbeat_ack' }); }); // ----------------------------------------------------------------------- @@ -801,11 +763,11 @@ 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { capabilities: { attachments: true }, }); @@ -816,95 +778,92 @@ describe("UserConnectionDO", () => { } }); - it("projects the latest connection capabilities onto every sessions.heartbeat row", () => { + it('projects the latest connection capabilities onto every sessions.heartbeat row', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); // Establish ownership and subscription first so the heartbeat broadcast // is targeted at this web socket. - sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")]); - sendSubscribe(doInstance, webWs, "s1"); + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); + sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { capabilities: { attachments: true }, }); expect(parseSent(webWs)).toMatchObject({ data: { sessions: [ - { id: "s1", capabilities: { attachments: true } }, - { id: "s2", capabilities: { attachments: true } }, + { id: 's1', capabilities: { attachments: true } }, + { id: 's2', capabilities: { attachments: true } }, ], }, }); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); expect(parseSent(webWs)).toMatchObject({ - data: { sessions: [{ id: "s1" }, { id: "s2" }] }, + data: { sessions: [{ id: 's1' }, { id: 's2' }] }, }); - const legacyRows = ( - parseSent(webWs) as { data: { sessions: Record[] } } - ).data.sessions; - expect( - legacyRows.every((row) => !Object.hasOwn(row, "capabilities")), - ).toBe(true); + const legacyRows = (parseSent(webWs) as { data: { sessions: Record[] } }) + .data.sessions; + expect(legacyRows.every(row => !Object.hasOwn(row, 'capabilities'))).toBe(true); webWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { capabilities: { attachments: false }, }); expect(parseSent(webWs)).toMatchObject({ data: { sessions: [ - { id: "s1", capabilities: { attachments: false } }, - { id: "s2", capabilities: { attachments: false } }, + { id: 's1', capabilities: { attachments: false } }, + { id: 's2', capabilities: { attachments: false } }, ], }, }); }); - 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)', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); // First heartbeat with capabilities - sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: true }, }); // Second heartbeat omits capabilities (CLI rollback / legacy) - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); const rows = doInstance.getActiveSessions(); expect(rows).toHaveLength(1); - expect(rows[0]).not.toHaveProperty("capabilities"); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: true }, }); - sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); // Latest heartbeat omits capabilities. - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); const sent = parseSent(webWs) as { data: Record }; - expect(sent.data).not.toHaveProperty("capabilities"); + 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)', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: true }, }); expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ @@ -912,7 +871,7 @@ describe("UserConnectionDO", () => { }); // CLI advertises attachments=false (e.g. feature gated, profile change) - sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: false }, }); @@ -921,17 +880,17 @@ 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); // Legacy heartbeat — no capabilities field - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); const legacy = doInstance.getActiveSessions(); - expect(legacy[0]).not.toHaveProperty("capabilities"); + expect(legacy[0]).not.toHaveProperty('capabilities'); // Upgraded CLI starts advertising attachments=true - sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: true }, }); @@ -939,18 +898,18 @@ 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: false }, }); expect(doInstance.getActiveSessions()[0].capabilities).toEqual({ attachments: false, }); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: true }, }); @@ -959,11 +918,11 @@ describe("UserConnectionDO", () => { }); }); - 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")], { + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { capabilities: { attachments: false }, }); @@ -973,14 +932,14 @@ 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', () => { const { doInstance, mockCtx } = setup(); // Pre-existing attachment with capabilities — simulates a socket that // was accepted before the DO was evicted. - const cliWs = createMockWs(["cli"], { - role: "cli", - connectionId: "cli-hiber", - sessions: [makeSession("s1")], + const cliWs = createMockWs(['cli'], { + role: 'cli', + connectionId: 'cli-hiber', + sessions: [makeSession('s1')], capabilities: { attachments: true }, }); mockCtx.addSocket(cliWs); @@ -995,32 +954,32 @@ describe("UserConnectionDO", () => { // Stale connection eviction // ------------------------------------------------------------------------- - describe("stale connection eviction", () => { - it("closes stale connection after timeout", async () => { + describe('stale connection eviction', () => { + it('closes stale connection after timeout', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); // Send heartbeat to register the connection and set lastHeartbeatAt - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // Fast-forward time so the connection appears stale - vi.spyOn(Date, "now") + vi.spyOn(Date, 'now') .mockReturnValueOnce(Date.now() + 31_000) // for ensureState check .mockReturnValue(Date.now() + 31_000); // for alarm's Date.now() await doInstance.alarm(); - expect(cliWs.close).toHaveBeenCalledWith(4408, "heartbeat timeout"); + expect(cliWs.close).toHaveBeenCalledWith(4408, 'heartbeat timeout'); }); - it("reschedules alarm if other live connections remain", async () => { + it('reschedules alarm if other live connections remain', async () => { const { doInstance, mockCtx, ctx } = setup(); - const staleCli = addCliSocket(mockCtx, "stale-1"); - const freshCli = addCliSocket(mockCtx, "fresh-1"); + const staleCli = addCliSocket(mockCtx, 'stale-1'); + const freshCli = addCliSocket(mockCtx, 'fresh-1'); // Both send heartbeats - sendHeartbeat(doInstance, staleCli, [makeSession("s1")]); - sendHeartbeat(doInstance, freshCli, [makeSession("s2")]); + sendHeartbeat(doInstance, staleCli, [makeSession('s1')]); + sendHeartbeat(doInstance, freshCli, [makeSession('s2')]); // Reset setAlarm call count ctx.storage.setAlarm.mockClear(); @@ -1028,31 +987,31 @@ describe("UserConnectionDO", () => { // Make stale-1 appear stale but fresh-1 stays fresh const now = Date.now(); const staleTime = now + 31_000; - vi.spyOn(Date, "now").mockReturnValue(staleTime); + vi.spyOn(Date, 'now').mockReturnValue(staleTime); // Manually set lastHeartbeatAt for fresh-1 to "just now" (staleTime) // by sending another heartbeat from fresh-1 - sendHeartbeat(doInstance, freshCli, [makeSession("s2")]); + sendHeartbeat(doInstance, freshCli, [makeSession('s2')]); ctx.storage.setAlarm.mockClear(); await doInstance.alarm(); // Stale one closed - expect(staleCli.close).toHaveBeenCalledWith(4408, "heartbeat timeout"); + expect(staleCli.close).toHaveBeenCalledWith(4408, 'heartbeat timeout'); // Fresh one alive expect(freshCli.close).not.toHaveBeenCalled(); // Alarm rescheduled because fresh-1 remains expect(ctx.storage.setAlarm).toHaveBeenCalled(); }); - it("does not evict connection with recent heartbeat", async () => { + it('does not evict connection with recent heartbeat', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // Time is within timeout window - vi.spyOn(Date, "now").mockReturnValue(Date.now() + 10_000); + vi.spyOn(Date, 'now').mockReturnValue(Date.now() + 10_000); await doInstance.alarm(); @@ -1064,122 +1023,122 @@ describe("UserConnectionDO", () => { // Subscribe / Unsubscribe // ------------------------------------------------------------------------- - describe("subscribe/unsubscribe", () => { - it("sends subscribe to owning CLI when web subscribes", () => { + describe('subscribe/unsubscribe', () => { + it('sends subscribe to owning CLI when web subscribes', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); // CLI owns s1 - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, 's1'); // CLI should receive subscribe expect(cliWs.send).toHaveBeenCalledTimes(1); - expect(parseSent(cliWs)).toEqual({ type: "subscribe", sessionId: "s1" }); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1", "busy", "Fix bug")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Fix bug')]); webWs.send.mockClear(); - sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, 's1'); expect(parseSent(webWs)).toEqual({ - type: "system", - event: "sessions.list", + type: 'system', + event: 'sessions.list', data: { sessions: [ { - id: "s1", - status: "busy", - title: "Fix bug", - connectionId: "cli-1", + 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', () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1"); - const cli2 = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cli1 = addCliSocket(mockCtx, 'cli-1'); + const cli2 = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); // No heartbeat sent, so no owner for 's1' // Trigger ensureState via a harmless message first - sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, 's1'); // Both CLIs should receive subscribe expect(cli1.send).toHaveBeenCalled(); expect(cli2.send).toHaveBeenCalled(); - expect(parseSent(cli1)).toEqual({ type: "subscribe", sessionId: "s1" }); - expect(parseSent(cli2)).toEqual({ type: "subscribe", sessionId: "s1" }); + expect(parseSent(cli1)).toEqual({ type: 'subscribe', sessionId: 's1' }); + expect(parseSent(cli2)).toEqual({ type: 'subscribe', sessionId: 's1' }); }); - it("duplicate subscribe is idempotent for attachment", () => { + it('duplicate subscribe is idempotent for attachment', () => { const { doInstance, mockCtx } = setup(); - addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendSubscribe(doInstance, webWs, "s1"); - sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, 's1'); const att = webWs.deserializeAttachment() as { subscribedSessions: string[]; }; - expect(att.subscribedSessions).toEqual(["s1"]); + expect(att.subscribedSessions).toEqual(['s1']); }); - it("unsubscribe sends to CLI when last subscriber leaves", () => { + it('unsubscribe sends to CLI when last subscriber leaves', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, 's1'); cliWs.send.mockClear(); - sendUnsubscribe(doInstance, webWs, "s1"); + sendUnsubscribe(doInstance, webWs, 's1'); expect(cliWs.send).toHaveBeenCalledTimes(1); expect(parseSent(cliWs)).toEqual({ - type: "unsubscribe", - sessionId: "s1", + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const web1 = addWebSocket(mockCtx, "web-1"); - const web2 = addWebSocket(mockCtx, "web-2"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const web1 = addWebSocket(mockCtx, 'web-1'); + const web2 = addWebSocket(mockCtx, 'web-2'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendSubscribe(doInstance, web1, "s1"); - sendSubscribe(doInstance, web2, "s1"); + sendSubscribe(doInstance, web1, 's1'); + sendSubscribe(doInstance, web2, 's1'); cliWs.send.mockClear(); // Unsubscribe first — CLI should NOT get unsubscribe - sendUnsubscribe(doInstance, web1, "s1"); + sendUnsubscribe(doInstance, web1, 's1'); expect(cliWs.send).not.toHaveBeenCalled(); // Unsubscribe second — CLI SHOULD get unsubscribe - sendUnsubscribe(doInstance, web2, "s1"); + sendUnsubscribe(doInstance, web2, 's1'); expect(cliWs.send).toHaveBeenCalledTimes(1); expect(parseSent(cliWs)).toEqual({ - type: "unsubscribe", - sessionId: "s1", + type: 'unsubscribe', + sessionId: 's1', }); }); }); @@ -1188,119 +1147,117 @@ describe("UserConnectionDO", () => { // Viewer liveness // ------------------------------------------------------------------------- - describe("viewer liveness", () => { - it("replies to a viewer ping with the matching nonce only", () => { + describe('viewer liveness', () => { + it('replies to a viewer ping with the matching nonce only', () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "viewer-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'viewer-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); ctx.storage.setAlarm.mockClear(); cliWs.send.mockClear(); webWs.send.mockClear(); - sendPing(doInstance, webWs, "nonce-1"); + sendPing(doInstance, webWs, 'nonce-1'); expect(webWs.send).toHaveBeenCalledTimes(1); - expect(parseSent(webWs)).toEqual({ type: "pong", nonce: "nonce-1" }); + expect(parseSent(webWs)).toEqual({ type: 'pong', nonce: 'nonce-1' }); expect(doInstance.getActiveSessions()).toEqual([ - { id: "s1", status: "busy", title: "Test", connectionId: "cli-1" }, + { id: 's1', status: 'busy', title: 'Test', connectionId: 'cli-1' }, ]); expect(cliWs.send).not.toHaveBeenCalled(); expect(ctx.storage.setAlarm).not.toHaveBeenCalled(); expect(webWs.deserializeAttachment()).toEqual({ - role: "web", - connectionId: "viewer-1", + role: 'web', + connectionId: 'viewer-1', subscribedSessions: [], }); }); }); - describe("viewer connection identity", () => { - it("replaces an older web viewer with the same connectionId and broadcasts only to its replacement", async () => { + describe('viewer connection identity', () => { + it('replaces an older web viewer with the same connectionId and broadcasts only to its replacement', async () => { const { doInstance, mockCtx } = setup(); - const oldWeb = connectWebSocket(doInstance, "viewer-1"); + const oldWeb = connectWebSocket(doInstance, 'viewer-1'); oldWeb.send.mockClear(); - const newWeb = connectWebSocket(doInstance, "viewer-1"); + const newWeb = connectWebSocket(doInstance, 'viewer-1'); newWeb.send.mockClear(); - expect(oldWeb.close).toHaveBeenCalledWith(1000, "replaced by reconnect"); + expect(oldWeb.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); await doInstance.notifySessionEvent({ - type: "session.deleted", + type: 'session.deleted', data: { - source: "v2", - sessionId: "s1", + source: 'v2', + sessionId: 's1', parentSessionId: null, organizationId: null, gitUrl: null, gitBranch: null, - createdOnPlatform: "web", - deletedAt: "2026-01-01T00:00:02.000Z", + createdOnPlatform: 'web', + deletedAt: '2026-01-01T00:00:02.000Z', }, }); expect(oldWeb.send).not.toHaveBeenCalled(); expect(newWeb.send).toHaveBeenCalledTimes(1); - expect( - mockCtx.sockets.filter((socket) => socket._tags.includes("web")), - ).toHaveLength(2); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - const oldWeb = connectWebSocket(doInstance, "viewer-1"); - sendSubscribe(doInstance, oldWeb, "s1"); + const oldWeb = connectWebSocket(doInstance, 'viewer-1'); + sendSubscribe(doInstance, oldWeb, 's1'); cliWs.send.mockClear(); - const newWeb = connectWebSocket(doInstance, "viewer-1"); + const newWeb = connectWebSocket(doInstance, 'viewer-1'); expect(cliWs.send).toHaveBeenCalledTimes(1); expect(parseSent(cliWs)).toEqual({ - type: "unsubscribe", - sessionId: "s1", + type: 'unsubscribe', + sessionId: 's1', }); expect(newWeb.deserializeAttachment()).toEqual({ - role: "web", - connectionId: "viewer-1", + role: 'web', + connectionId: 'viewer-1', subscribedSessions: [], }); }); - it("ignores messages from a viewer that has been replaced", () => { + it('ignores messages from a viewer that has been replaced', () => { const { doInstance } = setup(); - const oldWeb = connectWebSocket(doInstance, "viewer-1"); - connectWebSocket(doInstance, "viewer-1"); + const oldWeb = connectWebSocket(doInstance, 'viewer-1'); + connectWebSocket(doInstance, 'viewer-1'); oldWeb.send.mockClear(); - sendPing(doInstance, oldWeb, "stale-ping"); + sendPing(doInstance, oldWeb, 'stale-ping'); expect(oldWeb.send).not.toHaveBeenCalled(); }); - it("keeps distinct viewer identities connected for independent broadcasts", async () => { + it('keeps distinct viewer identities connected for independent broadcasts', async () => { const { doInstance } = setup(); - const firstWeb = connectWebSocket(doInstance, "viewer-1"); - const secondWeb = connectWebSocket(doInstance, "viewer-2"); + const firstWeb = connectWebSocket(doInstance, 'viewer-1'); + const secondWeb = connectWebSocket(doInstance, 'viewer-2'); firstWeb.send.mockClear(); secondWeb.send.mockClear(); await doInstance.notifySessionEvent({ - type: "session.deleted", + type: 'session.deleted', data: { - source: "v2", - sessionId: "s1", + source: 'v2', + sessionId: 's1', parentSessionId: null, organizationId: null, gitUrl: null, gitBranch: null, - createdOnPlatform: "web", - deletedAt: "2026-01-01T00:00:02.000Z", + createdOnPlatform: 'web', + deletedAt: '2026-01-01T00:00:02.000Z', }, }); @@ -1309,11 +1266,11 @@ 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "shared-id"); + const cliWs = addCliSocket(mockCtx, 'shared-id'); - connectWebSocket(doInstance, "shared-id"); + connectWebSocket(doInstance, 'shared-id'); expect(cliWs.close).not.toHaveBeenCalled(); }); @@ -1323,13 +1280,13 @@ describe("UserConnectionDO", () => { // CLI disconnect // ------------------------------------------------------------------------- - describe("CLI disconnect", () => { - it("cleans up session ownership and broadcasts cli.disconnected", async () => { + describe('CLI disconnect', () => { + it('cleans up session ownership and broadcasts cli.disconnected', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); webWs.send.mockClear(); // Remove from sockets before disconnect (simulates runtime closing) @@ -1340,40 +1297,39 @@ describe("UserConnectionDO", () => { expect(webWs.send).toHaveBeenCalled(); const msgs = allSent(webWs); const disconnectMsg = msgs.find( - (m: Record) => - m.type === "system" && m.event === "cli.disconnected", + (m: Record) => m.type === 'system' && m.event === 'cli.disconnected' ); expect(disconnectMsg).toEqual({ - type: "system", - event: "cli.disconnected", - data: { connectionId: "cli-1" }, + type: 'system', + event: 'cli.disconnected', + data: { connectionId: 'cli-1' }, }); // Session no longer routable - const web2 = addWebSocket(mockCtx, "web-2"); + const web2 = addWebSocket(mockCtx, 'web-2'); sendCommand(doInstance, web2, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); expect(parseSent(web2)).toMatchObject({ - type: "response", - error: "Session owner not found", + type: 'response', + error: 'Session owner not found', }); }); - it("sends error responses for pending commands on disconnect", async () => { + it('sends error responses for pending commands on disconnect', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // Send command from web sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); webWs.send.mockClear(); @@ -1384,27 +1340,26 @@ describe("UserConnectionDO", () => { // Web receives error response with original id const msgs = allSent(webWs); const errorResp = msgs.find( - (m: Record) => - m.type === "response" && m.id === "cmd-1", + (m: Record) => m.type === 'response' && m.id === 'cmd-1' ); expect(errorResp).toMatchObject({ - type: "response", - id: "cmd-1", - error: "CLI disconnected", + type: 'response', + id: 'cmd-1', + error: 'CLI disconnected', }); }); - it("reports owner change when an owner-fenced command target disconnects", async () => { + it('reports owner change when an owner-fenced command target disconnects', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); webWs.send.mockClear(); @@ -1412,60 +1367,57 @@ describe("UserConnectionDO", () => { await disconnectCli(doInstance, cliWs); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); }); - it("fails pending commands as soon as their target socket is replaced", () => { + it('fails pending commands as soon as their target socket is replaced', () => { const { doInstance, mockCtx } = setup(); - const firstCli = connectCliSocket(doInstance, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const firstCli = connectCliSocket(doInstance, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, firstCli, [makeSession("s1")]); + sendHeartbeat(doInstance, firstCli, [makeSession('s1')]); firstCli.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); webWs.send.mockClear(); - connectCliSocket(doInstance, "cli-1"); + connectCliSocket(doInstance, 'cli-1'); - expect(firstCli.close).toHaveBeenCalledWith( - 1000, - "replaced by reconnect", - ); + expect(firstCli.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); }); - it("sends error for connection-routed pending commands on CLI disconnect", async () => { + it('sends error for connection-routed pending commands on CLI disconnect', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, []); // Send command routed by connectionId (no sessionId) sendCommand(doInstance, webWs, { - id: "cmd-conn", - command: "send_message", - connectionId: "cli-1", + id: 'cmd-conn', + command: 'send_message', + connectionId: 'cli-1', }); webWs.send.mockClear(); @@ -1475,27 +1427,26 @@ describe("UserConnectionDO", () => { const msgs = allSent(webWs); const errorResp = msgs.find( - (m: Record) => - m.type === "response" && m.id === "cmd-conn", + (m: Record) => m.type === 'response' && m.id === 'cmd-conn' ); expect(errorResp).toMatchObject({ - type: "response", - id: "cmd-conn", - error: "CLI disconnected", + type: 'response', + id: 'cmd-conn', + error: 'CLI disconnected', }); }); - it("sends error for fallback-routed pending commands on CLI disconnect", async () => { + it('sends error for fallback-routed pending commands on CLI disconnect', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, []); // Send command with no sessionId or connectionId (fallback routing) sendCommand(doInstance, webWs, { - id: "cmd-fallback", - command: "send_message", + id: 'cmd-fallback', + command: 'send_message', }); webWs.send.mockClear(); @@ -1504,26 +1455,25 @@ describe("UserConnectionDO", () => { const msgs = allSent(webWs); const errorResp = msgs.find( - (m: Record) => - m.type === "response" && m.id === "cmd-fallback", + (m: Record) => m.type === 'response' && m.id === 'cmd-fallback' ); expect(errorResp).toMatchObject({ - type: "response", - id: "cmd-fallback", - error: "CLI disconnected", + type: 'response', + id: 'cmd-fallback', + error: 'CLI disconnected', }); }); - it("reconnecting CLI — old socket close does not destroy state", async () => { + it('reconnecting CLI — old socket close does not destroy state', async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cli1 = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cli1, [makeSession("s1")]); + sendHeartbeat(doInstance, cli1, [makeSession('s1')]); // CLI2 connects with same connectionId (simulates reconnect) - const cli2 = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cli2, [makeSession("s1")]); + const cli2 = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cli2, [makeSession('s1')]); // CLI1's close event fires (stale socket), but cli2 still holds the connectionId // DON'T remove cli2 from sockets — cli2 is the replacement @@ -1534,33 +1484,33 @@ describe("UserConnectionDO", () => { // State should NOT be cleaned up — cli2 is live // Verify by routing a command to s1 — should reach cli2 cli2.send.mockClear(); - sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, 's1'); expect(cli2.send).toHaveBeenCalled(); - expect(parseSent(cli2)).toEqual({ type: "subscribe", sessionId: "s1" }); + expect(parseSent(cli2)).toEqual({ type: 'subscribe', sessionId: 's1' }); }); - it("reconnecting CLI — commands sent to replacement socket are not spuriously failed", async () => { + it('reconnecting CLI — commands sent to replacement socket are not spuriously failed', async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cli1 = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cli1, [makeSession("s1")]); + sendHeartbeat(doInstance, cli1, [makeSession('s1')]); // cli2 connects with the same connectionId (reconnect). // In production, closeStaleSocket removes cli1 before cli2 is accepted. // Simulate that by removing cli1 from the socket list first. mockCtx.removeSocket(cli1); - const cli2 = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cli2, [makeSession("s1")]); + const cli2 = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cli2, [makeSession('s1')]); cli2.send.mockClear(); 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", + id: 'cmd-new', + command: 'send_message', + sessionId: 's1', }); expect(cli2.send).toHaveBeenCalled(); const correlationId = getCorrelationId(cli2); @@ -1572,40 +1522,40 @@ describe("UserConnectionDO", () => { // Web should NOT have received an error for cmd-new — it was sent to cli2, not cli1 const errorMsgs = allSent(webWs).filter( - (m) => m.type === "response" && m.id === "cmd-new" && m.error, + m => m.type === 'response' && m.id === 'cmd-new' && m.error ); expect(errorMsgs).toHaveLength(0); // cli2 responds successfully webWs.send.mockClear(); - sendCliResponse(doInstance, cli2, { id: correlationId, result: "ok" }); + sendCliResponse(doInstance, cli2, { id: correlationId, result: 'ok' }); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-new", - result: "ok", + type: 'response', + id: 'cmd-new', + result: 'ok', }); }); - it("reconnecting CLI — pending commands from old socket get error responses", async () => { + it('reconnecting CLI — pending commands from old socket get error responses', async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cli1 = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cli1, [makeSession("s1")]); + 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", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); webWs.send.mockClear(); // CLI2 connects with the same connectionId (reconnect) - const cli2 = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cli2, [makeSession("s1")]); + const cli2 = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cli2, [makeSession('s1')]); // cli1's close event fires — cmd-1 was sent on cli1's wire, cli2 never saw it mockCtx.removeSocket(cli1); @@ -1614,88 +1564,75 @@ describe("UserConnectionDO", () => { // Web should receive an error for the stranded command const msgs = allSent(webWs); const errorResp = msgs.find( - (m: Record) => - m.type === "response" && m.id === "cmd-1", + (m: Record) => m.type === 'response' && m.id === 'cmd-1' ); expect(errorResp).toMatchObject({ - type: "response", - id: "cmd-1", - error: "CLI disconnected", + type: 'response', + id: 'cmd-1', + error: 'CLI disconnected', }); }); - it("resets attention status for owned sessions before broadcasting cli.disconnected", async () => { + it('resets attention status for owned sessions before broadcasting cli.disconnected', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [ - makeSession("s-question", "question"), - makeSession("s-busy", "busy"), + makeSession('s-question', 'question'), + makeSession('s-busy', 'busy'), ]); webWs.send.mockClear(); const callOrder: string[] = []; - sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockImplementation( - async () => { - callOrder.push("reset"); - // Disconnect must not have been broadcast yet (ordering guarantee). - expect( - allSent(webWs).some( - (m) => m.type === "system" && m.event === "cli.disconnected", - ), - ).toBe(false); - }, - ); + sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockImplementation(async () => { + callOrder.push('reset'); + // Disconnect must not have been broadcast yet (ordering guarantee). + expect( + allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected') + ).toBe(false); + }); // Leave the socket in getWebSockets() — matches workerd during webSocketClose. await disconnectCli(doInstance, cliWs); - callOrder.push("disconnect"); + callOrder.push('disconnect'); - expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith( - expect.anything(), - { - kiloUserId: "usr_1", - sessionId: "s-question", - }, + expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith(expect.anything(), { + kiloUserId: 'usr_1', + sessionId: 's-question', + }); + expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith(expect.anything(), { + kiloUserId: 'usr_1', + sessionId: 's-busy', + }); + // Both owned sessions are delegated; attention-only filtering is on the metadata side. + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledTimes(2); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( + 'usr_1', + 's-question' ); - expect(sessionIngestMocks.getSessionIngestDO).toHaveBeenCalledWith( - expect.anything(), - { - kiloUserId: "usr_1", - sessionId: "s-busy", - }, + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( + 'usr_1', + 's-busy' + ); + expect(callOrder.filter(step => step === 'reset')).toHaveLength(2); + expect(callOrder.at(-1)).toBe('disconnect'); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true ); - // Both owned sessions are delegated; attention-only filtering is on the metadata side. - expect( - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, - ).toHaveBeenCalledTimes(2); - expect( - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, - ).toHaveBeenCalledWith("usr_1", "s-question"); - expect( - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, - ).toHaveBeenCalledWith("usr_1", "s-busy"); - expect(callOrder.filter((step) => step === "reset")).toHaveLength(2); - expect(callOrder.at(-1)).toBe("disconnect"); - expect( - allSent(webWs).some( - (m) => m.type === "system" && m.event === "cli.disconnected", - ), - ).toBe(true); }); - it("resets attention when the closing socket is still listed in getWebSockets (workerd)", async () => { + it('resets attention when the closing socket is still listed in getWebSockets (workerd)', async () => { // Production wrangler/workerd keeps the closing WebSocket in getWebSockets() // during webSocketClose. Matching connectionId without excluding self would // treat every disconnect as a stale reconnect and skip the attention reset. // Prior unit tests always called removeSocket first, so they never caught this. const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s-question", "question")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s-question', 'question')]); webWs.send.mockClear(); sessionIngestMocks.getSessionIngestDO.mockClear(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); @@ -1704,57 +1641,48 @@ describe("UserConnectionDO", () => { expect(mockCtx.sockets).toContain(cliWs); await disconnectCli(doInstance, cliWs); - expect( - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, - ).toHaveBeenCalledTimes(1); - expect( - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, - ).toHaveBeenCalledWith("usr_1", "s-question"); - expect( - allSent(webWs).some( - (m) => m.type === "system" && m.event === "cli.disconnected", - ), - ).toBe(true); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledTimes(1); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).toHaveBeenCalledWith( + 'usr_1', + 's-question' + ); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true + ); }); - it("does not reset attention when kiloUserId is missing on the attachment", async () => { + it('does not reset attention when kiloUserId is missing on the attachment', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); - const warn = vi - .spyOn(console, "warn") - .mockImplementation(() => undefined); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); - sendHeartbeat(doInstance, cliWs, [makeSession("s1", "question")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'question')]); webWs.send.mockClear(); mockCtx.removeSocket(cliWs); await disconnectCli(doInstance, cliWs); expect(sessionIngestMocks.getSessionIngestDO).not.toHaveBeenCalled(); - expect( - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, - ).not.toHaveBeenCalled(); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); expect(warn).toHaveBeenCalledWith( - "Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment", - { ownedSessionCount: 1 }, + 'Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment', + { ownedSessionCount: 1 } + ); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true ); - expect( - allSent(webWs).some( - (m) => m.type === "system" && m.event === "cli.disconnected", - ), - ).toBe(true); }); - it("does not reset attention for sessions owned by another live connection", async () => { + it('does not reset attention for sessions owned by another live connection', async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); - const cli2 = addCliSocket(mockCtx, "cli-2", [], undefined, "usr_1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cli1 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const cli2 = addCliSocket(mockCtx, 'cli-2', [], undefined, 'usr_1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cli1, [makeSession("s1", "question")]); + sendHeartbeat(doInstance, cli1, [makeSession('s1', 'question')]); // cli2 takes ownership of s1 - sendHeartbeat(doInstance, cli2, [makeSession("s1", "question")]); + sendHeartbeat(doInstance, cli2, [makeSession('s1', 'question')]); webWs.send.mockClear(); sessionIngestMocks.getSessionIngestDO.mockClear(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); @@ -1763,23 +1691,19 @@ describe("UserConnectionDO", () => { await disconnectCli(doInstance, cli1); // cli1 no longer owns s1, so no reset for that session - expect( - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, - ).not.toHaveBeenCalled(); - expect( - allSent(webWs).some( - (m) => m.type === "system" && m.event === "cli.disconnected", - ), - ).toBe(true); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true + ); }); - it("stale reconnect close does not reset attention status", async () => { + it('stale reconnect close does not reset attention status', async () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); - const cli2 = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); + const cli1 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const cli2 = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); - sendHeartbeat(doInstance, cli1, [makeSession("s1", "question")]); - sendHeartbeat(doInstance, cli2, [makeSession("s1", "question")]); + sendHeartbeat(doInstance, cli1, [makeSession('s1', 'question')]); + sendHeartbeat(doInstance, cli2, [makeSession('s1', 'question')]); sessionIngestMocks.getSessionIngestDO.mockClear(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockClear(); @@ -1787,36 +1711,30 @@ describe("UserConnectionDO", () => { await disconnectCli(doInstance, cli1); expect(sessionIngestMocks.getSessionIngestDO).not.toHaveBeenCalled(); - expect( - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, - ).not.toHaveBeenCalled(); + expect(sessionIngestMocks.resetAttentionStatusOnCliDisconnect).not.toHaveBeenCalled(); }); - it("still broadcasts cli.disconnected when attention reset fails", async () => { + it('still broadcasts cli.disconnected when attention reset fails', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1", [], undefined, "usr_1"); - const webWs = addWebSocket(mockCtx, "web-1"); - const error = vi - .spyOn(console, "error") - .mockImplementation(() => undefined); + const cliWs = addCliSocket(mockCtx, 'cli-1', [], undefined, 'usr_1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + const error = vi.spyOn(console, 'error').mockImplementation(() => undefined); - sendHeartbeat(doInstance, cliWs, [makeSession("s1", "question")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'question')]); webWs.send.mockClear(); sessionIngestMocks.resetAttentionStatusOnCliDisconnect.mockRejectedValueOnce( - new Error("db down"), + new Error('db down') ); // Leave socket listed (workerd close semantics). await disconnectCli(doInstance, cliWs); - expect( - allSent(webWs).some( - (m) => m.type === "system" && m.event === "cli.disconnected", - ), - ).toBe(true); + expect(allSent(webWs).some(m => m.type === 'system' && m.event === 'cli.disconnected')).toBe( + true + ); expect(error).toHaveBeenCalledWith( - "Failed to reset attention status on CLI disconnect", - expect.objectContaining({ error: "db down" }), + 'Failed to reset attention status on CLI disconnect', + expect.objectContaining({ error: 'db down' }) ); }); }); @@ -1825,37 +1743,37 @@ describe("UserConnectionDO", () => { // Web disconnect // ------------------------------------------------------------------------- - describe("web disconnect", () => { - it("removes from all subscription sets", () => { + describe('web disconnect', () => { + it('removes from all subscription sets', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1"), makeSession("s2")]); - sendSubscribe(doInstance, webWs, "s1"); - sendSubscribe(doInstance, webWs, "s2"); + sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); + sendSubscribe(doInstance, webWs, 's1'); + sendSubscribe(doInstance, webWs, 's2'); mockCtx.removeSocket(webWs); disconnectWeb(doInstance, webWs); // Verify: CLI events for s1 and s2 go nowhere (no crash) const cliEventMsg = JSON.stringify({ - type: "event", - sessionId: "s1", - event: "message.updated", + type: 'event', + sessionId: 's1', + event: 'message.updated', data: {}, }); doInstance.webSocketMessage(cliWs as never, cliEventMsg); // 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); - sendSubscribe(doInstance, webWs, "s1"); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendSubscribe(doInstance, webWs, 's1'); cliWs.send.mockClear(); mockCtx.removeSocket(webWs); @@ -1863,23 +1781,21 @@ describe("UserConnectionDO", () => { // CLI should get unsubscribe for s1 const msgs = allSent(cliWs); - const unsub = msgs.find( - (m: Record) => m.type === "unsubscribe", - ); - expect(unsub).toEqual({ type: "unsubscribe", sessionId: "s1" }); + const unsub = msgs.find((m: Record) => m.type === 'unsubscribe'); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); const correlationId = getCorrelationId(cliWs); @@ -1887,7 +1803,7 @@ describe("UserConnectionDO", () => { disconnectWeb(doInstance, webWs); // CLI sends response with correlationId, but the pending command is gone — no crash - sendCliResponse(doInstance, cliWs, { id: correlationId, result: "ok" }); + sendCliResponse(doInstance, cliWs, { id: correlationId, result: 'ok' }); }); }); @@ -1895,45 +1811,45 @@ describe("UserConnectionDO", () => { // Command routing // ------------------------------------------------------------------------- - describe("command routing", () => { - it("routes web command to correct CLI by sessionId", () => { + describe('command routing', () => { + it('routes web command to correct CLI by sessionId', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", - data: { text: "hello" }, + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + data: { text: 'hello' }, }); expect(cliWs.send).toHaveBeenCalledTimes(1); const sent = parseSent(cliWs) as Record; expect(sent).toMatchObject({ - type: "command", - command: "send_message", - sessionId: "s1", - data: { text: "hello" }, + type: 'command', + command: 'send_message', + sessionId: 's1', + data: { text: 'hello' }, }); - expect(typeof sent.id).toBe("string"); - expect(sent.id).not.toBe("cmd-1"); + expect(typeof sent.id).toBe('string'); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); const correlationId = getCorrelationId(cliWs); @@ -1946,23 +1862,23 @@ describe("UserConnectionDO", () => { expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', result: { success: true }, }); }); - it("sanitizes a relay-shaped CLI error before forwarding it to web", () => { + it('sanitizes a relay-shaped CLI error before forwarding it to web', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -1970,262 +1886,256 @@ describe("UserConnectionDO", () => { sendCliResponse(doInstance, cliWs, { id: correlationId, error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "cli", - message: "Command failed", + source: 'cli', + message: 'Command failed', }, }); }); - it("preserves CLI string errors for old-CLI compatibility", () => { + it('preserves CLI string errors for old-CLI compatibility', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "unknown command: list_models", + error: 'unknown command: list_models', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", - error: "unknown command: list_models", + type: 'response', + id: 'cmd-1', + error: 'unknown command: list_models', }); }); - it("accepts a pending response only from the targeted CLI socket", () => { + it('accepts a pending response only from the targeted CLI socket', () => { const { doInstance, mockCtx } = setup(); - const targetCli = addCliSocket(mockCtx, "cli-1"); - const otherCli = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + const targetCli = addCliSocket(mockCtx, 'cli-1'); + const otherCli = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, targetCli, [makeSession("s1")]); + sendHeartbeat(doInstance, targetCli, [makeSession('s1')]); sendHeartbeat(doInstance, otherCli, []); targetCli.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); const correlationId = getCorrelationId(targetCli); webWs.send.mockClear(); sendCliResponse(doInstance, otherCli, { id: correlationId, - result: "wrong-owner", + result: 'wrong-owner', }); expect(webWs.send).not.toHaveBeenCalled(); sendCliResponse(doInstance, targetCli, { id: correlationId, - result: "ok", + result: 'ok', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", - result: "ok", + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); sendCommand(doInstance, webWs, { - id: "cmd-2", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-2', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-2", + type: 'response', + id: 'cmd-2', error: { - source: "relay", - code: "CATALOG_REQUEST_PENDING", - message: "Model catalog request already pending", + source: 'relay', + code: 'CATALOG_REQUEST_PENDING', + message: 'Model catalog request already pending', }, }); - expect( - allSent(cliWs).filter((message) => message.type === "command"), - ).toHaveLength(1); + 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', () => { const now = 1_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); + vi.spyOn(Date, 'now').mockReturnValue(now); const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); vi.mocked(Date.now).mockReturnValue(now + 35_001); sendCommand(doInstance, webWs, { - id: "cmd-2", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-2', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "COMMAND_EXPIRED", - message: "Command expired", + source: 'relay', + code: 'COMMAND_EXPIRED', + message: 'Command expired', }, }); - expect( - allSent(cliWs).filter((message) => message.type === "command"), - ).toHaveLength(2); + 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', () => { const now = 1_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); + vi.spyOn(Date, 'now').mockReturnValue(now); const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); ctx.storage.setAlarm.mockClear(); vi.mocked(Date.now).mockReturnValue(now + 20_000); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); expect(ctx.storage.setAlarm).toHaveBeenCalledWith(now + 35_000); }); - it("expires pending commands during alarm processing", async () => { + it('expires pending commands during alarm processing', async () => { const now = 1_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); + vi.spyOn(Date, 'now').mockReturnValue(now); const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', }); const correlationId = getCorrelationId(cliWs); vi.mocked(Date.now).mockReturnValue(now + 34_000); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); webWs.send.mockClear(); vi.mocked(Date.now).mockReturnValue(now + 35_001); await doInstance.alarm(); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "COMMAND_EXPIRED", - message: "Command expired", + source: 'relay', + code: 'COMMAND_EXPIRED', + message: 'Command expired', }, }); - sendCliResponse(doInstance, cliWs, { id: correlationId, result: "late" }); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); for (let index = 0; index < 128; index++) { sendCommand(doInstance, webWs, { id: `cmd-${index}`, - command: "send_message", - sessionId: "s1", + command: 'send_message', + sessionId: 's1', }); } sendCommand(doInstance, webWs, { - id: "cmd-over-cap", - command: "send_message", - sessionId: "s1", + id: 'cmd-over-cap', + command: 'send_message', + sessionId: 's1', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-over-cap", + type: 'response', + id: 'cmd-over-cap', error: { - source: "relay", - code: "PENDING_COMMAND_LIMIT", - message: "Too many pending commands", + source: 'relay', + code: 'PENDING_COMMAND_LIMIT', + message: 'Too many pending commands', }, }); - expect( - allSent(cliWs).filter((message) => message.type === "command"), - ).toHaveLength(128); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2234,24 +2144,24 @@ describe("UserConnectionDO", () => { sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2262,28 +2172,28 @@ describe("UserConnectionDO", () => { }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "CATALOG_TOO_LARGE", - message: "Model catalog response is too large", + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', }, }); }); - it("rejects a multibyte list_models result over 512 KiB", () => { + it('rejects a multibyte list_models result over 512 KiB', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2294,71 +2204,71 @@ describe("UserConnectionDO", () => { }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "CATALOG_TOO_LARGE", - message: "Model catalog response is too large", + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', }, }); }); - it("returns error when CLI not found for session", () => { + it('returns error when CLI not found for session', () => { const { doInstance, mockCtx } = setup(); - const webWs = addWebSocket(mockCtx, "web-1"); + const webWs = addWebSocket(mockCtx, 'web-1'); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "unknown-session", + id: 'cmd-1', + command: 'send_message', + sessionId: 'unknown-session', }); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", - error: "Session owner not found", + type: 'response', + id: 'cmd-1', + error: 'Session owner not found', }); }); - it("rejects a stale expected session owner without forwarding", () => { + it('rejects a stale expected session owner without forwarding', () => { const { doInstance, mockCtx } = setup(); - const currentOwner = addCliSocket(mockCtx, "cli-1"); - const staleOwner = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + const currentOwner = addCliSocket(mockCtx, 'cli-1'); + const staleOwner = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, currentOwner, [makeSession("s1")]); + sendHeartbeat(doInstance, currentOwner, [makeSession('s1')]); sendHeartbeat(doInstance, staleOwner, []); currentOwner.send.mockClear(); staleOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", - connectionId: "cli-2", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + connectionId: 'cli-2', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); expect(currentOwner.send).not.toHaveBeenCalled(); expect(staleOwner.send).not.toHaveBeenCalled(); }); - it("routes command by connectionId to specific CLI", () => { + it('routes command by connectionId to specific CLI', () => { const { doInstance, mockCtx } = setup(); - const cli1 = addCliSocket(mockCtx, "cli-1"); - const cli2 = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cli1 = addCliSocket(mockCtx, 'cli-1'); + const cli2 = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); // Trigger ensureState sendHeartbeat(doInstance, cli1, []); @@ -2367,36 +2277,36 @@ describe("UserConnectionDO", () => { cli2.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - connectionId: "cli-2", + id: 'cmd-1', + command: 'send_message', + connectionId: 'cli-2', }); expect(cli1.send).not.toHaveBeenCalled(); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const web1 = addWebSocket(mockCtx, "web-1"); - const web2 = addWebSocket(mockCtx, "web-2"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const web1 = addWebSocket(mockCtx, 'web-1'); + const web2 = addWebSocket(mockCtx, 'web-2'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); // Both web sockets send commands with the same id sendCommand(doInstance, web1, { - id: "dup-id", - command: "send_message", - sessionId: "s1", + id: 'dup-id', + command: 'send_message', + sessionId: 's1', }); const corr1 = getCorrelationId(cliWs, 0); sendCommand(doInstance, web2, { - id: "dup-id", - command: "send_message", - sessionId: "s1", + id: 'dup-id', + command: 'send_message', + sessionId: 's1', }); const corr2 = getCorrelationId(cliWs, 1); @@ -2405,30 +2315,30 @@ describe("UserConnectionDO", () => { web1.send.mockClear(); web2.send.mockClear(); - sendCliResponse(doInstance, cliWs, { id: corr1, result: "result-1" }); - sendCliResponse(doInstance, cliWs, { id: corr2, result: "result-2" }); + sendCliResponse(doInstance, cliWs, { id: corr1, result: 'result-1' }); + sendCliResponse(doInstance, cliWs, { id: corr2, result: 'result-2' }); expect(parseSent(web1)).toEqual({ - type: "response", - id: "dup-id", - result: "result-1", + type: 'response', + id: 'dup-id', + result: 'result-1', }); expect(parseSent(web2)).toEqual({ - type: "response", - id: "dup-id", - result: "result-2", + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { id: "cmd-1", command: "send_message" }); + sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message' }); expect(cliWs.send).toHaveBeenCalledTimes(1); }); }); @@ -2437,183 +2347,180 @@ describe("UserConnectionDO", () => { // Command allowlist // ------------------------------------------------------------------------- - describe("command allowlist", () => { + describe('command allowlist', () => { const ALLOWED = [ - "send_message", - "interrupt", - "question_reply", - "question_reject", - "permission_respond", - "suggestion_accept", - "suggestion_dismiss", - "list_models", - "list_commands", - "send_command", - "create_session", - "exit_cli", + 'send_message', + 'interrupt', + 'question_reply', + 'question_reject', + 'permission_respond', + 'suggestion_accept', + 'suggestion_dismiss', + 'list_models', + 'list_commands', + 'send_command', + 'create_session', + 'exit_cli', ]; - it("forwards every allowed viewer command to the owning CLI", () => { + it('forwards every allowed viewer command to the owning CLI', () => { for (const command of ALLOWED) { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", + id: 'cmd-1', command, - sessionId: "s1", - data: - command === "exit_cli" - ? { protocolVersion: 1 } - : { hello: "world" }, + sessionId: 's1', + data: command === 'exit_cli' ? { protocolVersion: 1 } : { hello: 'world' }, }); expect(cliWs.send).toHaveBeenCalledTimes(1); const sent = parseSent(cliWs) as Record; expect(sent).toMatchObject({ - type: "command", + type: 'command', command, - sessionId: "s1", + sessionId: 's1', }); - expect(typeof sent.id).toBe("string"); - expect(sent.id).not.toBe("cmd-1"); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "eval", - sessionId: "s1", + id: 'cmd-1', + command: 'eval', + sessionId: 's1', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "COMMAND_NOT_ALLOWED", - message: "Command is not allowed", + source: 'relay', + code: 'COMMAND_NOT_ALLOWED', + message: 'Command is not allowed', }, }); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "shell", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'shell', + sessionId: 's1', + connectionId: 'cli-1', }); // No owner-fencing error — allowlist runs first. expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "COMMAND_NOT_ALLOWED", - message: "Command is not allowed", + source: 'relay', + code: 'COMMAND_NOT_ALLOWED', + message: 'Command is not allowed', }, }); 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', () => { const { doInstance, mockCtx } = setup(); - const webWs = addWebSocket(mockCtx, "web-1"); + const webWs = addWebSocket(mockCtx, 'web-1'); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "eval", - sessionId: "unknown-session", + id: 'cmd-1', + command: 'eval', + sessionId: 'unknown-session', }); // COMMAND_NOT_ALLOWED wins over "Session owner not found". expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "COMMAND_NOT_ALLOWED", - message: "Command is not allowed", + source: 'relay', + code: 'COMMAND_NOT_ALLOWED', + message: 'Command is not allowed', }, }); }); - it("does not allocate a pending entry or forward a disallowed command", () => { + it('does not allocate a pending entry or forward a disallowed command', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "eval", - sessionId: "s1", + 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"); + expect(sent).toMatchObject({ type: 'response', id: 'cmd-1' }); + expect(sent.id).toBe('cmd-1'); // The CLI must not have received any command envelope, and the DO must // 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" }); + 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', () => { const { doInstance, mockCtx } = setup(); - const currentOwner = addCliSocket(mockCtx, "cli-1"); - const staleOwner = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + const currentOwner = addCliSocket(mockCtx, 'cli-1'); + const staleOwner = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, currentOwner, [makeSession("s1")]); + sendHeartbeat(doInstance, currentOwner, [makeSession('s1')]); sendHeartbeat(doInstance, staleOwner, []); currentOwner.send.mockClear(); staleOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", - connectionId: "cli-2", + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + connectionId: 'cli-2', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); expect(currentOwner.send).not.toHaveBeenCalled(); @@ -2625,79 +2532,75 @@ describe("UserConnectionDO", () => { // list_commands dedupe and size cap // ------------------------------------------------------------------------- - describe("list_commands dedupe and size cap", () => { - it("rejects a duplicate in-flight list_commands request for the same viewer session and owner", () => { + describe('list_commands dedupe and size cap', () => { + it('rejects a duplicate in-flight list_commands request for the same viewer session and owner', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', }); sendCommand(doInstance, webWs, { - id: "cmd-2", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-2', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-2", + type: 'response', + id: 'cmd-2', error: { - source: "relay", - code: "CATALOG_REQUEST_PENDING", - message: "Model catalog request already pending", + source: 'relay', + code: 'CATALOG_REQUEST_PENDING', + message: 'Model catalog request already pending', }, }); - expect( - allSent(cliWs).filter((message) => message.type === "command"), - ).toHaveLength(1); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', + connectionId: 'cli-1', }); sendCommand(doInstance, webWs, { - id: "cmd-2", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-2', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', }); - expect( - allSent(cliWs).filter((message) => message.type === "command"), - ).toHaveLength(2); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2706,24 +2609,24 @@ describe("UserConnectionDO", () => { sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2734,28 +2637,28 @@ describe("UserConnectionDO", () => { }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "CATALOG_TOO_LARGE", - message: "Model catalog response is too large", + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', }, }); }); - it("rejects a multibyte list_commands result over 512 KiB", () => { + it('rejects a multibyte list_commands result over 512 KiB', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -2766,12 +2669,12 @@ describe("UserConnectionDO", () => { }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "CATALOG_TOO_LARGE", - message: "Model catalog response is too large", + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', }, }); }); @@ -2781,86 +2684,84 @@ describe("UserConnectionDO", () => { // Old CLI upgrade-required mapping // ------------------------------------------------------------------------- - describe("old CLI upgrade-required mapping", () => { + describe('old CLI upgrade-required mapping', () => { it('maps "unknown command: list_commands" to CLI_UPGRADE_REQUIRED with slash message', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "unknown command: list_commands", + error: 'unknown command: list_commands', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "CLI_UPGRADE_REQUIRED", - message: - "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', + message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', }, }); }); it('maps "unknown command: send_command" to CLI_UPGRADE_REQUIRED with slash message', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_command", - sessionId: "s1", - connectionId: "cli-1", - data: { command: "init" }, + id: 'cmd-1', + command: 'send_command', + sessionId: 's1', + connectionId: 'cli-1', + data: { command: 'init' }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "unknown command: send_command", + error: 'unknown command: send_command', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "CLI_UPGRADE_REQUIRED", - message: - "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', + message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', }, }); }); it('maps "unknown command: exit_cli" to CLI_UPGRADE_REQUIRED with slash message', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); const correlationId = getCorrelationId(cliWs); @@ -2868,150 +2769,149 @@ describe("UserConnectionDO", () => { sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "unknown command: exit_cli", + error: 'unknown command: exit_cli', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "CLI_UPGRADE_REQUIRED", - message: - "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', + message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', }, }); }); it('maps "unknown command: create_session" to CLI_UPGRADE_REQUIRED with create_session message', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "create_session", - connectionId: "cli-1", - data: { title: "New session" }, + id: 'cmd-1', + command: 'create_session', + connectionId: 'cli-1', + data: { title: 'New session' }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "unknown command: create_session", + error: 'unknown command: create_session', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "CLI_UPGRADE_REQUIRED", + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', message: - "Creating remote sessions from mobile requires a newer Kilo CLI. Update Kilo CLI and reconnect.", + 'Creating remote sessions from mobile requires a newer Kilo CLI. Update Kilo CLI and reconnect.', }, }); }); it('preserves "unknown command: list_models" because list_models is not in the upgrade-required set', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_models", - sessionId: "s1", + id: 'cmd-1', + command: 'list_models', + sessionId: 's1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "unknown command: list_models", + error: 'unknown command: list_models', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", - error: "unknown command: list_models", + type: 'response', + id: 'cmd-1', + error: 'unknown command: list_models', }); }); - it("preserves an unrelated CLI string error for send_command", () => { + it('preserves an unrelated CLI string error for send_command', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_command", - sessionId: "s1", - data: { command: "init" }, + id: 'cmd-1', + command: 'send_command', + sessionId: 's1', + data: { command: 'init' }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "session not ready", + error: 'session not ready', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", - error: "session not ready", + type: 'response', + id: 'cmd-1', + error: 'session not ready', }); }); it('does not match a longer error that merely starts with "unknown command: list_commands"', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'list_commands', + sessionId: 's1', + connectionId: 'cli-1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "unknown command: list_commands: try again", + error: 'unknown command: list_commands: try again', }); // Exact-match only — do not misclassify longer error strings. expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", - error: "unknown command: list_commands: try again", + type: 'response', + id: 'cmd-1', + error: 'unknown command: list_commands: try again', }); }); - it("preserves a longer exit_cli unknown-command error", () => { + it('preserves a longer exit_cli unknown-command error', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); const correlationId = getCorrelationId(cliWs); @@ -3019,13 +2919,13 @@ describe("UserConnectionDO", () => { sendCliResponse(doInstance, cliWs, { id: correlationId, - error: "unknown command: exit_cli: session not ready", + error: 'unknown command: exit_cli: session not ready', }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", - error: "unknown command: exit_cli: session not ready", + type: 'response', + id: 'cmd-1', + error: 'unknown command: exit_cli: session not ready', }); }); }); @@ -3036,142 +2936,134 @@ describe("UserConnectionDO", () => { // and must NOT be subject to the 512 KiB catalog response cap.) // ------------------------------------------------------------------------- - describe("send_command / create_session negative coverage", () => { - it("forwards two in-flight same-owner/same-session send_command requests without deduping", () => { + describe('send_command / create_session negative coverage', () => { + it('forwards two in-flight same-owner/same-session send_command requests without deduping', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_command", - sessionId: "s1", - connectionId: "cli-1", - data: { command: "init" }, + id: 'cmd-1', + command: 'send_command', + sessionId: 's1', + connectionId: 'cli-1', + data: { command: 'init' }, }); sendCommand(doInstance, webWs, { - id: "cmd-2", - command: "send_command", - sessionId: "s1", - connectionId: "cli-1", - data: { command: "plan" }, + id: 'cmd-2', + command: 'send_command', + sessionId: 's1', + connectionId: 'cli-1', + data: { command: 'plan' }, }); - const cliCommands = allSent(cliWs).filter( - (message) => message.type === "command", - ); + const cliCommands = allSent(cliWs).filter(message => message.type === 'command'); expect(cliCommands).toHaveLength(2); expect(cliCommands[0]).toMatchObject({ - type: "command", - command: "send_command", - sessionId: "s1", - data: { command: "init" }, + type: 'command', + command: 'send_command', + sessionId: 's1', + data: { command: 'init' }, }); expect(cliCommands[1]).toMatchObject({ - type: "command", - command: "send_command", - sessionId: "s1", - data: { command: "plan" }, + type: 'command', + command: 'send_command', + sessionId: 's1', + data: { command: 'plan' }, }); expect(cliCommands[0].id).not.toBe(cliCommands[1].id); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "create_session", - connectionId: "cli-1", - data: { title: "First session" }, + id: 'cmd-1', + command: 'create_session', + connectionId: 'cli-1', + data: { title: 'First session' }, }); sendCommand(doInstance, webWs, { - id: "cmd-2", - command: "create_session", - connectionId: "cli-1", - data: { title: "Second session" }, + id: 'cmd-2', + command: 'create_session', + connectionId: 'cli-1', + data: { title: 'Second session' }, }); - const cliCommands = allSent(cliWs).filter( - (message) => message.type === "command", - ); + const cliCommands = allSent(cliWs).filter(message => message.type === 'command'); expect(cliCommands).toHaveLength(2); expect(cliCommands[0]).toMatchObject({ - type: "command", - command: "create_session", - data: { title: "First session" }, + type: 'command', + command: 'create_session', + data: { title: 'First session' }, }); expect(cliCommands[1]).toMatchObject({ - type: "command", - command: "create_session", - data: { title: "Second session" }, + type: 'command', + command: 'create_session', + data: { title: 'Second session' }, }); expect(cliCommands[0].id).not.toBe(cliCommands[1].id); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "send_command", - sessionId: "s1", - connectionId: "cli-1", - data: { command: "init" }, + id: 'cmd-1', + command: 'send_command', + sessionId: 's1', + connectionId: 'cli-1', + data: { command: 'init' }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - const result = createResultWithSerializedBytes( - MAX_CATALOG_RESULT_BYTES + 1, - ); + const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', result, }); }); - it("relays a create_session result over 512 KiB unchanged", () => { + it('relays a create_session result over 512 KiB unchanged', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "create_session", - connectionId: "cli-1", - data: { title: "Big session" }, + id: 'cmd-1', + command: 'create_session', + connectionId: 'cli-1', + data: { title: 'Big session' }, }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - const result = createResultWithSerializedBytes( - MAX_CATALOG_RESULT_BYTES + 1, - ); + const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', result, }); }); @@ -3181,210 +3073,204 @@ describe("UserConnectionDO", () => { // exit_cli routing and relay policy // ------------------------------------------------------------------------- - describe("exit_cli routing and relay policy", () => { + describe('exit_cli routing and relay policy', () => { it.each([ - { label: "missing sessionId", input: { data: { protocolVersion: 1 } } }, - { label: "missing data", input: { sessionId: "s1" } }, + { 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: 'extra data field', + input: { sessionId: 's1', data: { protocolVersion: 1, extra: true } }, }, - { label: "null data", input: { sessionId: "s1", data: null } }, + { label: 'null data', input: { sessionId: 's1', data: null } }, { - label: "array data", - input: { sessionId: "s1", data: [{ protocolVersion: 1 }] }, + label: 'array data', + input: { sessionId: 's1', data: [{ protocolVersion: 1 }] }, }, { - label: "primitive data", - input: { sessionId: "s1", data: "protocolVersion=1" }, + label: 'primitive data', + input: { sessionId: 's1', data: 'protocolVersion=1' }, }, - ])("rejects $label before routing or pending allocation", ({ input }) => { + ])('rejects $label before routing or pending allocation', ({ input }) => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", + id: 'cmd-1', + command: 'exit_cli', ...input, }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "INVALID_COMMAND", - message: "Invalid command", + source: 'relay', + code: 'INVALID_COMMAND', + message: 'Invalid command', }, }); expect(cliWs.send).not.toHaveBeenCalled(); - expect(Reflect.get(doInstance, "pendingCommands")).toEqual(new Map()); + 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', () => { const { doInstance, mockCtx } = setup(); - const selectedOwner = addCliSocket(mockCtx, "cli-1"); - const otherCli = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + const selectedOwner = addCliSocket(mockCtx, 'cli-1'); + const otherCli = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, selectedOwner, [makeSession("s1")]); - sendHeartbeat(doInstance, otherCli, [makeSession("s2")]); + sendHeartbeat(doInstance, selectedOwner, [makeSession('s1')]); + sendHeartbeat(doInstance, otherCli, [makeSession('s2')]); selectedOwner.send.mockClear(); otherCli.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); - expect( - allSent(selectedOwner).filter((message) => message.type === "command"), - ).toEqual([ + expect(allSent(selectedOwner).filter(message => message.type === 'command')).toEqual([ expect.objectContaining({ - type: "command", - command: "exit_cli", - sessionId: "s1", + type: 'command', + command: 'exit_cli', + sessionId: 's1', data: { protocolVersion: 1 }, }), ]); 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', () => { const { doInstance, mockCtx } = setup(); - const currentOwner = addCliSocket(mockCtx, "cli-1"); - const staleOwner = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + const currentOwner = addCliSocket(mockCtx, 'cli-1'); + const staleOwner = addCliSocket(mockCtx, 'cli-2'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, currentOwner, [makeSession("s1")]); + sendHeartbeat(doInstance, currentOwner, [makeSession('s1')]); sendHeartbeat(doInstance, staleOwner, []); currentOwner.send.mockClear(); staleOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-2", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-2', data: { protocolVersion: 1 }, }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); expect(currentOwner.send).not.toHaveBeenCalled(); 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', () => { const { doInstance, mockCtx } = setup(); - const webWs = addWebSocket(mockCtx, "web-1"); + const webWs = addWebSocket(mockCtx, 'web-1'); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', data: { protocolVersion: 1 }, }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", - error: "Session owner not found", + type: 'response', + id: 'cmd-1', + error: 'Session owner not found', }); }); - it("does not dedupe concurrent exit_cli requests", () => { + it('does not dedupe concurrent exit_cli requests', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); sendCommand(doInstance, webWs, { - id: "cmd-2", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-2', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); - const commands = allSent(cliWs).filter( - (message) => message.type === "command", - ); + const commands = allSent(cliWs).filter(message => message.type === 'command'); expect(commands).toHaveLength(2); expect(commands[0].id).not.toBe(commands[1].id); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); const correlationId = getCorrelationId(cliWs); - const result = createResultWithSerializedBytes( - MAX_CATALOG_RESULT_BYTES + 1, - ); + const result = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); sendCliResponse(doInstance, cliWs, { id: correlationId, result }); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); const correlationId = getCorrelationId(cliWs); @@ -3393,44 +3279,42 @@ describe("UserConnectionDO", () => { // Exit's own effect: CLI heartbeat no longer lists the session. sendHeartbeat(doInstance, cliWs, []); - expect( - allSent(webWs).find((m) => m.type === "response" && m.id === "cmd-1"), - ).toEqual({ - type: "response", - id: "cmd-1", + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ + type: 'response', + id: 'cmd-1', result: {}, }); expect( allSent(webWs).some( - (m) => - m.type === "response" && - m.id === "cmd-1" && + m => + m.type === 'response' && + m.id === 'cmd-1' && isRecord(m.error) && - m.error.code === "SESSION_OWNER_CHANGED", - ), + m.error.code === 'SESSION_OWNER_CHANGED' + ) ).toBe(false); // CLI's late ACK must not produce a second cmd-1 response. const responsesBeforeLateAck = allSent(webWs).filter( - (m) => m.type === "response" && m.id === "cmd-1", + m => m.type === 'response' && m.id === 'cmd-1' ).length; sendCliResponse(doInstance, cliWs, { id: correlationId, result: {} }); - expect( - allSent(webWs).filter((m) => m.type === "response" && m.id === "cmd-1"), - ).toHaveLength(responsesBeforeLateAck); + expect(allSent(webWs).filter(m => m.type === 'response' && m.id === 'cmd-1')).toHaveLength( + responsesBeforeLateAck + ); }); - it("resolves exit_cli successfully when the owning socket closes", async () => { + it('resolves exit_cli successfully when the owning socket closes', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); webWs.send.mockClear(); @@ -3439,109 +3323,102 @@ describe("UserConnectionDO", () => { await disconnectCli(doInstance, cliWs); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', result: {}, }); }); - it.each(["list_models", "send_message"] as const)( - "still fails %s with SESSION_OWNER_CHANGED when heartbeat drops the session", - (command) => { + it.each(['list_models', 'send_message'] as const)( + 'still fails %s with SESSION_OWNER_CHANGED when heartbeat drops the session', + command => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", + id: 'cmd-1', command, - sessionId: "s1", - connectionId: "cli-1", + sessionId: 's1', + connectionId: 'cli-1', }); webWs.send.mockClear(); sendHeartbeat(doInstance, cliWs, []); - expect( - allSent(webWs).find((m) => m.type === "response" && m.id === "cmd-1"), - ).toEqual({ - type: "response", - id: "cmd-1", + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); - }, + } ); - it("still fails exit_cli with SESSION_OWNER_CHANGED on genuine takeover", () => { + it('still fails exit_cli with SESSION_OWNER_CHANGED on genuine takeover', () => { const { doInstance, mockCtx } = setup(); - const firstOwner = addCliSocket(mockCtx, "cli-1"); - const nextOwner = addCliSocket(mockCtx, "cli-2"); - const webWs = addWebSocket(mockCtx, "web-1"); + 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, firstOwner, [makeSession('s1')]); sendHeartbeat(doInstance, nextOwner, []); firstOwner.send.mockClear(); webWs.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); webWs.send.mockClear(); - sendHeartbeat(doInstance, nextOwner, [makeSession("s1")]); + sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); - expect( - allSent(webWs).find((m) => m.type === "response" && m.id === "cmd-1"), - ).toEqual({ - type: "response", - id: "cmd-1", + expect(allSent(webWs).find(m => m.type === 'response' && m.id === 'cmd-1')).toEqual({ + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); }); - 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', () => { const { doInstance, mockCtx } = setup(); - const firstCli = connectCliSocket(doInstance, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const firstCli = connectCliSocket(doInstance, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, firstCli, [makeSession("s1")]); + sendHeartbeat(doInstance, firstCli, [makeSession('s1')]); firstCli.send.mockClear(); sendCommand(doInstance, webWs, { - id: "cmd-1", - command: "exit_cli", - sessionId: "s1", - connectionId: "cli-1", + id: 'cmd-1', + command: 'exit_cli', + sessionId: 's1', + connectionId: 'cli-1', data: { protocolVersion: 1 }, }); webWs.send.mockClear(); - connectCliSocket(doInstance, "cli-1"); + connectCliSocket(doInstance, 'cli-1'); - expect(firstCli.close).toHaveBeenCalledWith( - 1000, - "replaced by reconnect", - ); + expect(firstCli.close).toHaveBeenCalledWith(1000, 'replaced by reconnect'); expect(parseSent(webWs)).toEqual({ - type: "response", - id: "cmd-1", + type: 'response', + id: 'cmd-1', error: { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }, }); }); @@ -3551,87 +3428,87 @@ describe("UserConnectionDO", () => { // CLI event forwarding // ------------------------------------------------------------------------- - describe("CLI event forwarding", () => { - it("forwards events to subscribed web sockets only", () => { + describe('CLI event forwarding', () => { + it('forwards events to subscribed web sockets only', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const subWeb = addWebSocket(mockCtx, "web-sub"); - const otherWeb = addWebSocket(mockCtx, "web-other"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const subWeb = addWebSocket(mockCtx, 'web-sub'); + const otherWeb = addWebSocket(mockCtx, 'web-other'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); - sendSubscribe(doInstance, subWeb, "s1"); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendSubscribe(doInstance, subWeb, 's1'); subWeb.send.mockClear(); otherWeb.send.mockClear(); // CLI sends event for s1 const eventMsg = JSON.stringify({ - type: "event", - sessionId: "s1", - event: "message.updated", - data: { id: "msg-1" }, + type: 'event', + sessionId: 's1', + event: 'message.updated', + data: { id: 'msg-1' }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(subWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(subWeb)).toEqual({ - type: "event", - sessionId: "s1", - event: "message.updated", - data: { id: "msg-1" }, + type: 'event', + sessionId: 's1', + event: 'message.updated', + data: { id: 'msg-1' }, }); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const parentWeb = addWebSocket(mockCtx, "web-parent"); - const childWeb = addWebSocket(mockCtx, "web-child"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const parentWeb = addWebSocket(mockCtx, 'web-parent'); + const childWeb = addWebSocket(mockCtx, 'web-child'); - sendHeartbeat(doInstance, cliWs, [makeSession("parent-session")]); - sendSubscribe(doInstance, parentWeb, "parent-session"); - sendSubscribe(doInstance, childWeb, "child-session-1"); + sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); + sendSubscribe(doInstance, parentWeb, 'parent-session'); + sendSubscribe(doInstance, childWeb, 'child-session-1'); parentWeb.send.mockClear(); childWeb.send.mockClear(); const eventMsg = JSON.stringify({ - type: "event", - sessionId: "child-session-1", - parentSessionId: "parent-session", - event: "message.updated", - data: { id: "msg-1" }, + type: 'event', + sessionId: 'child-session-1', + parentSessionId: 'parent-session', + event: 'message.updated', + data: { id: 'msg-1' }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(parentWeb.send).toHaveBeenCalledTimes(1); expect(childWeb.send).toHaveBeenCalledTimes(1); const expected = { - type: "event", - sessionId: "child-session-1", - parentSessionId: "parent-session", - event: "message.updated", - data: { id: "msg-1" }, + type: 'event', + sessionId: 'child-session-1', + parentSessionId: 'parent-session', + event: 'message.updated', + data: { id: 'msg-1' }, }; expect(parseSent(parentWeb)).toEqual(expected); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("parent-session")]); - sendSubscribe(doInstance, webWs, "parent-session"); - sendSubscribe(doInstance, webWs, "child-session-1"); + sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); + sendSubscribe(doInstance, webWs, 'parent-session'); + sendSubscribe(doInstance, webWs, 'child-session-1'); webWs.send.mockClear(); const eventMsg = JSON.stringify({ - type: "event", - sessionId: "child-session-1", - parentSessionId: "parent-session", - event: "message.updated", - data: { id: "msg-1" }, + type: 'event', + sessionId: 'child-session-1', + parentSessionId: 'parent-session', + event: 'message.updated', + data: { id: 'msg-1' }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); @@ -3639,102 +3516,102 @@ 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("parent-session")]); - sendSubscribe(doInstance, webWs, "parent-session"); + sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); + sendSubscribe(doInstance, webWs, 'parent-session'); webWs.send.mockClear(); // CLI sends event for a child session with parentSessionId const eventMsg = JSON.stringify({ - type: "event", - sessionId: "child-session-1", - parentSessionId: "parent-session", - event: "message.updated", - data: { id: "msg-child-1" }, + type: 'event', + sessionId: 'child-session-1', + parentSessionId: 'parent-session', + event: 'message.updated', + data: { id: 'msg-child-1' }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: "event", - sessionId: "child-session-1", - parentSessionId: "parent-session", - event: "message.updated", - data: { id: "msg-child-1" }, + type: 'event', + sessionId: 'child-session-1', + parentSessionId: 'parent-session', + event: 'message.updated', + data: { id: 'msg-child-1' }, }); }); - it("drops child event when neither sessionId nor parentSessionId has subscribers", () => { + it('drops child event when neither sessionId nor parentSessionId has subscribers', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("other-session")]); - sendSubscribe(doInstance, webWs, "other-session"); + sendHeartbeat(doInstance, cliWs, [makeSession('other-session')]); + sendSubscribe(doInstance, webWs, 'other-session'); webWs.send.mockClear(); // Child event with parent that nobody subscribes to const eventMsg = JSON.stringify({ - type: "event", - sessionId: "child-session-1", - parentSessionId: "unknown-parent", - event: "message.updated", - data: { id: "msg-child-1" }, + type: 'event', + sessionId: 'child-session-1', + parentSessionId: 'unknown-parent', + event: 'message.updated', + data: { id: 'msg-child-1' }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).not.toHaveBeenCalled(); }); - it("events without parentSessionId still route normally (backward compat)", () => { + it('events without parentSessionId still route normally (backward compat)', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); - sendSubscribe(doInstance, webWs, "s1"); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); const eventMsg = JSON.stringify({ - type: "event", - sessionId: "s1", - event: "message.updated", - data: { id: "msg-1" }, + type: 'event', + sessionId: 's1', + event: 'message.updated', + data: { id: 'msg-1' }, }); doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ - type: "event", - sessionId: "s1", - event: "message.updated", - data: { id: "msg-1" }, + type: 'event', + sessionId: 's1', + event: 'message.updated', + data: { id: 'msg-1' }, }); }); - it("child event does not include parentSessionId when not set", () => { + it('child event does not include parentSessionId when not set', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); - sendSubscribe(doInstance, webWs, "s1"); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); const eventMsg = JSON.stringify({ - type: "event", - sessionId: "s1", - event: "session.status", + type: 'event', + sessionId: 's1', + event: 'session.status', data: {}, }); doInstance.webSocketMessage(cliWs as never, eventMsg); const sent = parseSent(webWs); - expect(sent).not.toHaveProperty("parentSessionId"); + expect(sent).not.toHaveProperty('parentSessionId'); }); }); @@ -3742,12 +3619,12 @@ describe("UserConnectionDO", () => { // Broadcast resilience // ------------------------------------------------------------------------- - describe("broadcast resilience", () => { - it("one closed socket does not abort send to other web sockets", () => { + describe('broadcast resilience', () => { + it('one closed socket does not abort send to other web sockets', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const failWeb = addWebSocket(mockCtx, "web-fail"); - const okWeb = addWebSocket(mockCtx, "web-ok"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const failWeb = addWebSocket(mockCtx, 'web-fail'); + const okWeb = addWebSocket(mockCtx, 'web-ok'); // Both web sockets receive heartbeats via broadcast (no subscription needed). failWeb.send.mockClear(); @@ -3755,16 +3632,16 @@ describe("UserConnectionDO", () => { // Make failWeb throw on send failWeb.send.mockImplementation(() => { - throw new Error("socket closed"); + throw new Error('socket closed'); }); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // okWeb should still receive the message expect(okWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(okWeb)).toMatchObject({ - type: "system", - event: "sessions.heartbeat", + type: 'system', + event: 'sessions.heartbeat', }); }); }); @@ -3773,51 +3650,49 @@ describe("UserConnectionDO", () => { // Hibernation recovery (ensureState) // ------------------------------------------------------------------------- - describe("ensureState (hibernation recovery)", () => { - it("reconstructs sessionOwners and connectionSessions from CLI attachments", () => { + describe('ensureState (hibernation recovery)', () => { + it('reconstructs sessionOwners and connectionSessions from CLI attachments', () => { const { doInstance, mockCtx } = setup(); // Simulate hibernation: sockets exist with pre-set attachments - const sessions = [makeSession("s1"), makeSession("s2")]; - addCliSocket(mockCtx, "cli-1", sessions); - const webWs = addWebSocket(mockCtx, "web-1"); + const sessions = [makeSession('s1'), makeSession('s2')]; + addCliSocket(mockCtx, 'cli-1', sessions); + const webWs = addWebSocket(mockCtx, 'web-1'); // Trigger ensureState by calling any method (e.g., webSocketMessage with subscribe) - sendSubscribe(doInstance, webWs, "s1"); + sendSubscribe(doInstance, webWs, 's1'); // Verify state was reconstructed by routing a command - const web2 = addWebSocket(mockCtx, "web-2"); + const web2 = addWebSocket(mockCtx, 'web-2'); sendCommand(doInstance, web2, { - id: "cmd-1", - command: "send_message", - sessionId: "s1", + 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")); + 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", - ); + const cmdMsg = cliMsgs.find((m: Record) => m.type === 'command'); expect(cmdMsg).toMatchObject({ - type: "command", - command: "send_message", + type: 'command', + command: 'send_message', }); }); - it("reconstructs webSubscriptions from web attachments", () => { + it('reconstructs webSubscriptions from web attachments', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1", [makeSession("s1")]); + const cliWs = addCliSocket(mockCtx, 'cli-1', [makeSession('s1')]); // Web socket with pre-existing subscription (from hibernation) - const webWs = addWebSocket(mockCtx, "web-1", ["s1"]); + const webWs = addWebSocket(mockCtx, 'web-1', ['s1']); // Trigger ensureState by calling any method const triggerMsg = JSON.stringify({ - type: "event", - sessionId: "s1", - event: "test", + type: 'event', + sessionId: 's1', + event: 'test', data: {}, }); doInstance.webSocketMessage(cliWs as never, triggerMsg); @@ -3825,30 +3700,30 @@ describe("UserConnectionDO", () => { // webWs should have received the event because it was subscribed via attachment expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toMatchObject({ - type: "event", - sessionId: "s1", + type: 'event', + sessionId: 's1', }); }); - it("does not restore subscriptions from a viewer already replaced before hibernation", () => { + it('does not restore subscriptions from a viewer already replaced before hibernation', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1", [makeSession("s1")]); - const replacedWeb = addWebSocket(mockCtx, "web-old", ["s1"]); + const cliWs = addCliSocket(mockCtx, 'cli-1', [makeSession('s1')]); + const replacedWeb = addWebSocket(mockCtx, 'web-old', ['s1']); replacedWeb.serializeAttachment({ - role: "web", - connectionId: "web-old", - subscribedSessions: ["s1"], + role: 'web', + connectionId: 'web-old', + subscribedSessions: ['s1'], replaced: true, }); doInstance.webSocketMessage( cliWs as never, JSON.stringify({ - type: "event", - sessionId: "s1", - event: "test", + type: 'event', + sessionId: 's1', + event: 'test', data: {}, - }), + }) ); expect(replacedWeb.send).not.toHaveBeenCalled(); @@ -3859,47 +3734,47 @@ describe("UserConnectionDO", () => { // getActiveSessions RPC // ------------------------------------------------------------------------- - describe("getActiveSessions", () => { - it("returns sessions from live CLI connections", () => { + describe('getActiveSessions', () => { + it('returns sessions from live CLI connections', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [ - makeSession("s1", "busy", "Fix bug"), - makeSession("s2", "idle", "Review PR"), + makeSession('s1', 'busy', 'Fix bug'), + makeSession('s2', 'idle', 'Review PR'), ]); const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: "s1", status: "busy", title: "Fix bug", connectionId: "cli-1" }, - { id: "s2", status: "idle", title: "Review PR", connectionId: "cli-1" }, + { id: 's1', status: 'busy', title: 'Fix bug', connectionId: 'cli-1' }, + { id: 's2', status: 'idle', title: 'Review PR', connectionId: 'cli-1' }, ]); }); - it("includes the CLI-reported protocolVersion on each session row", () => { + it('includes the CLI-reported protocolVersion on each session row', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1", "busy", "Fix bug")], { - protocolVersion: "1", + sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Fix bug')], { + protocolVersion: '1', }); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // Remove from sockets (simulates close) mockCtx.removeSocket(cliWs); @@ -3908,86 +3783,84 @@ describe("UserConnectionDO", () => { expect(result).toEqual([]); }); - it("excludes child sessions reported with parentSessionId in heartbeat", () => { + it('excludes child sessions reported with parentSessionId in heartbeat', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [ - makeSession("root-1", "busy", "Root session"), - makeSession("child-1", "busy", "Child session", "root-1"), + makeSession('root-1', 'busy', 'Root session'), + makeSession('child-1', 'busy', 'Child session', 'root-1'), ]); 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); // First heartbeat: root + child sendHeartbeat(doInstance, cliWs, [ - makeSession("root-1", "busy", "Root session"), - makeSession("child-1", "busy", "Child session", "root-1"), + makeSession('root-1', 'busy', 'Root session'), + makeSession('child-1', 'busy', 'Child session', 'root-1'), ]); // Second heartbeat: only root (child finished) - sendHeartbeat(doInstance, cliWs, [ - makeSession("root-1", "idle", "Root session"), - ]); + sendHeartbeat(doInstance, cliWs, [makeSession('root-1', 'idle', 'Root session')]); 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)', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [ - makeSession("s1", "busy", "On a Mac", undefined, "darwin"), - makeSession("s2", "idle", "Other"), + makeSession('s1', 'busy', 'On a Mac', undefined, 'darwin'), + makeSession('s2', 'idle', 'Other'), ]); 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" }, + { 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)', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); // Legacy CLI heartbeat without platform. - sendHeartbeat(doInstance, cliWs, [makeSession("s1", "busy", "Legacy")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Legacy')]); const result = doInstance.getActiveSessions(); expect(result).toEqual([ - { id: "s1", status: "busy", title: "Legacy", connectionId: "cli-1" }, + { id: 's1', status: 'busy', title: 'Legacy', connectionId: 'cli-1' }, ]); - expect(result[0]).not.toHaveProperty("platform"); + expect(result[0]).not.toHaveProperty('platform'); }); }); @@ -3995,20 +3868,20 @@ describe("UserConnectionDO", () => { // getConnectedInstances RPC (W3) // ------------------------------------------------------------------------- - describe("getConnectedInstances", () => { - it("returns one row per CLI socket that has an `instance` attachment", () => { + describe('getConnectedInstances', () => { + it('returns one row per CLI socket that has an `instance` attachment', () => { const { doInstance, mockCtx } = setup(); // Use the hibernated-attachment pattern (no heartbeat) — the live // scan reads the `instance` directly from the attachment, which is // what the spec requires: a fresh value with no in-memory map. - addCliSocket(mockCtx, "cli-A", [], { - name: "laptop-A", - projectName: "kilo", - version: "0.1.2", + addCliSocket(mockCtx, 'cli-A', [], { + name: 'laptop-A', + 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); @@ -4017,47 +3890,45 @@ describe("UserConnectionDO", () => { 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" }, - ]), + { 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', () => { 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"); + 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`', () => { const { doInstance, mockCtx } = setup(); // Legacy CLI: pre-spawner heartbeat has no `instance`. - const cliWs = addCliSocket(mockCtx, "legacy-1"); + const cliWs = addCliSocket(mockCtx, 'legacy-1'); sendHeartbeat(doInstance, cliWs, []); const { instances } = doInstance.getConnectedInstances(); expect(instances).toEqual([]); }); - it("excludes web sockets", () => { + it('excludes web sockets', () => { const { doInstance, mockCtx } = setup(); addWebSocket(mockCtx); // A web socket with an `instance`-shaped attachment must still be skipped. - const webWithInstance = createMockWs(["web"], { - role: "web", - connectionId: "web-1", + const webWithInstance = createMockWs(['web'], { + role: 'web', + connectionId: 'web-1', subscribedSessions: [], } as never); mockCtx.addSocket(webWithInstance); @@ -4066,36 +3937,36 @@ 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)', () => { 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. - addCliSocket(mockCtx, "cli-h", [], { - name: "laptop-h", - projectName: "kilo", - version: "1.0.0", + addCliSocket(mockCtx, 'cli-h', [], { + name: 'laptop-h', + projectName: 'kilo', + version: '1.0.0', }); 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', () => { const { doInstance, mockCtx } = setup(); // Hibernated attachment carries capabilities — same source // getConnectedInstances already uses for instance/version. - const cliWs = createMockWs(["cli"], { - role: "cli", - connectionId: "cli-cap", + const cliWs = createMockWs(['cli'], { + role: 'cli', + connectionId: 'cli-cap', sessions: [], - instance: { name: "laptop-cap", projectName: "kilo" }, + instance: { name: 'laptop-cap', projectName: 'kilo' }, capabilities: { attachments: true }, }); mockCtx.addSocket(cliWs); @@ -4103,56 +3974,56 @@ describe("UserConnectionDO", () => { const { instances } = doInstance.getConnectedInstances(); expect(instances).toEqual([ { - connectionId: "cli-cap", - name: "laptop-cap", - projectName: "kilo", + connectionId: 'cli-cap', + name: 'laptop-cap', + projectName: 'kilo', capabilities: { attachments: true }, }, ]); }); - it("omits capabilities when the CLI attachment has none (legacy CLI)", () => { + it('omits capabilities when the CLI attachment has none (legacy CLI)', () => { const { doInstance, mockCtx } = setup(); - addCliSocket(mockCtx, "cli-legacy-cap", [], { - name: "laptop-legacy", - projectName: "kilo", + addCliSocket(mockCtx, 'cli-legacy-cap', [], { + name: 'laptop-legacy', + projectName: 'kilo', }); 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"); + expect(instances[0]).not.toHaveProperty('capabilities'); }); - it("persists `instance` in the WS attachment across heartbeats", () => { + it('persists `instance` in the WS attachment across heartbeats', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, [], { - protocolVersion: "1", - instance: { name: "laptop-1", projectName: "kilo", version: "0.1.0" }, + protocolVersion: '1', + 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", + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); // First heartbeat: with instance. sendHeartbeat(doInstance, cliWs, [], { - instance: { name: "laptop-1", projectName: "kilo" }, + instance: { name: 'laptop-1', projectName: 'kilo' }, }); // Second heartbeat: instance removed (legacy fallback). The DO must not // keep a stale `instance` value in the attachment. @@ -4167,7 +4038,7 @@ describe("UserConnectionDO", () => { // WS attachment size guardrail (W3) // ------------------------------------------------------------------------- - describe("WS attachment size", () => { + describe('WS attachment size', () => { // The Cloudflare `serializeAttachment` budget is ~2 KiB. A bounded // instance object (name+projectName+version, all max length) adds well // under 200 bytes; this test pins that contract so a future schema @@ -4179,42 +4050,38 @@ 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', () => { const worstCaseInstance = { - name: "x".repeat(64), - projectName: "x".repeat(64), - version: "x".repeat(32), + name: 'x'.repeat(64), + projectName: 'x'.repeat(64), + version: 'x'.repeat(32), }; // 4 sessions with realistic-but-large titles, git URLs, and branches. // (4 is a generous upper bound for a single CLI owning a live session // fleet; the actual HeartbeatSession shape imposes tighter per-field // limits at the protocol layer.) const sessions = Array.from({ length: 4 }, (_, i) => ({ - id: `ses_${String(i).padStart(26, "0")}`, - status: "busy", - title: "T".repeat(120), - gitUrl: "https://github.com/org/" + "x".repeat(60) + ".git", - gitBranch: "b".repeat(40), + id: `ses_${String(i).padStart(26, '0')}`, + status: 'busy', + title: 'T'.repeat(120), + gitUrl: 'https://github.com/org/' + 'x'.repeat(60) + '.git', + gitBranch: 'b'.repeat(40), })); const attachment = { - role: "cli" as const, - connectionId: "cli-1", + role: 'cli' as const, + connectionId: 'cli-1', sessions, - protocolVersion: "255.255.65535", - kiloUserId: "usr_" + "x".repeat(28), + protocolVersion: '255.255.65535', + kiloUserId: 'usr_' + 'x'.repeat(28), instance: worstCaseInstance, }; - const serialized = new TextEncoder().encode( - JSON.stringify(attachment), - ).byteLength; + const serialized = new TextEncoder().encode(JSON.stringify(attachment)).byteLength; expect(serialized).toBeLessThan(SERIALIZE_ATTACHMENT_BUDGET); // Sanity: the bounded instance alone is far below the headroom. - const instanceBytes = new TextEncoder().encode( - JSON.stringify(worstCaseInstance), - ).byteLength; + const instanceBytes = new TextEncoder().encode(JSON.stringify(worstCaseInstance)).byteLength; expect(instanceBytes).toBeLessThan(INSTANCE_HEADROOM); }); }); @@ -4223,29 +4090,23 @@ describe("UserConnectionDO", () => { // Owner-unique active sessions (W-followup) // ------------------------------------------------------------------------- - describe("owner-unique active sessions", () => { - it("emits owner-unique rows: ownership transfer with both CLIs live yields exactly one row under the new owner", () => { + describe('owner-unique active sessions', () => { + it('emits owner-unique rows: ownership transfer with both CLIs live yields exactly one row under the new owner', () => { const { doInstance, ctx, mockCtx } = setup(); - const oldOwner = addCliSocket(mockCtx, "cli-old"); - const newOwner = addCliSocket(mockCtx, "cli-new"); + const oldOwner = addCliSocket(mockCtx, 'cli-old'); + const newOwner = addCliSocket(mockCtx, 'cli-new'); // cli-old claims the session - sendHeartbeat(doInstance, oldOwner, [ - makeSession("ses_transfer", "busy", "Transfer me"), - ]); + sendHeartbeat(doInstance, oldOwner, [makeSession('ses_transfer', 'busy', 'Transfer me')]); // cli-new also claims the same session id while cli-old is still connected. // The DO routes the session to the new owner (sessionOwners.get === 'cli-new'). - sendHeartbeat(doInstance, newOwner, [ - makeSession("ses_transfer", "busy", "Transfer me"), - ]); + sendHeartbeat(doInstance, newOwner, [makeSession('ses_transfer', 'busy', 'Transfer me')]); // Both CLIs are still live sockets — the snapshot should see them both. - expect( - ctx.getWebSockets("cli").map((ws) => ws.deserializeAttachment()), - ).toEqual([ - expect.objectContaining({ role: "cli", connectionId: "cli-old" }), - expect.objectContaining({ role: "cli", connectionId: "cli-new" }), + expect(ctx.getWebSockets('cli').map(ws => ws.deserializeAttachment())).toEqual([ + expect.objectContaining({ role: 'cli', connectionId: 'cli-old' }), + expect.objectContaining({ role: 'cli', connectionId: 'cli-new' }), ]); const result = doInstance.getActiveSessions(); @@ -4253,10 +4114,10 @@ describe("UserConnectionDO", () => { // Exactly one row for the transferred session id, under the new owner. expect(result).toEqual([ { - id: "ses_transfer", - status: "busy", - title: "Transfer me", - connectionId: "cli-new", + id: 'ses_transfer', + status: 'busy', + title: 'Transfer me', + connectionId: 'cli-new', }, ]); }); @@ -4266,97 +4127,85 @@ describe("UserConnectionDO", () => { // Edge cases // ------------------------------------------------------------------------- - describe("edge cases", () => { - it("ignores non-JSON messages", () => { + describe('edge cases', () => { + it('ignores non-JSON messages', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); // Should not throw - doInstance.webSocketMessage(cliWs as never, "not-json"); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const warn = vi - .spyOn(console, "warn") - .mockImplementation(() => undefined); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); const malformed = '{"secret":"raw-secret-must-not-be-logged"'; doInstance.webSocketMessage(cliWs as never, malformed); - expect(warn).toHaveBeenCalledWith( - "Failed to parse WebSocket message as JSON", - { - role: "cli", - connectionId: "cli-1", - byteCount: new TextEncoder().encode(malformed).byteLength, - }, - ); - expect(JSON.stringify(warn.mock.calls)).not.toContain( - "raw-secret-must-not-be-logged", - ); + expect(warn).toHaveBeenCalledWith('Failed to parse WebSocket message as JSON', { + role: 'cli', + connectionId: 'cli-1', + byteCount: new TextEncoder().encode(malformed).byteLength, + }); + 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', () => { const { doInstance, mockCtx } = setup(); - const ws = createMockWs(["cli"], null); + const ws = createMockWs(['cli'], null); mockCtx.addSocket(ws); // Trigger ensureState first - doInstance.webSocketMessage( - ws as never, - JSON.stringify({ type: "heartbeat", sessions: [] }), - ); + doInstance.webSocketMessage(ws as never, JSON.stringify({ type: 'heartbeat', sessions: [] })); // Should not throw }); - it("ignores messages that fail Zod validation", () => { + it('ignores messages that fail Zod validation', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, []); // trigger ensureState // Invalid CLI message - const badMsg = JSON.stringify({ type: "invalid_type" }); + const badMsg = JSON.stringify({ type: 'invalid_type' }); doInstance.webSocketMessage(cliWs as never, badMsg); // Should not throw // Invalid web message - const webWs = addWebSocket(mockCtx, "web-1"); + const webWs = addWebSocket(mockCtx, 'web-1'); doInstance.webSocketMessage(webWs as never, badMsg); // Should not throw }); - it("logs malformed CLI message metadata without raw payload content", () => { + it('logs malformed CLI message metadata without raw payload content', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const warn = vi - .spyOn(console, "warn") - .mockImplementation(() => undefined); - const secret = "raw-secret-must-not-be-logged"; + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const secret = 'raw-secret-must-not-be-logged'; const malformed = JSON.stringify({ - type: "response", + type: 'response', id: 123, result: { secret }, }); doInstance.webSocketMessage(cliWs as never, malformed); - expect(warn).toHaveBeenCalledWith("CLI message parse failed", { - role: "cli", - connectionId: "cli-1", + expect(warn).toHaveBeenCalledWith('CLI message parse failed', { + role: 'cli', + connectionId: 'cli-1', byteCount: new TextEncoder().encode(malformed).byteLength, - issues: [{ path: ["id"], code: "invalid_type" }], + issues: [{ path: ['id'], code: 'invalid_type' }], }); expect(JSON.stringify(warn.mock.calls)).not.toContain(secret); }); - it("webSocketError triggers webSocketClose", async () => { + it('webSocketError triggers webSocketClose', async () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - const webWs = addWebSocket(mockCtx, "web-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); webWs.send.mockClear(); // Remove CLI so disconnect can clean up @@ -4365,20 +4214,16 @@ describe("UserConnectionDO", () => { // Should broadcast cli.disconnected const msgs = allSent(webWs); - expect( - msgs.some( - (m: Record) => m.event === "cli.disconnected", - ), - ).toBe(true); + 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', () => { const { doInstance, mockCtx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); sendHeartbeat(doInstance, cliWs, []); // Should not throw - sendCliResponse(doInstance, cliWs, { id: "nonexistent", result: "ok" }); + sendCliResponse(doInstance, cliWs, { id: 'nonexistent', result: 'ok' }); }); }); @@ -4386,15 +4231,14 @@ describe("UserConnectionDO", () => { // Session-ready push (delayed, Decision 9) // ------------------------------------------------------------------------- - describe("session-ready delayed push", () => { + describe('session-ready delayed push', () => { function setupWithIngestDO() { const mockCtx = createMockCtx(); const ctx = mockCtx.build(); const claimSessionReadyPush = vi.fn(async () => {}); sessionIngestMocks.getSessionIngestDO.mockReturnValue({ claimSessionReadyPush, - resetAttentionStatusOnCliDisconnect: - sessionIngestMocks.resetAttentionStatusOnCliDisconnect, + resetAttentionStatusOnCliDisconnect: sessionIngestMocks.resetAttentionStatusOnCliDisconnect, }); const doInstance = new UserConnectionDO(ctx as never, {} as never); return { doInstance, mockCtx, ctx, claimSessionReadyPush }; @@ -4403,86 +4247,83 @@ describe("UserConnectionDO", () => { function addCliSocketForUser( mockCtx: ReturnType, connectionId: string, - kiloUserId: string, + kiloUserId: string ): MockWS { const attachment = { - role: "cli" as const, + role: 'cli' as const, connectionId, sessions: [], kiloUserId, }; - const ws = createMockWs(["cli"], attachment); + const ws = createMockWs(['cli'], attachment); mockCtx.addSocket(ws); return ws; } - it("writes a pending readyPush entry on first sight and does not claim immediately", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + it('writes a pending readyPush entry on first sight and does not claim immediately', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); + vi.spyOn(Date, 'now').mockReturnValue(now); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_main")]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_main')]); await flushAsync(); expect(claimSessionReadyPush).not.toHaveBeenCalled(); - const entry = ctx.storage.store.get("readyPush:ses_main") as { + const entry = ctx.storage.store.get('readyPush:ses_main') as { kiloUserId: string; title: string; fireAt: number; attempts: number; }; expect(entry).toMatchObject({ - kiloUserId: "usr_1", - title: "Test", + kiloUserId: 'usr_1', + title: 'Test', fireAt: now + 5_000, attempts: 0, }); // Subsequent heartbeats for the same session must not re-schedule. - sendHeartbeat(doInstance, cliWs, [makeSession("ses_main")]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_main')]); await flushAsync(); expect(ctx.storage.put).toHaveBeenCalledTimes(1); }); - it("reconnect first-sight does not reset a pending readyPush fireAt or attempts", async () => { + it('reconnect first-sight does not reset a pending readyPush fireAt or attempts', async () => { const { doInstance, mockCtx, ctx } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); + vi.spyOn(Date, 'now').mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_reconnect")]); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_reconnect')]); await flushAsync(); const originalFireAt = now + 5_000; // Simulate a prior retry so attempts is non-zero. - await ctx.storage.put("readyPush:ses_reconnect", { - kiloUserId: "usr_1", - title: "Test", + await ctx.storage.put('readyPush:ses_reconnect', { + kiloUserId: 'usr_1', + title: 'Test', fireAt: originalFireAt, attempts: 2, }); const internal = doInstance as unknown as { readyPushFireAt: Map; }; - internal.readyPushFireAt.set("ses_reconnect", originalFireAt); + internal.readyPushFireAt.set('ses_reconnect', originalFireAt); // Disconnect clears sessionOwners; reconnect is a new "first sight". await disconnectCli(doInstance, cliWs); - expect(internal.readyPushFireAt.get("ses_reconnect")).toBe( - originalFireAt, - ); + expect(internal.readyPushFireAt.get('ses_reconnect')).toBe(originalFireAt); const later = now + 2_000; - vi.spyOn(Date, "now").mockReturnValue(later); + vi.spyOn(Date, 'now').mockReturnValue(later); ctx.storage.put.mockClear(); - const cliWs2 = addCliSocketForUser(mockCtx, "cli-2", "usr_1"); - sendHeartbeat(doInstance, cliWs2, [makeSession("ses_reconnect")]); + const cliWs2 = addCliSocketForUser(mockCtx, 'cli-2', 'usr_1'); + sendHeartbeat(doInstance, cliWs2, [makeSession('ses_reconnect')]); await flushAsync(); - const entry = ctx.storage.store.get("readyPush:ses_reconnect") as { + const entry = ctx.storage.store.get('readyPush:ses_reconnect') as { fireAt: number; attempts: number; }; @@ -4492,41 +4333,39 @@ describe("UserConnectionDO", () => { expect(ctx.storage.put).not.toHaveBeenCalled(); }); - it("never schedules for subagent sessions", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + it('never schedules for subagent sessions', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_main"), - makeSession("ses_sub", "busy", "Sub", "ses_main"), + makeSession('ses_main'), + makeSession('ses_sub', 'busy', 'Sub', 'ses_main'), ]); await flushAsync(); - expect(ctx.storage.store.has("readyPush:ses_main")).toBe(true); - expect(ctx.storage.store.has("readyPush:ses_sub")).toBe(false); + expect(ctx.storage.store.has('readyPush:ses_main')).toBe(true); + expect(ctx.storage.store.has('readyPush:ses_sub')).toBe(false); expect(claimSessionReadyPush).not.toHaveBeenCalled(); }); - it("does not schedule on sockets without a kiloUserId (legacy attachment)", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); - const cliWs = addCliSocket(mockCtx, "cli-1"); + it('does not schedule on sockets without a kiloUserId (legacy attachment)', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_main")]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_main')]); await flushAsync(); expect(ctx.storage.store.size).toBe(0); expect(claimSessionReadyPush).not.toHaveBeenCalled(); }); - it("arms the alarm for a readyPush even with no heartbeat candidates", async () => { + it('arms the alarm for a readyPush even with no heartbeat candidates', async () => { const { doInstance, mockCtx, ctx } = setupWithIngestDO(); // No CLI sockets — only a pending readyPush in KV, rebuilt via one-shot. const fireAt = Date.now() + 5_000; - await ctx.storage.put("readyPush:ses_orphan", { - kiloUserId: "usr_1", - title: "Orphan", + await ctx.storage.put('readyPush:ses_orphan', { + kiloUserId: 'usr_1', + title: 'Orphan', fireAt, attempts: 0, }); @@ -4542,10 +4381,10 @@ describe("UserConnectionDO", () => { ctx.storage.setAlarm.mockClear(); // Force a schedule by going through a heartbeat that schedules readyPush // then clear lastHeartbeat by disconnecting — still has readyPush mirror. - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_arm")]); + vi.spyOn(Date, 'now').mockReturnValue(now); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_arm')]); await flushAsync(); // Disconnect removes lastHeartbeatAt for this connection await disconnectCli(doInstance, cliWs); @@ -4561,27 +4400,24 @@ describe("UserConnectionDO", () => { expect(armedAt).toBe(now + 5_000); }); - it("arms immediately for overdue readyPush entries (never filtered by fireAt > now)", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); + it('arms immediately for overdue readyPush entries (never filtered by fireAt > now)', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); - await ctx.storage.put("readyPush:ses_overdue", { - kiloUserId: "usr_1", - title: "Old", + vi.spyOn(Date, 'now').mockReturnValue(now); + await ctx.storage.put('readyPush:ses_overdue', { + kiloUserId: 'usr_1', + title: 'Old', fireAt: now - 1_000, attempts: 0, }); // Seed mirror via first-sight path, then replace with overdue - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_overdue", "idle", "Old"), - ]); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_overdue', 'idle', 'Old')]); await flushAsync(); // Overwrite with overdue fireAt and mirror - await ctx.storage.put("readyPush:ses_overdue", { - kiloUserId: "usr_1", - title: "Old", + await ctx.storage.put('readyPush:ses_overdue', { + kiloUserId: 'usr_1', + title: 'Old', fireAt: now - 1_000, attempts: 0, }); @@ -4589,7 +4425,7 @@ describe("UserConnectionDO", () => { const internal = doInstance as unknown as { readyPushFireAt: Map; }; - internal.readyPushFireAt.set("ses_overdue", now - 1_000); + internal.readyPushFireAt.set('ses_overdue', now - 1_000); ctx.storage.setAlarm.mockClear(); claimSessionReadyPush.mockClear(); @@ -4599,13 +4435,13 @@ describe("UserConnectionDO", () => { // After fire, entry deleted; remaining schedule may or may not arm }); - it("one-shot mirror rebuild sets flag only after successful refresh", async () => { + it('one-shot mirror rebuild sets flag only after successful refresh', async () => { const { doInstance, mockCtx, ctx } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); - await ctx.storage.put("readyPush:ses_rebuild", { - kiloUserId: "usr_1", - title: "Rebuild", + vi.spyOn(Date, 'now').mockReturnValue(now); + await ctx.storage.put('readyPush:ses_rebuild', { + kiloUserId: 'usr_1', + title: 'Rebuild', fireAt: now + 2_000, attempts: 0, }); @@ -4619,204 +4455,190 @@ describe("UserConnectionDO", () => { expect(internal.readyPushRebuilt).toBe(false); // Kick schedule with empty mirror → async rebuild - const cliWs = addCliSocket(mockCtx, "cli-legacy"); // no kiloUserId → no new readyPush + const cliWs = addCliSocket(mockCtx, 'cli-legacy'); // no kiloUserId → no new readyPush sendHeartbeat(doInstance, cliWs, []); await flushAsync(); expect(internal.readyPushRebuilt).toBe(true); - expect(internal.readyPushFireAt.get("ses_rebuild")).toBe(now + 2_000); + expect(internal.readyPushFireAt.get('ses_rebuild')).toBe(now + 2_000); expect(ctx.storage.setAlarm).toHaveBeenCalled(); // Failed list leaves flag false internal.readyPushFireAt.clear(); internal.readyPushRebuilt = false; - ctx.storage.list.mockRejectedValueOnce(new Error("kv down")); + ctx.storage.list.mockRejectedValueOnce(new Error('kv down')); sendHeartbeat(doInstance, cliWs, []); await flushAsync(); expect(internal.readyPushRebuilt).toBe(false); }); - it("alarm passes changed heartbeat title when connectionSessions diverged, else undefined", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); + it('alarm passes changed heartbeat title when connectionSessions diverged, else undefined', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); + vi.spyOn(Date, 'now').mockReturnValue(now); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); const internal = doInstance as unknown as { readyPushFireAt: Map; }; // Establish ownership for both sessions first (first-sight schedules entries). sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_t", "idle", "Stored"), - makeSession("ses_same", "idle", "Same"), + makeSession('ses_t', 'idle', 'Stored'), + makeSession('ses_same', 'idle', 'Same'), ]); await flushAsync(); // Diverged title for ses_t; matching title for ses_same. Force both overdue. sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_t", "idle", "Generated"), - makeSession("ses_same", "idle", "Same"), + makeSession('ses_t', 'idle', 'Generated'), + makeSession('ses_same', 'idle', 'Same'), ]); - await ctx.storage.put("readyPush:ses_t", { - kiloUserId: "usr_1", - title: "Stored", + await ctx.storage.put('readyPush:ses_t', { + kiloUserId: 'usr_1', + title: 'Stored', fireAt: now - 1, attempts: 0, }); - await ctx.storage.put("readyPush:ses_same", { - kiloUserId: "usr_1", - title: "Same", + await ctx.storage.put('readyPush:ses_same', { + kiloUserId: 'usr_1', + title: 'Same', fireAt: now - 1, attempts: 0, }); - internal.readyPushFireAt.set("ses_t", now - 1); - internal.readyPushFireAt.set("ses_same", now - 1); + internal.readyPushFireAt.set('ses_t', now - 1); + internal.readyPushFireAt.set('ses_same', now - 1); claimSessionReadyPush.mockClear(); await doInstance.alarm(); - expect(claimSessionReadyPush).toHaveBeenCalledWith( - "usr_1", - "ses_t", - "Generated", - ); - expect(claimSessionReadyPush).toHaveBeenCalledWith( - "usr_1", - "ses_same", - undefined, - ); + expect(claimSessionReadyPush).toHaveBeenCalledWith('usr_1', 'ses_t', 'Generated'); + expect(claimSessionReadyPush).toHaveBeenCalledWith('usr_1', 'ses_same', undefined); }); - it("claim-then-delete: RPC rejection keeps the entry; drops after 3 attempts", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); + it('claim-then-delete: RPC rejection keeps the entry; drops after 3 attempts', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_retry")]); + vi.spyOn(Date, 'now').mockReturnValue(now); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_retry')]); await flushAsync(); 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", - title: "Test", + internal.readyPushFireAt.set('ses_retry', now - 1); + await ctx.storage.put('readyPush:ses_retry', { + kiloUserId: 'usr_1', + title: 'Test', fireAt: now - 1, attempts: 0, }); - claimSessionReadyPush.mockRejectedValueOnce(new Error("DO down")); + claimSessionReadyPush.mockRejectedValueOnce(new Error('DO down')); await doInstance.alarm(); - const after1 = ctx.storage.store.get("readyPush:ses_retry") as { + const after1 = ctx.storage.store.get('readyPush:ses_retry') as { attempts: number; fireAt: number; }; expect(after1.attempts).toBe(1); // Retry must back off so attempts span real time (not fireAt already ≤ now). expect(after1.fireAt).toBe(now + 5_000); - expect(internal.readyPushFireAt.get("ses_retry")).toBe(now + 5_000); + expect(internal.readyPushFireAt.get('ses_retry')).toBe(now + 5_000); // Advance past backoff so attempt 2 is due. - vi.spyOn(Date, "now").mockReturnValue(now + 5_000); - claimSessionReadyPush.mockRejectedValueOnce(new Error("DO down")); + vi.spyOn(Date, 'now').mockReturnValue(now + 5_000); + claimSessionReadyPush.mockRejectedValueOnce(new Error('DO down')); await doInstance.alarm(); - const after2 = ctx.storage.store.get("readyPush:ses_retry") as { + const after2 = ctx.storage.store.get('readyPush:ses_retry') as { attempts: number; fireAt: number; }; expect(after2.attempts).toBe(2); expect(after2.fireAt).toBe(now + 10_000); - vi.spyOn(Date, "now").mockReturnValue(now + 10_000); - claimSessionReadyPush.mockRejectedValueOnce(new Error("DO down")); + vi.spyOn(Date, 'now').mockReturnValue(now + 10_000); + claimSessionReadyPush.mockRejectedValueOnce(new Error('DO down')); await doInstance.alarm(); - expect(ctx.storage.store.has("readyPush:ses_retry")).toBe(false); - expect(internal.readyPushFireAt.has("ses_retry")).toBe(false); + expect(ctx.storage.store.has('readyPush:ses_retry')).toBe(false); + expect(internal.readyPushFireAt.has('ses_retry')).toBe(false); }); - it("claim rejection re-arms at now+backoff, not now", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); + it('claim rejection re-arms at now+backoff, not now', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_backoff")]); + vi.spyOn(Date, 'now').mockReturnValue(now); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_backoff')]); await flushAsync(); 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", - title: "Test", + internal.readyPushFireAt.set('ses_backoff', now - 1); + await ctx.storage.put('readyPush:ses_backoff', { + kiloUserId: 'usr_1', + title: 'Test', fireAt: now - 1, attempts: 0, }); - claimSessionReadyPush.mockRejectedValueOnce(new Error("transport blip")); + claimSessionReadyPush.mockRejectedValueOnce(new Error('transport blip')); await doInstance.alarm(); - const entry = ctx.storage.store.get("readyPush:ses_backoff") as { + const entry = ctx.storage.store.get('readyPush:ses_backoff') as { attempts: number; fireAt: number; }; expect(entry.attempts).toBe(1); expect(entry.fireAt).toBe(now + 5_000); - expect(internal.readyPushFireAt.get("ses_backoff")).toBe(now + 5_000); + expect(internal.readyPushFireAt.get('ses_backoff')).toBe(now + 5_000); // Same instant must not exhaust further attempts. claimSessionReadyPush.mockClear(); - claimSessionReadyPush.mockRejectedValue(new Error("still down")); + claimSessionReadyPush.mockRejectedValue(new Error('still down')); await doInstance.alarm(); expect(claimSessionReadyPush).not.toHaveBeenCalled(); expect( - (ctx.storage.store.get("readyPush:ses_backoff") as { attempts: number }) - .attempts, + (ctx.storage.store.get('readyPush:ses_backoff') as { attempts: number }).attempts ).toBe(1); }); - it("does not refire after a successful claim", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); + it('does not refire after a successful claim', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); - const cliWs = addCliSocketForUser(mockCtx, "cli-1", "usr_1"); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_ok")]); + vi.spyOn(Date, 'now').mockReturnValue(now); + const cliWs = addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_ok')]); await flushAsync(); 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", - title: "Test", + internal.readyPushFireAt.set('ses_ok', now - 1); + await ctx.storage.put('readyPush:ses_ok', { + kiloUserId: 'usr_1', + title: 'Test', fireAt: now - 1, attempts: 0, }); await doInstance.alarm(); expect(claimSessionReadyPush).toHaveBeenCalledTimes(1); - expect(ctx.storage.store.has("readyPush:ses_ok")).toBe(false); + expect(ctx.storage.store.has('readyPush:ses_ok')).toBe(false); claimSessionReadyPush.mockClear(); await doInstance.alarm(); expect(claimSessionReadyPush).not.toHaveBeenCalled(); }); - it("KV is source of truth for alarm even without the mirror", async () => { - const { doInstance, mockCtx, ctx, claimSessionReadyPush } = - setupWithIngestDO(); + it('KV is source of truth for alarm even without the mirror', async () => { + const { doInstance, mockCtx, ctx, claimSessionReadyPush } = setupWithIngestDO(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); + vi.spyOn(Date, 'now').mockReturnValue(now); // Ensure ensureState runs so DO is live - addCliSocketForUser(mockCtx, "cli-1", "usr_1"); - await ctx.storage.put("readyPush:ses_kv_only", { - kiloUserId: "usr_1", - title: "KV", + addCliSocketForUser(mockCtx, 'cli-1', 'usr_1'); + await ctx.storage.put('readyPush:ses_kv_only', { + kiloUserId: 'usr_1', + title: 'KV', fireAt: now - 1, attempts: 0, }); @@ -4827,41 +4649,37 @@ describe("UserConnectionDO", () => { internal.readyPushFireAt.clear(); await doInstance.alarm(); - expect(claimSessionReadyPush).toHaveBeenCalledWith( - "usr_1", - "ses_kv_only", - undefined, - ); - expect(ctx.storage.store.has("readyPush:ses_kv_only")).toBe(false); + expect(claimSessionReadyPush).toHaveBeenCalledWith('usr_1', 'ses_kv_only', undefined); + 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', () => { const { doInstance } = setupWithIngestDO(); const client = createMockWs(); const server = createMockWs(); vi.stubGlobal( - "WebSocketPair", + 'WebSocketPair', class { 0 = client; 1 = server; - }, + } ); vi.stubGlobal( - "Response", + 'Response', class { constructor(_body?: BodyInit | null, _init?: ResponseInit) {} - }, + } ); doInstance.fetch( - new Request("http://local/cli?connectionId=cli-1&kiloUserId=usr_1", { - headers: { Upgrade: "websocket" }, - }), + new Request('http://local/cli?connectionId=cli-1&kiloUserId=usr_1', { + headers: { Upgrade: 'websocket' }, + }) ); expect(server.deserializeAttachment()).toMatchObject({ - role: "cli", - kiloUserId: "usr_1", + role: 'cli', + kiloUserId: 'usr_1', }); }); }); @@ -4870,173 +4688,149 @@ describe("UserConnectionDO", () => { // notifySessionRenamed + rename catch-up (lazy KV per heartbeat) // ------------------------------------------------------------------------- - describe("notifySessionRenamed", () => { - it("delivers session.renamed to the owning CLI and always persists KV", async () => { + describe('notifySessionRenamed', () => { + it('delivers session.renamed to the owning CLI and always persists KV', async () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1", [makeSession("ses_r")]); + const cliWs = addCliSocket(mockCtx, 'cli-1', [makeSession('ses_r')]); // Establish ownership via heartbeat - sendHeartbeat(doInstance, cliWs, [makeSession("ses_r")]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_r')]); - const result = await doInstance.notifySessionRenamed( - "ses_r", - "Renamed Title", - ); + 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", + m => m.type === 'system' && m.event === 'session.renamed' ); expect(systemMsgs).toHaveLength(1); expect(systemMsgs[0].data).toEqual({ - sessionId: "ses_r", - title: "Renamed Title", + sessionId: 'ses_r', + title: 'Renamed Title', }); }); - it("returns delivered:false when no owner but still persists KV", async () => { + it('returns delivered:false when no owner but still persists KV', async () => { const { doInstance, ctx } = setup(); - const result = await doInstance.notifySessionRenamed( - "ses_missing", - "Offline Rename", - ); + const result = await doInstance.notifySessionRenamed('ses_missing', 'Offline Rename'); expect(result).toEqual({ delivered: false }); - expect(ctx.storage.store.get("rename:ses_missing")).toMatchObject({ - title: "Offline Rename", + expect(ctx.storage.store.get('rename:ses_missing')).toMatchObject({ + title: 'Offline Rename', }); }); - it("re-emits session.renamed on heartbeat title mismatch; deletes on match", async () => { + it('re-emits session.renamed on heartbeat title mismatch; deletes on match', async () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_catch", "idle", "Old Title"), - ]); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_catch', 'idle', 'Old Title')]); - await doInstance.notifySessionRenamed("ses_catch", "New Title"); + await doInstance.notifySessionRenamed('ses_catch', 'New Title'); cliWs.send.mockClear(); // Mismatch → re-emit - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_catch", "idle", "Old Title"), - ]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_catch', 'idle', 'Old Title')]); await flushAsync(); const reEmits = allSent(cliWs).filter( - (m) => m.type === "system" && m.event === "session.renamed", + m => m.type === 'system' && m.event === 'session.renamed' ); expect(reEmits).toHaveLength(1); expect(reEmits[0].data).toEqual({ - sessionId: "ses_catch", - title: "New Title", + sessionId: 'ses_catch', + title: 'New Title', }); - expect(ctx.storage.store.has("rename:ses_catch")).toBe(true); + expect(ctx.storage.store.has('rename:ses_catch')).toBe(true); // Match → delete entry cliWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_catch", "idle", "New Title"), - ]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_catch', 'idle', 'New Title')]); await flushAsync(); const afterMatch = allSent(cliWs).filter( - (m) => m.type === "system" && m.event === "session.renamed", + m => m.type === 'system' && m.event === 'session.renamed' ); expect(afterMatch).toHaveLength(0); - expect(ctx.storage.store.has("rename:ses_catch")).toBe(false); + expect(ctx.storage.store.has('rename:ses_catch')).toBe(false); }); - it("registers rename catch-up with waitUntil on heartbeat", async () => { + it('registers rename catch-up with waitUntil on heartbeat', async () => { const { doInstance, mockCtx, ctx } = setup(); - const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [makeSession("ses_wu")]); - await doInstance.notifySessionRenamed("ses_wu", "Catch Up Title"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_wu')]); + await doInstance.notifySessionRenamed('ses_wu', 'Catch Up Title'); const waitUntil = ctx.waitUntil as ReturnType; waitUntil.mockClear(); cliWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_wu", "idle", "Stale"), - ]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_wu', 'idle', 'Stale')]); expect(waitUntil).toHaveBeenCalled(); - const registered = waitUntil.mock.calls.map((c) => c[0]); + const registered = waitUntil.mock.calls.map(c => c[0]); expect( registered.some( - (p) => - p instanceof Promise || - typeof (p as PromiseLike)?.then === "function", - ), + p => p instanceof Promise || typeof (p as PromiseLike)?.then === 'function' + ) ).toBe(true); await flushAsync(); const reEmits = allSent(cliWs).filter( - (m) => m.type === "system" && m.event === "session.renamed", + m => m.type === 'system' && m.event === 'session.renamed' ); expect(reEmits).toHaveLength(1); expect(reEmits[0].data).toEqual({ - sessionId: "ses_wu", - title: "Catch Up Title", + sessionId: 'ses_wu', + title: 'Catch Up Title', }); }); - it("prunes rename entries older than TTL and does not re-emit them", async () => { + it('prunes rename entries older than TTL and does not re-emit them', async () => { const { doInstance, mockCtx, ctx } = setup(); const now = 1_700_000_000_000; const RENAME_ENTRY_TTL_MS = 7 * 24 * 60 * 60 * 1_000; - vi.spyOn(Date, "now").mockReturnValue(now); - const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_stale", "idle", "Old"), - ]); + vi.spyOn(Date, 'now').mockReturnValue(now); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_stale', 'idle', 'Old')]); - await ctx.storage.put("rename:ses_stale", { - title: "Never Applied", + await ctx.storage.put('rename:ses_stale', { + title: 'Never Applied', at: now - RENAME_ENTRY_TTL_MS - 1, }); cliWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_stale", "idle", "Old"), - ]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_stale', 'idle', 'Old')]); await flushAsync(); const reEmits = allSent(cliWs).filter( - (m) => m.type === "system" && m.event === "session.renamed", + m => m.type === 'system' && m.event === 'session.renamed' ); expect(reEmits).toHaveLength(0); - expect(ctx.storage.store.has("rename:ses_stale")).toBe(false); + expect(ctx.storage.store.has('rename:ses_stale')).toBe(false); }); - it("still re-emits a fresh rename entry within TTL", async () => { + it('still re-emits a fresh rename entry within TTL', async () => { const { doInstance, mockCtx, ctx } = setup(); const now = 1_700_000_000_000; - vi.spyOn(Date, "now").mockReturnValue(now); - const cliWs = addCliSocket(mockCtx, "cli-1"); - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_fresh", "idle", "Old"), - ]); + vi.spyOn(Date, 'now').mockReturnValue(now); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_fresh', 'idle', 'Old')]); - await ctx.storage.put("rename:ses_fresh", { - title: "Within TTL", + await ctx.storage.put('rename:ses_fresh', { + title: 'Within TTL', at: now - 60_000, }); cliWs.send.mockClear(); - sendHeartbeat(doInstance, cliWs, [ - makeSession("ses_fresh", "idle", "Old"), - ]); + sendHeartbeat(doInstance, cliWs, [makeSession('ses_fresh', 'idle', 'Old')]); await flushAsync(); const reEmits = allSent(cliWs).filter( - (m) => m.type === "system" && m.event === "session.renamed", + m => m.type === 'system' && m.event === 'session.renamed' ); expect(reEmits).toHaveLength(1); expect(reEmits[0].data).toEqual({ - sessionId: "ses_fresh", - title: "Within TTL", + sessionId: 'ses_fresh', + title: 'Within TTL', }); - expect(ctx.storage.store.has("rename:ses_fresh")).toBe(true); + expect(ctx.storage.store.has('rename:ses_fresh')).toBe(true); }); }); @@ -5050,32 +4844,30 @@ describe("UserConnectionDO", () => { // 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 () => { + 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"; + const correlationId = 'mut-1'; await ctx.storage.put(`pendingCommand/${correlationId}`, { - sessionId: "ses-a", - originalId: "original-req-1", - command: "send_message", + sessionId: 'ses-a', + originalId: 'original-req-1', + command: 'send_message', expectedOwnerConnectionId: undefined, - targetConnectionId: "cli-1", + targetConnectionId: 'cli-1', expiresAt: now + 35_000, - webConnectionId: "web-1", - state: "pending" as const, + 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"), - ]); + 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"); + const webWs = addWebSocket(mockCtx, 'web-1'); // Send the CLI response with the correlation id. await sendCliResponse(doInstance, cliWs, { @@ -5085,38 +4877,36 @@ describe("UserConnectionDO", () => { await flushAsync(); // The response must be routed to the web socket. - const responses = allSent(webWs).filter((m) => m.type === "response"); + const responses = allSent(webWs).filter(m => m.type === 'response'); expect(responses).toHaveLength(1); - expect(responses[0].id).toBe("original-req-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).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 () => { + 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"; + const correlationId = 'mut-2'; await ctx.storage.put(`pendingCommand/${correlationId}`, { - sessionId: "ses-b", - originalId: "original-req-2", - command: "send_message", - targetConnectionId: "cli-2", + sessionId: 'ses-b', + originalId: 'original-req-2', + command: 'send_message', + targetConnectionId: 'cli-2', expiresAt: now + 35_000, - webConnectionId: "web-gone", - state: "pending" as const, + 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"), - ]); + const cliWs = addCliSocket(mockCtx, 'cli-2'); + sendHeartbeat(doInstance, cliWs, [makeSession('ses-b', 'busy', 'Session B')]); // Send the CLI response. await sendCliResponse(doInstance, cliWs, { @@ -5128,130 +4918,118 @@ describe("UserConnectionDO", () => { // 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"); + expect((entry as Record).state).toBe('done'); }); - it("deduplicates a pending mutationId and returns an idempotent response for a non-catalog command", async () => { + 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"; + const mutationId = 'mut-dedup-pending'; await ctx.storage.put(`pendingCommand/${mutationId}`, { - sessionId: "ses-c", - originalId: "first-req", - command: "send_message", - targetConnectionId: "cli-3", + sessionId: 'ses-c', + originalId: 'first-req', + command: 'send_message', + targetConnectionId: 'cli-3', expiresAt: now + 35_000, - webConnectionId: "web-3", - state: "pending" as const, + 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"); + 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", + 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", - ); + 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", - ); + 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"); + 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 () => { + 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"; + const mutationId = 'mut-done-1'; await ctx.storage.put(`pendingCommand/${mutationId}`, { - sessionId: "ses-d", - originalId: "first-req-done", - command: "send_message", - targetConnectionId: "cli-4", + sessionId: 'ses-d', + originalId: 'first-req-done', + command: 'send_message', + targetConnectionId: 'cli-4', expiresAt: now + 35_000, - webConnectionId: "web-4", - state: "done" as const, + 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"); + 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", + 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", + 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", - ); + expect((entry as Record).originalId).toBe('retry-req-id'); // No command must reach the CLI. - const cliCommands = allSent(cliWs).filter((m) => m.type === "command"); + 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 () => { + 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"); + 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", + 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"); + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); expect(cliCommands).toHaveLength(1); expect(cliCommands[0].id).toBeTruthy(); expect(cliCommands[0].mutationId).toBeUndefined(); @@ -5269,31 +5047,27 @@ describe("UserConnectionDO", () => { }); await flushAsync(); - const responses = allSent(webWs).filter((m) => m.type === "response"); - expect(responses.some((r) => r.id === "no-mut-req" && r.result)).toBe( - 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 () => { + 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"); + 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", + id: 'disconnect-req', + command: 'send_message', + sessionId: 'ses-f', + connectionId: 'cli-6', }); await flushAsync(); - const cliCommands = allSent(cliWs).filter((m) => m.type === "command"); + const cliCommands = allSent(cliWs).filter(m => m.type === 'command'); expect(cliCommands).toHaveLength(1); const correlationId = cliCommands[0].id as string; @@ -5303,7 +5077,7 @@ describe("UserConnectionDO", () => { expect(entry).toBeDefined(); // Disconnect the web socket. - doInstance.webSocketClose(webWs as never, 1000, "", true); + doInstance.webSocketClose(webWs as never, 1000, '', true); await flushAsync(); // The durable entry must still exist. @@ -5311,31 +5085,29 @@ describe("UserConnectionDO", () => { expect(entry).toBeDefined(); }); - it("expires a durable pending entry and marks it done with COMMAND_EXPIRED_ERROR", async () => { + 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"; + 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", + sessionId: 'ses-g', + originalId: 'expired-req', + command: 'send_message', + targetConnectionId: 'cli-7', expiresAt: pastTime, - webConnectionId: "web-7", - state: "pending" as const, + 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"), - ]); + const cliWs = addCliSocket(mockCtx, 'cli-7'); + sendHeartbeat(doInstance2, cliWs, [makeSession('ses-g', 'busy', 'Session G')]); // Trigger alarm which calls expirePendingCommands. await doInstance2.alarm(); @@ -5344,46 +5116,44 @@ describe("UserConnectionDO", () => { // 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).state).toBe('done'); expect((entry as Record).error).toEqual({ - source: "relay", - code: "COMMAND_EXPIRED", - message: "Command expired", + source: 'relay', + code: 'COMMAND_EXPIRED', + message: 'Command expired', }); }); - it("counts durable entries toward the pending command cap (mutationId path)", async () => { + 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"); + 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", + sessionId: 'ses-h', originalId: `prefill-${i}`, - command: "send_message", - targetConnectionId: "cli-8", + command: 'send_message', + targetConnectionId: 'cli-8', expiresAt: now + 35_000, - webConnectionId: "web-8", - state: "pending" as const, + 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", + id: 'at-cap', + command: 'send_message', + sessionId: 'ses-h', + connectionId: 'cli-8', }); await flushAsync(); - const commandsAtCap = allSent(cliWs).filter((m) => m.type === "command"); + const commandsAtCap = allSent(cliWs).filter(m => m.type === 'command'); expect(commandsAtCap.length).toBeGreaterThanOrEqual(1); // The second command uses a fresh mutationId. The mutationId path @@ -5392,116 +5162,108 @@ describe("UserConnectionDO", () => { 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", + 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", - ); + 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", - ); + 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"); + 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 () => { + 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"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); sendCommand(doInstance, webWs, { - id: "err-1", - command: "list_models", - sessionId: "s1", + 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", + error: 'unknown command: list_models', }); await flushAsync(); // Live response carries the original string. const live = parseSent(webWs) as { error: unknown }; - expect(live.error).toBe("unknown command: list_models"); + 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", - ); + 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 () => { + 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"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); sendCommand(doInstance, webWs, { - id: "err-2", - command: "list_commands", - sessionId: "s1", - connectionId: "cli-1", + 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", + error: 'unknown command: list_commands', }); await flushAsync(); // 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.", + 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).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.", + 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 () => { + 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"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); sendCommand(doInstance, webWs, { - id: "dur-1", - command: "send_message", - sessionId: "s1", + id: 'dur-1', + command: 'send_message', + sessionId: 's1', }); const correlationId = getCorrelationId(cliWs); @@ -5517,7 +5279,7 @@ describe("UserConnectionDO", () => { // 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).state).toBe('done'); expect((entry as Record).result).toEqual({ durable: true, }); @@ -5526,32 +5288,28 @@ describe("UserConnectionDO", () => { // ------------------------------------------------------------------------- // Fix 5: rehydrated oversized catalog sends to live web socket // ------------------------------------------------------------------------- - it("sends the oversized-catalog error to a live web socket after rehydration", async () => { + 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"; + const correlationId = 'cat-rehydrated'; await ctx.storage.put(`pendingCommand/${correlationId}`, { - sessionId: "ses-cat", - originalId: "original-cat", - command: "list_models", + sessionId: 'ses-cat', + originalId: 'original-cat', + command: 'list_models', expectedOwnerConnectionId: undefined, - targetConnectionId: "cli-cat", + targetConnectionId: 'cli-cat', expiresAt: now + 35_000, - webConnectionId: "web-cat", - state: "pending" as const, + 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"); + 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, - ); + const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized, @@ -5559,21 +5317,21 @@ describe("UserConnectionDO", () => { await flushAsync(); // The live web socket must receive the oversized error. - const responses = allSent(webWs).filter((m) => m.type === "response"); + 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", + 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).state).toBe('done'); expect((entry as Record).error).toEqual({ - source: "relay", - code: "CATALOG_TOO_LARGE", - message: "Model catalog response is too large", + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', }); }); @@ -5582,14 +5340,14 @@ describe("UserConnectionDO", () => { // ------------------------------------------------------------------------- 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"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); sendCommand(doInstance, webWs, { - id: "cmd-disco", - command: "send_message", - sessionId: "s1", + id: 'cmd-disco', + command: 'send_message', + sessionId: 's1', }); webWs.send.mockClear(); @@ -5597,24 +5355,22 @@ describe("UserConnectionDO", () => { await disconnectCli(doInstance, cliWs); const msgs = allSent(webWs); - const errorResp = msgs.find( - (m) => m.type === "response" && m.id === "cmd-disco", - ); + 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"); + expect(errorResp!.error).toBe('CLI disconnected'); }); - it("stores the live CLI-disconnected error in the durable entry for retries", async () => { + 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"); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); - sendHeartbeat(doInstance, cliWs, [makeSession("s1")]); + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); sendCommand(doInstance, webWs, { - id: "cmd-dur-disco", - command: "send_message", - sessionId: "s1", + id: 'cmd-dur-disco', + command: 'send_message', + sessionId: 's1', }); const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); @@ -5625,8 +5381,8 @@ describe("UserConnectionDO", () => { // 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"); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).error).toBe('CLI disconnected'); }); }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 9abb803b913..9f916b86424 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -1,7 +1,7 @@ -import { DurableObject } from "cloudflare:workers"; +import { DurableObject } from 'cloudflare:workers'; -import type { Env } from "../env"; -import { getSessionIngestDO } from "./SessionIngestDO"; +import type { Env } from '../env'; +import { getSessionIngestDO } from './SessionIngestDO'; import { CLIOutboundMessageSchema, type CLIInboundMessage, @@ -10,7 +10,7 @@ import { SessionEventPayloadSchema, type WebInboundMessage, WebOutboundMessageSchema, -} from "../types/user-connection-protocol"; +} from '../types/user-connection-protocol'; type HeartbeatSession = { id: string; @@ -31,7 +31,7 @@ type ConnectionCapabilities = { attachments?: boolean }; type WSAttachment = | { - role: "cli"; + role: 'cli'; connectionId: string; sessions: HeartbeatSession[]; // Undefined means no protocolVersion has been reported yet — either the @@ -54,7 +54,7 @@ type WSAttachment = instance?: Instance; } | { - role: "web"; + role: 'web'; connectionId: string; subscribedSessions: string[]; replaced?: true; @@ -78,95 +78,91 @@ export const MAX_CATALOG_RESULT_BYTES = 512 * 1024; // 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([ - "send_message", - "interrupt", - "question_reply", - "question_reject", - "permission_respond", - "suggestion_accept", - "suggestion_dismiss", - "list_models", - "list_commands", - "send_command", - "create_session", - "exit_cli", + 'send_message', + 'interrupt', + 'question_reply', + 'question_reject', + 'permission_respond', + 'suggestion_accept', + 'suggestion_dismiss', + 'list_models', + 'list_commands', + 'send_command', + 'create_session', + 'exit_cli', ]); // In-flight dedupe and 512 KiB response cap apply to these catalog-style reads. -const CATALOG_DEDUPE_COMMANDS: ReadonlySet = new Set([ - "list_models", - "list_commands", -]); +const CATALOG_DEDUPE_COMMANDS: ReadonlySet = new Set(['list_models', 'list_commands']); // Operations that older CLIs reject with a precise "unknown command: " // string. Only these commands get mapped to a structured CLI_UPGRADE_REQUIRED // response; any other CLI error is preserved verbatim. const CLI_UPGRADE_REQUIRED_COMMANDS: ReadonlySet = new Set([ - "list_commands", - "send_command", - "create_session", - "exit_cli", + 'list_commands', + 'send_command', + 'create_session', + 'exit_cli', ]); const SESSION_OWNER_CHANGED_ERROR = { - source: "relay", - code: "SESSION_OWNER_CHANGED", - message: "Session owner changed", + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', }; const CATALOG_TOO_LARGE_ERROR = { - source: "relay", - code: "CATALOG_TOO_LARGE", - message: "Model catalog response is too large", + source: 'relay', + code: 'CATALOG_TOO_LARGE', + message: 'Model catalog response is too large', }; const CATALOG_REQUEST_PENDING_ERROR = { - source: "relay", - code: "CATALOG_REQUEST_PENDING", - message: "Model catalog request already pending", + source: 'relay', + code: 'CATALOG_REQUEST_PENDING', + message: 'Model catalog request already pending', }; const PENDING_COMMAND_LIMIT_ERROR = { - source: "relay", - code: "PENDING_COMMAND_LIMIT", - message: "Too many pending commands", + source: 'relay', + code: 'PENDING_COMMAND_LIMIT', + message: 'Too many pending commands', }; const COMMAND_EXPIRED_ERROR = { - source: "relay", - code: "COMMAND_EXPIRED", - message: "Command expired", + source: 'relay', + code: 'COMMAND_EXPIRED', + message: 'Command expired', }; const COMMAND_NOT_ALLOWED_ERROR = { - source: "relay", - code: "COMMAND_NOT_ALLOWED", - message: "Command is not allowed", + source: 'relay', + code: 'COMMAND_NOT_ALLOWED', + message: 'Command is not allowed', }; const INVALID_COMMAND_ERROR = { - source: "relay", - code: "INVALID_COMMAND", - message: "Invalid command", + source: 'relay', + code: 'INVALID_COMMAND', + message: 'Invalid command', }; const CLI_UPGRADE_REQUIRED_SLASH_ERROR = { - source: "relay", - code: "CLI_UPGRADE_REQUIRED", - message: - "Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.", + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', + message: 'Remote slash commands require a newer Kilo CLI. Update Kilo CLI and reconnect.', }; const CLI_UPGRADE_REQUIRED_CREATE_SESSION_ERROR = { - source: "relay", - code: "CLI_UPGRADE_REQUIRED", + source: 'relay', + code: 'CLI_UPGRADE_REQUIRED', message: - "Creating remote sessions from mobile requires a newer Kilo CLI. Update Kilo CLI and reconnect.", + 'Creating remote sessions from mobile requires a newer Kilo CLI. Update Kilo CLI and reconnect.', }; const CLI_COMMAND_ERROR = { - source: "cli", - message: "Command failed", + source: 'cli', + message: 'Command failed', }; type ReadyPushEntry = { @@ -192,14 +188,14 @@ type PendingCommandEntry = { // 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"; + 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 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; @@ -221,10 +217,7 @@ export class UserConnectionDO extends DurableObject { // Protocol version per CLI connection (from heartbeat); absent = legacy CLI private connectionProtocolVersion = new Map(); // Capabilities per CLI connection (from heartbeat); absent = legacy CLI - private connectionCapabilities = new Map< - string, - ConnectionCapabilities | undefined - >(); + private connectionCapabilities = new Map(); // Pending command responses: correlationId → originating web socket private pendingCommands = new Map< string, @@ -263,10 +256,9 @@ export class UserConnectionDO extends DurableObject { const attachment = ws.deserializeAttachment() as WSAttachment | null; if (!attachment) continue; - if (attachment.role === "cli") { + if (attachment.role === 'cli') { cliCount++; - const { connectionId, sessions, protocolVersion, capabilities } = - attachment; + const { connectionId, sessions, protocolVersion, capabilities } = attachment; this.connectionSessions.set(connectionId, sessions); this.connectionProtocolVersion.set(connectionId, protocolVersion); this.connectionCapabilities.set(connectionId, capabilities); @@ -289,7 +281,7 @@ export class UserConnectionDO extends DurableObject { } } - console.log("State reconstructed after hibernation", { + console.log('State reconstructed after hibernation', { cliSockets: cliCount, webSockets: webCount, sessions: sessionCount, @@ -302,48 +294,47 @@ export class UserConnectionDO extends DurableObject { } fetch(request: Request): Response { - const upgradeHeader = request.headers.get("Upgrade"); - if (upgradeHeader !== "websocket") { - return new Response("Expected WebSocket upgrade", { status: 426 }); + const upgradeHeader = request.headers.get('Upgrade'); + if (upgradeHeader !== 'websocket') { + return new Response('Expected WebSocket upgrade', { status: 426 }); } this.ensureState(); const url = new URL(request.url); - const role = url.pathname.endsWith("/cli") ? "cli" : "web"; + const role = url.pathname.endsWith('/cli') ? 'cli' : 'web'; const pair = new WebSocketPair(); const [client, server] = [pair[0], pair[1]]; - const connectionId = - url.searchParams.get("connectionId") ?? crypto.randomUUID(); + const connectionId = url.searchParams.get('connectionId') ?? crypto.randomUUID(); - if (role === "cli") { + if (role === 'cli') { // Close any stale socket from a previous connection with the same ID (CLI reconnect) const reconnect = this.closeStaleSocket(connectionId); - const kiloUserId = url.searchParams.get("kiloUserId") ?? undefined; + const kiloUserId = url.searchParams.get('kiloUserId') ?? undefined; const attachment: WSAttachment = { - role: "cli", + role: 'cli', connectionId, sessions: [], kiloUserId, }; - this.ctx.acceptWebSocket(server, ["cli"]); + this.ctx.acceptWebSocket(server, ['cli']); server.serializeAttachment(attachment); const now = Date.now(); this.lastHeartbeatAt.set(connectionId, now); this.scheduleNextAlarm(now); - console.log("CLI socket connected", { + console.log('CLI socket connected', { connectionId, reconnect, - totalCliSockets: this.ctx.getWebSockets("cli").length, + totalCliSockets: this.ctx.getWebSockets('cli').length, }); if (!reconnect) { this.broadcastToWeb({ - type: "system", - event: "cli.connected", + type: 'system', + event: 'cli.connected', data: { connectionId }, }); } @@ -351,24 +342,24 @@ export class UserConnectionDO extends DurableObject { this.replaceWebSocket(connectionId); const attachment: WSAttachment = { - role: "web", + role: 'web', connectionId, subscribedSessions: [], }; - this.ctx.acceptWebSocket(server, ["web"]); + this.ctx.acceptWebSocket(server, ['web']); server.serializeAttachment(attachment); const sessions = this.aggregateSessions(); - console.log("Web socket connected", { + console.log('Web socket connected', { connectionId, - totalWebSockets: this.ctx.getWebSockets("web").length, + totalWebSockets: this.ctx.getWebSockets('web').length, activeSessions: sessions.length, }); this.sendToWeb(server, { - type: "system", - event: "sessions.list", + type: 'system', + event: 'sessions.list', data: { sessions }, }); } @@ -376,27 +367,22 @@ export class UserConnectionDO extends DurableObject { return new Response(null, { status: 101, webSocket: client }); } - webSocketMessage( - ws: WebSocket, - message: string | ArrayBuffer, - ): void { + webSocketMessage(ws: WebSocket, message: string | ArrayBuffer): void { this.ensureState(); const attachment = ws.deserializeAttachment() as WSAttachment | null; if (!attachment) { - console.warn("WebSocket message from socket with no attachment"); + console.warn('WebSocket message from socket with no attachment'); return; } - const raw = - typeof message === "string" ? message : new TextDecoder().decode(message); - const binaryByteCount = - typeof message === "string" ? undefined : message.byteLength; + const raw = typeof message === 'string' ? message : new TextDecoder().decode(message); + const binaryByteCount = typeof message === 'string' ? undefined : message.byteLength; let parsed: unknown; try { parsed = JSON.parse(raw); } catch { - console.warn("Failed to parse WebSocket message as JSON", { + console.warn('Failed to parse WebSocket message as JSON', { role: attachment.role, connectionId: attachment.connectionId, byteCount: binaryByteCount ?? new TextEncoder().encode(raw).byteLength, @@ -404,14 +390,8 @@ export class UserConnectionDO extends DurableObject { return; } - if (attachment.role === "cli") { - this.handleCliMessage( - ws, - attachment, - parsed, - raw, - binaryByteCount, - ); + if (attachment.role === 'cli') { + this.handleCliMessage(ws, attachment, parsed, raw, binaryByteCount); } else if (!attachment.replaced) { this.handleWebMessage(ws, attachment, parsed); } @@ -421,14 +401,14 @@ export class UserConnectionDO extends DurableObject { ws: WebSocket, _code: number, _reason: string, - _wasClean: boolean, + _wasClean: boolean ): void | Promise { this.ensureState(); const attachment = ws.deserializeAttachment() as WSAttachment | null; if (!attachment) return; - if (attachment.role === "cli") { + if (attachment.role === 'cli') { // Await attention resets so cli.disconnected is not broadcast until the // stored status write has committed (mobile history refetch races otherwise). return this.handleCliDisconnect(ws, attachment); @@ -438,11 +418,11 @@ export class UserConnectionDO extends DurableObject { webSocketError(ws: WebSocket): void | Promise { const attachment = ws.deserializeAttachment() as WSAttachment | null; - console.error("WebSocket error", { - role: attachment?.role ?? "unknown", - connectionId: attachment?.connectionId ?? "unknown", + console.error('WebSocket error', { + role: attachment?.role ?? 'unknown', + connectionId: attachment?.connectionId ?? 'unknown', }); - return this.webSocketClose(ws, 0, "", false); + return this.webSocketClose(ws, 0, '', false); } async alarm(): Promise { @@ -461,13 +441,13 @@ export class UserConnectionDO extends DurableObject { for (const connectionId of staleConnectionIds) { // Find and close the stale CLI WebSocket - for (const ws of this.ctx.getWebSockets("cli")) { + 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)", { + if (att?.role === 'cli' && att.connectionId === connectionId) { + console.log('Closing stale CLI connection (heartbeat timeout)', { connectionId, }); - ws.close(4408, "heartbeat timeout"); + ws.close(4408, 'heartbeat timeout'); break; } } @@ -485,18 +465,18 @@ export class UserConnectionDO extends DurableObject { private handleCliMessage( ws: WebSocket, - attachment: WSAttachment & { role: "cli" }, + attachment: WSAttachment & { role: 'cli' }, parsed: unknown, raw: string, - binaryByteCount: number | undefined, + binaryByteCount: number | undefined ): void { const result = CLIOutboundMessageSchema.safeParse(parsed); if (!result.success) { - console.warn("CLI message parse failed", { - role: "cli", + console.warn('CLI message parse failed', { + role: 'cli', connectionId: attachment.connectionId, byteCount: binaryByteCount ?? new TextEncoder().encode(raw).byteLength, - issues: result.error.issues.map((issue) => ({ + issues: result.error.issues.map(issue => ({ path: issue.path, code: issue.code, })), @@ -506,42 +486,35 @@ export class UserConnectionDO extends DurableObject { const msg = result.data; switch (msg.type) { - case "heartbeat": + case 'heartbeat': this.handleHeartbeat( ws, attachment, msg.sessions, msg.protocolVersion, msg.capabilities, - msg.instance, + msg.instance ); break; - case "event": - this.handleCliEvent( - msg.sessionId, - msg.parentSessionId, - msg.event, - msg.data, - ); + case 'event': + this.handleCliEvent(msg.sessionId, msg.parentSessionId, msg.event, msg.data); break; - case "response": + case 'response': // Extend the DO lifetime for the async durable write (fix: reply // must persist before hibernation, but the surrounding handler stays // synchronous for existing callers). - this.ctx.waitUntil( - this.handleCliResponse(ws, msg.id, msg.result, msg.error), - ); + this.ctx.waitUntil(this.handleCliResponse(ws, msg.id, msg.result, msg.error)); break; } } private handleHeartbeat( ws: WebSocket, - attachment: WSAttachment & { role: "cli" }, + attachment: WSAttachment & { role: 'cli' }, sessions: HeartbeatSession[], protocolVersion: string | undefined, capabilities: ConnectionCapabilities | undefined, - instance: Instance | undefined, + instance: Instance | undefined ): void { const { connectionId } = attachment; const now = Date.now(); @@ -552,12 +525,9 @@ export class UserConnectionDO extends DurableObject { // Remove sessions this connection previously owned but no longer reports const previousSessions = this.connectionSessions.get(connectionId) ?? []; - const currentIds = new Set(sessions.map((s) => s.id)); + const currentIds = new Set(sessions.map(s => s.id)); for (const prev of previousSessions) { - if ( - !currentIds.has(prev.id) && - this.sessionOwners.get(prev.id) === connectionId - ) { + if (!currentIds.has(prev.id) && this.sessionOwners.get(prev.id) === connectionId) { this.sessionOwners.delete(prev.id); this.failPendingCommandsForOwnerChange(prev.id, undefined); } @@ -574,11 +544,7 @@ export class UserConnectionDO extends DurableObject { // push (Decision 9). The durable claim in SessionIngestDO makes reconnect // re-sights no-ops once the delayed claim lands. if (!previousOwner && !session.parentSessionId && attachment.kiloUserId) { - this.scheduleSessionReadyPush( - attachment.kiloUserId, - session.id, - session.title, - ); + this.scheduleSessionReadyPush(attachment.kiloUserId, session.id, session.title); } this.sessionOwners.set(session.id, connectionId); } @@ -589,26 +555,23 @@ export class UserConnectionDO extends DurableObject { // waitUntil so the DO does not hibernate before catch-up finishes. this.ctx.waitUntil( this.catchUpPendingRenames(ws, sessions).catch((error: unknown) => { - console.error("Failed to catch up pending renames (non-fatal)", { + console.error('Failed to catch up pending renames (non-fatal)', { error: error instanceof Error ? error.message : String(error), }); - }), + }) ); // Replay existing subscriptions for sessions newly owned by this CLI - const previousIds = new Set(previousSessions.map((s) => s.id)); + const previousIds = new Set(previousSessions.map(s => s.id)); for (const session of sessions) { - if ( - !previousIds.has(session.id) && - this.webSubscriptions.has(session.id) - ) { - this.sendToCli(ws, { type: "subscribe", sessionId: session.id }); + if (!previousIds.has(session.id) && this.webSubscriptions.has(session.id)) { + this.sendToCli(ws, { type: 'subscribe', sessionId: session.id }); } } // Persist to attachment for hibernation recovery const updatedAttachment: WSAttachment = { - role: "cli", + role: 'cli', connectionId, sessions, protocolVersion, @@ -622,20 +585,20 @@ export class UserConnectionDO extends DurableObject { // and non-subscribers both receive it: a removed session id is detectable // from its absence in the payload, so no subscriber special-case is needed. this.broadcastToWeb({ - type: "system", - event: "sessions.heartbeat", + type: 'system', + event: 'sessions.heartbeat', data: { connectionId, protocolVersion, capabilities, - sessions: sessions.map((session) => ({ + sessions: sessions.map(session => ({ ...session, ...(capabilities ? { capabilities } : {}), })), }, }); - this.sendToCli(ws, { type: "heartbeat_ack" }); + this.sendToCli(ws, { type: 'heartbeat_ack' }); } /** @@ -644,11 +607,7 @@ export class UserConnectionDO extends DurableObject { * No-op when a pending entry already exists (mirror or KV) so reconnect * first-sights do not slide fireAt or reset attempts. */ - private scheduleSessionReadyPush( - kiloUserId: string, - sessionId: string, - title: string, - ): void { + private scheduleSessionReadyPush(kiloUserId: string, sessionId: string, title: string): void { if (this.readyPushFireAt.has(sessionId)) return; const key = `${READY_PUSH_KEY_PREFIX}${sessionId}`; @@ -670,11 +629,11 @@ export class UserConnectionDO extends DurableObject { this.readyPushFireAt.set(sessionId, fireAt); this.scheduleNextAlarm(Date.now()); })().catch((error: unknown) => { - console.error("Failed to schedule session-ready push (non-fatal)", { + console.error('Failed to schedule session-ready push (non-fatal)', { sessionId, error: error instanceof Error ? error.message : String(error), }); - }), + }) ); } @@ -686,7 +645,7 @@ export class UserConnectionDO extends DurableObject { private async claimSessionReadyPush( kiloUserId: string, sessionId: string, - title?: string, + title?: string ): Promise { const stub = getSessionIngestDO(this.env, { kiloUserId, sessionId }); await stub.claimSessionReadyPush(kiloUserId, sessionId, title); @@ -699,14 +658,14 @@ export class UserConnectionDO extends DurableObject { prefix: READY_PUSH_KEY_PREFIX, }); } catch (error: unknown) { - console.error("Failed to list readyPush entries", { + console.error('Failed to list readyPush entries', { error: error instanceof Error ? error.message : String(error), }); return; } for (const [key, entry] of pending) { - if (!entry || typeof entry.fireAt !== "number") continue; + if (!entry || typeof entry.fireAt !== 'number') continue; if (entry.fireAt > now) continue; const sessionId = key.slice(READY_PUSH_KEY_PREFIX.length); @@ -716,9 +675,7 @@ export class UserConnectionDO extends DurableObject { // else undefined so notifications falls back to the live DB title. const heartbeatTitle = this.findHeartbeatTitle(sessionId); const freshest = - heartbeatTitle !== undefined && heartbeatTitle !== entry.title - ? heartbeatTitle - : undefined; + heartbeatTitle !== undefined && heartbeatTitle !== entry.title ? heartbeatTitle : undefined; try { await this.claimSessionReadyPush(entry.kiloUserId, sessionId, freshest); @@ -729,7 +686,7 @@ export class UserConnectionDO extends DurableObject { if (attempts >= READY_PUSH_MAX_ATTEMPTS) { await this.ctx.storage.delete(key); this.readyPushFireAt.delete(sessionId); - console.error("Dropping session-ready push after max attempts", { + console.error('Dropping session-ready push after max attempts', { sessionId, attempts, error: error instanceof Error ? error.message : String(error), @@ -738,7 +695,7 @@ export class UserConnectionDO extends DurableObject { const fireAt = now + READY_PUSH_RETRY_BACKOFF_MS; await this.ctx.storage.put(key, { ...entry, attempts, fireAt }); this.readyPushFireAt.set(sessionId, fireAt); - console.error("Session-ready push claim failed; will retry", { + console.error('Session-ready push claim failed; will retry', { sessionId, attempts, fireAt, @@ -758,10 +715,7 @@ export class UserConnectionDO extends DurableObject { return undefined; } - private async catchUpPendingRenames( - ws: WebSocket, - sessions: HeartbeatSession[], - ): Promise { + private async catchUpPendingRenames(ws: WebSocket, sessions: HeartbeatSession[]): Promise { const now = Date.now(); for (const session of sessions) { const key = `${RENAME_KEY_PREFIX}${session.id}`; @@ -769,7 +723,7 @@ export class UserConnectionDO extends DurableObject { try { entry = await this.ctx.storage.get(key); } catch (error: unknown) { - console.error("Failed to read rename catch-up entry", { + console.error('Failed to read rename catch-up entry', { sessionId: session.id, error: error instanceof Error ? error.message : String(error), }); @@ -778,14 +732,11 @@ export class UserConnectionDO extends DurableObject { if (!entry) continue; // Prune entries past TTL so finished/ignored renames do not re-emit forever. - if ( - typeof entry.at === "number" && - now - entry.at > RENAME_ENTRY_TTL_MS - ) { + if (typeof entry.at === 'number' && now - entry.at > RENAME_ENTRY_TTL_MS) { try { await this.ctx.storage.delete(key); } catch (error: unknown) { - console.error("Failed to delete expired rename entry", { + console.error('Failed to delete expired rename entry', { sessionId: session.id, error: error instanceof Error ? error.message : String(error), }); @@ -797,7 +748,7 @@ export class UserConnectionDO extends DurableObject { try { await this.ctx.storage.delete(key); } catch (error: unknown) { - console.error("Failed to delete matched rename entry", { + console.error('Failed to delete matched rename entry', { sessionId: session.id, error: error instanceof Error ? error.message : String(error), }); @@ -806,8 +757,8 @@ export class UserConnectionDO extends DurableObject { } this.sendToCli(ws, { - type: "system", - event: "session.renamed", + type: 'system', + event: 'session.renamed', data: { sessionId: session.id, title: entry.title }, }); } @@ -817,12 +768,10 @@ export class UserConnectionDO extends DurableObject { sessionId: string, parentSessionId: string | undefined, event: string, - data: unknown, + data: unknown ): void { const childSubs = this.webSubscriptions.get(sessionId); - const parentSubs = parentSessionId - ? this.webSubscriptions.get(parentSessionId) - : undefined; + const parentSubs = parentSessionId ? this.webSubscriptions.get(parentSessionId) : undefined; if (!childSubs && !parentSubs) return; const merged = new Set(); @@ -831,7 +780,7 @@ export class UserConnectionDO extends DurableObject { if (merged.size === 0) return; const msg: WebInboundMessage = { - type: "event", + type: 'event', sessionId, ...(parentSessionId ? { parentSessionId } : {}), event, @@ -846,7 +795,7 @@ export class UserConnectionDO extends DurableObject { respondingWs: WebSocket, id: string, result: unknown, - error: unknown, + error: unknown ): Promise { let entry = this.pendingCommands.get(id); let rehydrated = false; @@ -854,13 +803,12 @@ export class UserConnectionDO extends DurableObject { // Step 24: on in-memory miss, try to load the durable entry (D8 case 1). if (!entry) { const durable = await this.getDurablePendingCommand(id); - if (!durable || durable.state !== "pending") return; + if (!durable || durable.state !== 'pending') return; // Validate the responding CLI by attachment connectionId. - const respondingAttachment = - respondingWs.deserializeAttachment() as WSAttachment | null; + const respondingAttachment = respondingWs.deserializeAttachment() as WSAttachment | null; if ( - respondingAttachment?.role !== "cli" || + respondingAttachment?.role !== 'cli' || respondingAttachment.connectionId !== durable.targetConnectionId ) { return; @@ -896,7 +844,7 @@ export class UserConnectionDO extends DurableObject { if (!rehydrated) { this.pendingCommands.delete(id); this.sendToWeb(entry.ws, { - type: "response", + type: 'response', id: entry.originalId, error: CATALOG_TOO_LARGE_ERROR, }); @@ -905,14 +853,13 @@ export class UserConnectionDO extends DurableObject { // Persist the terminal outcome durably. const catWebConnectionId = rehydrated ? (await this.getDurablePendingCommand(id))!.webConnectionId - : (entry.ws.deserializeAttachment() as WSAttachment | null)?.role === - "web" + : (entry.ws.deserializeAttachment() as WSAttachment | null)?.role === 'web' ? ( entry.ws.deserializeAttachment() as WSAttachment & { - role: "web"; + role: 'web'; } ).connectionId - : "unknown"; + : 'unknown'; await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { sessionId: entry.sessionId, originalId: entry.originalId, @@ -921,17 +868,16 @@ export class UserConnectionDO extends DurableObject { targetConnectionId: entry.targetConnectionId, expiresAt: entry.expiresAt, webConnectionId: catWebConnectionId, - state: "done", + state: 'done', error: CATALOG_TOO_LARGE_ERROR, } satisfies PendingCommandEntry); // For a rehydrated entry, re-resolve and send now. if (rehydrated) { - const targetWeb = - this.findWebByConnectionId(catWebConnectionId); + const targetWeb = this.findWebByConnectionId(catWebConnectionId); if (targetWeb) { this.sendToWeb(targetWeb, { - type: "response", + type: 'response', id: entry.originalId, error: CATALOG_TOO_LARGE_ERROR, }); @@ -949,28 +895,23 @@ export class UserConnectionDO extends DurableObject { let stringError: string | null = null; let sanitizeAsFailed = false; - if ( - typeof error === "string" && - CLI_UPGRADE_REQUIRED_COMMANDS.has(entry.command) - ) { + if (typeof error === 'string' && CLI_UPGRADE_REQUIRED_COMMANDS.has(entry.command)) { if (error === `unknown command: ${entry.command}`) { structuredError = - entry.command === "create_session" + entry.command === 'create_session' ? CLI_UPGRADE_REQUIRED_CREATE_SESSION_ERROR : CLI_UPGRADE_REQUIRED_SLASH_ERROR; } else { stringError = error; } - } else if (typeof error === "string") { + } else if (typeof error === 'string') { stringError = error; } else if (error !== undefined) { sanitizeAsFailed = true; } const terminalError = - structuredError ?? - stringError ?? - (sanitizeAsFailed ? CLI_COMMAND_ERROR : undefined); + structuredError ?? stringError ?? (sanitizeAsFailed ? CLI_COMMAND_ERROR : undefined); // Step 26a: write terminal outcome durably, always. // For non-rehydrated entries, send the live response synchronously @@ -979,7 +920,7 @@ export class UserConnectionDO extends DurableObject { if (!rehydrated) { this.pendingCommands.delete(id); this.sendToWeb(entry.ws, { - type: "response", + type: 'response', id: entry.originalId, ...(result !== undefined ? { result } : {}), ...(structuredError !== null @@ -997,11 +938,9 @@ export class UserConnectionDO extends DurableObject { // delivered for the non-rehydrated case above. const webConnectionId = rehydrated ? (await this.getDurablePendingCommand(id))?.webConnectionId - : (entry.ws.deserializeAttachment() as WSAttachment | null)?.role === - "web" - ? (entry.ws.deserializeAttachment() as WSAttachment & { role: "web" }) - .connectionId - : "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, @@ -1009,8 +948,8 @@ export class UserConnectionDO extends DurableObject { expectedOwnerConnectionId: entry.expectedOwnerConnectionId, targetConnectionId: entry.targetConnectionId, expiresAt: entry.expiresAt, - webConnectionId: webConnectionId ?? "unknown", - state: "done", + webConnectionId: webConnectionId ?? 'unknown', + state: 'done', ...(result !== undefined ? { result } : {}), ...(terminalError !== undefined ? { error: terminalError } : {}), } satisfies PendingCommandEntry); @@ -1019,12 +958,10 @@ export class UserConnectionDO extends DurableObject { // id and send the response now (the non-rehydrated path already sent // it synchronously above). if (rehydrated) { - const targetWebWs = this.findWebByConnectionId( - webConnectionId ?? "unknown", - ); + const targetWebWs = this.findWebByConnectionId(webConnectionId ?? 'unknown'); if (targetWebWs) { this.sendToWeb(targetWebWs, { - type: "response", + type: 'response', id: entry.originalId, ...(result !== undefined ? { result } : {}), ...(structuredError !== null @@ -1045,39 +982,39 @@ export class UserConnectionDO extends DurableObject { private handleWebMessage( ws: WebSocket, - attachment: WSAttachment & { role: "web" }, - parsed: unknown, + attachment: WSAttachment & { role: 'web' }, + parsed: unknown ): void { const result = WebOutboundMessageSchema.safeParse(parsed); if (!result.success) { - console.warn("Invalid web message", { + console.warn('Invalid web message', { connectionId: attachment.connectionId, - errors: result.error.issues.map((i) => i.message), + errors: result.error.issues.map(i => i.message), }); return; } const msg = result.data; switch (msg.type) { - case "subscribe": + case 'subscribe': this.handleWebSubscribe(ws, attachment, msg.sessionId); break; - case "unsubscribe": + case 'unsubscribe': this.handleWebUnsubscribe(ws, attachment, msg.sessionId); break; - case "command": + case 'command': this.handleWebCommand(ws, msg); break; - case "ping": - this.sendToWeb(ws, { type: "pong", nonce: msg.nonce }); + case 'ping': + this.sendToWeb(ws, { type: 'pong', nonce: msg.nonce }); break; } } private handleWebSubscribe( ws: WebSocket, - attachment: WSAttachment & { role: "web" }, - sessionId: string, + attachment: WSAttachment & { role: 'web' }, + sessionId: string ): void { let subs = this.webSubscriptions.get(sessionId); if (!subs) { @@ -1093,8 +1030,8 @@ export class UserConnectionDO extends DurableObject { } this.sendToWeb(ws, { - type: "system", - event: "sessions.list", + type: 'system', + event: 'sessions.list', data: { sessions: this.aggregateSessions() }, }); @@ -1104,18 +1041,18 @@ export class UserConnectionDO extends DurableObject { // so it wasn't reported in the most recent heartbeat. const cliWs = this.findCliForSession(sessionId); if (cliWs) { - this.sendToCli(cliWs, { type: "subscribe", sessionId }); + this.sendToCli(cliWs, { type: 'subscribe', sessionId }); } else { - for (const ws of this.ctx.getWebSockets("cli")) { - this.sendToCli(ws, { type: "subscribe", sessionId }); + for (const ws of this.ctx.getWebSockets('cli')) { + this.sendToCli(ws, { type: 'subscribe', sessionId }); } } } private handleWebUnsubscribe( ws: WebSocket, - attachment: WSAttachment & { role: "web" }, - sessionId: string, + attachment: WSAttachment & { role: 'web' }, + sessionId: string ): void { const subs = this.webSubscriptions.get(sessionId); if (subs) { @@ -1126,7 +1063,7 @@ export class UserConnectionDO extends DurableObject { this.webSubscriptions.delete(sessionId); const cliWs = this.findCliForSession(sessionId); if (cliWs) { - this.sendToCli(cliWs, { type: "unsubscribe", sessionId }); + this.sendToCli(cliWs, { type: 'unsubscribe', sessionId }); } } } @@ -1148,7 +1085,7 @@ export class UserConnectionDO extends DurableObject { connectionId?: string; data?: unknown; mutationId?: string; - }, + } ): void { const now = Date.now(); this.expirePendingCommands(now); @@ -1157,7 +1094,7 @@ export class UserConnectionDO extends DurableObject { // ownership, allocate a pending slot, or forward to the CLI. if (!ALLOWED_VIEWER_COMMANDS.has(msg.command)) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: COMMAND_NOT_ALLOWED_ERROR, }); @@ -1165,17 +1102,17 @@ export class UserConnectionDO extends DurableObject { } if ( - msg.command === "exit_cli" && + msg.command === 'exit_cli' && (!msg.sessionId || - typeof msg.data !== "object" || + typeof msg.data !== 'object' || msg.data === null || Array.isArray(msg.data) || Object.keys(msg.data).length !== 1 || - !Object.hasOwn(msg.data, "protocolVersion") || - Reflect.get(msg.data, "protocolVersion") !== 1) + !Object.hasOwn(msg.data, 'protocolVersion') || + Reflect.get(msg.data, 'protocolVersion') !== 1) ) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: INVALID_COMMAND_ERROR, }); @@ -1187,12 +1124,9 @@ export class UserConnectionDO extends DurableObject { if (msg.sessionId && msg.connectionId) { targetCli = this.findCliByConnectionId(msg.connectionId); - if ( - this.sessionOwners.get(msg.sessionId) !== msg.connectionId || - !targetCli - ) { + if (this.sessionOwners.get(msg.sessionId) !== msg.connectionId || !targetCli) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: SESSION_OWNER_CHANGED_ERROR, }); @@ -1204,22 +1138,21 @@ export class UserConnectionDO extends DurableObject { targetCli = this.findCliForSession(msg.sessionId); } else { // Fall back to first available CLI - const cliSockets = this.ctx.getWebSockets("cli"); + const cliSockets = this.ctx.getWebSockets('cli'); targetCli = cliSockets[0]; } if (!targetCli) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, - error: "Session owner not found", + error: 'Session owner not found', }); return; } - const targetAttachment = - targetCli.deserializeAttachment() as WSAttachment | null; - if (targetAttachment?.role !== "cli") return; + const targetAttachment = targetCli.deserializeAttachment() as WSAttachment | null; + if (targetAttachment?.role !== 'cli') return; const expectedOwnerConnectionId = msg.sessionId && msg.connectionId ? msg.connectionId : undefined; const targetConnectionId = targetAttachment.connectionId; @@ -1227,24 +1160,24 @@ export class UserConnectionDO extends DurableObject { // 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?.role === 'web' && webAttachment.connectionId ? webAttachment.connectionId - : "unknown"; + : '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) => + entry => entry.ws === ws && entry.command === msg.command && entry.sessionId === msg.sessionId && - entry.targetConnectionId === targetConnectionId, + entry.targetConnectionId === targetConnectionId ) ) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: CATALOG_REQUEST_PENDING_ERROR, }); @@ -1259,18 +1192,18 @@ export class UserConnectionDO extends DurableObject { if (this.inflightMutations.has(mutationId)) { if (CATALOG_DEDUPE_COMMANDS.has(msg.command)) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: CATALOG_REQUEST_PENDING_ERROR, }); } else { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: { - source: "relay", - code: "COMMAND_ALREADY_PENDING", - message: "Command is already in flight", + source: 'relay', + code: 'COMMAND_ALREADY_PENDING', + message: 'Command is already in flight', }, }); } @@ -1289,7 +1222,7 @@ export class UserConnectionDO extends DurableObject { const total = this.pendingCommands.size + durableCount; if (total >= UserConnectionDO.MAX_PENDING_COMMANDS) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: PENDING_COMMAND_LIMIT_ERROR, }); @@ -1302,25 +1235,18 @@ export class UserConnectionDO extends DurableObject { targetConnectionId, expectedOwnerConnectionId, webConnectionId, - now, + now ); return; } - if (durable.state === "done") { - await this.updateDurablePendingCommandOriginalId( - mutationId, - msg.id, - ); + if (durable.state === 'done') { + await this.updateDurablePendingCommandOriginalId(mutationId, msg.id); this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, - ...(durable.result !== undefined - ? { result: durable.result } - : {}), - ...(durable.error !== undefined - ? { error: durable.error } - : {}), + ...(durable.result !== undefined ? { result: durable.result } : {}), + ...(durable.error !== undefined ? { error: durable.error } : {}), }); return; } @@ -1328,25 +1254,25 @@ export class UserConnectionDO extends DurableObject { // Entry is 'pending': dedupe. if (CATALOG_DEDUPE_COMMANDS.has(msg.command)) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: CATALOG_REQUEST_PENDING_ERROR, }); } else { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: { - source: "relay", - code: "COMMAND_ALREADY_PENDING", - message: "Command is already in flight", + source: 'relay', + code: 'COMMAND_ALREADY_PENDING', + message: 'Command is already in flight', }, }); } } finally { this.inflightMutations.delete(mutationId); } - })(), + })() ); return; } @@ -1354,7 +1280,7 @@ export class UserConnectionDO extends DurableObject { // In-memory cap check for the common path. if (this.pendingCommands.size >= UserConnectionDO.MAX_PENDING_COMMANDS) { this.sendToWeb(ws, { - type: "response", + type: 'response', id: msg.id, error: PENDING_COMMAND_LIMIT_ERROR, }); @@ -1368,7 +1294,7 @@ export class UserConnectionDO extends DurableObject { targetConnectionId, expectedOwnerConnectionId, webConnectionId, - now, + now ); } @@ -1389,7 +1315,7 @@ export class UserConnectionDO extends DurableObject { targetConnectionId: string, expectedOwnerConnectionId: string | undefined, webConnectionId: string, - now: number, + now: number ): Promise { const correlationId = msg.mutationId ?? crypto.randomUUID(); this.pendingCommands.set(correlationId, { @@ -1410,13 +1336,13 @@ export class UserConnectionDO extends DurableObject { targetConnectionId, expiresAt: now + UserConnectionDO.PENDING_COMMAND_TTL_MS, webConnectionId, - state: "pending" as const, + state: 'pending' as const, } satisfies PendingCommandEntry); this.scheduleNextAlarm(now); this.scheduleDurablePendingAlarm(); this.sendToCli(targetCli, { - type: "command", + type: 'command', id: correlationId, command: msg.command, data: msg.data, @@ -1439,7 +1365,7 @@ export class UserConnectionDO extends DurableObject { targetConnectionId: string, expectedOwnerConnectionId: string | undefined, webConnectionId: string, - now: number, + now: number ): void { const correlationId = crypto.randomUUID(); this.pendingCommands.set(correlationId, { @@ -1461,13 +1387,13 @@ export class UserConnectionDO extends DurableObject { targetConnectionId, expiresAt: now + UserConnectionDO.PENDING_COMMAND_TTL_MS, webConnectionId, - state: "pending" as const, - } satisfies PendingCommandEntry), + state: 'pending' as const, + } satisfies PendingCommandEntry) ); this.scheduleNextAlarm(now); this.scheduleDurablePendingAlarm(); this.sendToCli(targetCli, { - type: "command", + type: 'command', id: correlationId, command: msg.command, data: msg.data, @@ -1481,7 +1407,7 @@ export class UserConnectionDO extends DurableObject { private async handleCliDisconnect( disconnectedWs: WebSocket, - attachment: WSAttachment & { role: "cli" }, + attachment: WSAttachment & { role: 'cli' } ): Promise { const { connectionId } = attachment; @@ -1490,17 +1416,17 @@ export class UserConnectionDO extends DurableObject { // Exclude the closing socket: under wrangler/workerd, getWebSockets() still // includes it during webSocketClose, so matching self would always look "replaced" // and skip ownership cleanup + attention reset (DEF-5 E2E failure). - const replaced = this.ctx.getWebSockets("cli").some((ws) => { + const replaced = this.ctx.getWebSockets('cli').some(ws => { if (ws === disconnectedWs) return false; const att = ws.deserializeAttachment() as WSAttachment | null; - return att?.role === "cli" && att.connectionId === connectionId; + return att?.role === 'cli' && att.connectionId === connectionId; }); // Fail pending commands that targeted this specific socket this.failPendingCommandsForSocket(disconnectedWs, !replaced); if (replaced) { - console.log("Stale CLI socket closed (already replaced)", { + console.log('Stale CLI socket closed (already replaced)', { connectionId, }); return; @@ -1520,10 +1446,10 @@ export class UserConnectionDO extends DurableObject { this.connectionCapabilities.delete(connectionId); this.lastHeartbeatAt.delete(connectionId); - console.log("CLI socket disconnected", { + console.log('CLI socket disconnected', { connectionId, droppedSessions: ownedSessions.size, - remainingCliSockets: this.ctx.getWebSockets("cli").length, + remainingCliSockets: this.ctx.getWebSockets('cli').length, }); // Leave webSubscriptions intact — a reconnecting CLI can resume @@ -1532,14 +1458,11 @@ export class UserConnectionDO extends DurableObject { // departure refetch observes `retry` rather than a stuck `question`. // kiloUserId comes from the CLI attachment (authenticated /user/cli route); // without it we cannot safely target rows and must no-op. - await this.resetOwnedSessionAttentionOnDisconnect( - attachment.kiloUserId, - ownedSessions, - ); + await this.resetOwnedSessionAttentionOnDisconnect(attachment.kiloUserId, ownedSessions); this.broadcastToWeb({ - type: "system", - event: "cli.disconnected", + type: 'system', + event: 'cli.disconnected', data: { connectionId }, }); } @@ -1550,32 +1473,29 @@ export class UserConnectionDO extends DurableObject { */ private async resetOwnedSessionAttentionOnDisconnect( kiloUserId: string | undefined, - ownedSessions: ReadonlySet, + ownedSessions: ReadonlySet ): Promise { if (ownedSessions.size === 0) return; if (!kiloUserId) { console.warn( - "Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment", - { ownedSessionCount: ownedSessions.size }, + 'Skipping attention status reset on CLI disconnect: missing kiloUserId on attachment', + { ownedSessionCount: ownedSessions.size } ); return; } const results = await Promise.allSettled( - [...ownedSessions].map(async (sessionId) => { + [...ownedSessions].map(async sessionId => { const stub = getSessionIngestDO(this.env, { kiloUserId, sessionId }); await stub.resetAttentionStatusOnCliDisconnect(kiloUserId, sessionId); - }), + }) ); for (const result of results) { - if (result.status === "rejected") { - console.error("Failed to reset attention status on CLI disconnect", { - error: - result.reason instanceof Error - ? result.reason.message - : String(result.reason), + if (result.status === 'rejected') { + console.error('Failed to reset attention status on CLI disconnect', { + error: result.reason instanceof Error ? result.reason.message : String(result.reason), }); } } @@ -1583,8 +1503,7 @@ export class UserConnectionDO extends DurableObject { private handleWebDisconnect(ws: WebSocket): void { const attachment = ws.deserializeAttachment() as WSAttachment | null; - const connectionId = - attachment?.role === "web" ? attachment.connectionId : "unknown"; + const connectionId = attachment?.role === 'web' ? attachment.connectionId : 'unknown'; // Remove from all subscription sets let droppedSubscriptions = 0; @@ -1598,7 +1517,7 @@ export class UserConnectionDO extends DurableObject { // Tell owning CLI to stop forwarding const cliWs = this.findCliForSession(sessionId); if (cliWs) { - this.sendToCli(cliWs, { type: "unsubscribe", sessionId }); + this.sendToCli(cliWs, { type: 'unsubscribe', sessionId }); } } } @@ -1615,11 +1534,11 @@ export class UserConnectionDO extends DurableObject { } } - console.log("Web socket disconnected", { + console.log('Web socket disconnected', { connectionId, droppedSubscriptions, droppedCommands, - remainingWebSockets: this.ctx.getWebSockets("web").length, + remainingWebSockets: this.ctx.getWebSockets('web').length, }); } @@ -1652,9 +1571,9 @@ export class UserConnectionDO extends DurableObject { getConnectedInstances(): { instances: ConnectedInstanceRow[] } { this.ensureState(); const instances: ConnectedInstanceRow[] = []; - for (const ws of this.ctx.getWebSockets("cli")) { + for (const ws of this.ctx.getWebSockets('cli')) { const att = ws.deserializeAttachment() as WSAttachment | null; - if (att?.role !== "cli" || !att.instance) continue; + if (att?.role !== 'cli' || !att.instance) continue; instances.push({ connectionId: att.connectionId, name: att.instance.name, @@ -1666,13 +1585,11 @@ export class UserConnectionDO extends DurableObject { return { instances }; } - async notifySessionEvent( - event: SessionEventPayload, - ): Promise<{ delivered: number }> { + async notifySessionEvent(event: SessionEventPayload): Promise<{ delivered: number }> { this.ensureState(); const parsed = SessionEventPayloadSchema.parse(event); const msg: WebInboundMessage = { - type: "system", + type: 'system', event: parsed.type, data: parsed.data, }; @@ -1684,7 +1601,7 @@ export class UserConnectionDO extends DurableObject { ws.send(json); delivered++; } catch (err) { - console.warn("notifySessionEvent: skipping failed web socket:", err); + console.warn('notifySessionEvent: skipping failed web socket:', err); } } return { delivered }; @@ -1694,10 +1611,7 @@ export class UserConnectionDO extends DurableObject { * Persist a web→CLI rename under KV (offline catch-up) and deliver * `session.renamed` to the owning CLI socket when one is connected. */ - async notifySessionRenamed( - sessionId: string, - title: string, - ): Promise<{ delivered: boolean }> { + async notifySessionRenamed(sessionId: string, title: string): Promise<{ delivered: boolean }> { this.ensureState(); await this.ctx.storage.put(`${RENAME_KEY_PREFIX}${sessionId}`, { title, @@ -1714,8 +1628,8 @@ export class UserConnectionDO extends DurableObject { } this.sendToCli(cliWs, { - type: "system", - event: "session.renamed", + type: 'system', + event: 'session.renamed', data: { sessionId, title }, }); return { delivered: true }; @@ -1734,7 +1648,7 @@ export class UserConnectionDO extends DurableObject { try { ws.send(JSON.stringify(msg)); } catch (err) { - console.warn("sendToCli failed:", err); + console.warn('sendToCli failed:', err); } } @@ -1742,7 +1656,7 @@ export class UserConnectionDO extends DurableObject { try { ws.send(JSON.stringify(msg)); } catch (err) { - console.warn("sendToWeb failed:", err); + console.warn('sendToWeb failed:', err); } } @@ -1753,7 +1667,7 @@ export class UserConnectionDO extends DurableObject { try { ws.send(json); } catch (err) { - console.warn("broadcastToWeb: skipping failed socket:", err); + console.warn('broadcastToWeb: skipping failed socket:', err); } } } @@ -1761,13 +1675,13 @@ export class UserConnectionDO extends DurableObject { /** Close a stale CLI socket that has the same connectionId (from a previous connection). Returns true if one was found. */ private closeStaleSocket(connectionId: string): boolean { - for (const ws of this.ctx.getWebSockets("cli")) { + 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 socket for reconnect", { connectionId }); + if (att?.role === 'cli' && att.connectionId === connectionId) { + console.log('Closing stale CLI socket for reconnect', { connectionId }); this.failPendingCommandsForSocket(ws, false); // Preserve session ownership — the reconnecting CLI still owns these sessions - ws.close(1000, "replaced by reconnect"); + ws.close(1000, 'replaced by reconnect'); return true; } } @@ -1775,10 +1689,10 @@ export class UserConnectionDO extends DurableObject { } private replaceWebSocket(connectionId: string): void { - for (const ws of this.ctx.getWebSockets("web")) { + for (const ws of this.ctx.getWebSockets('web')) { const attachment = ws.deserializeAttachment() as WSAttachment | null; if ( - attachment?.role !== "web" || + attachment?.role !== 'web' || attachment.connectionId !== connectionId || attachment.replaced ) { @@ -1787,14 +1701,14 @@ export class UserConnectionDO extends DurableObject { ws.serializeAttachment({ ...attachment, replaced: true }); this.handleWebDisconnect(ws); - ws.close(1000, "replaced by reconnect"); + ws.close(1000, 'replaced by reconnect'); } } private activeWebSockets(): WebSocket[] { - return this.ctx.getWebSockets("web").filter((ws) => { + return this.ctx.getWebSockets('web').filter(ws => { const attachment = ws.deserializeAttachment() as WSAttachment | null; - return attachment?.role === "web" && !attachment.replaced; + return attachment?.role === 'web' && !attachment.replaced; }); } @@ -1805,12 +1719,9 @@ export class UserConnectionDO extends DurableObject { } private findCliByConnectionId(connectionId: string): WebSocket | undefined { - for (const ws of this.ctx.getWebSockets("cli")) { + for (const ws of this.ctx.getWebSockets('cli')) { const attachment = ws.deserializeAttachment() as WSAttachment | null; - if ( - attachment?.role === "cli" && - attachment.connectionId === connectionId - ) { + if (attachment?.role === 'cli' && attachment.connectionId === connectionId) { return ws; } } @@ -1818,10 +1729,10 @@ export class UserConnectionDO extends DurableObject { } private findWebByConnectionId(connectionId: string): WebSocket | undefined { - for (const ws of this.ctx.getWebSockets("web")) { + for (const ws of this.ctx.getWebSockets('web')) { const attachment = ws.deserializeAttachment() as WSAttachment | null; if ( - attachment?.role === "web" && + attachment?.role === 'web' && attachment.connectionId === connectionId && !attachment.replaced ) { @@ -1831,15 +1742,12 @@ export class UserConnectionDO extends DurableObject { return undefined; } - private failPendingCommandsForSocket( - targetWs: WebSocket, - cliGone: boolean, - ): void { + private failPendingCommandsForSocket(targetWs: WebSocket, cliGone: boolean): void { for (const [id, entry] of this.pendingCommands) { if (entry.targetCliWs !== targetWs) continue; // 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 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). @@ -1847,17 +1755,17 @@ export class UserConnectionDO extends DurableObject { ? undefined : isOwnerFenced ? SESSION_OWNER_CHANGED_ERROR - : "CLI disconnected"; + : 'CLI disconnected'; const durableError = liveError; this.sendToWeb( entry.ws, exited - ? { type: "response", id: entry.originalId, result: {} } + ? { type: 'response', id: entry.originalId, result: {} } : { - type: "response", + type: 'response', id: entry.originalId, error: liveError, - }, + } ); this.pendingCommands.delete(id); @@ -1871,50 +1779,45 @@ export class UserConnectionDO extends DurableObject { expectedOwnerConnectionId: entry.expectedOwnerConnectionId, targetConnectionId: entry.targetConnectionId, expiresAt: entry.expiresAt, - webConnectionId: - webAtt?.role === "web" ? webAtt.connectionId : "unknown", - state: "done", + webConnectionId: webAtt?.role === 'web' ? webAtt.connectionId : 'unknown', + state: 'done', ...(exited ? { result: {} } : { error: durableError }), - } satisfies PendingCommandEntry), + } satisfies PendingCommandEntry) ); } const attachment = targetWs.deserializeAttachment() as WSAttachment | null; - if (cliGone && attachment?.role === "cli") { + if (cliGone && attachment?.role === 'cli') { this.ctx.waitUntil( this.finishDurablePendingCommands( entry => entry.targetConnectionId === attachment.connectionId, - "CLI disconnected", - ), + 'CLI disconnected' + ) ); } } private failPendingCommandsForOwnerChange( sessionId: string, - nextOwnerConnectionId: string | undefined, + nextOwnerConnectionId: string | undefined ): void { for (const [id, entry] of this.pendingCommands) { - if ( - entry.sessionId !== sessionId || - entry.targetConnectionId === nextOwnerConnectionId - ) { + if (entry.sessionId !== sessionId || entry.targetConnectionId === nextOwnerConnectionId) { continue; } this.pendingCommands.delete(id); // `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; + const exited = entry.command === 'exit_cli' && nextOwnerConnectionId === undefined; this.sendToWeb( entry.ws, exited - ? { type: "response", id: entry.originalId, result: {} } + ? { type: 'response', id: entry.originalId, result: {} } : { - type: "response", + type: 'response', id: entry.originalId, error: SESSION_OWNER_CHANGED_ERROR, - }, + } ); // Step 26: pair the in-memory delete with a durable done transition. @@ -1927,21 +1830,19 @@ export class UserConnectionDO extends DurableObject { expectedOwnerConnectionId: entry.expectedOwnerConnectionId, targetConnectionId: entry.targetConnectionId, expiresAt: entry.expiresAt, - webConnectionId: - webAtt?.role === "web" ? webAtt.connectionId : "unknown", - state: "done", + webConnectionId: webAtt?.role === 'web' ? webAtt.connectionId : 'unknown', + state: 'done', ...(exited ? { result: {} } : { error: SESSION_OWNER_CHANGED_ERROR }), - } satisfies PendingCommandEntry), + } satisfies PendingCommandEntry) ); } this.ctx.waitUntil( this.finishDurablePendingCommands( entry => - entry.sessionId === sessionId && - entry.targetConnectionId !== nextOwnerConnectionId, - SESSION_OWNER_CHANGED_ERROR, - ), + entry.sessionId === sessionId && entry.targetConnectionId !== nextOwnerConnectionId, + SESSION_OWNER_CHANGED_ERROR + ) ); } @@ -1950,10 +1851,10 @@ export class UserConnectionDO extends DurableObject { // --------------------------------------------------------------------------- private async getDurablePendingCommand( - correlationId: string, + correlationId: string ): Promise { return this.ctx.storage.get( - `${PENDING_COMMAND_KEY_PREFIX}${correlationId}`, + `${PENDING_COMMAND_KEY_PREFIX}${correlationId}` ); } @@ -1966,30 +1867,27 @@ export class UserConnectionDO extends DurableObject { private async updateDurablePendingCommandOriginalId( correlationId: string, - newOriginalId: 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, - ); + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${correlationId}`, entry); } private async finishDurablePendingCommands( matches: (entry: PendingCommandEntry) => boolean, - error: unknown, + error: unknown ): Promise { const entries = await this.ctx.storage.list({ prefix: PENDING_COMMAND_KEY_PREFIX, }); await Promise.all( [...entries].flatMap(([key, entry]) => - entry.state === "pending" && matches(entry) - ? [this.ctx.storage.put(key, { ...entry, state: "done", error })] - : [], - ), + entry.state === 'pending' && matches(entry) + ? [this.ctx.storage.put(key, { ...entry, state: 'done', error })] + : [] + ) ); } @@ -2012,7 +1910,7 @@ export class UserConnectionDO extends DurableObject { if (currentAlarm === null || expiresAt < currentAlarm) { await this.ctx.storage.setAlarm(expiresAt); } - })(), + })() ); } @@ -2021,7 +1919,7 @@ export class UserConnectionDO extends DurableObject { if (entry.expiresAt > now) continue; this.pendingCommands.delete(id); this.sendToWeb(entry.ws, { - type: "response", + type: 'response', id: entry.originalId, error: COMMAND_EXPIRED_ERROR, }); @@ -2036,14 +1934,14 @@ export class UserConnectionDO extends DurableObject { prefix: PENDING_COMMAND_KEY_PREFIX, }); for (const [key, durable] of entries) { - if (!durable || typeof durable.expiresAt !== "number") continue; + if (!durable || typeof durable.expiresAt !== 'number') continue; if (durable.expiresAt > now) continue; - if (durable.state === "pending") { + if (durable.state === 'pending') { // Write terminal outcome durably (step 26a). await this.ctx.storage.put(key, { ...durable, - state: "done", + state: 'done', error: COMMAND_EXPIRED_ERROR, }); } else { @@ -2052,10 +1950,10 @@ export class UserConnectionDO extends DurableObject { } } })().catch((error: unknown) => { - console.error("Failed to sweep durable pending commands", { + console.error('Failed to sweep durable pending commands', { error: error instanceof Error ? error.message : String(error), }); - }), + }) ); } @@ -2070,7 +1968,7 @@ export class UserConnectionDO extends DurableObject { prefix: READY_PUSH_KEY_PREFIX, }); for (const [key, entry] of pending) { - if (!entry || typeof entry.fireAt !== "number") continue; + if (!entry || typeof entry.fireAt !== 'number') continue; const sessionId = key.slice(READY_PUSH_KEY_PREFIX.length); if (sessionId) this.readyPushFireAt.set(sessionId, entry.fireAt); } @@ -2078,11 +1976,11 @@ export class UserConnectionDO extends DurableObject { this.scheduleNextAlarm(Date.now()); } catch (error: unknown) { // Leave readyPushRebuilt false so the next schedule retries. - console.error("Failed to rebuild readyPush mirror", { + console.error('Failed to rebuild readyPush mirror', { error: error instanceof Error ? error.message : String(error), }); } - })(), + })() ); } @@ -2090,19 +1988,13 @@ export class UserConnectionDO extends DurableObject { for (const lastSeen of this.lastHeartbeatAt.values()) { const staleAt = lastSeen + UserConnectionDO.HEARTBEAT_TIMEOUT_MS; - if ( - staleAt > now && - (nextAlarmAt === undefined || staleAt < nextAlarmAt) - ) { + if (staleAt > now && (nextAlarmAt === undefined || staleAt < nextAlarmAt)) { nextAlarmAt = staleAt; } } for (const entry of this.pendingCommands.values()) { - if ( - entry.expiresAt > now && - (nextAlarmAt === undefined || entry.expiresAt < nextAlarmAt) - ) { + if (entry.expiresAt > now && (nextAlarmAt === undefined || entry.expiresAt < nextAlarmAt)) { nextAlarmAt = entry.expiresAt; } } @@ -2136,9 +2028,9 @@ export class UserConnectionDO extends DurableObject { // Build set of connectionIds that still have a live CLI WebSocket. // This guards against stale entries that persist if a close event is delayed. const liveConnectionIds = new Set(); - for (const ws of this.ctx.getWebSockets("cli")) { + for (const ws of this.ctx.getWebSockets('cli')) { const att = ws.deserializeAttachment() as WSAttachment | null; - if (att?.role === "cli") liveConnectionIds.add(att.connectionId); + if (att?.role === 'cli') liveConnectionIds.add(att.connectionId); } const result: Array< From 99a10a9ef4c34820a7b924e81ea75c0ca1de6542 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Mon, 3 Aug 2026 23:40:35 +0200 Subject: [PATCH 06/24] fix(chat): tolerate unknown bot capabilities --- .../src/use-bot-status.test.ts | 26 ++------------ .../kilo-chat-hooks/src/use-bot-status.ts | 9 ++--- packages/kilo-chat/src/events.ts | 4 +-- packages/kilo-chat/src/schemas.ts | 8 +++-- .../kilo-chat/test/events.bot-status.test.ts | 21 +++++++++++ .../kilo-chat/test/schemas.bot-status.test.ts | 29 +++++++++++++-- .../kilo-chat/test/schemas.capability.test.ts | 36 ++++++++++++++++++- 7 files changed, 96 insertions(+), 37 deletions(-) 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..458d78d08d3 100644 --- a/packages/kilo-chat-hooks/src/use-bot-status.test.ts +++ b/packages/kilo-chat-hooks/src/use-bot-status.test.ts @@ -210,36 +210,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); }); @@ -373,7 +352,6 @@ describe('useBotStatus', () => { useBotStatus(kiloChatClient, eventClient, null); const queryCall = testState.queryCalls[0]; - // enabled = sandboxId !== null && wsReady → false expect(queryCall?.enabled).toBe(false); }); }); diff --git a/packages/kilo-chat-hooks/src/use-bot-status.ts b/packages/kilo-chat-hooks/src/use-bot-status.ts index 76ae71ff2c8..906a51477e0 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 { @@ -32,12 +32,9 @@ export function useBotStatus( ): BotStatusRecord | null { const queryClient = useQueryClient(); - // WS-ready gate. onConnected fires synchronously if already connected, and on every reconnect. - const [wsReady, setWsReady] = useState(false); + // On reconnect, refetch to catch up on anything we missed while disconnected. useEffect(() => { return eventClient.onConnected(() => { - setWsReady(true); - // On reconnect, refetch to catch up on anything we missed while disconnected. if (sandboxId) { void queryClient.invalidateQueries({ queryKey: botStatusKey(sandboxId) }); } @@ -69,7 +66,7 @@ export function useBotStatus( }); return queryClient.getQueryData(botStatusKey(sandboxId)) ?? null; }, - enabled: sandboxId !== null && wsReady, + enabled: sandboxId !== null, refetchInterval: POLL_INTERVAL_MS, staleTime: 0, }); 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/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([]); + } + }); +}); From 0b953e9aade631d7db1da45586a90f9beb537ea8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 00:09:20 +0200 Subject: [PATCH 07/24] fix(events): detect sequence gaps --- .../src/__tests__/schemas.test.ts | 198 ++++++++++++++++++ packages/event-service/src/schemas.ts | 8 + packages/event-service/src/types.ts | 2 + .../src/__tests__/user-session-do.test.ts | 176 +++++++++++++++- .../event-service/src/do/user-session-do.ts | 79 ++++++- 5 files changed, 455 insertions(+), 8 deletions(-) create mode 100644 packages/event-service/src/__tests__/schemas.test.ts 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/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/services/event-service/src/__tests__/user-session-do.test.ts b/services/event-service/src/__tests__/user-session-do.test.ts index 8f51cdba971..be2bfdf5c20 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,122 @@ 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: 15000 }, + 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 (50). + // seq starts at 0, increments each push. After N pushes seq = N. + // Gap = seq - acked = N - 0. Need N > 50. + // Push 55 events for margin. + for (let i = 0; i < 55; i++) { + await stub.pushEvent('project:z', 'fill', {}); + // Brief pause to let the DO process the message and potential close + await new Promise(r => setTimeout(r, 20)); + // Check if we already closed + 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: 15000 }, + 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 many events without any ack — should never trigger zombie close + for (let i = 0; i < 60; i++) { + await stub.pushEvent('project:noack', 'fill', {}); + await new Promise(r => setTimeout(r, 10)); + if (closed) break; + } + + 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 +277,61 @@ 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: 15000 }, + 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 < 60; i++) { + await stub.pushEvent('project:old', 'fill', {}); + 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..c94bc8066cd 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 = 50; + +/** 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); } } From 7b4405527129469b22f74a7e1f64170b298feaba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 00:25:08 +0200 Subject: [PATCH 08/24] fix(events): resync consumers after gaps --- .../use-conversation-event-subscription.ts | 2 +- .../src/__tests__/client.test.ts | 286 ++++++++++++++++++ packages/event-service/src/client.ts | 53 ++++ .../src/use-bot-status.test.ts | 18 +- .../kilo-chat-hooks/src/use-bot-status.ts | 4 +- .../src/use-conversations.test.ts | 26 +- .../kilo-chat-hooks/src/use-conversations.ts | 8 +- 7 files changed, 379 insertions(+), 18 deletions(-) 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/packages/event-service/src/__tests__/client.test.ts b/packages/event-service/src/__tests__/client.test.ts index ce725aeb87b..f861d9b2283 100644 --- a/packages/event-service/src/__tests__/client.test.ts +++ b/packages/event-service/src/__tests__/client.test.ts @@ -855,4 +855,290 @@ 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('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(); + } + }); + }); }); diff --git a/packages/event-service/src/client.ts b/packages/event-service/src/client.ts index 7dc1b973145..53772489674 100644 --- a/packages/event-service/src/client.ts +++ b/packages/event-service/src/client.ts @@ -89,6 +89,14 @@ export class EventServiceClient { private hasConnectedBefore = false; private reconnectHandlers = new Set<() => void>(); private connectedHandlers = new Set<() => void>(); + private resyncHandlers = new Set<() => void>(); + // 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; @@ -207,11 +215,17 @@ export class EventServiceClient { 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(); @@ -344,6 +358,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) { @@ -378,6 +407,23 @@ 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). + if (message.seq !== undefined) { + if (message.seq > this.lastKnownSeq + 1) { + for (const handler of this.resyncHandlers) { + handler(); + } + } + if (message.seq > this.lastKnownSeq) { + this.lastKnownSeq = message.seq; + } + if (message.seq > this.highestAckedSeq) { + this.highestAckedSeq = message.seq; + } + } + if (message.type === 'event') { const handlers = this.eventHandlers.get(message.event); if (handlers) { @@ -401,6 +447,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); } 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 458d78d08d3..31b89d6699d 100644 --- a/packages/kilo-chat-hooks/src/use-bot-status.test.ts +++ b/packages/kilo-chat-hooks/src/use-bot-status.test.ts @@ -91,16 +91,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 +123,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, }; @@ -331,15 +331,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')); diff --git a/packages/kilo-chat-hooks/src/use-bot-status.ts b/packages/kilo-chat-hooks/src/use-bot-status.ts index 906a51477e0..26114d34244 100644 --- a/packages/kilo-chat-hooks/src/use-bot-status.ts +++ b/packages/kilo-chat-hooks/src/use-bot-status.ts @@ -32,9 +32,9 @@ export function useBotStatus( ): BotStatusRecord | null { const queryClient = useQueryClient(); - // On reconnect, refetch to catch up on anything we missed while disconnected. + // On resync (reconnect or sequence gap), refetch to catch up on missed events. useEffect(() => { - return eventClient.onConnected(() => { + return eventClient.onResync(() => { if (sandboxId) { void queryClient.invalidateQueries({ queryKey: botStatusKey(sandboxId) }); } 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); }), ]; From 6aca97a9f8e67c7c39fb5e97e3b0a48aa07eb861 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 00:54:46 +0200 Subject: [PATCH 09/24] fix(mobile): release events while inactive --- .../kilo-chat/conversation-list-screen.tsx | 4 +- .../kilo-chat/conversation-screen.tsx | 4 +- .../kilo-chat-provider.mounted.test.tsx | 166 ++++++++++++++++++ .../kilo-chat/kilo-chat-provider.tsx | 14 +- .../src/__tests__/client.test.ts | 105 +++++++++++ packages/event-service/src/client.ts | 32 ++++ .../src/use-bot-status.test.ts | 22 +++ .../kilo-chat-hooks/src/use-bot-status.ts | 5 +- 8 files changed, 345 insertions(+), 7 deletions(-) create mode 100644 apps/mobile/src/components/kilo-chat/kilo-chat-provider.mounted.test.tsx 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/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..d8d51edff65 --- /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'; + +// ── 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, +})); + +import { KiloChatProvider } from './kilo-chat-provider'; + +// ── 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/packages/event-service/src/__tests__/client.test.ts b/packages/event-service/src/__tests__/client.test.ts index f861d9b2283..9a63d8eb394 100644 --- a/packages/event-service/src/__tests__/client.test.ts +++ b/packages/event-service/src/__tests__/client.test.ts @@ -1141,4 +1141,109 @@ describe('EventServiceClient', () => { } }); }); + + 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); + }); + }); }); diff --git a/packages/event-service/src/client.ts b/packages/event-service/src/client.ts index 53772489674..8abf666afcd 100644 --- a/packages/event-service/src/client.ts +++ b/packages/event-service/src/client.ts @@ -90,6 +90,11 @@ export class EventServiceClient { 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). @@ -290,6 +295,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; } 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 31b89d6699d..bfc8f8ed925 100644 --- a/packages/kilo-chat-hooks/src/use-bot-status.test.ts +++ b/packages/kilo-chat-hooks/src/use-bot-status.test.ts @@ -354,4 +354,26 @@ describe('useBotStatus', () => { const queryCall = testState.queryCalls[0]; 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 26114d34244..238601691a6 100644 --- a/packages/kilo-chat-hooks/src/use-bot-status.ts +++ b/packages/kilo-chat-hooks/src/use-bot-status.ts @@ -28,7 +28,8 @@ export function reduceBotStatusOnEvent( export function useBotStatus( client: KiloChatClient, eventClient: EventServiceClient, - sandboxId: string | null + sandboxId: string | null, + active: boolean = true ): BotStatusRecord | null { const queryClient = useQueryClient(); @@ -67,7 +68,7 @@ export function useBotStatus( return queryClient.getQueryData(botStatusKey(sandboxId)) ?? null; }, enabled: sandboxId !== null, - refetchInterval: POLL_INTERVAL_MS, + refetchInterval: active ? POLL_INTERVAL_MS : undefined, staleTime: 0, }); From 5adcc079981f9dc7477b8cac3a2f21e21b059410 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 01:41:14 +0200 Subject: [PATCH 10/24] fix(mobile): bound network request deadlines --- apps/mobile/src/lib/trpc.test.ts | 143 ++++++++ apps/mobile/src/lib/trpc.ts | 24 +- .../src/__tests__/client.test.ts | 126 +++++++ .../src/__tests__/deadline.test.ts | 308 ++++++++++++++++++ packages/event-service/src/client.ts | 28 +- packages/event-service/src/deadline.ts | 114 +++++++ packages/event-service/src/index.ts | 6 + packages/kilo-chat/src/client.ts | 49 ++- packages/kilo-chat/test/client.test.ts | 270 +++++++++++++++ 9 files changed, 1041 insertions(+), 27 deletions(-) create mode 100644 apps/mobile/src/lib/trpc.test.ts create mode 100644 packages/event-service/src/__tests__/deadline.test.ts create mode 100644 packages/event-service/src/deadline.ts diff --git a/apps/mobile/src/lib/trpc.test.ts b/apps/mobile/src/lib/trpc.test.ts new file mode 100644 index 00000000000..0723aebe829 --- /dev/null +++ b/apps/mobile/src/lib/trpc.test.ts @@ -0,0 +1,143 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +// Mock expo-constants before any module that imports it. +vi.mock('expo-constants', () => ({ + default: { + expoConfig: { + extra: { + apiBaseUrl: 'https://api.example.com', + webBaseUrl: 'https://web.example.com', + appsFlyerDevKey: 'dev-key', + appsFlyerAppId: 'app-id', + cloudAgentWsUrl: 'wss://agent.example.com', + sessionIngestWsUrl: 'wss://ingest.example.com', + kiloChatUrl: 'https://chat.example.com', + eventServiceUrl: 'wss://events.example.com', + notificationsUrl: 'https://notifications.example.com', + posthogApiKey: 'ph-key', + }, + }, + }, +})); + +vi.mock('expo-secure-store', () => ({ + getItemAsync: vi.fn().mockResolvedValue(null), + setItemAsync: vi.fn().mockResolvedValue(undefined), + deleteItemAsync: vi.fn().mockResolvedValue(undefined), +})); + +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 abe2079abe9..a0ce4880c8b 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'; @@ -55,6 +56,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 + ); + return response; +}; + async function getAuthHeaders() { const token = await SecureStore.getItemAsync(AUTH_TOKEN_KEY); if (!token) { @@ -70,12 +90,12 @@ export const trpcClient = createTRPCClient({ true: httpLink({ url: trpcUrl, headers: getAuthHeaders, - fetch: e2eFetch, + fetch: deadlineFetch, }), false: httpBatchLink({ url: trpcUrl, headers: getAuthHeaders, - fetch: e2eFetch, + fetch: deadlineFetch, }), }), ], diff --git a/packages/event-service/src/__tests__/client.test.ts b/packages/event-service/src/__tests__/client.test.ts index 9a63d8eb394..a83fc31ef8a 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); @@ -1246,4 +1248,128 @@ describe('EventServiceClient', () => { 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..d929f35c542 --- /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(global, '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/client.ts b/packages/event-service/src/client.ts index 8abf666afcd..583155ca712 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 { 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/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/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(); + } + }); + }); }); From 641057d638ed7e6e379ea4dd1f2cd943cb49d99d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 01:44:31 +0200 Subject: [PATCH 11/24] test(mobile): fix deadline type checks --- .../event-service/src/__tests__/deadline.test.ts | 2 +- .../kilo-chat-hooks/src/use-bot-status.test.ts | 14 ++++++++++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/packages/event-service/src/__tests__/deadline.test.ts b/packages/event-service/src/__tests__/deadline.test.ts index d929f35c542..014a25eb2f9 100644 --- a/packages/event-service/src/__tests__/deadline.test.ts +++ b/packages/event-service/src/__tests__/deadline.test.ts @@ -209,7 +209,7 @@ describe('Hermes AbortError resilience', () => { it('does not leak timer or listener when fn rejects generic AbortError after deadline', async () => { vi.useFakeTimers(); try { - const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout'); + const clearTimeoutSpy = vi.spyOn(globalThis, 'clearTimeout'); const promise = withDeadline(100, signal => { return new Promise((_resolve, reject) => { 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 bfc8f8ed925..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) }; From 7fcfd4186f83d3957aa7228acbc360bb0ac7e626 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 01:53:11 +0200 Subject: [PATCH 12/24] fix(mobile): restore stream type contracts --- apps/mobile/src/lib/trpc.ts | 2 +- .../components/kiloChatLayoutCache.test.ts | 15 ++++++++------- apps/web/src/lib/session-ingest-client.ts | 1 + 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/apps/mobile/src/lib/trpc.ts b/apps/mobile/src/lib/trpc.ts index a0ce4880c8b..6bf6a34f6e3 100644 --- a/apps/mobile/src/lib/trpc.ts +++ b/apps/mobile/src/lib/trpc.ts @@ -70,7 +70,7 @@ export const deadlineFetch: typeof fetch = async (url, init) => { const res = await e2eFetch(url, { ...init, signal }); return res; }, - init?.signal + init?.signal ?? undefined ); return response; }; 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/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; }; /** From d5ce3b66c32d954690ef3d2d361f045de7113687 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 02:33:53 +0200 Subject: [PATCH 13/24] fix(cli): preserve durable reply outcomes --- .../kilo-chat-provider.mounted.test.tsx | 4 +- .../src/dos/UserConnectionDO.test.ts | 307 +++++++++++++++++- .../src/dos/UserConnectionDO.ts | 103 +++++- 3 files changed, 388 insertions(+), 26 deletions(-) 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 index d8d51edff65..d97e6f35193 100644 --- 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 @@ -4,6 +4,8 @@ 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(() => { @@ -72,8 +74,6 @@ vi.mock('./hooks/use-kilo-chat-token', () => ({ useKiloChatTokenResponseGetter: mocks.useKiloChatTokenResponseGetter, })); -import { KiloChatProvider } from './kilo-chat-provider'; - // ── Flat render helper ───────────────────────────────────────────────────── async function mountProvider(): Promise { diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 5eeb84eec49..74671333c9a 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -4870,7 +4870,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); // Send the CLI response with the correlation id. - await sendCliResponse(doInstance, cliWs, { + sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true }, }); @@ -4909,7 +4909,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses-b', 'busy', 'Session B')]); // Send the CLI response. - await sendCliResponse(doInstance, cliWs, { + sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true }, }); @@ -4944,7 +4944,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); // Send a second command with the same mutationId. - await sendCommand(doInstance, webWs, { + sendCommand(doInstance, webWs, { id: 'second-req', command: 'send_message', sessionId: 'ses-c', @@ -4987,7 +4987,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); // Send a retry with the same mutationId but new wire id. - await sendCommand(doInstance, webWs, { + sendCommand(doInstance, webWs, { id: 'retry-req-id', command: 'send_message', sessionId: 'ses-d', @@ -5012,6 +5012,50 @@ describe('UserConnectionDO', () => { 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. + 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(); @@ -5020,7 +5064,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-5'); cliWs.send.mockClear(); - await sendCommand(doInstance, webWs, { + sendCommand(doInstance, webWs, { id: 'no-mut-req', command: 'send_message', sessionId: 'ses-e', @@ -5041,7 +5085,7 @@ describe('UserConnectionDO', () => { expect(entry).toBeDefined(); // Send CLI response to verify normal resolution. - await sendCliResponse(doInstance, cliWs, { + sendCliResponse(doInstance, cliWs, { id: correlationId, result: { normal: true }, }); @@ -5059,7 +5103,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-6'); cliWs.send.mockClear(); - await sendCommand(doInstance, webWs, { + sendCommand(doInstance, webWs, { id: 'disconnect-req', command: 'send_message', sessionId: 'ses-f', @@ -5077,7 +5121,7 @@ describe('UserConnectionDO', () => { expect(entry).toBeDefined(); // Disconnect the web socket. - doInstance.webSocketClose(webWs as never, 1000, '', true); + await doInstance.webSocketClose(webWs as never, 1000, '', true); await flushAsync(); // The durable entry must still exist. @@ -5146,7 +5190,7 @@ describe('UserConnectionDO', () => { } // The first command (no mutationId) is just within the cap: 0 in-memory + 1 new = 1. - await sendCommand(doInstance, webWs, { + sendCommand(doInstance, webWs, { id: 'at-cap', command: 'send_message', sessionId: 'ses-h', @@ -5161,7 +5205,7 @@ describe('UserConnectionDO', () => { // command's durable write + 1 in-memory = 129 ≥ 128. Rejected. cliWs.send.mockClear(); webWs.send.mockClear(); - await sendCommand(doInstance, webWs, { + sendCommand(doInstance, webWs, { id: 'over-cap', command: 'send_message', sessionId: 'ses-h', @@ -5196,7 +5240,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - await sendCliResponse(doInstance, cliWs, { + sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_models', }); @@ -5227,7 +5271,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - await sendCliResponse(doInstance, cliWs, { + sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_commands', }); @@ -5310,7 +5354,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); - await sendCliResponse(doInstance, cliWs, { + sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized, }); @@ -5384,5 +5428,242 @@ describe('UserConnectionDO', () => { 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. + sendCliResponse(doInstance, cliWs, { + id: correlationId, + error: 'unknown command: list_commands', + }); + await flushAsync(); + + // 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')]); + + sendCliResponse(doInstance, cliWs, { + id: correlationId, + error: 'unknown command: list_models', + }); + await flushAsync(); + + 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); + sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: oversized, + }); + await flushAsync(); + + 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')]); + + sendCliResponse(doInstance, cliWs, { + id: correlationId, + error: { + source: 'relay', + code: 'SESSION_OWNER_CHANGED', + message: 'Session owner changed', + }, + }); + await flushAsync(); + + 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. + sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { partial: 'data' }, + error: 'something went wrong', + }); + await flushAsync(); + + 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); + sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: oversized, + }); + await flushAsync(); + + // 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(); + }); }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 9f916b86424..7f9a54d4dd3 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -799,11 +799,16 @@ export class UserConnectionDO extends DurableObject { ): 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; // Step 24: on in-memory miss, try to load the durable entry (D8 case 1). if (!entry) { const durable = await this.getDurablePendingCommand(id); if (!durable || durable.state !== 'pending') return; + rehydratedDurable = durable; // Validate the responding CLI by attachment connectionId. const respondingAttachment = respondingWs.deserializeAttachment() as WSAttachment | null; @@ -816,10 +821,24 @@ export class UserConnectionDO extends DurableObject { // 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); + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { + ...durable, + state: 'done', + ...(shaped.result !== undefined ? { result: shaped.result } : {}), + ...(shaped.error !== undefined ? { error: shaped.error } : {}), + }); + return; + } // Build an in-memory entry so the rest of handleCliResponse works. entry = { - ws: webWs!, + ws: webWs, sessionId: durable.sessionId, originalId: durable.originalId, command: durable.command, @@ -851,15 +870,22 @@ export class UserConnectionDO extends DurableObject { } // Persist the terminal outcome durably. - const catWebConnectionId = rehydrated - ? (await this.getDurablePendingCommand(id))!.webConnectionId - : (entry.ws.deserializeAttachment() as WSAttachment | null)?.role === 'web' - ? ( - entry.ws.deserializeAttachment() as WSAttachment & { - role: 'web'; - } - ).connectionId - : 'unknown'; + 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'; + } await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { sessionId: entry.sessionId, originalId: entry.originalId, @@ -937,7 +963,7 @@ export class UserConnectionDO extends DurableObject { // surrounding ctx.waitUntil; the live response has already been // delivered for the non-rehydrated case above. const webConnectionId = rehydrated - ? (await this.getDurablePendingCommand(id))?.webConnectionId + ? rehydratedDurable?.webConnectionId : (entry.ws.deserializeAttachment() as WSAttachment | null)?.role === 'web' ? (entry.ws.deserializeAttachment() as WSAttachment & { role: 'web' }).connectionId : 'unknown'; @@ -1846,6 +1872,61 @@ export class UserConnectionDO extends DurableObject { ); } + // --------------------------------------------------------------------------- + // 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 } : {}), + }; + } + // --------------------------------------------------------------------------- // Durable pending command storage (D8) // --------------------------------------------------------------------------- From b8330685a2944aa1f10d7fb6a23b8c24921f6a76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 04:23:04 +0200 Subject: [PATCH 14/24] fix(dev): align docker gateway port --- dev/local/scripts/start-tunnel.test.ts | 44 ++++ dev/local/scripts/start-tunnel.ts | 346 +++++++++++++------------ 2 files changed, 228 insertions(+), 162 deletions(-) create mode 100644 dev/local/scripts/start-tunnel.test.ts 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(); } From 5533d7722dc994bd7122afb45e5def305aadea29 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 05:38:24 +0200 Subject: [PATCH 15/24] fix(stream): harden durable replay handling --- .../src/cloud-agent-transport.test.ts | 101 +++++- .../src/cloud-agent-transport.ts | 14 +- .../src/__tests__/client.test.ts | 42 +++ packages/event-service/src/client.ts | 13 +- .../src/__tests__/user-session-do.test.ts | 78 ++++- .../event-service/src/do/user-session-do.ts | 2 +- .../plugins/kilo-chat/src/synced/events.ts | 4 +- .../plugins/kilo-chat/src/synced/schemas.ts | 8 +- .../src/dos/UserConnectionDO.test.ts | 314 +++++++++++++++++- .../src/dos/UserConnectionDO.ts | 96 +++++- .../src/types/user-connection-protocol.ts | 6 +- 11 files changed, 629 insertions(+), 49 deletions(-) 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 77203ca4829..41857e54b4f 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.test.ts @@ -1005,7 +1005,7 @@ describe('CloudAgentTransport page-seam', () => { }); describe('watermark cursor', () => { - it('uses fromId on first connect when the page carries a watermark', async () => { + 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, @@ -1021,7 +1021,10 @@ describe('CloudAgentTransport page-seam', () => { expect(webSocketConstructor).toHaveBeenCalledTimes(1); const wsUrl = String(webSocketConstructor.mock.calls[0]?.[0]); - expect(wsUrl).toContain('fromId=42'); + // 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(); @@ -1087,9 +1090,10 @@ describe('CloudAgentTransport page-seam', () => { transport.connect(); await flushMicrotasks(); - // Verify first connect used the watermark + // 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=10'); + expect(firstUrl).toContain('fromId=0'); // Send a wire event with eventId > 0 to advance the live cursor const establish = { @@ -1127,7 +1131,7 @@ describe('CloudAgentTransport page-seam', () => { } }); - it('reconnect on a watermarked page with only sentinel events uses the original watermark and skips the page refetch', async () => { + 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 { @@ -1166,13 +1170,15 @@ describe('CloudAgentTransport page-seam', () => { transport.connect(); await flushMicrotasks(); - // Verify first connect used the watermark as fromId. + // 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=42'); + 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 the watermark value. + // live cursor, so lastEventId stays at 0. const sentinel: CloudAgentEvent = { eventId: 0, executionId: null, @@ -1196,12 +1202,12 @@ describe('CloudAgentTransport page-seam', () => { await flushMicrotasks(); const reconnectUrl = String(webSocketConstructor.mock.calls.at(-1)?.[0]); - // Reconnect must still use the watermark (42) because sentinel events - // never advance the cursor past the watermark. - expect(reconnectUrl).toContain('fromId=42'); + // 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 (watermark 42) is not null, so + // No second page fetch — lastEventId (0) is not null, so // onReconnected skips the snapshot fallback path entirely. expect(fetchSnapshotPage).toHaveBeenCalledTimes(1); @@ -1210,6 +1216,77 @@ describe('CloudAgentTransport page-seam', () => { 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 { + type: 'session.status'; + sessionId: string; + status: { type: 'retry'; attempt: number }; + } => e.type === 'session.status' && (e.status as { type?: string }).type === 'retry' + ) + .map(e => (e.status as { attempt: number }).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(); + }); }); }); diff --git a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts index 5f762e3ed59..8dce2fea12c 100644 --- a/packages/cloud-agent-sdk/src/cloud-agent-transport.ts +++ b/packages/cloud-agent-sdk/src/cloud-agent-transport.ts @@ -149,12 +149,14 @@ function createCloudAgentTransport(config: CloudAgentTransportConfig): Transport return null; } if (page.kind === 'success') { - // Seed lastEventId from the page's event-log watermark so the - // first WebSocket connect uses `fromId`. The DO replays any - // events persisted between the page snapshot and the socket - // open, closing the gap. On reconnect, wire events advance - // lastEventId and the live cursor takes over. - lastEventId = page.watermarkEventId ?? null; + // 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 }; diff --git a/packages/event-service/src/__tests__/client.test.ts b/packages/event-service/src/__tests__/client.test.ts index a83fc31ef8a..901b7fbfaae 100644 --- a/packages/event-service/src/__tests__/client.test.ts +++ b/packages/event-service/src/__tests__/client.test.ts @@ -914,6 +914,48 @@ describe('EventServiceClient', () => { 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(); diff --git a/packages/event-service/src/client.ts b/packages/event-service/src/client.ts index 583155ca712..e009b332c4a 100644 --- a/packages/event-service/src/client.ts +++ b/packages/event-service/src/client.ts @@ -446,18 +446,21 @@ export class EventServiceClient { // 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) { - if (message.seq > this.lastKnownSeq + 1) { - for (const handler of this.resyncHandlers) { - handler(); - } - } + 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') { 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 be2bfdf5c20..fcd89001d55 100644 --- a/services/event-service/src/__tests__/user-session-do.test.ts +++ b/services/event-service/src/__tests__/user-session-do.test.ts @@ -206,7 +206,7 @@ describe('UserSessionDO', () => { it( 'closes socket with code 4002 when unacked gap exceeds threshold', - { timeout: 15000 }, + { timeout: 30000 }, async () => { const userId = 'user-zombie-close'; const { ws, stub } = await connectWs(userId); @@ -222,15 +222,16 @@ describe('UserSessionDO', () => { ws.addEventListener('close', resolve); }); - // Push enough events to exceed UNACKED_THRESHOLD (50). + // 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 > 50. - // Push 55 events for margin. - for (let i = 0; i < 55; i++) { + // 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', {}); - // Brief pause to let the DO process the message and potential close - await new Promise(r => setTimeout(r, 20)); - // Check if we already closed + 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; } @@ -241,7 +242,7 @@ describe('UserSessionDO', () => { it( 'never closes a socket that has not acknowledged (old client compat)', - { timeout: 15000 }, + { timeout: 25000 }, async () => { const userId = 'user-no-ack-no-close'; const { ws, stub } = await connectWs(userId); @@ -254,10 +255,13 @@ describe('UserSessionDO', () => { closed = true; }); - // Push many events without any ack — should never trigger zombie close - for (let i = 0; i < 60; i++) { + // 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', {}); - await new Promise(r => setTimeout(r, 10)); + if (i % 100 === 0) { + await new Promise(r => setTimeout(r, 10)); + } if (closed) break; } @@ -266,6 +270,48 @@ describe('UserSessionDO', () => { } ); + 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); @@ -280,7 +326,7 @@ describe('UserSessionDO', () => { it( 'defaults counters to zero and emits sequence one with old attachment', - { timeout: 15000 }, + { timeout: 25000 }, async () => { // Load a WebSocket whose server-side attachment has only the old // shape: `{ contexts: ['project:old'] }`. The DO fetches a new @@ -324,9 +370,11 @@ describe('UserSessionDO', () => { closed = true; }); - for (let i = 0; i < 60; i++) { + for (let i = 0; i < 1100; i++) { await stub.pushEvent('project:old', 'fill', {}); - await new Promise(r => setTimeout(r, 10)); + if (i % 100 === 0) { + await new Promise(r => setTimeout(r, 10)); + } if (closed) break; } diff --git a/services/event-service/src/do/user-session-do.ts b/services/event-service/src/do/user-session-do.ts index c94bc8066cd..e39ef0ee906 100644 --- a/services/event-service/src/do/user-session-do.ts +++ b/services/event-service/src/do/user-session-do.ts @@ -4,7 +4,7 @@ import { configureDevLogging, logger, withLogTags } from '../util/logger'; import type { ServerMessage } from '../types'; /** Number of unacked events before a socket is closed as a zombie. */ -const UNACKED_THRESHOLD = 50; +const UNACKED_THRESHOLD = 1000; /** Close code for a socket that fell too far behind on acknowledgements. */ const ZOMBIE_CLOSE_CODE = 4002; 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 74671333c9a..40c06af54e9 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 @@ -5665,5 +5669,313 @@ describe('UserConnectionDO', () => { // 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')]); + 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); + + sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); + await flushAsync(); + + // Live response must carry the full result unchanged. + const live = parseSent(webWs); + expect(live.result).toEqual(oversized); + + // Durable entry must have the truncated marker, not the raw oversized result. + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + const durableResult = (entry as Record).result; + expect(durableResult).toBeDefined(); + expect(durableResult).toMatchObject({ _truncated: true }); + }); + + 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')]); + 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); + + sendCliResponse(doInstance, cliWs, { id: correlationId, result: exact }); + await flushAsync(); + + 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); + sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); + await flushAsync(); + + const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); + expect((entry as Record).state).toBe('done'); + expect((entry as Record).result).toMatchObject({ _truncated: true }); + }); + + // ------------------------------------------------------------------------- + // 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')]); + sendCommand(doInstance, webWs, { + id: 'dedup-cmd', + command: 'send_message', + sessionId: 's1', + }); + const correlationId = getCorrelationId(cliWs); + webWs.send.mockClear(); + + // First response — delivered live. + 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. + sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true } }); + await flushAsync(); + + 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')]); + 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. + sendCliResponse(doInstance, cliWs, { id: correlationId, result: { late: true } }); + await flushAsync(); + + // 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); + 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); + 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', () => { + 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(); + + sendCliResponse(doInstance, cliWs, { + id: correlationId, + result: { late: true }, + }); + await flushAsync(); + + // 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')]); + 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); + sendCliResponse(doInstance, cliWs, { id: correlationId, result }); + await flushAsync(); + + 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); + }); }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 7f9a54d4dd3..af30678969d 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -75,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([ @@ -165,6 +175,12 @@ const CLI_COMMAND_ERROR = { message: 'Command failed', }; +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; @@ -243,6 +259,11 @@ export class UserConnectionDO extends DurableObject { // 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 { @@ -807,9 +828,29 @@ export class UserConnectionDO extends DurableObject { // Step 24: on in-memory miss, try to load the durable entry (D8 case 1). if (!entry) { const durable = await this.getDurablePendingCommand(id); - if (!durable || durable.state !== 'pending') return; + if (!durable || durable.state !== 'pending') { + // If the live path already delivered this response and the durable + // write is still in-flight, skip the rehydration delivery. + if (durable && this.completedCorrelationIds.has(id)) return; + return; + } + // Prevent duplicate delivery: a concurrently completed live response + // already sent the reply; the durable write will update the state soon. + if (this.completedCorrelationIds.has(id)) 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, + }); + return; + } + // Validate the responding CLI by attachment connectionId. const respondingAttachment = respondingWs.deserializeAttachment() as WSAttachment | null; if ( @@ -827,10 +868,11 @@ export class UserConnectionDO extends DurableObject { // so a mutationId retry receives the same shaped outcome as // live delivery. const shaped = this.shapeCommandError(durable.command, result, error); + const shapedResult = this.boundDurableResult(shaped.result); await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { ...durable, state: 'done', - ...(shaped.result !== undefined ? { result: shaped.result } : {}), + ...(shapedResult !== undefined ? { result: shapedResult } : {}), ...(shaped.error !== undefined ? { error: shaped.error } : {}), }); return; @@ -861,12 +903,19 @@ export class UserConnectionDO extends DurableObject { // For non-rehydrated entries, send the oversized error synchronously // before the durable write so the mock harness sees the send. 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); this.sendToWeb(entry.ws, { type: 'response', id: entry.originalId, error: CATALOG_TOO_LARGE_ERROR, }); + this.completedCorrelationIds.add(id); } // Persist the terminal outcome durably. @@ -898,6 +947,9 @@ export class UserConnectionDO extends DurableObject { error: CATALOG_TOO_LARGE_ERROR, } satisfies PendingCommandEntry); + // Clean up the synchronous completion marker. + this.completedCorrelationIds.delete(id); + // For a rehydrated entry, re-resolve and send now. if (rehydrated) { const targetWeb = this.findWebByConnectionId(catWebConnectionId); @@ -943,7 +995,15 @@ export class UserConnectionDO extends DurableObject { // For non-rehydrated entries, send the live response synchronously // before the durable write so the response reaches the client without // yielding to an async storage put. + // + // 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); this.sendToWeb(entry.ws, { type: 'response', @@ -957,11 +1017,16 @@ export class UserConnectionDO extends DurableObject { ? { error: CLI_COMMAND_ERROR } : {}), }); + // Mark completed synchronously so a rehydration path does not + // deliver a duplicate response before the durable write lands. + this.completedCorrelationIds.add(id); } // Persist the durable entry. Storage write is extended by the // surrounding ctx.waitUntil; the live response has already been // delivered for the non-rehydrated case above. + // Bound the result so the durable put never exceeds the DO value limit. + const durableResult = this.boundDurableResult(result); const webConnectionId = rehydrated ? rehydratedDurable?.webConnectionId : (entry.ws.deserializeAttachment() as WSAttachment | null)?.role === 'web' @@ -976,10 +1041,13 @@ export class UserConnectionDO extends DurableObject { expiresAt: entry.expiresAt, webConnectionId: webConnectionId ?? 'unknown', state: 'done', - ...(result !== undefined ? { result } : {}), + ...(durableResult !== undefined ? { result: durableResult } : {}), ...(terminalError !== undefined ? { error: terminalError } : {}), } satisfies PendingCommandEntry); + // Clean up the synchronous completion marker after the durable write. + this.completedCorrelationIds.delete(id); + // For a rehydrated entry, re-resolve the web socket by the persisted // id and send the response now (the non-rehydrated path already sent // it synchronously above). @@ -1214,6 +1282,15 @@ export class UserConnectionDO extends DurableObject { // 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)) { @@ -1927,6 +2004,19 @@ export class UserConnectionDO extends DurableObject { }; } + /** + * 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 error 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 }; + } + // --------------------------------------------------------------------------- // Durable pending command storage (D8) // --------------------------------------------------------------------------- diff --git a/services/session-ingest/src/types/user-connection-protocol.ts b/services/session-ingest/src/types/user-connection-protocol.ts index 8b612788bff..51aaa6a396b 100644 --- a/services/session-ingest/src/types/user-connection-protocol.ts +++ b/services/session-ingest/src/types/user-connection-protocol.ts @@ -83,7 +83,8 @@ export const CLIInboundMessageSchema = z.discriminatedUnion('type', [ // 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. - mutationId: z.string().optional(), + // Bounded at 128 chars to keep storage keys safe. + mutationId: z.string().max(128).optional(), }), z.object({ type: z.literal('system'), @@ -117,7 +118,8 @@ export const WebOutboundMessageSchema = z.discriminatedUnion('type', [ // 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. - mutationId: z.string().optional(), + // Bounded at 128 chars to keep storage keys safe. + mutationId: z.string().max(128).optional(), }), z.object({ type: z.literal('ping'), From 3c3a5f74d2a549dadc42762736366a740ced753c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 05:46:31 +0200 Subject: [PATCH 16/24] test(stream): fix replay regression typing --- .../src/cloud-agent-transport.test.ts | 30 +++++++++---------- 1 file changed, 14 insertions(+), 16 deletions(-) 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 41857e54b4f..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'; @@ -1271,13 +1272,14 @@ describe('CloudAgentTransport page-seam', () => { .filter( ( e - ): e is { - type: 'session.status'; - sessionId: string; - status: { type: 'retry'; attempt: number }; - } => e.type === 'session.status' && (e.status as { type?: string }).type === 'retry' + ): e is Extract & { + status: Extract; + } => { + if (e.type !== 'session.status') return false; + return e.status.type === 'retry'; + } ) - .map(e => (e.status as { attempt: number }).attempt) + .map(e => e.status.attempt) .sort((a, b) => a - b); // Every event 1–10 must be delivered. The dedupe with lastEventId=0 @@ -1456,16 +1458,12 @@ describe('CloudAgentTransport exactly-once event replay', () => { .filter( ( e - ): e is { - type: 'session.status'; - sessionId: string; - status: { - type: 'retry'; - attempt: number; - message: string; - next: number; - }; - } => e.type === 'session.status' && e.status.type === 'retry' + ): 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); From b06820cf659b4542ea625ec40334aa8026747da4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 07:54:37 +0200 Subject: [PATCH 17/24] fix(cli): fence durable terminal replies --- .../src/dos/UserConnectionDO.test.ts | 593 ++++++++++++++---- .../src/dos/UserConnectionDO.ts | 468 ++++++++------ 2 files changed, 759 insertions(+), 302 deletions(-) diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 40c06af54e9..416dfe9527a 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -340,14 +340,17 @@ function sendCommand( return doInstance.webSocketMessage(webWs as never, msg); } -/** 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 }); - return doInstance.webSocketMessage(cliWs as never, msg); + doInstance.webSocketMessage(cliWs as never, msg); + await flushAsync(); } function createResultWithSerializedBytes(targetBytes: number): { @@ -523,7 +526,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'); @@ -556,7 +559,7 @@ describe('UserConnectionDO', () => { message: 'Session owner changed', }, }); - sendCliResponse(doInstance, firstOwner, { + await sendCliResponse(doInstance, firstOwner, { id: correlationId, result: 'late', }); @@ -1532,7 +1535,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({ @@ -1789,7 +1792,7 @@ 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'); @@ -1807,7 +1810,7 @@ describe('UserConnectionDO', () => { 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' }); }); }); @@ -1843,7 +1846,7 @@ 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'); @@ -1859,7 +1862,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { success: true }, }); @@ -1872,7 +1875,7 @@ 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'); @@ -1887,7 +1890,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: { source: 'relay', @@ -1906,7 +1909,7 @@ 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'); @@ -1921,7 +1924,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_models', }); @@ -1933,7 +1936,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'); @@ -1950,13 +1953,13 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(targetCli); webWs.send.mockClear(); - sendCliResponse(doInstance, otherCli, { + await sendCliResponse(doInstance, otherCli, { id: correlationId, result: 'wrong-owner', }); expect(webWs.send).not.toHaveBeenCalled(); - sendCliResponse(doInstance, targetCli, { + await sendCliResponse(doInstance, targetCli, { id: correlationId, result: 'ok', }); @@ -2090,7 +2093,7 @@ 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); }); @@ -2128,7 +2131,7 @@ 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'); @@ -2145,7 +2148,7 @@ 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', @@ -2154,7 +2157,7 @@ describe('UserConnectionDO', () => { }); }); - 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'); @@ -2170,7 +2173,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), }); @@ -2186,7 +2189,7 @@ 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'); @@ -2202,7 +2205,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: createUtf8OversizedResult(), }); @@ -2290,7 +2293,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'); @@ -2319,8 +2322,8 @@ 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', @@ -2473,7 +2476,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'); @@ -2495,7 +2498,7 @@ 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); }); @@ -2593,7 +2596,7 @@ 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'); @@ -2610,7 +2613,7 @@ 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', @@ -2619,7 +2622,7 @@ describe('UserConnectionDO', () => { }); }); - 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'); @@ -2635,7 +2638,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), }); @@ -2651,7 +2654,7 @@ 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'); @@ -2667,7 +2670,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: createUtf8OversizedResult(), }); @@ -2689,7 +2692,7 @@ 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'); @@ -2705,7 +2708,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_commands', }); @@ -2721,7 +2724,7 @@ 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'); @@ -2738,7 +2741,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: send_command', }); @@ -2754,7 +2757,7 @@ 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'); @@ -2771,7 +2774,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: exit_cli', }); @@ -2787,7 +2790,7 @@ 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'); @@ -2803,7 +2806,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: create_session', }); @@ -2820,7 +2823,7 @@ 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'); @@ -2835,7 +2838,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_models', }); @@ -2847,7 +2850,7 @@ 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'); @@ -2863,7 +2866,7 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'session not ready', }); @@ -2875,7 +2878,7 @@ 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'); @@ -2891,7 +2894,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', }); @@ -2904,7 +2907,7 @@ 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'); @@ -2921,7 +2924,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', }); @@ -3019,7 +3022,7 @@ 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'); @@ -3037,7 +3040,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', @@ -3046,7 +3049,7 @@ 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'); @@ -3063,7 +3066,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', @@ -3235,7 +3238,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'); @@ -3253,7 +3256,7 @@ 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', @@ -3262,7 +3265,7 @@ describe('UserConnectionDO', () => { }); }); - 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'); @@ -3302,7 +3305,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 ); @@ -4221,13 +4224,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' }); }); }); @@ -4874,11 +4877,10 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); // Send the CLI response with the correlation id. - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true }, }); - await flushAsync(); // The response must be routed to the web socket. const responses = allSent(webWs).filter(m => m.type === 'response'); @@ -4913,11 +4915,10 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses-b', 'busy', 'Session B')]); // Send the CLI response. - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true }, }); - await flushAsync(); // The durable entry must be marked 'done'. const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); @@ -5089,11 +5090,10 @@ describe('UserConnectionDO', () => { expect(entry).toBeDefined(); // Send CLI response to verify normal resolution. - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { normal: true }, }); - await flushAsync(); const responses = allSent(webWs).filter(m => m.type === 'response'); expect(responses.some(r => r.id === 'no-mut-req' && r.result)).toBe(true); @@ -5244,11 +5244,10 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_models', }); - await flushAsync(); // Live response carries the original string. const live = parseSent(webWs) as { error: unknown }; @@ -5275,11 +5274,10 @@ describe('UserConnectionDO', () => { const correlationId = getCorrelationId(cliWs); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_commands', }); - await flushAsync(); // Live response carries the structured error. const live = parseSent(webWs) as { error: unknown }; @@ -5318,11 +5316,10 @@ describe('UserConnectionDO', () => { // 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. - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { durable: true }, }); - await flushAsync(); // The durable entry must now be written. const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); @@ -5358,11 +5355,10 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized, }); - await flushAsync(); // The live web socket must receive the oversized error. const responses = allSent(webWs).filter(m => m.type === 'response'); @@ -5457,11 +5453,10 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses-nw', 'busy', 'Session NW')]); // Send the CLI response with an upgrade-required error. - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_commands', }); - await flushAsync(); // The durable entry must be marked 'done' with the shaped structured error. const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); @@ -5494,11 +5489,10 @@ describe('UserConnectionDO', () => { const cliWs = addCliSocket(mockCtx, 'cli-nws'); sendHeartbeat(doInstance, cliWs, [makeSession('ses-nws', 'busy', 'Session NWS')]); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: 'unknown command: list_models', }); - await flushAsync(); const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); expect((entry as Record).state).toBe('done'); @@ -5527,11 +5521,10 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses-nwo', 'busy', 'Session NWO')]); const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized, }); - await flushAsync(); const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); expect((entry as Record).state).toBe('done'); @@ -5564,7 +5557,7 @@ describe('UserConnectionDO', () => { const cliWs = addCliSocket(mockCtx, 'cli-nwr'); sendHeartbeat(doInstance, cliWs, [makeSession('ses-nwr', 'busy', 'Session NWR')]); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, error: { source: 'relay', @@ -5572,7 +5565,6 @@ describe('UserConnectionDO', () => { message: 'Session owner changed', }, }); - await flushAsync(); const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); expect((entry as Record).state).toBe('done'); @@ -5603,12 +5595,11 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses-nwc', 'busy', 'Session NWC')]); // CLI sends both result and error. - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { partial: 'data' }, error: 'something went wrong', }); - await flushAsync(); const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); expect((entry as Record).state).toBe('done'); @@ -5641,11 +5632,10 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); const oversized = createResultWithSerializedBytes(MAX_CATALOG_RESULT_BYTES + 1); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized, }); - await flushAsync(); // The live web socket must receive the oversized error. const responses = allSent(webWs).filter(m => m.type === 'response'); @@ -5690,11 +5680,10 @@ describe('UserConnectionDO', () => { // Create a result just over the durable limit. const oversized = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES + 1); - sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); - await flushAsync(); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); // Live response must carry the full result unchanged. - const live = parseSent(webWs); + const live = parseSent(webWs) as { id: string; result: unknown }; expect(live.result).toEqual(oversized); // Durable entry must have the truncated marker, not the raw oversized result. @@ -5721,8 +5710,7 @@ describe('UserConnectionDO', () => { const exact = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES); - sendCliResponse(doInstance, cliWs, { id: correlationId, result: exact }); - await flushAsync(); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: exact }); const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); expect((entry as Record).result).toEqual(exact); @@ -5748,8 +5736,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('ses-big', 'busy', 'Session Big')]); const oversized = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES + 1); - sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); - await flushAsync(); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: oversized }); const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); expect((entry as Record).state).toBe('done'); @@ -5774,7 +5761,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); // First response — delivered live. - sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true } }); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true } }); const firstResponses = allSent(webWs).filter( m => m.type === 'response' && m.id === 'dedup-cmd' @@ -5783,8 +5770,7 @@ describe('UserConnectionDO', () => { expect(firstResponses[0].result).toEqual({ ok: true }); // Second response with same correlationId — must NOT deliver again. - sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true } }); - await flushAsync(); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { ok: true } }); const allResponses = allSent(webWs).filter( m => m.type === 'response' && m.id === 'dedup-cmd' @@ -5816,8 +5802,7 @@ describe('UserConnectionDO', () => { vi.spyOn(Date, 'now').mockReturnValue(baseTime + 35_001); // CLI sends response — should be dropped because it's expired. - sendCliResponse(doInstance, cliWs, { id: correlationId, result: { late: true } }); - await flushAsync(); + 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'); @@ -5827,7 +5812,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- // Fix: Oversized mutationId rejected // ------------------------------------------------------------------------- - it('rejects a mutationId longer than 128 characters via schema validation', async () => { + it('rejects a mutationId longer than 128 characters via schema validation', () => { const { doInstance, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); @@ -5926,11 +5911,10 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-re'); webWs.send.mockClear(); - sendCliResponse(doInstance, cliWs, { + await sendCliResponse(doInstance, cliWs, { id: correlationId, result: { late: true }, }); - await flushAsync(); // No response must reach the web socket. const responses = allSent(webWs).filter(m => m.type === 'response' && m.id === 'original-re'); @@ -5966,8 +5950,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); const result = createResultWithSerializedBytes(MAX_DURABLE_RESULT_BYTES); - sendCliResponse(doInstance, cliWs, { id: correlationId, result }); - await flushAsync(); + await sendCliResponse(doInstance, cliWs, { id: correlationId, result }); const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); expect((entry as Record).state).toBe('done'); @@ -5977,5 +5960,405 @@ describe('UserConnectionDO', () => { 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')]); + 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')]); + 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')]); + 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')]); + 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')]); + 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); + }); + + // ------------------------------------------------------------------------- + // 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 }); + }); }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index af30678969d..aae7f2a2a3c 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -524,7 +524,16 @@ export class UserConnectionDO extends DurableObject { // Extend the DO lifetime for the async durable write (fix: reply // must persist before hibernation, but the surrounding handler stays // synchronous for existing callers). - this.ctx.waitUntil(this.handleCliResponse(ws, msg.id, msg.result, msg.error)); + // 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; } } @@ -824,237 +833,280 @@ export class UserConnectionDO extends DurableObject { // 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) { - const durable = await this.getDurablePendingCommand(id); - if (!durable || durable.state !== 'pending') { - // If the live path already delivered this response and the durable - // write is still in-flight, skip the rehydration delivery. - if (durable && this.completedCorrelationIds.has(id)) return; - return; - } - // Prevent duplicate delivery: a concurrently completed live response - // already sent the reply; the durable write will update the state soon. + // 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; - rehydratedDurable = durable; + this.completedCorrelationIds.add(id); + ownsReservation = true; - // 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, - }); - 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 - ) { - 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); - const shapedResult = this.boundDurableResult(shaped.result); - await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { - ...durable, - state: 'done', - ...(shapedResult !== undefined ? { result: shapedResult } : {}), - ...(shaped.error !== undefined ? { error: shaped.error } : {}), - }); - 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); + const shapedResult = this.boundDurableResult(shaped.result); + await this.ctx.storage.put(`${PENDING_COMMAND_KEY_PREFIX}${id}`, { + ...durable, + state: 'done', + ...(shapedResult !== undefined ? { result: shapedResult } : {}), + ...(shaped.error !== undefined ? { error: shaped.error } : {}), + }); + // 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; + // 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; } - // 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, send the oversized error synchronously - // before the durable write so the mock harness sees the send. - if (!rehydrated) { - // Guard: do not override an already-expired terminal outcome. - const now = Date.now(); - if (entry.expiresAt <= now) { + // 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); - return; + // 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; } - this.pendingCommands.delete(id); - this.sendToWeb(entry.ws, { - type: 'response', - id: entry.originalId, - error: CATALOG_TOO_LARGE_ERROR, - }); - this.completedCorrelationIds.add(id); - } - // Persist the terminal outcome durably. - 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'; - } - 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); + // 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; + } - // Clean up the synchronous completion marker. - this.completedCorrelationIds.delete(id); + // Durable write succeeded — clear the marker. The durable 'done' + // state now guards against duplicate delivery. + this.completedCorrelationIds.delete(id); + ownsReservation = false; - // For a rehydrated entry, re-resolve and send now. - if (rehydrated) { - const targetWeb = this.findWebByConnectionId(catWebConnectionId); - if (targetWeb) { - this.sendToWeb(targetWeb, { + // 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; } - 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 { + 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) { - sanitizeAsFailed = true; - } - const terminalError = - structuredError ?? stringError ?? (sanitizeAsFailed ? CLI_COMMAND_ERROR : undefined); + const terminalError = + structuredError ?? stringError ?? (sanitizeAsFailed ? CLI_COMMAND_ERROR : undefined); - // Step 26a: write terminal outcome durably, always. - // For non-rehydrated entries, send the live response synchronously - // before the durable write so the response reaches the client without - // yielding to an async storage put. - // - // 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) { + // 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); - return; + // 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; } - this.pendingCommands.delete(id); - 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 } - : {}), - }); - // Mark completed synchronously so a rehydration path does not - // deliver a duplicate response before the durable write lands. - this.completedCorrelationIds.add(id); - } - // Persist the durable entry. Storage write is extended by the - // surrounding ctx.waitUntil; the live response has already been - // delivered for the non-rehydrated case above. - // Bound the result so the durable put never exceeds the DO value limit. - const durableResult = this.boundDurableResult(result); - const webConnectionId = rehydrated - ? rehydratedDurable?.webConnectionId - : (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 } : {}), - ...(terminalError !== undefined ? { error: terminalError } : {}), - } satisfies PendingCommandEntry); + // 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); + 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 } : {}), + ...(terminalError !== undefined ? { error: terminalError } : {}), + } 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; + } - // Clean up the synchronous completion marker after the durable write. - this.completedCorrelationIds.delete(id); + // 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; - // For a rehydrated entry, re-resolve the web socket by the persisted - // id and send the response now (the non-rehydrated path already sent - // it synchronously above). - if (rehydrated) { - const targetWebWs = this.findWebByConnectionId(webConnectionId ?? 'unknown'); - if (targetWebWs) { - this.sendToWeb(targetWebWs, { + 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 } : {}), @@ -1066,6 +1118,28 @@ export class UserConnectionDO extends DurableObject { ? { 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); } } } From 15d6133fe080bd0520614f510b4ec4aac49c48e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 14:13:26 +0200 Subject: [PATCH 18/24] fix(session-ingest): harden durable terminal delivery Persist terminal outcomes before delivery and prevent duplicate responses during CLI disconnect and ownership changes. --- .../src/dos/UserConnectionDO.test.ts | 776 +++++++++++++----- .../src/dos/UserConnectionDO.ts | 162 +++- 2 files changed, 704 insertions(+), 234 deletions(-) diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 416dfe9527a..5e3d3c5c589 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -323,8 +323,9 @@ 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: { @@ -335,9 +336,10 @@ function sendCommand( data?: unknown; mutationId?: string; } -) { +): Promise { const msg = JSON.stringify({ type: 'command', ...opts }); - return doInstance.webSocketMessage(webWs as never, msg); + doInstance.webSocketMessage(webWs as never, msg); + await flushAsync(); } /** Send a response from a CLI ws. Durable-before-send means the live @@ -474,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')]); @@ -487,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'); @@ -500,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'); @@ -513,7 +515,7 @@ describe('UserConnectionDO', () => { // Verify via command routing: command to s2 should fail (no owner) const webWs2 = addWebSocket(mockCtx, 'web-2'); - sendCommand(doInstance, webWs2, { + await sendCommand(doInstance, webWs2, { id: 'cmd-1', command: 'send_message', sessionId: 's2', @@ -536,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', @@ -546,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 @@ -570,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'); @@ -594,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'); @@ -629,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'); @@ -661,7 +666,7 @@ 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'); @@ -683,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'); @@ -698,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. @@ -728,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'); @@ -744,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'); @@ -752,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'); @@ -770,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'); @@ -785,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'); @@ -832,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'); @@ -848,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'); @@ -866,7 +871,7 @@ 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'); @@ -887,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'); @@ -905,7 +910,7 @@ 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'); @@ -925,7 +930,7 @@ describe('UserConnectionDO', () => { }); }); - 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'); @@ -939,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. @@ -1031,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'); @@ -1047,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'); @@ -1072,7 +1077,7 @@ describe('UserConnectionDO', () => { }); }); - 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'); @@ -1089,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'); @@ -1103,7 +1108,7 @@ describe('UserConnectionDO', () => { 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'); @@ -1123,7 +1128,7 @@ describe('UserConnectionDO', () => { }); }); - 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'); @@ -1155,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'); @@ -1212,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')]); @@ -1236,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'); @@ -1273,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'); @@ -1314,7 +1319,7 @@ describe('UserConnectionDO', () => { // Session no longer routable const web2 = addWebSocket(mockCtx, 'web-2'); - sendCommand(doInstance, web2, { + await sendCommand(doInstance, web2, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -1333,7 +1338,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); // Send command from web - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -1362,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', @@ -1384,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', @@ -1421,7 +1426,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', @@ -1451,7 +1456,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, []); // Send command with no sessionId or connectionId (fallback routing) - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-fallback', command: 'send_message', }); @@ -1514,7 +1519,7 @@ describe('UserConnectionDO', () => { webWs.send.mockClear(); // Web sends a command targeting s1 — should route to cli2 (the replacement) - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-new', command: 'send_message', sessionId: 's1', @@ -1553,7 +1558,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cli1, [makeSession('s1')]); // Web sends a command that gets forwarded to cli1 - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -1751,7 +1756,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'); @@ -1774,7 +1779,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'); @@ -1799,7 +1804,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', @@ -1819,7 +1824,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'); @@ -1827,7 +1832,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', @@ -1853,7 +1858,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', @@ -1882,7 +1887,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', @@ -1916,7 +1921,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -1945,7 +1950,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, targetCli, [makeSession('s1')]); sendHeartbeat(doInstance, otherCli, []); targetCli.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -1970,7 +1975,7 @@ describe('UserConnectionDO', () => { }); }); - 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'); @@ -1978,13 +1983,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', @@ -2003,7 +2008,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(); @@ -2013,7 +2018,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', @@ -2021,7 +2026,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', @@ -2040,7 +2045,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(); @@ -2048,7 +2053,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: 'send_message', sessionId: 's1', @@ -2070,7 +2075,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', @@ -2097,7 +2102,7 @@ describe('UserConnectionDO', () => { 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'); @@ -2106,14 +2111,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', @@ -2138,7 +2143,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -2164,7 +2169,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -2196,7 +2201,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -2221,11 +2226,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', @@ -2239,7 +2244,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'); @@ -2251,7 +2256,7 @@ describe('UserConnectionDO', () => { staleOwner.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -2271,7 +2276,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'); @@ -2283,7 +2288,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', @@ -2303,14 +2308,14 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); // Both web sockets send commands with the same id - sendCommand(doInstance, web1, { + await sendCommand(doInstance, web1, { id: 'dup-id', command: 'send_message', sessionId: 's1', }); const corr1 = getCorrelationId(cliWs, 0); - sendCommand(doInstance, web2, { + await sendCommand(doInstance, web2, { id: 'dup-id', command: 'send_message', sessionId: 's1', @@ -2337,7 +2342,7 @@ describe('UserConnectionDO', () => { }); }); - 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'); @@ -2345,7 +2350,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); }); }); @@ -2370,7 +2375,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'); @@ -2379,7 +2384,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command, sessionId: 's1', @@ -2398,7 +2403,7 @@ describe('UserConnectionDO', () => { } }); - 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'); @@ -2407,7 +2412,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'eval', sessionId: 's1', @@ -2425,7 +2430,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'); @@ -2434,7 +2439,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'shell', sessionId: 's1', @@ -2454,11 +2459,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', @@ -2485,7 +2490,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'eval', sessionId: 's1', @@ -2502,7 +2507,7 @@ describe('UserConnectionDO', () => { 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'); @@ -2514,7 +2519,7 @@ describe('UserConnectionDO', () => { staleOwner.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -2540,7 +2545,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'); @@ -2548,13 +2553,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', @@ -2573,20 +2578,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', @@ -2603,7 +2608,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2629,7 +2634,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2661,7 +2666,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2699,7 +2704,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2731,7 +2736,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_command', sessionId: 's1', @@ -2764,7 +2769,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -2797,7 +2802,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'create_session', connectionId: 'cli-1', @@ -2830,7 +2835,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_models', sessionId: 's1', @@ -2857,7 +2862,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_command', sessionId: 's1', @@ -2885,7 +2890,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'list_commands', sessionId: 's1', @@ -2914,7 +2919,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -2944,7 +2949,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'); @@ -2952,14 +2957,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', @@ -2985,7 +2990,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'); @@ -2993,13 +2998,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', @@ -3029,7 +3034,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'send_command', sessionId: 's1', @@ -3056,7 +3061,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, []); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'create_session', connectionId: 'cli-1', @@ -3101,7 +3106,7 @@ describe('UserConnectionDO', () => { label: 'primitive data', input: { sessionId: 's1', data: 'protocolVersion=1' }, }, - ])('rejects $label before routing or pending allocation', ({ input }) => { + ])('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'); @@ -3109,7 +3114,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, @@ -3128,7 +3133,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'); @@ -3139,7 +3144,7 @@ describe('UserConnectionDO', () => { selectedOwner.send.mockClear(); otherCli.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3158,7 +3163,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'); @@ -3170,7 +3175,7 @@ describe('UserConnectionDO', () => { staleOwner.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-1', command: 'exit_cli', sessionId: 's1', @@ -3191,11 +3196,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', @@ -3209,7 +3214,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'); @@ -3217,14 +3222,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', @@ -3246,7 +3251,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', @@ -3273,7 +3278,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', @@ -3317,7 +3322,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', @@ -3338,7 +3343,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'); @@ -3346,7 +3351,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', @@ -3368,7 +3373,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'); @@ -3378,7 +3383,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', @@ -3400,14 +3405,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', @@ -3436,7 +3441,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'); @@ -3466,7 +3471,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'); @@ -3500,7 +3505,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'); @@ -3523,7 +3528,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'); @@ -3552,7 +3557,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'); @@ -3574,7 +3579,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'); @@ -3600,7 +3605,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'); @@ -3627,7 +3632,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'); @@ -3658,7 +3663,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 @@ -3671,7 +3676,7 @@ describe('UserConnectionDO', () => { // Verify state was reconstructed by routing a command const web2 = addWebSocket(mockCtx, 'web-2'); - sendCommand(doInstance, web2, { + await sendCommand(doInstance, web2, { id: 'cmd-1', command: 'send_message', sessionId: 's1', @@ -3688,7 +3693,7 @@ describe('UserConnectionDO', () => { }); }); - 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')]); @@ -3712,7 +3717,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']); @@ -3742,7 +3747,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'); @@ -3758,7 +3763,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'); @@ -3778,7 +3783,7 @@ describe('UserConnectionDO', () => { ]); }); - 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')]); @@ -3790,7 +3795,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'); @@ -3810,7 +3815,7 @@ describe('UserConnectionDO', () => { ]); }); - 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'); @@ -3834,7 +3839,7 @@ describe('UserConnectionDO', () => { ]); }); - 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'); @@ -3856,7 +3861,7 @@ describe('UserConnectionDO', () => { ]); }); - 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'); @@ -3876,7 +3881,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 @@ -3907,7 +3912,7 @@ describe('UserConnectionDO', () => { ); }); - 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', @@ -3919,7 +3924,7 @@ describe('UserConnectionDO', () => { 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'); @@ -3929,7 +3934,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. @@ -3944,7 +3949,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. @@ -3965,7 +3970,7 @@ describe('UserConnectionDO', () => { ]); }); - 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. @@ -3989,7 +3994,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', @@ -4007,7 +4012,7 @@ describe('UserConnectionDO', () => { 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, [], { @@ -4025,7 +4030,7 @@ describe('UserConnectionDO', () => { }); }); - 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. @@ -4057,7 +4062,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), @@ -4098,7 +4103,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'); @@ -4135,7 +4140,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'); @@ -4143,7 +4148,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); @@ -4159,7 +4164,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); @@ -4169,7 +4174,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 @@ -4185,7 +4190,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); @@ -4660,7 +4665,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(); @@ -4949,7 +4954,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); // Send a second command with the same mutationId. - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'second-req', command: 'send_message', sessionId: 'ses-c', @@ -4992,7 +4997,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); // Send a retry with the same mutationId but new wire id. - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'retry-req-id', command: 'send_message', sessionId: 'ses-d', @@ -5040,7 +5045,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); // Retry with the same mutationId but new wire id. - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'retry-combined-id', command: 'send_message', sessionId: 'ses-dc', @@ -5069,7 +5074,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-5'); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'no-mut-req', command: 'send_message', sessionId: 'ses-e', @@ -5107,7 +5112,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-6'); cliWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'disconnect-req', command: 'send_message', sessionId: 'ses-f', @@ -5194,7 +5199,7 @@ describe('UserConnectionDO', () => { } // The first command (no mutationId) is just within the cap: 0 in-memory + 1 new = 1. - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'at-cap', command: 'send_message', sessionId: 'ses-h', @@ -5209,7 +5214,7 @@ describe('UserConnectionDO', () => { // command's durable write + 1 in-memory = 129 ≥ 128. Rejected. cliWs.send.mockClear(); webWs.send.mockClear(); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'over-cap', command: 'send_message', sessionId: 'ses-h', @@ -5236,7 +5241,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'err-1', command: 'list_models', sessionId: 's1', @@ -5265,7 +5270,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'err-2', command: 'list_commands', sessionId: 's1', @@ -5306,7 +5311,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'dur-1', command: 'send_message', sessionId: 's1', @@ -5388,7 +5393,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-disco', command: 'send_message', sessionId: 's1', @@ -5411,7 +5416,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-dur-disco', command: 'send_message', sessionId: 's1', @@ -5669,7 +5674,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'big-result', command: 'send_message', sessionId: 's1', @@ -5686,12 +5691,15 @@ describe('UserConnectionDO', () => { const live = parseSent(webWs) as { id: string; result: unknown }; expect(live.result).toEqual(oversized); - // Durable entry must have the truncated marker, not the raw oversized result. + // 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'); - const durableResult = (entry as Record).result; - expect(durableResult).toBeDefined(); - expect(durableResult).toMatchObject({ _truncated: true }); + 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 () => { @@ -5700,7 +5708,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'exact-result', command: 'send_message', sessionId: 's1', @@ -5740,7 +5748,13 @@ describe('UserConnectionDO', () => { const entry = await ctx.storage.get(`pendingCommand/${correlationId}`); expect((entry as Record).state).toBe('done'); - expect((entry as Record).result).toMatchObject({ _truncated: true }); + // 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', + }); }); // ------------------------------------------------------------------------- @@ -5752,7 +5766,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'dedup-cmd', command: 'send_message', sessionId: 's1', @@ -5790,7 +5804,7 @@ describe('UserConnectionDO', () => { vi.spyOn(Date, 'now').mockReturnValue(baseTime); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'expiry-cmd', command: 'send_message', sessionId: 's1', @@ -5812,7 +5826,7 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- // Fix: Oversized mutationId rejected // ------------------------------------------------------------------------- - it('rejects a mutationId longer than 128 characters via schema validation', () => { + 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'); @@ -5821,7 +5835,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); const longMutationId = 'x'.repeat(129); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'oversized-mut', command: 'send_message', sessionId: 's1', @@ -5850,7 +5864,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); const exactMutationId = 'x'.repeat(128); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'exact-mut', command: 'send_message', sessionId: 's1', @@ -5864,7 +5878,7 @@ describe('UserConnectionDO', () => { expect(cliCommands[0].mutationId).toBe(exactMutationId); }); - it('rejects a mutationId longer than 128 characters in the web message schema', () => { + 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'); @@ -5941,7 +5955,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'safety-check', command: 'send_message', sessionId: 's1', @@ -6045,7 +6059,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-fail-write', command: 'send_message', sessionId: 's1', @@ -6086,7 +6100,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-catalog', command: 'list_models', sessionId: 's1', @@ -6123,7 +6137,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'retry-after-fail', command: 'send_message', sessionId: 's1', @@ -6177,7 +6191,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'cmd-clear', command: 'send_message', sessionId: 's1', @@ -6213,7 +6227,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendCommand(doInstance, webWs, { + await sendCommand(doInstance, webWs, { id: 'order-cmd', command: 'send_message', sessionId: 's1', @@ -6360,5 +6374,361 @@ describe('UserConnectionDO', () => { 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(...args); + }); + // Spy on cliWs.send to track order. + const originalSend = cliWs.send; + cliWs.send = vi.fn((...args: unknown[]) => { + callOrder.push('send'); + return originalSend.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. + sendHeartbeat(doInstance, nextOwner, [makeSession('s1')]); + + // 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(...args); + }); + const originalSend = webWs.send; + webWs.send = vi.fn((...args: unknown[]) => { + callOrder.push('send'); + return originalSend.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); + }); }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index aae7f2a2a3c..192fd922af7 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -175,6 +175,12 @@ 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', @@ -887,12 +893,19 @@ export class UserConnectionDO extends DurableObject { // so a mutationId retry receives the same shaped outcome as // live delivery. const shaped = this.shapeCommandError(durable.command, result, error); - const shapedResult = this.boundDurableResult(shaped.result); + 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 } : {}), - ...(shaped.error !== undefined ? { error: shaped.error } : {}), + ...(shapedError !== undefined ? { error: shapedError } : {}), }); // Reservation held through the terminal write; clear after success. this.completedCorrelationIds.delete(id); @@ -1071,6 +1084,15 @@ export class UserConnectionDO extends DurableObject { 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' @@ -1087,7 +1109,7 @@ export class UserConnectionDO extends DurableObject { webConnectionId: webConnectionId ?? 'unknown', state: 'done', ...(durableResult !== undefined ? { result: durableResult } : {}), - ...(terminalError !== undefined ? { error: terminalError } : {}), + ...(durableError !== undefined ? { error: durableError } : {}), } satisfies PendingCommandEntry); } catch (_error: unknown) { // Clean the synchronous reservation on storage failure — the error @@ -1556,26 +1578,35 @@ export class UserConnectionDO extends DurableObject { targetCliWs: targetCli, }); this.ctx.waitUntil( - 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.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) + .then(() => { + 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) => { + console.error('Failed to persist durable pending command before forwarding', { + correlationId, + error: error instanceof Error ? error.message : String(error), + }); + }) ); - this.scheduleNextAlarm(now); - this.scheduleDurablePendingAlarm(); - this.sendToCli(targetCli, { - type: 'command', - id: correlationId, - command: msg.command, - data: msg.data, - ...(msg.sessionId ? { sessionId: msg.sessionId } : {}), - }); } // --------------------------------------------------------------------------- @@ -1920,8 +1951,12 @@ export class UserConnectionDO extends DurableObject { } private failPendingCommandsForSocket(targetWs: WebSocket, cliGone: boolean): void { + // Collect correlationIds handled in the in-memory sweep so the durable + // sweep in finishDurablePendingCommands does not double-deliver. + const handledIds = new Set(); for (const [id, entry] of this.pendingCommands) { if (entry.targetCliWs !== targetWs) continue; + handledIds.add(id); // 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'; @@ -1968,7 +2003,8 @@ export class UserConnectionDO extends DurableObject { this.ctx.waitUntil( this.finishDurablePendingCommands( entry => entry.targetConnectionId === attachment.connectionId, - 'CLI disconnected' + 'CLI disconnected', + handledIds ) ); } @@ -1978,10 +2014,12 @@ export class UserConnectionDO extends DurableObject { sessionId: string, nextOwnerConnectionId: string | undefined ): void { + const handledIds = new Set(); for (const [id, entry] of this.pendingCommands) { if (entry.sessionId !== sessionId || entry.targetConnectionId === nextOwnerConnectionId) { continue; } + handledIds.add(id); this.pendingCommands.delete(id); // `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. @@ -2018,7 +2056,8 @@ export class UserConnectionDO extends DurableObject { this.finishDurablePendingCommands( entry => entry.sessionId === sessionId && entry.targetConnectionId !== nextOwnerConnectionId, - SESSION_OWNER_CHANGED_ERROR + SESSION_OWNER_CHANGED_ERROR, + handledIds ) ); } @@ -2081,7 +2120,7 @@ export class UserConnectionDO extends DurableObject { /** * 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 error marker otherwise. + * threshold, or a size-exceeded marker otherwise. */ private boundDurableResult(result: unknown): unknown { if (result === undefined) return undefined; @@ -2091,6 +2130,21 @@ export class UserConnectionDO extends DurableObject { 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) // --------------------------------------------------------------------------- @@ -2122,17 +2176,46 @@ export class UserConnectionDO extends DurableObject { private async finishDurablePendingCommands( matches: (entry: PendingCommandEntry) => boolean, - error: unknown + error: unknown, + skipIds?: ReadonlySet ): Promise { const entries = await this.ctx.storage.list({ prefix: PENDING_COMMAND_KEY_PREFIX, }); await Promise.all( - [...entries].flatMap(([key, entry]) => - entry.state === 'pending' && matches(entry) - ? [this.ctx.storage.put(key, { ...entry, state: 'done', error })] - : [] - ) + [...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 []; + // Persist the terminal outcome durably BEFORE any live delivery. + // If storage.put fails, no live send occurs — the error propagates + // to the waitUntil boundary and the caller handles the rejection. + const putPromise = this.ctx.storage.put(key, { + ...entry, + state: 'done', + error, + }); + // 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); + if (webWs) { + return [ + putPromise.then(() => { + this.sendToWeb(webWs, { + type: 'response', + id: entry.originalId, + error, + }); + }), + ]; + } + return [putPromise]; + }) ); } @@ -2160,9 +2243,13 @@ export class UserConnectionDO extends DurableObject { } 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); + deliveredInMemory.add(id); this.sendToWeb(entry.ws, { type: 'response', id: entry.originalId, @@ -2189,6 +2276,19 @@ export class UserConnectionDO extends DurableObject { state: 'done', error: COMMAND_EXPIRED_ERROR, }); + // Deliver the terminal response to a live matching web socket + // that reconnected after the original socket disconnected (D8 case 3). + // Skip entries already delivered in the in-memory sweep above. + const correlationId = key.slice(PENDING_COMMAND_KEY_PREFIX.length); + if (deliveredInMemory.has(correlationId)) continue; + 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. await this.ctx.storage.delete(key); From fe321375b16ba4173c1b4b8c69a7f94e36f8c327 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 14:19:40 +0200 Subject: [PATCH 19/24] test(session-ingest): type durable delivery spies Keep durable write ordering assertions type-safe in the session ingest tests. --- services/session-ingest/src/dos/UserConnectionDO.test.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 5e3d3c5c589..a0986bfdf34 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -6564,13 +6564,13 @@ describe('UserConnectionDO', () => { const originalPut = ctx.storage.put; ctx.storage.put = vi.fn(async (...args: unknown[]) => { callOrder.push('put'); - return originalPut(...args); + 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.apply(cliWs, args); + return (originalSend as any).apply(cliWs, args); }); await sendCommand(doInstance, webWs, { @@ -6698,12 +6698,12 @@ describe('UserConnectionDO', () => { const originalPut = ctx.storage.put; ctx.storage.put = vi.fn(async (...args: unknown[]) => { callOrder.push('put'); - return originalPut(...args); + return (originalPut as any)(...args); }); const originalSend = webWs.send; webWs.send = vi.fn((...args: unknown[]) => { callOrder.push('send'); - return originalSend.apply(webWs, args); + return (originalSend as any).apply(webWs, args); }); // Disconnect the CLI — finishDurablePendingCommands handles the From 57e4530d5145cacc05d01194c8cddc8e48d21c37 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 19:14:48 +0200 Subject: [PATCH 20/24] fix(session-ingest): fence terminal command writes Keep terminal command outcomes durable while delayed initial writes settle. --- .../src/dos/UserConnectionDO.test.ts | 442 ++++++++++++++++++ .../src/dos/UserConnectionDO.ts | 238 ++++++++-- 2 files changed, 631 insertions(+), 49 deletions(-) diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index a0986bfdf34..586bf8bb45a 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -1406,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', @@ -3291,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', @@ -3361,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', @@ -3394,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', @@ -3423,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', @@ -6258,6 +6277,426 @@ describe('UserConnectionDO', () => { 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 // ------------------------------------------------------------------------- @@ -6658,7 +7097,10 @@ describe('UserConnectionDO', () => { // 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'); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 192fd922af7..b006fbd3436 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -254,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 @@ -565,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), + }); + }) + ); } } @@ -574,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 @@ -1567,6 +1585,7 @@ export class UserConnectionDO extends DurableObject { now: number ): void { const correlationId = crypto.randomUUID(); + const pendingCommandKey = `${PENDING_COMMAND_KEY_PREFIX}${correlationId}`; this.pendingCommands.set(correlationId, { ws, sessionId: msg.sessionId, @@ -1577,9 +1596,10 @@ export class UserConnectionDO extends DurableObject { expiresAt: now + UserConnectionDO.PENDING_COMMAND_TTL_MS, targetCliWs: targetCli, }); + this.pendingInitialCommandWrites.add(correlationId); this.ctx.waitUntil( this.ctx.storage - .put(`${PENDING_COMMAND_KEY_PREFIX}${correlationId}`, { + .put(pendingCommandKey, { sessionId: msg.sessionId, originalId: msg.id, command: msg.command, @@ -1589,7 +1609,14 @@ export class UserConnectionDO extends DurableObject { webConnectionId, state: 'pending' as const, } satisfies PendingCommandEntry) - .then(() => { + .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, { @@ -1606,6 +1633,12 @@ export class UserConnectionDO extends DurableObject { error: error instanceof Error ? error.message : String(error), }); }) + .finally(() => { + this.pendingInitialCommandWrites.delete(correlationId); + if (!this.terminalDuringInitialWrite.has(correlationId)) { + this.completedCorrelationIds.delete(correlationId); + } + }) ); } @@ -1630,8 +1663,10 @@ 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)', { @@ -1887,7 +1922,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; @@ -1950,13 +1992,17 @@ export class UserConnectionDO extends DurableObject { return undefined; } - private failPendingCommandsForSocket(targetWs: WebSocket, cliGone: boolean): void { + 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(); - for (const [id, entry] of this.pendingCommands) { - if (entry.targetCliWs !== targetWs) continue; + 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'; @@ -1969,6 +2015,67 @@ export class UserConnectionDO extends DurableObject { ? 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 @@ -1979,23 +2086,6 @@ export class UserConnectionDO extends DurableObject { error: liveError, } ); - this.pendingCommands.delete(id); - - // Step 26: pair the in-memory delete with a durable done transition. - const webAtt = entry.ws.deserializeAttachment() as WSAttachment | null; - this.ctx.waitUntil( - 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: webAtt?.role === 'web' ? webAtt.connectionId : 'unknown', - state: 'done', - ...(exited ? { result: {} } : { error: durableError }), - } satisfies PendingCommandEntry) - ); } const attachment = targetWs.deserializeAttachment() as WSAttachment | null; @@ -2010,20 +2100,85 @@ export class UserConnectionDO extends DurableObject { } } - private failPendingCommandsForOwnerChange( + private async failPendingCommandsForOwnerChange( sessionId: string, nextOwnerConnectionId: string | undefined - ): void { + ): Promise { const handledIds = new Set(); - for (const [id, entry] of this.pendingCommands) { - if (entry.sessionId !== sessionId || entry.targetConnectionId === nextOwnerConnectionId) { - continue; - } + 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 @@ -2034,22 +2189,6 @@ export class UserConnectionDO extends DurableObject { error: SESSION_OWNER_CHANGED_ERROR, } ); - - // Step 26: pair the in-memory delete with a durable done transition. - const webAtt = entry.ws.deserializeAttachment() as WSAttachment | null; - this.ctx.waitUntil( - 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: webAtt?.role === 'web' ? webAtt.connectionId : 'unknown', - state: 'done', - ...(exited ? { result: {} } : { error: SESSION_OWNER_CHANGED_ERROR }), - } satisfies PendingCommandEntry) - ); } this.ctx.waitUntil( @@ -2293,6 +2432,7 @@ export class UserConnectionDO extends DurableObject { // Already done and past TTL: clean up. await this.ctx.storage.delete(key); } + this.completedCorrelationIds.delete(key.slice(PENDING_COMMAND_KEY_PREFIX.length)); } })().catch((error: unknown) => { console.error('Failed to sweep durable pending commands', { From a8d7fb0c58273b94e653f001c9f2e193341eb332 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 19:48:06 +0200 Subject: [PATCH 21/24] fix(session-ingest): dedupe terminal sweeps Reserve terminal durable correlations across concurrent sweep paths. --- .../src/dos/UserConnectionDO.test.ts | 449 ++++++++++++++++++ .../src/dos/UserConnectionDO.ts | 107 ++++- 2 files changed, 531 insertions(+), 25 deletions(-) diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 586bf8bb45a..41a41e007eb 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -7173,4 +7173,453 @@ describe('UserConnectionDO', () => { 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).toBeGreaterThanOrEqual(0); + + // 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 b006fbd3436..9746d1ff713 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -1628,6 +1628,11 @@ export class UserConnectionDO extends DurableObject { }); }) .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), @@ -1636,7 +1641,14 @@ export class UserConnectionDO extends DurableObject { .finally(() => { this.pendingInitialCommandWrites.delete(correlationId); if (!this.terminalDuringInitialWrite.has(correlationId)) { - this.completedCorrelationIds.delete(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); + } } }) ); @@ -2329,30 +2341,42 @@ export class UserConnectionDO extends DurableObject { // Prevents a duplicate delivery when the in-memory storage.put and // finishDurablePendingCommands race within separate waitUntil tasks. if (skipIds?.has(correlationId)) return []; - // Persist the terminal outcome durably BEFORE any live delivery. - // If storage.put fails, no live send occurs — the error propagates - // to the waitUntil boundary and the caller handles the rejection. - const putPromise = this.ctx.storage.put(key, { - ...entry, - state: 'done', - error, - }); + // 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); - if (webWs) { - return [ - putPromise.then(() => { + 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]; }) ); @@ -2388,7 +2412,16 @@ export class UserConnectionDO extends DurableObject { 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, @@ -2409,17 +2442,42 @@ export class UserConnectionDO extends DurableObject { if (durable.expiresAt > now) continue; if (durable.state === 'pending') { - // Write terminal outcome durably (step 26a). - await this.ctx.storage.put(key, { - ...durable, - state: 'done', - error: COMMAND_EXPIRED_ERROR, - }); + 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). - // Skip entries already delivered in the in-memory sweep above. - const correlationId = key.slice(PENDING_COMMAND_KEY_PREFIX.length); - if (deliveredInMemory.has(correlationId)) continue; const webWs = this.findWebByConnectionId(durable.webConnectionId); if (webWs) { this.sendToWeb(webWs, { @@ -2432,7 +2490,6 @@ export class UserConnectionDO extends DurableObject { // Already done and past TTL: clean up. await this.ctx.storage.delete(key); } - this.completedCorrelationIds.delete(key.slice(PENDING_COMMAND_KEY_PREFIX.length)); } })().catch((error: unknown) => { console.error('Failed to sweep durable pending commands', { From fc4c6c89768ceca685da7c7a4f4c2357c7009858 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 20:03:44 +0200 Subject: [PATCH 22/24] chore: retrigger review From 3f12c37364a7fb4003226ace2c1dd08fededaf03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 4 Aug 2026 20:18:09 +0200 Subject: [PATCH 23/24] fix(session-ingest): release terminal sweep markers Drop expired durable correlation markers after their done entries expire. --- services/session-ingest/src/dos/UserConnectionDO.test.ts | 2 +- services/session-ingest/src/dos/UserConnectionDO.ts | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 41a41e007eb..a169a76b685 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -7465,7 +7465,7 @@ describe('UserConnectionDO', () => { terminalDuringInitialWrite: Map; } ).terminalDuringInitialWrite; - expect(terminalStash.size).toBeGreaterThanOrEqual(0); + expect(terminalStash.size).toBe(1); // Now reject the initial write. rejectInitialWrite!(new Error('simulated write failure')); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 9746d1ff713..d6bc20ac929 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -2487,8 +2487,12 @@ export class UserConnectionDO extends DurableObject { }); } } else { - // Already done and past TTL: clean up. + // 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) => { From 34885632f94d98e0178519269de0b6ea62444462 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Wed, 5 Aug 2026 16:49:24 +0200 Subject: [PATCH 24/24] test(dev): hold a tmux server open for the stale-environment test set-environment -g fails when no tmux server is running, which broke the test on CI runners without a leftover server (and on main). --- dev/local/tmux.test.ts | 8 ++++++++ 1 file changed, 8 insertions(+) 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. + } } } );