@vellumai/assistant 0.6.5 → 0.6.6

package/src/__tests__/token-estimate-pipeline.test.ts

@@ -0,0 +1,431 @@
+/**
+ * Tests for the `tokenEstimate` plugin pipeline (PR 22 of the
+ * agent-plugin-system plan).
+ *
+ * Covers:
+ * - The default plugin's terminal middleware matches
+ *   {@link estimatePromptTokensRaw} output exactly across a set of golden
+ *   inputs (empty history, text-only, tools, provider-specific image sizing).
+ * - Running the pipeline end-to-end with the default registered produces
+ *   the same numeric result as calling `estimatePromptTokensRaw` directly.
+ * - A custom plugin that short-circuits the chain can override the default,
+ *   proving the extension point works.
+ *
+ * These tests exercise the registry + runner directly. They do not touch
+ * `bootstrapPlugins` — the default registration path is covered by the
+ * bootstrap suite.
+ */
+
+import { afterEach, beforeEach, describe, expect, test } from "bun:test";
+
+import {
+  estimatePromptTokensRaw,
+  estimateToolsTokens,
+} from "../context/token-estimator.js";
+import type { TrustContext } from "../daemon/conversation-runtime-assembly.js";
+import {
+  defaultTokenEstimatePlugin,
+  defaultTokenEstimateTerminal,
+} from "../plugins/defaults/token-estimate.js";
+import { DEFAULT_TIMEOUTS, runPipeline } from "../plugins/pipeline.js";
+import {
+  getMiddlewaresFor,
+  registerPlugin,
+  resetPluginRegistryForTests,
+} from "../plugins/registry.js";
+import type {
+  EstimateArgs,
+  EstimateResult,
+  Middleware,
+  Plugin,
+  TurnContext,
+} from "../plugins/types.js";
+import type { Message, ToolDefinition } from "../providers/types.js";
+
+// ── Fixtures ─────────────────────────────────────────────────────────────
+
+const trust: TrustContext = {
+  sourceChannel: "vellum",
+  trustClass: "guardian",
+};
+
+function makeCtx(overrides: Partial<TurnContext> = {}): TurnContext {
+  return {
+    requestId: "req-token-estimate-test",
+    conversationId: "conv-token-estimate-test",
+    turnIndex: 0,
+    trust,
+    ...overrides,
+  };
+}
+
+const EMPTY_HISTORY: Message[] = [];
+
+const TEXT_HISTORY: Message[] = [
+  { role: "user", content: [{ type: "text", text: "hello there" }] },
+  {
+    role: "assistant",
+    content: [
+      { type: "text", text: "hi! how can I help you today?" },
+      { type: "text", text: "a second text block for good measure" },
+    ],
+  },
+];
+
+const TOOL_USE_HISTORY: Message[] = [
+  { role: "user", content: [{ type: "text", text: "what's in the log?" }] },
+  {
+    role: "assistant",
+    content: [
+      {
+        type: "tool_use",
+        id: "tu-1",
+        name: "bash",
+        input: { command: "tail -n 5 server.log" },
+      },
+    ],
+  },
+  {
+    role: "user",
+    content: [
+      {
+        type: "tool_result",
+        tool_use_id: "tu-1",
+        content: "line1\nline2\nline3",
+      },
+    ],
+  },
+];
+
+const SYSTEM_PROMPT = "You are a helpful assistant with a long preamble.";
+
+const SAMPLE_TOOLS: ToolDefinition[] = [
+  {
+    name: "bash",
+    description: "Execute a shell command and return its output.",
+    input_schema: {
+      type: "object",
+      properties: { command: { type: "string" } },
+      required: ["command"],
+    },
+  },
+  {
+    name: "file_read",
+    description: "Read a file from the workspace.",
+    input_schema: {
+      type: "object",
+      properties: { path: { type: "string" } },
+      required: ["path"],
+    },
+  },
+];
+
+// ── Helpers ──────────────────────────────────────────────────────────────
+
+function registerDefault(): void {
+  registerPlugin(defaultTokenEstimatePlugin);
+}
+
+function rawEstimate(
+  args: Pick<EstimateArgs, "history" | "systemPrompt" | "providerName"> & {
+    tools: ToolDefinition[];
+  },
+): number {
+  const toolTokenBudget =
+    args.tools.length > 0 ? estimateToolsTokens(args.tools) : 0;
+  return estimatePromptTokensRaw(args.history, args.systemPrompt, {
+    providerName: args.providerName,
+    toolTokenBudget,
+  });
+}
+
+async function runViaPipeline(args: EstimateArgs): Promise<EstimateResult> {
+  return runPipeline<EstimateArgs, EstimateResult>(
+    "tokenEstimate",
+    getMiddlewaresFor("tokenEstimate"),
+    // Mirror the production wiring in `daemon/conversation-agent-loop.ts`:
+    // the default plugin's middleware is a passthrough, so the terminal is
+    // wired in by the call site. Using the same terminal here means the
+    // tests exercise the exact composition shape that ships.
+    defaultTokenEstimateTerminal,
+    args,
+    makeCtx(),
+    DEFAULT_TIMEOUTS.tokenEstimate,
+  );
+}
+
+// ── Tests ────────────────────────────────────────────────────────────────
+
+beforeEach(() => {
+  resetPluginRegistryForTests();
+});
+
+afterEach(() => {
+  resetPluginRegistryForTests();
+});
+
+describe("tokenEstimate pipeline — default plugin parity", () => {
+  test("default matches estimatePromptTokensRaw on empty history", async () => {
+    registerDefault();
+    const args: EstimateArgs = {
+      history: EMPTY_HISTORY,
+      systemPrompt: undefined,
+      tools: [],
+      providerName: undefined,
+    };
+    const pipelineResult = await runViaPipeline(args);
+    expect(pipelineResult).toBe(rawEstimate(args));
+  });
+
+  test("default matches estimatePromptTokensRaw on text-only history", async () => {
+    registerDefault();
+    const args: EstimateArgs = {
+      history: TEXT_HISTORY,
+      systemPrompt: SYSTEM_PROMPT,
+      tools: [],
+      providerName: "anthropic",
+    };
+    const pipelineResult = await runViaPipeline(args);
+    expect(pipelineResult).toBe(rawEstimate(args));
+    // Sanity: the system prompt adds real token cost, so the number is
+    // strictly larger than the bare-history estimate.
+    expect(pipelineResult).toBeGreaterThan(
+      rawEstimate({
+        history: TEXT_HISTORY,
+        systemPrompt: undefined,
+        tools: [],
+        providerName: "anthropic",
+      }),
+    );
+  });
+
+  test("default matches estimatePromptTokensRaw with tool_use/tool_result blocks", async () => {
+    registerDefault();
+    const args: EstimateArgs = {
+      history: TOOL_USE_HISTORY,
+      systemPrompt: SYSTEM_PROMPT,
+      tools: SAMPLE_TOOLS,
+      providerName: "anthropic",
+    };
+    const pipelineResult = await runViaPipeline(args);
+    expect(pipelineResult).toBe(rawEstimate(args));
+  });
+
+  test("default folds tool definition tokens into the result", async () => {
+    registerDefault();
+    const baseArgs: EstimateArgs = {
+      history: TEXT_HISTORY,
+      systemPrompt: SYSTEM_PROMPT,
+      tools: [],
+      providerName: "anthropic",
+    };
+    const withoutTools = await runViaPipeline(baseArgs);
+    const withTools = await runViaPipeline({
+      ...baseArgs,
+      tools: SAMPLE_TOOLS,
+    });
+    // Tools contribute non-zero overhead; the pipeline result must grow.
+    const toolBudget = estimateToolsTokens(SAMPLE_TOOLS);
+    expect(toolBudget).toBeGreaterThan(0);
+    expect(withTools - withoutTools).toBe(toolBudget);
+  });
+
+  test("provider-specific image sizing flows through the default", async () => {
+    registerDefault();
+    // Two providers see different image token costs for the same content —
+    // the raw estimator is the source of truth, so the pipeline must agree
+    // under both provider names.
+    const imageHistory: Message[] = [
+      {
+        role: "user",
+        content: [
+          {
+            type: "image",
+            source: {
+              type: "base64",
+              media_type: "image/png",
+              // Small fake PNG-ish payload; the estimator's fallback path
+              // kicks in when parseImageDimensions fails, which is fine —
+              // the two providers still diverge on overhead.
+              data: "a".repeat(128),
+            },
+          },
+        ],
+      },
+    ];
+    const anthropicArgs: EstimateArgs = {
+      history: imageHistory,
+      systemPrompt: undefined,
+      tools: [],
+      providerName: "anthropic",
+    };
+    const openaiArgs: EstimateArgs = {
+      ...anthropicArgs,
+      providerName: "openai",
+    };
+    const anthropicResult = await runViaPipeline(anthropicArgs);
+    const openaiResult = await runViaPipeline(openaiArgs);
+    expect(anthropicResult).toBe(rawEstimate(anthropicArgs));
+    expect(openaiResult).toBe(rawEstimate(openaiArgs));
+  });
+});
+
+describe("tokenEstimate pipeline — custom override", () => {
+  test("custom plugin short-circuit returns a different value than the default", async () => {
+    // A plugin that completely replaces the default with a fixed value,
+    // proving plugins can substitute provider-native tokenizers (e.g.
+    // `countTokens`) without touching orchestrator code.
+    const FIXED = 424242;
+    const override: Middleware<EstimateArgs, EstimateResult> = async (
+      _args,
+      _next,
+      _ctx,
+    ) => FIXED;
+    const customPlugin: Plugin = {
+      manifest: {
+        name: "custom-token-estimate",
+        version: "1.0.0",
+        requires: { pluginRuntime: "v1", tokenEstimateApi: "v1" },
+      },
+      middleware: { tokenEstimate: override },
+    };
+
+    // Register the custom plugin FIRST so it sits outermost and short-
+    // circuits before the default's terminal runs.
+    registerPlugin(customPlugin);
+    registerDefault();
+
+    const args: EstimateArgs = {
+      history: TEXT_HISTORY,
+      systemPrompt: SYSTEM_PROMPT,
+      tools: SAMPLE_TOOLS,
+      providerName: "anthropic",
+    };
+    const pipelineResult = await runViaPipeline(args);
+    expect(pipelineResult).toBe(FIXED);
+    // And for contrast: the default alone would have given the raw value.
+    expect(pipelineResult).not.toBe(rawEstimate(args));
+  });
+
+  test("wrapper middleware that scales the downstream result composes with the default", async () => {
+    // A plugin that wraps the downstream estimate, doubling it. This
+    // exercises the onion composition: outer middleware sees the raw
+    // default result and returns its own modification.
+    const doubler: Middleware<EstimateArgs, EstimateResult> = async (
+      args,
+      next,
+      _ctx,
+    ) => {
+      const inner = await next(args);
+      return inner * 2;
+    };
+    const wrapperPlugin: Plugin = {
+      manifest: {
+        name: "doubling-token-estimate",
+        version: "1.0.0",
+        requires: { pluginRuntime: "v1", tokenEstimateApi: "v1" },
+      },
+      middleware: { tokenEstimate: doubler },
+    };
+
+    registerPlugin(wrapperPlugin);
+    registerDefault();
+
+    const args: EstimateArgs = {
+      history: TEXT_HISTORY,
+      systemPrompt: SYSTEM_PROMPT,
+      tools: SAMPLE_TOOLS,
+      providerName: "anthropic",
+    };
+    const pipelineResult = await runViaPipeline(args);
+    expect(pipelineResult).toBe(rawEstimate(args) * 2);
+  });
+});
+
+describe("tokenEstimate pipeline — default does not shadow late plugins", () => {
+  test("user middleware registered AFTER the default still runs", async () => {
+    // Regression test for the default-first shadowing hazard: defaults are
+    // registered before user plugins in `bootstrapPlugins()`, putting the
+    // default at the OUTERMOST onion position. If the default middleware
+    // runs the estimate directly instead of calling `next(args)`, any user
+    // plugin loaded afterward is invisible. The default is a passthrough —
+    // this test fails loudly if that invariant ever regresses.
+    registerDefault();
+    const observed: EstimateArgs[] = [];
+    const observer: Middleware<EstimateArgs, EstimateResult> = async (
+      args,
+      next,
+      _ctx,
+    ) => {
+      observed.push(args);
+      // Return a sentinel so we can distinguish the observer's result from
+      // the default's output.
+      await next(args);
+      return 999_999;
+    };
+    const userPlugin: Plugin = {
+      manifest: {
+        name: "late-registered-observer",
+        version: "1.0.0",
+        requires: { pluginRuntime: "v1", tokenEstimateApi: "v1" },
+      },
+      middleware: { tokenEstimate: observer },
+    };
+    registerPlugin(userPlugin);
+
+    const args: EstimateArgs = {
+      history: TEXT_HISTORY,
+      systemPrompt: SYSTEM_PROMPT,
+      tools: [],
+      providerName: "anthropic",
+    };
+    const result = await runViaPipeline(args);
+    expect(observed.length).toBe(1);
+    expect(result).toBe(999_999);
+  });
+});
+
+describe("tokenEstimate pipeline — args are immutable to middleware", () => {
+  test("frozen history/tools reject in-place mutation attempts", () => {
+    // The call site freezes shallow clones of `history` and `tools` before
+    // handing them to the pipeline. This mirrors the runtime protection
+    // that stops a misbehaving middleware from trimming `args.history` in
+    // place — which would silently drop prompt context from the
+    // orchestrator's live `runMessages` array before the provider call.
+    const frozenHistory = Object.freeze([...TEXT_HISTORY]);
+    const frozenTools = Object.freeze([...SAMPLE_TOOLS]);
+    expect(() => {
+      (frozenHistory as Message[]).pop();
+    }).toThrow(TypeError);
+    expect(() => {
+      (frozenTools as ToolDefinition[]).push({
+        name: "extra",
+        description: "",
+        input_schema: { type: "object", properties: {} },
+      });
+    }).toThrow(TypeError);
+  });
+});
+
+describe("tokenEstimate pipeline — empty registry fallback", () => {
+  test("without any plugin registered, the terminal receives the call", async () => {
+    // `runViaPipeline` uses a throwing terminal, so here we run the
+    // pipeline with an explicit terminal that returns a sentinel to prove
+    // that an empty middleware list falls through.
+    const SENTINEL = 12345;
+    const result = await runPipeline<EstimateArgs, EstimateResult>(
+      "tokenEstimate",
+      getMiddlewaresFor("tokenEstimate"),
+      async () => SENTINEL,
+      {
+        history: TEXT_HISTORY,
+        systemPrompt: SYSTEM_PROMPT,
+        tools: [],
+        providerName: "anthropic",
+      },
+      makeCtx(),
+      DEFAULT_TIMEOUTS.tokenEstimate,
+    );
+    expect(result).toBe(SENTINEL);
+  });
+});

