@pikku/core 0.12.20 → 0.12.22

package/src/wirings/ai-agent/ai-agent-memory.test.ts

@@ -0,0 +1,324 @@
+import { beforeEach, describe, test } from 'node:test'
+import assert from 'node:assert/strict'
+
+import { resetPikkuState } from '../../pikku-state.js'
+import {
+  buildWorkingMemoryPrompt,
+  createWorkingMemoryMiddleware,
+  deepMergeWorkingMemory,
+  loadContextMessages,
+  parseWorkingMemory,
+  resolveMemoryServices,
+  saveMessages,
+  stripWorkingMemoryTags,
+  trimMessages,
+} from './ai-agent-memory.js'
+import type { AIMessage, CoreAIAgent } from './ai-agent.types.js'
+
+beforeEach(() => {
+  resetPikkuState()
+})
+
+const makeMessage = (overrides: Partial<AIMessage>): AIMessage => ({
+  id: 'msg',
+  role: 'user',
+  content: 'hello',
+  createdAt: new Date('2026-01-01T00:00:00.000Z'),
+  ...overrides,
+})
+
+describe('ai-agent-memory', () => {
+  test('resolveMemoryServices prefers named storage and falls back to aiStorage', () => {
+    const singletonServices = {
+      aiStorage: { kind: 'default' },
+      customStorage: { kind: 'custom' },
+    } as any
+
+    const defaultAgent = {} as CoreAIAgent
+    const namedAgent = { memory: { storage: 'customStorage' } } as CoreAIAgent
+
+    assert.equal(
+      resolveMemoryServices(defaultAgent, singletonServices).storage,
+      singletonServices.aiStorage
+    )
+    assert.equal(
+      resolveMemoryServices(namedAgent, singletonServices).storage,
+      singletonServices.customStorage
+    )
+  })
+
+  test('deepMergeWorkingMemory merges nested objects and deletes null fields', () => {
+    const result = deepMergeWorkingMemory(
+      {
+        profile: {
+          name: 'Yasser',
+          stats: { commits: 1, reviews: 2 },
+        },
+        stale: true,
+      },
+      {
+        profile: {
+          stats: { reviews: 3 },
+        },
+        stale: null,
+        active: false,
+      }
+    )
+
+    assert.deepEqual(result, {
+      profile: {
+        name: 'Yasser',
+        stats: { commits: 1, reviews: 3 },
+      },
+      active: false,
+    })
+  })
+
+  test('deepMergeWorkingMemory ignores prototype-polluting keys', () => {
+    const before = ({} as any).isAdmin
+    deepMergeWorkingMemory(
+      {},
+      JSON.parse('{"constructor":{"prototype":{"isAdmin":true}}}')
+    )
+    assert.equal(({} as any).isAdmin, before)
+
+    deepMergeWorkingMemory({}, JSON.parse('{"__proto__":{"isAdmin":true}}'))
+    assert.equal(({} as any).isAdmin, before)
+
+    deepMergeWorkingMemory({}, { prototype: { isAdmin: true } })
+    assert.equal(({} as any).isAdmin, before)
+  })
+
+  test('buildWorkingMemoryPrompt includes schema fields and current state', () => {
+    const prompt = buildWorkingMemoryPrompt(
+      { city: 'Berlin' },
+      {
+        properties: {
+          city: { type: 'string', description: 'Current city' },
+          age: { type: 'number' },
+        },
+      }
+    )
+
+    assert.match(prompt, /Working memory fields:/)
+    assert.match(prompt, /city \(string\) - Current city/)
+    assert.match(prompt, /age \(number\)/)
+    assert.match(prompt, /Current working memory:\n\{"city":"Berlin"\}/)
+    assert.match(prompt, /<working_memory>/)
+  })
+
+  test('buildWorkingMemoryPrompt marks empty memory explicitly', () => {
+    const prompt = buildWorkingMemoryPrompt(null)
+    assert.match(prompt, /Current working memory: \(empty\)/)
+  })
+
+  test('loadContextMessages injects working memory system prompt when enabled', async () => {
+    const storage = {
+      getWorkingMemory: async () => ({ city: 'Berlin' }),
+    } as any
+
+    const messages = await loadContextMessages(
+      { workingMemory: true },
+      storage,
+      { message: 'hello', threadId: 'thread-1', resourceId: 'resource-1' },
+      { properties: { city: { type: 'string' } } }
+    )
+
+    assert.equal(messages.length, 1)
+    assert.equal(messages[0].role, 'system')
+    assert.match(String(messages[0].content), /Current working memory:/)
+  })
+
+  test('loadContextMessages returns no messages without storage', async () => {
+    const messages = await loadContextMessages(
+      { workingMemory: true },
+      undefined,
+      { message: 'hello', threadId: 'thread-1', resourceId: 'resource-1' }
+    )
+    assert.deepEqual(messages, [])
+  })
+
+  test('saveMessages persists tool calls and strips working memory tags from returned text', async () => {
+    const savedMessages: AIMessage[][] = []
+
+    const storage = {
+      saveMessages: async (_threadId: string, messages: AIMessage[]) => {
+        savedMessages.push(messages)
+      },
+    } as any
+
+    const resultText = await saveMessages(
+      storage,
+      'thread-1',
+      'resource-1',
+      { workingMemory: true },
+      makeMessage({ id: 'user-1' }),
+      {
+        text: 'Done <working_memory>{"city":"Berlin"}</working_memory>',
+        steps: [
+          {
+            toolCalls: [
+              {
+                name: 'tool-1',
+                args: { q: 'hello' },
+                result: 'world',
+              },
+            ],
+          },
+        ],
+      }
+    )
+
+    assert.equal(resultText, 'Done')
+    assert.equal(savedMessages.length, 1)
+    assert.equal(savedMessages[0].length, 4)
+    assert.equal(savedMessages[0][0].role, 'user')
+    assert.equal(savedMessages[0][1].role, 'assistant')
+    assert.equal(savedMessages[0][2].role, 'tool')
+    assert.equal(savedMessages[0][3].role, 'assistant')
+  })
+
+  test('saveMessages returns text unchanged when no storage is configured', async () => {
+    const passthrough = await saveMessages(
+      undefined,
+      'thread-1',
+      'resource-1',
+      { workingMemory: true },
+      null,
+      { text: 'plain text', steps: [] }
+    )
+    assert.equal(passthrough, 'plain text')
+  })
+
+  test('createWorkingMemoryMiddleware merges, validates, and persists working memory', async () => {
+    const savedWorkingMemory: unknown[] = []
+    const warnings: string[] = []
+
+    const storage = {
+      getWorkingMemory: async () => ({ city: 'Paris', removeMe: 'x' }),
+      saveWorkingMemory: async (
+        threadId: string,
+        scope: string,
+        value: unknown
+      ) => {
+        savedWorkingMemory.push({ threadId, scope, value })
+      },
+    } as any
+
+    const mw = createWorkingMemoryMiddleware({
+      storage,
+      threadId: 'thread-1',
+      workingMemorySchemaName: 'WorkingMemory',
+      schemaService: { validateSchema: async () => {} } as any,
+      logger: { warn: (message: string) => warnings.push(message) } as any,
+    })
+
+    const result = (await mw.modifyOutput!({} as any, {
+      text: 'Done <working_memory>{"city":"Berlin","removeMe":null}</working_memory>',
+      messages: [],
+      usage: {} as any,
+    })) as { text: string }
+
+    assert.equal(result.text, 'Done')
+    assert.deepEqual(savedWorkingMemory, [
+      {
+        threadId: 'thread-1',
+        scope: 'thread',
+        value: { city: 'Berlin' },
+      },
+    ])
+    assert.deepEqual(warnings, [])
+  })
+
+  test('createWorkingMemoryMiddleware warns and skips persistence on schema failure', async () => {
+    const savedWorkingMemory: unknown[] = []
+    const warnings: string[] = []
+
+    const storage = {
+      getWorkingMemory: async () => ({}),
+      saveWorkingMemory: async (...args: unknown[]) => {
+        savedWorkingMemory.push(args)
+      },
+    } as any
+
+    const mw = createWorkingMemoryMiddleware({
+      storage,
+      threadId: 'thread-1',
+      workingMemorySchemaName: 'WorkingMemory',
+      schemaService: {
+        validateSchema: async () => {
+          throw new Error('bad schema')
+        },
+      } as any,
+      logger: { warn: (message: string) => warnings.push(message) } as any,
+    })
+
+    const result = (await mw.modifyOutput!({} as any, {
+      text: '<working_memory>{"city":"Berlin"}</working_memory>',
+      messages: [],
+      usage: {} as any,
+    })) as { text: string }
+
+    assert.equal(result.text, '')
+    assert.deepEqual(warnings, ['Working memory validation failed: bad schema'])
+    assert.deepEqual(savedWorkingMemory, [])
+  })
+
+  test('trimMessages sanitizes orphaned tool messages and keeps first user/system boundary', () => {
+    const messages: AIMessage[] = [
+      makeMessage({
+        id: 'assistant-1',
+        role: 'assistant',
+        content: 'tool call text',
+        toolCalls: [{ id: 'tc-1', name: 'tool', args: { x: 1 } }],
+      }),
+      makeMessage({
+        id: 'tool-1',
+        role: 'tool',
+        toolResults: [{ id: 'other', name: 'tool', result: 'wrong' }],
+      }),
+      makeMessage({
+        id: 'user-1',
+        role: 'user',
+        content: 'real user message',
+      }),
+      makeMessage({
+        id: 'assistant-2',
+        role: 'assistant',
+        content: 'final answer',
+      }),
+    ]
+
+    const trimmed = trimMessages(messages, 1000)
+    assert.equal(trimmed[0].role, 'user')
+    assert.equal(trimmed[1].role, 'assistant')
+    assert.equal(trimmed[1].content, 'final answer')
+  })
+
+  test('trimMessages keeps latest message even when over token budget', () => {
+    const trimmed = trimMessages(
+      [makeMessage({ content: 'a'.repeat(1000) })],
+      1
+    )
+    assert.equal(trimmed.length, 1)
+  })
+
+  test('parseWorkingMemory and stripWorkingMemoryTags handle valid and invalid payloads', () => {
+    assert.deepEqual(
+      parseWorkingMemory(
+        'a <working_memory>{"city":"Berlin"}</working_memory> b'
+      ),
+      { city: 'Berlin' }
+    )
+    assert.equal(
+      parseWorkingMemory('<working_memory>{oops}</working_memory>'),
+      null
+    )
+    assert.equal(
+      stripWorkingMemoryTags(
+        'before <working_memory>{"city":"Berlin"}</working_memory> after'
+      ),
+      'before  after'.trim()
+    )
+  })
+})

