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.

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<Array<{ entityRefs?: unknown[] }>>()[1]?.entityRefs, undefined)
+
   const globalList = await app.inject({ method: 'GET', url: '/_web/ai/global-chats' })
   assert.deepEqual(
     globalList.json<Array<{ id: string }>>().map((chat) => chat.id),

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')

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<chatlab_entity_refs>'))
+  })
+
   it('falls back to plain text when tool blocks lack persisted toolCallId/result (legacy rows)', () => {
     const history: SimpleHistoryMessage[] = [
       {

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
         )

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,
     }
   }