package/src/__tests__/tool-error-pipeline.test.ts

@@ -0,0 +1,244 @@
+/**
+ * Tests for the `toolError` pipeline (PR 19).
+ *
+ * Covers:
+ * - Default plugin nudges on the first error turn and keeps nudging up to the
+ *   `maxConsecutiveErrorNudges` cap.
+ * - Default plugin suppresses the nudge once the cap is exceeded (the error is
+ *   likely unrecoverable — burning tokens on more nudges is wasteful).
+ * - Default plugin uses the canonical {@link DEFAULT_TOOL_ERROR_NUDGE_TEXT}.
+ * - Default plugin skips when `hasToolError` is false, regardless of the
+ *   consecutive counter (no error this turn → nothing to nudge).
+ * - Swapping in a user plugin that provides its own `toolError` middleware
+ *   changes the nudge text end-to-end through `runPipeline`.
+ */
+
+import { beforeEach, describe, expect, test } from "bun:test";
+
+import type { TrustContext } from "../daemon/conversation-runtime-assembly.js";
+import {
+  DEFAULT_TOOL_ERROR_NUDGE_TEXT,
+  defaultToolErrorPlugin,
+  defaultToolErrorTerminal,
+} from "../plugins/defaults/tool-error.js";
+import { runPipeline } from "../plugins/pipeline.js";
+import {
+  getMiddlewaresFor,
+  registerPlugin,
+  resetPluginRegistryForTests,
+} from "../plugins/registry.js";
+import {
+  type Middleware,
+  type Plugin,
+  type ToolErrorArgs,
+  type ToolErrorDecision,
+  type TurnContext,
+} from "../plugins/types.js";
+
+const trust: TrustContext = {
+  sourceChannel: "vellum",
+  trustClass: "guardian",
+};
+
+function makeCtx(): TurnContext {
+  return {
+    requestId: "req-tool-error-test",
+    conversationId: "conv-tool-error-test",
+    turnIndex: 0,
+    trust,
+  };
+}
+
+async function runToolErrorPipeline(
+  args: ToolErrorArgs,
+): Promise<ToolErrorDecision> {
+  // Mirror the production call site in `agent/loop.ts`: the pipeline terminal
+  // is `defaultToolErrorTerminal`, not a no-op. The default plugin's
+  // middleware is a passthrough that calls `next(args)`, so the decision
+  // logic lives in the terminal.
+  return runPipeline<ToolErrorArgs, ToolErrorDecision>(
+    "toolError",
+    getMiddlewaresFor("toolError"),
+    async (pipelineArgs) => defaultToolErrorTerminal(pipelineArgs),
+    args,
+    makeCtx(),
+    500,
+  );
+}
+
+describe("toolError pipeline", () => {
+  describe("default plugin", () => {
+    beforeEach(() => {
+      resetPluginRegistryForTests();
+      registerPlugin(defaultToolErrorPlugin);
+    });
+
+    test("nudges on first error turn with canonical text", async () => {
+      const decision = await runToolErrorPipeline({
+        hasToolError: true,
+        consecutiveErrorTurns: 1,
+        maxConsecutiveErrorNudges: 3,
+      });
+      expect(decision.action).toBe("nudge");
+      if (decision.action === "nudge") {
+        expect(decision.nudgeText).toBe(DEFAULT_TOOL_ERROR_NUDGE_TEXT);
+      }
+    });
+
+    test("keeps nudging up to and including the cap", async () => {
+      // Cap of 3: turns 1, 2, and 3 all nudge. Turn 4 is past the cap.
+      for (let turn = 1; turn <= 3; turn++) {
+        const decision = await runToolErrorPipeline({
+          hasToolError: true,
+          consecutiveErrorTurns: turn,
+          maxConsecutiveErrorNudges: 3,
+        });
+        expect(decision.action).toBe("nudge");
+      }
+    });
+
+    test("suppresses the nudge once the consecutive counter exceeds the cap", async () => {
+      const decision = await runToolErrorPipeline({
+        hasToolError: true,
+        consecutiveErrorTurns: 4,
+        maxConsecutiveErrorNudges: 3,
+      });
+      expect(decision.action).toBe("skip");
+    });
+
+    test("skips when there is no tool error this turn, regardless of counter", async () => {
+      // Counter is non-zero (the previous turn errored) but this turn succeeded,
+      // so nothing to nudge about.
+      const decision = await runToolErrorPipeline({
+        hasToolError: false,
+        consecutiveErrorTurns: 2,
+        maxConsecutiveErrorNudges: 3,
+      });
+      expect(decision.action).toBe("skip");
+    });
+
+    test("honors a caller-supplied cap of zero (never nudges)", async () => {
+      // Some call-sites may want to disable nudging entirely by passing cap = 0.
+      // The decision logic uses `<=`, so counter 0 with cap 0 does nudge; counter
+      // 1 with cap 0 suppresses. The cap is inclusive.
+      const turn1 = await runToolErrorPipeline({
+        hasToolError: true,
+        consecutiveErrorTurns: 1,
+        maxConsecutiveErrorNudges: 0,
+      });
+      expect(turn1.action).toBe("skip");
+    });
+  });
+
+  describe("user-supplied plugin", () => {
+    beforeEach(() => {
+      resetPluginRegistryForTests();
+    });
+
+    test("swapping in a plugin changes the nudge text", async () => {
+      const customText = "<system_notice>Custom error hint.</system_notice>";
+      const customMiddleware: Middleware<
+        ToolErrorArgs,
+        ToolErrorDecision
+      > = async (args) => {
+        if (args.hasToolError) {
+          return { action: "nudge", nudgeText: customText };
+        }
+        return { action: "skip" };
+      };
+      const customPlugin: Plugin = {
+        manifest: {
+          name: "custom-tool-error",
+          version: "0.0.1",
+          requires: { pluginRuntime: "v1", toolErrorApi: "v1" },
+        },
+        middleware: { toolError: customMiddleware },
+      };
+      registerPlugin(customPlugin);
+
+      const decision = await runToolErrorPipeline({
+        hasToolError: true,
+        consecutiveErrorTurns: 1,
+        maxConsecutiveErrorNudges: 3,
+      });
+      expect(decision.action).toBe("nudge");
+      if (decision.action === "nudge") {
+        expect(decision.nudgeText).toBe(customText);
+      }
+    });
+
+    test("swapping in a plugin can suppress nudges even when the default would nudge", async () => {
+      const suppressingMiddleware: Middleware<
+        ToolErrorArgs,
+        ToolErrorDecision
+      > = async () => ({ action: "skip" });
+      const plugin: Plugin = {
+        manifest: {
+          name: "no-nudge",
+          version: "0.0.1",
+          requires: { pluginRuntime: "v1", toolErrorApi: "v1" },
+        },
+        middleware: { toolError: suppressingMiddleware },
+      };
+      registerPlugin(plugin);
+
+      const decision = await runToolErrorPipeline({
+        hasToolError: true,
+        consecutiveErrorTurns: 1,
+        maxConsecutiveErrorNudges: 3,
+      });
+      expect(decision.action).toBe("skip");
+    });
+
+    test("terminal still produces the legacy nudge when no plugin is registered", async () => {
+      // No registerPlugin call — the registry is empty for this slot. Since
+      // `agent/loop.ts` now passes `defaultToolErrorTerminal` as the pipeline
+      // terminal (rather than an inline `() => skip`), direct AgentLoop
+      // callers that skip `bootstrapPlugins()` still get the legacy nudge.
+      const decision = await runToolErrorPipeline({
+        hasToolError: true,
+        consecutiveErrorTurns: 1,
+        maxConsecutiveErrorNudges: 3,
+      });
+      expect(decision.action).toBe("nudge");
+      if (decision.action === "nudge") {
+        expect(decision.nudgeText).toBe(DEFAULT_TOOL_ERROR_NUDGE_TEXT);
+      }
+    });
+
+    test("user plugin registered AFTER the default still runs (no shadowing)", async () => {
+      // Production registration order: defaults load first via the side-effect
+      // imports in `defaults/index.ts`, then user plugins register on top via
+      // `bootstrapPlugins()`. The user's middleware ends up at a deeper onion
+      // layer than the default. If the default's middleware were to bypass
+      // `next` and call the decision logic directly, the user middleware
+      // would never run — this test guards against that regression.
+      registerPlugin(defaultToolErrorPlugin);
+
+      let userMiddlewareRan = false;
+      const userMiddleware: Middleware<
+        ToolErrorArgs,
+        ToolErrorDecision
+      > = async (args, next) => {
+        userMiddlewareRan = true;
+        return next(args);
+      };
+      registerPlugin({
+        manifest: {
+          name: "late-user-plugin",
+          version: "0.0.1",
+          requires: { pluginRuntime: "v1", toolErrorApi: "v1" },
+        },
+        middleware: { toolError: userMiddleware },
+      });
+
+      await runToolErrorPipeline({
+        hasToolError: true,
+        consecutiveErrorTurns: 1,
+        maxConsecutiveErrorNudges: 3,
+      });
+
+      expect(userMiddlewareRan).toBe(true);
+    });
+  });
+});