@pikku/core 0.12.72 → 0.12.74

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

@@ -37,10 +37,12 @@ import {
   resolveOwnerResourceId,
   agentSessionScope,
   assertResourceOwner,
+  assertResourcePrincipalOwner,
   assertAgentAuthorized,
   ToolApprovalRequired,
   ToolCredentialRequired,
   APPROVAL_REQUIRED,
+  CREDENTIAL_REQUIRED,
   type RunAIAgentParams,
   type StreamAIAgentOptions,
   type StreamContext,
@@ -48,10 +50,21 @@ import {
 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'
+import {
+  getInFlightTools,
+  isAbortError,
+  persistOrphanedToolResults,
+  registerInterruptibleRun,
+  signalRunInterrupt,
+  trackInterruptNote,
+  trackToolExecution,
+  type AgentInterruption,
+  type AgentInterruptResult,
+} from './ai-agent-interrupt.js'
 
 type PersistingChannel = AIStreamChannel & {
   fullText: string
-  flush: () => Promise<void>
+  flush: (opts?: { interrupted?: boolean }) => Promise<void>
   totalUsage: { inputTokens: number; outputTokens: number; model?: string }
 }
 
@@ -75,7 +88,7 @@ function createPersistingChannel(
     outputTokens: 0,
   }
 
-  const flushStep = async () => {
+  const flushStep = async (opts?: { interrupted?: boolean }) => {
     if (!storage) return
     const text = stepText
     const generativeUI = stepGenerativeUI
@@ -97,6 +110,7 @@ function createPersistingChannel(
               ]
             : text || undefined,
         toolCalls: calls.length > 0 ? calls : undefined,
+        ...(opts?.interrupted ? { interrupted: true } : {}),
         createdAt: new Date(),
       }
       const messages: AIMessage[] = [assistantMsg]
@@ -112,14 +126,6 @@ function createPersistingChannel(
     }
   }
 
-  /**
-   * `send` is synchronous and so cannot await the flush. A rejection would have
-   * nothing to propagate to and would take the process down as an unhandled
-   * rejection — a model reusing a toolCallId, which is a primary key in AI
-   * storage, is enough to trigger it. Persistence from inside `send` is
-   * therefore best-effort: the run carries on, and the awaited `flush()` on the
-   * suspend paths still surfaces failures to its caller.
-   */
   const flushDetached = () => {
     void flushStep().catch((error) => {
       logger?.error('Failed to persist agent messages', {
@@ -145,11 +151,14 @@ function createPersistingChannel(
     close: () => parent.close(),
     sendBinary: (data) => parent.sendBinary(data),
     send: (event: AIStreamEvent) => {
+      // Accumulated whether or not storage is configured: `fullText` is what
+      // the client was streamed, and an interrupted run has to be able to
+      // report the fragment it got through even with persistence turned off.
+      if (event.type === 'text-delta') fullText += event.text
       if (storage) {
         switch (event.type) {
           case 'text-delta':
             stepText += event.text
-            fullText += event.text
             break
           case 'tool-call':
             stepToolCalls.push({
@@ -187,6 +196,7 @@ function createPersistingChannel(
     setState: (s) => parent.setState(s),
     getState: () => parent.getState(),
     clearState: () => parent.clearState(),
+    remote: (funcName: string, data?: unknown) => parent.remote(funcName, data),
   }
   return channel
 }
@@ -240,7 +250,6 @@ type StepLoopParams = {
   streamChannel: AIStreamChannel
   persistingChannel: PersistingChannel
   channel: AIStreamChannel
-  runId: string
   aiMiddlewares: PikkuAIMiddlewareHooks[]
 }
 
@@ -259,7 +268,6 @@ async function runStreamStepLoop(
     agentRunner,
     streamChannel,
     channel,
-    runId,
     aiMiddlewares,
   } = params
 
@@ -300,13 +308,8 @@ async function runStreamStepLoop(
 
     if (stepResult.toolCalls.length === 0) break
 
-    const approvalsNeeded = checkForApprovals(
-      stepResult,
-      runnerParams.tools,
-      runId
-    )
+    const approvalsNeeded = checkForApprovals(stepResult, runnerParams.tools)
     if (approvalsNeeded.length > 0) {
-      // For each approval, call approvalDescriptionFn if available
       for (const approval of approvalsNeeded) {
         const toolDef = runnerParams.tools.find(
           (t) => t.name === approval.toolName
@@ -314,17 +317,14 @@ async function runStreamStepLoop(
         if (toolDef?.approvalDescriptionFn && !approval.reason) {
           try {
             approval.reason = await toolDef.approvalDescriptionFn(approval.args)
-          } catch {
-            // If description generation fails, continue without it
-          }
+          } catch {}
         }
       }
       return { outcome: 'approval', approvals: approvalsNeeded }
     }
 
-    const credentialRequests = checkForCredentialRequests(stepResult, runId)
+    const credentialRequests = checkForCredentialRequests(stepResult)
     if (credentialRequests.length > 0) {
-      // Append step messages so the tool result is preserved for resume
       appendStepMessages(runnerParams, stepResult)
       return { outcome: 'credential', credentialRequests }
     }
@@ -337,8 +337,7 @@ async function runStreamStepLoop(
 
 export function checkForApprovals(
   stepResult: AIAgentStepResult,
-  tools: AIAgentRunnerParams['tools'],
-  runId: string
+  tools: AIAgentRunnerParams['tools']
 ): ToolApprovalRequired[] {
   const approvals: ToolApprovalRequired[] = []
   for (const tc of stepResult.toolCalls) {
@@ -351,12 +350,6 @@ export function checkForApprovals(
       continue
     }
 
-    // The approval marker is only trusted from a framework tool declared with
-    // `forwardsApproval` (the sub-agent delegating tools) AND carrying the
-    // non-forgeable `APPROVAL_REQUIRED` Symbol brand. A plain tool's output — or
-    // a delegating tool's LLM-shaped `result.object`, which an attacker can
-    // influence — is plain JSON and can never carry the Symbol, so it can never
-    // forge an approval/suspension.
     if (!toolDef?.forwardsApproval) {
       continue
     }
@@ -417,15 +410,15 @@ export function checkForApprovals(
 }
 
 export function checkForCredentialRequests(
-  stepResult: AIAgentStepResult,
-  runId: string
+  stepResult: AIAgentStepResult
 ): ToolCredentialRequired[] {
   const requests: ToolCredentialRequired[] = []
   for (const tr of stepResult.toolResults) {
+    // knowledge: decisions/security/ai-agent-credential-requests-are-symbol-branded.md
     if (
       tr.result &&
       typeof tr.result === 'object' &&
-      '__credentialRequired' in (tr.result as object)
+      CREDENTIAL_REQUIRED in (tr.result as object)
     ) {
       const r = tr.result as {
         credentialName: string
@@ -579,10 +572,6 @@ function handleCredentialRequests(
       pendingApprovals,
     })
 
-    // The __credentialRequired tool result is suppressed from the stream, so
-    // credential-request events (with the runId needed for /resume) are the
-    // client's signal to show Connect/Ignore buttons — mirroring how
-    // approval-request suspensions work.
     for (const req of requests) {
       channel.send({
         type: 'credential-request',
@@ -626,7 +615,6 @@ export async function streamAIAgent(
   }
 
   const streamContext: StreamContext = { channel, options }
-  // delegateState is attached after prepareAgentRun resolves the agent config
 
   const {
     agent,
@@ -710,6 +698,13 @@ export async function streamAIAgent(
   })
   options?.onRunCreated?.(runId)
 
+  // Registered before the first model call and released in `finally`, so the
+  // window in which an interrupt can land is exactly the window in which there
+  // is something to interrupt.
+  const interruptHandle = registerInterruptibleRun(runId)
+  runnerParams.abortSignal = interruptHandle.signal
+  runnerParams.tools = trackToolExecution(runnerParams.tools, interruptHandle)
+
   if (storage) {
     await storage.saveMessages(threadId, [userMessage])
   }
@@ -725,6 +720,10 @@ export async function streamAIAgent(
           event,
           allEvents,
           state,
+          // Sends downstream directly, so a hook can hand back the fast event
+          // now and push the slow one when it is ready.
+          emit: next,
+          signal: interruptHandle.signal,
         })
         if (result == null) return
         if (Array.isArray(result)) {
@@ -765,12 +764,12 @@ export async function streamAIAgent(
         ).channel as AIStreamChannel)
       : persistingChannel
 
-  // Tool results carrying the __credentialRequired marker must never reach
-  // the client or persisted history: the run suspends with credential-request
-  // events instead (mirroring approvals), and leaving the tool call
-  // unresulted is what lets the client resume it after connecting.
   const credentialFilteredChannel: AIStreamChannel = {
     ...wrappedChannel,
+    // Returns what the inner send returns. Middleware runs asynchronously, so
+    // swallowing the promise here makes every `await channel.send(...)` upstream
+    // a no-op — including the one that waits for a buffering hook to flush
+    // before the channel closes.
     send: (event: AIStreamEvent) => {
       if (
         event.type === 'tool-result' &&
@@ -780,15 +779,11 @@ export async function streamAIAgent(
       ) {
         return
       }
-      wrappedChannel.send(event)
+      return wrappedChannel.send(event)
     },
   }
 
-  // In delegate mode (default), suppress parent's text from reaching the client
-  // AFTER a sub-agent has been called. If the parent responds directly (no delegation),
-  // its text goes through normally. Sub-agent text bypasses this path entirely
-  // (goes through subChannel → channel directly).
-  const isDelegateMode = agent.agentMode !== 'supervise' && meta.agents?.length
+  const isDelegateMode = agent.agentMode !== 'supervise' && meta?.agents?.length
   const delegateState = { delegated: false }
   if (isDelegateMode) {
     streamContext.delegateState = delegateState
@@ -802,7 +797,7 @@ export async function streamAIAgent(
             (event.type === 'text-delta' || event.type === 'reasoning-delta')
           )
             return
-          credentialFilteredChannel.send(event)
+          return credentialFilteredChannel.send(event)
         },
         delegateState,
       }
@@ -817,7 +812,6 @@ export async function streamAIAgent(
       streamChannel: outputChannel,
       persistingChannel,
       channel,
-      runId,
       aiMiddlewares,
     })
 
@@ -852,10 +846,41 @@ export async function streamAIAgent(
       runId
     )
 
-    channel.send({ type: 'done' })
+    // Through the middleware rather than straight at the channel, and awaited.
+    // `done` is the only signal a stream hook gets that the reply is over, and
+    // the ones that buffer need it: `voiceOutput` speaks a trailing fragment
+    // that never reached a full stop, and waits for audio it has already paid
+    // to synthesize. Sent raw, that work is discarded by the `close()` below.
+    await outputChannel.send({ type: 'done' })
     channel.close()
     return persistingChannel.fullText
   } catch (err) {
+    // An interrupt is not a failure: the truncated text is real output the user
+    // already heard part of, so it is persisted and marked rather than dropped.
+    const interruption = interruptHandle.interruption
+    if (interruption || isAbortError(err)) {
+      await persistingChannel.flush({ interrupted: true })
+      await aiRunState.updateRun(runId, { status: 'interrupted' })
+      // Deliberately not awaited. A tool still running must not hold the stream
+      // open — the whole point of barge-in is that the agent stops now — so the
+      // note lands later and the next run on this thread waits for it instead.
+      if (storage) {
+        trackInterruptNote(
+          threadId,
+          persistOrphanedToolResults(interruptHandle, storage, threadId)
+        )
+      }
+      channel.send({
+        type: 'interrupted',
+        runId,
+        text: persistingChannel.fullText,
+        reason: interruption?.reason ?? 'user',
+      })
+      channel.send({ type: 'done' })
+      channel.close()
+      return persistingChannel.fullText
+    }
+
     for (const mw of aiMiddlewares) {
       if (mw.onError) {
         try {
@@ -864,9 +889,7 @@ export async function streamAIAgent(
             stepNumber: -1,
             messages: runnerParams.messages,
           })
-        } catch {
-          // onError hooks must not affect error flow
-        }
+        } catch {}
       }
     }
     await aiRunState.updateRun(runId, {
@@ -880,9 +903,74 @@ export async function streamAIAgent(
     channel.send({ type: 'done' })
     channel.close()
     return persistingChannel.fullText
+  } finally {
+    interruptHandle.release()
   }
 }
 
+/**
+ * Stop an in-flight run on behalf of a caller, after checking the run is theirs.
+ *
+ * Unlike {@link resumeAIAgent} this needs no channel: it does not continue the
+ * stream, it ends one. The `interrupted` event and the truncated message are
+ * emitted by the original {@link streamAIAgent} call on its own channel, so this
+ * can be reached over a plain RPC while the stream is held open elsewhere —
+ * which is exactly the shape a voice barge-in needs.
+ *
+ * Ownership is the whole gate, deliberately without {@link assertAgentAuthorized}.
+ * Resuming re-enters the agent and approves a tool call, so a revoked grant must
+ * block it; stopping is the opposite — if a caller's access to an agent was
+ * revoked mid-run, being able to stop their own run is harmless and arguably
+ * what you want.
+ *
+ * Returns `false` when there was nothing to stop. Racing a run that finishes on
+ * its own is the normal case in voice, so it is not an error.
+ */
+export async function interruptAIAgent(
+  input: { runId: string; reason?: AgentInterruption['reason'] },
+  params: RunAIAgentParams
+): Promise<AgentInterruptResult> {
+  const { aiRunState, logger } = getSingletonServices()
+  if (!aiRunState) {
+    throw new Error('AIRunStateService not available in singletonServices')
+  }
+
+  const run = await aiRunState.getRun(input.runId)
+  if (!run) {
+    throw new Error(`No run found for runId ${input.runId}`)
+  }
+  assertResourceOwner(
+    resolveOwnerResourceId(
+      params,
+      agentSessionScope(run.agentName),
+      run.resourceId
+    ),
+    run.resourceId,
+    'run'
+  )
+
+  const stopped = signalRunInterrupt(input.runId, {
+    reason: input.reason ?? 'user',
+  })
+
+  // A run still marked `running` that this process cannot abort is executing on
+  // another instance. Returning a bare `false` there is indistinguishable from
+  // "it already finished", so the one deployment shape this registry does not
+  // cover fails silently — an agent that simply refuses to shut up, with nothing
+  // in the logs. Say so instead. See `signalRunInterrupt` for the fix (fan the
+  // interrupt out over eventHub so every instance tries locally).
+  if (!stopped && run.status === 'running') {
+    logger?.warn(
+      `Could not interrupt run ${input.runId}: it is running in another process. ` +
+        'Interrupts are process-local; fan them out over eventHub to support multiple instances.'
+    )
+  }
+
+  // Reported after the abort so it names what is *still* running: aborting the
+  // model call does nothing to a tool already executing.
+  return { stopped, inFlightTools: getInFlightTools(input.runId) }
+}
+
 export async function resumeAIAgent(
   input: {
     runId: string
@@ -903,12 +991,9 @@ export async function resumeAIAgent(
   if (!run) {
     throw new Error(`No run found for runId ${input.runId}`)
   }
-  assertResourceOwner(
-    resolveOwnerResourceId(
-      params,
-      agentSessionScope(run.agentName),
-      run.resourceId
-    ),
+  assertResourcePrincipalOwner(
+    params,
+    agentSessionScope(run.agentName),
     run.resourceId,
     'run'
   )
@@ -924,10 +1009,6 @@ export async function resumeAIAgent(
 
   const { agent, packageName, resolvedName } = resolveAgent(run.agentName)
 
-  // Gate before resolving the approval: recording it is a persisted side
-  // effect, so an unauthorized caller must not reach it. Run ownership alone is
-  // not enough — a grant revoked while the run was suspended must stop the
-  // caller from approving its pending tool calls.
   await assertAgentAuthorized(agent, params, packageName)
 
   await aiRunState.resolveApproval(
@@ -981,12 +1062,10 @@ export async function resumeAIAgent(
       result: denialResult,
     })
 
-    // Check remaining pending approvals
     const updatedRun = await aiRunState.getRun(run.runId)
     const remaining = updatedRun?.pendingApprovals ?? []
 
     if (remaining.length > 0) {
-      // Still waiting for more approvals - don't continue step loop
       channel.send({ type: 'done' })
       channel.close()
       return
@@ -1134,12 +1213,10 @@ export async function resumeAIAgent(
     })
   }
 
-  // Check remaining pending approvals after processing this one
   const updatedRun = await aiRunState.getRun(run.runId)
   const remaining = updatedRun?.pendingApprovals ?? []
 
   if (remaining.length > 0) {
-    // Still waiting for more approvals - don't continue step loop
     channel.send({ type: 'done' })
     channel.close()
     return
@@ -1224,6 +1301,12 @@ async function continueAfterToolResult(
     }
   }
 
+  // Resuming is as interruptible as the first turn. It is the same person
+  // listening to the same voice, and after an approval it is where most of the
+  // reply actually gets spoken — an approved delete is followed by the agent
+  // talking about it, and that is a normal thing to talk over.
+  const interruptHandle = registerInterruptibleRun(run.runId)
+
   const streamMiddleware = aiMiddlewares
     .filter((mw) => mw.modifyOutputStream)
     .map((mw) => {
@@ -1235,6 +1318,10 @@ async function continueAfterToolResult(
           event,
           allEvents,
           state,
+          // Sends downstream directly, so a hook can hand back the fast event
+          // now and push the slow one when it is ready.
+          emit: next,
+          signal: interruptHandle.signal,
         })
         if (result == null) return
         if (Array.isArray(result)) {
@@ -1297,11 +1384,15 @@ async function continueAfterToolResult(
     tools: resumeTools,
     maxSteps: 1,
     toolChoice: (agent.toolChoice ?? 'auto') as 'auto' | 'required' | 'none',
+    providerOptions: agent.providerOptions,
     outputSchema: meta?.outputSchema
       ? pikkuState(packageName, 'misc', 'schemas').get(meta.outputSchema)
       : undefined,
   }
 
+  runnerParams.abortSignal = interruptHandle.signal
+  runnerParams.tools = trackToolExecution(runnerParams.tools, interruptHandle)
+
   try {
     const loopResult = await runStreamStepLoop({
       agent,
@@ -1311,7 +1402,6 @@ async function continueAfterToolResult(
       streamChannel: wrappedChannel,
       persistingChannel,
       channel,
-      runId: run.runId,
       aiMiddlewares,
     })
 
@@ -1346,9 +1436,36 @@ async function continueAfterToolResult(
       run.runId
     )
 
-    channel.send({ type: 'done' })
+    // Through the middleware, for the same reason as the first turn — and it
+    // matters more here. After an approval, most of what gets spoken is the
+    // agent talking about what it just did, so this is the half of the reply a
+    // dropped flush would silence.
+    await wrappedChannel.send({ type: 'done' })
     channel.close()
   } catch (err) {
+    // Same reasoning as the first turn: an interrupt is not a failure, and the
+    // part of the reply the user already heard is real output.
+    const interruption = interruptHandle.interruption
+    if (interruption || isAbortError(err)) {
+      await persistingChannel.flush({ interrupted: true })
+      await aiRunState.updateRun(run.runId, { status: 'interrupted' })
+      if (storage) {
+        trackInterruptNote(
+          run.threadId,
+          persistOrphanedToolResults(interruptHandle, storage, run.threadId)
+        )
+      }
+      channel.send({
+        type: 'interrupted',
+        runId: run.runId,
+        text: persistingChannel.fullText,
+        reason: interruption?.reason ?? 'user',
+      })
+      channel.send({ type: 'done' })
+      channel.close()
+      return
+    }
+
     for (const mw of aiMiddlewares) {
       if (mw.onError) {
         try {
@@ -1357,9 +1474,7 @@ async function continueAfterToolResult(
             stepNumber: -1,
             messages: runnerParams.messages,
           })
-        } catch {
-          // onError hooks must not affect error flow
-        }
+        } catch {}
       }
     }
     await aiRunState.updateRun(run.runId, {
@@ -1372,5 +1487,7 @@ async function continueAfterToolResult(
     })
     channel.send({ type: 'done' })
     channel.close()
+  } finally {
+    interruptHandle.release()
   }
 }