fix(chat): sanitize UI messages before model conversion

Author: albegosu

pages/index.vue

@@ -424,6 +424,11 @@ interface KbToolOutput {
   latencyMs?: number
 }
 
+/** Persisted on reload so citations render without synthetic tool parts in the API payload. */
+interface ChatMessageMetadata {
+  sources?: ConverseSource[]
+}
+
 const { t } = useI18n()
 const store = useDocumentsStore()
 const input = ref('')
@@ -543,27 +548,14 @@ async function loadConversation(id: string) {
           ? (m.parts as Array<{ type: string }>)
           : [{ type: 'text', text: m.content }]
 
-      // Inject a synthetic tool-output part for assistant messages that have
-      // stored sources but no tool call in their parts (e.g. older messages).
       const storedSources = Array.isArray(m.sources) ? (m.sources as ConverseSource[]) : []
-      const hasTool = rawParts.some((p) => p.type.startsWith('tool-'))
-      if (m.role === 'assistant' && storedSources.length > 0 && !hasTool) {
-        rawParts.push({
-          type: 'tool-searchKnowledgeBase',
-          state: 'output-available',
-          output: {
-            context: '',
-            sources: storedSources,
-            results: storedSources,
-            count: storedSources.length,
-          },
-        } as { type: string })
-      }
+      const partsForApi = rawParts.filter((p) => !p.type.startsWith('tool-') && p.type !== 'dynamic-tool')
 
       return {
         id: m.id,
         role: m.role as 'user' | 'assistant' | 'system',
-        parts: rawParts,
+        parts: partsForApi,
+        metadata: storedSources.length ? { sources: storedSources } satisfies ChatMessageMetadata : undefined,
       }
     }) as unknown as UIMessage[]
     expandedSet.clear()
