diff --git a/apps/cli/src/ai/agent.ts b/apps/cli/src/ai/agent.ts index e4e15b9c..1c2d2a19 100644 --- a/apps/cli/src/ai/agent.ts +++ b/apps/cli/src/ai/agent.ts @@ -102,7 +102,7 @@ export async function runServerAgent(options: RunAgentOptions): Promise { systemPrompt, }) - if (compressionConfig?.enabled) { + if (compressionConfig?.enabled && historyLeafMessageId === undefined) { const llmAdapter = createCompressionLlmAdapter({ piModel, apiKey: llmConfig.apiKey, @@ -127,6 +127,11 @@ export async function runServerAgent(options: RunAgentOptions): Promise { }, }) } + } else if (compressionConfig?.enabled && historyLeafMessageId !== undefined) { + aiLogger?.info?.('Compression', 'Skipping compression for edited branch request', { + conversationId, + historyLeafMessageId, + }) } if (abortSignal?.aborted) { diff --git a/apps/cli/src/http/routes/ai.ts b/apps/cli/src/http/routes/ai.ts index 528af13c..57057e52 100644 --- a/apps/cli/src/http/routes/ai.ts +++ b/apps/cli/src/http/routes/ai.ts @@ -522,28 +522,72 @@ export function registerAiRoutes( }) server.post<{ - Params: { id: string } - Body: { - content: string - assistantContent: string - contentBlocks?: unknown[] - tokenUsage?: { promptTokens: number; completionTokens: number; totalTokens: number } + Params: { id: string; messageId: string } + }>('/_web/ai/conversations/:id/messages/:messageId/delete-from', async (request, reply) => { + const { id, messageId } = request.params + if (!id || typeof id !== 'string' || !messageId || typeof messageId !== 'string') { + return reply.code(400).send({ error: 'conversation id and messageId are required' }) } - }>('/_web/ai/messages/:id/branches', async (request) => { - return convManager.createMessageBranch( - request.params.id, - request.body.content, - request.body.assistantContent, - request.body.contentBlocks as any, - request.body.tokenUsage - ) + convManager.deleteMessagesFrom(id, messageId) + return { success: true } }) server.post<{ Params: { id: string } - Body: { messageId: string } - }>('/_web/ai/conversations/:id/branches/switch', async (request) => { - return convManager.switchMessageBranch(request.params.id, request.body.messageId) + Body: { upToMessageId: string; title?: string } + }>('/_web/ai/conversations/:id/fork', async (request, reply) => { + const { upToMessageId, title } = request.body + if (!upToMessageId || typeof upToMessageId !== 'string') { + return reply.code(400).send({ error: 'upToMessageId is required' }) + } + return convManager.forkConversation(request.params.id, upToMessageId, title) + }) + + server.put<{ + Params: { messageId: string } + Body: { content: string } + }>('/_web/ai/messages/:messageId/content', async (request, reply) => { + const { content } = request.body + if (!content || typeof content !== 'string') { + return reply.code(400).send({ error: 'content is required' }) + } + convManager.updateMessageContent(request.params.messageId, content) + return { success: true } + }) + + server.post<{ + Params: { id: string; messageId: string } + }>('/_web/ai/conversations/:id/messages/:messageId/delete-relink', async (request, reply) => { + const { id, messageId } = request.params + if (!id || !messageId) { + return reply.code(400).send({ error: 'conversation id and messageId are required' }) + } + convManager.deleteAndRelinkMessage(id, messageId) + return { success: true } + }) + + server.post<{ + Params: { id: string } + Body: { + afterMessageId: string + role: 'user' | 'assistant' | 'summary' + content: string + contentBlocks?: unknown[] + tokenUsage?: { promptTokens: number; completionTokens: number; totalTokens: number } + } + }>('/_web/ai/conversations/:id/messages/insert-after', async (request, reply) => { + const { afterMessageId, role, content, contentBlocks, tokenUsage } = request.body + if (!afterMessageId || typeof afterMessageId !== 'string') { + return reply.code(400).send({ error: 'afterMessageId is required' }) + } + return convManager.insertMessageAfter( + request.params.id, + afterMessageId, + role, + content, + contentBlocks as any, + tokenUsage + ) }) server.get<{ Params: { id: string } }>('/_web/ai/conversations/:id/token-usage', async (request) => { diff --git a/apps/desktop/main/ai/conversations.ts b/apps/desktop/main/ai/conversations.ts index e3edfeee..40ef11a7 100644 --- a/apps/desktop/main/ai/conversations.ts +++ b/apps/desktop/main/ai/conversations.ts @@ -9,14 +9,7 @@ import { AIConversationManager } from '@openchatlab/node-runtime' import { getPathProvider } from '../path-context' import { aiLogger } from './logger' -export type { - AIConversation, - AIMessage, - AIMessageRole, - ContentBlock, - TokenUsageData, - MessageBranchResult, -} from '@openchatlab/node-runtime' +export type { AIConversation, AIMessage, AIMessageRole, ContentBlock, TokenUsageData } from '@openchatlab/node-runtime' let manager: AIConversationManager | null = null @@ -100,24 +93,31 @@ export function deleteMessage(messageId: string) { return getManager().deleteMessage(messageId) } -export function createMessageBranch( - originalUserMessageId: string, - newUserContent: string, - assistantContent: string, +export function deleteMessagesFrom(conversationId: string, messageId: string) { + return getManager().deleteMessagesFrom(conversationId, messageId) +} + +export function forkConversation(sourceConversationId: string, upToMessageId: string, title?: string) { + return getManager().forkConversation(sourceConversationId, upToMessageId, title) +} + +export function updateMessageContent(messageId: string, newContent: string) { + return getManager().updateMessageContent(messageId, newContent) +} + +export function deleteAndRelinkMessage(conversationId: string, messageId: string) { + return getManager().deleteAndRelinkMessage(conversationId, messageId) +} + +export function insertMessageAfter( + conversationId: string, + afterMessageId: string, + role: import('@openchatlab/node-runtime').AIMessageRole, + content: string, contentBlocks?: import('@openchatlab/node-runtime').ContentBlock[], tokenUsage?: import('@openchatlab/node-runtime').TokenUsageData ) { - return getManager().createMessageBranch( - originalUserMessageId, - newUserContent, - assistantContent, - contentBlocks, - tokenUsage - ) -} - -export function switchMessageBranch(conversationId: string, messageId: string) { - return getManager().switchMessageBranch(conversationId, messageId) + return getManager().insertMessageAfter(conversationId, afterMessageId, role, content, contentBlocks, tokenUsage) } export function setPendingDebugContext(conversationId: string, debugContext: string) { diff --git a/apps/desktop/main/ipc/ai.ts b/apps/desktop/main/ipc/ai.ts index 0dd4b04c..505f2053 100644 --- a/apps/desktop/main/ipc/ai.ts +++ b/apps/desktop/main/ipc/ai.ts @@ -284,40 +284,72 @@ export function registerAIHandlers({ win }: IpcContext): void { } }) + ipcMain.handle('ai:deleteMessagesFrom', async (_, conversationId: string, messageId: string) => { + try { + return aiConversations.deleteMessagesFrom(conversationId, messageId) + } catch (error) { + console.error('Failed to delete messages:', error) + throw error + } + }) + ipcMain.handle( - 'ai:createMessageBranch', - async ( - _, - originalUserMessageId: string, - newUserContent: string, - assistantContent: string, - contentBlocks?: aiConversations.ContentBlock[], - tokenUsage?: aiConversations.TokenUsageData - ) => { + 'ai:forkConversation', + async (_, sourceConversationId: string, upToMessageId: string, title?: string) => { try { - return aiConversations.createMessageBranch( - originalUserMessageId, - newUserContent, - assistantContent, - contentBlocks, - tokenUsage - ) + return aiConversations.forkConversation(sourceConversationId, upToMessageId, title) } catch (error) { - console.error('Failed to create AI message branch:', error) + console.error('Failed to fork conversation:', error) throw error } } ) - ipcMain.handle('ai:switchMessageBranch', async (_, conversationId: string, messageId: string) => { + ipcMain.handle('ai:updateMessageContent', async (_, messageId: string, newContent: string) => { try { - return aiConversations.switchMessageBranch(conversationId, messageId) + return aiConversations.updateMessageContent(messageId, newContent) } catch (error) { - console.error('Failed to switch AI message branch:', error) + console.error('Failed to update message content:', error) throw error } }) + ipcMain.handle('ai:deleteAndRelinkMessage', async (_, conversationId: string, messageId: string) => { + try { + return aiConversations.deleteAndRelinkMessage(conversationId, messageId) + } catch (error) { + console.error('Failed to delete and relink message:', error) + throw error + } + }) + + ipcMain.handle( + 'ai:insertMessageAfter', + async ( + _, + conversationId: string, + afterMessageId: string, + role: aiConversations.AIMessageRole, + content: string, + contentBlocks?: aiConversations.ContentBlock[], + tokenUsage?: aiConversations.TokenUsageData + ) => { + try { + return aiConversations.insertMessageAfter( + conversationId, + afterMessageId, + role, + content, + contentBlocks, + tokenUsage + ) + } catch (error) { + console.error('Failed to insert message after:', error) + throw error + } + } + ) + // ==================== 脱敏规则 ==================== ipcMain.handle('ai:getDefaultDesensitizeRules', (_, locale: string) => { @@ -993,7 +1025,7 @@ export function registerAIHandlers({ win }: IpcContext): void { const piModel = buildPiModel(activeAIConfig) // 上下文压缩前置步骤（在 Agent 创建之前执行） - if (compressionConfig?.enabled && context.conversationId) { + if (compressionConfig?.enabled && context.conversationId && context.historyLeafMessageId === undefined) { try { // 获取助手 systemPrompt 用于 token 计算 const tempAssistantConfig = assistantId @@ -1046,6 +1078,11 @@ export function registerAIHandlers({ win }: IpcContext): void { error: String(error), }) } + } else if (compressionConfig?.enabled && context.historyLeafMessageId !== undefined) { + aiLogger.info('IPC', `Skipping compression for edited branch request: ${requestId}`, { + conversationId: context.conversationId, + historyLeafMessageId: context.historyLeafMessageId, + }) } const pp = context.preprocessConfig diff --git a/apps/desktop/preload/apis/ai.ts b/apps/desktop/preload/apis/ai.ts index dde5a49d..03689352 100644 --- a/apps/desktop/preload/apis/ai.ts +++ b/apps/desktop/preload/apis/ai.ts @@ -64,25 +64,12 @@ export interface AIMessage { content: string timestamp: number parentId?: string | null - siblingGroupId?: string - branchIndex?: number - branch?: { - index: number - total: number - prevMessageId: string | null - nextMessageId: string | null - } dataKeywords?: string[] dataMessageCount?: number contentBlocks?: ContentBlock[] tokenUsage?: TokenUsageData } -export interface MessageBranchResult { - userMessage: AIMessage - assistantMessage: AIMessage -} - // LLM API 类型 export interface LLMProvider { id: string @@ -515,27 +502,41 @@ export const aiApi = { ) }, - createMessageBranch: ( - originalUserMessageId: string, - newUserContent: string, - assistantContent: string, + deleteMessagesFrom: (conversationId: string, messageId: string): Promise => { + return ipcRenderer.invoke('ai:deleteMessagesFrom', conversationId, messageId) + }, + + forkConversation: (sourceConversationId: string, upToMessageId: string, title?: string): Promise => { + return ipcRenderer.invoke('ai:forkConversation', sourceConversationId, upToMessageId, title) + }, + + updateMessageContent: (messageId: string, newContent: string): Promise => { + return ipcRenderer.invoke('ai:updateMessageContent', messageId, newContent) + }, + + deleteAndRelinkMessage: (conversationId: string, messageId: string): Promise => { + return ipcRenderer.invoke('ai:deleteAndRelinkMessage', conversationId, messageId) + }, + + insertMessageAfter: ( + conversationId: string, + afterMessageId: string, + role: AIMessageRole, + content: string, contentBlocks?: ContentBlock[], tokenUsage?: TokenUsageData - ): Promise => { + ): Promise => { return ipcRenderer.invoke( - 'ai:createMessageBranch', - originalUserMessageId, - newUserContent, - assistantContent, + 'ai:insertMessageAfter', + conversationId, + afterMessageId, + role, + content, contentBlocks, tokenUsage ) }, - switchMessageBranch: (conversationId: string, messageId: string): Promise => { - return ipcRenderer.invoke('ai:switchMessageBranch', conversationId, messageId) - }, - /** * 获取 AI 对话的所有消息 */ diff --git a/apps/desktop/preload/index.d.ts b/apps/desktop/preload/index.d.ts index 18c9a826..64323b10 100644 --- a/apps/desktop/preload/index.d.ts +++ b/apps/desktop/preload/index.d.ts @@ -345,25 +345,12 @@ interface AIMessage { content: string timestamp: number parentId?: string | null - siblingGroupId?: string - branchIndex?: number - branch?: { - index: number - total: number - prevMessageId: string | null - nextMessageId: string | null - } dataKeywords?: string[] dataMessageCount?: number contentBlocks?: AIContentBlock[] tokenUsage?: AITokenUsageData } -interface MessageBranchResult { - userMessage: AIMessage - assistantMessage: AIMessage -} - interface AiApi { searchMessages: ( sessionId: string, @@ -425,14 +412,18 @@ interface AiApi { contentBlocks?: AIContentBlock[], tokenUsage?: AITokenUsageData ) => Promise - createMessageBranch: ( - originalUserMessageId: string, - newUserContent: string, - assistantContent: string, + deleteMessagesFrom: (conversationId: string, messageId: string) => Promise + forkConversation: (sourceConversationId: string, upToMessageId: string, title?: string) => Promise + updateMessageContent: (messageId: string, newContent: string) => Promise + deleteAndRelinkMessage: (conversationId: string, messageId: string) => Promise + insertMessageAfter: ( + conversationId: string, + afterMessageId: string, + role: 'user' | 'assistant', + content: string, contentBlocks?: AIContentBlock[], tokenUsage?: AITokenUsageData - ) => Promise - switchMessageBranch: (conversationId: string, messageId: string) => Promise + ) => Promise getMessages: (conversationId: string) => Promise getConversationTokenUsage: (conversationId: string) => Promise deleteMessage: (messageId: string) => Promise @@ -1250,7 +1241,6 @@ export { SearchMessageResult, AIConversation, AIMessage, - MessageBranchResult, LLMProviderInfo, AIServiceConfigDisplay, LLMChatMessage, diff --git a/packages/node-runtime/src/ai/__tests__/conversations.test.ts b/packages/node-runtime/src/ai/__tests__/conversations.test.ts index bfb45c8a..3df8bc29 100644 --- a/packages/node-runtime/src/ai/__tests__/conversations.test.ts +++ b/packages/node-runtime/src/ai/__tests__/conversations.test.ts @@ -30,7 +30,7 @@ function cleanup(dir: string): void { } } -describe('AIConversationManager message branches', () => { +describe('AIConversationManager legacy migration', () => { it('migrates legacy flat messages into an active path', () => { const dir = createTempDir() try { @@ -86,117 +86,296 @@ describe('AIConversationManager message branches', () => { } }) - it('does not rerun legacy backfill over existing root branches', () => { + it('repairs partially migrated legacy rows that already have tree columns', () => { const dir = createTempDir() try { - let manager = createManager(dir) - const conversation = manager.createConversation('session-1', 'Root branch', 'general_cn') - const root = manager.addMessage(conversation.id, 'user', 'original root') - manager.addMessage(conversation.id, 'assistant', 'original answer') - const branch = manager.createMessageBranch(root.id, 'edited root', 'edited answer') - assert.equal(branch.userMessage.parentId, null) - manager.close() + const db = createTestDatabase(join(dir, 'conversations.db')) + db.exec(` + CREATE TABLE ai_conversation ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + title TEXT, + assistant_id TEXT DEFAULT 'general_cn', + active_message_id TEXT, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ); + CREATE TABLE ai_message ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL, + timestamp INTEGER NOT NULL, + data_keywords TEXT, + data_message_count INTEGER, + content_blocks TEXT, + token_usage TEXT, + debug_context TEXT, + parent_id TEXT, + sibling_group_id TEXT, + branch_index INTEGER DEFAULT 0 + ); + `) + db.prepare('INSERT INTO ai_conversation VALUES (?, ?, ?, ?, NULL, ?, ?)').run( + 'conv-partial', + 'session-1', + 'Partial', + 'general_cn', + 1, + 3 + ) + db.prepare('INSERT INTO ai_message VALUES (?, ?, ?, ?, ?, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL)').run( + 'pm1', + 'conv-partial', + 'user', + 'first', + 1 + ) + db.prepare('INSERT INTO ai_message VALUES (?, ?, ?, ?, ?, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL)').run( + 'pm2', + 'conv-partial', + 'assistant', + 'second', + 2 + ) + db.close() - manager = createManager(dir) - const messages = manager.getMessages(conversation.id) + const manager = createManager(dir) + const messages = manager.getMessages('conv-partial') assert.deepEqual( messages.map((message) => message.content), - ['edited root', 'edited answer'] + ['first', 'second'] ) assert.equal(messages[0]?.parentId, null) - assert.equal(manager.getConversation(conversation.id)?.activeMessageId, branch.assistantMessage.id) + assert.equal(messages[1]?.parentId, 'pm1') + assert.equal(manager.getConversation('conv-partial')?.activeMessageId, 'pm2') manager.close() } finally { cleanup(dir) } }) +}) - it('creates editable branches without losing the original path', () => { +describe('AIConversationManager message editing', () => { + it('updateMessageContent updates message text in place', () => { const dir = createTempDir() try { const manager = createManager(dir) - const conversation = manager.createConversation('session-1', 'Branch', 'general_cn') - const user1 = manager.addMessage(conversation.id, 'user', 'original question') - manager.addMessage(conversation.id, 'assistant', 'original answer', undefined, undefined, undefined, { - promptTokens: 1, - completionTokens: 2, - totalTokens: 3, - }) - manager.addMessage(conversation.id, 'user', 'follow up') - const oldLeaf = manager.addMessage( - conversation.id, - 'assistant', - 'follow answer', - undefined, - undefined, - undefined, - { - promptTokens: 4, - completionTokens: 5, - totalTokens: 9, - } + const conv = manager.createConversation('s1', 'Test', 'general_cn') + const msg = manager.addMessage(conv.id, 'user', 'original text') + manager.addMessage(conv.id, 'assistant', 'reply') + + manager.updateMessageContent(msg.id, 'edited text') + + const messages = manager.getMessages(conv.id) + assert.equal(messages[0]?.content, 'edited text') + assert.equal(messages[1]?.content, 'reply') + manager.close() + } finally { + cleanup(dir) + } + }) + + it('updateMessageContent throws for non-existent message', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + assert.throws(() => manager.updateMessageContent('non-existent', 'text'), /Message not found/) + manager.close() + } finally { + cleanup(dir) + } + }) + + it('deleteAndRelinkMessage removes a message and rewires children', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + const conv = manager.createConversation('s1', 'Test', 'general_cn') + const userMsg = manager.addMessage(conv.id, 'user', 'question') + const aiMsg = manager.addMessage(conv.id, 'assistant', 'answer') + const followUp = manager.addMessage(conv.id, 'user', 'follow up') + manager.addMessage(conv.id, 'assistant', 'follow answer') + + manager.deleteAndRelinkMessage(conv.id, aiMsg.id) + + const messages = manager.getMessages(conv.id) + assert.equal(messages.length, 3) + assert.deepEqual( + messages.map((m) => m.content), + ['question', 'follow up', 'follow answer'] ) + // follow up's parent should now be userMsg (was aiMsg) + assert.equal(messages[1]?.parentId, userMsg.id) + assert.equal(messages[2]?.parentId, followUp.id) + manager.close() + } finally { + cleanup(dir) + } + }) + + it('deleteAndRelinkMessage updates activeMessageId when removing the active leaf', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + const conv = manager.createConversation('s1', 'Test', 'general_cn') + const userMsg = manager.addMessage(conv.id, 'user', 'question') + const aiMsg = manager.addMessage(conv.id, 'assistant', 'answer') + assert.equal(manager.getConversation(conv.id)?.activeMessageId, aiMsg.id) + + manager.deleteAndRelinkMessage(conv.id, aiMsg.id) + + assert.equal(manager.getConversation(conv.id)?.activeMessageId, userMsg.id) + const messages = manager.getMessages(conv.id) + assert.equal(messages.length, 1) + assert.equal(messages[0]?.content, 'question') + manager.close() + } finally { + cleanup(dir) + } + }) - const branch = manager.createMessageBranch('' + user1.id, 'edited question', 'edited answer', undefined, { - promptTokens: 10, - completionTokens: 11, - totalTokens: 21, - }) + it('insertMessageAfter inserts a message in the middle of a chain', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + const conv = manager.createConversation('s1', 'Test', 'general_cn') + const userMsg = manager.addMessage(conv.id, 'user', 'question') + const followUp = manager.addMessage(conv.id, 'user', 'follow up') + manager.addMessage(conv.id, 'assistant', 'follow answer') + + const inserted = manager.insertMessageAfter(conv.id, userMsg.id, 'assistant', 'inserted answer') + const messages = manager.getMessages(conv.id) + assert.equal(messages.length, 4) assert.deepEqual( - manager.getMessages(conversation.id).map((message) => message.content), - ['edited question', 'edited answer'] + messages.map((m) => m.content), + ['question', 'inserted answer', 'follow up', 'follow answer'] ) - assert.equal(branch.userMessage.branch?.total, 2) - assert.equal(manager.getHistoryForAgent(conversation.id, undefined, branch.userMessage.parentId).length, 0) - assert.deepEqual(manager.getConversationTokenUsage(conversation.id), { - promptTokens: 10, - completionTokens: 11, - totalTokens: 21, - }) - - manager.switchMessageBranch(conversation.id, user1.id) + assert.equal(inserted.parentId, userMsg.id) + // followUp's parent should be updated to the inserted message + assert.equal(messages[2]?.parentId, inserted.id) + manager.close() + } finally { + cleanup(dir) + } + }) + + it('insertMessageAfter appends to the end when no child exists', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + const conv = manager.createConversation('s1', 'Test', 'general_cn') + const userMsg = manager.addMessage(conv.id, 'user', 'question') + + const inserted = manager.insertMessageAfter(conv.id, userMsg.id, 'assistant', 'new answer') + + const messages = manager.getMessages(conv.id) + assert.equal(messages.length, 2) assert.deepEqual( - manager.getMessages(conversation.id).map((message) => message.content), - ['original question', 'original answer', 'follow up', 'follow answer'] + messages.map((m) => m.content), + ['question', 'new answer'] ) - assert.equal(manager.getConversation(conversation.id)?.activeMessageId, oldLeaf.id) - assert.deepEqual(manager.getConversationTokenUsage(conversation.id), { - promptTokens: 5, - completionTokens: 7, - totalTokens: 12, - }) + assert.equal(inserted.parentId, userMsg.id) manager.close() } finally { cleanup(dir) } }) +}) - it('keeps summaries on inactive branches when adding a new summary', () => { +describe('AIConversationManager deleteMessagesFrom', () => { + it('deletes the target message and all subsequent messages', () => { const dir = createTempDir() try { const manager = createManager(dir) - const conversation = manager.createConversation('session-1', 'Summary branch', 'general_cn') - manager.addMessage(conversation.id, 'user', 'intro') - manager.addMessage(conversation.id, 'assistant', 'answer') - manager.addSummaryMessage(conversation.id, 'original summary', { - bufferBoundaryTimestamp: 2, - compressedMessageCount: 2, - }) - const followUp = manager.addMessage(conversation.id, 'user', 'follow up') - manager.addMessage(conversation.id, 'assistant', 'follow answer') - - manager.createMessageBranch(followUp.id, 'edited follow up', 'edited follow answer') - manager.addSummaryMessage(conversation.id, 'edited summary', { - bufferBoundaryTimestamp: 4, - compressedMessageCount: 2, - }) - manager.switchMessageBranch(conversation.id, followUp.id) + const conv = manager.createConversation('s1', 'Test', 'general_cn') + const user1 = manager.addMessage(conv.id, 'user', 'q1') + manager.addMessage(conv.id, 'assistant', 'a1') + manager.addMessage(conv.id, 'user', 'q2') + manager.addMessage(conv.id, 'assistant', 'a2') + manager.deleteMessagesFrom(conv.id, user1.id) + + const messages = manager.getMessages(conv.id) + assert.equal(messages.length, 0) + assert.equal(manager.getConversation(conv.id)?.activeMessageId, null) + manager.close() + } finally { + cleanup(dir) + } + }) + + it('deletes from a mid-chain message, preserving earlier ones', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + const conv = manager.createConversation('s1', 'Test', 'general_cn') + const user1 = manager.addMessage(conv.id, 'user', 'q1') + const ai1 = manager.addMessage(conv.id, 'assistant', 'a1') + manager.addMessage(conv.id, 'user', 'q2') + manager.addMessage(conv.id, 'assistant', 'a2') + + manager.deleteMessagesFrom(conv.id, ai1.id) + + const messages = manager.getMessages(conv.id) + assert.equal(messages.length, 1) + assert.equal(messages[0]?.content, 'q1') + assert.equal(manager.getConversation(conv.id)?.activeMessageId, user1.id) + manager.close() + } finally { + cleanup(dir) + } + }) +}) + +describe('AIConversationManager forkConversation', () => { + it('creates a new conversation with copied messages up to the specified point', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + const conv = manager.createConversation('s1', 'Original', 'general_cn') + manager.addMessage(conv.id, 'user', 'q1') + manager.addMessage(conv.id, 'assistant', 'a1') + const q2 = manager.addMessage(conv.id, 'user', 'q2') + manager.addMessage(conv.id, 'assistant', 'a2') + + const forked = manager.forkConversation(conv.id, q2.id, 'Forked') + + assert.notEqual(forked.id, conv.id) + assert.equal(forked.title, 'Forked') + assert.equal(forked.sessionId, 's1') + + const forkedMessages = manager.getMessages(forked.id) + assert.equal(forkedMessages.length, 2) assert.deepEqual( - manager.getMessages(conversation.id).map((message) => message.content), - ['intro', 'answer', 'original summary', 'follow up', 'follow answer'] + forkedMessages.map((m) => m.content), + ['q1', 'a1'] ) + + // Original conversation should be untouched + const originalMessages = manager.getMessages(conv.id) + assert.equal(originalMessages.length, 4) + manager.close() + } finally { + cleanup(dir) + } + }) + + it('uses default fork title when none provided', () => { + const dir = createTempDir() + try { + const manager = createManager(dir) + const conv = manager.createConversation('s1', 'MyChat', 'general_cn') + manager.addMessage(conv.id, 'user', 'q1') + const a1 = manager.addMessage(conv.id, 'assistant', 'a1') + + const forked = manager.forkConversation(conv.id, a1.id) + + assert.equal(forked.title, 'MyChat (fork)') + const forkedMessages = manager.getMessages(forked.id) + assert.equal(forkedMessages.length, 1) + assert.equal(forkedMessages[0]?.content, 'q1') manager.close() } finally { cleanup(dir) diff --git a/packages/node-runtime/src/ai/conversations.ts b/packages/node-runtime/src/ai/conversations.ts index 3cb5d6ea..0fef4349 100644 --- a/packages/node-runtime/src/ai/conversations.ts +++ b/packages/node-runtime/src/ai/conversations.ts @@ -56,14 +56,6 @@ export interface AIMessage { content: string timestamp: number parentId?: string | null - siblingGroupId?: string - branchIndex?: number - branch?: { - index: number - total: number - prevMessageId: string | null - nextMessageId: string | null - } dataKeywords?: string[] dataMessageCount?: number contentBlocks?: ContentBlock[] @@ -85,11 +77,6 @@ interface AIMessageRow { tokenUsage: string | null } -export interface MessageBranchResult { - userMessage: AIMessage - assistantMessage: AIMessage -} - export interface ConversationManagerLogger { warn(category: string, message: string, extra?: Record): void } @@ -202,7 +189,7 @@ export class AIConversationManager { db.exec('ALTER TABLE ai_conversation ADD COLUMN active_message_id TEXT') } - if (needsMessageTreeBackfill || needsConversationBackfill) { + if (needsMessageTreeBackfill || needsConversationBackfill || this.hasUnbackfilledMessageTree(db)) { this.backfillMessageTree(db) } @@ -215,6 +202,29 @@ export class AIConversationManager { } } + private hasUnbackfilledMessageTree(db: Database.Database): boolean { + const row = db + .prepare( + `SELECT 1 + FROM ai_conversation c + WHERE EXISTS ( + SELECT 1 FROM ai_message m + WHERE m.conversation_id = c.id + ) + AND ( + c.active_message_id IS NULL + OR EXISTS ( + SELECT 1 FROM ai_message m + WHERE m.conversation_id = c.id + AND (m.sibling_group_id IS NULL OR m.branch_index IS NULL) + ) + ) + LIMIT 1` + ) + .get() + return !!row + } + private backfillMessageTree(db: Database.Database): void { const conversations = db.prepare('SELECT id FROM ai_conversation').all() as Array<{ id: string }> const updateMessage = db.prepare( @@ -254,7 +264,7 @@ export class AIConversationManager { return `${prefix}_${Date.now()}_${Math.random().toString(36).slice(2, 8)}` } - private parseMessageRow(row: AIMessageRow, branch?: AIMessage['branch']): AIMessage { + private parseMessageRow(row: AIMessageRow): AIMessage { return { id: row.id, conversationId: row.conversationId, @@ -262,9 +272,6 @@ export class AIConversationManager { content: row.content, timestamp: row.timestamp, parentId: row.parentId ?? null, - siblingGroupId: row.siblingGroupId ?? row.id, - branchIndex: row.branchIndex ?? 0, - branch, dataKeywords: row.dataKeywords ? JSON.parse(row.dataKeywords) : undefined, dataMessageCount: row.dataMessageCount ?? undefined, contentBlocks: row.contentBlocks ? JSON.parse(row.contentBlocks) : undefined, @@ -306,77 +313,38 @@ export class AIConversationManager { return null } + private getAllMessageRows(conversationId: string): AIMessageRow[] { + return this.getDb() + .prepare( + `SELECT id, conversation_id as conversationId, role, content, timestamp, + parent_id as parentId, sibling_group_id as siblingGroupId, branch_index as branchIndex, + data_keywords as dataKeywords, data_message_count as dataMessageCount, + content_blocks as contentBlocks, token_usage as tokenUsage + FROM ai_message WHERE conversation_id = ? ORDER BY timestamp ASC, id ASC` + ) + .all(conversationId) as AIMessageRow[] + } + private getActivePathRows(conversationId: string, leafMessageId?: string | null): AIMessageRow[] { - const db = this.getDb() if (leafMessageId === null) return [] + + const allRows = this.getAllMessageRows(conversationId) + if (allRows.length === 0) return [] + + const rowMap = new Map(allRows.map((row) => [row.id, row])) let currentId = leafMessageId ?? this.getActiveMessageId(conversationId) - const rows: AIMessageRow[] = [] + const path: AIMessageRow[] = [] const seen = new Set() while (currentId && !seen.has(currentId)) { seen.add(currentId) - const row = this.getMessageRow(currentId) - if (!row || row.conversationId !== conversationId) break - rows.push(row) + const row = rowMap.get(currentId) + if (!row) break + path.push(row) currentId = row.parentId } - if (rows.length === 0) { - return db - .prepare( - `SELECT id, conversation_id as conversationId, role, content, timestamp, - parent_id as parentId, sibling_group_id as siblingGroupId, branch_index as branchIndex, - data_keywords as dataKeywords, data_message_count as dataMessageCount, - content_blocks as contentBlocks, token_usage as tokenUsage - FROM ai_message WHERE conversation_id = ? ORDER BY timestamp ASC, id ASC` - ) - .all(conversationId) as AIMessageRow[] - } - - return rows.reverse() - } - - private getBranchMeta(row: AIMessageRow): AIMessage['branch'] | undefined { - const groupId = row.siblingGroupId ?? row.id - const siblings = this.getDb() - .prepare( - `SELECT id FROM ai_message - WHERE conversation_id = ? AND sibling_group_id = ? - ORDER BY branch_index ASC, timestamp ASC, id ASC` - ) - .all(row.conversationId, groupId) as Array<{ id: string }> - - if (siblings.length <= 1) return undefined - const index = siblings.findIndex((item) => item.id === row.id) - const safeIndex = index >= 0 ? index : 0 - return { - index: safeIndex, - total: siblings.length, - prevMessageId: safeIndex > 0 ? siblings[safeIndex - 1]!.id : null, - nextMessageId: safeIndex < siblings.length - 1 ? siblings[safeIndex + 1]!.id : null, - } - } - - private getDeepestLeafId(messageId: string): string { - const db = this.getDb() - let currentId = messageId - const seen = new Set() - - while (!seen.has(currentId)) { - seen.add(currentId) - const child = db - .prepare( - `SELECT id FROM ai_message - WHERE parent_id = ? - ORDER BY branch_index DESC, timestamp DESC, id DESC - LIMIT 1` - ) - .get(currentId) as { id: string } | undefined - if (!child?.id) break - currentId = child.id - } - - return currentId + return path.length > 0 ? path.reverse() : allRows } // ==================== 生命周期 ==================== @@ -585,8 +553,6 @@ export class AIConversationManager { content, timestamp: now, parentId, - siblingGroupId, - branchIndex, dataKeywords, dataMessageCount, contentBlocks, @@ -595,8 +561,7 @@ export class AIConversationManager { } getMessages(conversationId: string): AIMessage[] { - const rows = this.getActivePathRows(conversationId) - return rows.map((row) => this.parseMessageRow(row, this.getBranchMeta(row))) + return this.getActivePathRows(conversationId).map((row) => this.parseMessageRow(row)) } deleteMessage(messageId: string): boolean { @@ -605,115 +570,202 @@ export class AIConversationManager { return result.changes > 0 } - createMessageBranch( - originalUserMessageId: string, - newUserContent: string, - assistantContent: string, - contentBlocks?: ContentBlock[], - tokenUsage?: TokenUsageData - ): MessageBranchResult { + deleteMessagesFrom(conversationId: string, messageId: string): void { const db = this.getDb() - const original = this.getMessageRow(originalUserMessageId) - if (!original || original.role !== 'user') { - throw new Error('Original user message not found') + const target = this.getMessageRow(messageId) + if (!target || target.conversationId !== conversationId) { + throw new Error('Message not found in conversation') } + const activePath = this.getActivePathRows(conversationId) + const targetIndex = activePath.findIndex((row) => row.id === messageId) + if (targetIndex < 0) { + throw new Error('Message not on active path') + } + + const idsToDelete = activePath.slice(targetIndex).map((row) => row.id) + const placeholders = idsToDelete.map(() => '?').join(', ') + db.prepare(`DELETE FROM ai_message WHERE id IN (${placeholders})`).run(...idsToDelete) + + const newLeafId = targetIndex > 0 ? activePath[targetIndex - 1]!.id : null const now = Math.floor(Date.now() / 1000) - const userMessageId = this.generateId('msg') - const assistantMessageId = this.generateId('msg') - const siblingGroupId = original.siblingGroupId ?? original.id - const branchRow = db - .prepare( - `SELECT COALESCE(MAX(branch_index), -1) + 1 as branchIndex - FROM ai_message WHERE conversation_id = ? AND sibling_group_id = ?` - ) - .get(original.conversationId, siblingGroupId) as { branchIndex: number } + db.prepare('UPDATE ai_conversation SET active_message_id = ?, updated_at = ? WHERE id = ?').run( + newLeafId, + now, + conversationId + ) + } - const pendingDebug = this.pendingDebugContextMap.get(original.conversationId) - if (pendingDebug) { - this.pendingDebugContextMap.delete(original.conversationId) + forkConversation(sourceConversationId: string, upToMessageId: string, title?: string): AIConversation { + const db = this.getDb() + const source = this.getConversation(sourceConversationId) + if (!source) { + throw new Error('Source conversation not found') } - const tx = db.transaction(() => { - db.prepare( - `INSERT INTO ai_message ( - id, conversation_id, role, content, timestamp, data_keywords, data_message_count, - content_blocks, token_usage, debug_context, parent_id, sibling_group_id, branch_index - ) - VALUES (?, ?, 'user', ?, ?, NULL, NULL, NULL, NULL, NULL, ?, ?, ?)` - ).run( - userMessageId, - original.conversationId, - newUserContent, - now, - original.parentId, - siblingGroupId, - branchRow.branchIndex - ) + const activePath = this.getActivePathRows(sourceConversationId) + const cutIndex = activePath.findIndex((row) => row.id === upToMessageId) + if (cutIndex < 0) { + throw new Error('Message not on active path') + } + + const messagesToCopy = activePath.slice(0, cutIndex) + const now = Math.floor(Date.now() / 1000) + const newConvId = this.generateId('conv') + const forkTitle = title || `${source.title || 'Untitled'} (fork)` + + db.prepare( + `INSERT INTO ai_conversation (id, session_id, title, assistant_id, active_message_id, created_at, updated_at) + VALUES (?, ?, ?, ?, NULL, ?, ?)` + ).run(newConvId, source.sessionId, forkTitle, source.assistantId, now, now) + + const idMap = new Map() + let lastNewId: string | null = null + + for (const row of messagesToCopy) { + const newMsgId = this.generateId('msg') + idMap.set(row.id, newMsgId) + const newParentId = row.parentId ? (idMap.get(row.parentId) ?? null) : null db.prepare( `INSERT INTO ai_message ( id, conversation_id, role, content, timestamp, data_keywords, data_message_count, content_blocks, token_usage, debug_context, parent_id, sibling_group_id, branch_index ) - VALUES (?, ?, 'assistant', ?, ?, NULL, NULL, ?, ?, ?, ?, ?, 0)` + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, 0)` ).run( - assistantMessageId, - original.conversationId, - assistantContent, - now, - contentBlocks ? JSON.stringify(contentBlocks) : null, - tokenUsage ? JSON.stringify(tokenUsage) : null, - pendingDebug ?? null, - userMessageId, - assistantMessageId + newMsgId, + newConvId, + row.role, + row.content, + row.timestamp, + row.dataKeywords, + row.dataMessageCount, + row.contentBlocks, + row.tokenUsage, + newParentId, + newMsgId ) + lastNewId = newMsgId + } - db.prepare('UPDATE ai_conversation SET active_message_id = ?, updated_at = ? WHERE id = ?').run( - assistantMessageId, - now, - original.conversationId - ) - }) + if (lastNewId) { + db.prepare('UPDATE ai_conversation SET active_message_id = ? WHERE id = ?').run(lastNewId, newConvId) + } - tx() - - const userRow = this.getMessageRow(userMessageId)! - const assistantRow = this.getMessageRow(assistantMessageId)! return { - userMessage: this.parseMessageRow(userRow, this.getBranchMeta(userRow)), - assistantMessage: this.parseMessageRow(assistantRow, this.getBranchMeta(assistantRow)), + id: newConvId, + sessionId: source.sessionId, + title: forkTitle, + assistantId: source.assistantId, + activeMessageId: lastNewId, + createdAt: now, + updatedAt: now, } } - switchMessageBranch(conversationId: string, messageId: string): AIMessage[] { + updateMessageContent(messageId: string, newContent: string): void { + const db = this.getDb() + const result = db.prepare('UPDATE ai_message SET content = ? WHERE id = ?').run(newContent, messageId) + if (result.changes === 0) throw new Error('Message not found') + } + + deleteAndRelinkMessage(conversationId: string, messageId: string): void { + const db = this.getDb() const target = this.getMessageRow(messageId) if (!target || target.conversationId !== conversationId) { - throw new Error('Message branch not found') + throw new Error('Message not found in conversation') } - const leafId = this.getDeepestLeafId(messageId) - const now = Math.floor(Date.now() / 1000) - this.getDb() - .prepare('UPDATE ai_conversation SET active_message_id = ?, updated_at = ? WHERE id = ?') - .run(leafId, now, conversationId) + db.prepare('UPDATE ai_message SET parent_id = ? WHERE parent_id = ? AND conversation_id = ?').run( + target.parentId, + messageId, + conversationId + ) - return this.getMessages(conversationId) + const conv = this.getConversation(conversationId) + if (conv?.activeMessageId === messageId) { + const now = Math.floor(Date.now() / 1000) + db.prepare('UPDATE ai_conversation SET active_message_id = ?, updated_at = ? WHERE id = ?').run( + target.parentId, + now, + conversationId + ) + } + + db.prepare('DELETE FROM ai_message WHERE id = ?').run(messageId) + } + + insertMessageAfter( + conversationId: string, + afterMessageId: string, + role: AIMessageRole, + content: string, + contentBlocks?: ContentBlock[], + tokenUsage?: TokenUsageData + ): AIMessage { + const db = this.getDb() + const now = Math.floor(Date.now() / 1000) + const id = this.generateId('msg') + + const pendingDebug = role === 'assistant' ? this.pendingDebugContextMap.get(conversationId) : undefined + if (pendingDebug) { + this.pendingDebugContextMap.delete(conversationId) + } + + const childRow = db + .prepare('SELECT id FROM ai_message WHERE parent_id = ? AND conversation_id = ? LIMIT 1') + .get(afterMessageId, conversationId) as { id: string } | undefined + + db.prepare( + `INSERT INTO ai_message ( + id, conversation_id, role, content, timestamp, data_keywords, data_message_count, + content_blocks, token_usage, debug_context, parent_id, sibling_group_id, branch_index + ) + VALUES (?, ?, ?, ?, ?, NULL, NULL, ?, ?, ?, ?, ?, 0)` + ).run( + id, + conversationId, + role, + content, + now, + contentBlocks ? JSON.stringify(contentBlocks) : null, + tokenUsage ? JSON.stringify(tokenUsage) : null, + pendingDebug ?? null, + afterMessageId, + id + ) + + if (childRow) { + db.prepare('UPDATE ai_message SET parent_id = ? WHERE id = ?').run(id, childRow.id) + } + + db.prepare('UPDATE ai_conversation SET updated_at = ? WHERE id = ?').run(now, conversationId) + + return { + id, + conversationId, + role, + content, + timestamp: now, + parentId: afterMessageId, + contentBlocks, + tokenUsage, + } } getConversationTokenUsage(conversationId: string): TokenUsageData { - const messages = this.getMessages(conversationId) - return messages.reduce( - (total, message) => { - if (message.tokenUsage) { - total.promptTokens += message.tokenUsage.promptTokens - total.completionTokens += message.tokenUsage.completionTokens - total.totalTokens += message.tokenUsage.totalTokens - } - return total - }, - { promptTokens: 0, completionTokens: 0, totalTokens: 0 } - ) + const rows = this.getActivePathRows(conversationId) + const result: TokenUsageData = { promptTokens: 0, completionTokens: 0, totalTokens: 0 } + for (const row of rows) { + if (row.tokenUsage) { + const usage = JSON.parse(row.tokenUsage) as TokenUsageData + result.promptTokens += usage.promptTokens + result.completionTokens += usage.completionTokens + result.totalTokens += usage.totalTokens + } + } + return result } // ==================== Debug context ==================== diff --git a/packages/node-runtime/src/ai/index.ts b/packages/node-runtime/src/ai/index.ts index 1b20bb4d..2e199522 100644 --- a/packages/node-runtime/src/ai/index.ts +++ b/packages/node-runtime/src/ai/index.ts @@ -32,7 +32,6 @@ export type { ContentBlock, TokenUsageData, ConversationManagerLogger, - MessageBranchResult, } from './conversations' // Tokenizer diff --git a/packages/node-runtime/src/index.ts b/packages/node-runtime/src/index.ts index 2538d296..5fc79f90 100644 --- a/packages/node-runtime/src/index.ts +++ b/packages/node-runtime/src/index.ts @@ -79,7 +79,7 @@ export { } from './nlp' // AI 助手/技能解析器 + 对话管理 -export type { AssistantConfig, AssistantSummary, SkillDef, SkillSummary, MessageBranchResult } from './ai' +export type { AssistantConfig, AssistantSummary, SkillDef, SkillSummary } from './ai' export { parseAssistantFile, serializeAssistant, parseSkillFile, extractSkillId } from './ai' export { AIConversationManager } from './ai' export { countTokens, countMessagesTokens } from './ai' diff --git a/src/components/AIChat/ChatExplorer.vue b/src/components/AIChat/ChatExplorer.vue index 17feb419..62d3a39b 100644 --- a/src/components/AIChat/ChatExplorer.vue +++ b/src/components/AIChat/ChatExplorer.vue @@ -55,7 +55,6 @@ const { selectedAssistantId, sendMessage, editMessageAndRegenerate, - switchMessageBranch, loadConversation, startNewConversation, loadMoreSourceMessages, @@ -246,8 +245,10 @@ async function handleSend(payload: { content: string; mentionedMembers: Mentione conversationListRef.value?.refresh() } -async function handleEditMessage(payload: { messageId: string; content: string }) { - const result = await editMessageAndRegenerate(payload.messageId, payload.content) +async function handleEditMessage(payload: { messageId: string; content: string; overwriteSubsequent?: boolean }) { + const result = await editMessageAndRegenerate(payload.messageId, payload.content, { + overwriteSubsequent: payload.overwriteSubsequent, + }) if (!result.success) { if (result.reason === 'busy') { showRunningTaskToast() @@ -258,16 +259,6 @@ async function handleEditMessage(payload: { messageId: string; content: string } conversationListRef.value?.refresh() } -async function handleSwitchMessageBranch(messageId: string | null) { - if (!messageId) return - const ok = await switchMessageBranch(messageId) - if (!ok) { - showLockedActionToast() - return - } - scrollToBottom(true) -} - // 切换数据源面板 function toggleSourcePanel() { isSourcePanelCollapsed.value = !isSourcePanelCollapsed.value @@ -442,11 +433,8 @@ watch( :timestamp="pair.user.timestamp" :is-streaming="pair.user.isStreaming" :content-blocks="pair.user.contentBlocks" - :branch="pair.user.branch" :editable="!isAIThinking" @edit="handleEditMessage" - @branch-prev="handleSwitchMessageBranch" - @branch-next="handleSwitchMessageBranch" /> () const emit = defineEmits<{ - edit: [payload: { messageId: string; content: string }] - branchPrev: [messageId: string | null] - branchNext: [messageId: string | null] + edit: [payload: { messageId: string; content: string; overwriteSubsequent?: boolean }] }>() // 格式化时间 @@ -49,7 +41,7 @@ const isEditing = ref(false) const editContent = ref(props.content) const editTextareaRef = ref(null) const canEdit = computed(() => isUser.value && props.editable && !props.isStreaming && !!props.messageId) -const hasBranchSwitcher = computed(() => isUser.value && !!props.branch && props.branch.total > 1) +const overwriteSubsequent = ref(false) // 创建 markdown-it 实例 const md = new MarkdownIt({ @@ -103,17 +95,29 @@ watch( } ) +function syncEditTextareaHeight() { + const el = editTextareaRef.value + if (!el) return + el.style.height = 'auto' + const maxHeight = 384 + const nextHeight = Math.min(el.scrollHeight, maxHeight) + el.style.height = `${nextHeight}px` + el.style.overflowY = el.scrollHeight > maxHeight ? 'auto' : 'hidden' +} + async function startEditing() { if (!canEdit.value) return editContent.value = props.content isEditing.value = true await nextTick() + syncEditTextareaHeight() editTextareaRef.value?.focus() } function cancelEditing() { isEditing.value = false editContent.value = props.content + overwriteSubsequent.value = false } function submitEditing() { @@ -124,7 +128,8 @@ function submitEditing() { return } isEditing.value = false - emit('edit', { messageId: props.messageId, content }) + emit('edit', { messageId: props.messageId, content, overwriteSubsequent: overwriteSubsequent.value }) + overwriteSubsequent.value = false } // 过滤无内容的文本/思考块，避免显示空气泡 @@ -352,9 +357,12 @@ async function handleCopyMarkdown() {