@pikku/core 0.12.0 → 0.12.1

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

@@ -9,13 +9,17 @@ import type {
   CoreAIAgent,
   AIAgentMemoryConfig,
 } from './ai-agent.types.js'
-import { pikkuState } from '../../pikku-state.js'
+import { pikkuState, getSingletonServices } from '../../pikku-state.js'
 import {
   combineChannelMiddleware,
   wrapChannelWithMiddleware,
 } from '../channel/channel-middleware-runner.js'
 import { randomUUID } from 'crypto'
 import type { AIStorageService } from '../../services/ai-storage-service.js'
+import type {
+  AIAgentRunnerParams,
+  AIAgentStepResult,
+} from '../../services/ai-agent-runner-service.js'
 
 import {
   parseWorkingMemory,
@@ -35,6 +39,10 @@ import {
   type StreamAIAgentOptions,
   type StreamContext,
 } from './ai-agent-prepare.js'
+import {
+  createAssistantUIChannel,
+  parseAssistantUIInput,
+} from './ai-agent-assistant-ui.js'
 import { resolveModelConfig } from './ai-agent-model-config.js'
 import type { AIRunStateService } from '../../services/ai-run-state-service.js'
 import type { AIAgentRunnerService } from '../../services/ai-agent-runner-service.js'
@@ -42,6 +50,7 @@ import type { AIAgentRunnerService } from '../../services/ai-agent-runner-servic
 type PersistingChannel = AIStreamChannel & {
   fullText: string
   flush: () => Promise<void>
+  totalUsage: { inputTokens: number; outputTokens: number; model?: string }
 }
 
 function createPersistingChannel(
@@ -53,6 +62,14 @@ function createPersistingChannel(
   let stepText = ''
   let stepToolCalls: AIToolCall[] = []
   let stepToolResults: AIToolResult[] = []
+  const totalUsage: {
+    inputTokens: number
+    outputTokens: number
+    model?: string
+  } = {
+    inputTokens: 0,
+    outputTokens: 0,
+  }
 
   const flushStep = async () => {
     if (!storage) return
@@ -92,6 +109,9 @@ function createPersistingChannel(
     get fullText() {
       return fullText
     },
+    get totalUsage() {
+      return totalUsage
+    },
     flush: flushStep,
     close: () => parent.close(),
     send: (event: AIStreamEvent) => {
@@ -119,6 +139,9 @@ function createPersistingChannel(
             })
             break
           case 'usage':
+            totalUsage.inputTokens += event.tokens.input
+            totalUsage.outputTokens += event.tokens.output
+            if (event.model) totalUsage.model = event.model
             flushStep()
             break
           case 'done':
@@ -144,6 +167,7 @@ async function postStreamCleanup(
   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--) {
@@ -152,7 +176,10 @@ async function postStreamCleanup(
       const result = await mw.modifyOutput(singletonServices, {
         text: outputText,
         messages: outputMessages,
-        usage: { inputTokens: 0, outputTokens: 0 },
+        usage: {
+          inputTokens: usage.inputTokens,
+          outputTokens: usage.outputTokens,
+        },
       })
       outputText = result.text
       outputMessages = result.messages
@@ -183,18 +210,254 @@ async function postStreamCleanup(
     }
   }
 
-  await aiRunState.updateRun(runId, { status: 'completed' })
+  await aiRunState.updateRun(runId, {
+    status: 'completed',
+    ...(usage.model
+      ? {
+          usage: {
+            inputTokens: usage.inputTokens,
+            outputTokens: usage.outputTokens,
+            model: usage.model,
+          },
+        }
+      : {}),
+  })
+}
+
+type StepLoopParams = {
+  agent: CoreAIAgent
+  runnerParams: AIAgentRunnerParams
+  maxSteps: number
+  agentRunner: AIAgentRunnerService
+  persistingChannel: PersistingChannel
+  channel: AIStreamChannel
+  runId: string
+  aiMiddlewares: PikkuAIMiddlewareHooks[]
+}
+
+type StepLoopResult =
+  | { outcome: 'done' }
+  | { outcome: 'approval'; approval: ToolApprovalRequired }
+
+async function runStreamStepLoop(
+  params: StepLoopParams
+): Promise<StepLoopResult> {
+  const {
+    agent,
+    runnerParams,
+    maxSteps,
+    agentRunner,
+    persistingChannel,
+    channel,
+    runId,
+    aiMiddlewares,
+  } = params
+
+  const singletonServices = getSingletonServices()
+
+  for (let step = 0; step < maxSteps; step++) {
+    if (agent.prepareStep) {
+      let stopped = false
+      await agent.prepareStep({
+        stepNumber: step,
+        messages: runnerParams.messages,
+        tools: runnerParams.tools,
+        toolChoice: runnerParams.toolChoice,
+        model: runnerParams.model,
+        stop: () => {
+          stopped = true
+        },
+      })
+      if (stopped) break
+    }
+
+    channel.send({ type: 'step-start', stepNumber: step })
+
+    const stepResult = await agentRunner.stream(runnerParams, persistingChannel)
+
+    for (const mw of aiMiddlewares) {
+      if (mw.afterStep) {
+        await mw.afterStep(singletonServices, {
+          stepNumber: step,
+          text: stepResult.text,
+          toolCalls: stepResult.toolCalls,
+          toolResults: stepResult.toolResults,
+          usage: stepResult.usage,
+          finishReason: stepResult.finishReason,
+        })
+      }
+    }
+
+    if (stepResult.toolCalls.length === 0) break
+
+    const approvalNeeded = checkForApproval(
+      stepResult,
+      runnerParams.tools,
+      runId
+    )
+    if (approvalNeeded) {
+      return { outcome: 'approval', approval: approvalNeeded }
+    }
+
+    appendStepMessages(runnerParams, stepResult)
+  }
+
+  return { outcome: 'done' }
+}
+
+function checkForApproval(
+  stepResult: AIAgentStepResult,
+  tools: AIAgentRunnerParams['tools'],
+  runId: string
+): ToolApprovalRequired | null {
+  for (const tc of stepResult.toolCalls) {
+    const toolDef = tools.find((t) => t.name === tc.toolName)
+
+    if (toolDef?.needsApproval) {
+      return new ToolApprovalRequired(tc.toolCallId, tc.toolName, tc.args)
+    }
+
+    const tr = stepResult.toolResults.find(
+      (r) => r.toolCallId === tc.toolCallId
+    )
+    if (
+      tr?.result &&
+      typeof tr.result === 'object' &&
+      '__approvalRequired' in (tr.result as object)
+    ) {
+      const r = tr.result as {
+        toolName: string
+        args: unknown
+        reason?: string
+        displayToolName?: string
+        displayArgs?: unknown
+        agentRunId?: string
+      }
+      return new ToolApprovalRequired(
+        tc.toolCallId,
+        r.toolName,
+        r.args,
+        r.reason,
+        r.displayToolName,
+        r.displayArgs,
+        r.agentRunId
+      )
+    }
+  }
+  return null
+}
+
+function appendStepMessages(
+  runnerParams: AIAgentRunnerParams,
+  stepResult: AIAgentStepResult
+): void {
+  const assistantMsg: AIMessage = {
+    id: randomUUID(),
+    role: 'assistant',
+    content: stepResult.text || undefined,
+    toolCalls:
+      stepResult.toolCalls.length > 0
+        ? stepResult.toolCalls.map((tc) => ({
+            id: tc.toolCallId,
+            name: tc.toolName,
+            args: tc.args as Record<string, unknown>,
+          }))
+        : undefined,
+    createdAt: new Date(),
+  }
+  runnerParams.messages.push(assistantMsg)
+
+  if (stepResult.toolResults.length > 0) {
+    const toolMsg: AIMessage = {
+      id: randomUUID(),
+      role: 'tool',
+      toolResults: stepResult.toolResults.map((tr) => ({
+        id: tr.toolCallId,
+        name: tr.toolName,
+        result:
+          typeof tr.result === 'string' ? tr.result : JSON.stringify(tr.result),
+      })),
+      createdAt: new Date(),
+    }
+    runnerParams.messages.push(toolMsg)
+  }
+}
+
+function handleApproval(
+  err: ToolApprovalRequired,
+  runId: string,
+  channel: AIStreamChannel,
+  aiRunState: AIRunStateService,
+  persistingChannel: PersistingChannel
+): Promise<void> {
+  return (async () => {
+    await persistingChannel.flush()
+
+    const pendingApproval = err.agentRunId
+      ? {
+          type: 'agent-call' as const,
+          toolCallId: err.toolCallId,
+          agentName: err.toolName,
+          agentRunId: err.agentRunId,
+          displayToolName: err.displayToolName ?? err.toolName,
+          displayArgs: err.displayArgs ?? err.args,
+        }
+      : {
+          type: 'tool-call' as const,
+          toolCallId: err.toolCallId,
+          toolName: err.toolName,
+          args: err.args,
+        }
+
+    await aiRunState.updateRun(runId, {
+      status: 'suspended',
+      suspendReason: 'approval',
+      pendingApprovals: [pendingApproval],
+    })
+
+    const approvalEvent = {
+      type: 'approval-request' as const,
+      toolCallId: err.toolCallId,
+      toolName: err.displayToolName ?? err.toolName,
+      args: err.displayArgs ?? err.args,
+      reason: err.reason,
+      runId,
+    }
+    channel.send(approvalEvent as any)
+    channel.send({ type: 'done' })
+    channel.close()
+  })()
 }
 
 export async function streamAIAgent(
   agentName: string,
-  input: { message: string; threadId: string; resourceId: string },
+  input:
+    | { message: string; threadId: string; resourceId: string }
+    | Record<string, unknown>,
   channel: AIStreamChannel,
   params: RunAIAgentParams,
   agentSessionMap?: Map<string, string>,
   options?: StreamAIAgentOptions
 ): Promise<void> {
   const sessionMap = agentSessionMap ?? new Map<string, string>()
+
+  const { agent: resolvedAgentForProtocol } = resolveAgent(agentName)
+  const inputRecord = input as Record<string, unknown>
+  const useUIMessageStream =
+    resolvedAgentForProtocol.protocol === 'ui-message-stream' ||
+    Array.isArray(inputRecord.messages)
+  let normalizedInput: { message: string; threadId: string; resourceId: string }
+  if (useUIMessageStream) {
+    normalizedInput = parseAssistantUIInput(inputRecord)
+    channel = createAssistantUIChannel(channel)
+  } else {
+    normalizedInput = input as {
+      message: string
+      threadId: string
+      resourceId: string
+    }
+  }
+
   const streamContext: StreamContext = { channel, options }
 
   const {
@@ -207,11 +470,18 @@ export async function streamAIAgent(
     threadId,
     userMessage,
     runnerParams,
+    maxSteps,
     missingRpcs,
     workingMemorySchemaName,
-  } = await prepareAgentRun(agentName, input, params, sessionMap, streamContext)
+  } = await prepareAgentRun(
+    agentName,
+    normalizedInput,
+    params,
+    sessionMap,
+    streamContext
+  )
 
-  const { singletonServices } = params
+  const singletonServices = getSingletonServices()
   const { aiRunState } = singletonServices
   if (!aiRunState) {
     throw new Error('AIRunStateService not available in singletonServices')
@@ -221,7 +491,7 @@ export async function streamAIAgent(
     await aiRunState.createRun({
       agentName,
       threadId,
-      resourceId: input.resourceId,
+      resourceId: normalizedInput.resourceId,
       status: 'suspended',
       suspendReason: 'rpc-missing',
       missingRpcs,
@@ -254,7 +524,7 @@ export async function streamAIAgent(
   const runId = await aiRunState.createRun({
     agentName,
     threadId,
-    resourceId: input.resourceId,
+    resourceId: normalizedInput.resourceId,
     status: 'running',
     usage: { inputTokens: 0, outputTokens: 0, model: agent.model },
     createdAt: new Date(),
@@ -311,7 +581,27 @@ export async function streamAIAgent(
   )
 
   try {
-    await agentRunner.stream(runnerParams, persistingChannel)
+    const loopResult = await runStreamStepLoop({
+      agent,
+      runnerParams,
+      maxSteps,
+      agentRunner,
+      persistingChannel,
+      channel,
+      runId,
+      aiMiddlewares,
+    })
+
+    if (loopResult.outcome === 'approval') {
+      await handleApproval(
+        loopResult.approval,
+        runId,
+        channel,
+        aiRunState,
+        persistingChannel
+      )
+      return
+    }
 
     await postStreamCleanup(
       persistingChannel,
@@ -325,46 +615,23 @@ export async function streamAIAgent(
       aiRunState,
       runId
     )
-  } catch (err) {
-    if (err instanceof ToolApprovalRequired) {
-      await persistingChannel.flush()
-
-      const pendingApproval = err.agentRunId
-        ? {
-            type: 'agent-call' as const,
-            toolCallId: err.toolCallId,
-            agentName: err.toolName,
-            agentRunId: err.agentRunId,
-            displayToolName: err.displayToolName ?? err.toolName,
-            displayArgs: err.displayArgs ?? err.args,
-          }
-        : {
-            type: 'tool-call' as const,
-            toolCallId: err.toolCallId,
-            toolName: err.toolName,
-            args: err.args,
-          }
-
-      await aiRunState.updateRun(runId, {
-        status: 'suspended',
-        suspendReason: 'approval',
-        pendingApprovals: [pendingApproval],
-      })
 
-      const approvalEvent = {
-        type: 'approval-request' as const,
-        toolCallId: err.toolCallId,
-        toolName: err.displayToolName ?? err.toolName,
-        args: err.displayArgs ?? err.args,
-        reason: err.reason,
-        runId,
+    channel.send({ type: 'done' })
+    channel.close()
+  } catch (err) {
+    for (const mw of aiMiddlewares) {
+      if (mw.onError) {
+        try {
+          await mw.onError(singletonServices, {
+            error: err instanceof Error ? err : new Error(String(err)),
+            stepNumber: -1,
+            messages: runnerParams.messages,
+          })
+        } catch {
+          // onError hooks must not affect error flow
+        }
       }
-      channel.send(approvalEvent as any)
-      channel.send({ type: 'done' })
-      channel.close()
-      return
     }
-
     await aiRunState.updateRun(runId, { status: 'failed' })
     channel.send({
       type: 'error',
@@ -377,6 +644,7 @@ export async function streamAIAgent(
 
 export async function resumeAIAgent(
   input: {
+    runId: string
     toolCallId: string
     approved: boolean
   },
@@ -384,18 +652,25 @@ export async function resumeAIAgent(
   params: RunAIAgentParams,
   options?: StreamAIAgentOptions
 ): Promise<void> {
-  const { singletonServices } = params
+  const singletonServices = getSingletonServices()
   const { aiRunState } = singletonServices
   if (!aiRunState) {
     throw new Error('AIRunStateService not available in singletonServices')
   }
 
-  const found = await aiRunState.findRunByToolCallId(input.toolCallId)
-  if (!found) {
-    throw new Error(`No suspended run found for toolCallId ${input.toolCallId}`)
+  const run = await aiRunState.getRun(input.runId)
+  if (!run) {
+    throw new Error(`No run found for runId ${input.runId}`)
   }
 
-  const { run, approval: pending } = found
+  const pending = run.pendingApprovals?.find(
+    (p) => p.toolCallId === input.toolCallId
+  )
+  if (!pending) {
+    throw new Error(
+      `No pending approval for toolCallId ${input.toolCallId} on run ${input.runId}`
+    )
+  }
 
   await aiRunState.resolveApproval(
     input.toolCallId,
@@ -476,7 +751,11 @@ export async function resumeAIAgent(
     })
 
     await resumeAIAgent(
-      { toolCallId: subPending.toolCallId, approved: true },
+      {
+        runId: pending.agentRunId,
+        toolCallId: subPending.toolCallId,
+        approved: true,
+      },
       subChannel,
       params,
       options
@@ -510,14 +789,16 @@ export async function resumeAIAgent(
       channel,
       options: { ...options, requiresToolApproval: false },
     }
+    const aiMiddlewaresForResume: PikkuAIMiddlewareHooks[] =
+      agent.aiMiddleware ?? []
     const { tools } = buildToolDefs(
-      singletonServices,
       params,
       new Map<string, string>(),
       run.resourceId,
       resolvedName,
       packageName,
-      streamContext
+      streamContext,
+      aiMiddlewaresForResume
     )
 
     const matchingTool = tools.find((t) => t.name === pending.toolName)
@@ -529,14 +810,10 @@ export async function resumeAIAgent(
 
     const toolArgs =
       typeof pending.args === 'string' ? JSON.parse(pending.args) : pending.args
-    const { toolApprovalReason: _ignored, ...cleanArgs } = toolArgs as Record<
-      string,
-      unknown
-    >
 
     let toolResult: unknown
     try {
-      toolResult = await matchingTool.execute(cleanArgs)
+      toolResult = await matchingTool.execute(toolArgs)
     } catch (execErr) {
       toolResult = `Error: ${execErr instanceof Error ? execErr.message : String(execErr)}`
     }
@@ -598,7 +875,7 @@ async function continueAfterToolResult(
   aiRunState: AIRunStateService,
   options?: StreamAIAgentOptions
 ): Promise<void> {
-  const { singletonServices } = params
+  const singletonServices = getSingletonServices()
   const agentsMeta = pikkuState(packageName, 'agent', 'agentsMeta')
   const meta = agentsMeta[resolvedName]
   const workingMemorySchemaName = meta?.workingMemorySchema ?? null
@@ -684,24 +961,25 @@ async function continueAfterToolResult(
 
   const streamContext: StreamContext = { channel, options }
   const resumeTools = buildToolDefs(
-    singletonServices,
     params,
     new Map<string, string>(),
     run.resourceId,
     resolvedName,
     packageName,
-    streamContext
+    streamContext,
+    aiMiddlewares
   ).tools
 
   const resolved = resolveModelConfig(resolvedName, agent)
+  const maxSteps = resolved.maxSteps ?? 10
 
-  const runnerParams = {
+  const runnerParams: AIAgentRunnerParams = {
     model: resolved.model,
     temperature: resolved.temperature,
     instructions: modifiedInstructions,
     messages: modifiedMessages,
     tools: resumeTools,
-    maxSteps: resolved.maxSteps ?? 10,
+    maxSteps: 1,
     toolChoice: (agent.toolChoice ?? 'auto') as 'auto' | 'required' | 'none',
     outputSchema: meta?.outputSchema
       ? pikkuState(packageName, 'misc', 'schemas').get(meta.outputSchema)
@@ -709,7 +987,27 @@ async function continueAfterToolResult(
   }
 
   try {
-    await agentRunner.stream(runnerParams, persistingChannel)
+    const loopResult = await runStreamStepLoop({
+      agent,
+      runnerParams,
+      maxSteps,
+      agentRunner,
+      persistingChannel,
+      channel,
+      runId: run.runId,
+      aiMiddlewares,
+    })
+
+    if (loopResult.outcome === 'approval') {
+      await handleApproval(
+        loopResult.approval,
+        run.runId,
+        channel,
+        aiRunState,
+        persistingChannel
+      )
+      return
+    }
 
     await postStreamCleanup(
       persistingChannel,
@@ -723,46 +1021,23 @@ async function continueAfterToolResult(
       aiRunState,
       run.runId
     )
-  } catch (err) {
-    if (err instanceof ToolApprovalRequired) {
-      await persistingChannel.flush()
-
-      const pendingApproval = err.agentRunId
-        ? {
-            type: 'agent-call' as const,
-            toolCallId: err.toolCallId,
-            agentName: err.toolName,
-            agentRunId: err.agentRunId,
-            displayToolName: err.displayToolName ?? err.toolName,
-            displayArgs: err.displayArgs ?? err.args,
-          }
-        : {
-            type: 'tool-call' as const,
-            toolCallId: err.toolCallId,
-            toolName: err.toolName,
-            args: err.args,
-          }
-
-      await aiRunState.updateRun(run.runId, {
-        status: 'suspended',
-        suspendReason: 'approval',
-        pendingApprovals: [pendingApproval],
-      })
 
-      const approvalEvent = {
-        type: 'approval-request' as const,
-        toolCallId: err.toolCallId,
-        toolName: err.displayToolName ?? err.toolName,
-        args: err.displayArgs ?? err.args,
-        reason: err.reason,
-        runId: run.runId,
+    channel.send({ type: 'done' })
+    channel.close()
+  } catch (err) {
+    for (const mw of aiMiddlewares) {
+      if (mw.onError) {
+        try {
+          await mw.onError(singletonServices, {
+            error: err instanceof Error ? err : new Error(String(err)),
+            stepNumber: -1,
+            messages: runnerParams.messages,
+          })
+        } catch {
+          // onError hooks must not affect error flow
+        }
       }
-      channel.send(approvalEvent as any)
-      channel.send({ type: 'done' })
-      channel.close()
-      return
     }
-
     await aiRunState.updateRun(run.runId, { status: 'failed' })
     channel.send({
       type: 'error',