fix(stage-ui): improve LLM KV-cache compatibility for XML context injection (#1593)

2026-04-08 23:24:53 +08:00
parent 7b1041fa8e
commit d212da146a
4 changed files with 192 additions and 17 deletions
@@ -13,7 +13,7 @@ import { ref, toRaw } from 'vue'
 import { useAnalytics } from '../composables'
 import { useLlmmarkerParser } from '../composables/llm-marker-parser'
 import { categorizeResponse, createStreamingCategorizer } from '../composables/response-categoriser'
-import { buildContextPromptMessage } from './chat/context-prompt'
+import { formatContextPromptText } from './chat/context-prompt'
 import { createDatetimeContext, createMinecraftContext } from './chat/context-providers'
 import { useChatContextStore } from './chat/context-store'
 import { createChatHooks } from './chat/hooks'
@@ -279,7 +279,7 @@ export const useChatOrchestratorStore = defineStore('chat-orchestrator', () => {
        ],
      })

-      let newMessages = sessionMessagesForSend.map((msg) => {
+      const newMessages = sessionMessagesForSend.map((msg) => {
        const { context: _context, id: _id, createdAt: _createdAt, ...withoutContext } = msg
        const rawMessage = toRaw(withoutContext)

@@ -292,16 +292,28 @@ export const useChatOrchestratorStore = defineStore('chat-orchestrator', () => {
      })

      const contextsSnapshot = chatContext.getContextsSnapshot()
-      const contextPromptMessage = buildContextPromptMessage(contextsSnapshot)
-      if (contextPromptMessage) {
-        const system = newMessages.slice(0, 1)
-        const afterSystem = newMessages.slice(1, newMessages.length)
+      const contextPromptText = formatContextPromptText(contextsSnapshot)
+      if (contextPromptText) {
+        // Merge context into the latest user message instead of inserting a
+        // separate user message, which would create consecutive same-role
+        // messages forbidden by some providers (e.g. Anthropic → 400 error).
+        // Appending at the end keeps the static history prefix stable for
+        // LLM KV-cache reuse.
+        // See: https://github.com/moeru-ai/airi/issues/1539
+        const lastMessage = newMessages.at(-1)
+        if (lastMessage && lastMessage.role === 'user') {
+          // Append context after the user's content, separated by a newline.
+          // Keeping it at the end of the last message preserves the static
+          // history prefix for LLM KV-cache reuse.
+          const existingParts = typeof lastMessage.content === 'string'
+            ? [{ type: 'text' as const, text: lastMessage.content }]
+            : lastMessage.content

-        newMessages = [
-          ...system,
-          contextPromptMessage,
-          ...afterSystem,
-        ]
+          lastMessage.content = [
+            ...existingParts,
+            { type: 'text' as const, text: `\n${contextPromptText}` },
+          ]
+        }

        contextObservability.recordLifecycle({
          phase: 'prompt-context-built',
@@ -309,7 +321,7 @@ export const useChatOrchestratorStore = defineStore('chat-orchestrator', () => {
          sessionId,
          details: {
            contexts: contextsSnapshot,
-            promptMessage: contextPromptMessage,
+            promptText: contextPromptText,
          },
        })
      }
@@ -319,7 +331,7 @@ export const useChatOrchestratorStore = defineStore('chat-orchestrator', () => {
        sessionId,
        message: sendingMessage,
        contexts: contextsSnapshot,
-        promptMessage: contextPromptMessage,
+        promptMessage: undefined,
        composedMessage: newMessages as Message[],
      })
      contextObservability.recordLifecycle({
@@ -0,0 +1,104 @@
+import type { ContextSnapshot } from './context-prompt'
+
+import { ContextUpdateStrategy } from '@proj-airi/server-sdk'
+import { describe, expect, it } from 'vitest'
+
+import { buildContextPromptMessage, formatContextPromptText } from './context-prompt'
+
+function makeContext(overrides: Record<string, unknown> = {}): ContextSnapshot {
+  return {
+    'system:datetime': [
+      {
+        id: 'volatile-random-id',
+        contextId: 'system:datetime',
+        strategy: ContextUpdateStrategy.ReplaceSelf,
+        text: 'Current datetime: 2026-04-07T12:34:00.000Z',
+        createdAt: 1743940440000,
+        metadata: {
+          source: {
+            id: 'system:datetime',
+            kind: 'plugin' as const,
+            plugin: { id: 'airi:system:datetime' },
+          },
+        },
+        ...overrides,
+      },
+    ],
+  }
+}
+
+describe('formatContextPromptText', () => {
+  it('returns empty string for empty snapshot', () => {
+    expect(formatContextPromptText({})).toBe('')
+  })
+
+  // https://github.com/moeru-ai/airi/issues/1539
+  it('issue #1539: excludes id, createdAt, and metadata from serialized output', () => {
+    const text = formatContextPromptText(makeContext())
+
+    expect(text).not.toContain('volatile-random-id')
+    expect(text).not.toContain('1743940440000')
+    expect(text).not.toContain('airi:system:datetime')
+  })
+
+  // https://github.com/moeru-ai/airi/issues/1539
+  it('issue #1539: only includes text content in XML format', () => {
+    const text = formatContextPromptText(makeContext())
+
+    expect(text).toContain('<context>')
+    expect(text).toContain('</context>')
+    expect(text).toContain('<module name="system:datetime">')
+    expect(text).toContain('Current datetime: 2026-04-07T12:34:00.000Z')
+  })
+
+  // https://github.com/moeru-ai/airi/issues/1539
+  it('issue #1539: produces identical output regardless of volatile fields', () => {
+    const a = formatContextPromptText(makeContext({ id: 'aaa', createdAt: 1 }))
+    const b = formatContextPromptText(makeContext({ id: 'bbb', createdAt: 2 }))
+
+    expect(a).toBe(b)
+  })
+
+  it('formats multiple modules', () => {
+    const snapshot: ContextSnapshot = {
+      'system:datetime': [
+        {
+          id: 'a',
+          contextId: 'system:datetime',
+          strategy: ContextUpdateStrategy.ReplaceSelf,
+          text: 'Current datetime: 2026-04-07T12:34:00.000Z',
+          createdAt: 0,
+        },
+      ],
+      'system:minecraft': [
+        {
+          id: 'b',
+          contextId: 'system:minecraft',
+          strategy: ContextUpdateStrategy.ReplaceSelf,
+          text: 'Bot is online',
+          createdAt: 0,
+        },
+      ],
+    }
+
+    const text = formatContextPromptText(snapshot)
+
+    expect(text).toContain('<module name="system:datetime">')
+    expect(text).toContain('<module name="system:minecraft">')
+    expect(text).toContain('Bot is online')
+  })
+})
+
+describe('buildContextPromptMessage', () => {
+  it('returns null for empty snapshot', () => {
+    expect(buildContextPromptMessage({})).toBeNull()
+  })
+
+  it('returns a user message with context text', () => {
+    const msg = buildContextPromptMessage(makeContext())
+
+    expect(msg).not.toBeNull()
+    expect(msg!.role).toBe('user')
+    expect(msg!.content).toBeInstanceOf(Array)
+  })
+})
@@ -2,15 +2,32 @@ import type { UserMessage } from '@xsai/shared-chat'

 import type { ContextMessage } from '../../types/chat'

+import { toXml } from 'xast-util-to-xml'
+import { x } from 'xastscript'
+
 export type ContextSnapshot = Record<string, ContextMessage[]>

+/**
+ * Build an xast tree from context snapshot.
+ * Only the `text` field is included — volatile metadata (random IDs,
+ * millisecond timestamps) is excluded to keep the output deterministic
+ * and friendly to LLM KV-cache prefix matching.
+ * See: https://github.com/moeru-ai/airi/issues/1539
+ */
+function buildContextTree(contextsSnapshot: ContextSnapshot) {
+  const modules = Object.entries(contextsSnapshot).map(([key, messages]) =>
+    x('module', { name: key }, messages.map(m => x(null, m.text))),
+  )
+
+  return x('context', modules)
+}
+
 export function formatContextPromptText(contextsSnapshot: ContextSnapshot) {
-  if (Object.keys(contextsSnapshot).length === 0)
+  const entries = Object.entries(contextsSnapshot)
+  if (entries.length === 0)
    return ''

-  return ''
-    + 'These are the contextual information retrieved or on-demand updated from other modules, you may use them as context for chat, or reference of the next action, tool call, etc.:\n'
-    + `${Object.entries(contextsSnapshot).map(([key, value]) => `Module ${key}: ${JSON.stringify(value)}`).join('\n')}\n`
+  return toXml(buildContextTree(contextsSnapshot))
 }

 export function buildContextPromptMessage(contextsSnapshot: ContextSnapshot): UserMessage | null {
@@ -0,0 +1,42 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import { createDatetimeContext } from './datetime'
+
+describe('createDatetimeContext', () => {
+  it('returns a context message with datetime text', () => {
+    const ctx = createDatetimeContext()
+
+    expect(ctx.contextId).toBe('system:datetime')
+    expect(ctx.text).toContain('Current datetime:')
+    expect(ctx.strategy).toBe('replace-self')
+  })
+
+  it('includes ISO string in text', () => {
+    vi.useFakeTimers()
+    try {
+      vi.setSystemTime(new Date('2026-04-07T12:34:56.789Z'))
+      const ctx = createDatetimeContext()
+
+      expect(ctx.text).toContain('2026-04-07T12:34:56.789Z')
+    }
+    finally {
+      vi.useRealTimers()
+    }
+  })
+
+  it('produces different text at different times', () => {
+    vi.useFakeTimers()
+    try {
+      vi.setSystemTime(new Date('2026-04-07T12:34:00.000Z'))
+      const a = createDatetimeContext()
+
+      vi.setSystemTime(new Date('2026-04-07T12:35:00.000Z'))
+      const b = createDatetimeContext()
+
+      expect(a.text).not.toBe(b.text)
+    }
+    finally {
+      vi.useRealTimers()
+    }
+  })
+})