@pikku/core 0.12.80 → 0.12.82

package/src/wirings/ai-agent/ai-agent-stream-output-hooks.test.ts

@@ -0,0 +1,353 @@
+import { beforeEach, describe, test } from 'node:test'
+import assert from 'node:assert/strict'
+
+import { resetPikkuState, pikkuState } from '../../pikku-state.js'
+import { streamAIAgent } from './ai-agent-stream.js'
+import type { CoreAIAgent, PikkuAIMiddlewareHooks } from './ai-agent.types.js'
+import type { AIAgentStepResult } from '../../services/ai-agent-runner-service.js'
+
+beforeEach(() => {
+  resetPikkuState()
+})
+
+const addTestAgent = (agentName: string) => {
+  const agent: CoreAIAgent = {
+    name: agentName,
+    description: 'test agent',
+    instructions: 'be helpful',
+    model: 'test/test-model',
+  }
+
+  pikkuState(null, 'agent', 'agentsMeta')[agentName] = {
+    ...agent,
+    inputSchema: null,
+    outputSchema: null,
+    workingMemorySchema: null,
+  }
+  pikkuState(null, 'agent', 'agents').set(agentName, agent)
+}
+
+const makeStepResult = (
+  overrides?: Partial<AIAgentStepResult>
+): AIAgentStepResult => ({
+  text: '',
+  toolCalls: [],
+  toolResults: [],
+  usage: { inputTokens: 0, outputTokens: 0 },
+  finishReason: 'stop',
+  ...overrides,
+})
+
+describe('streamAIAgent output hooks', () => {
+  test('does not run modifyOutput on a streamed run, and warns the hook it is inert there', async () => {
+    addTestAgent('stream-modify-output-agent')
+
+    const warnings: unknown[][] = []
+    const modifyOutputCalls: unknown[] = []
+    const sideEffects: string[] = []
+
+    const middleware: PikkuAIMiddlewareHooks = {
+      modifyOutput: async (_services, ctx) => {
+        modifyOutputCalls.push(ctx)
+        sideEffects.push(ctx.text)
+        return { text: `${ctx.text} [redacted]`, messages: ctx.messages }
+      },
+    }
+    const agent = pikkuState(null, 'agent', 'agents').get(
+      'stream-modify-output-agent'
+    )!
+    agent.aiMiddleware = [middleware] as any
+    pikkuState(null, 'agent', 'agents').set('stream-modify-output-agent', agent)
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: (...args: unknown[]) => warnings.push(args),
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async (_params: any, channel: any) => {
+          channel.send({ type: 'text-delta', text: 'Hello' })
+          return makeStepResult({ text: 'Hello', finishReason: 'stop' })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-modify-output',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    const result = await streamAIAgent(
+      'stream-modify-output-agent',
+      {
+        message: 'hello',
+        threadId: 'thread-modify-output',
+        resourceId: 'resource-modify-output',
+      },
+      {
+        channelId: 'channel-modify-output',
+        openingData: undefined,
+        state: 'open',
+        send: () => {},
+        close: () => {},
+      },
+      {}
+    )
+
+    // It does not run at all — nothing on this path could act on what it
+    // returns, and the one hook that used to rely on the side effect (working
+    // memory) now persists from its own modifyOutputStream.
+    assert.equal(modifyOutputCalls.length, 0)
+    assert.deepEqual(sideEffects, [])
+    assert.equal(result, 'Hello')
+
+    // And the author of that hook has to be told, or the gap is silent.
+    assert.equal(
+      warnings.filter((args) =>
+        args.some(
+          (arg) =>
+            typeof arg === 'string' &&
+            arg.includes('modifyOutput') &&
+            arg.includes('stream-modify-output-agent')
+        )
+      ).length,
+      1
+    )
+  })
+
+  test('persists working memory from a streamed run', async () => {
+    addTestAgent('stream-working-memory-agent')
+
+    const savedWorkingMemory: unknown[] = []
+
+    const agent = pikkuState(null, 'agent', 'agents').get(
+      'stream-working-memory-agent'
+    )!
+    agent.memory = { workingMemory: true } as any
+    pikkuState(null, 'agent', 'agents').set(
+      'stream-working-memory-agent',
+      agent
+    )
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async (_params: any, channel: any) => {
+          channel.send({
+            type: 'text-delta',
+            text: 'Noted <working_memory>{"city":"Berlin"}</working_memory>',
+          })
+          return makeStepResult({ text: 'Noted', finishReason: 'stop' })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-working-memory',
+        updateRun: async () => {},
+      },
+      aiStorage: {
+        createThread: async () => {},
+        getMessages: async () => [],
+        saveMessages: async () => {},
+        getWorkingMemory: async () => ({}),
+        saveWorkingMemory: async (
+          threadId: string,
+          scope: string,
+          value: unknown
+        ) => {
+          savedWorkingMemory.push({ threadId, scope, value })
+        },
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await streamAIAgent(
+      'stream-working-memory-agent',
+      {
+        message: 'remember I live in Berlin',
+        threadId: 'thread-working-memory',
+        resourceId: 'resource-working-memory',
+      },
+      {
+        channelId: 'channel-working-memory',
+        openingData: undefined,
+        state: 'open',
+        send: () => {},
+        close: () => {},
+      },
+      {}
+    )
+
+    // The block never reaches modifyOutput on this path: the middleware's own
+    // stream hook strips it before the persisting channel accumulates the text.
+    // Persisting has to happen from the stream hook, where the raw text is.
+    assert.deepEqual(savedWorkingMemory, [
+      {
+        threadId: 'thread-working-memory',
+        scope: 'thread',
+        value: { city: 'Berlin' },
+      },
+    ])
+  })
+
+  test('a failing tool on a streamed run is persisted as a failure, not as text that reads like one', async () => {
+    addTestAgent('stream-tool-error-agent')
+
+    const savedMessages: any[] = []
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async (_params: any, channel: any) => {
+          channel.send({
+            type: 'tool-call',
+            toolCallId: 'call-1',
+            toolName: 'lookup',
+            args: { city: 'Berlin' },
+          })
+          channel.send({
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            toolName: 'lookup',
+            result: 'Error: upstream refused',
+            error: 'upstream refused',
+          })
+          channel.send({
+            type: 'tool-call',
+            toolCallId: 'call-2',
+            toolName: 'echo',
+            args: {},
+          })
+          channel.send({
+            type: 'tool-result',
+            toolCallId: 'call-2',
+            toolName: 'echo',
+            result: 'Error: this is just what the tool said',
+          })
+          return makeStepResult({ text: 'done', finishReason: 'stop' })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-tool-error',
+        updateRun: async () => {},
+      },
+      aiStorage: {
+        createThread: async () => {},
+        getMessages: async () => [],
+        saveMessages: async (_threadId: string, messages: any[]) => {
+          savedMessages.push(...messages)
+        },
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await streamAIAgent(
+      'stream-tool-error-agent',
+      {
+        message: 'look it up',
+        threadId: 'thread-tool-error',
+        resourceId: 'resource-tool-error',
+      },
+      {
+        channelId: 'channel-tool-error',
+        openingData: undefined,
+        state: 'open',
+        send: () => {},
+        close: () => {},
+      },
+      {}
+    )
+
+    const toolResults = savedMessages
+      .filter((message) => message.role === 'tool')
+      .flatMap((message) => message.toolResults ?? [])
+
+    assert.deepEqual(
+      toolResults.map((r: any) => [r.name, r.error]),
+      [
+        ['lookup', 'upstream refused'],
+        ['echo', undefined],
+      ]
+    )
+  })
+
+  test('does not warn about modifyOutput when the middleware also handles the stream', async () => {
+    addTestAgent('stream-both-hooks-agent')
+
+    const warnings: unknown[][] = []
+
+    const middleware: PikkuAIMiddlewareHooks = {
+      modifyOutput: async (_services, ctx) => ({
+        text: ctx.text,
+        messages: ctx.messages,
+      }),
+      modifyOutputStream: async (_services, ctx) => ctx.event,
+    }
+    const agent = pikkuState(null, 'agent', 'agents').get(
+      'stream-both-hooks-agent'
+    )!
+    agent.aiMiddleware = [middleware] as any
+    pikkuState(null, 'agent', 'agents').set('stream-both-hooks-agent', agent)
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: (...args: unknown[]) => warnings.push(args),
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async (_params: any, channel: any) => {
+          channel.send({ type: 'text-delta', text: 'Hi' })
+          return makeStepResult({ text: 'Hi', finishReason: 'stop' })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-both-hooks',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await streamAIAgent(
+      'stream-both-hooks-agent',
+      {
+        message: 'hello',
+        threadId: 'thread-both-hooks',
+        resourceId: 'resource-both-hooks',
+      },
+      {
+        channelId: 'channel-both-hooks',
+        openingData: undefined,
+        state: 'open',
+        send: () => {},
+        close: () => {},
+      },
+      {}
+    )
+
+    assert.deepEqual(
+      warnings.filter((args) =>
+        args.some(
+          (arg) => typeof arg === 'string' && arg.includes('modifyOutput')
+        )
+      ),
+      []
+    )
+  })
+})

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

@@ -1,6 +1,7 @@
 import type {
   AIStreamChannel,
   AIStreamEvent,
+  AIAgentStep,
   AIMessage,
   AIToolCall,
   AIToolResult,
@@ -9,6 +10,10 @@ import type {
   CoreAIAgent,
   AIAgentMemoryConfig,
 } from './ai-agent.types.js'
+import {
+  finalizeAgentRun,
+  lastUserMessageText,
+} from './ai-agent-finalize.js'
 import { pikkuState, getSingletonServices } from '../../pikku-state.js'
 import { applyInputMiddleware } from './ai-agent-turn.js'
 import { AIProviderNotConfiguredError } from '../../errors/errors.js'
@@ -68,6 +73,8 @@ type PersistingChannel = AIStreamChannel & {
   fullText: string
   flush: (opts?: { interrupted?: boolean }) => Promise<void>
   totalUsage: { inputTokens: number; outputTokens: number; model?: string }
+  /** Every tool the run called, kept for the whole run rather than per step. */
+  runToolCalls: NonNullable<AIAgentStep['toolCalls']>
 }
 
 function createPersistingChannel(
@@ -89,6 +96,9 @@ function createPersistingChannel(
     inputTokens: 0,
     outputTokens: 0,
   }
+  // Survives the per-step flush below, which clears its own buffers: the run
+  // record needs every call the run made, not just the last step's.
+  const runToolCalls: NonNullable<AIAgentStep['toolCalls']> = []
 
   const flushStep = async (opts?: { interrupted?: boolean }) => {
     if (!storage) return
@@ -137,6 +147,8 @@ function createPersistingChannel(
     })
   }
 
+  const runToolCallIndex = new Map<string, number>()
+
   const channel: PersistingChannel = {
     channelId: parent.channelId,
     openingData: parent.openingData,
@@ -149,6 +161,9 @@ function createPersistingChannel(
     get totalUsage() {
       return totalUsage
     },
+    get runToolCalls() {
+      return runToolCalls
+    },
     flush: flushStep,
     close: () => parent.close(),
     sendBinary: (data) => parent.sendBinary(data),
@@ -157,6 +172,26 @@ function createPersistingChannel(
       // the client was streamed, and an interrupted run has to be able to
       // report the fragment it got through even with persistence turned off.
       if (event.type === 'text-delta') fullText += event.text
+      if (event.type === 'tool-call') {
+        runToolCallIndex.set(event.toolCallId, runToolCalls.length)
+        runToolCalls.push({
+          name: event.toolName,
+          args: event.args as Record<string, unknown>,
+          result: '',
+        })
+      }
+      if (event.type === 'tool-result') {
+        const index = runToolCallIndex.get(event.toolCallId)
+        const result =
+          typeof event.result === 'string'
+            ? event.result
+            : JSON.stringify(event.result)
+        const call = index === undefined ? undefined : runToolCalls[index]
+        if (call) {
+          call.result = result
+          if (event.error) call.error = event.error
+        }
+      }
       if (storage) {
         switch (event.type) {
           case 'text-delta':
@@ -177,6 +212,7 @@ function createPersistingChannel(
                 typeof event.result === 'string'
                   ? event.result
                   : JSON.stringify(event.result),
+              ...(event.error ? { error: event.error } : {}),
             })
             break
           case 'generative-ui':
@@ -203,44 +239,67 @@ function createPersistingChannel(
   return channel
 }
 
+/**
+ * Agents already warned about, so a per-request hook does not become a
+ * per-request log line.
+ */
+const warnedUnstreamedOutputHooks = new Set<string>()
+
+/**
+ * `modifyOutput` does not run on a streamed run at all. Nothing here could act
+ * on what it returns — the text has already reached the client, and
+ * `createPersistingChannel` flushes each step to storage as it goes, so by the
+ * time the run ends the transcript is already written.
+ *
+ * Rewriting on this path belongs to `modifyOutputStream`, which genuinely
+ * works: the stream middleware wraps the persisting channel, so what is stored
+ * and accumulated is already what the client was sent. A middleware that
+ * rewrites in `modifyOutput` only — a redaction hook, typically — is therefore
+ * silently ineffective when the agent is streamed, and is told so once.
+ */
+const warnUnstreamedOutputHooks = (
+  agentName: string,
+  aiMiddlewares: PikkuAIMiddlewareHooks[],
+  logger?: { warn: (...args: any[]) => void }
+) => {
+  if (warnedUnstreamedOutputHooks.has(agentName)) return
+  const unstreamed = aiMiddlewares.some(
+    (mw) => mw.modifyOutput && !mw.modifyOutputStream
+  )
+  if (!unstreamed) return
+  warnedUnstreamedOutputHooks.add(agentName)
+  logger?.warn(
+    `Agent '${agentName}' has AI middleware with modifyOutput but no modifyOutputStream — modifyOutput does not apply to streamed runs. Implement modifyOutputStream to affect a streamed reply.`
+  )
+}
+
 async function postStreamCleanup(
   persistingChannel: PersistingChannel,
-  aiMiddlewares: PikkuAIMiddlewareHooks[],
-  singletonServices: any,
-  messages: AIMessage[],
   aiRunState: AIRunStateService,
-  runId: string
-): Promise<void> {
-  const usage = persistingChannel.totalUsage
-  let outputText = persistingChannel.fullText
-  let outputMessages = messages
-  for (let i = aiMiddlewares.length - 1; i >= 0; i--) {
-    const mw = aiMiddlewares[i]
-    if (mw.modifyOutput) {
-      const result = await mw.modifyOutput(singletonServices, {
-        text: outputText,
-        messages: outputMessages,
-        usage: {
-          inputTokens: usage.inputTokens,
-          outputTokens: usage.outputTokens,
-        },
-      })
-      outputText = result.text
-      outputMessages = result.messages
-    }
+  runId: string,
+  run: {
+    agentName: string
+    threadId: string
+    resourceId?: string
+    input: string
   }
-
-  await aiRunState.updateRun(runId, {
-    status: 'completed',
-    ...(usage.model
-      ? {
-          usage: {
-            inputTokens: usage.inputTokens,
-            outputTokens: usage.outputTokens,
-            model: usage.model,
-          },
-        }
-      : {}),
+): Promise<void> {
+  await finalizeAgentRun(aiRunState, {
+    runId,
+    agentName: run.agentName,
+    threadId: run.threadId,
+    resourceId: run.resourceId,
+    input: run.input,
+    // Already what the client received: the stream middleware wraps the
+    // persisting channel, so both were accumulated post-rewrite.
+    text: persistingChannel.fullText,
+    steps: [
+      {
+        usage: persistingChannel.totalUsage,
+        toolCalls: persistingChannel.runToolCalls,
+      },
+    ],
+    usage: persistingChannel.totalUsage,
   })
 }
 
@@ -723,6 +782,8 @@ export async function streamAIAgent(
     await storage.saveMessages(threadId, [persistedUserMessage])
   }
 
+  warnUnstreamedOutputHooks(agentName, aiMiddlewares, singletonServices.logger)
+
   const streamMiddleware = aiMiddlewares
     .filter((mw) => mw.modifyOutputStream)
     .map((mw) => {
@@ -849,14 +910,12 @@ export async function streamAIAgent(
       return persistingChannel.fullText
     }
 
-    await postStreamCleanup(
-      persistingChannel,
-      aiMiddlewares,
-      singletonServices,
-      runnerParams.messages,
-      aiRunState,
-      runId
-    )
+    await postStreamCleanup(persistingChannel, aiRunState, runId, {
+      agentName,
+      threadId,
+      resourceId: input.resourceId,
+      input: lastUserMessageText(runnerParams.messages),
+    })
 
     // knowledge: decisions/internals/the-agent-done-event-goes-through-the-middleware-and-is-awaited.md
     await outputChannel.send({ type: 'done' })
@@ -1184,16 +1243,17 @@ export async function resumeAIAgent(
       typeof pending.args === 'string' ? JSON.parse(pending.args) : pending.args
 
     let toolResult: unknown
-    let isError = false
+    let toolError: string | undefined
     try {
       toolResult = await matchingTool.execute(toolArgs)
     } catch (execErr: any) {
       if (execErr?.payload?.error === 'missing_credential') {
         toolResult = execErr.payload
+        toolError = 'missing_credential'
       } else {
-        toolResult = `Error: ${execErr instanceof Error ? execErr.message : String(execErr)}`
+        toolError = execErr instanceof Error ? execErr.message : String(execErr)
+        toolResult = `Error: ${toolError}`
       }
-      isError = true
     }
 
     const resultStr =
@@ -1220,7 +1280,7 @@ export async function resumeAIAgent(
       toolCallId: input.toolCallId,
       toolName: pending.toolName,
       result: toolResult,
-      ...(isError ? { isError: true } : {}),
+      ...(toolError ? { error: toolError } : {}),
     })
   }
 
@@ -1313,6 +1373,12 @@ async function continueAfterToolResult(
   // knowledge: decisions/internals/a-resumed-agent-turn-is-as-interruptible-as-the-first.md
   const interruptHandle = registerInterruptibleRun(run.runId)
 
+  warnUnstreamedOutputHooks(
+    run.agentName,
+    aiMiddlewares,
+    singletonServices.logger
+  )
+
   const streamMiddleware = aiMiddlewares
     .filter((mw) => mw.modifyOutputStream)
     .map((mw) => {
@@ -1434,14 +1500,10 @@ async function continueAfterToolResult(
       return
     }
 
-    await postStreamCleanup(
-      persistingChannel,
-      aiMiddlewares,
-      singletonServices,
-      runnerParams.messages,
-      aiRunState,
-      run.runId
-    )
+    await postStreamCleanup(persistingChannel, aiRunState, run.runId, {
+      ...run,
+      input: lastUserMessageText(runnerParams.messages),
+    })
 
     // knowledge: decisions/internals/the-agent-done-event-goes-through-the-middleware-and-is-awaited.md
     await wrappedChannel.send({ type: 'done' })

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

@@ -0,0 +1,67 @@
+import { describe, test } from 'node:test'
+import assert from 'node:assert/strict'
+
+import { toAccumulatedStep } from './ai-agent-turn.js'
+import type { AIAgentStepResult } from '../../services/ai-agent-runner-service.js'
+
+const stepResult = (
+  overrides?: Partial<AIAgentStepResult>
+): AIAgentStepResult => ({
+  text: '',
+  toolCalls: [],
+  toolResults: [],
+  usage: { inputTokens: 0, outputTokens: 0 },
+  finishReason: 'stop',
+  ...overrides,
+})
+
+describe('toAccumulatedStep', () => {
+  test('carries a tool failure as its own field, not only as rendered text', () => {
+    const step = toAccumulatedStep(
+      stepResult({
+        toolCalls: [
+          { toolCallId: 'call-1', toolName: 'lookupOrder', args: { id: 7 } },
+        ],
+        toolResults: [
+          {
+            toolCallId: 'call-1',
+            toolName: 'lookupOrder',
+            result: 'Error: order service unreachable',
+            error: 'order service unreachable',
+          },
+        ],
+      })
+    )
+
+    assert.deepEqual(step.toolCalls, [
+      {
+        name: 'lookupOrder',
+        args: { id: 7 },
+        result: 'Error: order service unreachable',
+        error: 'order service unreachable',
+      },
+    ])
+  })
+
+  test('leaves error unset on a tool that returned normally, even if it says Error', () => {
+    // A tool is allowed to return the word "Error" — which is exactly why
+    // "did this fail" cannot be answered by matching on the result text.
+    const step = toAccumulatedStep(
+      stepResult({
+        toolCalls: [
+          { toolCallId: 'call-1', toolName: 'searchLogs', args: { q: 'x' } },
+        ],
+        toolResults: [
+          {
+            toolCallId: 'call-1',
+            toolName: 'searchLogs',
+            result: 'Error: connection refused (1 match)',
+          },
+        ],
+      })
+    )
+
+    assert.equal(step.toolCalls[0].error, undefined)
+    assert.ok(!('error' in step.toolCalls[0]))
+  })
+})