diff --git a/src/renderer/stores/chatStore.ts b/src/renderer/stores/chatStore.ts index 2695d5c..45f6bbb 100644 --- a/src/renderer/stores/chatStore.ts +++ b/src/renderer/stores/chatStore.ts @@ -165,6 +165,15 @@ const abortControllers = new Map(); /** Per-conversation ask-user resolvers */ const askUserResolvers = new Map void>(); +function cancelConversationRun(convId: string): void { + const resolver = askUserResolvers.get(convId); + if (resolver) { + resolver('__ABORT__'); + askUserResolvers.delete(convId); + } + abortControllers.get(convId)?.abort(); +} + async function runAgentWithModel( convId: string, assistantMsgId: string, @@ -505,16 +514,26 @@ export const useChatStore = create((set, get) => ({ }, deleteConversation: (id) => { + cancelConversationRun(id); set((s) => { const conversations = s.conversations.filter((c) => c.id !== id); const activeConversationId = s.activeConversationId === id ? (conversations[0]?.id ?? null) : s.activeConversationId; - if (activeConversationId) { - dbAPI.setConfig('activeConversationId', activeConversationId).catch(console.error); - } - return { conversations, activeConversationId }; + const streamingConversationIds = s.streamingConversationIds.filter((convId) => convId !== id); + const isStreaming = activeConversationId + ? streamingConversationIds.includes(activeConversationId) + : false; + dbAPI.setConfig('activeConversationId', activeConversationId ?? '').catch(console.error); + return { + conversations, + activeConversationId, + isStreaming, + streamingConversationIds, + pendingQuestion: s.pendingQuestion?.convId === id ? null : s.pendingQuestion, + agentTurnInfo: s.agentTurnInfo?.convId === id ? null : s.agentTurnInfo, + }; }); dbAPI.deleteConversation(id).catch(console.error); }, @@ -594,7 +613,8 @@ export const useChatStore = create((set, get) => ({ } const updatedConv = get().conversations.find((c) => c.id === convId); - const messagesForAgent = updatedConv ? updatedConv.messages.slice(0, -1) : [userMsg]; + if (!updatedConv) return; + const messagesForAgent = updatedConv.messages.slice(0, -1); await runAgentWithModel(convId!, assistantMsg.id, messagesForAgent, set, get); }, @@ -690,17 +710,12 @@ export const useChatStore = create((set, get) => ({ stopStreaming: () => { const { activeConversationId, pendingQuestion } = get(); - // If there's a pending AG2U question, resolve it with __ABORT__ so the agent loop unblocks - if (activeConversationId && pendingQuestion?.convId === activeConversationId) { - const resolver = askUserResolvers.get(activeConversationId); - if (resolver) { - resolver('__ABORT__'); - askUserResolvers.delete(activeConversationId); - } - set({ pendingQuestion: null }); - } if (activeConversationId) { - abortControllers.get(activeConversationId)?.abort(); + cancelConversationRun(activeConversationId); + } + // Clear the active conversation's pending AG2UI question after unblocking its agent. + if (pendingQuestion?.convId === activeConversationId) { + set({ pendingQuestion: null }); } // isStreaming/streamingConversationIds will be cleaned up in runAgentWithModel's finally block }, diff --git a/tests/chat-store-delete-cancellation.test.cjs b/tests/chat-store-delete-cancellation.test.cjs new file mode 100644 index 0000000..4c093fa --- /dev/null +++ b/tests/chat-store-delete-cancellation.test.cjs @@ -0,0 +1,288 @@ +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const Module = require('node:module'); +const path = require('node:path'); +const test = require('node:test'); +const ts = require('typescript'); + +const chatStorePath = path.join(__dirname, '..', 'src', 'renderer', 'stores', 'chatStore.ts'); +const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); + +async function waitFor(predicate, message) { + const deadline = Date.now() + 1000; + while (Date.now() < deadline) { + if (predicate()) return; + await sleep(5); + } + assert.fail(message); +} + +async function settlesWithin(promise, timeoutMs) { + let timer; + const timeout = new Promise((resolve) => { + timer = setTimeout(() => resolve(false), timeoutMs); + }); + try { + return await Promise.race([promise.then(() => true), timeout]); + } finally { + clearTimeout(timer); + } +} + +async function withChatStore({ ids, runReactAgent, upsertMessage = async () => {} }, run) { + const originalTsLoader = require.extensions['.ts']; + const originalLoad = Module._load; + const originalLocalStorage = global.localStorage; + const deletedConversationIds = []; + const configWrites = []; + let memoryCalls = 0; + + global.localStorage = { + getItem: () => null, + setItem: () => {}, + removeItem: () => {}, + }; + + require.extensions['.ts'] = (mod, filename) => { + const source = fs.readFileSync(filename, 'utf8'); + const output = ts.transpileModule(source, { + compilerOptions: { + module: ts.ModuleKind.CommonJS, + target: ts.ScriptTarget.ES2022, + esModuleInterop: true, + }, + }).outputText; + mod._compile(output, filename); + }; + + const model = { + id: 'local-model', + name: 'Local model', + provider: 'ollama', + modelId: 'test-model', + apiKey: '', + baseUrl: 'http://localhost', + }; + const configState = { + models: [model], + activeModelId: model.id, + language: 'zh-CN', + memoryShortTermRounds: 5, + setActiveModel: () => {}, + }; + + const dbAPI = { + getConversations: async () => [], + getMessages: async () => [], + getConfig: async () => null, + setConfig: async (key, value) => { + configWrites.push([key, value]); + }, + upsertConversation: async () => {}, + upsertMessage, + updateConversationTitle: async () => {}, + deleteConversation: async (id) => { + deletedConversationIds.push(id); + }, + deleteMessage: async () => {}, + }; + + const mocks = { + uuid: { v4: () => ids.shift() }, + './configStore': { useConfigStore: { getState: () => configState } }, + './subagentStore': { useSubagentStore: { getState: () => ({ clearTodos: () => {} }) } }, + '../services/reactAgent': { runReactAgent }, + '../services/memoryService': { + consolidateSessionMemory: async () => { + memoryCalls += 1; + }, + }, + '../services/dbAPI': { dbAPI, isElectron: () => true }, + './suggestionsStore': { + useSuggestionsStore: { getState: () => ({ setSuggestions: () => {} }) }, + parseSuggestionsFromMessage: () => [], + }, + '../i18n': { getLocale: () => ({ sidebar: { newConversationTitle: 'New conversation' } }) }, + }; + + Module._load = function loadWithMocks(request, parent, isMain) { + if (parent?.filename === chatStorePath && Object.hasOwn(mocks, request)) { + return mocks[request]; + } + return originalLoad.call(this, request, parent, isMain); + }; + + try { + delete require.cache[chatStorePath]; + const { useChatStore } = require(chatStorePath); + await run({ + useChatStore, + configWrites, + deletedConversationIds, + getMemoryCalls: () => memoryCalls, + }); + } finally { + delete require.cache[chatStorePath]; + Module._load = originalLoad; + if (originalTsLoader) require.extensions['.ts'] = originalTsLoader; + else delete require.extensions['.ts']; + global.localStorage = originalLocalStorage; + } +} + +test('deleting a background run cancels only that conversation', async () => { + const signals = new Map(); + const settledRuns = []; + + async function* runReactAgent(messages, _model, signal) { + const label = messages.at(-1)?.content; + signals.set(label, signal); + if (!signal.aborted) { + await new Promise((resolve) => signal.addEventListener('abort', resolve, { once: true })); + } + settledRuns.push(label); + } + + await withChatStore( + { + ids: ['conv-a', 'user-a', 'assistant-a', 'conv-b', 'user-b', 'assistant-b'], + runReactAgent, + }, + async ({ useChatStore, deletedConversationIds, getMemoryCalls }) => { + const runA = useChatStore.getState().sendMessage('A', []); + await waitFor(() => signals.has('A'), 'conversation A did not start'); + + assert.equal(useChatStore.getState().createConversation(), 'conv-b'); + const runB = useChatStore.getState().sendMessage('B', []); + await waitFor(() => signals.has('B'), 'conversation B did not start'); + + useChatStore.getState().deleteConversation('conv-a'); + const afterDelete = useChatStore.getState(); + const deletionSnapshot = { + aAborted: signals.get('A').aborted, + bAborted: signals.get('B').aborted, + conversations: afterDelete.conversations.map((conversation) => conversation.id), + activeConversationId: afterDelete.activeConversationId, + isStreaming: afterDelete.isStreaming, + streamingConversationIds: afterDelete.streamingConversationIds, + }; + + useChatStore.getState().stopStreaming(); + useChatStore.getState().setActiveConversation('conv-a'); + useChatStore.getState().stopStreaming(); + useChatStore.getState().setActiveConversation('conv-b'); + await Promise.all([runA, runB]); + + assert.deepEqual(deletionSnapshot, { + aAborted: true, + bAborted: false, + conversations: ['conv-b'], + activeConversationId: 'conv-b', + isStreaming: true, + streamingConversationIds: ['conv-b'], + }); + assert.deepEqual(deletedConversationIds, ['conv-a']); + assert.deepEqual(new Set(settledRuns), new Set(['A', 'B'])); + assert.equal(getMemoryCalls(), 0); + + const finalState = useChatStore.getState(); + assert.deepEqual(finalState.conversations.map((conversation) => conversation.id), ['conv-b']); + assert.deepEqual(finalState.streamingConversationIds, []); + assert.equal(finalState.isStreaming, false); + }, + ); +}); + +test('deleting an ask_user run resolves it with the abort sentinel', async () => { + let signalSeen; + let answerSeen; + + async function* runReactAgent(_messages, _model, signal, onAskUser) { + signalSeen = signal; + answerSeen = await onAskUser('ask-1', 'Choose one'); + } + + await withChatStore( + { + ids: ['conv-question', 'user-question', 'assistant-question'], + runReactAgent, + }, + async ({ useChatStore, getMemoryCalls }) => { + const sendPromise = useChatStore.getState().sendMessage('question', []); + await waitFor(() => useChatStore.getState().pendingQuestion != null, 'ask_user did not become pending'); + + useChatStore.getState().deleteConversation('conv-question'); + const settledAfterDelete = await settlesWithin(sendPromise, 500); + const afterDelete = useChatStore.getState(); + const deletionSnapshot = { + signalAborted: signalSeen.aborted, + settledAfterDelete, + answerSeen, + pendingQuestion: afterDelete.pendingQuestion, + isStreaming: afterDelete.isStreaming, + streamingConversationIds: afterDelete.streamingConversationIds, + }; + + if (!settledAfterDelete) { + useChatStore.getState().answerQuestion('__TEST_CLEANUP__'); + useChatStore.getState().setActiveConversation('conv-question'); + useChatStore.getState().stopStreaming(); + await sendPromise; + } + + assert.deepEqual(deletionSnapshot, { + signalAborted: true, + settledAfterDelete: true, + answerSeen: '__ABORT__', + pendingQuestion: null, + isStreaming: false, + streamingConversationIds: [], + }); + assert.equal(getMemoryCalls(), 0); + }, + ); +}); + +test('deleting before agent startup prevents a late run', async () => { + let releaseUserPersistence; + let userPersistenceStarted = false; + let agentCalls = 0; + const userPersistenceGate = new Promise((resolve) => { + releaseUserPersistence = resolve; + }); + + const runReactAgent = () => { + agentCalls += 1; + return (async function* emptyAgentRun() {})(); + }; + + await withChatStore( + { + ids: ['conv-before-start', 'user-before-start', 'assistant-before-start'], + runReactAgent, + upsertMessage: async (message) => { + if (message.id === 'user-before-start') { + userPersistenceStarted = true; + await userPersistenceGate; + } + }, + }, + async ({ useChatStore, configWrites, deletedConversationIds, getMemoryCalls }) => { + const sendPromise = useChatStore.getState().sendMessage('go', []); + await waitFor(() => userPersistenceStarted, 'user persistence did not start'); + + useChatStore.getState().deleteConversation('conv-before-start'); + releaseUserPersistence(); + await sendPromise; + + const state = useChatStore.getState(); + assert.equal(agentCalls, 0); + assert.equal(getMemoryCalls(), 0); + assert.deepEqual(deletedConversationIds, ['conv-before-start']); + assert.deepEqual(configWrites.at(-1), ['activeConversationId', '']); + assert.deepEqual(state.conversations, []); + assert.deepEqual(state.streamingConversationIds, []); + assert.equal(state.isStreaming, false); + }, + ); +});