diff --git a/apps/cli/src/ai/agent-stream-runner.ts b/apps/cli/src/ai/agent-stream-runner.ts index eb30eed6..5fad67e0 100644 --- a/apps/cli/src/ai/agent-stream-runner.ts +++ b/apps/cli/src/ai/agent-stream-runner.ts @@ -160,6 +160,7 @@ export function createCliRunAgentStream( maxToolResultTokens: resolveCrossChatToolResultTokenBudget(contextWindow), memoryService: options.memoryService, aiChatId, + reportMemoryChange: (memoryId) => onEvent({ type: 'memory_change', memoryId }), }) await runCrossChatAgent({ userMessage, diff --git a/apps/cli/src/ai/cross-chat-tool-adapter.ts b/apps/cli/src/ai/cross-chat-tool-adapter.ts index 1529e385..58644895 100644 --- a/apps/cli/src/ai/cross-chat-tool-adapter.ts +++ b/apps/cli/src/ai/cross-chat-tool-adapter.ts @@ -23,6 +23,7 @@ export function createCliCrossChatTools(options: { maxToolResultTokens: number memoryService: AIMemoryService aiChatId: string + reportMemoryChange?: (memoryId: string) => void }): AgentTool[] { const context: Omit = { locale: options.locale, @@ -30,6 +31,7 @@ export function createCliCrossChatTools(options: { analysisService: options.analysisService, memoryService: options.memoryService, aiChatId: options.aiChatId, + reportMemoryChange: options.reportMemoryChange, maxToolResultTokens: options.maxToolResultTokens, countTokens, preprocessMessagesBySession: (sessionId, messages) => diff --git a/apps/desktop/main/ai/agent-stream-runner.ts b/apps/desktop/main/ai/agent-stream-runner.ts index f908aecf..7cf11424 100644 --- a/apps/desktop/main/ai/agent-stream-runner.ts +++ b/apps/desktop/main/ai/agent-stream-runner.ts @@ -155,6 +155,7 @@ export function createElectronRunAgentStream( maxToolResultTokens: resolveCrossChatToolResultTokenBudget(resolvedContextWindow), memoryService: crossChatOptions.memoryService, aiChatId, + reportMemoryChange: (memoryId) => onEvent({ type: 'memory_change', memoryId }), }) await runCrossChatAgent({ userMessage, diff --git a/apps/desktop/main/ai/cross-chat-tool-adapter.ts b/apps/desktop/main/ai/cross-chat-tool-adapter.ts index 97f76114..6395bc8b 100644 --- a/apps/desktop/main/ai/cross-chat-tool-adapter.ts +++ b/apps/desktop/main/ai/cross-chat-tool-adapter.ts @@ -23,6 +23,7 @@ export function createElectronCrossChatTools(options: { maxToolResultTokens: number memoryService: AIMemoryService aiChatId: string + reportMemoryChange?: (memoryId: string) => void }): AgentTool[] { const context: Omit = { locale: options.locale, @@ -30,6 +31,7 @@ export function createElectronCrossChatTools(options: { analysisService: options.analysisService, memoryService: options.memoryService, aiChatId: options.aiChatId, + reportMemoryChange: options.reportMemoryChange, maxToolResultTokens: options.maxToolResultTokens, countTokens, preprocessMessagesBySession: (sessionId, messages) => 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 d80dc66f..4c57e32e 100644 --- a/packages/http-routes/src/routes/web/ai-agent-stream.ts +++ b/packages/http-routes/src/routes/web/ai-agent-stream.ts @@ -1,20 +1,28 @@ import type { FastifyInstance } from 'fastify' +import { randomUUID } from 'node:crypto' import type { AgentStreamRequest, AiRouteContext } from '../../context/ai' import { chatTopicWorkCoordinator } from '@openchatlab/node-runtime' +import type { MemoryProvenanceCoordinator } from './memory-provenance-coordinator' const activeAgentAborts = new Map() -type AiAgentStreamRouteContext = Pick +type AiAgentStreamRouteContext = Pick -export function registerAiAgentStreamRoutes(server: FastifyInstance, ctx: AiAgentStreamRouteContext): void { +export function registerAiAgentStreamRoutes( + server: FastifyInstance, + ctx: AiAgentStreamRouteContext, + memoryProvenanceCoordinator?: MemoryProvenanceCoordinator +): void { if (!ctx.runAgentStream) return const runAgentStream = ctx.runAgentStream server.post<{ Body: AgentStreamRequest }>('/_web/ai/agent/stream', async (request, reply) => { - const requestId = `req_${Date.now()}_${Math.random().toString(36).slice(2, 8)}` + 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) reply.raw.writeHead(200, { 'Content-Type': 'text/event-stream', @@ -50,6 +58,11 @@ export function registerAiAgentStreamRoutes(server: FastifyInstance, ctx: AiAgen request.body, (chunk) => { if (chunk.type === 'done') emittedDone = true + if (chunk.type === 'memory_change' && chunk.memoryId) { + memoryProvenanceCoordinator?.record(requestId, chunk.memoryId) + safeSendSSE(chunk.type, { ...chunk, provenanceToken: requestId }) + return + } safeSendSSE(chunk.type, chunk) }, abortController.signal @@ -58,6 +71,7 @@ export function registerAiAgentStreamRoutes(server: FastifyInstance, ctx: AiAgen const msg = error instanceof Error ? error.message : String(error) safeSendSSE('error', { type: 'error', error: { name: 'ServerError', message: msg } }) } finally { + memoryProvenanceCoordinator?.complete(requestId) 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 263e0d35..8887e998 100644 --- a/packages/http-routes/src/routes/web/ai-memories.test.ts +++ b/packages/http-routes/src/routes/web/ai-memories.test.ts @@ -4,23 +4,29 @@ import { tmpdir } from 'node:os' import { join } from 'node:path' import { describe, it } from 'node:test' import Fastify from 'fastify' -import { AIMemoryService } from '@openchatlab/node-runtime' +import { AIChatManager, AIMemoryService } from '@openchatlab/node-runtime' import { registerAiMemoryRoutes } from './ai-memories' +import { MemoryProvenanceCoordinator } from './memory-provenance-coordinator' const sqliteNativeBinding = process.env.CHATLAB_TEST_SQLITE_NATIVE_BINDING function createFixture() { const dir = mkdtempSync(join(tmpdir(), 'chatlab-ai-memory-routes-')) const service = new AIMemoryService(dir, { nativeBinding: sqliteNativeBinding }) + const aiChatManager = new AIChatManager(dir, { nativeBinding: sqliteNativeBinding }) + const memoryProvenanceCoordinator = new MemoryProvenanceCoordinator() const app = Fastify() - registerAiMemoryRoutes(app, { aiMemoryService: service }) + registerAiMemoryRoutes(app, { aiMemoryService: service, aiChatManager }, memoryProvenanceCoordinator) return { app, service, + aiChatManager, + memoryProvenanceCoordinator, async close() { await app.close() service.close() + aiChatManager.close() try { rmSync(dir, { recursive: true, force: true, maxRetries: 5, retryDelay: 50 }) } catch { @@ -199,6 +205,147 @@ describe('AI memory routes', () => { } }) + it('links only memories and persisted message roles owned by the requested global conversation', async () => { + const fixture = createFixture() + try { + const chat = fixture.aiChatManager.createGlobalAIChat('source', 'general_cn') + const otherChat = fixture.aiChatManager.createGlobalAIChat('other', 'general_cn') + const userMessage = fixture.aiChatManager.addMessage(chat.id, 'user', '请记住') + const assistantMessage = fixture.aiChatManager.addMessage(chat.id, 'assistant', '已经记住') + const memory = fixture.service.create({ + scopeType: 'global', + scopeId: null, + content: '先给结论', + sourceType: 'user', + sourceAIChatId: chat.id, + }) + const provenanceToken = 'test-turn' + fixture.memoryProvenanceCoordinator.begin(provenanceToken, chat.id, userMessage.parentId ?? null) + fixture.memoryProvenanceCoordinator.record(provenanceToken, memory.id) + fixture.memoryProvenanceCoordinator.complete(provenanceToken) + + 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: [memory.id], + }, + }) + assert.equal(linked.statusCode, 200) + assert.deepEqual(linked.json(), { linkedMemoryIds: [memory.id], skippedMemoryIds: [] }) + + const listed = await fixture.app.inject({ method: 'GET', url: '/_web/ai/memories' }) + assert.equal(listed.json()[0].sourceStatus, 'message') + assert.equal(listed.json()[0].sourceMessageId, userMessage.id) + + const forged = await fixture.app.inject({ + method: 'POST', + url: '/_web/ai/memories/link-sources', + payload: { + provenanceToken, + aiChatId: otherChat.id, + userMessageId: userMessage.id, + assistantMessageId: assistantMessage.id, + memoryIds: [memory.id], + }, + }) + assert.equal(forged.statusCode, 400) + assert.equal(fixture.service.get(memory.id)?.sourceMessageId, userMessage.id) + } finally { + await fixture.close() + } + }) + + it('binds source backfill to one completed agent turn and its exact memory change set', async () => { + const fixture = createFixture() + try { + const chat = fixture.aiChatManager.createGlobalAIChat('source', 'general_cn') + const previousUser = fixture.aiChatManager.addMessage(chat.id, 'user', 'previous') + const previousAssistant = fixture.aiChatManager.addMessage(chat.id, 'assistant', 'previous answer') + const memory = fixture.service.create({ + scopeType: 'global', + scopeId: null, + content: 'changed this turn', + sourceType: 'user', + sourceAIChatId: chat.id, + }) + const untrackedMemory = fixture.service.create({ + scopeType: 'global', + scopeId: null, + content: 'not changed this turn', + sourceType: 'user', + sourceAIChatId: chat.id, + }) + const provenanceToken = 'bound-turn' + fixture.memoryProvenanceCoordinator.begin(provenanceToken, chat.id, previousAssistant.id) + fixture.memoryProvenanceCoordinator.record(provenanceToken, memory.id) + fixture.memoryProvenanceCoordinator.complete(provenanceToken) + + const currentUser = fixture.aiChatManager.addMessage(chat.id, 'user', 'current') + const currentAssistant = fixture.aiChatManager.addMessage(chat.id, 'assistant', 'current answer') + const wrongChangeSet = await fixture.app.inject({ + method: 'POST', + url: '/_web/ai/memories/link-sources', + payload: { + provenanceToken, + aiChatId: chat.id, + userMessageId: currentUser.id, + assistantMessageId: currentAssistant.id, + memoryIds: [memory.id, untrackedMemory.id], + }, + }) + assert.equal(wrongChangeSet.statusCode, 400) + assert.equal(fixture.service.get(memory.id)?.sourceMessageId, null) + + const oldMessages = await fixture.app.inject({ + method: 'POST', + url: '/_web/ai/memories/link-sources', + payload: { + provenanceToken, + aiChatId: chat.id, + userMessageId: previousUser.id, + assistantMessageId: previousAssistant.id, + memoryIds: [memory.id], + }, + }) + assert.equal(oldMessages.statusCode, 400) + assert.equal(fixture.service.get(memory.id)?.sourceMessageId, null) + + const linked = await fixture.app.inject({ + method: 'POST', + url: '/_web/ai/memories/link-sources', + payload: { + provenanceToken, + aiChatId: chat.id, + userMessageId: currentUser.id, + assistantMessageId: currentAssistant.id, + memoryIds: [memory.id], + }, + }) + assert.equal(linked.statusCode, 200) + assert.equal(fixture.service.get(memory.id)?.sourceMessageId, currentUser.id) + + const replay = await fixture.app.inject({ + method: 'POST', + url: '/_web/ai/memories/link-sources', + payload: { + provenanceToken, + aiChatId: chat.id, + userMessageId: currentUser.id, + assistantMessageId: currentAssistant.id, + memoryIds: [memory.id], + }, + }) + assert.equal(replay.statusCode, 400) + } 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/ai-memories.ts b/packages/http-routes/src/routes/web/ai-memories.ts index 8c3050df..318d2fe5 100644 --- a/packages/http-routes/src/routes/web/ai-memories.ts +++ b/packages/http-routes/src/routes/web/ai-memories.ts @@ -1,8 +1,14 @@ import type { FastifyInstance, FastifyReply } from 'fastify' -import type { AIMemoryScope, AIMemoryScopeType } from '@openchatlab/shared-types' +import type { AIMemoryScope, AIMemoryScopeType, LinkAIMemorySourcesInput } from '@openchatlab/shared-types' +import { + linkAIMemorySources, + listAIMemoriesWithSourceStatus, + resolveAIMemorySourceStatus, +} from '@openchatlab/node-runtime' import type { AiRouteContext } from '../../context/ai' +import type { MemoryProvenanceCoordinator } from './memory-provenance-coordinator' -type AiMemoryRouteContext = Pick +type AiMemoryRouteContext = Pick interface MemoryScopeInput { scopeType?: string @@ -25,9 +31,22 @@ function sendBadRequest(reply: FastifyReply, error: unknown) { return reply.code(400).send({ error: error instanceof Error ? error.message : String(error) }) } -export function registerAiMemoryRoutes(server: FastifyInstance, ctx: AiMemoryRouteContext): void { +export function registerAiMemoryRoutes( + server: FastifyInstance, + ctx: AiMemoryRouteContext, + memoryProvenanceCoordinator?: MemoryProvenanceCoordinator +): void { const memoryService = ctx.aiMemoryService if (!memoryService) return + const aiChatManager = ctx.aiChatManager + const withSourceStatus = (entry: ReturnType) => ({ + ...entry, + sourceStatus: aiChatManager + ? resolveAIMemorySourceStatus(entry, aiChatManager) + : entry.sourceAIChatId + ? ('unavailable' as const) + : ('none' as const), + }) server.get<{ Querystring: MemoryScopeInput @@ -36,9 +55,45 @@ export function registerAiMemoryRoutes(server: FastifyInstance, ctx: AiMemoryRou const { scopeType, scopeId } = request.query if (!scopeType) { if (scopeId != null) throw new Error('scopeType is required when scopeId is provided') - return memoryService.list() + return aiChatManager + ? listAIMemoriesWithSourceStatus(memoryService, aiChatManager) + : memoryService.list().map((entry) => ({ + ...entry, + sourceStatus: entry.sourceAIChatId ? ('unavailable' as const) : ('none' as const), + })) } - return memoryService.list(parseScope({ scopeType, scopeId })) + const scope = parseScope({ scopeType, scopeId }) + return aiChatManager + ? listAIMemoriesWithSourceStatus(memoryService, aiChatManager, scope) + : memoryService.list(scope).map((entry) => ({ + ...entry, + sourceStatus: entry.sourceAIChatId ? ('unavailable' as const) : ('none' as const), + })) + } catch (error) { + return sendBadRequest(reply, error) + } + }) + + server.post<{ + Body: Partial + }>('/_web/ai/memories/link-sources', async (request, reply) => { + if (!aiChatManager) { + return reply.code(503).send({ error: 'AI conversation storage is unavailable' }) + } + try { + const body = request.body ?? {} + const input = { + provenanceToken: body.provenanceToken ?? '', + aiChatId: body.aiChatId ?? '', + userMessageId: body.userMessageId ?? '', + assistantMessageId: body.assistantMessageId ?? '', + memoryIds: body.memoryIds ?? [], + } + if (!memoryProvenanceCoordinator) throw new Error('Memory provenance coordination is unavailable') + const provenance = memoryProvenanceCoordinator.validate(input.provenanceToken, input.aiChatId, input.memoryIds) + const result = linkAIMemorySources(memoryService, aiChatManager, input, provenance) + memoryProvenanceCoordinator.consume(input.provenanceToken) + return result } catch (error) { return sendBadRequest(reply, error) } @@ -49,13 +104,15 @@ export function registerAiMemoryRoutes(server: FastifyInstance, ctx: AiMemoryRou }>('/_web/ai/memories', async (request, reply) => { try { const scope = parseScope(request.body ?? {}) - return memoryService.create({ - ...scope, - content: request.body?.content ?? '', - sourceType: 'user', - sourceAIChatId: null, - sourceMessageId: null, - }) + return withSourceStatus( + memoryService.create({ + ...scope, + content: request.body?.content ?? '', + sourceType: 'user', + sourceAIChatId: null, + sourceMessageId: null, + }) + ) } catch (error) { return sendBadRequest(reply, error) } @@ -69,12 +126,14 @@ export function registerAiMemoryRoutes(server: FastifyInstance, ctx: AiMemoryRou return reply.code(404).send({ error: 'AI memory not found' }) } try { - return memoryService.update(request.params.id, { - content: request.body?.content ?? '', - sourceType: 'user', - sourceAIChatId: null, - sourceMessageId: null, - }) + return withSourceStatus( + memoryService.update(request.params.id, { + content: request.body?.content ?? '', + sourceType: 'user', + sourceAIChatId: null, + sourceMessageId: null, + }) + ) } catch (error) { return sendBadRequest(reply, error) } diff --git a/packages/http-routes/src/routes/web/memory-provenance-coordinator.ts b/packages/http-routes/src/routes/web/memory-provenance-coordinator.ts new file mode 100644 index 00000000..3f6a56dc --- /dev/null +++ b/packages/http-routes/src/routes/web/memory-provenance-coordinator.ts @@ -0,0 +1,66 @@ +interface PendingMemoryProvenance { + aiChatId: string + expectedParentMessageId: string | null + memoryIds: Set + completed: boolean + expiresAt: number +} + +const DEFAULT_TTL_MS = 10 * 60 * 1000 + +export class MemoryProvenanceCoordinator { + private readonly pending = new Map() + + constructor(private readonly ttlMs = DEFAULT_TTL_MS) {} + + begin(token: string, aiChatId: string, expectedParentMessageId: string | null): void { + this.cleanupExpired() + this.pending.set(token, { + aiChatId, + expectedParentMessageId, + memoryIds: new Set(), + completed: false, + expiresAt: Date.now() + this.ttlMs, + }) + } + + record(token: string, memoryId: string): void { + const changeSet = this.pending.get(token) + if (!changeSet || changeSet.completed || Date.now() > changeSet.expiresAt) return + changeSet.memoryIds.add(memoryId) + } + + complete(token: string): void { + const changeSet = this.pending.get(token) + if (!changeSet) return + changeSet.completed = true + } + + validate(token: string, aiChatId: string, memoryIds: string[]): { expectedParentMessageId: string | null } { + this.cleanupExpired() + const changeSet = this.pending.get(token) + if (!changeSet || !changeSet.completed || changeSet.aiChatId !== aiChatId) { + throw new Error('Memory provenance token is invalid or expired') + } + + const requestedIds = new Set(memoryIds) + if ( + requestedIds.size !== changeSet.memoryIds.size || + [...requestedIds].some((memoryId) => !changeSet.memoryIds.has(memoryId)) + ) { + throw new Error('Memory IDs do not match the agent turn change set') + } + return { expectedParentMessageId: changeSet.expectedParentMessageId } + } + + consume(token: string): void { + this.pending.delete(token) + } + + private cleanupExpired(): void { + const now = Date.now() + for (const [token, changeSet] of this.pending) { + if (now > changeSet.expiresAt) this.pending.delete(token) + } + } +} diff --git a/packages/http-routes/src/routes/web/register-ai.ts b/packages/http-routes/src/routes/web/register-ai.ts index 4801eeb7..371afb87 100644 --- a/packages/http-routes/src/routes/web/register-ai.ts +++ b/packages/http-routes/src/routes/web/register-ai.ts @@ -14,6 +14,7 @@ import { registerSemanticIndexRoutes } from './ai-semantic-index' import { registerAiSkillRoutes } from './ai-skills' import { registerAiSummaryRoutes } from './ai-summaries' import { registerAiToolRoutes } from './ai-tools' +import { MemoryProvenanceCoordinator } from './memory-provenance-coordinator' export type AiRoutesContext = AiRouteContext & Pick & @@ -46,15 +47,16 @@ function assertRequiredAiContext(ctx: AiRouteContext): void { /** Register all AI-related Web routes while preserving optional-feature behavior. */ export function registerAiRoutes(server: FastifyInstance, ctx: AiRoutesContext, options?: AiRouteOptions): void { if (options?.requireAi) assertRequiredAiContext(ctx) + const memoryProvenanceCoordinator = new MemoryProvenanceCoordinator() registerAiAssistantRoutes(server, ctx) registerAiSkillRoutes(server, ctx) registerAiLlmRoutes(server, ctx) registerAiLlmStreamRoutes(server, ctx) - registerAiAgentStreamRoutes(server, ctx) + registerAiAgentStreamRoutes(server, ctx, memoryProvenanceCoordinator) registerAiToolRoutes(server, ctx) registerAiChatRoutes(server, ctx) - registerAiMemoryRoutes(server, ctx) + registerAiMemoryRoutes(server, ctx, memoryProvenanceCoordinator) registerAiChatTopicRoutes(server, ctx) registerAiSummaryRoutes(server, ctx) registerAiLogRoutes(server, ctx) diff --git a/packages/node-runtime/src/ai/__tests__/memory-provenance.test.ts b/packages/node-runtime/src/ai/__tests__/memory-provenance.test.ts new file mode 100644 index 00000000..ef93a8dc --- /dev/null +++ b/packages/node-runtime/src/ai/__tests__/memory-provenance.test.ts @@ -0,0 +1,262 @@ +import assert from 'node:assert/strict' +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { describe, it } from 'node:test' +import { AIChatManager } from '../chats' +import { AIMemoryService } from '../memory' +import { linkAIMemorySources, listAIMemoriesWithSourceStatus, resolveAIMemorySourceStatus } from '../memory-provenance' + +const sqliteNativeBinding = process.env.CHATLAB_TEST_SQLITE_NATIVE_BINDING + +function createFixture() { + const dir = mkdtempSync(join(tmpdir(), 'chatlab-ai-memory-provenance-')) + const memoryService = new AIMemoryService(dir, { nativeBinding: sqliteNativeBinding, now: () => 1234 }) + const aiChatManager = new AIChatManager(dir, { nativeBinding: sqliteNativeBinding }) + return { + dir, + memoryService, + aiChatManager, + close() { + memoryService.close() + aiChatManager.close() + try { + rmSync(dir, { recursive: true, force: true, maxRetries: 5, retryDelay: 50 }) + } catch { + // Windows can hold SQLite WAL handles briefly after close; temp cleanup is best-effort. + } + }, + } +} + +describe('AI memory provenance', () => { + it('links user and AI memories to the matching persisted messages without changing update time', () => { + const fixture = createFixture() + try { + const chat = fixture.aiChatManager.createGlobalAIChat('source', 'general_cn') + const userMessage = fixture.aiChatManager.addMessage(chat.id, 'user', '请记住我喜欢先看结论') + const assistantMessage = fixture.aiChatManager.addMessage(chat.id, 'assistant', '已经记住') + const userMemory = fixture.memoryService.create({ + scopeType: 'global', + scopeId: null, + content: '回答时先给结论', + sourceType: 'user', + sourceAIChatId: chat.id, + }) + const aiMemory = fixture.memoryService.create({ + scopeType: 'self', + scopeId: null, + content: '用户可能在准备搬家', + sourceType: 'ai', + sourceAIChatId: chat.id, + }) + + const result = linkAIMemorySources(fixture.memoryService, fixture.aiChatManager, { + provenanceToken: 'test-turn', + aiChatId: chat.id, + userMessageId: userMessage.id, + assistantMessageId: assistantMessage.id, + memoryIds: [userMemory.id, aiMemory.id, userMemory.id], + }) + + assert.deepEqual(result, { linkedMemoryIds: [userMemory.id, aiMemory.id], skippedMemoryIds: [] }) + assert.equal(fixture.memoryService.get(userMemory.id)?.sourceMessageId, userMessage.id) + assert.equal(fixture.memoryService.get(aiMemory.id)?.sourceMessageId, assistantMessage.id) + assert.equal(fixture.memoryService.get(userMemory.id)?.updatedAt, 1234) + assert.equal(fixture.memoryService.get(aiMemory.id)?.updatedAt, 1234) + + assert.deepEqual( + linkAIMemorySources(fixture.memoryService, fixture.aiChatManager, { + provenanceToken: 'test-turn', + aiChatId: chat.id, + userMessageId: userMessage.id, + assistantMessageId: assistantMessage.id, + memoryIds: [userMemory.id, aiMemory.id], + }), + { linkedMemoryIds: [userMemory.id, aiMemory.id], skippedMemoryIds: [] } + ) + } finally { + fixture.close() + } + }) + + it('rejects cross-conversation memories and invalid source message roles before writing anything', () => { + const fixture = createFixture() + try { + const chatA = fixture.aiChatManager.createGlobalAIChat('a', 'general_cn') + const chatB = fixture.aiChatManager.createGlobalAIChat('b', 'general_cn') + const userA = fixture.aiChatManager.addMessage(chatA.id, 'user', 'a user') + const assistantA = fixture.aiChatManager.addMessage(chatA.id, 'assistant', 'a assistant') + const memoryA = fixture.memoryService.create({ + scopeType: 'global', + scopeId: null, + content: 'memory a', + sourceType: 'user', + sourceAIChatId: chatA.id, + }) + const memoryB = fixture.memoryService.create({ + scopeType: 'global', + scopeId: null, + content: 'memory b', + sourceType: 'user', + sourceAIChatId: chatB.id, + }) + + assert.throws( + () => + linkAIMemorySources(fixture.memoryService, fixture.aiChatManager, { + provenanceToken: 'test-turn', + aiChatId: chatA.id, + userMessageId: userA.id, + assistantMessageId: assistantA.id, + memoryIds: [memoryA.id, memoryB.id], + }), + /does not belong/i + ) + assert.equal(fixture.memoryService.get(memoryA.id)?.sourceMessageId, null) + assert.equal(fixture.memoryService.get(memoryB.id)?.sourceMessageId, null) + + assert.throws( + () => + linkAIMemorySources(fixture.memoryService, fixture.aiChatManager, { + provenanceToken: 'test-turn', + aiChatId: chatA.id, + userMessageId: assistantA.id, + assistantMessageId: userA.id, + memoryIds: [memoryA.id], + }), + /source message/i + ) + assert.equal(fixture.memoryService.get(memoryA.id)?.sourceMessageId, null) + } finally { + fixture.close() + } + }) + + it('treats a memory deleted before backfill as an idempotent no-op', () => { + const fixture = createFixture() + try { + const chat = fixture.aiChatManager.createGlobalAIChat('source', 'general_cn') + const userMessage = fixture.aiChatManager.addMessage(chat.id, 'user', 'remember') + const assistantMessage = fixture.aiChatManager.addMessage(chat.id, 'assistant', 'done') + + assert.deepEqual( + linkAIMemorySources(fixture.memoryService, fixture.aiChatManager, { + provenanceToken: 'test-turn', + aiChatId: chat.id, + userMessageId: userMessage.id, + assistantMessageId: assistantMessage.id, + memoryIds: ['already-deleted'], + }), + { linkedMemoryIds: [], skippedMemoryIds: ['already-deleted'] } + ) + } finally { + fixture.close() + } + }) + + it('rejects non-parent messages and batch conflicts without partially linking memories', () => { + const fixture = createFixture() + try { + const chat = fixture.aiChatManager.createGlobalAIChat('source', 'general_cn') + const firstUser = fixture.aiChatManager.addMessage(chat.id, 'user', 'first') + const firstAssistant = fixture.aiChatManager.addMessage(chat.id, 'assistant', 'first answer') + const conflictingMemory = fixture.memoryService.create({ + scopeType: 'global', + scopeId: null, + content: 'already linked', + sourceType: 'user', + sourceAIChatId: chat.id, + sourceMessageId: firstUser.id, + }) + const secondUser = fixture.aiChatManager.addMessage(chat.id, 'user', 'second') + const secondAssistant = fixture.aiChatManager.addMessage(chat.id, 'assistant', 'second answer') + const unlinkedMemory = fixture.memoryService.create({ + scopeType: 'global', + scopeId: null, + content: 'must remain unlinked', + sourceType: 'user', + sourceAIChatId: chat.id, + }) + + assert.throws( + () => + linkAIMemorySources(fixture.memoryService, fixture.aiChatManager, { + provenanceToken: 'test-turn', + aiChatId: chat.id, + userMessageId: firstUser.id, + assistantMessageId: secondAssistant.id, + memoryIds: [unlinkedMemory.id], + }), + /directly reply/i + ) + + assert.throws( + () => + linkAIMemorySources(fixture.memoryService, fixture.aiChatManager, { + provenanceToken: 'test-turn', + aiChatId: chat.id, + userMessageId: secondUser.id, + assistantMessageId: secondAssistant.id, + memoryIds: [unlinkedMemory.id, conflictingMemory.id], + }), + /already linked/i + ) + assert.equal(fixture.memoryService.get(unlinkedMemory.id)?.sourceMessageId, null) + assert.equal(fixture.memoryService.get(conflictingMemory.id)?.sourceMessageId, firstUser.id) + assert.equal(firstAssistant.parentId, firstUser.id) + } finally { + fixture.close() + } + }) + + it('reports conversation, message, and unavailable source states without deleting the memory', () => { + const fixture = createFixture() + try { + const chat = fixture.aiChatManager.createGlobalAIChat('source', 'general_cn') + const message = fixture.aiChatManager.addMessage(chat.id, 'user', 'remember') + const none = fixture.memoryService.create({ + scopeType: 'global', + scopeId: null, + content: 'manual', + sourceType: 'user', + }) + const conversation = fixture.memoryService.create({ + scopeType: 'global', + scopeId: null, + content: 'conversation only', + sourceType: 'user', + sourceAIChatId: chat.id, + }) + const linked = fixture.memoryService.create({ + scopeType: 'global', + scopeId: null, + content: 'linked', + sourceType: 'user', + sourceAIChatId: chat.id, + sourceMessageId: message.id, + }) + + assert.equal(resolveAIMemorySourceStatus(none, fixture.aiChatManager), 'none') + assert.equal(resolveAIMemorySourceStatus(conversation, fixture.aiChatManager), 'conversation') + assert.equal(resolveAIMemorySourceStatus(linked, fixture.aiChatManager), 'message') + + fixture.aiChatManager.deleteMessage(message.id) + assert.equal(resolveAIMemorySourceStatus(linked, fixture.aiChatManager), 'unavailable') + fixture.aiChatManager.deleteAIChat(chat.id) + assert.equal(resolveAIMemorySourceStatus(conversation, fixture.aiChatManager), 'unavailable') + assert.equal(fixture.memoryService.list().length, 3) + assert.deepEqual( + Object.fromEntries( + listAIMemoriesWithSourceStatus(fixture.memoryService, fixture.aiChatManager).map((entry) => [ + entry.content, + entry.sourceStatus, + ]) + ), + { manual: 'none', 'conversation only': 'unavailable', linked: 'unavailable' } + ) + } finally { + fixture.close() + } + }) +}) diff --git a/packages/node-runtime/src/ai/agent/event-handler.ts b/packages/node-runtime/src/ai/agent/event-handler.ts index ff4d63b7..c1ca43bf 100644 --- a/packages/node-runtime/src/ai/agent/event-handler.ts +++ b/packages/node-runtime/src/ai/agent/event-handler.ts @@ -38,6 +38,7 @@ export interface AgentStreamChunk { | 'tool_start' | 'tool_update' | 'tool_result' + | 'memory_change' | 'status' | 'compression_done' | 'route' @@ -56,6 +57,8 @@ export interface AgentStreamChunk { toolProgress?: ToolProgress toolResult?: unknown toolIsError?: boolean + memoryId?: string + provenanceToken?: string error?: unknown isFinished?: boolean usage?: TokenUsage diff --git a/packages/node-runtime/src/ai/chats.ts b/packages/node-runtime/src/ai/chats.ts index a0e21f4b..cb3bb287 100644 --- a/packages/node-runtime/src/ai/chats.ts +++ b/packages/node-runtime/src/ai/chats.ts @@ -811,6 +811,11 @@ export class AIChatManager { return this.getActivePathRows(aiChatId).map((row) => this.parseMessageRow(row)) } + getMessage(messageId: string): AIMessage | null { + const row = this.getMessageRow(messageId) + return row ? this.parseMessageRow(row) : null + } + deleteMessage(messageId: string): boolean { const db = this.getDb() const result = db.prepare('DELETE FROM ai_message WHERE id = ?').run(messageId) diff --git a/packages/node-runtime/src/ai/index.ts b/packages/node-runtime/src/ai/index.ts index 5c23cd45..d517e601 100644 --- a/packages/node-runtime/src/ai/index.ts +++ b/packages/node-runtime/src/ai/index.ts @@ -49,6 +49,7 @@ export type { } from './chats' export { AI_MEMORY_CONTENT_MAX_CHARS, AIMemoryService, buildGlobalMemoryPrompt } from './memory' export type { AIMemoryServiceOptions, CreateAIMemoryInput, UpdateAIMemoryInput } from './memory' +export { linkAIMemorySources, listAIMemoriesWithSourceStatus, resolveAIMemorySourceStatus } from './memory-provenance' // Tokenizer export { countTokens, countMessagesTokens, initTokenizer } from './tokenizer' diff --git a/packages/node-runtime/src/ai/memory-provenance.ts b/packages/node-runtime/src/ai/memory-provenance.ts new file mode 100644 index 00000000..f5bd9d83 --- /dev/null +++ b/packages/node-runtime/src/ai/memory-provenance.ts @@ -0,0 +1,113 @@ +import type { + AIMemoryEntry, + AIMemoryManagementEntry, + AIMemoryScope, + AIMemorySourceStatus, + LinkAIMemorySourcesInput, + LinkAIMemorySourcesResult, +} from '@openchatlab/shared-types' +import type { AIChatManager, AIMessage } from './chats' +import type { AIMemoryService } from './memory' + +export function listAIMemoriesWithSourceStatus( + memoryService: AIMemoryService, + aiChatManager: AIChatManager, + scope?: AIMemoryScope +): AIMemoryManagementEntry[] { + return memoryService.list(scope).map((entry) => ({ + ...entry, + sourceStatus: resolveAIMemorySourceStatus(entry, aiChatManager), + })) +} + +export function linkAIMemorySources( + memoryService: AIMemoryService, + aiChatManager: AIChatManager, + input: LinkAIMemorySourcesInput, + options?: { expectedParentMessageId?: string | null } +): LinkAIMemorySourcesResult { + const aiChatId = requireId(input.aiChatId, 'aiChatId') + const conversation = aiChatManager.getAIChat(aiChatId) + if (!conversation || conversation.kind !== 'global') { + throw new Error('Global AI conversation not found') + } + + const userMessage = requireSourceMessage(aiChatManager, input.userMessageId, aiChatId, 'user') + const assistantMessage = requireSourceMessage(aiChatManager, input.assistantMessageId, aiChatId, 'assistant') + if (assistantMessage.parentId !== userMessage.id) { + throw new Error('Assistant source message must directly reply to the user source message') + } + if ( + options && + options.expectedParentMessageId !== undefined && + userMessage.parentId !== options.expectedParentMessageId + ) { + throw new Error('User source message does not belong to the requested agent turn') + } + const memoryIds = normalizeMemoryIds(input.memoryIds) + const memories = memoryIds.map((id) => ({ id, entry: memoryService.get(id) })) + + for (const { entry } of memories) { + if (entry && entry.sourceAIChatId !== aiChatId) { + throw new Error('AI memory does not belong to the requested conversation') + } + if (entry?.sourceMessageId) { + const expectedSourceMessageId = entry.sourceType === 'user' ? userMessage.id : assistantMessage.id + if (entry.sourceMessageId !== expectedSourceMessageId) { + throw new Error('AI memory is already linked to another source message') + } + } + } + + const skippedMemoryIds: string[] = [] + const links: Array<{ id: string; sourceAIChatId: string; sourceMessageId: string }> = [] + for (const { id, entry } of memories) { + if (!entry) { + skippedMemoryIds.push(id) + continue + } + + const sourceMessage = entry.sourceType === 'user' ? userMessage : assistantMessage + links.push({ id, sourceAIChatId: aiChatId, sourceMessageId: sourceMessage.id }) + } + + const linkedMemoryIds = memoryService.linkSourceMessages(links) + return { linkedMemoryIds, skippedMemoryIds } +} + +export function resolveAIMemorySourceStatus(entry: AIMemoryEntry, aiChatManager: AIChatManager): AIMemorySourceStatus { + if (!entry.sourceAIChatId) return 'none' + + const conversation = aiChatManager.getAIChat(entry.sourceAIChatId) + if (!conversation || conversation.kind !== 'global') return 'unavailable' + if (!entry.sourceMessageId) return 'conversation' + + const message = aiChatManager.getMessage(entry.sourceMessageId) + const expectedRole = entry.sourceType === 'user' ? 'user' : 'assistant' + return message?.aiChatId === entry.sourceAIChatId && message.role === expectedRole ? 'message' : 'unavailable' +} + +function requireSourceMessage( + aiChatManager: AIChatManager, + messageId: string, + aiChatId: string, + role: AIMessage['role'] +): AIMessage { + const normalizedMessageId = requireId(messageId, `${role}MessageId`) + const message = aiChatManager.getMessage(normalizedMessageId) + if (!message || message.aiChatId !== aiChatId || message.role !== role) { + throw new Error(`${role} source message does not belong to the requested conversation`) + } + return message +} + +function normalizeMemoryIds(memoryIds: string[]): string[] { + if (!Array.isArray(memoryIds)) throw new Error('memoryIds must be an array') + return [...new Set(memoryIds.map((id) => requireId(id, 'memoryId')))] +} + +function requireId(value: string, field: string): string { + const normalized = typeof value === 'string' ? value.trim() : '' + if (!normalized) throw new Error(`${field} is required`) + return normalized +} diff --git a/packages/node-runtime/src/ai/memory.ts b/packages/node-runtime/src/ai/memory.ts index b018928a..dbc37416 100644 --- a/packages/node-runtime/src/ai/memory.ts +++ b/packages/node-runtime/src/ai/memory.ts @@ -371,6 +371,29 @@ export class AIMemoryService { }) } + linkSourceMessages(links: Array<{ id: string; sourceAIChatId: string; sourceMessageId: string }>): string[] { + if (links.length === 0) return [] + return this.runWrite('link sources', () => { + const db = this.getDb() + const update = db.prepare( + `UPDATE ai_memory + SET source_message_id = ? + WHERE id = ? + AND source_ai_chat_id = ? + AND (source_message_id IS NULL OR source_message_id = ?)` + ) + return db.transaction(() => { + const linkedIds: string[] = [] + for (const link of links) { + const result = update.run(link.sourceMessageId, link.id, link.sourceAIChatId, link.sourceMessageId) + if (result.changes === 0) throw new Error('AI memory source changed before it could be linked') + linkedIds.push(link.id) + } + return linkedIds + })() + }) + } + forget(id: string): boolean { return this.runWrite('forget', () => this.getDb().prepare('DELETE FROM ai_memory WHERE id = ?').run(id).changes > 0) } diff --git a/packages/node-runtime/src/index.ts b/packages/node-runtime/src/index.ts index 24101e63..0847fc3c 100644 --- a/packages/node-runtime/src/index.ts +++ b/packages/node-runtime/src/index.ts @@ -319,6 +319,7 @@ export { runCrossChatAgent, } from './ai' export type { CrossChatAgentLogger, RunCrossChatAgentOptions } from './ai' +export { linkAIMemorySources, listAIMemoriesWithSourceStatus, resolveAIMemorySourceStatus } from './ai' export type { AnalysisPlanIntent, AnalysisPlanner, diff --git a/packages/shared-types/index.ts b/packages/shared-types/index.ts index 7127ee84..335f907b 100644 --- a/packages/shared-types/index.ts +++ b/packages/shared-types/index.ts @@ -138,6 +138,7 @@ export type AIEntityRef = export type AIMemoryScopeType = 'global' | 'self' | 'contact' | 'group' export type AIMemorySourceType = 'user' | 'ai' +export type AIMemorySourceStatus = 'none' | 'conversation' | 'message' | 'unavailable' export interface AIMemoryScope { scopeType: AIMemoryScopeType @@ -154,6 +155,23 @@ export interface AIMemoryEntry extends AIMemoryScope { updatedAt: number } +export interface AIMemoryManagementEntry extends AIMemoryEntry { + sourceStatus: AIMemorySourceStatus +} + +export interface LinkAIMemorySourcesInput { + provenanceToken: string + aiChatId: string + userMessageId: string + assistantMessageId: string + memoryIds: string[] +} + +export interface LinkAIMemorySourcesResult { + linkedMemoryIds: string[] + skippedMemoryIds: string[] +} + export { CHATLAB_FORMAT_VERSION, CHATLAB_SUPPORTED_FORMAT_VERSIONS, diff --git a/packages/tools/src/definitions/memory-tools.test.ts b/packages/tools/src/definitions/memory-tools.test.ts index ceded48b..dba28706 100644 --- a/packages/tools/src/definitions/memory-tools.test.ts +++ b/packages/tools/src/definitions/memory-tools.test.ts @@ -116,7 +116,11 @@ function entityResolution(refs: AIEntityRef[]): CrossChatEntityResolution { } function createContext( - options: { entityRefs?: AIEntityRef[]; resolvedEntityRefs?: AIEntityRef[] } = {} + options: { + entityRefs?: AIEntityRef[] + resolvedEntityRefs?: AIEntityRef[] + reportMemoryChange?: (memoryId: string) => void + } = {} ): CrossChatToolExecutionContext { const analysisService = { resolveEntities: async (refs: AIEntityRef[]) => entityResolution(refs), @@ -128,6 +132,7 @@ function createContext( analysisService, memoryService: new FakeMemoryService(), aiChatId: 'global-chat-1', + reportMemoryChange: options.reportMemoryChange, preprocessMessagesBySession: (_sessionId, messages) => messages, preprocessSummariesBySession: (_sessionId, summaries) => summaries, preprocessModelLabel: (value) => value, @@ -157,7 +162,8 @@ describe('global memory tool registry', () => { describe('global memory tools', () => { it('writes user memory with runtime-owned source metadata and updates only the same scope', async () => { - const context = createContext() + const changedMemoryIds: string[] = [] + const context = createContext({ reportMemoryChange: (memoryId) => changedMemoryIds.push(memoryId) }) const created = await memoryWriteTool.handler( { scope_type: 'contact', @@ -182,6 +188,7 @@ describe('global memory tools', () => { context ) assert.equal((updated.data as AIMemoryEntry).content, '用户纠正为高中同学') + assert.deepEqual(changedMemoryIds, [createdEntry.id, createdEntry.id]) await assert.rejects( async () => diff --git a/packages/tools/src/definitions/memory-tools.ts b/packages/tools/src/definitions/memory-tools.ts index 09f139f0..dca41c57 100644 --- a/packages/tools/src/definitions/memory-tools.ts +++ b/packages/tools/src/definitions/memory-tools.ts @@ -128,6 +128,7 @@ async function memoryWriteHandler( }) } + context.reportMemoryChange?.(entry.id) return { content: JSON.stringify(entry), data: entry } } diff --git a/packages/tools/src/types.ts b/packages/tools/src/types.ts index bbf508b7..00bdb151 100644 --- a/packages/tools/src/types.ts +++ b/packages/tools/src/types.ts @@ -438,6 +438,8 @@ export interface CrossChatToolExecutionContext { analysisService: CrossChatAnalysisToolService memoryService: AIMemoryToolService aiChatId: string + /** Reports a successful memory create/update without exposing it to the model context. */ + reportMemoryChange?: (memoryId: string) => void preprocessMessagesBySession: ( sessionId: string, messages: CrossChatMessageSource[] diff --git a/src/App.vue b/src/App.vue index cab2ff3f..c811f87c 100644 --- a/src/App.vue +++ b/src/App.vue @@ -366,11 +366,14 @@ onMounted(async () => { // 同步原生 modal 在异步锁屏就绪前保持底层文档 inert；chunk 加载失败时不会自动放行。 syncBootstrapMask(bootstrapDialogRef.value) window.addEventListener('keydown', handleGlobalKeydown) - const platform = navigator.platform.toLowerCase() - if (platform.includes('win')) { - document.documentElement.classList.add('platform-windows') - } else if (platform.includes('linux')) { - document.documentElement.classList.add('platform-linux') + if (IS_ELECTRON) { + document.documentElement.classList.add('platform-electron') + const platform = navigator.platform.toLowerCase() + if (platform.includes('win')) { + document.documentElement.classList.add('platform-windows') + } else if (platform.includes('linux')) { + document.documentElement.classList.add('platform-linux') + } } if (IS_ELECTRON) { @@ -417,7 +420,7 @@ onUnmounted(() => { diff --git a/src/services/ai-stream/types.ts b/src/services/ai-stream/types.ts index e3eb6ae3..fa8ce639 100644 --- a/src/services/ai-stream/types.ts +++ b/src/services/ai-stream/types.ts @@ -17,6 +17,7 @@ export interface AgentStreamChunk { | 'tool_start' | 'tool_update' | 'tool_result' + | 'memory_change' | 'status' | 'compression_done' | 'route' @@ -35,6 +36,8 @@ export interface AgentStreamChunk { toolProgress?: ToolProgress toolResult?: unknown toolIsError?: boolean + memoryId?: string + provenanceToken?: string error?: unknown isFinished?: boolean usage?: TokenUsage diff --git a/src/services/ai/fetch.ts b/src/services/ai/fetch.ts index 48ef0b37..1d87572f 100644 --- a/src/services/ai/fetch.ts +++ b/src/services/ai/fetch.ts @@ -20,11 +20,12 @@ import type { ToolCatalogEntry, ToolExecuteResult, DesensitizeRule, - AIMemoryEntry, + AIMemoryManagementEntry, AIMemoryScope, CreateUserAIMemoryInput, ClearAIMemoriesInput, } from './types' +import type { LinkAIMemorySourcesInput, LinkAIMemorySourcesResult } from '@openchatlab/shared-types' import { get, post, put, del, fetchWithAuth } from '../utils/http' const NOT_AVAILABLE_WEB = 'This feature is not available in web mode' @@ -64,19 +65,19 @@ export class FetchAIAdapter implements AIAdapter { } // ===== 长期记忆 ===== - async getAIMemories(scope?: AIMemoryScope): Promise { - if (!scope) return get('/ai/memories') + async getAIMemories(scope?: AIMemoryScope): Promise { + if (!scope) return get('/ai/memories') const params = new URLSearchParams({ scopeType: scope.scopeType }) if (scope.scopeId) params.set('scopeId', scope.scopeId) - return get(`/ai/memories?${params.toString()}`) + return get(`/ai/memories?${params.toString()}`) } - async createAIMemory(input: CreateUserAIMemoryInput): Promise { - return post('/ai/memories', input) + async createAIMemory(input: CreateUserAIMemoryInput): Promise { + return post('/ai/memories', input) } - async updateAIMemory(id: string, content: string): Promise { - return put(`/ai/memories/${id}`, { content }) + async updateAIMemory(id: string, content: string): Promise { + return put(`/ai/memories/${id}`, { content }) } async deleteAIMemory(id: string): Promise<{ success: boolean }> { @@ -87,6 +88,10 @@ export class FetchAIAdapter implements AIAdapter { return post<{ success: boolean; cleared: number }>('/ai/memories/clear', input) } + async linkAIMemorySources(input: LinkAIMemorySourcesInput): Promise { + return post('/ai/memories/link-sources', input) + } + // ===== 消息 ===== async getMessages(aiChatId: string): Promise { return get(`/ai/chats/${aiChatId}/messages`) diff --git a/src/services/ai/types.ts b/src/services/ai/types.ts index c316f2b3..d076f7c7 100644 --- a/src/services/ai/types.ts +++ b/src/services/ai/types.ts @@ -1,6 +1,13 @@ import type { TimeFilter } from '@/types/base' import type { ChartPayload, ChatEvidencePayload } from '@openchatlab/core' -import type { AIEntityRef, AIMemoryEntry, AIMemoryScope, CrossChatEvidencePayload } from '@openchatlab/shared-types' +import type { + AIEntityRef, + AIMemoryManagementEntry, + AIMemoryScope, + CrossChatEvidencePayload, + LinkAIMemorySourcesInput, + LinkAIMemorySourcesResult, +} from '@openchatlab/shared-types' import type { PlanContentBlock } from './planBlocks' export interface AIChat { @@ -14,7 +21,7 @@ export interface AIChat { updatedAt: number } -export type { AIEntityRef, AIMemoryEntry, AIMemoryScope } from '@openchatlab/shared-types' +export type { AIEntityRef, AIMemoryManagementEntry, AIMemoryScope } from '@openchatlab/shared-types' export type CreateUserAIMemoryInput = AIMemoryScope & { content: string } export type ClearAIMemoriesInput = AIMemoryScope | { all: true } @@ -133,11 +140,12 @@ export interface AIAdapter { deleteAIChat(aiChatId: string): Promise // ===== 长期记忆 ===== - getAIMemories(scope?: AIMemoryScope): Promise - createAIMemory(input: CreateUserAIMemoryInput): Promise - updateAIMemory(id: string, content: string): Promise + getAIMemories(scope?: AIMemoryScope): Promise + createAIMemory(input: CreateUserAIMemoryInput): Promise + updateAIMemory(id: string, content: string): Promise deleteAIMemory(id: string): Promise<{ success: boolean }> clearAIMemories(input: ClearAIMemoriesInput): Promise<{ success: boolean; cleared: number }> + linkAIMemorySources(input: LinkAIMemorySourcesInput): Promise // ===== 消息 ===== getMessages(aiChatId: string): Promise diff --git a/src/stores/aiChat.ts b/src/stores/aiChat.ts index 88d0bf90..f107ed2e 100644 --- a/src/stores/aiChat.ts +++ b/src/stores/aiChat.ts @@ -50,6 +50,7 @@ import { } from './aiChatContentBlocks' import { createToolLifecycleTracker, type ToolStatus } from './aiToolLifecycle' import { applyQueuedStreamTextDeltas, createAIStreamTextBatcher } from './aiChatStreamBatcher' +import { reportError } from '@/services/log-report' export type { ToolStatus } from './aiToolLifecycle' @@ -1041,6 +1042,8 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { // 在 try 外部声明，以便 catch 块能正确引用当前轮次的用户消息 let currentUserMessage: ChatMessage | undefined let lastDoneUsage: TokenUsage | undefined + const changedMemoryIds = new Set() + let memoryProvenanceToken: string | null = null targetBuffer.assistantId = state.selectedAssistantId targetBuffer.loaded = true @@ -1318,6 +1321,11 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { } break + case 'memory_change': + if (chunk.memoryId) changedMemoryIds.add(chunk.memoryId) + if (chunk.provenanceToken) memoryProvenanceToken = chunk.provenanceToken + break + case 'plan': if (chunk.plan) { appendPlanToBlocks(chunk.plan) @@ -1433,7 +1441,9 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { resolvedAIChatId, userMessage, targetBuffer.messages[aiMessageIndex], - lastDoneUsage + lastDoneUsage, + [...changedMemoryIds], + memoryProvenanceToken ) if (!savedMessages) return { success: false, reason: 'error' } Object.assign(userMessage, savedMessages.userMessage) @@ -1457,7 +1467,9 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { resolvedAIChatId, userMessage, targetBuffer.messages[aiMessageIndex], - lastDoneUsage + lastDoneUsage, + [...changedMemoryIds], + memoryProvenanceToken ) if (!savedMessages) return { success: false, reason: 'error' } Object.assign(userMessage, savedMessages.userMessage) @@ -1471,7 +1483,9 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { resolvedAIChatId, userMessage, targetBuffer.messages[aiMessageIndex], - lastDoneUsage + lastDoneUsage, + [...changedMemoryIds], + memoryProvenanceToken ) if (!savedMessages) return { success: false, reason: 'error' } Object.assign(userMessage, savedMessages.userMessage) @@ -1505,7 +1519,14 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { // 优先使用当前轮次的用户消息，避免多轮对话取到第一条历史消息 const userMsg = currentUserMessage || targetBuffer.messages.findLast((m) => m.role === 'user') if (userMsg) { - await saveAIChatMessages(resolvedAIChatId, userMsg, lastMessage, lastDoneUsage) + await saveAIChatMessages( + resolvedAIChatId, + userMsg, + lastMessage, + lastDoneUsage, + [...changedMemoryIds], + memoryProvenanceToken + ) } } @@ -1525,7 +1546,9 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { aiChatId: string | null, userMsg: ChatMessage, aiMsg: ChatMessage, - tokenUsage?: TokenUsage + tokenUsage?: TokenUsage, + changedMemoryIds: string[] = [], + memoryProvenanceToken: string | null = null ): Promise<{ userMessage: ChatMessage; assistantMessage: ChatMessage } | null> { try { if (!aiChatId) { @@ -1539,6 +1562,20 @@ export const useAIChatStore = defineStore('aiChatRuntime', () => { { content: userMsg.content, entityRefs: userMsg.entityRefs }, { content: aiMsg.content, contentBlocks: serializableContentBlocks, tokenUsage } ) + if (changedMemoryIds.length > 0) { + try { + if (!memoryProvenanceToken) throw new Error('Memory provenance token is missing') + await useAIService().linkAIMemorySources({ + provenanceToken: memoryProvenanceToken, + aiChatId, + userMessageId: savedUserMessage.id, + assistantMessageId: savedAssistantMessage.id, + memoryIds: changedMemoryIds, + }) + } catch (error) { + reportError('AI memory source linking failed', error instanceof Error ? error.stack : undefined) + } + } return { userMessage: toRuntimeMessage(savedUserMessage), assistantMessage: toRuntimeMessage(savedAssistantMessage),