From 4f6bd0188b2e5cb8a8894e8046348b2dfdf73cc3 Mon Sep 17 00:00:00 2001 From: digua Date: Sat, 22 Aug 2026 00:13:35 +0800 Subject: [PATCH] fix(ai): restrict entity refs to user messages Prevent generic assistant or summary writes from persisting structured entity references. Keep compressed-summary references on the dedicated internal path and ignore malformed assistant references during replay and compression. --- .../src/routes/web/ai-chats.test.ts | 14 +++++ .../src/ai/__tests__/chats.test.ts | 25 ++++++++- .../src/ai/agent/__tests__/history.test.ts | 28 ++++++++++ packages/node-runtime/src/ai/agent/history.ts | 7 ++- packages/node-runtime/src/ai/chats.ts | 51 +++++++++++++++---- 5 files changed, 114 insertions(+), 11 deletions(-) diff --git a/packages/http-routes/src/routes/web/ai-chats.test.ts b/packages/http-routes/src/routes/web/ai-chats.test.ts index 710b083c..46680de4 100644 --- a/packages/http-routes/src/routes/web/ai-chats.test.ts +++ b/packages/http-routes/src/routes/web/ai-chats.test.ts @@ -56,6 +56,20 @@ test('global AI chat routes keep global history separate and persist entity refe assert.equal(messageResponse.statusCode, 200) assert.deepEqual(messageResponse.json<{ entityRefs: unknown[] }>().entityRefs, refs) + const assistantResponse = await app.inject({ + method: 'POST', + url: `/_web/ai/chats/${globalChat.id}/messages`, + payload: { role: 'assistant', content: 'I will compare them', entityRefs: refs }, + }) + assert.equal(assistantResponse.statusCode, 200) + assert.equal(assistantResponse.json<{ entityRefs?: unknown[] }>().entityRefs, undefined) + + const messageList = await app.inject({ + method: 'GET', + url: `/_web/ai/chats/${globalChat.id}/messages`, + }) + assert.equal(messageList.json>()[1]?.entityRefs, undefined) + const globalList = await app.inject({ method: 'GET', url: '/_web/ai/global-chats' }) assert.deepEqual( globalList.json>().map((chat) => chat.id), diff --git a/packages/node-runtime/src/ai/__tests__/chats.test.ts b/packages/node-runtime/src/ai/__tests__/chats.test.ts index 3642f8c4..d4cb3858 100644 --- a/packages/node-runtime/src/ai/__tests__/chats.test.ts +++ b/packages/node-runtime/src/ai/__tests__/chats.test.ts @@ -396,10 +396,33 @@ describe('AIChatManager global chats and entity references', () => { undefined, entityRefs ) - manager.addMessage(globalChat.id, 'assistant', 'I will compare them.') + const assistantMessage = manager.addMessage( + globalChat.id, + 'assistant', + 'I will compare them.', + undefined, + undefined, + undefined, + undefined, + entityRefs + ) assert.deepEqual(manager.getMessages(globalChat.id)[0]?.entityRefs, entityRefs) assert.deepEqual(manager.getHistoryForAgent(globalChat.id)[0]?.entityRefs, entityRefs) + assert.equal(assistantMessage.entityRefs, undefined) + assert.equal(manager.getMessages(globalChat.id)[1]?.entityRefs, undefined) + assert.equal(manager.getHistoryForAgent(globalChat.id)[1]?.entityRefs, undefined) + + const insertedAssistant = manager.insertMessageAfter( + globalChat.id, + userMessage.id, + 'assistant', + 'Updated comparison', + undefined, + undefined, + entityRefs + ) + assert.equal(insertedAssistant.entityRefs, undefined) const fork = manager.forkAIChat(globalChat.id, userMessage.id) assert.equal(fork.kind, 'global') diff --git a/packages/node-runtime/src/ai/agent/__tests__/history.test.ts b/packages/node-runtime/src/ai/agent/__tests__/history.test.ts index 6aed3b1e..c230e99f 100644 --- a/packages/node-runtime/src/ai/agent/__tests__/history.test.ts +++ b/packages/node-runtime/src/ai/agent/__tests__/history.test.ts @@ -69,6 +69,34 @@ describe('toPiHistoryMessages — legacy plain history', () => { assert.ok(text?.type === 'text' && text.text.includes('"contactKey":"qq:10001"')) }) + it('does not replay entity references attached to assistant messages', () => { + const out = toPiHistoryMessages([ + { + role: 'assistant', + content: 'comparison result', + entityRefs: [{ type: 'contact', contactKey: 'qq:10001', displayName: 'Alice' }], + }, + ]) + + assert.equal(out[0]?.role, 'assistant') + assert.deepEqual(out[0]?.content, [{ type: 'text', text: 'comparison result' }]) + }) + + it('replays entity references retained by a compressed summary', () => { + const out = toPiHistoryMessages([ + { + role: 'summary', + content: 'compressed context', + entityRefs: [{ type: 'contact', contactKey: 'qq:10001', displayName: 'Alice' }], + }, + ]) + + assert.equal(out[0]?.role, 'assistant') + if (out[0]?.role !== 'assistant') return + const text = out[0].content[0] + assert.ok(text?.type === 'text' && text.text.startsWith('compressed context\n\n')) + }) + it('falls back to plain text when tool blocks lack persisted toolCallId/result (legacy rows)', () => { const history: SimpleHistoryMessage[] = [ { diff --git a/packages/node-runtime/src/ai/agent/history.ts b/packages/node-runtime/src/ai/agent/history.ts index 7f8284f4..1337b700 100644 --- a/packages/node-runtime/src/ai/agent/history.ts +++ b/packages/node-runtime/src/ai/agent/history.ts @@ -155,7 +155,12 @@ export function toPiHistoryMessages(messages: SimpleHistoryMessage[], options?: } else { out.push( makeAssistantMessage( - [{ type: 'text', text: appendEntityRefsForModel(msg.content || '', msg.entityRefs) }], + [ + { + type: 'text', + text: appendEntityRefsForModel(msg.content || '', msg.role === 'summary' ? msg.entityRefs : undefined), + }, + ], 'stop', options?.modelInfo ) diff --git a/packages/node-runtime/src/ai/chats.ts b/packages/node-runtime/src/ai/chats.ts index 458fe483..2eae3257 100644 --- a/packages/node-runtime/src/ai/chats.ts +++ b/packages/node-runtime/src/ai/chats.ts @@ -664,6 +664,28 @@ export class AIChatManager { contentBlocks?: ContentBlock[], tokenUsage?: TokenUsageData, entityRefs?: AIEntityRef[] + ): AIMessage { + return this.addMessageWithEntityRefs( + aiChatId, + role, + content, + dataKeywords, + dataMessageCount, + contentBlocks, + tokenUsage, + role === 'user' ? entityRefs : undefined + ) + } + + private addMessageWithEntityRefs( + aiChatId: string, + role: AIMessageRole, + content: string, + dataKeywords?: string[], + dataMessageCount?: number, + contentBlocks?: ContentBlock[], + tokenUsage?: TokenUsageData, + entityRefs?: AIEntityRef[] ): AIMessage { const db = this.getDb() const now = Math.floor(Date.now() / 1000) @@ -671,6 +693,7 @@ export class AIChatManager { const parentId = this.getActiveMessageId(aiChatId) const siblingGroupId = id const branchIndex = 0 + const persistedEntityRefs = entityRefs?.length ? entityRefs : undefined const pendingDebug = role === 'assistant' ? this.pendingDebugContextMap.get(aiChatId) : undefined if (pendingDebug) { @@ -697,7 +720,7 @@ export class AIChatManager { parentId, siblingGroupId, branchIndex, - entityRefs?.length ? JSON.stringify(entityRefs) : null + persistedEntityRefs?.length ? JSON.stringify(persistedEntityRefs) : null ) db.prepare('UPDATE ai_chat SET active_message_id = ?, updated_at = ? WHERE id = ?').run(id, now, aiChatId) @@ -713,7 +736,7 @@ export class AIChatManager { dataMessageCount, contentBlocks, tokenUsage, - entityRefs: entityRefs?.length ? entityRefs : undefined, + entityRefs: persistedEntityRefs?.length ? persistedEntityRefs : undefined, } } @@ -798,7 +821,7 @@ export class AIChatManager { row.tokenUsage, newParentId, newMsgId, - row.entityRefs + row.role === 'user' || row.role === 'summary' ? row.entityRefs : null ) lastNewId = newMsgId } @@ -863,6 +886,7 @@ export class AIChatManager { const db = this.getDb() const now = Math.floor(Date.now() / 1000) const id = this.generateId('msg') + const persistedEntityRefs = role === 'user' ? entityRefs : undefined const pendingDebug = role === 'assistant' ? this.pendingDebugContextMap.get(aiChatId) : undefined if (pendingDebug) { @@ -890,7 +914,7 @@ export class AIChatManager { pendingDebug ?? null, afterMessageId, id, - entityRefs?.length ? JSON.stringify(entityRefs) : null + persistedEntityRefs?.length ? JSON.stringify(persistedEntityRefs) : null ) if (childRow) { @@ -909,7 +933,7 @@ export class AIChatManager { parentId: afterMessageId, contentBlocks, tokenUsage, - entityRefs: entityRefs?.length ? entityRefs : undefined, + entityRefs: persistedEntityRefs?.length ? persistedEntityRefs : undefined, } } @@ -1000,7 +1024,7 @@ export class AIChatManager { role: m.role, content: m.content, contentBlocks: m.contentBlocks, - entityRefs: m.entityRefs, + entityRefs: m.role === 'user' ? m.entityRefs : undefined, })), ] } else { @@ -1008,7 +1032,7 @@ export class AIChatManager { role: m.role, content: m.content, contentBlocks: m.contentBlocks, - entityRefs: m.entityRefs, + entityRefs: m.role === 'user' ? m.entityRefs : undefined, })) } @@ -1039,7 +1063,16 @@ export class AIChatManager { }, ] - return this.addMessage(aiChatId, 'summary', content, undefined, undefined, contentBlocks, undefined, entityRefs) + return this.addMessageWithEntityRefs( + aiChatId, + 'summary', + content, + undefined, + undefined, + contentBlocks, + undefined, + entityRefs + ) } getLatestSummary(aiChatId: string): AIMessage | null { @@ -1086,7 +1119,7 @@ export class AIChatManager { content: row.content, timestamp: row.timestamp, contentBlocks: row.contentBlocks ? JSON.parse(row.contentBlocks) : undefined, - entityRefs: row.entityRefs ? JSON.parse(row.entityRefs) : undefined, + entityRefs: row.role === 'user' && row.entityRefs ? JSON.parse(row.entityRefs) : undefined, } }