package/src/wirings/ai-agent/ai-agent-memory.ts

@@ -9,6 +9,7 @@ import type {
 import type { AIStorageService } from '../../services/ai-storage-service.js'
 import type { Logger } from '../../services/logger.js'
 import type { SchemaService } from '../../services/schema-service.js'
+import type { PikkuAIMiddlewareHooks } from './ai-agent.types.js'
 
 export function resolveMemoryServices(
   agent: CoreAIAgent,
@@ -23,12 +24,22 @@ export function resolveMemoryServices(
   return { storage }
 }
 
+export function isWorkingMemoryEnabled(
+  memoryConfig: AIAgentMemoryConfig | undefined,
+  storage: AIStorageService | undefined
+): boolean {
+  return !!memoryConfig?.workingMemory && !!storage
+}
+
 export function deepMergeWorkingMemory(
   existing: Record<string, unknown>,
   updates: Record<string, unknown>
 ): Record<string, unknown> {
   const result = { ...existing }
   for (const key of Object.keys(updates)) {
+    if (key === '__proto__' || key === 'constructor' || key === 'prototype') {
+      continue
+    }
     const value = updates[key]
     if (value === null) {
       delete result[key]
@@ -79,7 +90,9 @@ export function buildWorkingMemoryPrompt(
 
   parts.push(
     'When you learn new information, output a partial JSON update in <working_memory> tags. ' +
-      'Only include changed fields. Set a field to null to delete it.'
+      'Only include durable facts you have actually derived or the user has confirmed. ' +
+      'Do not output templates, placeholders, or narration. ' +
+      'Only include changed fields. Leave unknown fields untouched. Set a field to null to delete it.'
   )
 
   return parts.join('\n\n')
@@ -93,8 +106,15 @@ export async function loadContextMessages(
 ): Promise<AIMessage[]> {
   const contextMessages: AIMessage[] = []
 
-  if (memoryConfig?.workingMemory && storage) {
-    const workingMem = await storage.getWorkingMemory(input.threadId, 'thread')
+  const workingMemoryStorage = isWorkingMemoryEnabled(memoryConfig, storage)
+    ? storage
+    : undefined
+
+  if (workingMemoryStorage) {
+    const workingMem = await workingMemoryStorage.getWorkingMemory(
+      input.threadId,
+      'thread'
+    )
     const prompt = buildWorkingMemoryPrompt(workingMem, workingMemoryJsonSchema)
     contextMessages.push({
       id: `wm-${randomUUID()}`,
@@ -115,6 +135,7 @@ export async function saveMessages(
   userMessage: AIMessage | null | undefined,
   result: {
     text: string
+    uiSpec?: unknown
     steps: {
       toolCalls?: {
         name: string
@@ -122,15 +143,11 @@ export async function saveMessages(
         result: string
       }[]
     }[]
-  },
-  options?: {
-    workingMemoryJsonSchema?: Record<string, unknown>
-    workingMemorySchemaName?: string | null
-    logger?: Logger
-    schemaService?: SchemaService
   }
 ): Promise<string> {
-  let responseText = result.text
+  const responseText = memoryConfig?.workingMemory
+    ? extractWorkingMemory(result.text).cleanedText
+    : result.text
 
   if (storage) {
     const newMessages: AIMessage[] = userMessage ? [userMessage] : []
@@ -161,39 +178,24 @@ export async function saveMessages(
       }
     }
 
+    const assistantContent =
+      result.uiSpec != null
+        ? [
+            ...(responseText
+              ? [{ type: 'text' as const, text: responseText }]
+              : []),
+            { type: 'generative-ui' as const, spec: result.uiSpec },
+          ]
+        : responseText
+
     newMessages.push({
       id: randomUUID(),
       role: 'assistant',
-      content: responseText,
+      content: assistantContent || undefined,
       createdAt: new Date(),
     })
 
     await storage.saveMessages(threadId, newMessages)
-
-    if (memoryConfig?.workingMemory) {
-      const parsed = parseWorkingMemory(responseText)
-      if (parsed) {
-        const existing =
-          (await storage.getWorkingMemory(threadId, 'thread')) ?? {}
-        const merged = deepMergeWorkingMemory(existing, parsed)
-
-        if (options?.schemaService && options?.workingMemorySchemaName) {
-          try {
-            await options.schemaService.validateSchema(
-              options.workingMemorySchemaName,
-              merged
-            )
-          } catch (err) {
-            options.logger?.warn(
-              `Working memory validation failed: ${err instanceof Error ? err.message : String(err)}`
-            )
-          }
-        }
-
-        await storage.saveWorkingMemory(threadId, 'thread', merged)
-        responseText = stripWorkingMemoryTags(responseText)
-      }
-    }
   }
 
   return responseText
@@ -308,3 +310,152 @@ export function parseWorkingMemory(
 export function stripWorkingMemoryTags(text: string): string {
   return text.replace(/<working_memory>[\s\S]*?<\/working_memory>/g, '').trim()
 }
+
+export function extractWorkingMemory(text: string): {
+  workingMemory: Record<string, unknown> | null
+  cleanedText: string
+} {
+  const workingMemory = parseWorkingMemory(text)
+  return {
+    workingMemory,
+    cleanedText: workingMemory ? stripWorkingMemoryTags(text) : text,
+  }
+}
+
+function getPendingWorkingMemoryPrefixLength(text: string): number {
+  const prefixes = ['<working_memory>', '</working_memory>']
+  const lastLt = text.lastIndexOf('<')
+  if (lastLt === -1) return 0
+
+  const suffix = text.slice(lastLt)
+  for (const prefix of prefixes) {
+    if (prefix.startsWith(suffix)) {
+      return suffix.length
+    }
+  }
+
+  return 0
+}
+
+export function stripWorkingMemoryForStreaming(text: string): string {
+  let output = ''
+  let cursor = 0
+
+  while (cursor < text.length) {
+    const openIndex = text.indexOf('<working_memory>', cursor)
+    if (openIndex === -1) {
+      const tail = text.slice(cursor)
+      const pendingPrefixLength = getPendingWorkingMemoryPrefixLength(tail)
+      output +=
+        pendingPrefixLength > 0 ? tail.slice(0, -pendingPrefixLength) : tail
+      return output
+    }
+
+    output += text.slice(cursor, openIndex)
+    const closeIndex = text.indexOf('</working_memory>', openIndex)
+    if (closeIndex === -1) {
+      return output
+    }
+
+    cursor = closeIndex + '</working_memory>'.length
+  }
+
+  return output
+}
+
+export function createWorkingMemoryMiddleware(options: {
+  storage?: AIStorageService
+  threadId: string
+  workingMemorySchemaName?: string | null
+  logger?: Logger
+  schemaService?: SchemaService
+}): PikkuAIMiddlewareHooks<{
+  rawText?: string
+  emittedVisibleText?: string
+}> {
+  return {
+    modifyOutputStream: async (_services, { event, state }) => {
+      if (event.type !== 'text-delta') return event
+
+      const rawText = `${state.rawText ?? ''}${event.text}`
+      state.rawText = rawText
+
+      const visibleText = stripWorkingMemoryForStreaming(rawText)
+      const emittedVisibleText = state.emittedVisibleText ?? ''
+
+      if (!visibleText.startsWith(emittedVisibleText)) {
+        state.emittedVisibleText = visibleText
+        return { ...event, text: visibleText }
+      }
+
+      const delta = visibleText.slice(emittedVisibleText.length)
+      state.emittedVisibleText = visibleText
+      if (!delta) return null
+      return { ...event, text: delta }
+    },
+    modifyOutput: async (_services, { text, messages, usage }) => {
+      const { workingMemory, cleanedText } = extractWorkingMemory(text)
+      if (workingMemory && options.storage) {
+        const existing =
+          (await options.storage.getWorkingMemory(
+            options.threadId,
+            'thread'
+          )) ?? {}
+        const merged = deepMergeWorkingMemory(existing, workingMemory)
+
+        let valid = true
+        if (options.schemaService && options.workingMemorySchemaName) {
+          try {
+            await options.schemaService.validateSchema(
+              options.workingMemorySchemaName,
+              merged
+            )
+          } catch (err) {
+            valid = false
+            options.logger?.warn(
+              `Working memory validation failed: ${err instanceof Error ? err.message : String(err)}`
+            )
+          }
+        }
+
+        // Only persist when the merged value passes schema validation —
+        // saving invalid data would poison subsequent getWorkingMemory reads.
+        if (valid) {
+          await options.storage.saveWorkingMemory(
+            options.threadId,
+            'thread',
+            merged
+          )
+        }
+      }
+
+      return {
+        text: cleanedText,
+        messages,
+      }
+    },
+  }
+}
+
+export function getWorkingMemoryMiddleware(
+  memoryConfig: AIAgentMemoryConfig | undefined,
+  storage: AIStorageService | undefined,
+  options: {
+    threadId: string
+    workingMemorySchemaName?: string | null
+    logger?: Logger
+    schemaService?: SchemaService
+  }
+): PikkuAIMiddlewareHooks[] {
+  if (!isWorkingMemoryEnabled(memoryConfig, storage)) return []
+
+  return [
+    createWorkingMemoryMiddleware({
+      storage,
+      threadId: options.threadId,
+      workingMemorySchemaName: options.workingMemorySchemaName,
+      logger: options.logger,
+      schemaService: options.schemaService,
+    }),
+  ]
+}