@vellumai/assistant 0.8.8 → 0.8.9-staging.2

package/src/__tests__/agent-loop-mutable-latest-user-message.test.ts

@@ -1,16 +1,17 @@
 /**
- * Verifies that the per-turn `mutableLatestUserMessage` plumbed into
- * `AgentLoop.run()` surfaces on every `SendMessageOptions.config` the loop
- * emits — this is the memory-v3-live cache-anchor signal (the latest user
- * message carries a volatile `<memory>` block, so the provider must anchor its
- * long-TTL cache breakpoint on the most recent STABLE message instead).
+ * Verifies that the `memory-v3-live` cache-anchor signal surfaces on every
+ * `SendMessageOptions.config` the loop emits. When the flag is on, the latest
+ * user message carries a volatile `<memory>` block, so the loop sets
+ * `providerConfig.mutableLatestUserMessage` to tell the provider to anchor its
+ * long-TTL cache breakpoint on the most recent STABLE message instead.
  *
- * Default behavior (option unset) must remain unchanged — the field is omitted
- * from `providerConfig` rather than carrying `undefined`/`false`, so the wire
- * is byte-identical to today when the flag is off.
+ * The loop reads the flag directly where it assembles `providerConfig`, so the
+ * signal is sourced from the flag rather than a run option. When the flag is
+ * off the field is omitted (not `false`/`undefined`) so the wire stays
+ * byte-identical to today.
  */
 
-import { describe, expect, mock, test } from "bun:test";
+import { afterEach, describe, expect, mock, test } from "bun:test";
 
 import { makeMockLogger } from "./helpers/mock-logger.js";
 
@@ -19,6 +20,10 @@ mock.module("../util/logger.js", () => ({
 }));
 
 import { AgentLoop } from "../agent/loop.js";
