@pikku/core 0.12.0 → 0.12.1

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

@@ -4,7 +4,12 @@ import assert from 'node:assert/strict'
 import { resetPikkuState, pikkuState } from '../../pikku-state.js'
 import { streamAIAgent } from './ai-agent-stream.js'
 import { ToolApprovalRequired } from './ai-agent-prepare.js'
-import type { CoreAIAgent, AIStreamEvent } from './ai-agent.types.js'
+import type {
+  CoreAIAgent,
+  AIStreamEvent,
+  PikkuAIMiddlewareHooks,
+} from './ai-agent.types.js'
+import type { AIAgentStepResult } from '../../services/ai-agent-runner-service.js'
 
 beforeEach(() => {
   resetPikkuState()
@@ -30,6 +35,17 @@ const addTestAgent = (agentName: string) => {
   })
 }
 
+const makeStepResult = (
+  overrides?: Partial<AIAgentStepResult>
+): AIAgentStepResult => ({
+  text: '',
+  toolCalls: [],
+  toolResults: [],
+  usage: { inputTokens: 0, outputTokens: 0 },
+  finishReason: 'stop',
+  ...overrides,
+})
+
 describe('streamAIAgent', () => {
   test('marks run as failed and emits error/done when stream runner throws', async () => {
     addTestAgent('failing-stream-agent')
@@ -37,7 +53,7 @@ describe('streamAIAgent', () => {
     const updates: Array<{ runId: string; patch: unknown }> = []
     const events: AIStreamEvent[] = []
 
-    const singletonServices = {
+    const mockServices = {
       logger: {
         info: () => {},
         warn: () => {},
@@ -57,6 +73,8 @@ describe('streamAIAgent', () => {
       },
     } as any
 
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
     await streamAIAgent(
       'failing-stream-agent',
       {
@@ -73,23 +91,76 @@ describe('streamAIAgent', () => {
         },
         close: () => {},
       },
-      { singletonServices }
+      {}
     )
 
     assert.deepEqual(updates, [{ runId: 'run-1', patch: { status: 'failed' } }])
     assert.deepEqual(
       events.map((event) => event.type),
-      ['error', 'done']
+      ['step-start', 'error', 'done']
+    )
+  })
+
+  test('completes normally when stream returns no tool calls', async () => {
+    addTestAgent('simple-stream-agent')
+
+    const updates: Array<{ runId: string; patch: unknown }> = []
+    const events: AIStreamEvent[] = []
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async (): Promise<AIAgentStepResult> => {
+          return makeStepResult({ text: 'Hello world' })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-simple',
+        updateRun: async (runId: string, patch: unknown) => {
+          updates.push({ runId, patch })
+        },
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await streamAIAgent(
+      'simple-stream-agent',
+      {
+        message: 'hello',
+        threadId: 'thread-simple',
+        resourceId: 'resource-simple',
+      },
+      {
+        channelId: 'channel-simple',
+        openingData: undefined,
+        state: 'open',
+        send: (event: AIStreamEvent) => {
+          events.push(event)
+        },
+        close: () => {},
+      },
+      {}
     )
+
+    assert.deepEqual(updates, [
+      { runId: 'run-simple', patch: { status: 'completed' } },
+    ])
+    assert.ok(events.some((e) => e.type === 'done'))
   })
 
-  test('suspends run and emits approval request when tool approval is required (no throw)', async () => {
+  test('suspends run when tool needs approval (needsApproval flag)', async () => {
     addTestAgent('approval-stream-agent')
 
     const updates: Array<{ runId: string; patch: unknown }> = []
     const events: AIStreamEvent[] = []
 
-    const singletonServices = {
+    const mockServices = {
       logger: {
         info: () => {},
         warn: () => {},
@@ -97,8 +168,16 @@ describe('streamAIAgent', () => {
         debug: () => {},
       },
       aiAgentRunner: {
-        stream: async () => {
-          throw new ToolApprovalRequired('tool-call-1', 'tool-x', { id: 1 })
+        stream: async (): Promise<AIAgentStepResult> => {
+          return makeStepResult({
+            toolCalls: [
+              {
+                toolCallId: 'tool-call-1',
+                toolName: 'tool-x',
+                args: { id: 1 },
+              },
+            ],
+          })
         },
       },
       aiRunState: {
@@ -109,8 +188,25 @@ describe('streamAIAgent', () => {
       },
     } as any
 
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    const agentName = 'approval-stream-agent'
+    const agent = pikkuState(null, 'agent', 'agents').get(agentName)!
+    agent.tools = ['tool-x']
+    pikkuState(null, 'agent', 'agents').set(agentName, agent)
+    pikkuState(null, 'agent', 'agentsMeta')[agentName].tools = ['tool-x']
+    pikkuState(null, 'rpc', 'meta')['tool-x'] = 'tool-x'
+    pikkuState(null, 'function', 'meta')['tool-x'] = {
+      description: 'test tool',
+      requiresApproval: true,
+    }
+    pikkuState(null, 'misc', 'schemas').set('ToolXInput', {
+      type: 'object',
+      properties: { id: { type: 'number' } },
+    })
+
     await streamAIAgent(
-      'approval-stream-agent',
+      agentName,
       {
         message: 'hello',
         threadId: 'thread-2',
@@ -125,7 +221,7 @@ describe('streamAIAgent', () => {
         },
         close: () => {},
       },
-      { singletonServices },
+      {},
       undefined,
       { requiresToolApproval: 'all' }
     )
@@ -157,18 +253,20 @@ describe('streamAIAgent', () => {
     assert.deepEqual(approvalEvent.args, { id: 1 })
 
     assert.deepEqual(
-      events.map((event) => event.type),
+      events
+        .filter((e) => e.type === 'approval-request' || e.type === 'done')
+        .map((e) => e.type),
       ['approval-request', 'done']
     )
   })
 
-  test('suspends with agent-call type when agentRunId is present', async () => {
+  test('suspends with agent-call type when sub-agent returns approval sentinel', async () => {
     addTestAgent('parent-agent')
 
     const updates: Array<{ runId: string; patch: unknown }> = []
     const events: AIStreamEvent[] = []
 
-    const singletonServices = {
+    const mockServices = {
       logger: {
         info: () => {},
         warn: () => {},
@@ -176,16 +274,30 @@ describe('streamAIAgent', () => {
         debug: () => {},
       },
       aiAgentRunner: {
-        stream: async () => {
-          throw new ToolApprovalRequired(
-            'tool-call-2',
-            'sub-agent',
-            { message: 'hi' },
-            undefined,
-            'deleteTodo',
-            { todoId: '2' },
-            'sub-run-1'
-          )
+        stream: async (): Promise<AIAgentStepResult> => {
+          return makeStepResult({
+            toolCalls: [
+              {
+                toolCallId: 'tool-call-2',
+                toolName: 'sub-agent',
+                args: { message: 'hi' },
+              },
+            ],
+            toolResults: [
+              {
+                toolCallId: 'tool-call-2',
+                toolName: 'sub-agent',
+                result: {
+                  __approvalRequired: true,
+                  toolName: 'sub-agent',
+                  args: { message: 'hi' },
+                  displayToolName: 'deleteTodo',
+                  displayArgs: { todoId: '2' },
+                  agentRunId: 'sub-run-1',
+                },
+              },
+            ],
+          })
         },
       },
       aiRunState: {
@@ -196,6 +308,8 @@ describe('streamAIAgent', () => {
       },
     } as any
 
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
     await streamAIAgent(
       'parent-agent',
       {
@@ -212,7 +326,7 @@ describe('streamAIAgent', () => {
         },
         close: () => {},
       },
-      { singletonServices }
+      {}
     )
 
     assert.deepEqual(updates, [
@@ -243,4 +357,296 @@ describe('streamAIAgent', () => {
     assert.equal(approvalEvent.toolName, 'deleteTodo')
     assert.deepEqual(approvalEvent.args, { todoId: '2' })
   })
+
+  test('loops through multiple steps when tools are called', async () => {
+    addTestAgent('multi-step-agent')
+
+    const updates: Array<{ runId: string; patch: unknown }> = []
+    const events: AIStreamEvent[] = []
+    let callCount = 0
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async (): Promise<AIAgentStepResult> => {
+          callCount++
+          if (callCount === 1) {
+            return makeStepResult({
+              text: 'Calling tool...',
+              toolCalls: [
+                { toolCallId: 'tc-1', toolName: 'myTool', args: { q: 'test' } },
+              ],
+              toolResults: [
+                {
+                  toolCallId: 'tc-1',
+                  toolName: 'myTool',
+                  result: 'tool result',
+                },
+              ],
+              finishReason: 'tool-calls',
+            })
+          }
+          return makeStepResult({ text: 'Final answer', finishReason: 'stop' })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-multi',
+        updateRun: async (runId: string, patch: unknown) => {
+          updates.push({ runId, patch })
+        },
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await streamAIAgent(
+      'multi-step-agent',
+      {
+        message: 'hello',
+        threadId: 'thread-multi',
+        resourceId: 'resource-multi',
+      },
+      {
+        channelId: 'channel-multi',
+        openingData: undefined,
+        state: 'open',
+        send: (event: AIStreamEvent) => {
+          events.push(event)
+        },
+        close: () => {},
+      },
+      {}
+    )
+
+    assert.equal(callCount, 2)
+    assert.deepEqual(updates, [
+      { runId: 'run-multi', patch: { status: 'completed' } },
+    ])
+    assert.ok(events.some((e) => e.type === 'done'))
+  })
+
+  test('afterStep hook is called for each stream step', async () => {
+    addTestAgent('afterstep-stream-agent')
+
+    const afterStepCalls: unknown[] = []
+    const middleware: PikkuAIMiddlewareHooks = {
+      afterStep: async (_services, ctx) => {
+        afterStepCalls.push(ctx)
+      },
+    }
+    const agent = pikkuState(null, 'agent', 'agents').get(
+      'afterstep-stream-agent'
+    )!
+    agent.aiMiddleware = [middleware] as any
+    pikkuState(null, 'agent', 'agents').set('afterstep-stream-agent', agent)
+
+    let callCount = 0
+    const events: AIStreamEvent[] = []
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async (): Promise<AIAgentStepResult> => {
+          callCount++
+          if (callCount === 1) {
+            return makeStepResult({
+              text: 'step1',
+              toolCalls: [
+                { toolCallId: 'tc-1', toolName: 'myTool', args: { q: 'x' } },
+              ],
+              toolResults: [
+                { toolCallId: 'tc-1', toolName: 'myTool', result: 'ok' },
+              ],
+              usage: { inputTokens: 10, outputTokens: 5 },
+              finishReason: 'tool-calls',
+            })
+          }
+          return makeStepResult({
+            text: 'final',
+            finishReason: 'stop',
+          })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-as',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await streamAIAgent(
+      'afterstep-stream-agent',
+      { message: 'hi', threadId: 't-as', resourceId: 'r-as' },
+      {
+        channelId: 'c-as',
+        openingData: undefined,
+        state: 'open',
+        send: (event: AIStreamEvent) => events.push(event),
+        close: () => {},
+      },
+      {}
+    )
+
+    assert.equal(afterStepCalls.length, 2)
+    assert.equal((afterStepCalls[0] as any).stepNumber, 0)
+    assert.equal((afterStepCalls[0] as any).finishReason, 'tool-calls')
+    assert.equal((afterStepCalls[1] as any).stepNumber, 1)
+    assert.equal((afterStepCalls[1] as any).text, 'final')
+  })
+
+  test('onError hook is called when stream throws', async () => {
+    addTestAgent('onerror-stream-agent')
+
+    const onErrorCalls: unknown[] = []
+    const middleware: PikkuAIMiddlewareHooks = {
+      onError: async (_services, ctx) => {
+        onErrorCalls.push(ctx)
+      },
+    }
+    const agent = pikkuState(null, 'agent', 'agents').get(
+      'onerror-stream-agent'
+    )!
+    agent.aiMiddleware = [middleware] as any
+    pikkuState(null, 'agent', 'agents').set('onerror-stream-agent', agent)
+
+    const events: AIStreamEvent[] = []
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async () => {
+          throw new Error('stream boom')
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-oe',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+
+    await streamAIAgent(
+      'onerror-stream-agent',
+      { message: 'hi', threadId: 't-oe', resourceId: 'r-oe' },
+      {
+        channelId: 'c-oe',
+        openingData: undefined,
+        state: 'open',
+        send: (event: AIStreamEvent) => events.push(event),
+        close: () => {},
+      },
+      {}
+    )
+
+    assert.equal(onErrorCalls.length, 1)
+    assert.equal((onErrorCalls[0] as any).error.message, 'stream boom')
+    assert.ok(events.some((e) => e.type === 'error'))
+    assert.ok(events.some((e) => e.type === 'done'))
+  })
+
+  test('beforeToolCall and afterToolCall hooks fire for tool executions', async () => {
+    addTestAgent('toolhooks-stream-agent')
+
+    const hookCalls: { hook: string; ctx: any }[] = []
+    const middleware: PikkuAIMiddlewareHooks = {
+      beforeToolCall: async (_services, ctx) => {
+        hookCalls.push({ hook: 'before', ctx })
+        return { args: { ...ctx.args, injected: true } }
+      },
+      afterToolCall: async (_services, ctx) => {
+        hookCalls.push({ hook: 'after', ctx })
+        return { result: `modified:${ctx.result}` }
+      },
+    }
+    const agentName = 'toolhooks-stream-agent'
+    const agent = pikkuState(null, 'agent', 'agents').get(agentName)!
+    agent.aiMiddleware = [middleware] as any
+    agent.tools = ['myTool']
+    pikkuState(null, 'agent', 'agents').set(agentName, agent)
+    pikkuState(null, 'agent', 'agentsMeta')[agentName].tools = ['myTool']
+    pikkuState(null, 'rpc', 'meta')['myTool'] = 'myTool'
+    pikkuState(null, 'function', 'meta')['myTool'] = {
+      description: 'test tool',
+    }
+
+    const events: AIStreamEvent[] = []
+    let callCount = 0
+
+    const mockServices = {
+      logger: {
+        info: () => {},
+        warn: () => {},
+        error: () => {},
+        debug: () => {},
+      },
+      aiAgentRunner: {
+        stream: async (
+          params: any,
+          _channel: any
+        ): Promise<AIAgentStepResult> => {
+          callCount++
+          if (callCount === 1) {
+            const tool = params.tools.find((t: any) => t.name === 'myTool')
+            if (tool) {
+              try {
+                await tool.execute({ q: 'test' })
+              } catch {
+                // runPikkuFunc not fully mocked; hooks still fire
+              }
+            }
+            return makeStepResult({
+              text: 'done',
+              finishReason: 'stop',
+            })
+          }
+          return makeStepResult({ text: 'done', finishReason: 'stop' })
+        },
+      },
+      aiRunState: {
+        createRun: async () => 'run-th',
+        updateRun: async () => {},
+      },
+    } as any
+
+    pikkuState(null, 'package', 'singletonServices', mockServices)
+    pikkuState(null, 'package', 'createWireServices', () => ({}))
+
+    await streamAIAgent(
+      agentName,
+      { message: 'hi', threadId: 't-th', resourceId: 'r-th' },
+      {
+        channelId: 'c-th',
+        openingData: undefined,
+        state: 'open',
+        send: (event: AIStreamEvent) => events.push(event),
+        close: () => {},
+      },
+      {}
+    )
+
+    assert.equal(hookCalls.length, 2)
+    assert.equal(hookCalls[0].hook, 'before')
+    assert.equal(hookCalls[0].ctx.toolName, 'myTool')
+    assert.deepEqual(hookCalls[0].ctx.args, { q: 'test' })
+    assert.equal(hookCalls[1].hook, 'after')
+    assert.equal(hookCalls[1].ctx.toolName, 'myTool')
+    assert.deepEqual(hookCalls[1].ctx.args, { q: 'test', injected: true })
+    assert.ok(hookCalls[1].ctx.durationMs >= 0)
+  })
 })