348 lines
18 KiB
TypeScript
348 lines
18 KiB
TypeScript
import assert from 'node:assert/strict'
|
|
import { testTerminalHostPlugin } from './test-terminal-host-plugin'
|
|
import { mkdir, mkdtemp, readFile, rm } from 'node:fs/promises'
|
|
import { tmpdir } from 'node:os'
|
|
import { join } from 'node:path'
|
|
import { pathToFileURL } from 'node:url'
|
|
import { build } from 'esbuild'
|
|
import type { ChatMessage } from '../src/shared/types'
|
|
|
|
interface PersistenceTestState {
|
|
userData: string
|
|
handlers: Map<string, (...args: unknown[]) => unknown>
|
|
promptCalls: number
|
|
abortCalls: number
|
|
cancelOutputEmitted: boolean
|
|
useOmp: boolean
|
|
fallbackHistory: ChatMessage[] | null
|
|
fallbackWireHistory?: Array<{ role: string; content: string }>
|
|
}
|
|
|
|
const TEST_STATE_KEY = '__YCODE_CHAT_PERSISTENCE_TEST__'
|
|
|
|
function wait(ms: number): Promise<void> {
|
|
return new Promise((resolve) => setTimeout(resolve, ms))
|
|
}
|
|
|
|
async function waitForMessages(
|
|
invoke: (channel: string, ...args: unknown[]) => Promise<unknown>,
|
|
sessionId: string,
|
|
done: (messages: ChatMessage[]) => boolean
|
|
): Promise<ChatMessage[]> {
|
|
for (let attempt = 0; attempt < 100; attempt++) {
|
|
const messages = (await invoke('session:messages', sessionId)) as ChatMessage[]
|
|
if (done(messages)) return messages
|
|
await wait(10)
|
|
}
|
|
const last = (await invoke('session:messages', sessionId)) as ChatMessage[]
|
|
assert.fail(`Timed out waiting for persisted messages: ${JSON.stringify(last)}`)
|
|
}
|
|
|
|
async function readDiskMessages(userData: string, sessionId: string): Promise<ChatMessage[]> {
|
|
const stored = JSON.parse(await readFile(join(userData, 'sessions.json'), 'utf8')) as {
|
|
sessions: Array<{ id: string; messages: ChatMessage[] }>
|
|
}
|
|
return stored.sessions.find((item) => item.id === sessionId)?.messages ?? []
|
|
}
|
|
|
|
async function waitForDiskAssistant(
|
|
userData: string,
|
|
sessionId: string,
|
|
matches: (message: ChatMessage) => boolean
|
|
): Promise<ChatMessage> {
|
|
for (let attempt = 0; attempt < 100; attempt++) {
|
|
const message = (await readDiskMessages(userData, sessionId)).findLast((entry) => entry.role === 'assistant')
|
|
if (message && matches(message)) return message
|
|
await wait(10)
|
|
}
|
|
const messages = await readDiskMessages(userData, sessionId)
|
|
assert.fail(`Timed out waiting for assistant checkpoint: ${JSON.stringify(messages.at(-1))}`)
|
|
}
|
|
|
|
function assistant(messages: ChatMessage[]): ChatMessage {
|
|
const result = messages.findLast((message) => message.role === 'assistant')
|
|
assert.ok(result, 'session should contain an assistant message')
|
|
return result
|
|
}
|
|
|
|
async function main(): Promise<void> {
|
|
const tempRoot = await mkdtemp(join(tmpdir(), 'ycode-chat-persistence-'))
|
|
const state: PersistenceTestState = {
|
|
userData: join(tempRoot, 'userData'),
|
|
handlers: new Map(),
|
|
promptCalls: 0,
|
|
abortCalls: 0,
|
|
cancelOutputEmitted: false,
|
|
useOmp: true,
|
|
fallbackHistory: null
|
|
}
|
|
await mkdir(state.userData, { recursive: true })
|
|
|
|
const globals = globalThis as typeof globalThis & Record<string, unknown>
|
|
const previousState = globals[TEST_STATE_KEY]
|
|
globals[TEST_STATE_KEY] = state
|
|
|
|
try {
|
|
const ipcEntry = join(process.cwd(), 'src', 'main', 'ipc.ts')
|
|
const sessionsEntry = join(process.cwd(), 'src', 'main', 'sessions.ts')
|
|
const outfile = join(tempRoot, 'ipc-chat-persistence.mjs')
|
|
await build({
|
|
stdin: {
|
|
contents: `export { registerIpc } from ${JSON.stringify(ipcEntry)};\nexport * as sessions from ${JSON.stringify(sessionsEntry)};`,
|
|
resolveDir: process.cwd(),
|
|
sourcefile: 'chat-persistence-regression-entry.ts',
|
|
loader: 'ts'
|
|
},
|
|
outfile,
|
|
bundle: true,
|
|
external: ['node-pty'],
|
|
platform: 'node',
|
|
format: 'esm',
|
|
banner: { js: "import { createRequire as createIpcTestRequire } from 'node:module'; const require = createIpcTestRequire(import.meta.url);" },
|
|
logLevel: 'silent',
|
|
plugins: [
|
|
testTerminalHostPlugin,
|
|
{
|
|
name: 'chat-persistence-mocks',
|
|
setup(pluginBuild) {
|
|
pluginBuild.onResolve({ filter: /^electron$/ }, () => ({ path: 'electron', namespace: 'mock-electron' }))
|
|
pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-electron' }, () => ({
|
|
loader: 'js',
|
|
contents: `
|
|
const state = globalThis.${TEST_STATE_KEY};
|
|
export const app = {
|
|
getPath: (name) => name === 'userData' ? state.userData : state.userData,
|
|
getName: () => 'YCode',
|
|
getVersion: () => '0.1.0',
|
|
isPackaged: false, isReady: () => true,
|
|
getAppPath: () => state.userData,
|
|
getLoginItemSettings: () => ({ openAtLogin: false }),
|
|
setLoginItemSettings: () => {}
|
|
};
|
|
export const safeStorage = {
|
|
isEncryptionAvailable: () => false,
|
|
encryptString: (value) => Buffer.from(value),
|
|
decryptString: (value) => value.toString()
|
|
};
|
|
export const ipcMain = {
|
|
handle: (name, handler) => state.handlers.set(name, handler)
|
|
};
|
|
export const dialog = { showOpenDialog: async () => ({ canceled: true, filePaths: [] }) };
|
|
export const shell = { openPath: async () => '' };
|
|
export const powerSaveBlocker = { start: () => 1, stop: () => true, isStarted: () => false };
|
|
export class BrowserWindow {}
|
|
export class WebContentsView {}
|
|
`
|
|
}))
|
|
pluginBuild.onResolve({ filter: /^\.\/omp$/ }, () => ({ path: 'omp', namespace: 'mock-omp' }))
|
|
pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-omp' }, () => ({
|
|
loader: 'js',
|
|
contents: `
|
|
import { EventEmitter } from 'node:events';
|
|
const state = globalThis.${TEST_STATE_KEY};
|
|
const toolId = 'persisted-tool';
|
|
const fullOutput = '${'long-output-'.repeat(60)}PERSISTED-OUTPUT-SENTINEL';
|
|
const emit = (kernel, frame) => kernel.emit('event', frame);
|
|
export class OmpKernel extends EventEmitter {
|
|
constructor() { super(); }
|
|
setThinkingLevel() { return Promise.resolve(); }
|
|
setHostTools() { return Promise.resolve({ type: 'response', success: true, data: { toolNames: [] } }); }
|
|
send() { return Promise.resolve({ type: 'response', success: true, data: { messages: [] } }); }
|
|
prompt(message) {
|
|
state.promptCalls++;
|
|
if (message.split('\\n', 1)[0] === 'persist-cancel') {
|
|
setImmediate(() => {
|
|
emit(this, { type: 'message_update', assistantMessageEvent: { type: 'text_delta', delta: 'cancel-before' } });
|
|
emit(this, { type: 'tool_execution_start', toolCall: { id: toolId, name: 'shell', args: { command: 'long-running' } } });
|
|
emit(this, { type: 'tool_execution_update', toolCall: { id: toolId, partialResult: 'cancelled tool output retained' } });
|
|
state.cancelOutputEmitted = true;
|
|
});
|
|
} else {
|
|
setImmediate(() => {
|
|
emit(this, { type: 'message_end', message: { role: 'assistant', usage: { input: 1200, output: 10, totalTokens: 1210 } } });
|
|
emit(this, { type: 'message_update', assistantMessageEvent: { type: 'text_delta', delta: 'before' } });
|
|
emit(this, { type: 'tool_execution_start', toolCall: { id: toolId, name: 'read_file', args: { path: 'large.txt' } } });
|
|
emit(this, { type: 'tool_execution_update', toolCall: { id: toolId, partialResult: fullOutput } });
|
|
emit(this, { type: 'message_update', assistantMessageEvent: { type: 'text_delta', delta: 'after' } });
|
|
emit(this, { type: 'tool_execution_end', toolCall: { id: toolId, name: 'read_file', result: { content: [{ type: 'text', text: fullOutput }], isError: false }, success: true } });
|
|
emit(this, { type: 'message_end', message: { role: 'assistant', usage: { contextTokens: 12800, input: 12000, cacheRead: 800, output: 200, totalTokens: 13000 } } });
|
|
emit(this, { type: 'agent_end', yielded: true });
|
|
});
|
|
}
|
|
return Promise.resolve({ type: 'response', success: true });
|
|
}
|
|
abort() {
|
|
state.abortCalls++;
|
|
setImmediate(() => emit(this, { type: 'agent_end', yielded: true }));
|
|
return Promise.resolve({ type: 'response', success: true });
|
|
}
|
|
dispose() {}
|
|
}
|
|
export const detectOmpPath = () => state.useOmp ? 'mock-omp' : null;
|
|
export const getOmpInfo = () => ({ available: true, exePath: 'mock-omp', profile: 'ycode', model: 'ycode/test' });
|
|
`
|
|
}))
|
|
pluginBuild.onResolve({ filter: /^\.\/agent\/loop$/ }, () => ({ path: 'agent-loop', namespace: 'mock-loop' }))
|
|
pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-loop' }, () => ({
|
|
loader: 'js',
|
|
contents: `
|
|
const state = globalThis.${TEST_STATE_KEY};
|
|
export async function runTurn(options) {
|
|
state.fallbackHistory = options.history;
|
|
state.fallbackWireHistory = options.wireHistory;
|
|
const content = 'fallback response';
|
|
options.emit({ runId: options.runId, type: 'text', delta: content });
|
|
options.emit({ runId: options.runId, type: 'done', content, tools: [] });
|
|
return { content, tools: [] };
|
|
}
|
|
`
|
|
}))
|
|
pluginBuild.onResolve({ filter: /^\.\/gateway$/ }, () => ({ path: 'gateway', namespace: 'mock-gateway' }))
|
|
pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-gateway' }, () => ({
|
|
loader: 'js',
|
|
contents: 'export class Gateway { constructor() {} }'
|
|
}))
|
|
}
|
|
}
|
|
]
|
|
})
|
|
|
|
const ipc = await import(pathToFileURL(outfile).href)
|
|
;(ipc.registerIpc as (getWindow: () => unknown) => void)(() => ({
|
|
webContents: { send: () => undefined }
|
|
}))
|
|
|
|
const invoke = (channel: string, ...args: unknown[]): Promise<unknown> => {
|
|
const handler = state.handlers.get(channel)
|
|
if (!handler) throw new Error(`Missing IPC handler: ${channel}`)
|
|
return Promise.resolve(handler({}, ...args))
|
|
}
|
|
|
|
await invoke('models:save', { id: 'usage-snapshot-model', name: 'Usage snapshot model', providerId: 'custom',
|
|
model: 'usage-test', baseUrl: 'http://127.0.0.1:1/v1', authMode: 'none', maxInputTokens: 64000 })
|
|
await invoke('models:select', 'usage-snapshot-model')
|
|
const normalSession = (await invoke('session:create', 'ordered persistence')) as { id: string }
|
|
await invoke('chat:send', { sessionId: normalSession.id, text: 'persist-order' })
|
|
const normalMessages = await waitForMessages(invoke, normalSession.id, (messages) => assistant(messages).status === 'done')
|
|
const renamedSession = ((await invoke('session:list')) as Array<{ id: string; title: string }>)
|
|
.find((entry) => entry.id === normalSession.id)
|
|
assert.equal(renamedSession?.title, 'persist-order', 'the first send should set the session title')
|
|
const normalAssistant = assistant(normalMessages)
|
|
const expectedParts = [
|
|
{ type: 'text', text: 'before' },
|
|
{ type: 'tool', toolId: 'persisted-tool' },
|
|
{ type: 'text', text: 'after' }
|
|
]
|
|
assert.deepEqual(normalAssistant.parts, expectedParts, 'session:messages should retain text/tool/text order')
|
|
assert.equal(normalAssistant.status, 'done')
|
|
assert.equal(normalAssistant.tools?.length, 1)
|
|
const normalTool = normalAssistant.tools?.[0]
|
|
assert.ok(normalTool, 'completed tool should be present')
|
|
assert.ok((normalTool.output?.length ?? 0) > 600, 'full tool output should exceed the summary limit')
|
|
assert.ok(normalTool.output?.endsWith('PERSISTED-OUTPUT-SENTINEL'), 'full tool output sentinel should be retained')
|
|
assert.ok((normalTool.summary?.length ?? Infinity) < (normalTool.output?.length ?? 0),
|
|
'the summary may be truncated while full output remains intact')
|
|
|
|
const diskNormalAssistant = assistant(await readDiskMessages(state.userData, normalSession.id))
|
|
assert.deepEqual(normalAssistant.contextUsage, {
|
|
usage: { inputTokens: 12800, outputTokens: 200, totalTokens: 13000, cachedTokens: 800 }, contextWindow: 64000,
|
|
modelId:'usage-snapshot-model',requestId:normalAssistant.requestId,source:'upstream',sampledAt:normalAssistant.contextUsage?.sampledAt
|
|
}, 'latest tool-round input is a sample, not the sum of all round usage')
|
|
assert.ok((normalAssistant.contextUsage?.sampledAt??0)>=normalAssistant.createdAt,'the persisted sample identifies its observation time')
|
|
assert.deepEqual(diskNormalAssistant.contextUsage, normalAssistant.contextUsage,
|
|
'the checkpoint whitelist must persist the upstream sample and frozen denominator together')
|
|
assert.deepEqual(diskNormalAssistant.parts, expectedParts, 'disk JSON should retain text/tool/text order')
|
|
assert.ok((diskNormalAssistant.tools?.[0].output?.length ?? 0) > 600,
|
|
'disk JSON should retain output beyond the summary limit')
|
|
assert.equal(diskNormalAssistant.tools?.[0].output, normalTool.output,
|
|
'disk JSON should retain the complete tool output')
|
|
assert.ok(diskNormalAssistant.tools?.[0].output?.endsWith('PERSISTED-OUTPUT-SENTINEL'))
|
|
|
|
const fallbackSession = (await invoke('session:create', 'fallback history')) as { id: string }
|
|
const sessions = ipc.sessions as {
|
|
appendMessage(sessionId: string, message: ChatMessage): Promise<void>
|
|
}
|
|
await sessions.appendMessage(fallbackSession.id, {
|
|
id: 'fallback-user',
|
|
role: 'user',
|
|
content: 'old request',
|
|
quotes: [{ id: 'saved-table', sourceId: 'table-source', kind: 'table', markdown: '| name | value |\n| --- | --- |\n| A | 42 |', preview: 'A 42' }],
|
|
createdAt: Date.now(),
|
|
status: 'done'
|
|
})
|
|
await sessions.appendMessage(fallbackSession.id, {
|
|
id: 'fallback-assistant',
|
|
role: 'assistant',
|
|
content: 'old answer',
|
|
parts: [{ type: 'text', text: 'old answer' }],
|
|
createdAt: Date.now(),
|
|
status: 'done'
|
|
})
|
|
state.useOmp = false
|
|
await invoke('chat:send', { sessionId: fallbackSession.id, text: 'fallback-history' })
|
|
const fallbackMessages = await waitForMessages(
|
|
invoke,
|
|
fallbackSession.id,
|
|
(messages) => assistant(messages).content === 'fallback response' && assistant(messages).status === 'done'
|
|
)
|
|
assert.equal(assistant(fallbackMessages).content, 'fallback response')
|
|
assert.deepEqual(state.fallbackHistory?.map((message) => message.content), ['old request', 'old answer'],
|
|
'the fallback loop should receive previously persisted user and assistant messages')
|
|
assert.equal(state.fallbackWireHistory?.[0].content, 'old request\n\n| name | value |\n| --- | --- |\n| A | 42 |',
|
|
'the next fallback request must include the exact previously quoted table snapshot')
|
|
state.useOmp = true
|
|
|
|
const cancelSession = (await invoke('session:create', 'cancelled persistence')) as { id: string }
|
|
const cancelledSend = (await invoke('chat:send', {
|
|
sessionId: cancelSession.id,
|
|
text: 'persist-cancel'
|
|
})) as { runId: string }
|
|
for (let attempt = 0; attempt < 100 && !state.cancelOutputEmitted; attempt++) await wait(10)
|
|
assert.equal(state.cancelOutputEmitted, true, 'fake OMP should emit a running tool and partial output')
|
|
const streamingCheckpoint = await waitForDiskAssistant(state.userData, cancelSession.id, (message) =>
|
|
message.status === 'streaming' &&
|
|
message.content === 'cancel-before' &&
|
|
message.tools?.[0].status === 'running' &&
|
|
message.tools?.[0].output === 'cancelled tool output retained'
|
|
)
|
|
assert.deepEqual(streamingCheckpoint.parts, [
|
|
{ type: 'text', text: 'cancel-before' },
|
|
{ type: 'tool', toolId: 'persisted-tool' }
|
|
], 'the streaming checkpoint should persist text/tool order before cancellation')
|
|
const cancelResult = (await invoke('chat:cancel', cancelledSend.runId)) as { cancelled: boolean }
|
|
assert.equal(cancelResult.cancelled, true)
|
|
|
|
const cancelledMessages = await waitForMessages(
|
|
invoke,
|
|
cancelSession.id,
|
|
(messages) => assistant(messages).status === 'cancelled'
|
|
)
|
|
const cancelledAssistant = assistant(cancelledMessages)
|
|
assert.deepEqual(cancelledAssistant.parts, [
|
|
{ type: 'text', text: 'cancel-before' },
|
|
{ type: 'tool', toolId: 'persisted-tool' }
|
|
])
|
|
assert.equal(cancelledAssistant.status, 'cancelled')
|
|
assert.equal(cancelledAssistant.tools?.[0].status, 'cancelled')
|
|
assert.equal(cancelledAssistant.tools?.[0].output, 'cancelled tool output retained')
|
|
|
|
const diskCancelledAssistant = assistant(await readDiskMessages(state.userData, cancelSession.id))
|
|
assert.equal(diskCancelledAssistant.status, 'cancelled', 'disk JSON should persist the assistant cancellation')
|
|
assert.equal(diskCancelledAssistant.tools?.[0].status, 'cancelled', 'disk JSON should persist terminal tool status')
|
|
assert.equal(diskCancelledAssistant.tools?.[0].output, 'cancelled tool output retained',
|
|
'disk JSON should preserve output from a running tool cancelled mid-turn')
|
|
assert.equal(state.promptCalls, 2)
|
|
assert.equal(state.abortCalls, 1)
|
|
|
|
console.log('chat-persistence-regression: all checks passed')
|
|
} finally {
|
|
globals[TEST_STATE_KEY] = previousState
|
|
await rm(tempRoot, { recursive: true, force: true })
|
|
}
|
|
}
|
|
|
|
void main().catch((error: unknown) => {
|
|
console.error('chat-persistence-regression failed:', error)
|
|
process.exitCode = 1
|
|
})
|