@vellumai/assistant 0.6.5 → 0.6.6

package/src/agent/loop.ts

@@ -6,8 +6,23 @@ import {
   estimateToolsTokens,
   getCalibrationProviderKey,
 } from "../context/token-estimator.js";
-import { truncateOversizedToolResults } from "../context/tool-result-truncation.js";
-import { getHookManager } from "../hooks/manager.js";
+import { calculateMaxToolResultChars } from "../context/tool-result-truncation.js";
+import { defaultEmptyResponseTerminal } from "../plugins/defaults/empty-response.js";
+import { defaultToolErrorTerminal } from "../plugins/defaults/tool-error.js";
+import { defaultToolResultTruncateTerminal } from "../plugins/defaults/tool-result-truncate.js";
+import { DEFAULT_TIMEOUTS, runPipeline } from "../plugins/pipeline.js";
+import { getMiddlewaresFor } from "../plugins/registry.js";
+import type {
+  EmptyResponseArgs,
+  EmptyResponseDecision,
+  LLMCallArgs,
+  LLMCallResult,
+  ToolErrorArgs,
+  ToolErrorDecision,
+  ToolResultTruncateArgs,
+  ToolResultTruncateResult,
+  TurnContext,
+} from "../plugins/types.js";
 import type {
   ContentBlock,
   Message,
@@ -20,7 +35,7 @@ import {
   applyStreamingSubstitution,
   applySubstitutions,
 } from "../tools/sensitive-output-placeholders.js";
-import { ProviderError } from "../util/errors.js";
+import { AssistantError, ErrorCode, ProviderError } from "../util/errors.js";
 import { getLogger } from "../util/logger.js";
 import { isRetryableNetworkError } from "../util/retry.js";
 
@@ -75,6 +90,10 @@ export type AgentEvent =
       };
       status?: string;
       contentBlocks?: ContentBlock[];
