import assert from 'node:assert/strict'; import { mkdtemp, rm } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import path from 'node:path'; import test from 'node:test'; import { closeConnection, initializeDatabase, sessionsDb } from '@/modules/database/index.js'; import { chatRunRegistry } from '@/modules/websocket/services/chat-run-registry.service.js'; import type { NormalizedMessage } from '@/shared/types.js'; /** * Minimal stand-in for a websocket connection: collects every JSON frame the * gateway writer forwards so assertions can inspect the outbound protocol. */ class FakeConnection { readyState = 1; // WS_OPEN_STATE frames: NormalizedMessage[] = []; send(data: string): void { this.frames.push(JSON.parse(data) as NormalizedMessage); } } async function withIsolatedDatabase(runTest: () => void | Promise): Promise { const previousDatabasePath = process.env.DATABASE_PATH; const tempDirectory = await mkdtemp(path.join(tmpdir(), 'chat-run-registry-')); const databasePath = path.join(tempDirectory, 'auth.db'); closeConnection(); process.env.DATABASE_PATH = databasePath; await initializeDatabase(); try { await runTest(); } finally { chatRunRegistry.clearAll(); closeConnection(); if (previousDatabasePath === undefined) { delete process.env.DATABASE_PATH; } else { process.env.DATABASE_PATH = previousDatabasePath; } await rm(tempDirectory, { recursive: true, force: true }); } } test('live events are remapped to the app session id and sequenced', async () => { await withIsolatedDatabase(() => { sessionsDb.createAppSession('app-run-1', 'claude', '/workspace/demo'); const connection = new FakeConnection(); const run = chatRunRegistry.startRun({ appSessionId: 'app-run-1', provider: 'claude', providerSessionId: null, connection, userId: 'user-1', }); assert.ok(run); run.writer.send({ kind: 'stream_delta', provider: 'claude', sessionId: 'provider-id-9', content: 'hello' }); run.writer.send({ kind: 'text', provider: 'claude', sessionId: 'provider-id-9', content: 'hello world' }); assert.equal(connection.frames.length, 2); assert.equal(connection.frames[0]?.sessionId, 'app-run-1'); assert.equal(connection.frames[0]?.seq, 1); assert.equal(connection.frames[1]?.sessionId, 'app-run-1'); assert.equal(connection.frames[1]?.seq, 2); }); }); test('session_created is swallowed and persisted as the provider-id mapping', async () => { await withIsolatedDatabase(() => { sessionsDb.createAppSession('app-run-2', 'cursor', '/workspace/demo'); const connection = new FakeConnection(); const run = chatRunRegistry.startRun({ appSessionId: 'app-run-2', provider: 'cursor', providerSessionId: null, connection, userId: null, }); assert.ok(run); run.writer.send({ kind: 'session_created', provider: 'cursor', sessionId: 'cursor-native-7', newSessionId: 'cursor-native-7', }); // Never forwarded to the client... assert.equal(connection.frames.length, 0); // ...but recorded in the registry and persisted in the database. assert.equal(run.providerSessionId, 'cursor-native-7'); assert.equal(sessionsDb.getSessionById('app-run-2')?.provider_session_id, 'cursor-native-7'); }); }); test('complete marks the run finished and duplicate completes are dropped', async () => { await withIsolatedDatabase(() => { sessionsDb.createAppSession('app-run-3', 'codex', '/workspace/demo'); const connection = new FakeConnection(); const run = chatRunRegistry.startRun({ appSessionId: 'app-run-3', provider: 'codex', providerSessionId: null, connection, userId: null, }); assert.ok(run); run.writer.send({ kind: 'complete', provider: 'codex', sessionId: 'native-3', exitCode: 0 }); // Late duplicate from a killed runtime's exit handler. run.writer.send({ kind: 'complete', provider: 'codex', sessionId: 'native-3', exitCode: 1 }); const completes = connection.frames.filter((frame) => frame.kind === 'complete'); assert.equal(completes.length, 1); assert.equal(completes[0]?.actualSessionId, 'app-run-3'); assert.equal(chatRunRegistry.isProcessing('app-run-3'), false); // completeRun is also a no-op once the run already completed. chatRunRegistry.completeRun('app-run-3', { exitCode: 1 }); assert.equal(connection.frames.filter((frame) => frame.kind === 'complete').length, 1); }); }); test('listRunningRuns returns only currently running app sessions', async () => { await withIsolatedDatabase(() => { sessionsDb.createAppSession('app-run-7', 'claude', '/workspace/demo'); sessionsDb.createAppSession('app-run-8', 'codex', '/workspace/demo'); const connection = new FakeConnection(); const completedRun = chatRunRegistry.startRun({ appSessionId: 'app-run-7', provider: 'claude', providerSessionId: null, connection, userId: null, }); assert.ok(completedRun); const runningRun = chatRunRegistry.startRun({ appSessionId: 'app-run-8', provider: 'codex', providerSessionId: null, connection, userId: null, }); assert.ok(runningRun); chatRunRegistry.completeRun('app-run-7', { exitCode: 0 }); const runningSessions = chatRunRegistry.listRunningRuns(); assert.deepEqual(runningSessions.map((session) => session.sessionId), ['app-run-8']); assert.equal(runningSessions[0]?.provider, 'codex'); }); }); test('replayEvents returns only events after the requested seq', async () => { await withIsolatedDatabase(() => { sessionsDb.createAppSession('app-run-4', 'claude', '/workspace/demo'); const connection = new FakeConnection(); const run = chatRunRegistry.startRun({ appSessionId: 'app-run-4', provider: 'claude', providerSessionId: null, connection, userId: null, }); assert.ok(run); run.writer.send({ kind: 'stream_delta', provider: 'claude', sessionId: 'x', content: 'a' }); run.writer.send({ kind: 'stream_delta', provider: 'claude', sessionId: 'x', content: 'b' }); run.writer.send({ kind: 'stream_delta', provider: 'claude', sessionId: 'x', content: 'c' }); const replayed = chatRunRegistry.replayEvents('app-run-4', 1); assert.deepEqual(replayed.map((event) => event.content), ['b', 'c']); assert.deepEqual(replayed.map((event) => event.seq), [2, 3]); }); }); test('attachConnection reroutes the live stream to a new socket', async () => { await withIsolatedDatabase(() => { sessionsDb.createAppSession('app-run-5', 'gemini', '/workspace/demo'); const firstConnection = new FakeConnection(); const run = chatRunRegistry.startRun({ appSessionId: 'app-run-5', provider: 'gemini', providerSessionId: null, connection: firstConnection, userId: null, }); assert.ok(run); run.writer.send({ kind: 'stream_delta', provider: 'gemini', sessionId: 'g', content: 'before' }); const secondConnection = new FakeConnection(); assert.equal(chatRunRegistry.attachConnection('app-run-5', secondConnection), true); run.writer.send({ kind: 'stream_delta', provider: 'gemini', sessionId: 'g', content: 'after' }); assert.deepEqual(firstConnection.frames.map((frame) => frame.content), ['before']); assert.deepEqual(secondConnection.frames.map((frame) => frame.content), ['after']); }); }); test('startRun rejects a second concurrent run for the same session', async () => { await withIsolatedDatabase(() => { sessionsDb.createAppSession('app-run-6', 'opencode', '/workspace/demo'); const connection = new FakeConnection(); const first = chatRunRegistry.startRun({ appSessionId: 'app-run-6', provider: 'opencode', providerSessionId: null, connection, userId: null, }); assert.ok(first); const second = chatRunRegistry.startRun({ appSessionId: 'app-run-6', provider: 'opencode', providerSessionId: null, connection, userId: null, }); assert.equal(second, null); // After the run finishes a new one is allowed again. chatRunRegistry.completeRun('app-run-6', { exitCode: 0 }); const third = chatRunRegistry.startRun({ appSessionId: 'app-run-6', provider: 'opencode', providerSessionId: null, connection, userId: null, }); assert.ok(third); }); });