335 lines
17 KiB
TypeScript
335 lines
17 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
|
|
resumeAts: Array<string | undefined>
|
|
}
|
|
|
|
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,
|
|
resumeAts: []
|
|
}
|
|
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: [] }), showSaveDialog: async () => ({ canceled: true }) };
|
|
export const webContents = { getAllWebContents: () => [], fromId: () => undefined };
|
|
export const shell = { openPath: async () => '' };
|
|
export const powerSaveBlocker = { start: () => 1, stop: () => true, isStarted: () => false };
|
|
export class BrowserWindow {}
|
|
export class WebContentsView {}
|
|
`
|
|
}))
|
|
pluginBuild.onResolve({ filter: /^\.\/agent\/claude\/claude-kernel$/ }, () => ({ path: 'claude-kernel', namespace: 'mock-claude-kernel' }))
|
|
pluginBuild.onLoad({ filter: /.*/, namespace: 'mock-claude-kernel' }, () => ({
|
|
loader: 'js',
|
|
contents: `
|
|
const state = globalThis.${TEST_STATE_KEY};
|
|
const toolId = 'persisted-tool';
|
|
const fullOutput = '${'long-output-'.repeat(60)}PERSISTED-OUTPUT-SENTINEL';
|
|
const kernels = new Map();
|
|
class MockClaudeKernel {
|
|
constructor(options) { this.options = options; this.busy = false; this.active = null; }
|
|
async ensureStarted() {}
|
|
async setPolicy() {}
|
|
async runTurn(request) {
|
|
state.promptCalls++;
|
|
this.busy = true;
|
|
const cancelled = request.prompt.startsWith('persist-cancel');
|
|
if (cancelled) {
|
|
request.emit({ runId: request.runId, type: 'text', delta: 'cancel-before' });
|
|
request.emit({ runId: request.runId, type: 'tool_call', call: {
|
|
id: toolId, name: 'Read', args: { path: 'large.txt' }, argsRaw: '{"path":"large.txt"}',
|
|
status: 'running', output: 'cancelled tool output retained'
|
|
} });
|
|
state.cancelOutputEmitted = true;
|
|
return new Promise((resolve, reject) => {
|
|
const abort = () => {
|
|
state.abortCalls++;
|
|
this.busy = false;
|
|
reject(Object.assign(new Error('回合已取消'), { name: 'AbortError' }));
|
|
};
|
|
if (request.signal.aborted) abort();
|
|
else request.signal.addEventListener('abort', abort, { once: true });
|
|
});
|
|
}
|
|
const tool = {
|
|
id: toolId, name: 'Read', args: { path: 'large.txt' }, argsRaw: '{"path":"large.txt"}',
|
|
status: 'ok', output: fullOutput, summary: fullOutput.slice(0, 600), result: { content: fullOutput }
|
|
};
|
|
request.emit({ runId: request.runId, type: 'text', delta: 'before' });
|
|
request.emit({ runId: request.runId, type: 'tool_call', call: { ...tool, status: 'running', output: undefined, summary: undefined, result: undefined } });
|
|
request.emit({ runId: request.runId, type: 'tool_result', call: tool });
|
|
request.emit({ runId: request.runId, type: 'text', delta: 'after' });
|
|
request.emit({ runId: request.runId, type: 'usage', usage: { inputTokens: 12800, outputTokens: 200, totalTokens: 13000, cachedTokens: 800 } });
|
|
this.busy = false;
|
|
return { content: 'beforeafter', tools: [tool], status: 'done', entryId: 'mock-claude-entry' };
|
|
}
|
|
async rewindFiles() {}
|
|
async close() { this.busy = false; }
|
|
}
|
|
export function getClaudeKernel(options) {
|
|
state.resumeAts.push(options.resumeAt);
|
|
let kernel = kernels.get(options.sessionId);
|
|
if (!kernel) { kernel = new MockClaudeKernel(options); kernels.set(options.sessionId, kernel); }
|
|
return kernel;
|
|
}
|
|
export const cachedClaudeKernel = (sessionId) => kernels.get(sessionId);
|
|
export const releaseClaudeKernel = async (sessionId, force) => {
|
|
if (!force) return;
|
|
const kernel = kernels.get(sessionId);
|
|
if (!kernel) return;
|
|
kernels.delete(sessionId);
|
|
await kernel.close();
|
|
};
|
|
export const setSessionVisible = () => {};
|
|
export const removeSessionView = () => {};
|
|
export const disposeClaudeKernels = async () => { for (const kernel of kernels.values()) await kernel.close(); kernels.clear(); };
|
|
export const claudeRuntimeInfo = () => ({ available: false, exePath: null, sessionDir: state.userData });
|
|
`
|
|
}))
|
|
}
|
|
}
|
|
]
|
|
})
|
|
|
|
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.claudeEntryId, 'mock-claude-entry')
|
|
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 = await waitForDiskAssistant(state.userData, normalSession.id,
|
|
(message) => message.status === 'done' && message.claudeEntryId === 'mock-claude-entry')
|
|
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.equal(diskNormalAssistant.claudeEntryId, 'mock-claude-entry',
|
|
'disk JSON should retain the Claude fork point')
|
|
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'))
|
|
|
|
await invoke('chat:send', { sessionId: normalSession.id, text: 'persist-second' })
|
|
const secondMessages = await waitForMessages(invoke, normalSession.id,
|
|
(messages) => messages.filter((message) => message.role === 'assistant').length === 2 &&
|
|
assistant(messages).status === 'done')
|
|
const secondUser = secondMessages.filter((message) => message.role === 'user').at(-1)
|
|
assert.ok(secondUser)
|
|
await invoke('chat:send', {
|
|
sessionId: normalSession.id,
|
|
text: 'persist-edited',
|
|
fromMessageId: secondUser.id,
|
|
action: 'edit',
|
|
rollbackArtifacts: false
|
|
})
|
|
await waitForMessages(invoke, normalSession.id,
|
|
(messages) => messages.filter((message) => message.role === 'assistant').length === 2 &&
|
|
assistant(messages).status === 'done')
|
|
assert.equal(state.resumeAts.at(-1), normalAssistant.claudeEntryId,
|
|
'editing a prior turn should fork at the last kept Claude chain entry')
|
|
|
|
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, 'the mock Claude kernel 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, 4)
|
|
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
|
|
})
|