+      riskLevel?: string;
+      riskReason?: string;
+      isContainerized?: boolean;
+      riskScopeOptions?: Array<{ pattern: string; label: string }>;
     }
   | { type: "tool_use_preview_start"; toolUseId: string; toolName: string }
   | {
@@ -125,6 +144,61 @@ const DEFAULT_CONFIG: AgentLoopConfig = {
 const MAX_CONSECUTIVE_ERROR_NUDGES = 3;
 const MAX_EMPTY_RESPONSE_RETRIES = 1;
 
+/**
+ * Build a minimal {@link TurnContext} for pipeline invocations inside the
+ * agent loop. Real production call sites thread a full `TurnContext` into
+ * `AgentLoop.run()` (see the `turnContext` parameter on
+ * {@link AgentLoop.run}); this helper is the fallback used only by unit
+ * tests that construct `AgentLoop` directly without an orchestrator.
+ *
+ * When the orchestrator-supplied context is present, {@link resolveLoopTurnContext}
+ * is used instead of this helper so the pipeline sees the real
+ * `conversationId`, trust, and `contextWindowManager`. In the fallback path
+ * the returned context is still useful for pipeline logging: `requestId`
+ * surfaces in every structured record, and `turnIndex` reflects the
+ * current tool-use iteration.
+ */
+function buildLoopTurnContext(
+  requestId: string | undefined,
+  turnIndex: number,
+): TurnContext {
+  return {
+    requestId: requestId ?? "agent-loop",
+    // Loop-scoped pipelines do not currently carry a conversation ID; the
+    // outer orchestrator owns that dimension. Use a fixed sentinel so log
+    // consumers can filter loop-origin records out of conversation queries.
+    conversationId: "agent-loop",
+    turnIndex,
+    trust: {
+      sourceChannel: "vellum",
+      trustClass: "unknown",
+    },
+  };
+}
+
+/**
+ * Produce a `TurnContext` for a pipeline call inside {@link AgentLoop.run}.
+ *
+ * When the orchestrator supplied a `turnContext`, clone it and overwrite
+ * `requestId` + `turnIndex` with the loop-scoped values so plugin log
+ * records correctly attribute the call to the current tool-use iteration
+ * while preserving the real `conversationId`, trust context, and
+ * `contextWindowManager` the orchestrator assembled for the turn. Without
+ * an orchestrator context (unit tests that instantiate `AgentLoop` with no
+ * `turnContext`), fall back to {@link buildLoopTurnContext}'s synthesized
+ * placeholder.
+ */
+function resolveLoopTurnContext(
+  base: TurnContext | undefined,
+  requestId: string | undefined,
+  turnIndex: number,
+): TurnContext {
+  if (base) {
+    return { ...base, requestId: requestId ?? base.requestId, turnIndex };
+  }
+  return buildLoopTurnContext(requestId, turnIndex);
+}
+
 /**
  * User-config HTTP status codes that should never page the on-call: billing
  * exhaustion (402), invalid credentials (401), and forbidden/plan-gated (403).
@@ -167,6 +241,42 @@ export interface ResolvedSystemPrompt {
   model?: string;
 }
 
+/**
+ * Callback shape the loop uses to execute a tool invocation.
+ *
+ * The trailing `turnContext` is optional so in-process tests that wire the
+ * callback without an orchestrator keep working. Production sites (the
+ * `Conversation`'s `createToolExecutor`) forward the supplied context into
+ * `ToolExecutor.execute` so the `toolExecute` pipeline sees the orchestrator's
+ * real conversation identity/trust/contextWindowManager instead of the
+ * synthesized placeholder `ToolExecutor` would otherwise build from the
+ * `ToolContext` alone.
+ */
+export type LoopToolExecutor = (
+  name: string,
+  input: Record<string, unknown>,
+  onOutput?: (chunk: string) => void,
+  toolUseId?: string,
+  turnContext?: TurnContext,
+) => Promise<{
+  content: string;
+  isError: boolean;
+  diff?: {
+    filePath: string;
+    oldContent: string;
+    newContent: string;
+    isNewFile: boolean;
+  };
+  status?: string;
+  contentBlocks?: ContentBlock[];
+  sensitiveBindings?: SensitiveOutputBinding[];
+  yieldToUser?: boolean;
+  riskLevel?: string;
+  riskReason?: string;
+  isContainerized?: boolean;
+  riskScopeOptions?: Array<{ pattern: string; label: string }>;
+}>;
+
 export class AgentLoop {
   private provider: Provider;
   private systemPrompt: string;
@@ -176,52 +286,14 @@ export class AgentLoop {
   private resolveSystemPrompt:
     | ((history: Message[]) => ResolvedSystemPrompt)
     | null;
-  private toolExecutor:
-    | ((
-        name: string,
-        input: Record<string, unknown>,
-        onOutput?: (chunk: string) => void,
-        toolUseId?: string,
-      ) => Promise<{
-        content: string;
-        isError: boolean;
-        diff?: {
-          filePath: string;
-          oldContent: string;
-          newContent: string;
-          isNewFile: boolean;
-        };
-        status?: string;
-        contentBlocks?: ContentBlock[];
-        sensitiveBindings?: SensitiveOutputBinding[];
-        yieldToUser?: boolean;
-      }>)
-    | null;
+  private toolExecutor: LoopToolExecutor | null;
 
   constructor(
     provider: Provider,
     systemPrompt: string,
     config?: Partial<AgentLoopConfig>,
     tools?: ToolDefinition[],
-    toolExecutor?: (
-      name: string,
-      input: Record<string, unknown>,
-      onOutput?: (chunk: string) => void,
-      toolUseId?: string,
-    ) => Promise<{
-      content: string;
-      isError: boolean;
-      diff?: {
-        filePath: string;
-        oldContent: string;
-        newContent: string;
-        isNewFile: boolean;
-      };
-      status?: string;
-      contentBlocks?: ContentBlock[];
-      sensitiveBindings?: SensitiveOutputBinding[];
-      yieldToUser?: boolean;
-    }>,
+    toolExecutor?: LoopToolExecutor,
     resolveTools?: (history: Message[]) => ToolDefinition[],
     resolveSystemPrompt?: (history: Message[]) => ResolvedSystemPrompt,
   ) {
@@ -234,6 +306,21 @@ export class AgentLoop {
     this.toolExecutor = toolExecutor ?? null;
   }
 
+  /**
+   * Resolve the tool definitions sent to the provider for the given turn.
+   *
+   * Mirrors the logic of {@link getToolTokenBudget} but returns the tool
+   * array itself — callers that need to thread the tool set into a plugin
+   * pipeline (e.g. `tokenEstimate`, where the pipeline's args include
+   * `tools`) use this rather than re-implementing the dynamic-vs-static
+   * resolver fork.
+   */
+  getResolvedTools(history?: Message[]): ToolDefinition[] {
+    return history && this.resolveTools
+      ? this.resolveTools(history)
+      : this.tools;
+  }
+
   /**
    * Estimate token cost of the tool definitions sent to the provider.
    *
@@ -243,9 +330,7 @@ export class AgentLoop {
    * without a resolver), falls back to the static `this.tools`.
    */
   getToolTokenBudget(history?: Message[]): number {
-    const tools =
-      history && this.resolveTools ? this.resolveTools(history) : this.tools;
-    return estimateToolsTokens(tools);
+    return estimateToolsTokens(this.getResolvedTools(history));
   }
 
   async run(
@@ -253,8 +338,19 @@ export class AgentLoop {
     onEvent: (event: AgentEvent) => void | Promise<void>,
     signal?: AbortSignal,
     requestId?: string,
-    onCheckpoint?: (checkpoint: CheckpointInfo) => CheckpointDecision,
+    onCheckpoint?: (
+      checkpoint: CheckpointInfo,
+    ) => CheckpointDecision | Promise<CheckpointDecision>,
     callSite?: LLMCallSite,
+    /**
+     * Optional per-turn context supplied by the orchestrator. Every pipeline
+     * invocation inside the loop clones from this value (overwriting only
+     * `turnIndex`/`requestId`) so middleware sees the real conversation
+     * identity, trust class, and `contextWindowManager` rather than the
+     * `"agent-loop"` sentinel used when the loop is instantiated standalone
+     * in unit tests.
+     */
+    turnContext?: TurnContext,
   ): Promise<Message[]> {
     const history = [...messages];
     const initialHistoryLength = messages.length;
@@ -354,22 +450,6 @@ export class AgentLoop {
           providerConfig.callSite = callSite;
         }
 
-        const preLlmResult = await getHookManager().trigger("pre-llm-call", {
-          systemPrompt: turnSystemPrompt,
-          messages: history,
-          toolCount: currentTools.length,
-        });
-
-        if (preLlmResult.blocked) {
-          onEvent({
-            type: "error",
-            error: new Error(
-              `LLM call blocked by hook "${preLlmResult.blockedBy}"`,
-            ),
-          });
-          break;
-        }
-
         // Rate-limit consecutive LLM calls to prevent spin when tools return instantly
         const minInterval = this.config.minTurnIntervalMs ?? 0;
         if (minInterval > 0 && lastLlmCallTime > 0) {
@@ -411,11 +491,26 @@ export class AgentLoop {
           stripOldImageBlocks(history),
         );
 
-        const response = await this.provider.sendMessage(
-          providerHistory,
-          currentTools.length > 0 ? currentTools : undefined,
-          turnSystemPrompt,
-          {
+        // Wrap the provider call in the `llmCall` pipeline so middleware
+        // contributed by plugins may observe, rewrite, short-circuit, or
+        // post-process every LLM request. The terminal below is the real
+        // `provider.sendMessage(...)` call; middleware that call `next(args)`
+        // eventually reach it. The default `defaultLlmCallPlugin` contributes
+        // only a passthrough middleware that forwards to `next(args)` —
+        // registered at module load, it sits at the outermost layer in the
+        // onion, so short-circuiting there would silently disable every
+        // user-registered `llmCall` middleware. Timeout is `null`
+        // (`DEFAULT_TIMEOUTS.llmCall`) — the provider layer already enforces
+        // its own HTTP-level budgets.
+        //
+        // The `onEvent` wrapping is kept inside `args.options` so substitution
+        // and streaming behavior exactly match the pre-pipeline call site.
+        const llmCallArgs: LLMCallArgs = {
+          provider: this.provider,
+          messages: providerHistory,
+          tools: currentTools.length > 0 ? currentTools : undefined,
+          systemPrompt: turnSystemPrompt,
+          options: {
             config: providerConfig,
             onEvent: (event) => {
               if (event.type === "text_delta") {
@@ -466,6 +561,36 @@ export class AgentLoop {
             },
             signal,
           },
+        };
+
+        // Per-turn pipeline context. When the orchestrator threaded a full
+        // `turnContext` into `run()`, use it (overwriting `turnIndex` with
+        // the loop-scoped tool-use iteration) so middleware sees the real
+        // conversation identity, trust, and `contextWindowManager`. The
+        // synthesized fallback is only reached by standalone unit-test
+        // instantiations that never plumb a context through.
+        const turnCtx = resolveLoopTurnContext(
+          turnContext,
+          requestId,
+          toolUseTurns,
+        );
+
+        const response: LLMCallResult = await runPipeline<
+          LLMCallArgs,
+          LLMCallResult
+        >(
+          "llmCall",
+          getMiddlewaresFor("llmCall"),
+          (args) =>
+            args.provider.sendMessage(
+              args.messages,
+              args.tools,
+              args.systemPrompt,
+              args.options,
+            ),
+          llmCallArgs,
+          turnCtx,
+          DEFAULT_TIMEOUTS.llmCall,
         );
 
         const providerDurationMs = Date.now() - providerStart;
@@ -484,14 +609,6 @@ export class AgentLoop {
           estimatedInputTokens: preSendEstimatedTokens,
         });
 
-        void getHookManager().trigger("post-llm-call", {
-          model: response.model,
-          inputTokens: response.usage.inputTokens,
-          outputTokens: response.usage.outputTokens,
-          contentBlockCount: response.content.length,
-          durationMs: providerDurationMs,
-        });
-
         // Flush any buffered streaming text from the substitution pipeline
         if (streamingPending.length > 0) {
           const flushed = applySubstitutions(streamingPending, substitutionMap);
@@ -554,6 +671,12 @@ export class AgentLoop {
         // invocations passed in via `messages`) must NOT suppress the
         // nudge — those turns completed long ago and have no bearing on
         // whether the current tool-use chain has delivered text yet.
+        //
+        // The actual decision (nudge vs. accept vs. error) is delegated to
+        // the `emptyResponse` plugin pipeline. The pipeline returns a
+        // decision; the loop carries out the side-effect (pushing the nudge
+        // or surfacing the error). See `plugins/defaults/empty-response.ts`
+        // for the default decision logic.
         const hasVisibleText = response.content.some(
           (block) => block.type === "text" && block.text.trim().length > 0,
         );
@@ -571,13 +694,37 @@ export class AgentLoop {
           }
           return false;
         })();
-        if (
-          !hasVisibleText &&
-          toolUseBlocks.length === 0 &&
-          toolUseTurns > 0 &&
-          !priorAssistantHadVisibleText &&
-          emptyResponseRetries < MAX_EMPTY_RESPONSE_RETRIES
-        ) {
+
+        const emptyResponseArgs: EmptyResponseArgs = {
+          responseContent: response.content,
+          toolUseBlocksLength: toolUseBlocks.length,
+          toolUseTurns,
+          emptyResponseRetries,
+          maxEmptyResponseRetries: MAX_EMPTY_RESPONSE_RETRIES,
+          priorAssistantHadVisibleText,
+        };
+        const emptyResponseCtx = resolveLoopTurnContext(
+          turnContext,
+          requestId,
+          toolUseTurns,
+        );
+        const emptyResponseDecision: EmptyResponseDecision = await runPipeline(
+          "emptyResponse",
+          getMiddlewaresFor("emptyResponse"),
+          async (args) => defaultEmptyResponseTerminal(args),
+          emptyResponseArgs,
+          emptyResponseCtx,
+          DEFAULT_TIMEOUTS.emptyResponse,
+        );
+
+        if (emptyResponseDecision.action === "nudge") {
+          // Fall back to the canonical nudge text if the plugin returned
+          // `action: "nudge"` but forgot `nudgeText`. Keeps a misbehaving
+          // plugin from silently breaking the loop invariant that the
+          // model sees a coherent prompt.
+          const nudgeText =
+            emptyResponseDecision.nudgeText ??
+            "<system_notice>Your previous response was empty. You must respond to the user with a summary of what you found or did. Do not use any tools — just respond with text.</system_notice>";
           emptyResponseRetries++;
           rlog.warn(
             { turn: toolUseTurns, retry: emptyResponseRetries },
@@ -585,16 +732,25 @@ export class AgentLoop {
           );
           history.push({
             role: "user",
-            content: [
-              {
-                type: "text",
-                text: "<system_notice>Your previous response was empty. You must respond to the user with a summary of what you found or did. Do not use any tools — just respond with text.</system_notice>",
-              },
-            ],
+            content: [{ type: "text", text: nudgeText }],
           });
           continue;
         }
 
+        if (emptyResponseDecision.action === "error") {
+          rlog.error(
+            { turn: toolUseTurns, retries: emptyResponseRetries },
+            "emptyResponse pipeline requested error surface",
+          );
+          throw new AssistantError(
+            "Model returned empty response after tool results",
+            ErrorCode.INTERNAL_ERROR,
+          );
+        }
+
+        // action === "accept" — fall through. Emit a dedicated log line for
+        // the specific "empty turn after tool results, retries exhausted"
+        // case so ops dashboards that grep on this line keep working.
         if (
           !hasVisibleText &&
           toolUseBlocks.length === 0 &&
@@ -664,6 +820,14 @@ export class AgentLoop {
                 });
               },
               toolUse.id,
+              // Forward the loop's resolved `TurnContext` through the
+              // executor callback so `ToolExecutor.execute` can thread the
+              // real orchestrator context into the `toolExecute` pipeline.
+              // Standalone tests that don't wire a `turnContext` into
+              // `AgentLoop.run()` pass `undefined` here and the executor
+              // falls back to the synthesized placeholder — preserving the
+              // existing unit-test behavior.
+              turnCtx,
             );
 
             return { toolUse, result };
@@ -727,12 +891,52 @@ export class AgentLoop {
           }),
         );
 
-        // Pre-emptively truncate oversized tool results to prevent context overflow
-        const { blocks: resultBlocks, truncatedCount } =
-          truncateOversizedToolResults(
-            rawResultBlocks,
-            this.config.maxInputTokens ?? 180_000,
+        // Pre-emptively truncate oversized tool results to prevent context
+        // overflow. The work is delegated to the `toolResultTruncate`
+        // plugin pipeline so downstream plugins can swap in a smarter
+        // truncation strategy (e.g. a summariser) while the default
+        // middleware preserves the historical tail-drop behaviour.
+        const contextWindowTokens = this.config.maxInputTokens ?? 180_000;
+        const maxChars = calculateMaxToolResultChars(contextWindowTokens);
+        const truncateMiddlewares = getMiddlewaresFor("toolResultTruncate");
+
+        let truncatedCount = 0;
+        const truncatedBlocks: ContentBlock[] = [];
+        for (const block of rawResultBlocks) {
+          if (block.type !== "tool_result") {
+            truncatedBlocks.push(block);
+            continue;
+          }
+          const toolBlock = block as ToolResultContent;
+          if (
+            typeof toolBlock.content !== "string" ||
+            toolBlock.content.length <= maxChars
+          ) {
+            truncatedBlocks.push(block);
+            continue;
+          }
+          const pipelineResult = await runPipeline<
+            ToolResultTruncateArgs,
+            ToolResultTruncateResult
+          >(
+            "toolResultTruncate",
+            truncateMiddlewares,
+            async (args) => defaultToolResultTruncateTerminal(args),
+            { content: toolBlock.content, maxChars },
+            turnCtx,
+            DEFAULT_TIMEOUTS.toolResultTruncate,
           );
+          if (pipelineResult.truncated) {
+            truncatedCount++;
+            truncatedBlocks.push({
+              ...toolBlock,
+              content: pipelineResult.content,
+            });
+          } else {
+            truncatedBlocks.push(block);
+          }
+        }
+        const resultBlocks = truncatedBlocks;
         if (truncatedCount > 0) {
           log.warn(
             `Truncated ${truncatedCount} oversized tool result(s) to prevent context overflow`,
@@ -758,6 +962,10 @@ export class AgentLoop {
             diff: result.diff,
             status: result.status,
             contentBlocks: result.contentBlocks,
+            riskLevel: result.riskLevel,
+            riskReason: result.riskReason,
+            isContainerized: result.isContainerized,
+            riskScopeOptions: result.riskScopeOptions,
           });
         }
 
@@ -779,29 +987,59 @@ export class AgentLoop {
         // When any tool returned an error, nudge the LLM to retry with
         // corrected parameters instead of ending its turn. Skip the nudge
         // after MAX_CONSECUTIVE_ERROR_NUDGES consecutive error turns
-        // (the error is likely unrecoverable at that point).
+        // (the error is likely unrecoverable at that point). The nudge
+        // decision is delegated to the `toolError` plugin pipeline so user
+        // plugins can change the text, observe the event, or suppress it.
         const hasToolError = toolResults.some(({ result }) => result.isError);
         if (hasToolError) {
           consecutiveErrorTurns++;
         } else {
           consecutiveErrorTurns = 0;
         }
-        if (
-          hasToolError &&
-          consecutiveErrorTurns <= MAX_CONSECUTIVE_ERROR_NUDGES
-        ) {
+        const toolErrorArgs: ToolErrorArgs = {
+          hasToolError,
+          consecutiveErrorTurns,
+          maxConsecutiveErrorNudges: MAX_CONSECUTIVE_ERROR_NUDGES,
+        };
+        const toolErrorCtx: TurnContext = resolveLoopTurnContext(
+          turnContext,
+          requestId,
+          toolUseTurns - 1,
+        );
+        const toolErrorDecision = await runPipeline<
+          ToolErrorArgs,
+          ToolErrorDecision
+        >(
+          "toolError",
+          getMiddlewaresFor("toolError"),
+          // Terminal: the canonical nudge decision. The default plugin's
+          // middleware is a passthrough (so later-registered user plugins
+          // aren't shadowed), so this terminal is what actually produces
+          // the decision when no user plugin overrides it. Wiring the
+          // decision here — rather than inside the default plugin's
+          // middleware — also preserves the legacy nudge for direct
+          // AgentLoop callers (tests, benchmarks) that skip
+          // `bootstrapPlugins()` and therefore never register the default.
+          async (args) => defaultToolErrorTerminal(args),
+          toolErrorArgs,
+          toolErrorCtx,
+          DEFAULT_TIMEOUTS.toolError,
+        );
+        if (toolErrorDecision.action === "nudge") {
           resultBlocks.push({
             type: "text",
-            text: "<system_notice>One or more tool calls returned an error. If the error looks recoverable (e.g. missing or invalid parameters), fix the parameters and retry. If the error is clearly unrecoverable (e.g. a service is down, a resource does not exist, or a permission is permanently denied), report it to the user.</system_notice>",
+            text: toolErrorDecision.nudgeText,
           });
         }
 
         // Add tool results as a user message and continue the loop
         history.push({ role: "user", content: resultBlocks });
 
-        // Invoke checkpoint callback after tool results are in history
+        // Invoke checkpoint callback after tool results are in history.
+        // The callback may be async — the mid-loop budget check delegates
+        // to the `tokenEstimate` plugin pipeline, which is asynchronous.
         if (onCheckpoint) {
-          const decision = onCheckpoint({
+          const decision = await onCheckpoint({
             turnIndex: toolUseTurns - 1, // 0-based (toolUseTurns was already incremented)
             toolCount: toolUseBlocks.length,
             hasToolUse: true,