@pikku/core 0.11.2 → 0.12.1

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

@@ -0,0 +1,317 @@
+import { beforeEach, describe, test } from 'node:test'
+import assert from 'node:assert/strict'
+
+import { resetPikkuState, pikkuState } from '../../pikku-state.js'
+import { runAIAgent } from './ai-agent-runner.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-model',
+  }
+
+  pikkuState(null, 'agent', 'agentsMeta')[agentName] = {
+    ...agent,
+    inputSchema: null,
+    outputSchema: null,
+    workingMemorySchema: null,
+  }
+  pikkuState(null, 'agent', 'agents').set(agentName, agent)
+  pikkuState(null, 'models', 'config', {
+    models: { 'test-model': 'test/test-model' },
+  })
+}
+
+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')
+
+    const updates: unknown[] = []
+    const expectedError = new Error('runner failed')
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        run: async () => {
+          throw expectedError
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-1',
+        updateRun: async (_runId: string, patch: unknown) => {
+          updates.push(patch)
+        },
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await assert.rejects(
+      () =>
+        runAIAgent(
+          'failing-agent',
+          {
+            message: 'hello',
+            threadId: 'thread-1',
+            resourceId: 'resource-1',
+          },
+          {}
+        ),
+      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

@@ -0,0 +1,251 @@
+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,
+  input: AIAgentInput,
+  params: RunAIAgentParams,
+  agentSessionMap?: Map<string, string>
+): Promise<AIAgentOutput> {
+  const sessionMap = agentSessionMap ?? new Map<string, string>()
+
+  const {
+    agent,
+    agentRunner,
+    storage,
+    memoryConfig,
+    threadId,
+    userMessage,
+    runnerParams,
+    maxSteps,
+    missingRpcs,
+    workingMemoryJsonSchema,
+    workingMemorySchemaName,
+  } = await prepareAgentRun(agentName, input, params, sessionMap)
+
+  const singletonServices = getSingletonServices()
+  const { aiRunState } = singletonServices
+  if (!aiRunState) {
+    throw new Error('AIRunStateService not available in singletonServices')
+  }
+
+  if (missingRpcs.length > 0) {
+    const runId = await aiRunState.createRun({
+      agentName,
+      threadId,
+      resourceId: input.resourceId,
+      status: 'suspended',
+      suspendReason: 'rpc-missing',
+      missingRpcs,
+      usage: { inputTokens: 0, outputTokens: 0, model: agent.model },
+      createdAt: new Date(),
+      updatedAt: new Date(),
+    })
+    return {
+      runId,
+      text: '',
+      threadId,
+      steps: [],
+      usage: { inputTokens: 0, outputTokens: 0 },
+    }
+  }
+
+  const aiMiddlewares: PikkuAIMiddlewareHooks[] = agent.aiMiddleware ?? []
+
+  let modifiedMessages = runnerParams.messages
+  let modifiedInstructions = runnerParams.instructions
+  for (const mw of aiMiddlewares) {
+    if (mw.modifyInput) {
+      const result = await mw.modifyInput(singletonServices, {
+        messages: modifiedMessages,
+        instructions: modifiedInstructions,
+      })
+      modifiedMessages = result.messages
+      modifiedInstructions = result.instructions
+    }
+  }
+  runnerParams.messages = modifiedMessages
+  runnerParams.instructions = modifiedInstructions
+
+  const runId = await aiRunState.createRun({
+    agentName,
+    threadId,
+    resourceId: input.resourceId,
+    status: 'running',
+    usage: { inputTokens: 0, outputTokens: 0, model: agent.model },
+    createdAt: new Date(),
+    updatedAt: new Date(),
+  })
+
+  try {
+    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,
+      threadId,
+      input.resourceId,
+      memoryConfig,
+      userMessage,
+      result,
+      {
+        workingMemoryJsonSchema,
+        workingMemorySchemaName,
+        logger: singletonServices.logger,
+        schemaService: singletonServices.schema,
+      }
+    )
+
+    let outputText = responseText
+    let outputMessages = runnerParams.messages
+    for (let i = aiMiddlewares.length - 1; i >= 0; i--) {
+      const mw = aiMiddlewares[i]
+      if (mw.modifyOutput) {
+        const modResult = await mw.modifyOutput(singletonServices, {
+          text: outputText,
+          messages: outputMessages,
+          usage: totalUsage,
+        })
+        outputText = modResult.text
+        outputMessages = modResult.messages
+      }
+    }
+
+    await aiRunState.updateRun(runId, {
+      status: 'completed',
+      usage: { ...totalUsage, model: agent.model },
+    })
+
+    return {
+      runId,
+      text: outputText,
+      object: finalObject,
+      threadId,
+      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',
+    })
+    throw error
+  }
+}