diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3e36b28b67..91e02e89f3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,7 +4,6 @@ on: push: branches: [main] pull_request: - branches: [main] # Cancel in-progress runs when a new commit is pushed to the same PR concurrency: @@ -381,6 +380,10 @@ jobs: - name: Run tests run: pnpm --filter ${{ matrix.workspace.name }} test + - name: Test native heartbeat attachment compatibility + if: matrix.workspace.name == 'cloudflare-session-ingest' + run: pnpm --filter cloudflare-session-ingest run test:integration test/integration/user-connection-attachment.test.ts + notify-main-failure: if: ${{ always() && github.ref == 'refs/heads/main' && contains(join(needs.*.result, ','), 'failure') }} needs: diff --git a/.github/workflows/extension-ci.yml b/.github/workflows/extension-ci.yml index fd39099b50..2df86dba1d 100644 --- a/.github/workflows/extension-ci.yml +++ b/.github/workflows/extension-ci.yml @@ -4,7 +4,6 @@ on: push: branches: [main] pull_request: - branches: [main] concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} diff --git a/.github/workflows/kilo-app-ci.yml b/.github/workflows/kilo-app-ci.yml index 7bbb23c1c2..c7e5210024 100644 --- a/.github/workflows/kilo-app-ci.yml +++ b/.github/workflows/kilo-app-ci.yml @@ -16,7 +16,6 @@ on: - 'pnpm-lock.yaml' - '.github/workflows/kilo-app-ci.yml' pull_request: - branches: [main] paths: - 'apps/mobile/**' - 'packages/trpc/**' diff --git a/scripts/changed-dependencies.test.mjs b/scripts/changed-dependencies.test.mjs index b5ef49bccd..0831ae8f06 100644 --- a/scripts/changed-dependencies.test.mjs +++ b/scripts/changed-dependencies.test.mjs @@ -1,5 +1,8 @@ import assert from 'node:assert/strict'; +import { readFileSync } from 'node:fs'; +import { posix } from 'node:path'; import test from 'node:test'; +import { load } from 'js-yaml'; import { changedDependencyWorkspaces } from './changed-dependencies.mjs'; @@ -256,3 +259,86 @@ test('a missing dependency cannot silently skip tests', () => { assert.throws(() => changedDependencyWorkspaces(snapshot(), after), /Missing lockfile entry/); }); + +function readWorkflow(file) { + return load(readFileSync(new URL(`../.github/workflows/${file}`, import.meta.url), 'utf8')); +} + +function matchesPatterns(value, patterns) { + return patterns.some(pattern => posix.matchesGlob(value, pattern)); +} + +function admitsEvent(event, branch, path) { + return ( + event !== undefined && + (!event?.branches || matchesPatterns(branch, event.branches)) && + (!event?.['branches-ignore'] || !matchesPatterns(branch, event['branches-ignore'])) && + (!event?.paths || matchesPatterns(path, event.paths)) && + (!event?.['paths-ignore'] || !matchesPatterns(path, event['paths-ignore'])) + ); +} + +for (const [file, paths, filterName] of [ + ['ci.yml', ['apps/web/src/routers/active-sessions-router.ts'], 'kilocode_backend'], + [ + 'kilo-app-ci.yml', + [ + 'apps/mobile/src/app/index.tsx', + 'packages/trpc/src/mobile.ts', + 'apps/web/src/routers/active-sessions-router.ts', + ], + ], + [ + 'extension-ci.yml', + [ + 'apps/extension/tests/e2e/agents-fixture.ts', + 'apps/web/src/routers/active-sessions-router.ts', + ], + 'extension', + ], +]) { + test(`${file} admits main and stack PRs, keeps main-only pushes and relevant paths`, () => { + const workflow = readWorkflow(file); + for (const branch of ['main', 'mobile-ux-ad6d-s1', 'mobile-ux-ad6d-s5', 'feature/other']) { + for (const path of paths) { + assert.equal( + admitsEvent(workflow.on.pull_request, branch, path), + true, + `PR ${branch}: ${path}` + ); + assert.equal( + admitsEvent(workflow.on.push, branch, path), + branch === 'main', + `push ${branch}: ${path}` + ); + } + } + if (filterName) { + const step = workflow.jobs.changes.steps.find(step => step.id === 'filter'); + const patterns = load(step.with.filters)[filterName]; + for (const path of paths) assert.equal(matchesPatterns(path, patterns), true, path); + assert.equal(matchesPatterns('docs/unrelated.md', patterns), false); + } else { + assert.equal( + admitsEvent(workflow.on.pull_request, 'mobile-ux-ad6d-s1', 'docs/unrelated.md'), + false + ); + assert.equal(admitsEvent(workflow.on.push, 'main', 'docs/unrelated.md'), false); + assert.ok(Object.hasOwn(workflow.on, 'workflow_call')); + } + }); +} + +test('the ingest workspace runs only the explicit native attachment regression', () => { + const workflow = readWorkflow('ci.yml'); + const path = 'services/session-ingest/test/integration/user-connection-attachment.test.ts'; + assert.equal(admitsEvent(workflow.on.pull_request, 'mobile-ux-ad6d-s1', path), true); + const step = workflow.jobs['workspace-tests'].steps.find(step => + step.run?.includes('test:integration') + ); + assert.equal(step?.if, "matrix.workspace.name == 'cloudflare-session-ingest'"); + assert.equal( + step?.run, + 'pnpm --filter cloudflare-session-ingest run test:integration test/integration/user-connection-attachment.test.ts' + ); +}); diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 19705a1690..f5b10de258 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -35,6 +35,7 @@ import { MAX_DURABLE_RESULT_BYTES, UserConnectionDO, } from './UserConnectionDO'; +import type { Instance } from '../types/user-connection-protocol'; // --------------------------------------------------------------------------- // Mock WebSocket @@ -258,7 +259,7 @@ function addCliSocket( title: string; platform?: string; }> = [], - instance?: { name: string; projectName: string; version?: string }, + instance?: Instance, kiloUserId?: string ): MockWS { const attachment: { @@ -301,8 +302,8 @@ function sendHeartbeat( }>, options: { protocolVersion?: string; - capabilities?: { attachments?: boolean }; - instance?: { name: string; projectName: string; version?: string }; + capabilities?: { attachments?: boolean; sessionClone?: boolean }; + instance?: Instance; } = {} ) { const msg = JSON.stringify({ @@ -4039,6 +4040,202 @@ describe('UserConnectionDO', () => { // getConnectedInstances RPC (W3) // ------------------------------------------------------------------------- + describe('heartbeat attachment compatibility', () => { + // Measured by the native Workers regression, not inferred from JSON size. + const capacityMessage = + "A WebSocket 'attachment' cannot be larger than 16384 bytes.'attachment' was 16472 bytes."; + const legacyInstance = { + name: 'current-host', + projectName: 'current-project', + version: '1.0.0', + }; + const heartbeat = { + type: 'heartbeat', + protocolVersion: '1', + capabilities: { attachments: true, sessionClone: true }, + instance: { + ...legacyInstance, + kind: 'remote' as const, + startedAt: '2026-08-28T12:34:56.789Z', + gitBranch: 'feature/identity', + }, + sessions: [ + { + id: 'current-session', + status: 'busy', + title: 'Current title', + gitUrl: 'https://github.com/org/project.git', + gitBranch: 'session-branch', + parentSessionId: 'parent-session', + platform: 'darwin', + prLink: { + platform: 'github', + prUrl: 'https://github.com/org/project/pull/1', + prNumber: 1, + }, + }, + ], + }; + + it('retries capacity with every current legacy field, then broadcasts and acknowledges', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket( + mockCtx, + 'cli-1', + [makeSession('previous-session', 'idle')], + { name: 'previous-host', projectName: 'previous-project', version: '0.0.0' }, + 'usr_1' + ); + const webWs = addWebSocket(mockCtx); + doInstance.getActiveSessions(); + const now = Date.now() + 1_000; + vi.spyOn(Date, 'now').mockReturnValue(now); + const write = vi.spyOn(cliWs, 'serializeAttachment').mockImplementationOnce(() => { + throw new Error(capacityMessage); + }); + + await doInstance.webSocketMessage(cliWs as never, JSON.stringify(heartbeat)); + + expect(write).toHaveBeenCalledTimes(2); + expect(cliWs.deserializeAttachment()).toEqual({ + role: 'cli', + connectionId: 'cli-1', + sessions: heartbeat.sessions, + heartbeatAt: now, + protocolVersion: '1', + capabilities: heartbeat.capabilities, + kiloUserId: 'usr_1', + instance: legacyInstance, + }); + expect(doInstance.hasActiveCliSession('current-session')).toBe(true); + expect(doInstance.hasActiveCliSession('previous-session')).toBe(false); + expect(doInstance.getConnectedInstances()).toEqual({ + instances: [ + { connectionId: 'cli-1', ...legacyInstance, capabilities: heartbeat.capabilities }, + ], + }); + expect(allSent(webWs)).toEqual([ + { + type: 'system', + event: 'sessions.heartbeat', + data: { + connectionId: 'cli-1', + protocolVersion: '1', + capabilities: heartbeat.capabilities, + sessions: [{ ...heartbeat.sessions[0], capabilities: heartbeat.capabilities }], + }, + }, + ]); + expect(allSent(cliWs)).toEqual([{ type: 'heartbeat_ack' }]); + }); + + it.each([{ kind: 'remote' }, { startedAt: '2026-08-28T12:34:56.789Z' }, { gitBranch: '' }])( + 'permits a capacity retry when only %j is present', + async metadata => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + doInstance.getActiveSessions(); + vi.spyOn(cliWs, 'serializeAttachment').mockImplementationOnce(() => { + throw new Error(capacityMessage); + }); + await doInstance.webSocketMessage( + cliWs as never, + JSON.stringify({ + ...heartbeat, + instance: { ...legacyInstance, ...metadata }, + }) + ); + expect(cliWs.deserializeAttachment()).toHaveProperty('instance', legacyInstance); + expect(cliWs.deserializeAttachment()).toHaveProperty('sessions', heartbeat.sessions); + expect(allSent(cliWs)).toEqual([{ type: 'heartbeat_ack' }]); + } + ); + + it.each([ + { + label: 'unrelated error', + error: new Error('unrelated persistence failure'), + instance: heartbeat.instance, + }, + { + label: 'wrong error class', + error: new TypeError(capacityMessage), + instance: heartbeat.instance, + }, + { + label: 'legacy-only capacity failure', + error: new Error(capacityMessage), + instance: legacyInstance, + }, + { + label: 'instance-free capacity failure', + error: new Error(capacityMessage), + instance: undefined, + }, + { + label: 'failed retry', + error: new Error(capacityMessage), + instance: heartbeat.instance, + retryError: new Error('retry failed'), + }, + { + label: 'capacity failure on retry', + error: new Error(capacityMessage), + instance: heartbeat.instance, + retryError: new Error(capacityMessage), + }, + ])( + 'rethrows $label unchanged without broadcasting or acknowledging', + async ({ error, instance, retryError }) => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1', [], heartbeat.instance); + const webWs = addWebSocket(mockCtx); + doInstance.getActiveSessions(); + const before = structuredClone(cliWs.deserializeAttachment()); + const write = vi + .spyOn(cliWs, 'serializeAttachment') + .mockImplementation(() => { + throw retryError ?? error; + }) + .mockImplementationOnce(() => { + throw error; + }); + + // sendHeartbeat discards its promise; await the production handler itself. + await expect( + doInstance.webSocketMessage(cliWs as never, JSON.stringify({ ...heartbeat, instance })) + ).rejects.toBe(retryError ?? error); + + expect(write).toHaveBeenCalledTimes(retryError ? 2 : 1); + expect(cliWs.deserializeAttachment()).toEqual(before); + expect(allSent(cliWs)).toEqual([]); + expect(allSent(webWs)).toEqual([]); + } + ); + + it.each([undefined, legacyInstance])( + 'writes a metadata-free heartbeat once: %j', + async instance => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1', [], heartbeat.instance); + doInstance.getActiveSessions(); + const write = vi.spyOn(cliWs, 'serializeAttachment'); + await doInstance.webSocketMessage( + cliWs as never, + JSON.stringify({ ...heartbeat, instance }) + ); + expect(write).toHaveBeenCalledTimes(1); + const attachment = cliWs.deserializeAttachment() as { + instance?: Instance; + sessions: unknown; + }; + expect(attachment.instance).toEqual(instance); + expect(attachment.sessions).toEqual(heartbeat.sessions); + expect(allSent(cliWs)).toEqual([{ type: 'heartbeat_ack' }]); + } + ); + }); + describe('getConnectedInstances', () => { it('returns one row per CLI socket that has an `instance` attachment', async () => { const { doInstance, mockCtx } = setup(); @@ -4129,28 +4326,31 @@ describe('UserConnectionDO', () => { ]); }); - it('includes capabilities when the CLI attachment advertises them', async () => { + it('projects metadata and both capabilities from a live hibernated attachment', async () => { const { doInstance, mockCtx } = setup(); - // Hibernated attachment carries capabilities — same source - // getConnectedInstances already uses for instance/version. + const instance: Instance = { + name: 'laptop-cap', + projectName: 'kilo', + version: '1.0.0', + kind: 'cli', + startedAt: '2026-08-28T12:34:56.789Z', + gitBranch: '', + }; + const capabilities = { attachments: true, sessionClone: true }; const cliWs = createMockWs(['cli'], { role: 'cli', connectionId: 'cli-cap', sessions: [], - instance: { name: 'laptop-cap', projectName: 'kilo' }, - capabilities: { attachments: true }, + instance, + capabilities, }); mockCtx.addSocket(cliWs); - const { instances } = doInstance.getConnectedInstances(); - expect(instances).toEqual([ - { - connectionId: 'cli-cap', - name: 'laptop-cap', - projectName: 'kilo', - capabilities: { attachments: true }, - }, - ]); + expect(doInstance.getConnectedInstances()).toEqual({ + instances: [{ connectionId: 'cli-cap', ...instance, capabilities }], + }); + cliWs.readyState = WebSocket.CLOSED; + expect(doInstance.getConnectedInstances()).toEqual({ instances: [] }); }); it('omits capabilities when the CLI attachment has none (legacy CLI)', async () => { @@ -4171,22 +4371,32 @@ describe('UserConnectionDO', () => { expect(instances[0]).not.toHaveProperty('capabilities'); }); - it('persists `instance` in the WS attachment across heartbeats', async () => { - const { doInstance, mockCtx } = setup(); + it('refreshes instance metadata and preserves it across hibernation', async () => { + const { doInstance, ctx, mockCtx } = setup(); const cliWs = addCliSocket(mockCtx, 'cli-1'); - sendHeartbeat(doInstance, cliWs, [], { - protocolVersion: '1', - instance: { name: 'laptop-1', projectName: 'kilo', version: '0.1.0' }, - }); - - const att = cliWs.deserializeAttachment() as { - instance?: { name: string }; - }; - expect(att.instance).toEqual({ + const instance: Instance = { name: 'laptop-1', projectName: 'kilo', version: '0.1.0', + kind: 'remote', + startedAt: '2026-08-28T12:34:56.789Z', + gitBranch: 'first-branch', + }; + const capabilities = { attachments: true, sessionClone: true }; + sendHeartbeat(doInstance, cliWs, [], { protocolVersion: '1', instance, capabilities }); + const refreshed = { ...instance, gitBranch: 'current-branch' }; + sendHeartbeat(doInstance, cliWs, [], { + protocolVersion: '1', + instance: refreshed, + capabilities, }); + + expect(cliWs.deserializeAttachment()).toMatchObject({ instance: refreshed, capabilities }); + for (const relay of [doInstance, new UserConnectionDO(ctx as never, {} as never)]) { + expect(relay.getConnectedInstances()).toEqual({ + instances: [{ connectionId: 'cli-1', ...refreshed, capabilities }], + }); + } }); it('drops `instance` from the attachment on a subsequent heartbeat that omits it', async () => { @@ -4534,27 +4744,34 @@ describe('UserConnectionDO', () => { // ------------------------------------------------------------------------- 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 - // change cannot silently push us over the budget. - const SERIALIZE_ATTACHMENT_BUDGET = 2048; - // Bounded `instance` = 64 + 64 + 32 chars content + JSON framing ≈ 200 - // bytes; we allow a 25% safety margin so a future protocol bump to the - // instance shape (e.g. adding `pid`) cannot silently blow the 2 KiB - // attachment budget. - const INSTANCE_HEADROOM = 250; - - it('keeps the combined CLI attachment comfortably under 2 KiB with a worst-case instance', async () => { - const worstCaseInstance = { + // These are JSON fixture guards, not proof of native attachment capacity. + // The Workers regression calibrates the actual production heartbeat write. + const LEGACY_JSON_BUDGET = 2048; + // 184 bounded UTF-16 units can each need six JSON bytes, plus 30 bytes for + // kind/timestamp and 81 bytes of framing: at most 1215, below 1280. + const INSTANCE_HEADROOM = 1280; + + it('keeps maximally escaped instance metadata within its JSON bound', () => { + const instance: Instance = { + name: '\u0000'.repeat(64), + projectName: '\u0000'.repeat(64), + version: '\u0000'.repeat(32), + kind: 'remote', + startedAt: '2026-08-28T12:34:56.789Z', + gitBranch: '\u0000'.repeat(24), + }; + expect(new TextEncoder().encode(JSON.stringify(instance)).byteLength).toBeLessThan( + INSTANCE_HEADROOM + ); + }); + + it('keeps the full representative legacy attachment below 2 KiB of JSON', async () => { + const legacyInstance = { 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.) + // Preserve four representative sessions, not a production capacity limit. const sessions = Array.from({ length: 4 }, (_, i) => ({ id: `ses_${String(i).padStart(26, '0')}`, status: 'busy', @@ -4567,17 +4784,15 @@ describe('UserConnectionDO', () => { role: 'cli' as const, connectionId: 'cli-1', sessions, + heartbeatAt: 1_788_000_000_000, protocolVersion: '255.255.65535', + capabilities: { attachments: true, sessionClone: true }, kiloUserId: 'usr_' + 'x'.repeat(28), - instance: worstCaseInstance, + instance: legacyInstance, }; 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; - expect(instanceBytes).toBeLessThan(INSTANCE_HEADROOM); + expect(serialized).toBeLessThan(LEGACY_JSON_BUDGET); }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 09b4dde489..935065099d 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -82,11 +82,10 @@ type WSAttachment = // Type re-export so test files and other internal callers can reference the // connection-row shape from a single place. -export type ConnectedInstanceRow = { +// Instance metadata stays optional for old producers and hibernated attachments. +// Remove that compatibility only after every supported old form has retired. +export type ConnectedInstanceRow = Instance & { connectionId: string; - name: string; - projectName: string; - version?: string; // Latest capabilities from the CLI socket attachment. Omitted when the // attachment has no capabilities (legacy CLI / pre-field build) so the // response stays byte-identical for those clients. @@ -675,7 +674,32 @@ export class UserConnectionDO extends DurableObject { kiloUserId: attachment.kiloUserId, ...(instance ? { instance } : {}), }; - ws.serializeAttachment(updatedAttachment); + try { + ws.serializeAttachment(updatedAttachment); + } catch (error) { + if ( + !(error instanceof Error) || + error.name !== 'Error' || + !/^A WebSocket 'attachment' cannot be larger than 16384 bytes\.'attachment' was \d+ bytes\.$/.test( + error.message + ) || + !instance || + (instance.kind === undefined && + instance.startedAt === undefined && + instance.gitBranch === undefined) + ) { + throw error; + } + // The native regression verifies this capacity error and failed-write atomicity. + // Retry the current heartbeat in the old metadata-free form, never a stale one. + // Remove only after old producers/attachments retire and enriched heartbeats + // have proven native capacity safety. + const legacyInstance = { ...instance }; + delete legacyInstance.kind; + delete legacyInstance.startedAt; + delete legacyInstance.gitBranch; + ws.serializeAttachment({ ...updatedAttachment, instance: legacyInstance }); + } // Broadcast the heartbeat to every one of the user's web sockets. Subscribers // and non-subscribers both receive it: a removed session id is detectable @@ -1925,8 +1949,8 @@ export class UserConnectionDO extends DurableObject { * * No in-memory map is consulted: hibernation/restart can never produce a * stale row because we only read from sockets that are alive right now. - * The 2KB `serializeAttachment` budget comfortably accommodates a bounded - * instance object (well under 200 bytes). + * Old attachments can omit metadata; the heartbeat write handles native + * capacity without discarding their legacy instance identity. */ getConnectedInstances(): { instances: ConnectedInstanceRow[] } { this.ensureState(); @@ -1944,6 +1968,9 @@ export class UserConnectionDO extends DurableObject { name: att.instance.name, projectName: att.instance.projectName, ...(att.instance.version ? { version: att.instance.version } : {}), + ...(att.instance.kind !== undefined ? { kind: att.instance.kind } : {}), + ...(att.instance.startedAt !== undefined ? { startedAt: att.instance.startedAt } : {}), + ...(att.instance.gitBranch !== undefined ? { gitBranch: att.instance.gitBranch } : {}), ...(att.capabilities ? { capabilities: att.capabilities } : {}), }); } diff --git a/services/session-ingest/src/routes/api.test.ts b/services/session-ingest/src/routes/api.test.ts index 803602c49c..6df688126f 100644 --- a/services/session-ingest/src/routes/api.test.ts +++ b/services/session-ingest/src/routes/api.test.ts @@ -2262,6 +2262,16 @@ describe('api routes', () => { instances: [ { connectionId: 'cli-A', name: 'laptop-A', projectName: 'kilo', version: '0.1.2' }, { connectionId: 'cli-B', name: 'laptop-B', projectName: 'kilo' }, + { + connectionId: 'cli-remote', + name: 'remote-host', + projectName: 'kilo', + version: '1.0.0', + kind: 'remote', + startedAt: '2026-08-28T12:34:56.789Z', + gitBranch: 'feature/identity', + capabilities: { attachments: true, sessionClone: true }, + }, ], })); vi.mocked(getUserConnectionDO).mockReturnValue({ @@ -2280,10 +2290,34 @@ describe('api routes', () => { instances: [ { connectionId: 'cli-A', name: 'laptop-A', projectName: 'kilo', version: '0.1.2' }, { connectionId: 'cli-B', name: 'laptop-B', projectName: 'kilo' }, + { + connectionId: 'cli-remote', + name: 'remote-host', + projectName: 'kilo', + version: '1.0.0', + kind: 'remote', + startedAt: '2026-08-28T12:34:56.789Z', + gitBranch: 'feature/identity', + capabilities: { attachments: true, sessionClone: true }, + }, ], }); }); + it('keeps a failed instance lookup distinct from an empty list', async () => { + vi.mocked(getUserConnectionDO).mockReturnValue({ + getConnectedInstances: vi.fn(async () => { + throw new Error('instance lookup failed'); + }), + } as never); + const res = await makeApiApp().fetch( + new Request('http://local/instances/active', { method: 'GET' }), + makeTestEnv() + ); + expect(res.status).toBe(500); + expect(await res.text()).toBe('Internal Server Error'); + }); + it('returns 200 with an empty `instances` array when no CLIs are connected', async () => { vi.mocked(getUserConnectionDO).mockReturnValue({ getConnectedInstances: vi.fn(async () => ({ instances: [] })), diff --git a/services/session-ingest/src/types/user-connection-protocol.test.ts b/services/session-ingest/src/types/user-connection-protocol.test.ts index 0b9f9625b9..c08c994b5d 100644 --- a/services/session-ingest/src/types/user-connection-protocol.test.ts +++ b/services/session-ingest/src/types/user-connection-protocol.test.ts @@ -63,6 +63,66 @@ describe('CLIOutboundMessageSchema', () => { } }); + it.each(['cli', 'remote'])('preserves full %s metadata and both capabilities', kind => { + const msg = { + type: 'heartbeat', + protocolVersion: '1', + capabilities: { attachments: true, sessionClone: true }, + instance: { + name: 'laptop-1', + projectName: 'kilo', + version: '0.1.2', + kind, + startedAt: '2026-08-28T12:34:56.789Z', + gitBranch: 'feature/identity', + }, + sessions: [{ id: 'ses_1', status: 'busy', title: 'Remote session', platform: 'darwin' }], + }; + expect(CLIOutboundMessageSchema.parse(msg)).toEqual(msg); + }); + + it.each([ + ['kind', 'terminal'], + ['kind', null], + ['startedAt', '2026-08-28T12:34:56Z'], + ['startedAt', '2026-08-28T12:34:56.78Z'], + ['startedAt', '2026-08-28T12:34:56.7890Z'], + ['startedAt', '2026-08-28T12:34:56.789+00:00'], + ['startedAt', '2026-02-30T12:34:56.789Z'], + ['startedAt', null], + ['gitBranch', null], + ])('rejects invalid instance %s: %s', (field, value) => { + expect( + CLIOutboundMessageSchema.safeParse({ + type: 'heartbeat', + sessions: [], + instance: { name: 'host', projectName: 'project', [field]: value }, + }).success + ).toBe(false); + }); + + it.each([ + ['ASCII', 'a'.repeat(24), 'a'.repeat(25)], + ['escaped characters', '\\"'.repeat(12), '\\"'.repeat(12) + '\\'], + ['CJK', '界'.repeat(24), '界'.repeat(25)], + ['surrogate pairs', '\u{10400}'.repeat(12), '\u{10400}'.repeat(12) + 'a'], + ])('bounds %s branches by UTF-16 units, not JSON bytes', (_label, valid, invalid) => { + const heartbeat = { + type: 'heartbeat', + sessions: [], + instance: { name: 'host', projectName: 'project', gitBranch: valid }, + }; + expect(CLIOutboundMessageSchema.parse(JSON.parse(JSON.stringify(heartbeat)))).toEqual( + heartbeat + ); + expect( + CLIOutboundMessageSchema.safeParse({ + ...heartbeat, + instance: { ...heartbeat.instance, gitBranch: invalid }, + }).success + ).toBe(false); + }); + it('rejects instance with empty name', () => { const msg = { type: 'heartbeat', diff --git a/services/session-ingest/src/types/user-connection-protocol.ts b/services/session-ingest/src/types/user-connection-protocol.ts index be6c711d8f..8df21625d5 100644 --- a/services/session-ingest/src/types/user-connection-protocol.ts +++ b/services/session-ingest/src/types/user-connection-protocol.ts @@ -6,7 +6,7 @@ import { z } from 'zod'; // -- CLI → DO (CLIOutbound) --------------------------------------------------- -// Identity of the CLI process (kilo remote spawner) attached to this WebSocket. +// Identity of the CLI process (terminal or kilo remote) attached to this WebSocket. // Newer CLIs include this on every heartbeat; legacy CLIs that predate the // `kilo remote` spawner omit it entirely. The DO persists the latest value // in the WebSocket attachment and uses it for `getConnectedInstances()`. @@ -14,6 +14,11 @@ const instanceSchema = z.object({ name: z.string().min(1).max(64), projectName: z.string().min(1).max(64), version: z.string().max(32).optional(), + // Old producers and hibernated attachments omit this metadata. Keep it optional + // until all supported metadata-free producers and attachments have retired. + kind: z.enum(['cli', 'remote']).optional(), + startedAt: z.string().datetime({ precision: 3 }).length(24).optional(), + gitBranch: z.string().max(24).optional(), }); export type Instance = z.infer; diff --git a/services/session-ingest/test/integration/user-connection-attachment.test.ts b/services/session-ingest/test/integration/user-connection-attachment.test.ts new file mode 100644 index 0000000000..28e983b784 --- /dev/null +++ b/services/session-ingest/test/integration/user-connection-attachment.test.ts @@ -0,0 +1,181 @@ +import { env, runInDurableObject } from 'cloudflare:test'; +import { expect, it, vi } from 'vitest'; +import type { UserConnectionDO } from '../../src/dos/UserConnectionDO'; + +it('preserves the current heartbeat through native attachment capacity fallback', async () => { + const stub = env.USER_CONNECTION_DO.get(env.USER_CONNECTION_DO.newUniqueId()); + const evidence = await runInDurableObject(stub, async (instance: UserConnectionDO, state) => { + const [client, server] = Object.values(new WebSocketPair()); + state.acceptWebSocket(server, ['cli']); + client.accept(); + const now = Date.now(); + const clock = vi.spyOn(Date, 'now').mockReturnValue(now); + const heartbeat = (titleLength: number) => ({ + type: 'heartbeat', + sessions: [ + { + id: 'ses_native_capacity', + status: 'busy', + title: 'T'.repeat(titleLength), + gitUrl: 'https://github.com/org/project.git', + gitBranch: 'session-branch', + parentSessionId: 'ses_parent', + platform: 'darwin', + prLink: { + platform: 'github', + prUrl: 'https://github.com/org/project/pull/1', + prNumber: 1, + }, + }, + ], + protocolVersion: '1', + capabilities: { attachments: true, sessionClone: true }, + instance: { name: 'host', projectName: 'project', version: '1.0.0' }, + }); + server.serializeAttachment({ + role: 'cli', + connectionId: 'native-capacity', + sessions: [], + heartbeatAt: now, + kiloUserId: 'usr_native_capacity', + }); + const nativeSerialize = server.serializeAttachment.bind(server); + const failures: unknown[] = []; + // Observe the real serializer without replacing its capacity behavior. + const write = vi.spyOn(server, 'serializeAttachment').mockImplementation(value => { + const before: unknown = server.deserializeAttachment(); + try { + nativeSerialize(value); + } catch (error) { + failures.push(error); + expect(server.deserializeAttachment()).toEqual(before); + throw error; + } + }); + try { + // Calibrate through the production parser and heartbeat write, not a + // hand-built fixture or JSON length. + let fitting = 0; + let rejected = 64 * 1024; + await instance.webSocketMessage(server, JSON.stringify(heartbeat(fitting))); + await expect( + instance.webSocketMessage(server, JSON.stringify(heartbeat(rejected))) + ).rejects.toThrow("A WebSocket 'attachment' cannot be larger than 16384 bytes."); + while (fitting + 1 < rejected) { + const length = Math.floor((fitting + rejected) / 2); + try { + await instance.webSocketMessage(server, JSON.stringify(heartbeat(length))); + fitting = length; + } catch (error) { + expect(error).toMatchObject({ + name: 'Error', + message: expect.stringContaining( + "A WebSocket 'attachment' cannot be larger than 16384 bytes." + ), + }); + rejected = length; + } + } + const current = heartbeat(fitting); + await instance.webSocketMessage(server, JSON.stringify(current)); + const legacy = { + role: 'cli', + connectionId: 'native-capacity', + sessions: current.sessions, + heartbeatAt: now, + protocolVersion: '1', + capabilities: current.capabilities, + kiloUserId: 'usr_native_capacity', + instance: current.instance, + }; + expect(server.deserializeAttachment()).toEqual(legacy); + + // Install a previous heartbeat to detect retries that reuse stale fields. + clock.mockReturnValue(now - 5_000); + await instance.webSocketMessage( + server, + JSON.stringify({ + ...current, + sessions: [ + { + id: 'ses_previous', + status: 'idle', + title: 'Previous', + parentSessionId: 'ses_parent', + }, + ], + protocolVersion: '0', + capabilities: { attachments: false, sessionClone: false }, + instance: { name: 'previous-host', projectName: 'previous-project', version: '0.0.0' }, + }) + ); + clock.mockReturnValue(now); + const enriched = { + ...current, + instance: { + ...current.instance, + kind: 'remote', + startedAt: '2026-08-28T12:34:56.789Z', + gitBranch: 'b'.repeat(24), + }, + }; + failures.length = 0; + write.mockClear(); + await instance.webSocketMessage(server, JSON.stringify(enriched)); + expect(failures).toHaveLength(1); + const [failure] = failures; + if (!(failure instanceof Error)) throw new Error('Expected a native capacity exception'); + const capacityEvidence = { + name: failure.name, + constructor: failure.constructor.name, + message: failure.message, + }; + expect(write).toHaveBeenCalledTimes(2); + expect(server.deserializeAttachment()).toEqual(legacy); + expect(instance.hasActiveCliSession('ses_native_capacity')).toBe(true); + expect(instance.hasActiveCliSession('ses_previous')).toBe(false); + expect(instance.getConnectedInstances()).toEqual({ + instances: [ + { + connectionId: legacy.connectionId, + ...current.instance, + capabilities: current.capabilities, + }, + ], + }); + + // A later fitting heartbeat must advertise metadata again, not keep the fallback. + await instance.webSocketMessage(server, JSON.stringify({ ...enriched, sessions: [] })); + expect(server.deserializeAttachment()).toEqual({ + ...legacy, + instance: enriched.instance, + sessions: [], + }); + expect(instance.getConnectedInstances()).toEqual({ + instances: [ + { + connectionId: legacy.connectionId, + ...enriched.instance, + capabilities: current.capabilities, + }, + ], + }); + return capacityEvidence; + } finally { + write.mockRestore(); + clock.mockRestore(); + // Empty the owned connection before close; this test needs no Postgres disconnect work. + await instance.webSocketMessage(server, JSON.stringify({ type: 'heartbeat', sessions: [] })); + await state.storage.deleteAlarm(); + client.close(); + server.close(); + } + }); + // Pin the measured native exception separately from the production classifier. + expect(evidence).toEqual({ + name: 'Error', + constructor: 'Error', + message: + "A WebSocket 'attachment' cannot be larger than 16384 bytes.'attachment' was 16472 bytes.", + }); +}); diff --git a/services/session-ingest/test/test-worker.ts b/services/session-ingest/test/test-worker.ts index 6052f5c1ef..e0f8829780 100644 --- a/services/session-ingest/test/test-worker.ts +++ b/services/session-ingest/test/test-worker.ts @@ -1,5 +1,7 @@ export { SessionIngestDO } from '../src/dos/SessionIngestDO'; export { SessionAccessCacheDO } from '../src/dos/SessionAccessCacheDO'; +export { UserConnectionDO } from '../src/dos/UserConnectionDO'; +export { ConnectionTicketDO } from '../src/dos/connection-ticket-do'; export default { fetch(): Response { diff --git a/services/session-ingest/wrangler.test.jsonc b/services/session-ingest/wrangler.test.jsonc index e076a4cca4..bb5172a556 100644 --- a/services/session-ingest/wrangler.test.jsonc +++ b/services/session-ingest/wrangler.test.jsonc @@ -24,6 +24,10 @@ "name": "SESSION_ACCESS_CACHE_DO", "class_name": "SessionAccessCacheDO", }, + { + "name": "USER_CONNECTION_DO", + "class_name": "UserConnectionDO", + }, { "name": "CONNECTION_TICKET_DO", "class_name": "ConnectionTicketDO", @@ -40,6 +44,10 @@ "tag": "v2", "new_sqlite_classes": ["ConnectionTicketDO"], }, + { + "tag": "v3", + "new_sqlite_classes": ["UserConnectionDO"], + }, ], "r2_buckets": [