diff --git a/apps/cli/src/ai/chat-command.test.ts b/apps/cli/src/ai/chat-command.test.ts index 08ee814b..e79ebbb1 100644 --- a/apps/cli/src/ai/chat-command.test.ts +++ b/apps/cli/src/ai/chat-command.test.ts @@ -34,23 +34,23 @@ function createAIChatManager( chats.set(id, chat) return chat }, - addMessage: ( - aiChatId: string, - role: string, - content: string, - _dataKeywords?: string[], - _dataMessageCount?: number, - contentBlocks?: ContentBlock[], - tokenUsage?: TokenUsageData - ) => { + addMessagePair: (aiChatId, userMessage, assistantMessage) => { + const savedUserMessage = { id: `msg_${messages.length + 1}`, aiChatId, role: 'user' as const, timestamp: 1 } + messages.push({ aiChatId, role: 'user', content: userMessage.content }) + const savedAssistantMessage = { + id: `msg_${messages.length + 1}`, + aiChatId, + role: 'assistant' as const, + timestamp: 1, + } messages.push({ aiChatId, - role, - content, - ...(contentBlocks ? { contentBlocks } : {}), - ...(tokenUsage ? { tokenUsage } : {}), + role: 'assistant', + content: assistantMessage.content, + ...(assistantMessage.contentBlocks ? { contentBlocks: assistantMessage.contentBlocks } : {}), + ...(assistantMessage.tokenUsage ? { tokenUsage: assistantMessage.tokenUsage } : {}), }) - return { id: `msg_${messages.length}`, aiChatId, role, content, timestamp: 1 } + return { userMessage: savedUserMessage, assistantMessage: savedAssistantMessage } }, __messages: messages, } as unknown as AIChatManager diff --git a/apps/cli/src/ai/chat-command.ts b/apps/cli/src/ai/chat-command.ts index 2ea8e405..537a6d5f 100644 --- a/apps/cli/src/ai/chat-command.ts +++ b/apps/cli/src/ai/chat-command.ts @@ -317,15 +317,14 @@ export async function runChatTurn( write(stdout, '\n') } - deps.aiChatManager.addMessage(target.aiChatId, 'user', options.question) - deps.aiChatManager.addMessage( + deps.aiChatManager.addMessagePair( target.aiChatId, - 'assistant', - answer, - undefined, - undefined, - hasReplayContentBlocks ? contentBlocks : undefined, - tokenUsage ?? undefined + { content: options.question }, + { + content: answer, + contentBlocks: hasReplayContentBlocks ? contentBlocks : undefined, + tokenUsage: tokenUsage ?? undefined, + } ) return { diff --git a/packages/http-routes/src/routes/web/ai-chats.test.ts b/packages/http-routes/src/routes/web/ai-chats.test.ts index 3085fd00..e50fc7d4 100644 --- a/packages/http-routes/src/routes/web/ai-chats.test.ts +++ b/packages/http-routes/src/routes/web/ai-chats.test.ts @@ -48,21 +48,22 @@ test('global AI chat routes keep global history separate and persist entity refe sessionType: 'group' as const, }, ] - const messageResponse = await app.inject({ + const messagePairResponse = await app.inject({ method: 'POST', - url: `/_web/ai/chats/${globalChat.id}/messages`, - payload: { role: 'user', content: 'Compare them', entityRefs: refs }, + url: `/_web/ai/chats/${globalChat.id}/message-pair`, + payload: { + userMessage: { content: 'Compare them', entityRefs: refs }, + assistantMessage: { content: 'I will compare them' }, + }, }) - assert.equal(messageResponse.statusCode, 200) - assert.deepEqual(messageResponse.json<{ entityRefs: unknown[] }>().entityRefs, refs) - - const assistantResponse = await app.inject({ - method: 'POST', - url: `/_web/ai/chats/${globalChat.id}/messages`, - payload: { role: 'assistant', content: 'I will compare them', entityRefs: refs }, - }) - assert.equal(assistantResponse.statusCode, 200) - assert.equal(assistantResponse.json<{ entityRefs?: unknown[] }>().entityRefs, undefined) + assert.equal(messagePairResponse.statusCode, 200) + const messagePair = messagePairResponse.json<{ + userMessage: { id: string; entityRefs: unknown[] } + assistantMessage: { parentId: string; entityRefs?: unknown[] } + }>() + assert.deepEqual(messagePair.userMessage.entityRefs, refs) + assert.equal(messagePair.assistantMessage.parentId, messagePair.userMessage.id) + assert.equal(messagePair.assistantMessage.entityRefs, undefined) const messageList = await app.inject({ method: 'GET', @@ -78,7 +79,7 @@ test('global AI chat routes keep global history separate and persist entity refe sessionType: 'group' as const, }, ] - const userMessageId = messageResponse.json<{ id: string }>().id + const userMessageId = messagePair.userMessage.id const replaceResponse = await app.inject({ method: 'PUT', url: `/_web/ai/messages/${userMessageId}/content`, diff --git a/packages/http-routes/src/routes/web/ai-chats.ts b/packages/http-routes/src/routes/web/ai-chats.ts index c8df04b8..9cfc428e 100644 --- a/packages/http-routes/src/routes/web/ai-chats.ts +++ b/packages/http-routes/src/routes/web/ai-chats.ts @@ -1,7 +1,7 @@ import type { FastifyInstance } from 'fastify' import type { AiRouteContext } from '../../context/ai' import { countMessagesTokens } from '@openchatlab/node-runtime' -import type { AIEntityRef } from '@openchatlab/node-runtime' +import type { AIEntityRef, ContentBlock, TokenUsageData } from '@openchatlab/node-runtime' type AiChatRouteContext = Pick @@ -56,6 +56,20 @@ export function registerAiChatRoutes(server: FastifyInstance, ctx: AiChatRouteCo // ==================== Message CRUD ==================== + server.post<{ + Params: { id: string } + Body: { + userMessage: { content: string; entityRefs?: AIEntityRef[] } + assistantMessage: { + content: string + contentBlocks?: ContentBlock[] + tokenUsage?: TokenUsageData + } + } + }>('/_web/ai/chats/:id/message-pair', async (request) => { + return cm.addMessagePair(request.params.id, request.body.userMessage, request.body.assistantMessage) + }) + server.post<{ Params: { id: string } Body: { diff --git a/packages/node-runtime/src/ai/__tests__/chats.test.ts b/packages/node-runtime/src/ai/__tests__/chats.test.ts index 07daf7e2..6a921a85 100644 --- a/packages/node-runtime/src/ai/__tests__/chats.test.ts +++ b/packages/node-runtime/src/ai/__tests__/chats.test.ts @@ -434,7 +434,54 @@ describe('AIChatManager global chats and entity references', () => { }) }) -describe('AIChatManager message editing', () => { +describe('AIChatManager message persistence and editing', () => { + it('persists a user and assistant message atomically', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + const chat = manager.createAIChat('s1', 'Atomic turn', 'general_cn') + const refs: AIEntityRef[] = [{ type: 'contact', contactKey: 'qq:10001', displayName: 'Alice' }] + const usage = { + promptTokens: 1, + completionTokens: 2, + totalTokens: 3, + cacheReadTokens: 0, + cacheWriteTokens: 0, + } + const saved = manager.addMessagePair( + chat.id, + { content: 'question', entityRefs: refs }, + { content: 'answer', contentBlocks: [{ type: 'text', text: 'answer' }], tokenUsage: usage } + ) + + assert.equal(saved.assistantMessage.parentId, saved.userMessage.id) + assert.deepEqual(saved.userMessage.entityRefs, refs) + assert.deepEqual(saved.assistantMessage.contentBlocks, [{ type: 'text', text: 'answer' }]) + assert.deepEqual(manager.getAIChatTokenUsage(chat.id), usage) + + const messagesBeforeFailure = manager.getMessages(chat.id) + manager.executeAiSQL(` + CREATE TRIGGER reject_assistant_message + BEFORE INSERT ON ai_message + WHEN NEW.role = 'assistant' + BEGIN + SELECT RAISE(ABORT, 'injected assistant failure'); + END + `) + + assert.throws( + () => manager.addMessagePair(chat.id, { content: 'orphan' }, { content: 'must roll back' }), + /injected assistant failure/ + ) + assert.deepEqual(manager.getMessages(chat.id), messagesBeforeFailure) + assert.equal(manager.getAIChat(chat.id)?.activeMessageId, saved.assistantMessage.id) + assert.deepEqual(manager.getAIChatTokenUsage(chat.id), usage) + manager.close() + } finally { + cleanup(dir) + } + }) + it('updateMessageContent atomically replaces or clears user entity references', () => { const dir = createTempDir() try { diff --git a/packages/node-runtime/src/ai/chats.ts b/packages/node-runtime/src/ai/chats.ts index f1f77719..9b13549f 100644 --- a/packages/node-runtime/src/ai/chats.ts +++ b/packages/node-runtime/src/ai/chats.ts @@ -671,6 +671,35 @@ export class AIChatManager { ) } + addMessagePair( + aiChatId: string, + userMessage: { content: string; entityRefs?: AIEntityRef[] }, + assistantMessage: { content: string; contentBlocks?: ContentBlock[]; tokenUsage?: TokenUsageData } + ): { userMessage: AIMessage; assistantMessage: AIMessage } { + const db = this.getDb() + return db.transaction(() => ({ + userMessage: this.addMessage( + aiChatId, + 'user', + userMessage.content, + undefined, + undefined, + undefined, + undefined, + userMessage.entityRefs + ), + assistantMessage: this.addMessage( + aiChatId, + 'assistant', + assistantMessage.content, + undefined, + undefined, + assistantMessage.contentBlocks, + assistantMessage.tokenUsage + ), + }))() + } + private addMessageWithEntityRefs( aiChatId: string, role: AIMessageRole, diff --git a/src/services/ai/fetch.ts b/src/services/ai/fetch.ts index 60fb1205..0b49fce5 100644 --- a/src/services/ai/fetch.ts +++ b/src/services/ai/fetch.ts @@ -92,6 +92,17 @@ export class FetchAIAdapter implements AIAdapter { return get(`/ai/chats/${aiChatId}/messages`) } + async addMessagePair( + aiChatId: string, + userMessage: { content: string; entityRefs?: AIEntityRef[] }, + assistantMessage: { content: string; contentBlocks?: ContentBlock[]; tokenUsage?: TokenUsageData } + ): Promise<{ userMessage: AIMessage; assistantMessage: AIMessage }> { + return post<{ userMessage: AIMessage; assistantMessage: AIMessage }>(`/ai/chats/${aiChatId}/message-pair`, { + userMessage, + assistantMessage, + }) + } + async addMessage( aiChatId: string, role: AIMessageRole, diff --git a/src/services/ai/types.ts b/src/services/ai/types.ts index cc8cb449..da818663 100644 --- a/src/services/ai/types.ts +++ b/src/services/ai/types.ts @@ -141,6 +141,11 @@ export interface AIAdapter { // ===== 消息 ===== getMessages(aiChatId: string): Promise + addMessagePair( + aiChatId: string, + userMessage: { content: string; entityRefs?: AIEntityRef[] }, + assistantMessage: { content: string; contentBlocks?: ContentBlock[]; tokenUsage?: TokenUsageData } + ): Promise<{ userMessage: AIMessage; assistantMessage: AIMessage }> addMessage( aiChatId: string, role: AIMessageRole, diff --git a/src/stores/aiChat.ts b/src/stores/aiChat.ts index bc4a9733..0f22e4d1 100644 --- a/src/stores/aiChat.ts +++ b/src/stores/aiChat.ts @@ -1434,13 +1434,12 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { targetBuffer.messages[aiMessageIndex], lastDoneUsage ) - if (savedMessages) { - Object.assign(userMessage, savedMessages.userMessage) - targetBuffer.messages[aiMessageIndex] = { - ...targetBuffer.messages[aiMessageIndex], - ...savedMessages.assistantMessage, - isStreaming: false, - } + if (!savedMessages) return { success: false, reason: 'error' } + Object.assign(userMessage, savedMessages.userMessage) + targetBuffer.messages[aiMessageIndex] = { + ...targetBuffer.messages[aiMessageIndex], + ...savedMessages.assistantMessage, + isStreaming: false, } } else if (!hasStreamError) { const blocks = targetBuffer.messages[aiMessageIndex].contentBlocks || [] @@ -1459,13 +1458,12 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { targetBuffer.messages[aiMessageIndex], lastDoneUsage ) - if (savedMessages) { - Object.assign(userMessage, savedMessages.userMessage) - targetBuffer.messages[aiMessageIndex] = { - ...targetBuffer.messages[aiMessageIndex], - ...savedMessages.assistantMessage, - isStreaming: false, - } + if (!savedMessages) return { success: false, reason: 'error' } + Object.assign(userMessage, savedMessages.userMessage) + targetBuffer.messages[aiMessageIndex] = { + ...targetBuffer.messages[aiMessageIndex], + ...savedMessages.assistantMessage, + isStreaming: false, } } else { const savedMessages = await saveAIChatMessages( @@ -1474,13 +1472,12 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { targetBuffer.messages[aiMessageIndex], lastDoneUsage ) - if (savedMessages) { - Object.assign(userMessage, savedMessages.userMessage) - targetBuffer.messages[aiMessageIndex] = { - ...targetBuffer.messages[aiMessageIndex], - ...savedMessages.assistantMessage, - isStreaming: false, - } + if (!savedMessages) return { success: false, reason: 'error' } + Object.assign(userMessage, savedMessages.userMessage) + targetBuffer.messages[aiMessageIndex] = { + ...targetBuffer.messages[aiMessageIndex], + ...savedMessages.assistantMessage, + isStreaming: false, } } @@ -1534,26 +1531,13 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { return null } - const savedUserMessage = await useAIService().addMessage( - aiChatId, - 'user', - userMsg.content, - undefined, - undefined, - undefined, - undefined, - userMsg.entityRefs - ) const serializableContentBlocks = toSerializableContentBlocks(aiMsg.contentBlocks) - const savedAssistantMessage = await useAIService().addMessage( - aiChatId, - 'assistant', - aiMsg.content, - undefined, - undefined, - serializableContentBlocks, - tokenUsage - ) + const { userMessage: savedUserMessage, assistantMessage: savedAssistantMessage } = + await useAIService().addMessagePair( + aiChatId, + { content: userMsg.content, entityRefs: userMsg.entityRefs }, + { content: aiMsg.content, contentBlocks: serializableContentBlocks, tokenUsage } + ) return { userMessage: toRuntimeMessage(savedUserMessage), assistantMessage: toRuntimeMessage(savedAssistantMessage), diff --git a/src/stores/aiChatPersistence.test.ts b/src/stores/aiChatPersistence.test.ts new file mode 100644 index 00000000..062edbf6 --- /dev/null +++ b/src/stores/aiChatPersistence.test.ts @@ -0,0 +1,111 @@ +import assert from 'node:assert/strict' +import test from 'node:test' +import type { AgentStreamChunk } from '@/services/ai-stream/types' +import { createPinia, setActivePinia } from 'pinia' +import { ref } from 'vue' + +test('reports failure when an accepted turn cannot be persisted atomically', async (t) => { + let messagePairWrites = 0 + let singleMessageWrites = 0 + const persistenceError = new Error('injected persistence failure') + const aiService = { + createAIChat: async (sessionId: string, title: string, assistantId: string) => ({ + id: 'chat-1', + sessionId, + kind: 'session' as const, + title, + assistantId, + createdAt: 1, + updatedAt: 1, + }), + addMessagePair: async () => { + messagePairWrites++ + throw persistenceError + }, + addMessage: async (aiChatId: string, role: 'user' | 'assistant', content: string) => { + singleMessageWrites++ + if (role === 'assistant') throw persistenceError + return { id: 'persisted-user', aiChatId, role, content, timestamp: 1 } + }, + } + + await Promise.all([ + t.mock.module('@/stores/session', { + namedExports: { useSessionStore: () => ({ sessions: [] }) }, + }), + t.mock.module('@/stores/settings', { + namedExports: { + useSettingsStore: () => ({ + aiPreprocessConfig: { + dataCleaning: false, + mergeConsecutive: false, + mergeWindowSeconds: 180, + blacklistKeywords: [], + denoise: false, + desensitize: false, + desensitizeRules: [], + anonymizeNames: false, + }, + }), + }, + }), + t.mock.module('@/stores/assistant', { + namedExports: { + useAssistantStore: () => ({ + isLoaded: true, + loadAssistants: async () => undefined, + selectAssistant: () => undefined, + clearSelection: () => undefined, + }), + }, + }), + t.mock.module('@/stores/skill', { + namedExports: { useSkillStore: () => ({ activeSkillId: null, activeSkill: ref(null) }) }, + }), + t.mock.module('@/stores/llm', { + namedExports: { useLLMStore: () => ({}) }, + }), + t.mock.module('@/services', { + namedExports: { + useAIService: () => aiService, + useDataService: () => ({}), + useLLMService: () => ({ hasConfig: async () => true }), + }, + }), + t.mock.module('@/services/ai-stream/service', { + namedExports: { + useAgentStreamService: () => ({ + runStream: (_params: unknown, onChunk?: (chunk: AgentStreamChunk) => void) => { + onChunk?.({ type: 'content', content: 'model answer' }) + onChunk?.({ type: 'done', isFinished: true }) + return { + requestId: 'stream-1', + promise: Promise.resolve({ + success: true, + result: { content: 'model answer', toolsUsed: [], toolRounds: 0 }, + }), + } + }, + }), + }, + }), + ]) + + t.mock.method(console, 'error', () => undefined) + setActivePinia(createPinia()) + const { useAIChatStore } = await import('./aiChat') + const store = useAIChatStore() + const { chatKey } = store.ensureSessionState({ + sessionId: 'session-1', + sessionName: 'Test session', + chatType: 'private', + locale: 'zh-CN', + }) + store.selectAssistantForSession(chatKey, 'general_cn') + + const result = await store.sendMessage(chatKey, 'atomic question') + + assert.deepEqual(result, { success: false, reason: 'error' }) + assert.equal(messagePairWrites, 1) + assert.equal(singleMessageWrites, 0) +})