mirror of
https://github.com/siteboon/claudecodeui.git
synced 2026-06-13 01:22:06 +08:00
feat(chat): unify session gateway with stable IDs and a single WS protocol
The frontend previously juggled placeholder IDs, provider-native IDs, and session_created handoffs, which caused race conditions and provider-specific branching. This introduces app-allocated session IDs, a chat run registry with event replay, delta sidebar updates, and one kind-based websocket contract so the UI can treat every provider the same while JSONL remains the source of truth.
This commit is contained in:
Haileyesus
207
server/modules/websocket/tests/chat-run-registry.test.ts
Normal file
207
server/modules/websocket/tests/chat-run-registry.test.ts
Normal file
@@ -0,0 +1,207 @@
|
||||
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<void>): Promise<void> {
|
||||
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('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);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user