@pikku/core 0.12.0 → 0.12.2

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

@@ -3,7 +3,8 @@ import assert from 'node:assert/strict'
 
 import { resetPikkuState, pikkuState } from '../../pikku-state.js'
 import { runAIAgent } from './ai-agent-runner.js'
-import type { CoreAIAgent } from './ai-agent.types.js'
+import type { CoreAIAgent, PikkuAIMiddlewareHooks } from './ai-agent.types.js'
+import type { AIAgentStepResult } from '../../services/ai-agent-runner-service.js'
 
 beforeEach(() => {
   resetPikkuState()
@@ -29,6 +30,17 @@ const addTestAgent = (agentName: string) => {
   })
 }
 
+const makeStepResult = (
+  overrides?: Partial<AIAgentStepResult>
+): AIAgentStepResult => ({
+  text: '',
+  toolCalls: [],
+  toolResults: [],
+  usage: { inputTokens: 0, outputTokens: 0 },
+  finishReason: 'stop',
+  ...overrides,
+})
+
 describe('runAIAgent', () => {
   test('marks run as failed when runner throws', async () => {
     addTestAgent('failing-agent')
@@ -36,7 +48,7 @@ describe('runAIAgent', () => {
     const updates: unknown[] = []
     const expectedError = new Error('runner failed')
 
-    const singletonServices = {
+    const mockServices = {
       logger: {
         info: () => {},
         warn: () => {},
@@ -56,6 +68,8 @@ describe('runAIAgent', () => {
       },
     } as any
 
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
     await assert.rejects(
       () =>
         runAIAgent(
@@ -65,11 +79,239 @@ describe('runAIAgent', () => {
             threadId: 'thread-1',
             resourceId: 'resource-1',
           },
-          { singletonServices }
+          {}
         ),
       expectedError
     )
 
     assert.deepEqual(updates, [{ status: 'failed' }])
   })
+
+  test('loops through multiple steps and accumulates usage', async () => {
+    addTestAgent('multi-step-agent')
+
+    let callCount = 0
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        run: async (): Promise<AIAgentStepResult> => {
+          callCount++
+          if (callCount === 1) {
+            return makeStepResult({
+              text: '',
+              toolCalls: [
+                { toolCallId: 'tc-1', toolName: 'myTool', args: { q: 'test' } },
+              ],
+              toolResults: [
+                {
+                  toolCallId: 'tc-1',
+                  toolName: 'myTool',
+                  result: 'tool output',
+                },
+              ],
+              usage: { inputTokens: 100, outputTokens: 50 },
+              finishReason: 'tool-calls',
+            })
+          }
+          return makeStepResult({
+            text: 'Final answer',
+            usage: { inputTokens: 80, outputTokens: 30 },
+            finishReason: 'stop',
+          })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-multi',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    const result = await runAIAgent(
+      'multi-step-agent',
+      {
+        message: 'hello',
+        threadId: 'thread-multi',
+        resourceId: 'resource-multi',
+      },
+      {}
+    )
+
+    assert.equal(callCount, 2)
+    assert.equal(result.text, 'Final answer')
+    assert.equal(result.usage.inputTokens, 180)
+    assert.equal(result.usage.outputTokens, 80)
+    assert.equal(result.steps.length, 2)
+  })
+
+  test('afterStep hook is called for each step', async () => {
+    addTestAgent('afterstep-agent')
+
+    const afterStepCalls: unknown[] = []
+    const middleware: PikkuAIMiddlewareHooks = {
+      afterStep: async (_services, ctx) => {
+        afterStepCalls.push(ctx)
+      },
+    }
+    const agent = pikkuState(null, 'agent', 'agents').get('afterstep-agent')!
+    agent.aiMiddleware = [middleware] as any
+    pikkuState(null, 'agent', 'agents').set('afterstep-agent', agent)
+
+    let callCount = 0
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        run: async (): Promise<AIAgentStepResult> => {
+          callCount++
+          if (callCount === 1) {
+            return makeStepResult({
+              text: 'thinking...',
+              toolCalls: [
+                { toolCallId: 'tc-1', toolName: 'myTool', args: { q: 'test' } },
+              ],
+              toolResults: [
+                { toolCallId: 'tc-1', toolName: 'myTool', result: 'output' },
+              ],
+              usage: { inputTokens: 100, outputTokens: 50 },
+              finishReason: 'tool-calls',
+            })
+          }
+          return makeStepResult({
+            text: 'done',
+            usage: { inputTokens: 80, outputTokens: 30 },
+            finishReason: 'stop',
+          })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-afterstep',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await runAIAgent(
+      'afterstep-agent',
+      { message: 'hi', threadId: 't1', resourceId: 'r1' },
+      {}
+    )
+
+    assert.equal(afterStepCalls.length, 2)
+    const first = afterStepCalls[0] as any
+    assert.equal(first.stepNumber, 0)
+    assert.equal(first.finishReason, 'tool-calls')
+    assert.equal(first.toolCalls.length, 1)
+    const second = afterStepCalls[1] as any
+    assert.equal(second.stepNumber, 1)
+    assert.equal(second.finishReason, 'stop')
+    assert.equal(second.text, 'done')
+  })
+
+  test('onError hook is called when runner throws', async () => {
+    addTestAgent('onerror-agent')
+
+    const onErrorCalls: unknown[] = []
+    const middleware: PikkuAIMiddlewareHooks = {
+      onError: async (_services, ctx) => {
+        onErrorCalls.push(ctx)
+      },
+    }
+    const agent = pikkuState(null, 'agent', 'agents').get('onerror-agent')!
+    agent.aiMiddleware = [middleware] as any
+    pikkuState(null, 'agent', 'agents').set('onerror-agent', agent)
+
+    const expectedError = new Error('boom')
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        run: async () => {
+          throw expectedError
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-err',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await assert.rejects(
+      () =>
+        runAIAgent(
+          'onerror-agent',
+          { message: 'hi', threadId: 't2', resourceId: 'r2' },
+          {}
+        ),
+      expectedError
+    )
+
+    assert.equal(onErrorCalls.length, 1)
+    const call = onErrorCalls[0] as any
+    assert.equal(call.error.message, 'boom')
+    assert.equal(call.stepNumber, -1)
+  })
+
+  test('onError hook failure does not suppress original error', async () => {
+    addTestAgent('onerror-throws-agent')
+
+    const middleware: PikkuAIMiddlewareHooks = {
+      onError: async () => {
+        throw new Error('hook error')
+      },
+    }
+    const agent = pikkuState(null, 'agent', 'agents').get(
+      'onerror-throws-agent'
+    )!
+    agent.aiMiddleware = [middleware] as any
+    pikkuState(null, 'agent', 'agents').set('onerror-throws-agent', agent)
+
+    const expectedError = new Error('original error')
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        run: async () => {
+          throw expectedError
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-safe',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await assert.rejects(
+      () =>
+        runAIAgent(
+          'onerror-throws-agent',
+          { message: 'hi', threadId: 't3', resourceId: 'r3' },
+          {}
+        ),
+      expectedError
+    )
+  })
 })

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

@@ -1,11 +1,16 @@
 import type {
   AIAgentInput,
   AIAgentOutput,
+  AIAgentStep,
+  AIMessage,
   PikkuAIMiddlewareHooks,
 } from './ai-agent.types.js'
+import type { AIAgentStepResult } from '../../services/ai-agent-runner-service.js'
 
 import { saveMessages } from './ai-agent-memory.js'
 import { prepareAgentRun, type RunAIAgentParams } from './ai-agent-prepare.js'
+import { getSingletonServices } from '../../pikku-state.js'
+import { randomUUID } from 'crypto'
 
 export async function runAIAgent(
   agentName: string,
@@ -23,12 +28,13 @@ export async function runAIAgent(
     threadId,
     userMessage,
     runnerParams,
+    maxSteps,
     missingRpcs,
     workingMemoryJsonSchema,
     workingMemorySchemaName,
   } = await prepareAgentRun(agentName, input, params, sessionMap)
 
-  const { singletonServices } = params
+  const singletonServices = getSingletonServices()
   const { aiRunState } = singletonServices
   if (!aiRunState) {
     throw new Error('AIRunStateService not available in singletonServices')
@@ -83,7 +89,102 @@ export async function runAIAgent(
   })
 
   try {
-    const result = await agentRunner.run(runnerParams)
+    const accumulatedSteps: AIAgentStep[] = []
+    const totalUsage = { inputTokens: 0, outputTokens: 0 }
+    let lastStepResult: AIAgentStepResult | null = null
+
+    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
+      }
+
+      const stepResult = await agentRunner.run(runnerParams)
+      lastStepResult = stepResult
+
+      totalUsage.inputTokens += stepResult.usage.inputTokens
+      totalUsage.outputTokens += stepResult.usage.outputTokens
+
+      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,
+          })
+        }
+      }
+
+      accumulatedSteps.push({
+        usage: stepResult.usage,
+        toolCalls: stepResult.toolCalls.map((tc) => {
+          const tr = stepResult.toolResults.find(
+            (r) => r.toolCallId === tc.toolCallId
+          )
+          return {
+            name: tc.toolName,
+            args: tc.args as Record<string, unknown>,
+            result:
+              typeof tr?.result === 'string'
+                ? tr.result
+                : JSON.stringify(tr?.result ?? ''),
+          }
+        }),
+      })
+
+      if (stepResult.toolCalls.length === 0) break
+
+      const assistantMsg: AIMessage = {
+        id: randomUUID(),
+        role: 'assistant',
+        content: stepResult.text || undefined,
+        toolCalls: stepResult.toolCalls.map((tc) => ({
+          id: tc.toolCallId,
+          name: tc.toolName,
+          args: tc.args as Record<string, unknown>,
+        })),
+        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)
+      }
+    }
+
+    const finalText = lastStepResult?.text ?? ''
+    const finalObject = lastStepResult?.object
+
+    const result = {
+      text: finalText,
+      steps: accumulatedSteps,
+    }
 
     const responseText = await saveMessages(
       storage,
@@ -108,7 +209,7 @@ export async function runAIAgent(
         const modResult = await mw.modifyOutput(singletonServices, {
           text: outputText,
           messages: outputMessages,
-          usage: result.usage,
+          usage: totalUsage,
         })
         outputText = modResult.text
         outputMessages = modResult.messages
@@ -117,18 +218,31 @@ export async function runAIAgent(
 
     await aiRunState.updateRun(runId, {
       status: 'completed',
-      usage: { ...result.usage, model: agent.model },
+      usage: { ...totalUsage, model: agent.model },
     })
 
     return {
       runId,
       text: outputText,
-      object: result.object,
+      object: finalObject,
       threadId,
-      steps: result.steps,
-      usage: result.usage,
+      steps: accumulatedSteps,
+      usage: totalUsage,
     }
   } catch (error) {
+    for (const mw of aiMiddlewares) {
+      if (mw.onError) {
+        try {
+          await mw.onError(singletonServices, {
+            error: error instanceof Error ? error : new Error(String(error)),
+            stepNumber: -1,
+            messages: runnerParams.messages,
+          })
+        } catch {
+          // onError hooks must not affect error flow
+        }
+      }
+    }
     await aiRunState.updateRun(runId, {
       status: 'failed',
     })