@@ -691,7 +683,9 @@ const displayMessages = computed<DisplayMessage[]>(() => {
       m.role === 'assistant' && !text.trim() && toolOutput != null && toolCalled
     const displayText = assistantNoTextButRetrieved ? t('chat.noModelReply') : text
 
-    const sources = normalizeToolSources(toolOutput?.sources)
+    const metaSources = (m.metadata as ChatMessageMetadata | undefined)?.sources
+    const sources = normalizeToolSources(toolOutput?.sources ?? metaSources)
+    const hadKbSearch = toolCalled || (metaSources?.length ?? 0) > 0
     // The assistant "cited" only if its text actually contains a [n] reference
     // matching one of the returned sources. Retrieving chunks alone isn't enough.
     const cited = m.role === 'assistant'
@@ -704,7 +698,7 @@ const displayMessages = computed<DisplayMessage[]>(() => {
       html: m.role === 'assistant' ? renderMarkdown(displayText) : '',
       sources,
       results: toolOutput?.results,
-      searched: m.role === 'assistant' ? (toolCalled ? true : (text ? false : null)) : null,
+      searched: m.role === 'assistant' ? (hadKbSearch ? true : (text ? false : null)) : null,
       cited,
       latencyMs: toolOutput?.latencyMs,
     }

server/utils/agent-messages.ts

@@ -0,0 +1,57 @@
+import type { UIMessage } from 'ai'
+
+const COMPLETED_TOOL_STATES = new Set([
+  'output-available',
+  'output-error',
+  'output-denied',
+  'output',
+])
+
+function isToolPart(part: { type: string }): boolean {
+  return part.type.startsWith('tool-') || part.type === 'dynamic-tool'
+}
+
+function isCompletedToolPart(part: { type: string; state?: string; toolCallId?: string }): boolean {
+  if (!isToolPart(part)) return false
+  const state = part.state
+  if (!state || !COMPLETED_TOOL_STATES.has(state)) return false
+  return typeof part.toolCallId === 'string' && part.toolCallId.length > 0
+}
+
+/**
+ * Normalize UI messages before convertToModelMessages / streamText.
+ *
+ * Incomplete tool parts (mid-stream) and synthetic tool parts rehydrated from DB
+ * (missing toolCallId) produce invalid ModelMessage[] and trigger:
+ * "The messages do not match the ModelMessage[] schema."
+ */
+export function sanitizeUIMessagesForAgent(messages: UIMessage[]): UIMessage[] {
+  const out: UIMessage[] = []
+
+  for (const message of messages) {
+    if (message.role === 'user' || message.role === 'system') {
+      out.push(message)
+      continue
+    }
+
+    if (message.role !== 'assistant') continue
+
+    const parts = (message.parts ?? []).filter((part) => {
+      if (part.type === 'text' || part.type === 'reasoning') return true
+      if (isToolPart(part)) return isCompletedToolPart(part as { type: string; state?: string; toolCallId?: string })
+      return false
+    })
+
+    const hasText = parts.some(
+      (p) => p.type === 'text' && (p as { text: string }).text.trim().length > 0,
+    )
+    const hasCompletedTool = parts.some((p) => isToolPart(p))
+
+    if (!hasText && !hasCompletedTool) continue
+    if (!parts.length) continue
+
+    out.push({ ...message, parts })
+  }
+
+  return out
+}

server/utils/agent.service.ts

@@ -24,6 +24,7 @@ import {
 } from './documents.service'
 import { truncate } from './text'
 import { parseMemoryCommand, getMessageText, type MemoryCommand } from './agent-commands'
+import { sanitizeUIMessagesForAgent } from './agent-messages'
 
 async function getLlmCfg() {
   const config = useRuntimeConfig()
@@ -234,26 +235,8 @@ export async function agentStreamText(
     else if (cfg.ragResponseLang === 'en') system += '\n- Respond always in English.'
   }
 
-  const validMessages = messages.filter((m) => {
-    if (m.role !== 'assistant') return true
-    const textParts = (m.parts ?? []).filter(
-      (p): p is { type: 'text'; text: string } => p.type === 'text',
-    )
-    const hasText = textParts.some((p) => p.text.trim().length > 0)
-    const hasCompletedTool = (m.parts ?? []).some((p) => {
-      if (!p.type.startsWith('tool-')) return false
-      const state = (p as { state?: string }).state
-      return (
-        state === 'output-available' ||
-        state === 'output-error' ||
-        state === 'output-denied' ||
-        state === 'output'
-      )
-    })
-    return hasText || hasCompletedTool
-  })
-
-  const modelMessages = await convertToModelMessages(validMessages)
+  const sanitizedMessages = sanitizeUIMessagesForAgent(messages)
+  const modelMessages = await convertToModelMessages(sanitizedMessages)
   const model = await getLlmModel(input.modelOverride)
 
   return streamText({

tests/agent.spec.ts

@@ -1,5 +1,6 @@
 import { describe, it, expect } from 'vitest'
 import { parseMemoryCommand, getMessageText } from '../server/utils/agent-commands'
+import { sanitizeUIMessagesForAgent } from '../server/utils/agent-messages'
 import type { UIMessage } from 'ai'
 
 describe('parseMemoryCommand', () => {
@@ -55,3 +56,60 @@ describe('getMessageText', () => {
     expect(getMessageText(m)).toBe('')
   })
 })
+
+describe('sanitizeUIMessagesForAgent', () => {
+  it('drops assistant messages with only in-progress tool parts', () => {
+    const messages = [
+      { id: '1', role: 'user', parts: [{ type: 'text', text: 'hi' }] },
+      {
+        id: '2',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'tool-searchKnowledgeBase',
+            state: 'input-streaming',
+            toolCallId: 'call-1',
+          },
+        ],
+      },
+    ] as unknown as UIMessage[]
+    expect(sanitizeUIMessagesForAgent(messages)).toEqual([messages[0]])
+  })
+
+  it('drops synthetic tool parts without toolCallId', () => {
+    const messages = [
+      { id: '1', role: 'user', parts: [{ type: 'text', text: 'q' }] },
+      {
+        id: '2',
+        role: 'assistant',
+        parts: [
+          { type: 'text', text: 'answer with [1]' },
+          {
+            type: 'tool-searchKnowledgeBase',
+            state: 'output-available',
+            output: { sources: [{ chunkId: 'c1' }] },
+          },
+        ],
+      },
+    ] as unknown as UIMessage[]
+    const out = sanitizeUIMessagesForAgent(messages)
+    expect(out).toHaveLength(2)
+    expect(out[1].parts).toEqual([{ type: 'text', text: 'answer with [1]' }])
+  })
+
+  it('keeps completed tool parts that include toolCallId', () => {
+    const toolPart = {
+      type: 'tool-searchKnowledgeBase',
+      state: 'output-available',
+      toolCallId: 'call-1',
+      input: { query: 'test' },
+      output: { context: 'ctx', sources: [], results: [], count: 0 },
+    }
+    const messages = [
+      { id: '1', role: 'user', parts: [{ type: 'text', text: 'q' }] },
+      { id: '2', role: 'assistant', parts: [toolPart] },
+    ] as unknown as UIMessage[]
+    const out = sanitizeUIMessagesForAgent(messages)
+    expect(out[1].parts).toContainEqual(toolPart)
+  })
+})