@pikku/core 0.12.20 → 0.12.22

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

@@ -23,11 +23,10 @@ import type {
 } from '../../services/ai-agent-runner-service.js'
 
 import {
-  parseWorkingMemory,
-  deepMergeWorkingMemory,
   resolveMemoryServices,
   loadContextMessages,
   trimMessages,
+  getWorkingMemoryMiddleware,
 } from './ai-agent-memory.js'
 import {
   prepareAgentRun,
@@ -58,6 +57,7 @@ function createPersistingChannel(
 ): PersistingChannel {
   let fullText = ''
   let stepText = ''
+  let stepGenerativeUI: unknown | null = null
   let stepToolCalls: AIToolCall[] = []
   let stepToolResults: AIToolResult[] = []
   const totalUsage: {
@@ -72,16 +72,24 @@ function createPersistingChannel(
   const flushStep = async () => {
     if (!storage) return
     const text = stepText
+    const generativeUI = stepGenerativeUI
     const calls = stepToolCalls
     const results = stepToolResults
     stepText = ''
+    stepGenerativeUI = null
     stepToolCalls = []
     stepToolResults = []
-    if (text || calls.length > 0) {
+    if (text || generativeUI != null || calls.length > 0) {
       const assistantMsg: AIMessage = {
         id: randomUUID(),
         role: 'assistant',
-        content: text || undefined,
+        content:
+          generativeUI != null
+            ? [
+                ...(text ? [{ type: 'text' as const, text }] : []),
+                { type: 'generative-ui' as const, spec: generativeUI },
+              ]
+            : text || undefined,
         toolCalls: calls.length > 0 ? calls : undefined,
         createdAt: new Date(),
       }
@@ -137,6 +145,9 @@ function createPersistingChannel(
                   : JSON.stringify(event.result),
             })
             break
+          case 'generative-ui':
+            stepGenerativeUI = event.spec
+            break
           case 'usage':
             totalUsage.inputTokens += event.tokens.input
             totalUsage.outputTokens += event.tokens.output
@@ -150,6 +161,9 @@ function createPersistingChannel(
       }
       parent.send(event)
     },
+    setState: (s) => parent.setState(s),
+    getState: () => parent.getState(),
+    clearState: () => parent.clearState(),
   }
   return channel
 }
@@ -158,10 +172,6 @@ async function postStreamCleanup(
   persistingChannel: PersistingChannel,
   aiMiddlewares: PikkuAIMiddlewareHooks[],
   singletonServices: any,
-  storage: AIStorageService | undefined,
-  memoryConfig: any,
-  threadId: string,
-  workingMemorySchemaName: string | null,
   messages: AIMessage[],
   aiRunState: AIRunStateService,
   runId: string
@@ -185,30 +195,6 @@ async function postStreamCleanup(
     }
   }
 
-  if (storage && memoryConfig?.workingMemory && outputText) {
-    const parsed = parseWorkingMemory(outputText)
-    if (parsed) {
-      const existing =
-        (await storage.getWorkingMemory(threadId, 'thread')) ?? {}
-      const merged = deepMergeWorkingMemory(existing, parsed)
-
-      if (singletonServices.schema && workingMemorySchemaName) {
-        try {
-          await singletonServices.schema.validateSchema(
-            workingMemorySchemaName,
-            merged
-          )
-        } catch (err) {
-          singletonServices.logger.warn(
-            `Working memory validation failed: ${err instanceof Error ? err.message : String(err)}`
-          )
-        }
-      }
-
-      await storage.saveWorkingMemory(threadId, 'thread', merged)
-    }
-  }
-
   await aiRunState.updateRun(runId, {
     status: 'completed',
     ...(usage.model
@@ -228,6 +214,7 @@ type StepLoopParams = {
   runnerParams: AIAgentRunnerParams
   maxSteps: number
   agentRunner: AIAgentRunnerService
+  streamChannel: AIStreamChannel
   persistingChannel: PersistingChannel
   channel: AIStreamChannel
   runId: string
@@ -247,7 +234,7 @@ async function runStreamStepLoop(
     runnerParams,
     maxSteps,
     agentRunner,
-    persistingChannel,
+    streamChannel,
     channel,
     runId,
     aiMiddlewares,
@@ -273,7 +260,7 @@ async function runStreamStepLoop(
 
     channel.send({ type: 'step-start', stepNumber: step })
 
-    const stepResult = await agentRunner.stream(runnerParams, persistingChannel)
+    const stepResult = await agentRunner.stream(runnerParams, streamChannel)
 
     for (const mw of aiMiddlewares) {
       if (mw.afterStep) {
@@ -433,10 +420,24 @@ export function appendStepMessages(
   runnerParams: AIAgentRunnerParams,
   stepResult: AIAgentStepResult
 ): void {
+  const structuredOutput =
+    stepResult.object && typeof stepResult.object === 'object'
+      ? (stepResult.object as Record<string, unknown>)
+      : null
+  const assistantContent =
+    structuredOutput?.ui != null
+      ? [
+          ...(stepResult.text
+            ? [{ type: 'text' as const, text: stepResult.text }]
+            : []),
+          { type: 'generative-ui' as const, spec: structuredOutput.ui },
+        ]
+      : stepResult.text || undefined
+
   const assistantMsg: AIMessage = {
     id: randomUUID(),
     role: 'assistant',
-    content: stepResult.text || undefined,
+    content: assistantContent,
     toolCalls:
       stepResult.toolCalls.length > 0
         ? stepResult.toolCalls.map((tc) => ({
@@ -445,6 +446,9 @@ export function appendStepMessages(
             args: tc.args as Record<string, unknown>,
           }))
         : undefined,
+    ...(stepResult.reasoningContent
+      ? { reasoningContent: stepResult.reasoningContent }
+      : {}),
     createdAt: new Date(),
   }
   runnerParams.messages.push(assistantMsg)
@@ -614,7 +618,15 @@ export async function streamAIAgent(
     return ''
   }
 
-  const aiMiddlewares: PikkuAIMiddlewareHooks[] = agent.aiMiddleware ?? []
+  const aiMiddlewares: PikkuAIMiddlewareHooks[] = [
+    ...getWorkingMemoryMiddleware(memoryConfig, storage, {
+      threadId,
+      workingMemorySchemaName,
+      logger: singletonServices.logger,
+      schemaService: singletonServices.schema,
+    }),
+    ...(agent.aiMiddleware ?? []),
+  ]
 
   let modifiedMessages = runnerParams.messages
   let modifiedInstructions = runnerParams.instructions
@@ -680,14 +692,16 @@ export async function streamAIAgent(
     }
   )
 
+  const persistingChannel = createPersistingChannel(channel, storage, threadId)
+
   const wrappedChannel =
     allChannelMiddleware.length > 0
       ? (wrapChannelWithMiddleware(
-          { channel },
+          { channel: persistingChannel },
           singletonServices,
           allChannelMiddleware
         ).channel as AIStreamChannel)
-      : channel
+      : persistingChannel
 
   // In delegate mode (default), suppress parent's text from reaching the client
   // AFTER a sub-agent has been called. If the parent responds directly (no delegation),
@@ -713,18 +727,13 @@ export async function streamAIAgent(
       }
     : wrappedChannel
 
-  const persistingChannel = createPersistingChannel(
-    outputChannel,
-    storage,
-    threadId
-  )
-
   try {
     const loopResult = await runStreamStepLoop({
       agent,
       runnerParams,
       maxSteps,
       agentRunner,
+      streamChannel: outputChannel,
       persistingChannel,
       channel,
       runId,
@@ -757,10 +766,6 @@ export async function streamAIAgent(
       persistingChannel,
       aiMiddlewares,
       singletonServices,
-      storage,
-      memoryConfig,
-      threadId,
-      workingMemorySchemaName,
       runnerParams.messages,
       aiRunState,
       runId
@@ -1090,7 +1095,15 @@ async function continueAfterToolResult(
 
   const instructions = await buildInstructions(resolvedName, packageName)
 
-  const aiMiddlewares: PikkuAIMiddlewareHooks[] = agent.aiMiddleware ?? []
+  const aiMiddlewares: PikkuAIMiddlewareHooks[] = [
+    ...getWorkingMemoryMiddleware(memoryConfig, storage, {
+      threadId: run.threadId,
+      workingMemorySchemaName,
+      logger: singletonServices.logger,
+      schemaService: singletonServices.schema,
+    }),
+    ...(agent.aiMiddleware ?? []),
+  ]
   let modifiedMessages = trimmedMessages
   let modifiedInstructions = instructions
   for (const mw of aiMiddlewares) {
@@ -1137,20 +1150,20 @@ async function continueAfterToolResult(
     }
   )
 
+  const persistingChannel = createPersistingChannel(
+    channel,
+    storage,
+    run.threadId
+  )
+
   const wrappedChannel =
     allChannelMiddleware.length > 0
       ? (wrapChannelWithMiddleware(
-          { channel },
+          { channel: persistingChannel },
           singletonServices,
           allChannelMiddleware
         ).channel as AIStreamChannel)
-      : channel
-
-  const persistingChannel = createPersistingChannel(
-    wrappedChannel,
-    storage,
-    run.threadId
-  )
+      : persistingChannel
 
   const streamContext: StreamContext = { channel, options }
   const resumeTools = (
@@ -1166,14 +1179,6 @@ async function continueAfterToolResult(
   ).tools
 
   const resolved = resolveModelConfig(resolvedName, agent)
-  console.log(
-    '[DEBUG resume] resolvedName:',
-    resolvedName,
-    'agent.model:',
-    agent.model,
-    'resolved.model:',
-    resolved.model
-  )
   const maxSteps = resolved.maxSteps ?? 10
 
   const runnerParams: AIAgentRunnerParams = {
@@ -1195,6 +1200,7 @@ async function continueAfterToolResult(
       runnerParams,
       maxSteps,
       agentRunner,
+      streamChannel: wrappedChannel,
       persistingChannel,
       channel,
       runId: run.runId,
@@ -1227,10 +1233,6 @@ async function continueAfterToolResult(
       persistingChannel,
       aiMiddlewares,
       singletonServices,
-      storage,
-      memoryConfig,
-      run.threadId,
-      workingMemorySchemaName,
       runnerParams.messages,
       aiRunState,
       run.runId

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

@@ -29,6 +29,15 @@ export type AIContentPart =
       mediaType: string
       filename?: string
     }
+  | {
+      type: 'data'
+      name: string
+      data: unknown
+    }
+  | {
+      type: 'generative-ui'
+      spec: unknown
+    }
 
 export interface AIToolCall {
   id: string
@@ -48,6 +57,7 @@ export interface AIMessage {
   content?: string | AIContentPart[]
   toolCalls?: AIToolCall[]
   toolResults?: AIToolResult[]
+  reasoningContent?: string
   createdAt: Date
 }
 
@@ -299,6 +309,19 @@ export type AIStreamEvent =
       agent?: string
       session?: string
     }
+  | {
+      type: 'data'
+      name: string
+      data: unknown
+      agent?: string
+      session?: string
+    }
+  | {
+      type: 'generative-ui'
+      spec: unknown
+      agent?: string
+      session?: string
+    }
   | {
       type: 'suspended'
       reason: 'rpc-missing'
@@ -367,6 +390,7 @@ export interface AgentRunRow {
 export interface AgentRunService {
   listThreads(options?: {
     agentName?: string
+    resourceId?: string
     limit?: number
     offset?: number
   }): Promise<AIThread[]>

package/src/wirings/channel/channel-handler.test.ts

@@ -0,0 +1,272 @@
+import { beforeEach, describe, test } from 'node:test'
+import assert from 'node:assert/strict'
+
+import { addFunction } from '../../function/function-runner.js'
+import { pikkuState, resetPikkuState } from '../../pikku-state.js'
+import { processMessageHandlers } from './channel-handler.js'
+
+const createLogger = () => {
+  const errors: string[] = []
+  return {
+    errors,
+    logger: {
+      debug: () => {},
+      info: () => {},
+      warn: () => {},
+      error: (...args: unknown[]) => {
+        errors.push(args.map((arg) => String(arg)).join(' '))
+      },
+    },
+  }
+}
+
+const createChannelHandler = () => {
+  const sent: unknown[] = []
+  const channel = {
+    channelId: 'channel-1',
+    openingData: {},
+    state: 'open' as const,
+    send: (message: unknown) => {
+      sent.push(message)
+    },
+    sendBinary: () => {},
+    close: () => {},
+    setState: () => {},
+    getState: () => undefined,
+    clearState: () => {},
+  }
+
+  return {
+    sent,
+    handler: {
+      send: (message: unknown) => channel.send(message),
+      sendBinary: () => {},
+      getChannel: () => channel,
+    },
+  }
+}
+
+const registerFunction = (
+  funcName: string,
+  func: (services: any, data: any, wire: any) => Promise<unknown> | unknown
+) => {
+  addFunction(funcName, { func } as never)
+  pikkuState(null, 'function', 'meta')[funcName] = {
+    name: funcName,
+    sessionless: true,
+    permissions: [],
+  } as never
+}
+
+beforeEach(() => {
+  resetPikkuState()
+})
+
+describe('processMessageHandlers', () => {
+  test('routes json messages using onMessageWiring and appends the routing key to the result', async () => {
+    const { logger } = createLogger()
+    const { handler } = createChannelHandler()
+
+    pikkuState(null, 'channel', 'meta').chat = {
+      name: 'chat',
+      route: '/chat',
+      input: null,
+      connect: null,
+      disconnect: null,
+      message: null,
+      messageWirings: {
+        kind: {
+          ping: {
+            pikkuFuncId: 'pingFunc',
+            packageName: null,
+          },
+        },
+      },
+    } as never
+    registerFunction('pingFunc', async (_services, data) => ({
+      reply: `pong:${data.value}`,
+    }))
+
+    const { onMessage } = processMessageHandlers(
+      { logger } as never,
+      {
+        name: 'chat',
+        route: '/chat',
+        auth: false,
+        onMessageWiring: {
+          kind: {
+            ping: {
+              func: { func: async () => ({ ok: true }) } as never,
+            },
+          },
+        },
+      } as never,
+      handler as never
+    )
+
+    const result = await onMessage('{"kind":"ping","value":"hello"}')
+
+    assert.deepEqual(result, {
+      reply: 'pong:hello',
+      kind: 'ping',
+    })
+  })
+
+  test('falls back to the default handler for parsed json with no matching route', async () => {
+    const { logger } = createLogger()
+    const { handler } = createChannelHandler()
+
+    pikkuState(null, 'channel', 'meta').chat = {
+      name: 'chat',
+      route: '/chat',
+      input: null,
+      connect: null,
+      disconnect: null,
+      message: {
+        pikkuFuncId: 'defaultFunc',
+        packageName: null,
+      },
+      messageWirings: {
+        kind: {
+          ping: {
+            pikkuFuncId: 'pingFunc',
+            packageName: null,
+          },
+        },
+      },
+    } as never
+    registerFunction('defaultFunc', async (_services, data) => ({
+      seen: data,
+    }))
+
+    const { onMessage } = processMessageHandlers(
+      { logger } as never,
+      {
+        name: 'chat',
+        route: '/chat',
+        auth: false,
+        onMessage: { func: async () => ({ ok: true }) } as never,
+        onMessageWiring: {
+          kind: {
+            ping: {
+              func: { func: async () => ({ ok: true }) } as never,
+            },
+          },
+        },
+      } as never,
+      handler as never
+    )
+
+    const result = await onMessage('{"kind":"unknown","value":"hello"}')
+
+    assert.deepEqual(result, {
+      seen: { kind: 'unknown', value: 'hello' },
+    })
+  })
+
+  test('uses wrapper auth for default messages and sends unauthorized errors without a session', async () => {
+    const { errors, logger } = createLogger()
+    const { sent, handler } = createChannelHandler()
+
+    pikkuState(null, 'channel', 'meta').chat = {
+      name: 'chat',
+      route: '/chat',
+      input: null,
+      connect: null,
+      disconnect: null,
+      message: {
+        pikkuFuncId: 'defaultFunc',
+        packageName: null,
+      },
+      messageWirings: {},
+    } as never
+    registerFunction('defaultFunc', async () => ({ ok: true }))
+
+    const { onMessage } = processMessageHandlers(
+      { logger } as never,
+      {
+        name: 'chat',
+        route: '/chat',
+        auth: false,
+        onMessage: {
+          auth: true,
+          func: { func: async () => ({ ok: true }) } as never,
+        },
+      } as never,
+      handler as never
+    )
+
+    const result = await onMessage('hello')
+
+    assert.equal(result, undefined)
+    assert.deepEqual(sent, ['Unauthorized for the default message route'])
+    assert.match(
+      errors[0] || '',
+      /requires a session for the default message route/
+    )
+  })
+
+  test('logs when no message handler can process the payload', async () => {
+    const { errors, logger } = createLogger()
+    const { handler } = createChannelHandler()
+
+    pikkuState(null, 'channel', 'meta').chat = {
+      name: 'chat',
+      route: '/chat',
+      input: null,
+      connect: null,
+      disconnect: null,
+      message: null,
+      messageWirings: {},
+    } as never
+
+    const { onMessage } = processMessageHandlers(
+      { logger } as never,
+      {
+        name: 'chat',
+        route: '/chat',
+        auth: false,
+      } as never,
+      handler as never
+    )
+
+    const result = await onMessage({ raw: true })
+
+    assert.equal(result, undefined)
+    assert.match(
+      errors[0] || '',
+      /No handler found for message in channel chat/
+    )
+    assert.match(errors[1] || '', /Channel config name: chat/)
+  })
+
+  test('blocks binary messages when a session is required and missing', async () => {
+    const { errors, logger } = createLogger()
+    const { handler } = createChannelHandler()
+
+    pikkuState(null, 'channel', 'meta').chat = {
+      name: 'chat',
+      route: '/chat',
+      input: null,
+      connect: null,
+      disconnect: null,
+      message: null,
+      messageWirings: {},
+    } as never
+
+    const { onBinaryMessage } = processMessageHandlers(
+      { logger } as never,
+      {
+        name: 'chat',
+        route: '/chat',
+        onBinaryMessage: async () => new Uint8Array([1, 2, 3]),
+      } as never,
+      handler as never
+    )
+
+    const result = await onBinaryMessage?.(new Uint8Array([1, 2, 3]))
+
+    assert.equal(result, undefined)
+    assert.match(errors[0] || '', /requires a session for binary message/)
+  })
+})

package/src/wirings/channel/channel-handler.ts

@@ -90,7 +90,7 @@ export const processMessageHandlers = (
         ? `route '${routingProperty}:${routerValue}'`
         : 'the default message route'
       logger.error(
-        `Channel ${channelConfig.name} with id ${channelHandler.getChannel().channelId} requires a session for ${routeMessage}`
+        `Channel ${channelConfig.name} with id ${channelHandler.getChannel().channelId} requires a session for ${routeMessage}. No session is attached to this websocket connection. Ensure auth middleware establishes a session during websocket upgrade, and configure sessionStore when the runtime needs persisted sessions.`
       )
       // TODO: Send error message back breaks typescript, but should be implemented somehow
       channelHandler.getChannel().send(`Unauthorized for ${routeMessage}`)