From cb7ba1f3ccf425e835478bdb3c4347131098dcf4 Mon Sep 17 00:00:00 2001 From: digua Date: Wed, 26 Aug 2026 22:14:04 +0800 Subject: [PATCH] fix(ai): refresh memory source parent after compression Capture the conversation active message when the agent turn completes so an automatically persisted compression summary becomes the expected parent for the pending user message. Add an integration regression covering summary persistence, memory creation, message-pair persistence, and source linking. --- .../src/routes/web/ai-agent-stream.ts | 6 +- .../src/routes/web/ai-memories.test.ts | 76 ++++++++++++++++++- .../web/memory-provenance-coordinator.test.ts | 12 +-- .../web/memory-provenance-coordinator.ts | 7 +- 4 files changed, 85 insertions(+), 16 deletions(-) diff --git a/packages/http-routes/src/routes/web/ai-agent-stream.ts b/packages/http-routes/src/routes/web/ai-agent-stream.ts index 4c57e32e..8fa21127 100644 --- a/packages/http-routes/src/routes/web/ai-agent-stream.ts +++ b/packages/http-routes/src/routes/web/ai-agent-stream.ts @@ -21,8 +21,7 @@ export function registerAiAgentStreamRoutes( const requestId = `req_${randomUUID()}` const abortController = new AbortController() activeAgentAborts.set(requestId, abortController) - const expectedParentMessageId = ctx.aiChatManager?.getAIChat(request.body.aiChatId)?.activeMessageId ?? null - memoryProvenanceCoordinator?.begin(requestId, request.body.aiChatId, expectedParentMessageId) + memoryProvenanceCoordinator?.begin(requestId, request.body.aiChatId) reply.raw.writeHead(200, { 'Content-Type': 'text/event-stream', @@ -71,7 +70,8 @@ export function registerAiAgentStreamRoutes( const msg = error instanceof Error ? error.message : String(error) safeSendSSE('error', { type: 'error', error: { name: 'ServerError', message: msg } }) } finally { - memoryProvenanceCoordinator?.complete(requestId) + const expectedParentMessageId = ctx.aiChatManager?.getAIChat(request.body.aiChatId)?.activeMessageId ?? null + memoryProvenanceCoordinator?.complete(requestId, expectedParentMessageId) releaseInteractiveWork() if (!emittedDone) safeSendSSE('done', { type: 'done', isFinished: true }) activeAgentAborts.delete(requestId) diff --git a/packages/http-routes/src/routes/web/ai-memories.test.ts b/packages/http-routes/src/routes/web/ai-memories.test.ts index 8887e998..8ba246b8 100644 --- a/packages/http-routes/src/routes/web/ai-memories.test.ts +++ b/packages/http-routes/src/routes/web/ai-memories.test.ts @@ -5,6 +5,7 @@ import { join } from 'node:path' import { describe, it } from 'node:test' import Fastify from 'fastify' import { AIChatManager, AIMemoryService } from '@openchatlab/node-runtime' +import { registerAiAgentStreamRoutes } from './ai-agent-stream' import { registerAiMemoryRoutes } from './ai-memories' import { MemoryProvenanceCoordinator } from './memory-provenance-coordinator' @@ -220,9 +221,9 @@ describe('AI memory routes', () => { sourceAIChatId: chat.id, }) const provenanceToken = 'test-turn' - fixture.memoryProvenanceCoordinator.begin(provenanceToken, chat.id, userMessage.parentId ?? null) + fixture.memoryProvenanceCoordinator.begin(provenanceToken, chat.id) fixture.memoryProvenanceCoordinator.record(provenanceToken, memory.id) - fixture.memoryProvenanceCoordinator.complete(provenanceToken) + fixture.memoryProvenanceCoordinator.complete(provenanceToken, userMessage.parentId ?? null) const linked = await fixture.app.inject({ method: 'POST', @@ -281,9 +282,9 @@ describe('AI memory routes', () => { sourceAIChatId: chat.id, }) const provenanceToken = 'bound-turn' - fixture.memoryProvenanceCoordinator.begin(provenanceToken, chat.id, previousAssistant.id) + fixture.memoryProvenanceCoordinator.begin(provenanceToken, chat.id) fixture.memoryProvenanceCoordinator.record(provenanceToken, memory.id) - fixture.memoryProvenanceCoordinator.complete(provenanceToken) + fixture.memoryProvenanceCoordinator.complete(provenanceToken, previousAssistant.id) const currentUser = fixture.aiChatManager.addMessage(chat.id, 'user', 'current') const currentAssistant = fixture.aiChatManager.addMessage(chat.id, 'assistant', 'current answer') @@ -346,6 +347,73 @@ describe('AI memory routes', () => { } }) + it('links a memory to the current turn after automatic compression changes the active parent', async () => { + const fixture = createFixture() + try { + const chat = fixture.aiChatManager.createGlobalAIChat('source', 'general_cn') + fixture.aiChatManager.addMessage(chat.id, 'user', 'previous') + fixture.aiChatManager.addMessage(chat.id, 'assistant', 'previous answer') + let memoryId = '' + let summaryMessageId = '' + + registerAiAgentStreamRoutes( + fixture.app, + { + aiChatManager: fixture.aiChatManager, + runAgentStream: async (_request, onEvent) => { + const summary = fixture.aiChatManager.addSummaryMessage(chat.id, 'compressed history', { + bufferBoundaryTimestamp: Date.now(), + compressedMessageCount: 2, + }) + summaryMessageId = summary.id + const memory = fixture.service.create({ + scopeType: 'global', + scopeId: null, + content: 'changed after compression', + sourceType: 'ai', + sourceAIChatId: chat.id, + }) + memoryId = memory.id + onEvent({ type: 'memory_change', memoryId }) + onEvent({ type: 'done', isFinished: true }) + }, + }, + fixture.memoryProvenanceCoordinator + ) + + const streamResponse = await fixture.app.inject({ + method: 'POST', + url: '/_web/ai/agent/stream', + payload: { aiChatId: chat.id }, + }) + assert.equal(streamResponse.statusCode, 200) + const provenanceToken = String(streamResponse.headers['x-request-id']) + + const { userMessage, assistantMessage } = fixture.aiChatManager.addMessagePair( + chat.id, + { content: 'current question' }, + { content: 'current answer' } + ) + assert.equal(userMessage.parentId, summaryMessageId) + + const linked = await fixture.app.inject({ + method: 'POST', + url: '/_web/ai/memories/link-sources', + payload: { + provenanceToken, + aiChatId: chat.id, + userMessageId: userMessage.id, + assistantMessageId: assistantMessage.id, + memoryIds: [memoryId], + }, + }) + assert.equal(linked.statusCode, 200) + assert.equal(fixture.service.get(memoryId)?.sourceMessageId, assistantMessage.id) + } finally { + await fixture.close() + } + }) + it('rejects invalid scope and content instead of changing stored memories', async () => { const fixture = createFixture() try { diff --git a/packages/http-routes/src/routes/web/memory-provenance-coordinator.test.ts b/packages/http-routes/src/routes/web/memory-provenance-coordinator.test.ts index 119bfd0c..04cc5ccb 100644 --- a/packages/http-routes/src/routes/web/memory-provenance-coordinator.test.ts +++ b/packages/http-routes/src/routes/web/memory-provenance-coordinator.test.ts @@ -8,10 +8,10 @@ describe('MemoryProvenanceCoordinator', () => { t.mock.method(Date, 'now', () => now) const coordinator = new MemoryProvenanceCoordinator() - coordinator.begin('turn-token', 'global-chat', 'previous-message') + coordinator.begin('turn-token', 'global-chat') now = 24 * 60 * 60 * 1000 coordinator.record('turn-token', 'memory-1') - coordinator.complete('turn-token') + coordinator.complete('turn-token', 'previous-message') assert.deepEqual(coordinator.validate('turn-token', 'global-chat', ['memory-1']), { expectedParentMessageId: 'previous-message', @@ -21,8 +21,8 @@ describe('MemoryProvenanceCoordinator', () => { it('removes completed turns that did not change any memories', () => { const coordinator = new MemoryProvenanceCoordinator() - coordinator.begin('empty-turn', 'global-chat', null) - coordinator.complete('empty-turn') + coordinator.begin('empty-turn', 'global-chat') + coordinator.complete('empty-turn', null) assert.throws(() => coordinator.validate('empty-turn', 'global-chat', []), /Memory provenance token is invalid/) }) @@ -30,9 +30,9 @@ describe('MemoryProvenanceCoordinator', () => { it('keeps completed changes until they are consumed', () => { const coordinator = new MemoryProvenanceCoordinator() - coordinator.begin('changed-turn', 'global-chat', null) + coordinator.begin('changed-turn', 'global-chat') coordinator.record('changed-turn', 'memory-1') - coordinator.complete('changed-turn') + coordinator.complete('changed-turn', null) assert.deepEqual(coordinator.validate('changed-turn', 'global-chat', ['memory-1']), { expectedParentMessageId: null, diff --git a/packages/http-routes/src/routes/web/memory-provenance-coordinator.ts b/packages/http-routes/src/routes/web/memory-provenance-coordinator.ts index 0658cb44..716ebfb9 100644 --- a/packages/http-routes/src/routes/web/memory-provenance-coordinator.ts +++ b/packages/http-routes/src/routes/web/memory-provenance-coordinator.ts @@ -8,10 +8,10 @@ interface PendingMemoryProvenance { export class MemoryProvenanceCoordinator { private readonly pending = new Map() - begin(token: string, aiChatId: string, expectedParentMessageId: string | null): void { + begin(token: string, aiChatId: string): void { this.pending.set(token, { aiChatId, - expectedParentMessageId, + expectedParentMessageId: null, memoryIds: new Set(), completed: false, }) @@ -23,13 +23,14 @@ export class MemoryProvenanceCoordinator { changeSet.memoryIds.add(memoryId) } - complete(token: string): void { + complete(token: string, expectedParentMessageId: string | null): void { const changeSet = this.pending.get(token) if (!changeSet) return if (changeSet.memoryIds.size === 0) { this.pending.delete(token) return } + changeSet.expectedParentMessageId = expectedParentMessageId changeSet.completed = true }