@vellumai/assistant 0.6.5 → 0.6.6

package/src/__tests__/llm-call-pipeline.test.ts

@@ -0,0 +1,285 @@
+/**
+ * Unit tests for the `llmCall` pipeline wrapping (PR 15).
+ *
+ * Exercises the three behaviors the plan calls out:
+ *
+ * 1. The default `llmCall` pipeline delegates to `provider.sendMessage(...)`
+ *    and returns its response unchanged.
+ * 2. A spy middleware registered for `llmCall` observes the full argument
+ *    payload before the provider is called.
+ * 3. A short-circuit middleware synthesizes a `ProviderResponse` and prevents
+ *    the real `provider.sendMessage` from running.
+ */
+
+import { afterAll, beforeEach, describe, expect, test } from "bun:test";
+
+import type { TrustContext } from "../daemon/conversation-runtime-assembly.js";
+import { defaultLlmCallPlugin } from "../plugins/defaults/llm-call.js";
+import { DEFAULT_TIMEOUTS, runPipeline } from "../plugins/pipeline.js";
+import {
+  getMiddlewaresFor,
+  registerPlugin,
+  resetPluginRegistryForTests,
+} from "../plugins/registry.js";
+import type {
+  LLMCallArgs,
+  LLMCallResult,
+  Middleware,
+  Plugin,
+  TurnContext,
+} from "../plugins/types.js";
+import type {
+  Message,
+  Provider,
+  ProviderResponse,
+  ToolDefinition,
+} from "../providers/types.js";
+
+// ─── Fixtures ───────────────────────────────────────────────────────────────
+
+const trust: TrustContext = {
+  sourceChannel: "vellum",
+  trustClass: "guardian",
+};
+
+function makeCtx(overrides: Partial<TurnContext> = {}): TurnContext {
+  return {
+    requestId: "req-test",
+    conversationId: "conv-test",
+    turnIndex: 0,
+    trust,
+    ...overrides,
+  };
+}
+
+function makeResponse(
+  overrides: Partial<ProviderResponse> = {},
+): ProviderResponse {
+  return {
+    content: [{ type: "text", text: "hello from provider" }],
+    model: "fake-model",
+    usage: {
+      inputTokens: 10,
+      outputTokens: 5,
+    },
+    stopReason: "end_turn",
+    ...overrides,
+  };
+}
+
+type FakeProviderCall = {
+  messages: Message[];
+  tools?: ToolDefinition[];
+  systemPrompt?: string;
+};
+
+function makeFakeProvider(
+  response: ProviderResponse = makeResponse(),
+): Provider & { calls: FakeProviderCall[] } {
+  const calls: FakeProviderCall[] = [];
+  return {
+    name: "fake-provider",
+    calls,
+    async sendMessage(messages, tools, systemPrompt, _options) {
+      calls.push({ messages, tools, systemPrompt });
+      return response;
+    },
+  };
+}
+
+function makeArgs(
+  provider: Provider,
+  overrides: Partial<LLMCallArgs> = {},
+): LLMCallArgs {
+  return {
+    provider,
+    messages: [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+    tools: undefined,
+    systemPrompt: "you are a helpful assistant",
+    options: { config: {} },
+    ...overrides,
+  };
+}
+
+// The terminal passed into `runPipeline` matches the one in `agent/loop.ts`:
+// it delegates straight to `args.provider.sendMessage(...)` with no
+// transformation. Keeping it identical here means the test exercises the
+// exact call shape the real loop uses.
+const terminal = (args: LLMCallArgs): Promise<LLMCallResult> =>
+  args.provider.sendMessage(
+    args.messages,
+    args.tools,
+    args.systemPrompt,
+    args.options,
+  );
+
+// ─── Tests ──────────────────────────────────────────────────────────────────
+
+describe("llmCall pipeline", () => {
+  beforeEach(() => {
+    resetPluginRegistryForTests();
+  });
+
+  // Clear the registry on the way out too so later test files in the same
+  // `bun test` run don't inherit `llmCall` middleware from our final test.
+  // Bun runs files sequentially within a process; `beforeEach` only clears
+  // at the start of each case, leaving whatever the final test registered
+  // in place for the next file.
+  afterAll(() => {
+    resetPluginRegistryForTests();
+  });
+
+  test("default pipeline invokes provider.sendMessage and returns its response", async () => {
+    registerPlugin(defaultLlmCallPlugin);
+
+    const expected = makeResponse({ model: "expected-model" });
+    const provider = makeFakeProvider(expected);
+    const args = makeArgs(provider);
+
+    const result = await runPipeline<LLMCallArgs, LLMCallResult>(
+      "llmCall",
+      getMiddlewaresFor("llmCall"),
+      terminal,
+      args,
+      makeCtx(),
+      DEFAULT_TIMEOUTS.llmCall,
+    );
+
+    expect(result).toBe(expected);
+    expect(provider.calls).toHaveLength(1);
+    expect(provider.calls[0]!.messages).toBe(args.messages);
+    expect(provider.calls[0]!.systemPrompt).toBe("you are a helpful assistant");
+  });
+
+  test("spy middleware records the full invocation arguments", async () => {
+    const observed: LLMCallArgs[] = [];
+    const spyPlugin: Plugin = {
+      manifest: {
+        name: "spy-llm",
+        version: "0.0.1",
+        requires: { pluginRuntime: "v1" },
+      },
+      middleware: {
+        llmCall: async (args, next, _ctx) => {
+          observed.push(args);
+          return next(args);
+        },
+      },
+    };
+
+    registerPlugin(spyPlugin);
+    registerPlugin(defaultLlmCallPlugin);
+
+    const provider = makeFakeProvider();
+    const tools: ToolDefinition[] = [
+      {
+        name: "echo",
+        description: "echoes its input",
+        input_schema: { type: "object" },
+      },
+    ];
+    const args = makeArgs(provider, { tools });
+
+    await runPipeline<LLMCallArgs, LLMCallResult>(
+      "llmCall",
+      getMiddlewaresFor("llmCall"),
+      terminal,
+      args,
+      makeCtx(),
+      DEFAULT_TIMEOUTS.llmCall,
+    );
+
+    expect(observed).toHaveLength(1);
+    expect(observed[0]!.provider).toBe(provider);
+    expect(observed[0]!.messages).toBe(args.messages);
+    expect(observed[0]!.tools).toBe(tools);
+    expect(observed[0]!.systemPrompt).toBe("you are a helpful assistant");
+    expect(provider.calls).toHaveLength(1);
+  });
+
+  test("default registered first does not shadow later-registered user middleware", async () => {
+    // Regression for a shadowing bug where `defaultLlmCallPlugin` called
+    // `provider.sendMessage` directly instead of `next(args)`. Because the
+    // default registers at module load (before `bootstrapPlugins()` loads
+    // user plugins), it sat at the outermost layer in production — any
+    // user-registered `llmCall` middleware would have been silently skipped.
+    // This test locks in the fix by registering the default FIRST (matching
+    // production ordering) and asserting a user-registered spy still runs.
+    const observed: LLMCallArgs[] = [];
+    const spyPlugin: Plugin = {
+      manifest: {
+        name: "spy-llm-after-default",
+        version: "0.0.1",
+        requires: { pluginRuntime: "v1" },
+      },
+      middleware: {
+        llmCall: async (args, next, _ctx) => {
+          observed.push(args);
+          return next(args);
+        },
+      },
+    };
+
+    registerPlugin(defaultLlmCallPlugin);
+    registerPlugin(spyPlugin);
+
+    const provider = makeFakeProvider();
+    const args = makeArgs(provider);
+
+    await runPipeline<LLMCallArgs, LLMCallResult>(
+      "llmCall",
+      getMiddlewaresFor("llmCall"),
+      terminal,
+      args,
+      makeCtx(),
+      DEFAULT_TIMEOUTS.llmCall,
+    );
+
+    expect(observed).toHaveLength(1);
+    expect(observed[0]!.provider).toBe(provider);
+    expect(provider.calls).toHaveLength(1);
+  });
+
+  test("short-circuit middleware prevents the real provider call", async () => {
+    const synthetic = makeResponse({
+      model: "synthetic-model",
+      content: [{ type: "text", text: "synthesized" }],
+    });
+
+    const shortCircuit: Middleware<LLMCallArgs, LLMCallResult> = async (
+      _args,
+      _next,
+      _ctx,
+    ) => synthetic;
+
+    const shortCircuitPlugin: Plugin = {
+      manifest: {
+        name: "short-circuit-llm",
+        version: "0.0.1",
+        requires: { pluginRuntime: "v1" },
+      },
+      middleware: { llmCall: shortCircuit },
+    };
+
+    registerPlugin(shortCircuitPlugin);
+    registerPlugin(defaultLlmCallPlugin);
+
+    const provider = makeFakeProvider();
+    const args = makeArgs(provider);
+
+    const result = await runPipeline<LLMCallArgs, LLMCallResult>(
+      "llmCall",
+      getMiddlewaresFor("llmCall"),
+      terminal,
+      args,
+      makeCtx(),
+      DEFAULT_TIMEOUTS.llmCall,
+    );
+
+    expect(result).toBe(synthetic);
+    // The short-circuit middleware never calls `next`, so the terminal and
+    // every downstream middleware (including the default) are skipped and
+    // the provider is never contacted.
+    expect(provider.calls).toHaveLength(0);
+  });
+});

package/src/__tests__/llm-schema.test.ts

@@ -70,7 +70,7 @@ describe("LLMSchema", () => {
     const parsed = LLMSchema.parse({});
     expect(parsed.default).toEqual({
       provider: "anthropic",
-      model: "claude-opus-4-7",
+      model: "claude-sonnet-4-6",
       maxTokens: 64000,
       effort: "max",
       speed: "standard",

package/src/__tests__/media-generate-image.test.ts

@@ -10,14 +10,17 @@ import type { ToolContext } from "../tools/types.js";
 // Mock dependencies for the tool wrapper
 // ---------------------------------------------------------------------------
 
-let mockApiKey: string | undefined = "test-gemini-key";
+let mockGeminiKey: string | undefined = "test-gemini-key";
+let mockOpenAIKey: string | undefined = "test-openai-key";
 let mockImageGenMode: "your-own" | "managed" = "your-own";
+let mockImageGenProvider: "gemini" | "openai" = "gemini";
 let mockGenerateResult = {
   images: [{ mimeType: "image/png", dataBase64: "generated-data" }],
   text: "A beautiful image",
   resolvedModel: "gemini-3.1-flash-image-preview",
 };
 let mockGenerateError: Error | null = null;
+let lastGenerateProvider: unknown = null;
 let lastGenerateCredentials: unknown = null;
 
 mock.module("../config/loader.js", () => ({
@@ -31,7 +34,7 @@ mock.module("../config/loader.js", () => ({
       },
       "image-generation": {
         mode: mockImageGenMode,
-        provider: "gemini",
+        provider: mockImageGenProvider,
         model: "gemini-3.1-flash-image-preview",
       },
       "web-search": { mode: "your-own", provider: "inference-provider-native" },
@@ -41,27 +44,33 @@ mock.module("../config/loader.js", () => ({
 
 mock.module("../security/secure-keys.js", () => ({
   getSecureKeyAsync: async (account: string) => {
-    if (account === "gemini") return mockApiKey;
+    if (account === "gemini") return mockGeminiKey;
+    if (account === "openai") return mockOpenAIKey;
     return undefined;
   },
   getProviderKeyAsync: async (provider: string) => {
-    if (provider === "gemini") return mockApiKey;
+    if (provider === "gemini") return mockGeminiKey;
+    if (provider === "openai") return mockOpenAIKey;
     return undefined;
   },
 }));
 
-mock.module("../media/gemini-image-service.js", () => ({
+mock.module("../media/image-service.js", () => ({
   generateImage: async (
+    provider: unknown,
     credentials: unknown,
     _request: Record<string, unknown>,
   ) => {
+    lastGenerateProvider = provider;
     lastGenerateCredentials = credentials;
     if (mockGenerateError) throw mockGenerateError;
     return mockGenerateResult;
   },
-  mapGeminiError: (error: unknown) => {
-    if (error instanceof Error) return `Mock error: ${error.message}`;
-    return "Mock unknown error";
+  mapImageGenError: (provider: unknown, error: unknown) => {
+    const providerLabel = provider === "openai" ? "OpenAI" : "Gemini";
+    if (error instanceof Error)
+      return `Mock ${providerLabel} error: ${error.message}`;
+    return `Mock ${providerLabel} unknown error`;
   },
 }));
 
@@ -97,14 +106,17 @@ const CONFIG_DIR = join(
 // ---------------------------------------------------------------------------
 
 beforeEach(() => {
-  mockApiKey = "test-gemini-key";
+  mockGeminiKey = "test-gemini-key";
+  mockOpenAIKey = "test-openai-key";
   mockImageGenMode = "your-own";
+  mockImageGenProvider = "gemini";
   mockGenerateResult = {
     images: [{ mimeType: "image/png", dataBase64: "generated-data" }],
     text: "A beautiful image",
     resolvedModel: "gemini-3.1-flash-image-preview",
   };
   mockGenerateError = null;
+  lastGenerateProvider = null;
   lastGenerateCredentials = null;
   mockManagedBaseUrl = undefined;
   mockManagedProxyContext = {
@@ -127,7 +139,7 @@ describe("image-studio skill script wrapper", () => {
   });
 
   test("returns error when no API key and no managed proxy", async () => {
-    mockApiKey = undefined;
+    mockGeminiKey = undefined;
 
     const result = await run({ prompt: "a cat" }, fakeContext);
 
@@ -148,6 +160,7 @@ describe("image-studio skill script wrapper", () => {
 
     expect(result.isError).toBe(false);
     expect(result.content).toContain("Generated 1 image");
+    expect(lastGenerateProvider).toBe("gemini");
     expect(lastGenerateCredentials).toEqual({
       type: "managed-proxy",
       assistantApiKey: "managed-key-123",
@@ -157,7 +170,7 @@ describe("image-studio skill script wrapper", () => {
 
   test("managed mode returns error when managed proxy is unavailable", async () => {
     mockImageGenMode = "managed";
-    mockApiKey = "direct-key"; // should be ignored in managed mode
+    mockGeminiKey = "direct-key"; // should be ignored in managed mode
     mockManagedBaseUrl = undefined;
 
     const result = await run({ prompt: "a cat" }, fakeContext);
@@ -168,7 +181,7 @@ describe("image-studio skill script wrapper", () => {
 
   test("your-own mode uses direct API key", async () => {
     mockImageGenMode = "your-own";
-    mockApiKey = "direct-key";
+    mockGeminiKey = "direct-key";
     mockManagedBaseUrl = "https://platform.example.com/v1/runtime-proxy/gemini";
     mockManagedProxyContext = {
       enabled: true,
@@ -178,12 +191,94 @@ describe("image-studio skill script wrapper", () => {
 
     await run({ prompt: "a cat" }, fakeContext);
 
+    expect(lastGenerateProvider).toBe("gemini");
     expect(lastGenerateCredentials).toEqual({
       type: "direct",
       apiKey: "direct-key",
     });
   });
 
+  test("openai provider dispatches to OpenAI with its key", async () => {
+    mockImageGenProvider = "openai";
+    mockOpenAIKey = "openai-direct-key";
+
+    const result = await run({ prompt: "a robot" }, fakeContext);
+
+    expect(result.isError).toBe(false);
+    expect(lastGenerateProvider).toBe("openai");
+    expect(lastGenerateCredentials).toEqual({
+      type: "direct",
+      apiKey: "openai-direct-key",
+    });
+  });
+
+  test("openai provider returns OpenAI-specific error hint when no key", async () => {
+    mockImageGenProvider = "openai";
+    mockOpenAIKey = undefined;
+
+    const result = await run({ prompt: "a robot" }, fakeContext);
+
+    expect(result.isError).toBe(true);
+    expect(result.content).toContain("OpenAI");
+    expect(result.content).not.toContain("No Gemini API key");
+  });
+
+  test("explicit model override routes to owning provider (gemini config → openai call)", async () => {
+    // Config says the user's default provider is gemini, but the LLM
+    // explicitly requests a gpt-* model. The tool must dispatch to OpenAI
+    // and resolve OpenAI credentials, not fall back to Gemini's default.
+    mockImageGenProvider = "gemini";
+    mockOpenAIKey = "openai-direct-key";
+
+    const result = await run(
+      { prompt: "a robot", model: "gpt-image-2" },
+      fakeContext,
+    );
+
+    expect(result.isError).toBe(false);
+    expect(lastGenerateProvider).toBe("openai");
+    expect(lastGenerateCredentials).toEqual({
+      type: "direct",
+      apiKey: "openai-direct-key",
+    });
+  });
+
+  test("explicit model override routes to owning provider (openai config → gemini call)", async () => {
+    // The inverse: config says openai but the LLM asks for a gemini-* model.
+    mockImageGenProvider = "openai";
+    mockGeminiKey = "gemini-direct-key";
+
+    const result = await run(
+      { prompt: "a cat", model: "gemini-3-pro-image-preview" },
+      fakeContext,
+    );
+
+    expect(result.isError).toBe(false);
+    expect(lastGenerateProvider).toBe("gemini");
+    expect(lastGenerateCredentials).toEqual({
+      type: "direct",
+      apiKey: "gemini-direct-key",
+    });
+  });
+
+  test("cross-provider override surfaces owning provider's credential error", async () => {
+    // Config: gemini (with a gemini key). LLM asks for gpt-image-2 but the
+    // OpenAI key is missing. The error hint must reference OpenAI, not
+    // Gemini, because the dispatch target is OpenAI.
+    mockImageGenProvider = "gemini";
+    mockGeminiKey = "test-gemini-key";
+    mockOpenAIKey = undefined;
+
+    const result = await run(
+      { prompt: "a robot", model: "gpt-image-2" },
+      fakeContext,
+    );
+
+    expect(result.isError).toBe(true);
+    expect(result.content).toContain("OpenAI");
+    expect(result.content).not.toContain("No Gemini API key");
+  });
+
   test("returns generated image with contentBlocks", async () => {
     const result = await run({ prompt: "a sunset" }, fakeContext);
 
@@ -225,7 +320,17 @@ describe("image-studio skill script wrapper", () => {
     const result = await run({ prompt: "a cat" }, fakeContext);
 
     expect(result.isError).toBe(true);
-    expect(result.content).toContain("Mock error: API failure");
+    expect(result.content).toContain("Mock Gemini error: API failure");
+  });
+
+  test("openai generation error uses OpenAI-specific mapping", async () => {
+    mockImageGenProvider = "openai";
+    mockGenerateError = new Error("openai failure");
+
+    const result = await run({ prompt: "a cat" }, fakeContext);
+
+    expect(result.isError).toBe(true);
+    expect(result.content).toContain("Mock OpenAI error: openai failure");
   });
 
   test("reads source images from file paths on disk", async () => {
@@ -330,6 +435,7 @@ describe("image-studio TOOLS.json manifest", () => {
     expect(props.model.enum).toEqual([
       "gemini-3.1-flash-image-preview",
       "gemini-3-pro-image-preview",
+      "gpt-image-2",
     ]);
     expect(props.variants.type).toBe("number");
   });