@pikku/core 0.12.79 → 0.12.82

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

@@ -43,6 +43,17 @@ import {
   trackToolExecution,
 } from './ai-agent-interrupt.js'
 import { pikkuState, getSingletonServices } from '../../pikku-state.js'
+import {
+  applyInputMiddleware,
+  describeApprovals,
+  notifyAfterStep,
+  toAccumulatedStep,
+} from './ai-agent-turn.js'
+import {
+  applyOutputMiddleware,
+  finalizeAgentRun,
+  lastUserMessageText,
+} from './ai-agent-finalize.js'
 import { resolveModelConfig } from './ai-agent-model-config.js'
 import { AIProviderNotConfiguredError } from '../../errors/errors.js'
 import { randomUUID } from './ai-agent-utils.js'
@@ -147,26 +158,20 @@ export async function runAIAgent(
   // One bag per run, shared by every middleware — see PikkuAIMiddlewareHooks.
   const sharedNotes: Record<string, unknown> = {}
 
-  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,
-        shared: sharedNotes,
-      })
-      modifiedMessages = result.messages
-      modifiedInstructions = result.instructions
-    }
-  }
+  const { messages: modifiedMessages, instructions: modifiedInstructions } =
+    await applyInputMiddleware(
+      aiMiddlewares,
+      singletonServices,
+      {
+        messages: runnerParams.messages,
+        instructions: runnerParams.instructions,
+      },
+      sharedNotes
+    )
   runnerParams.messages = modifiedMessages
   runnerParams.instructions = modifiedInstructions
 
-  // History records what the model was asked, which for a spoken turn is the
-  // transcript rather than the base64 audio that arrived — see the same note on
-  // the streaming path. Identity-checked, because a middleware may legitimately
-  // replace the message list with something unrelated to this turn.
+  // knowledge: decisions/internals/thread-history-records-the-transcript-not-the-audio.md
   const lastModified = modifiedMessages[modifiedMessages.length - 1]
   const persistedUserMessage =
     lastModified?.id === userMessage.id ? lastModified : userMessage
@@ -181,10 +186,7 @@ export async function runAIAgent(
     updatedAt: new Date(),
   })
 
-  // Registered on the same terms as `streamAIAgent`. A run created here is
-  // visible to `interruptAIAgent` through `aiRunState` either way, so skipping
-  // this would leave that call finding the run, passing the ownership check and
-  // then failing to stop it — reported as if it were running on another host.
+  // knowledge: decisions/internals/a-non-streaming-agent-run-registers-with-airunstate-too.md
   const interruptHandle = registerInterruptibleRun(runId)
   runnerParams.abortSignal = interruptHandle.signal
   runnerParams.tools = trackToolExecution(runnerParams.tools, interruptHandle)
@@ -216,52 +218,15 @@ export async function runAIAgent(
       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,
-          })
-        }
-      }
+      await notifyAfterStep(aiMiddlewares, singletonServices, step, stepResult)
 
