@vellumai/assistant 0.12.2 → 0.12.3-staging.2

package/src/mcp/startup.ts

@@ -20,8 +20,14 @@ import { registerMcpTools, unregisterAllMcpTools } from "../tools/registry.js";
 import { getLogger } from "../util/logger.js";
 import { buildEffectiveMcpConfig } from "./effective-config.js";
 import { getMcpServerManager, stopMcpServerManager } from "./manager.js";
+import {
+  EMPTY_MCP_STARTUP_SNAPSHOT,
+  type McpStartupSnapshot,
+} from "./tool-caps.js";
 import { loadWorkspaceMcpConfig } from "./workspace-mcp-config.js";
 
+export type { McpStartupSnapshot } from "./tool-caps.js";
+
 const log = getLogger("mcp-startup");
 
 /**
@@ -39,23 +45,27 @@ const log = getLogger("mcp-startup");
  *
  * @param workspaceMcpConfig The projected workspace MCP config. Callers
  * that omit it read `/workspace/mcp.json`.
- * @returns The number of tools registered across all servers.
+ * @returns A count snapshot of configured/connected servers and the tools
+ * that reached this process's registry after caps.
  */
 export async function startConfiguredMcpServers(
   workspaceMcpConfig?: McpConfig,
-): Promise<number> {
-  let registered = 0;
+): Promise<McpStartupSnapshot> {
+  let snapshot: McpStartupSnapshot = EMPTY_MCP_STARTUP_SNAPSHOT;
   try {
     const mcpConfig = buildEffectiveMcpConfig(
       workspaceMcpConfig ?? loadWorkspaceMcpConfig(),
     );
-    if (Object.keys(mcpConfig.servers).length === 0) {
-      return 0;
+    const configuredServerCount = Object.keys(mcpConfig.servers).length;
+    snapshot = { ...EMPTY_MCP_STARTUP_SNAPSHOT, configuredServerCount };
+    if (configuredServerCount === 0) {
+      return snapshot;
     }
 
     const manager = getMcpServerManager();
-    const serverToolInfos = await manager.start(mcpConfig);
-    for (const { serverId, serverConfig, tools } of serverToolInfos) {
+    const started = await manager.start(mcpConfig);
+    let registered = 0;
+    for (const { serverId, serverConfig, tools } of started.servers) {
       const mcpTools = createMcpToolsFromServer(
         tools,
         serverId,
@@ -64,16 +74,26 @@ export async function startConfiguredMcpServers(
       );
       registered += registerMcpTools(serverId, mcpTools).length;
     }
+    snapshot = {
+      configuredServerCount: started.configuredServerCount,
+      connectedServerCount: started.connectedServerCount,
+      discoveredToolCount: started.discoveredToolCount,
+      registeredToolCount: registered,
+      droppedToolCount: started.droppedToolCount,
+      truncatedServerIds: started.truncatedServerIds,
+      errorServerCount: started.errorServerCount,
+      needsAuthServerCount: started.needsAuthServerCount,
+    };
   } catch (err) {
     log.error(
-      { err },
+      { err, ...snapshot },
       "MCP server initialization failed, continuing without MCP tools",
     );
-    return registered;
+    return snapshot;
   }
 
-  log.info({ toolCount: registered }, "MCP tools registered");
-  return registered;
+  log.info(snapshot, "MCP tools registered");
+  return snapshot;
 }
 
 /**
@@ -91,9 +111,9 @@ export async function startConfiguredMcpServers(
  * The window between the teardown and the reconnect has no MCP tools
  * registered, so an `mcp__*` call landing inside it fails as an unknown tool.
  *
- * @returns The number of tools registered after the reconnect.
+ * @returns A count snapshot of the tools registered after the reconnect.
  */
-export async function restartConfiguredMcpServers(): Promise<number> {
+export async function restartConfiguredMcpServers(): Promise<McpStartupSnapshot> {
   try {
     await stopMcpServerManager();
   } catch (err) {

package/src/mcp/tool-caps.ts

@@ -0,0 +1,184 @@
+/**
+ * Code-owned MCP tool-count caps and the selection used when they bind.
+ *
+ * The shipped global cap is {@link MCP_GLOBAL_MAX_TOOLS}. A workspace may
+ * raise or lower it with `tools.mcpGlobalMaxTools` in config.json. The
+ * per-server cap stays code-owned. These numbers are a context-window
+ * budget, not a permission gate. Allowlists and risk still decide what
+ * a turn may call.
+ *
+ * The global cap is applied by a deterministic round-robin across
+ * servers sorted by id. Insertion order must not empty a later server
+ * while earlier ones keep every tool.
+ */
+
+import {
+  MCP_GLOBAL_MAX_TOOLS,
+  MCP_MAX_TOOLS_PER_SERVER,
+} from "../config/schemas/mcp.js";
+import type { ToolsConfig } from "../config/schemas/tools.js";
+
+export interface McpServerToolSet<T> {
+  serverId: string;
+  tools: readonly T[];
+}
+
+export interface McpToolCapDecision {
+  serverId: string;
+  discovered: number;
+  kept: number;
+  droppedByPerServerCap: number;
+  droppedByGlobalCap: number;
+}
+
+export interface McpToolCapResult<T> {
+  servers: Array<{ serverId: string; tools: T[] }>;
+  discoveredToolCount: number;
+  keptToolCount: number;
+  droppedToolCount: number;
+  globalCap: number;
+  perServerCap: number;
+  decisions: McpToolCapDecision[];
+}
+
+export function resolveMcpGlobalMaxTools(
+  tools?: Pick<ToolsConfig, "mcpGlobalMaxTools">,
+): number {
+  return tools?.mcpGlobalMaxTools ?? MCP_GLOBAL_MAX_TOOLS;
+}
+
+export function applyMcpToolCaps<T>(
+  servers: ReadonlyArray<McpServerToolSet<T>>,
+  options?: { globalMax?: number; perServerMax?: number },
+): McpToolCapResult<T> {
+  const globalCap = options?.globalMax ?? MCP_GLOBAL_MAX_TOOLS;
+  const perServerCap = options?.perServerMax ?? MCP_MAX_TOOLS_PER_SERVER;
+  const discoveredToolCount = servers.reduce(
+    (sum, server) => sum + server.tools.length,
+    0,
+  );
+
+  const afterPerServer = servers.map((server) => {
+    const droppedByPerServerCap = Math.max(
+      0,
+      server.tools.length - perServerCap,
+    );
+    return {
+      serverId: server.serverId,
+      tools: server.tools.slice(0, perServerCap),
+      discovered: server.tools.length,
+      droppedByPerServerCap,
+    };
+  });
+
+  const afterPerServerCount = afterPerServer.reduce(
+    (sum, server) => sum + server.tools.length,
+    0,
+  );
+  if (afterPerServerCount <= globalCap) {
+    const decisions = afterPerServer.map((server) => ({
+      serverId: server.serverId,
+      discovered: server.discovered,
+      kept: server.tools.length,
+      droppedByPerServerCap: server.droppedByPerServerCap,
+      droppedByGlobalCap: 0,
+    }));
+    return {
+      servers: afterPerServer.map((server) => ({
+        serverId: server.serverId,
+        tools: [...server.tools],
+      })),
+      discoveredToolCount,
+      keptToolCount: afterPerServerCount,
+      droppedToolCount: discoveredToolCount - afterPerServerCount,
+      globalCap,
+      perServerCap,
+      decisions,
+    };
+  }
+
+  const order = [...afterPerServer].sort((a, b) =>
+    a.serverId.localeCompare(b.serverId),
+  );
+  const selected = new Map<string, T[]>(
+    order.map((server) => [server.serverId, []]),
+  );
+  const cursors = new Map<string, number>(
+    order.map((server) => [server.serverId, 0]),
+  );
+  let kept = 0;
+  let progressed = true;
+  while (kept < globalCap && progressed) {
+    progressed = false;
+    for (const server of order) {
+      if (kept >= globalCap) {
+        break;
+      }
+      const cursor = cursors.get(server.serverId) ?? 0;
+      const tool = server.tools[cursor];
+      if (tool !== undefined) {
+        selected.get(server.serverId)?.push(tool);
+        cursors.set(server.serverId, cursor + 1);
+        kept += 1;
+        progressed = true;
+      }
+    }
+  }
+
+  const decisions = afterPerServer.map((server) => {
+    const keptTools = selected.get(server.serverId) ?? [];
+    return {
+      serverId: server.serverId,
+      discovered: server.discovered,
+      kept: keptTools.length,
+      droppedByPerServerCap: server.droppedByPerServerCap,
+      droppedByGlobalCap: server.tools.length - keptTools.length,
+    };
+  });
+
+  return {
+    servers: afterPerServer.map((server) => ({
+      serverId: server.serverId,
+      tools: [...(selected.get(server.serverId) ?? [])],
+    })),
+    discoveredToolCount,
+    keptToolCount: kept,
+    droppedToolCount: discoveredToolCount - kept,
+    globalCap,
+    perServerCap,
+    decisions,
+  };
+}
+
+export interface McpStartupSnapshot {
+  configuredServerCount: number;
+  connectedServerCount: number;
+  discoveredToolCount: number;
+  registeredToolCount: number;
+  droppedToolCount: number;
+  truncatedServerIds: string[];
+  errorServerCount: number;
+  needsAuthServerCount: number;
+}
+
+export const EMPTY_MCP_STARTUP_SNAPSHOT: McpStartupSnapshot = {
+  configuredServerCount: 0,
+  connectedServerCount: 0,
+  discoveredToolCount: 0,
+  registeredToolCount: 0,
+  droppedToolCount: 0,
+  truncatedServerIds: [],
+  errorServerCount: 0,
+  needsAuthServerCount: 0,
+};
+
+export function truncatedServerIdsFromCaps(
+  decisions: readonly McpToolCapDecision[],
+): string[] {
+  return decisions
+    .filter(
+      (decision) =>
+        decision.droppedByPerServerCap > 0 || decision.droppedByGlobalCap > 0,
+    )
+    .map((decision) => decision.serverId);
+}

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

@@ -2,9 +2,12 @@ import { describe, expect, test } from "bun:test";
 
 import {
   DEFAULT_IMAGE_MODEL,
+  DEFAULT_OPENROUTER_IMAGE_MODEL,
   describeImageModels,
   IMAGE_MODELS,
+  qualifyImageModelForOpenRouter,
   resolveImageModel,
+  resolveRequestedImageModel,
 } from "../image-models.js";
 import { providerForModel } from "../image-service.js";
 
@@ -55,3 +58,58 @@ describe("image model registry", () => {
     }
   });
 });
+
+describe("resolveRequestedImageModel", () => {
+  test("qualifies built-in aliases for OpenRouter", () => {
+    expect(resolveRequestedImageModel("fast", "openrouter")).toEqual({
+      model: "google/gemini-3.1-flash-image-preview",
+    });
+    expect(resolveRequestedImageModel("openai", "openrouter")).toEqual({
+      model: "openai/gpt-image-2",
+    });
+    expect(qualifyImageModelForOpenRouter("quality")).toBe(
+      "google/gemini-3-pro-image-preview",
+    );
+  });
+
+  test("qualifies bare built-in IDs for OpenRouter", () => {
+    expect(
+      resolveRequestedImageModel("gemini-3.1-flash-image-preview", "openrouter"),
+    ).toEqual({ model: DEFAULT_OPENROUTER_IMAGE_MODEL });
+    expect(resolveRequestedImageModel("gpt-image-2", "openrouter")).toEqual({
+      model: "openai/gpt-image-2",
+    });
+  });
+
+  test("accepts arbitrary trimmed OpenRouter slugs", () => {
+    expect(
+      resolveRequestedImageModel("  google/gemini-3.5-flash  ", "openrouter"),
+    ).toEqual({ model: "google/gemini-3.5-flash" });
+    expect(
+      resolveRequestedImageModel("openai/gpt-image-2", "openrouter"),
+    ).toEqual({ model: "openai/gpt-image-2" });
+  });
+
+  test("rejects unknown models for built-in providers", () => {
+    const gemini = resolveRequestedImageModel(
+      "google/gemini-3.5-flash",
+      "gemini",
+    );
+    expect(gemini.model).toBeUndefined();
+    expect(gemini.error).toContain("Unknown model");
+
+    const openai = resolveRequestedImageModel("black-forest-labs/flux", "openai");
+    expect(openai.model).toBeUndefined();
+    expect(openai.error).toContain("Unknown model");
+
+    expect(resolveRequestedImageModel("fast", "vellum")).toEqual({
+      model: "gemini-3.1-flash-image-preview",
+    });
+  });
+
+  test("ignores empty or non-string input", () => {
+    expect(resolveRequestedImageModel("", "openrouter")).toEqual({});
+    expect(resolveRequestedImageModel("   ", "openrouter")).toEqual({});
+    expect(resolveRequestedImageModel(undefined, "gemini")).toEqual({});
+  });
+});

package/src/media/image-credentials.ts

@@ -37,6 +37,9 @@ export function resolveImageGenRouting(
         : svc.model;
     return { backendProvider: providerForImageModelPrefix(model), managed };
   }
+  if (svc.provider === "openrouter") {
+    return { backendProvider: "openrouter", managed: false };
+  }
   return {
     backendProvider: providerForModel(
       modelOverride,
@@ -101,5 +104,7 @@ function providerKeyHint(provider: ImageGenProvider): string {
       return "No Gemini API key configured. Please set your Gemini API key in Settings → Models & Services.";
     case "openai":
       return "No OpenAI API key configured. Please set your OpenAI API key in Settings → Models & Services.";
+    case "openrouter":
+      return "No OpenRouter API key configured. Please set your OpenRouter API key in Settings → Models & Services.";
   }
 }

package/src/media/image-models.ts

@@ -55,6 +55,58 @@ export function resolveImageModel(model: string): ImageModelEntry | undefined {
   return IMAGE_MODELS.find((m) => m.alias === model || m.id === model);
 }
 
+const OPENROUTER_VENDOR_PREFIX: Record<ImageModelProvider, string> = {
+  gemini: "google",
+  openai: "openai",
+};
+
+export const DEFAULT_OPENROUTER_IMAGE_MODEL =
+  "google/gemini-3.1-flash-image-preview";
+
+/**
+ * Qualify a built-in alias or concrete ID for OpenRouter's Images API.
+ * Unknown values are returned trimmed so live catalog slugs pass through.
+ */
+export function qualifyImageModelForOpenRouter(model: string): string {
+  const trimmed = model.trim();
+  const entry = resolveImageModel(trimmed);
+  if (!entry) {
+    return trimmed;
+  }
+  return `${OPENROUTER_VENDOR_PREFIX[entry.provider]}/${entry.id}`;
+}
+
+/**
+ * Resolve a caller-supplied image model against the configured provider.
+ *
+ * OpenRouter accepts any trimmed slug from its live catalog. Built-in
+ * aliases and bare IDs are qualified (`fast` ->
+ * `google/gemini-3.1-flash-image-preview`). Gemini, OpenAI, and Vellum
+ * still require a registry match.
+ */
+export function resolveRequestedImageModel(
+  model: unknown,
+  configuredProvider: string,
+): { model?: string; error?: string } {
+  if (typeof model !== "string" || !model) {
+    return {};
+  }
+  const trimmed = model.trim();
+  if (!trimmed) {
+    return {};
+  }
+  if (configuredProvider === "openrouter") {
+    return { model: qualifyImageModelForOpenRouter(trimmed) };
+  }
+  const entry = resolveImageModel(trimmed);
+  if (!entry) {
+    return {
+      error: `Unknown model "${trimmed}". Available models and aliases:\n${describeImageModels()}`,
+    };
+  }
+  return { model: entry.id };
+}
+
 /**
  * One line per model: "alias -> id (label)". Used in CLI help text and in
  * unknown-model error messages so the available set is always current.

package/src/media/image-service.ts

@@ -3,6 +3,10 @@ import {
   mapGeminiError,
 } from "./gemini-image-service.js";
 import { generateImageOpenAI, mapOpenAIError } from "./openai-image-service.js";
+import {
+  generateImageOpenRouter,
+  mapOpenRouterError,
+} from "./openrouter-image-service.js";
 import type {
   ImageGenCredentials,
   ImageGenerationRequest,
@@ -23,6 +27,8 @@ export function generateImage(
       return generateImageOpenAI(credentials, request);
     case "gemini":
       return generateImageGemini(credentials, request);
+    case "openrouter":
+      return generateImageOpenRouter(credentials, request);
     default: {
       const _exhaustive: never = provider;
       throw new Error(`Unknown image generation provider: ${_exhaustive}`);
@@ -42,6 +48,8 @@ export function mapImageGenError(
       return mapOpenAIError(error);
     case "gemini":
       return mapGeminiError(error);
+    case "openrouter":
+      return mapOpenRouterError(error);
     default: {
       const _exhaustive: never = provider;
       throw new Error(`Unknown image generation provider: ${_exhaustive}`);

package/src/media/openrouter-image-service.test.ts

@@ -0,0 +1,133 @@
+import { afterEach, describe, expect, mock, test } from "bun:test";
+
+import {
+  generateImageOpenRouter,
+  mapOpenRouterError,
+} from "./openrouter-image-service.js";
+
+const originalFetch = globalThis.fetch;
+
+afterEach(() => {
+  globalThis.fetch = originalFetch;
+});
+
+describe("generateImageOpenRouter", () => {
+  test("generates images through the dedicated OpenRouter Images API", async () => {
+    let request: Request | undefined;
+    globalThis.fetch = mock(async (input, init) => {
+      request = new Request(input, init);
+      return Response.json({
+        data: [
+          { b64_json: "png-data", media_type: "image/png" },
+          { b64_json: "svg-data", media_type: "image/svg+xml" },
+        ],
+      });
+    }) as unknown as typeof fetch;
+
+    const result = await generateImageOpenRouter(
+      { type: "direct", apiKey: "or-key" },
+      {
+        prompt: "a friendly robot",
+        mode: "generate",
+        model: "google/gemini-3.1-flash-image-preview",
+        variants: 2,
+      },
+    );
+
+    expect(request?.url).toBe("https://openrouter.ai/api/v1/images");
+    expect(request?.headers.get("authorization")).toBe("Bearer or-key");
+    expect(await request?.json()).toEqual({
+      model: "google/gemini-3.1-flash-image-preview",
+      prompt: "a friendly robot",
+      n: 2,
+    });
+    expect(result).toEqual({
+      images: [
+        { mimeType: "image/png", dataBase64: "png-data" },
+        { mimeType: "image/svg+xml", dataBase64: "svg-data" },
+      ],
+      text: undefined,
+      resolvedModel: "google/gemini-3.1-flash-image-preview",
+    });
+  });
+
+  test("sends source images as data URL references for edits", async () => {
+    let body: Record<string, unknown> | undefined;
+    globalThis.fetch = mock(async (_input, init) => {
+      body = JSON.parse(String(init?.body)) as Record<string, unknown>;
+      return Response.json({ data: [] });
+    }) as unknown as typeof fetch;
+
+    await generateImageOpenRouter(
+      { type: "direct", apiKey: "or-key" },
+      {
+        prompt: "make it watercolor",
+        mode: "edit",
+        sourceImages: [{ mimeType: "image/jpeg", dataBase64: "abc123" }],
+      },
+    );
+
+    expect(body?.input_references).toEqual([
+      {
+        type: "image_url",
+        image_url: { url: "data:image/jpeg;base64,abc123" },
+      },
+    ]);
+  });
+
+  test("qualifies a bare built-in model ID at the backend boundary", async () => {
+    let body: Record<string, unknown> | undefined;
+    globalThis.fetch = mock(async (_input, init) => {
+      body = JSON.parse(String(init?.body)) as Record<string, unknown>;
+      return Response.json({ data: [] });
+    }) as unknown as typeof fetch;
+
+    const result = await generateImageOpenRouter(
+      { type: "direct", apiKey: "or-key" },
+      {
+        prompt: "a cat",
+        mode: "generate",
+        model: "gemini-3.1-flash-image-preview",
+      },
+    );
+
+    expect(body?.model).toBe("google/gemini-3.1-flash-image-preview");
+    expect(result.resolvedModel).toBe("google/gemini-3.1-flash-image-preview");
+  });
+
+  test("defaults an empty model to the OpenRouter Gemini Flash slug", async () => {
+    let body: Record<string, unknown> | undefined;
+    globalThis.fetch = mock(async (_input, init) => {
+      body = JSON.parse(String(init?.body)) as Record<string, unknown>;
+      return Response.json({ data: [] });
+    }) as unknown as typeof fetch;
+
+    await generateImageOpenRouter(
+      { type: "direct", apiKey: "or-key" },
+      { prompt: "a cat", mode: "generate" },
+    );
+
+    expect(body?.model).toBe("google/gemini-3.1-flash-image-preview");
+  });
+
+  test("surfaces the provider's invalid-request message", async () => {
+    globalThis.fetch = mock(async () =>
+      Response.json(
+        { error: { message: "Model does not support image output" } },
+        { status: 400 },
+      ),
+    ) as unknown as typeof fetch;
+
+    try {
+      await generateImageOpenRouter(
+        { type: "direct", apiKey: "or-key" },
+        { prompt: "a cat", mode: "generate", model: "text-only/model" },
+      );
+      throw new Error("expected request to fail");
+    } catch (error) {
+      expect(mapOpenRouterError(error)).toBe(
+        "The OpenRouter image request was invalid: Model does not support image output",
+      );
+    }
+  });
+});