+import {
+  clearCachedOverrides,
+  setCachedOverrides,
+} from "../config/feature-flag-cache.js";
 import type {
   Message,
   Provider,
@@ -75,13 +80,20 @@ function makeRecordingProvider(responses: ProviderResponse[]): {
   return { provider, configs: () => configs };
 }
 
-describe("AgentLoop.run — mutableLatestUserMessage plumbing", () => {
-  test("forwards mutableLatestUserMessage=true to providerConfig on every LLM call (multi-turn)", async () => {
+describe("AgentLoop.run — mutableLatestUserMessage from memory-v3-live", () => {
+  afterEach(() => {
+    clearCachedOverrides();
+  });
+
+  test("sets mutableLatestUserMessage on every LLM call when memory-v3-live is on (multi-turn)", async () => {
+    // GIVEN memory-v3-live is enabled
+    setCachedOverrides({ "memory-v3-live": true }, { fromGateway: true });
+
+    // AND a provider that records the config of each LLM call across a tool round-trip
     const { provider, configs } = makeRecordingProvider([
       toolUseResponse("t1", "echo", { value: "first" }),
       textResponse("done"),
     ]);
-
     const dummyTools: ToolDefinition[] = [
       {
         name: "echo",
@@ -92,50 +104,53 @@ describe("AgentLoop.run — mutableLatestUserMessage plumbing", () => {
         },
       },
     ];
-
-    const toolExecutor = async () => ({ content: "ok", isError: false });
-
-    const loop = new AgentLoop(provider, "system", {
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
       config: { maxTokens: 1024 },
       tools: dummyTools,
-      toolExecutor: toolExecutor,
+      toolExecutor: async () => ({ content: "ok", isError: false }),
     });
 
-    await loop.run([userMessage], () => {}, {
+    // WHEN the loop runs
+    await loop.run({
+      messages: [userMessage],
+      onEvent: () => {},
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
       callSite: "mainAgent",
-      mutableLatestUserMessage: true,
     });
 
-    // Two sends — initial + one tool round-trip. Every send carries the flag.
+    // THEN every send (initial + tool round-trip) carries the cache-anchor signal
     expect(configs()).toHaveLength(2);
     for (const cfg of configs()) {
       expect(cfg?.mutableLatestUserMessage).toBe(true);
     }
   });
 
-  test("omits mutableLatestUserMessage when unset (default behavior unchanged)", async () => {
-    const { provider, configs } = makeRecordingProvider([textResponse("hi")]);
-    const loop = new AgentLoop(provider, "system", {
-      config: { maxTokens: 1024 },
-    });
-
-    await loop.run([userMessage], () => {});
+  test("omits mutableLatestUserMessage when memory-v3-live is off (flag-off byte-identity)", async () => {
+    // GIVEN memory-v3-live is off (no override; registry default is false)
+    clearCachedOverrides();
 
-    expect(configs()).toHaveLength(1);
-    expect("mutableLatestUserMessage" in (configs()[0] ?? {})).toBe(false);
-  });
-
-  test("omits mutableLatestUserMessage when explicitly false (flag-off byte-identity)", async () => {
+    // AND a provider that records the config of each LLM call
     const { provider, configs } = makeRecordingProvider([textResponse("hi")]);
-    const loop = new AgentLoop(provider, "system", {
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
       config: { maxTokens: 1024 },
     });
 
-    await loop.run([userMessage], () => {}, {
+    // WHEN the loop runs
+    await loop.run({
+      messages: [userMessage],
+      onEvent: () => {},
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
       callSite: "mainAgent",
-      mutableLatestUserMessage: false,
     });
 
+    // THEN the field is omitted entirely, not carried as false/undefined
+    expect(configs()).toHaveLength(1);
     expect("mutableLatestUserMessage" in (configs()[0] ?? {})).toBe(false);
   });
 });

package/src/__tests__/agent-loop-output-hooks.test.ts

@@ -0,0 +1,357 @@
+/**
+ * Tests for the `pre-model-call` and `assistant-message` plugin hooks: a plugin
+ * can edit the outbound request, transform the finalized assistant message
+ * (persisted + streamed), and defer the live stream so the transformed text is
+ * emitted once. All hooks here use neutral transforms (redaction / uppercasing).
+ */
+import { beforeEach, describe, expect, test } from "bun:test";
+
+import type { AgentEvent } from "../agent/loop.js";
+import { AgentLoop } from "../agent/loop.js";
+import type {
+  AssistantMessageContext,
+  PreModelCallContext,
+} from "../plugin-api/types.js";
+import { resetPluginRegistryAndRegisterDefaults } from "../plugins/defaults/index.js";
+import { registerPlugin } from "../plugins/registry.js";
+import type {
+  ContentBlock,
+  Message,
+  ProviderResponse,
+} from "../providers/types.js";
+import { createMockProvider, textResponse } from "./helpers/mock-provider.js";
+
+const userMessage: Message = {
+  role: "user",
+  content: [{ type: "text", text: "Hello" }],
+};
+
+function collect(events: AgentEvent[]): (event: AgentEvent) => void {
+  return (event) => events.push(event);
+}
+
+function streamedText(events: AgentEvent[]): string {
+  return events
+    .filter(
+      (e): e is Extract<AgentEvent, { type: "text_delta" }> =>
+        e.type === "text_delta",
+    )
+    .map((e) => e.text)
+    .join("");
+}
+
+function textOf(content: ReadonlyArray<ContentBlock>): string {
+  let out = "";
+  for (const block of content) if (block.type === "text") out += block.text;
+  return out;
+}
+
+function lastAssistant(history: Message[]): Message {
+  for (let i = history.length - 1; i >= 0; i--) {
+    if (history[i].role === "assistant") return history[i];
+  }
+  throw new Error("no assistant message in history");
+}
+
+function registerOutputHookPlugin(hooks: {
+  preModelCall?: (ctx: PreModelCallContext) => void;
+  assistantMessage?: (ctx: AssistantMessageContext) => void;
+}): void {
+  registerPlugin({
+    manifest: { name: "test-output-hooks", version: "0.0.0" },
+    hooks: {
+      ...(hooks.preModelCall
+        ? {
+            "pre-model-call": async (ctx: PreModelCallContext) => {
+              hooks.preModelCall!(ctx);
+            },
+          }
+        : {}),
+      ...(hooks.assistantMessage
+        ? {
+            "assistant-message": async (ctx: AssistantMessageContext) => {
+              hooks.assistantMessage!(ctx);
+            },
+          }
+        : {}),
+    },
+  });
+}
+
+describe("agent loop output hooks", () => {
+  beforeEach(() => {
+    resetPluginRegistryAndRegisterDefaults();
+  });
+
+  test("assistant-message transforms the persisted message content", async () => {
+    registerOutputHookPlugin({
+      assistantMessage: (ctx) => {
+        ctx.content = ctx.content.map((b) =>
+          b.type === "text"
+            ? { type: "text", text: b.text.replace("secret", "[redacted]") }
+            : b,
+        );
+      },
+    });
+    const { provider } = createMockProvider([textResponse("my secret value")]);
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+    });
+    const events: AgentEvent[] = [];
+    const { history } = await loop.run({
+      messages: [userMessage],
+      onEvent: collect(events),
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
+    expect(textOf(lastAssistant(history).content)).toBe("my [redacted] value");
+  });
+
+  test("deferred output: the real stream is suppressed and the transformed text is emitted once", async () => {
+    registerOutputHookPlugin({
+      preModelCall: (ctx) => {
+        ctx.deferAssistantOutput = true;
+      },
+      assistantMessage: (ctx) => {
+        ctx.content = [{ type: "text", text: "[filtered]" }];
+      },
+    });
+    const { provider } = createMockProvider([textResponse("my secret value")]);
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+    });
+    const events: AgentEvent[] = [];
+    const { history } = await loop.run({
+      messages: [userMessage],
+      onEvent: collect(events),
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
+    // Her real words never streamed; only the transformed text did, once.
+    expect(streamedText(events)).not.toContain("secret");
+    expect(streamedText(events)).toBe("[filtered]");
+    expect(textOf(lastAssistant(history).content)).toBe("[filtered]");
+  });
+
+  test("without defer, the real text still streams while storage is transformed", async () => {
+    registerOutputHookPlugin({
+      assistantMessage: (ctx) => {
+        ctx.content = [{ type: "text", text: "[stored]" }];
+      },
+    });
+    const { provider } = createMockProvider([textResponse("live text")]);
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+    });
+    const events: AgentEvent[] = [];
+    const { history } = await loop.run({
+      messages: [userMessage],
+      onEvent: collect(events),
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
+    expect(streamedText(events)).toBe("live text"); // streamed live, untransformed
+    expect(textOf(lastAssistant(history).content)).toBe("[stored]"); // storage transformed
+  });
+
+  test("transforms text but leaves tool_use blocks intact", async () => {
+    const toolAndText: ProviderResponse = {
+      content: [
+        { type: "text", text: "calling" },
+        { type: "tool_use", id: "t1", name: "noop", input: {} },
+      ],
+      model: "mock-model",
+      usage: { inputTokens: 1, outputTokens: 1 },
+      stopReason: "tool_use",
+    };
+    registerOutputHookPlugin({
+      assistantMessage: (ctx) => {
+        ctx.content = ctx.content.map((b) =>
+          b.type === "text" ? { type: "text", text: b.text.toUpperCase() } : b,
+        );
+      },
+    });
+    const { provider } = createMockProvider([
+      toolAndText,
+      textResponse("done"),
+    ]);
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+      tools: [
+        {
+          name: "noop",
+          description: "",
+          input_schema: { type: "object", properties: {} },
+        },
+      ],
+      toolExecutor: async () => ({ content: "ok", isError: false }),
+    });
+    const { history } = await loop.run({
+      messages: [userMessage],
+      onEvent: collect([]),
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
+    const toolTurn = history.find(
+      (m) =>
+        m.role === "assistant" && m.content.some((b) => b.type === "tool_use"),
+    )!;
+    expect(toolTurn.content.find((b) => b.type === "text")).toEqual({
+      type: "text",
+      text: "CALLING",
+    });
+    expect(toolTurn.content.find((b) => b.type === "tool_use")).toMatchObject({
+      type: "tool_use",
+      id: "t1",
+      name: "noop",
+    });
+  });
+
+  test("pre-model-call can edit the outbound system prompt", async () => {
+    registerOutputHookPlugin({
+      preModelCall: (ctx) => {
+        ctx.systemPrompt = `${ctx.systemPrompt ?? ""} [EDITED]`;
+      },
+    });
+    const { provider, calls } = createMockProvider([textResponse("hi")]);
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "base prompt",
+      conversationId: "test-conversation",
+    });
+    await loop.run({
+      messages: [userMessage],
+      onEvent: collect([]),
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
+    expect(calls[0].systemPrompt).toContain("[EDITED]");
+  });
+
+  test("fail-open: a hook that mutates in place AND then throws cannot corrupt the persisted content", async () => {
+    // The hook mutates the array it receives before throwing. If the loop
+    // handed the hook the real `assistantMessage.content` array, the
+    // mid-mutation would survive into history. The loop must clone the
+    // content before invoking the hook.
+    registerOutputHookPlugin({
+      assistantMessage: (ctx) => {
+        ctx.content.push({ type: "text", text: "[INJECTED]" });
+        throw new Error("boom");
+      },
+    });
+    const { provider } = createMockProvider([textResponse("untouched")]);
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+    });
+    const { history } = await loop.run({
+      messages: [userMessage],
+      onEvent: collect([]),
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
+    const finalContent = lastAssistant(history).content;
+    expect(textOf(finalContent)).toBe("untouched");
+    expect(
+      finalContent.some((b) => b.type === "text" && b.text === "[INJECTED]"),
+    ).toBe(false);
+  });
+
+  test("max_tokens turn: the hook fires and the deferred final text is emitted", async () => {
+    // Without this, an output-filter plugin misses the truncated reply; with
+    // defer set, the live stream was suppressed and the client would see
+    // nothing at all.
+    const seen: { calls: number } = { calls: 0 };
+    registerOutputHookPlugin({
+      preModelCall: (ctx) => {
+        ctx.deferAssistantOutput = true;
+      },
+      assistantMessage: (ctx) => {
+        seen.calls += 1;
+        ctx.content = ctx.content.map((b) =>
+          b.type === "text" ? { type: "text", text: b.text.toUpperCase() } : b,
+        );
+      },
+    });
+    const truncated: ProviderResponse = {
+      content: [{ type: "text", text: "partial answer" }],
+      model: "mock-model",
+      usage: { inputTokens: 1, outputTokens: 1 },
+      stopReason: "max_tokens",
+    };
+    const { provider } = createMockProvider([truncated]);
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+    });
+    const events: AgentEvent[] = [];
+    const { history } = await loop.run({
+      messages: [userMessage],
+      onEvent: collect(events),
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
+    expect(seen.calls).toBe(1);
+    expect(textOf(lastAssistant(history).content)).toBe("PARTIAL ANSWER");
+    // Real stream suppressed; transformed final text emitted once.
+    expect(streamedText(events)).toBe("PARTIAL ANSWER");
+  });
+
+  test("deferred final text passes through sensitive-output substitution", async () => {
+    // A tool returns a sensitive binding; the next assistant turn uses the
+    // placeholder (the persisted message must keep it — model never sees real
+    // values on reload), but the deferred final stream must show the real
+    // value, just like the normal live stream would.
+    registerOutputHookPlugin({
+      preModelCall: (ctx) => {
+        ctx.deferAssistantOutput = true;
+      },
+    });
+    const placeholder = "VELLUM_ASSISTANT_INVITE_CODE_TESTXXXX";
+    const real = "real-invite-code-xyz";
+    const toolThenText = [
+      {
+        content: [{ type: "tool_use", id: "t1", name: "issue", input: {} }],
+        model: "mock-model",
+        usage: { inputTokens: 1, outputTokens: 1 },
+        stopReason: "tool_use",
+      } as ProviderResponse,
+      textResponse(`Your code is ${placeholder}.`),
+    ];
+    const { provider } = createMockProvider(toolThenText);
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+      tools: [
+        {
+          name: "issue",
+          description: "",
+          input_schema: { type: "object", properties: {} },
+        },
+      ],
+      toolExecutor: async () => ({
+        content: `code=${placeholder}`,
+        isError: false,
+        sensitiveBindings: [{ kind: "invite_code", placeholder, value: real }],
+      }),
+    });
+    const events: AgentEvent[] = [];
+    const { history } = await loop.run({
+      messages: [userMessage],
+      onEvent: collect(events),
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
+    // Persisted message keeps the placeholder — model must never see real
+    // values on reload.
+    expect(textOf(lastAssistant(history).content)).toContain(placeholder);
+    expect(textOf(lastAssistant(history).content)).not.toContain(real);
+    // Streamed final text shows the substituted real value, matching the
+    // behavior of the normal live stream.
+    expect(streamedText(events)).toContain(real);
+    expect(streamedText(events)).not.toContain(placeholder);
+  });
+});

package/src/__tests__/agent-loop-override-profile.test.ts

@@ -110,13 +110,19 @@ describe("AgentLoop.run — overrideProfile plumbing", () => {
       _input: Record<string, unknown>,
     ) => ({ content: "ok", isError: false });
 
-    const loop = new AgentLoop(provider, "system", {
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
       config: { maxTokens: 1024 },
       tools: dummyTools,
       toolExecutor: toolExecutor,
     });
 
-    await loop.run([userMessage], () => {}, {
+    await loop.run({
+      messages: [userMessage],
+      onEvent: () => {},
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
       callSite: "mainAgent",
       overrideProfile: "fast",
     });
@@ -130,11 +136,18 @@ describe("AgentLoop.run — overrideProfile plumbing", () => {
 
   test("omits overrideProfile from providerConfig when unset (default behavior unchanged)", async () => {
     const { provider, configs } = makeRecordingProvider([textResponse("hi")]);
-    const loop = new AgentLoop(provider, "system", {
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
       config: { maxTokens: 1024 },
     });
 
-    await loop.run([userMessage], () => {});
+    await loop.run({
+      messages: [userMessage],
+      onEvent: () => {},
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
 
     // Single send, no overrideProfile field at all.
     expect(configs()).toHaveLength(1);
@@ -148,11 +161,17 @@ describe("AgentLoop.run — overrideProfile plumbing", () => {
     // receives so a non-existent profile silently falls back at the
     // provider layer (covered by provider-send-message-override-profile.test.ts).
     const { provider, configs } = makeRecordingProvider([textResponse("hi")]);
-    const loop = new AgentLoop(provider, "system", {
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
       config: { maxTokens: 1024 },
     });
 
-    await loop.run([userMessage], () => {}, {
+    await loop.run({
+      messages: [userMessage],
+      onEvent: () => {},
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
       callSite: "mainAgent",
       overrideProfile: "does-not-exist",
     });

package/src/__tests__/agent-loop-provider-error-recording.test.ts

@@ -83,14 +83,19 @@ describe("AgentLoop provider_error event emission", () => {
     const { provider, calls } = makeThrowingProvider("anthropic", () => thrown);
 
     const events: AgentEvent[] = [];
-    const loop = new AgentLoop(provider, "you are a helpful assistant");
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "you are a helpful assistant",
+      conversationId: "test-conversation",
+    });
 
-    await loop.run(
-      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
-      (e) => {
+    await loop.run({
+      messages: [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      onEvent: (e) => {
         events.push(e);
       },
-    );
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
 
     expect(calls).toHaveLength(1);
 
@@ -122,14 +127,19 @@ describe("AgentLoop provider_error event emission", () => {
     const { provider } = makeThrowingProvider("gemini", () => thrown);
 
     const events: AgentEvent[] = [];
-    const loop = new AgentLoop(provider, "system");
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+    });
 
-    await loop.run(
-      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
-      (e) => {
+    await loop.run({
+      messages: [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      onEvent: (e) => {
         events.push(e);
       },
-    );
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
 
     const providerErrorIdx = events.findIndex(
       (e) => e.type === "provider_error",
@@ -148,14 +158,19 @@ describe("AgentLoop provider_error event emission", () => {
     const { provider } = makeThrowingProvider("openai", () => thrown);
 
     const events: AgentEvent[] = [];
-    const loop = new AgentLoop(provider, "system");
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+    });
 
-    await loop.run(
-      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
-      (e) => {
+    await loop.run({
+      messages: [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      onEvent: (e) => {
         events.push(e);
       },
-    );
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+    });
 
     const providerErrorEvent = events.find((e) => e.type === "provider_error");
     expect(providerErrorEvent).toBeDefined();
@@ -179,15 +194,20 @@ describe("AgentLoop provider_error event emission", () => {
     });
 
     const events: AgentEvent[] = [];
-    const loop = new AgentLoop(provider, "system");
+    const loop = new AgentLoop({
+      provider: provider,
+      systemPrompt: "system",
+      conversationId: "test-conversation",
+    });
 
-    await loop.run(
-      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
-      (e) => {
+    await loop.run({
+      messages: [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      onEvent: (e) => {
         events.push(e);
       },
-      { signal: controller.signal },
-    );
+      trust: { sourceChannel: "vellum", trustClass: "unknown" },
+      signal: controller.signal,
+    });
 
     const providerErrorEvent = events.find((e) => e.type === "provider_error");
     // Cancellation should never produce a recording row — there's no