-      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 ?? ''),
-          }
-        }),
-      })
+      accumulatedSteps.push(toAccumulatedStep(stepResult))
 
       if (stepResult.toolCalls.length === 0) break
 
       const approvalsNeeded = checkForApprovals(stepResult, runnerParams.tools)
       if (approvalsNeeded.length > 0) {
-        for (const approval of approvalsNeeded) {
-          const toolDef = runnerParams.tools.find(
-            (t) => t.name === approval.toolName
-          )
-          if (toolDef?.approvalDescriptionFn && !approval.reason) {
-            try {
-              approval.reason = await toolDef.approvalDescriptionFn(
-                approval.args
-              )
-            } catch {}
-          }
-        }
+        await describeApprovals(approvalsNeeded, runnerParams.tools)
 
         const pendingApprovals = approvalsNeeded.map((a) =>
           a.agentRunId
@@ -343,20 +308,13 @@ export async function runAIAgent(
       steps: accumulatedSteps,
     }
 
-    let outputText = finalText
-    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
-      }
-    }
+    const { text: outputText, steps: outputSteps } =
+      await applyOutputMiddleware(aiMiddlewares, singletonServices, {
+        text: finalText,
+        messages: runnerParams.messages,
+        steps: result.steps,
+        usage: totalUsage,
+      })
 
     await saveMessages(
       storage,
@@ -365,14 +323,20 @@ export async function runAIAgent(
       memoryConfig,
       persistedUserMessage,
       {
-        ...result,
         text: outputText,
+        steps: outputSteps,
         uiSpec: structuredOutput.uiSpec,
       }
     )
 
-    await aiRunState.updateRun(runId, {
-      status: 'completed',
+    await finalizeAgentRun(aiRunState, {
+      runId,
+      agentName,
+      threadId,
+      resourceId: input.resourceId,
+      input: lastUserMessageText(runnerParams.messages),
+      text: outputText,
+      steps: outputSteps,
       usage: { ...totalUsage, model: agent.model },
     })
 
@@ -381,14 +345,11 @@ export async function runAIAgent(
       text: outputText,
       object: finalObject,
       threadId,
-      steps: accumulatedSteps,
+      steps: outputSteps,
       usage: totalUsage,
     }
   } catch (error) {
-    // An interrupt is not a failure, so it skips the `onError` hooks and never
-    // becomes an `errorMessage`. Unlike the streaming path there is no partial
-    // reply to hand back — nothing was delivered — so the caller gets a typed
-    // throw it can tell apart from a provider outage instead of a result.
+    // knowledge: decisions/internals/an-agent-interrupt-is-not-a-failure.md
     const interruption = interruptHandle.interruption
     if (interruption || isAbortError(error)) {
       await aiRunState.updateRun(runId, { status: 'interrupted' })
@@ -470,11 +431,24 @@ export async function resumeAIAgentSync(
 
   const savedPendingApprovals = [...(run.pendingApprovals ?? [])]
 
+  // The read above is not a claim — concurrent resumes all see the same pending
+  // list. `resolveApproval` is the claim, and only the caller it returns true
+  // for may run the tool.
+  const claimedIds = new Set<string>()
   for (const { toolCallId, approved } of approvals) {
-    await aiRunState.resolveApproval(
+    const claimed = await aiRunState.resolveApproval(
       toolCallId,
       approved ? 'approved' : 'denied'
     )
+    if (claimed) {
+      claimedIds.add(toolCallId)
+    }
+  }
+
+  if (approvals.length > 0 && claimedIds.size === 0) {
+    throw new Error(
+      `Approvals for run ${runId} were already resolved by another caller`
+    )
   }
 
   const { tools } = await buildToolDefs(
@@ -492,13 +466,17 @@ export async function resumeAIAgentSync(
     toolName: string
     args: any
     result: string
+    error?: string
   }[] = []
 
   for (const pending of savedPendingApprovals) {
     if (pending.type !== 'tool-call') continue
 
     const toolCallId = pending.toolCallId
+    if (!claimedIds.has(toolCallId)) continue
+
     let resultStr: string
+    let toolError: string | undefined
 
     if (rejectedIds.has(toolCallId)) {
       resultStr =
@@ -525,7 +503,8 @@ export async function resumeAIAgentSync(
         if (err?.payload?.error === 'missing_credential') {
           resultStr = JSON.stringify(err.payload)
         } else {
-          resultStr = `Error: ${err instanceof Error ? err.message : String(err)}`
+          toolError = err instanceof Error ? err.message : String(err)
+          resultStr = `Error: ${toolError}`
         }
       }
     } else {
@@ -540,6 +519,7 @@ export async function resumeAIAgentSync(
           ? JSON.parse(pending.args)
           : pending.args,
       result: resultStr,
+      ...(toolError ? { error: toolError } : {}),
     })
   }
 
@@ -569,7 +549,21 @@ export async function resumeAIAgentSync(
     memoryConfig,
     agentRunner,
     params,
-    aiRunState
+    aiRunState,
+    // The approved tools were executed here, before the model was re-entered,
+    // so they belong to the run's step record — otherwise a tool that failed
+    // after approval leaves no trace on the run at all.
+    toolCallMessages.length > 0
+      ? {
+          usage: { inputTokens: 0, outputTokens: 0 },
+          toolCalls: toolCallMessages.map((tc) => ({
+            name: tc.toolName,
+            args: (tc.args ?? {}) as Record<string, unknown>,
+            result: tc.result,
+            ...(tc.error ? { error: tc.error } : {}),
+          })),
+        }
+      : undefined
   )
 }
 
@@ -582,7 +576,8 @@ async function continueAfterToolResultSync(
   memoryConfig: AIAgentMemoryConfig | undefined,
   agentRunner: AIAgentRunnerService,
   params: RunAIAgentParams,
-  aiRunState: AIRunStateService
+  aiRunState: AIRunStateService,
+  resumedToolStep?: AIAgentStep
 ): Promise<AIAgentOutput> {
   const singletonServices = getSingletonServices()
   const agentsMeta = pikkuState(packageName, 'agent', 'agentsMeta')
@@ -623,19 +618,13 @@ async function continueAfterToolResultSync(
   // One bag per run, shared by every middleware — see PikkuAIMiddlewareHooks.
   const sharedNotes: Record<string, unknown> = {}
 
-  let modifiedMessages = trimmedMessages
-  let modifiedInstructions = instructions
-  for (const mw of aiMiddlewares) {
-    if (mw.modifyInput) {
-      const result = await mw.modifyInput(singletonServices, {
-        messages: modifiedMessages,
-        instructions: modifiedInstructions,
-        shared: sharedNotes,
-      })
-      modifiedMessages = result.messages
-      modifiedInstructions = result.instructions
-    }
-  }
+  const { messages: modifiedMessages, instructions: modifiedInstructions } =
+    await applyInputMiddleware(
+      aiMiddlewares,
+      singletonServices,
+      { messages: trimmedMessages, instructions: instructions },
+      sharedNotes
+    )
 
   const { tools: resumeTools } = await buildToolDefs(
     params,
@@ -673,6 +662,13 @@ async function continueAfterToolResultSync(
   runnerParams.tools = trackToolExecution(runnerParams.tools, interruptHandle)
 
   try {
+    // Kept out of `accumulatedSteps` deliberately: that array drives
+    // `saveMessages`, and the approved tool's messages were already written to
+    // the thread before the model was re-entered. It belongs to the run's step
+    // record, not to persistence.
+    const withResumedStep = (steps: AIAgentStep[]): AIAgentStep[] =>
+      resumedToolStep ? [resumedToolStep, ...steps] : steps
+
     const accumulatedSteps: AIAgentStep[] = []
     const totalUsage = { inputTokens: 0, outputTokens: 0 }
     let lastStepResult: AIAgentStepResult | null = null
@@ -684,52 +680,15 @@ async function continueAfterToolResultSync(
       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,
-          })
-        }
-      }
+      await notifyAfterStep(aiMiddlewares, singletonServices, step, stepResult)
 
-      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 ?? ''),
-          }
-        }),
-      })
+      accumulatedSteps.push(toAccumulatedStep(stepResult))
 
       if (stepResult.toolCalls.length === 0) break
 
       const approvalsNeeded = checkForApprovals(stepResult, runnerParams.tools)
       if (approvalsNeeded.length > 0) {
-        for (const approval of approvalsNeeded) {
-          const toolDef = runnerParams.tools.find(
-            (t) => t.name === approval.toolName
-          )
-          if (toolDef?.approvalDescriptionFn && !approval.reason) {
-            try {
-              approval.reason = await toolDef.approvalDescriptionFn(
-                approval.args
-              )
-            } catch {}
-          }
-        }
+        await describeApprovals(approvalsNeeded, runnerParams.tools)
 
         const pendingApprovals = approvalsNeeded.map((a) =>
           a.agentRunId
@@ -788,7 +747,7 @@ async function continueAfterToolResultSync(
           runId: run.runId,
           text: suspendedText,
           threadId: run.threadId,
-          steps: accumulatedSteps,
+          steps: withResumedStep(accumulatedSteps),
           usage: totalUsage,
           status: 'suspended',
           pendingApprovals: approvalsNeeded.map((a) => ({
@@ -812,35 +771,36 @@ async function continueAfterToolResultSync(
       steps: accumulatedSteps,
     }
 
-    let outputText = finalText
-    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
-      }
-    }
+    const { text: outputText, steps: outputSteps } =
+      await applyOutputMiddleware(aiMiddlewares, singletonServices, {
+        text: finalText,
+        messages: runnerParams.messages,
+        steps: withResumedStep(result.steps),
+        usage: totalUsage,
+      })
 
     await saveMessages(
       storage,
       run.threadId,
       run.resourceId,
       memoryConfig,
-      null as any,
+      null,
       {
-        ...result,
         text: outputText,
+        // The approved tool's messages were written before the model was
+        // re-entered, so only the steps this leg generated are persisted here.
+        steps: accumulatedSteps,
       }
     )
 
-    await aiRunState.updateRun(run.runId, {
-      status: 'completed',
+    await finalizeAgentRun(aiRunState, {
+      runId: run.runId,
+      agentName: resolvedName,
+      threadId: run.threadId,
+      resourceId: run.resourceId,
+      input: lastUserMessageText(runnerParams.messages),
+      text: outputText,
+      steps: outputSteps,
       usage: { ...totalUsage, model: agent.model },
     })
 
@@ -849,7 +809,7 @@ async function continueAfterToolResultSync(
       text: outputText,
       object: finalObject,
       threadId: run.threadId,
-      steps: accumulatedSteps,
+      steps: outputSteps,
       usage: totalUsage,
     }
   } catch (error) {