@vellumai/assistant 0.12.1 → 0.12.2-staging.2

package/src/providers/inference/resolve-auth.ts

@@ -6,7 +6,7 @@
  *   - platform             → build managed proxy URL and fetch the platform API key
  *   - none                 → pass through with no auth headers
  *   - oauth_subscription   → fetch OAuth token from vault (with auto-refresh) → inject as bearer header
- *   - service_account      → reject (v2 not yet shipped)
+ *   - service_account      → sign JWT and exchange for bearer token at token_uri (Google Vertex AI)
  */
 
 import { getSecureKeyAsync } from "../../security/secure-keys.js";
@@ -119,10 +119,28 @@ export async function resolveAuth(
       };
     }
 
-    case "service_account":
+    case "service_account": {
+      const { getValidServiceAccountToken } =
+        await import("./service-account-token.js");
+      const result = await getValidServiceAccountToken(auth.credential);
+      if (!result.ok) {
+        if (result.reason === "exchange_failed") {
+          return { ok: false, error: { code: "platform_unavailable" } };
+        }
+        // not_found or invalid_config: the credential needs to be supplied or fixed.
+        return {
+          ok: false,
+          error: { code: "credential_not_found", credential: auth.credential },
+        };
+      }
       return {
-        ok: false,
-        error: { code: "not_implemented", authType: auth.type },
+        ok: true,
+        resolved: {
+          kind: "header",
+          headers: { Authorization: `Bearer ${result.token}` },
+          ...(safeBaseUrl ? { baseUrl: safeBaseUrl } : {}),
+        },
       };
+    }
   }
 }

package/src/providers/inference/service-account-token.ts

@@ -0,0 +1,248 @@
+/**
+ * Token acquisition for Google service-account credentials (Vertex AI).
+ *
+ * Reads the service-account JSON from the vault, signs a short-lived JWT
+ * with RS256, and exchanges it for a bearer access token at the token_uri
+ * embedded in the key file. Caches the resulting token and its expiry as a
+ * single JSON blob in the vault so the write is treated as one unit.
+ *
+ * A per-credential mutex prevents concurrent callers for the same credential
+ * from racing to refresh. Callers for different credentials are independent.
+ */
+
+import { createSign } from "node:crypto";
+
+import {
+  getSecureKeyAsync as _realGetKey,
+  setSecureKeyAsync as _realSetKey,
+} from "../../security/secure-keys.js";
+import { getLogger } from "../../util/logger.js";
+
+// Mutable references — swapped in tests via _injectVaultAccessors so tests
+// never need mock.module (which contaminates the entire bun worker process).
+let _getKey: typeof _realGetKey = _realGetKey;
+let _setKey: typeof _realSetKey = _realSetKey;
+
+const log = getLogger("service-account-token");
+
+/** Refresh 5 minutes before expiry to avoid using a nearly-expired token. */
+const REFRESH_MARGIN_SECONDS = 300;
+
+/** Google Cloud Platform scope required for Vertex AI inference. */
+const GCP_SCOPE = "https://www.googleapis.com/auth/cloud-platform";
+
+export interface ServiceAccountKey {
+  client_email: string;
+  private_key: string;
+  token_uri: string;
+}
+
+export type ServiceAccountTokenResult =
+  | { ok: true; token: string }
+  | { ok: false; reason: "not_found" | "invalid_config" | "exchange_failed" };
+
+/**
+ * Parse and validate a service-account JSON string.
+ * Returns the key object when valid, or null when the JSON is missing, invalid,
+ * or lacks required fields (client_email, private_key, token_uri).
+ */
+export function parseServiceAccountKey(json: string): ServiceAccountKey | null {
+  let parsed: unknown;
+  try {
+    parsed = JSON.parse(json);
+  } catch {
+    return null;
+  }
+  if (!parsed || typeof parsed !== "object") {
+    return null;
+  }
+  const key = parsed as Record<string, unknown>;
+  if (
+    typeof key["client_email"] !== "string" ||
+    !key["client_email"] ||
+    typeof key["private_key"] !== "string" ||
+    !key["private_key"] ||
+    typeof key["token_uri"] !== "string" ||
+    !key["token_uri"]
+  ) {
+    return null;
+  }
+  // All required fields validated as non-empty strings above.
+  return parsed as ServiceAccountKey;
+}
+
+/** Per-credential in-flight mutex. Prevents concurrent exchange races per account. */
+const exchangesInFlight = new Map<string, Promise<ServiceAccountTokenResult>>();
+
+/**
+ * Return a valid bearer access token for a Google service account, fetching
+ * and caching one if the stored token is absent or about to expire.
+ *
+ * @param credential - Vault key under which the service-account JSON is stored
+ *   (e.g. `"credential/my-vertex-ai"`). The derived token is cached as a JSON
+ *   blob at `<credential>/token_cache`.
+ */
+export async function getValidServiceAccountToken(
+  credential: string,
+): Promise<ServiceAccountTokenResult> {
+  const cacheKey = `${credential}/token_cache`;
+  const cached = await _getKey(cacheKey);
+  if (cached) {
+    try {
+      const blob = JSON.parse(cached) as {
+        access_token?: string;
+        expires_at?: number;
+      };
+      if (blob.access_token) {
+        if (!blob.expires_at) {
+          return { ok: true, token: blob.access_token };
+        }
+        const now = Date.now() / 1000;
+        if (now < blob.expires_at - REFRESH_MARGIN_SECONDS) {
+          return { ok: true, token: blob.access_token };
+        }
+      }
+    } catch {
+      // Corrupted cache entry — fall through to re-exchange.
+    }
+  }
+
+  const inFlight = exchangesInFlight.get(credential);
+  if (inFlight) {
+    return await inFlight;
+  }
+
+  const promise = doExchange(credential, cacheKey);
+  exchangesInFlight.set(credential, promise);
+  try {
+    return await promise;
+  } finally {
+    exchangesInFlight.delete(credential);
+  }
+}
+
+async function doExchange(
+  credential: string,
+  cacheKey: string,
+): Promise<ServiceAccountTokenResult> {
+  const keyJson = await _getKey(credential);
+  if (!keyJson) {
+    return { ok: false, reason: "not_found" };
+  }
+
+  const key = parseServiceAccountKey(keyJson);
+  if (!key) {
+    log.error(
+      { credential },
+      "Service account credential is not valid JSON or missing required fields",
+    );
+    return { ok: false, reason: "invalid_config" };
+  }
+
+  const jwt = buildJwt(key.client_email, key.private_key, key.token_uri);
+
+  let accessToken: string;
+  let expiresIn: number;
+  try {
+    const resp = await fetch(key.token_uri, {
+      method: "POST",
+      headers: { "Content-Type": "application/x-www-form-urlencoded" },
+      body: new URLSearchParams({
+        grant_type: "urn:ietf:params:oauth:grant-type:jwt-bearer",
+        assertion: jwt,
+      }).toString(),
+    });
+
+    if (!resp.ok) {
+      const body = await resp.text();
+      log.error(
+        { status: resp.status, body, credential },
+        "Service account token exchange failed",
+      );
+      return { ok: false, reason: "exchange_failed" };
+    }
+
+    const data = (await resp.json()) as {
+      access_token?: string;
+      expires_in?: number;
+    };
+
+    if (!data.access_token) {
+      log.error({ credential }, "Token exchange response missing access_token");
+      return { ok: false, reason: "exchange_failed" };
+    }
+
+    accessToken = data.access_token;
+    expiresIn = data.expires_in ?? 3600;
+  } catch (err) {
+    log.error({ err, credential }, "Service account token exchange threw");
+    return { ok: false, reason: "exchange_failed" };
+  }
+
+  const newExpiresAt = Math.floor(Date.now() / 1000 + expiresIn);
+  const ok = await _setKey(
+    cacheKey,
+    JSON.stringify({ access_token: accessToken, expires_at: newExpiresAt }),
+  );
+  if (!ok) {
+    log.warn(
+      { credential },
+      "Failed to cache service account token — token still valid for this request",
+    );
+  } else {
+    log.info({ credential }, "Service account token exchanged and cached");
+  }
+  return { ok: true, token: accessToken };
+}
+
+function buildJwt(
+  clientEmail: string,
+  privateKey: string,
+  tokenUri: string,
+): string {
+  const now = Math.floor(Date.now() / 1000);
+  const header = Buffer.from(
+    JSON.stringify({ alg: "RS256", typ: "JWT" }),
+  ).toString("base64url");
+  const payload = Buffer.from(
+    JSON.stringify({
+      iss: clientEmail,
+      scope: GCP_SCOPE,
+      aud: tokenUri,
+      exp: now + 3600,
+      iat: now,
+    }),
+  ).toString("base64url");
+
+  const signingInput = `${header}.${payload}`;
+  const sign = createSign("SHA256");
+  sign.update(signingInput);
+  sign.end();
+  const signature = sign.sign(privateKey, "base64url");
+
+  return `${signingInput}.${signature}`;
+}
+
+// ---------------------------------------------------------------------------
+// Test helpers
+// ---------------------------------------------------------------------------
+
+/** @internal Test-only: reset all in-flight exchange mutexes. */
+export function _resetServiceAccountMutex(): void {
+  exchangesInFlight.clear();
+}
+
+/** @internal Test-only: swap vault accessors so tests never need mock.module. */
+export function _injectVaultAccessors(
+  get: typeof _realGetKey,
+  set: typeof _realSetKey,
+): void {
+  _getKey = get;
+  _setKey = set;
+}
+
+/** @internal Test-only: restore real vault accessors. */
+export function _resetVaultAccessors(): void {
+  _getKey = _realGetKey;
+  _setKey = _realSetKey;
+}

package/src/providers/model-catalog.ts

@@ -76,6 +76,13 @@ export interface CatalogModel {
   supportsAudioInput?: boolean;
   supportsToolUse?: boolean;
   supportsEffort?: boolean;
+  /**
+   * Whether this provider/model serving surface accepts a forced OpenAI
+   * chat-completions tool choice while thinking is enabled. Omit unless the
+   * combination is known incompatible. Daemon-only: not projected into the
+   * client catalog (see scripts/sync-llm-catalog.ts).
+   */
+  supportsForcedToolChoiceWithThinking?: boolean;
   pricing?: CatalogModelPricing;
   /**
    * Upper bound for `reasoning_effort` accepted by this model's upstream API.
@@ -856,6 +863,9 @@ const RAW_PROVIDER_CATALOG: ProviderCatalogEntry[] = [
           cacheReadPer1mTokens: 0.03,
         },
       },
+      // Limited to grandfathered accounts: other API keys get HTTP 404 "no
+      // longer available to new users", so this model is user-selectable
+      // only and no intent column may resolve to it.
       {
         id: "gemini-2.5-flash-lite",
         displayName: "Gemini 2.5 Flash Lite",
@@ -1816,6 +1826,7 @@ const RAW_PROVIDER_CATALOG: ProviderCatalogEntry[] = [
         supportsCaching: true,
         supportsVision: true,
         supportsToolUse: true,
+        supportsForcedToolChoiceWithThinking: false,
         pricing: {
           inputPer1mTokens: 0.95,
           outputPer1mTokens: 4.0,
@@ -2628,6 +2639,29 @@ export function modelSupportedEfforts(
   );
 }
 
+/**
+ * Whether a provider/model serving surface accepts a forced OpenAI
+ * chat-completions tool choice while thinking is enabled. Unknown providers
+ * and models fail open so custom routes retain their existing request shape
+ * and can rely on the bounded provider-error retry if needed.
+ */
+export function supportsForcedToolChoiceWithThinking(
+  providerId: string,
+  modelId: string,
+): boolean {
+  const provider = PROVIDER_CATALOG.find((entry) => entry.id === providerId);
+  if (!provider) {
+    return true;
+  }
+  const stripDateSuffix = (id: string): string => id.replace(/-\d{8}$/, "");
+  const normalizedModelId = stripDateSuffix(modelId);
+  return !provider.models.some(
+    (model) =>
+      model.supportsForcedToolChoiceWithThinking === false &&
+      (model.id === modelId || stripDateSuffix(model.id) === normalizedModelId),
+  );
+}
+
 /**
  * Return the catalog provider that owns a model ID, if known. When multiple
  * providers list the same ID (e.g. OpenRouter and the Vercel AI Gateway share

package/src/providers/model-intents.ts

@@ -11,10 +11,15 @@ const PROVIDER_DEFAULT_MODELS: Record<string, string> = Object.fromEntries(
 );
 
 // `cost-optimized` is the cheapest model a provider serves, `latency-optimized`
-// the fastest to first token. On anthropic, ollama, fireworks and openai the
-// same model is both, so those columns repeat by construction rather than by
-// oversight (openai's balanced repeats it too: the default-profile templates
-// split those tiers by reasoning effort, not model).
+// the fastest to first token. On anthropic, ollama, fireworks, openai and
+// gemini the same model is both, so those columns repeat by construction
+// rather than by oversight (openai's balanced repeats it too: the
+// default-profile templates split those tiers by reasoning effort, not model).
+//
+// "Serves" means serves to every API key: Google's `gemini-2.5-flash-lite`
+// undercuts every 3.x Flash-Lite but is limited to grandfathered accounts
+// (other keys get HTTP 404 "no longer available to new users"), so the
+// gemini column skips it.
 const PROVIDER_MODEL_INTENTS: Record<string, Record<ModelIntent, string>> = {
   anthropic: {
     balanced: "claude-sonnet-4-6",
@@ -32,7 +37,7 @@ const PROVIDER_MODEL_INTENTS: Record<string, Record<ModelIntent, string>> = {
   },
   gemini: {
     balanced: "gemini-3-flash-preview",
-    "cost-optimized": "gemini-2.5-flash-lite",
+    "cost-optimized": "gemini-3.1-flash-lite",
     "latency-optimized": "gemini-3.1-flash-lite",
     "quality-optimized": "gemini-3.1-pro-preview",
     "vision-optimized": "gemini-3-flash-preview",

package/src/providers/openai/__tests__/chat-completions-output-token-limit.test.ts

@@ -0,0 +1,141 @@
+/**
+ * Request-shape coverage for the output-token-limit wire key.
+ *
+ * The shared OpenAI-compatible transport defaults to `max_completion_tokens`.
+ * OpenRouter's parameter router matches `max_tokens` on `require_parameters`
+ * routes, except GPT-5.1/5.2 Codex and Chat which advertise only
+ * `max_completion_tokens`.
+ */
+
+import { describe, expect, test } from "bun:test";
+
+import {
+  openRouterOutputTokenLimitField,
+  OpenRouterProvider,
+} from "../../openrouter/client.js";
+import { OpenAIChatCompletionsProvider } from "../chat-completions-provider.js";
+
+type MockChunk = {
+  choices: Array<{
+    delta: { content?: string | null };
+    finish_reason?: string | null;
+  }>;
+  usage?: { prompt_tokens: number; completion_tokens: number };
+};
+
+const OK_CHUNKS: MockChunk[] = [
+  {
+    choices: [{ delta: { content: "ok" }, finish_reason: "stop" }],
+    usage: { prompt_tokens: 1, completion_tokens: 1 },
+  },
+];
+
+function makeStream(chunks: MockChunk[]): AsyncIterable<MockChunk> {
+  return {
+    async *[Symbol.asyncIterator]() {
+      for (const c of chunks) {
+        yield c;
+      }
+    },
+  };
+}
+
+function stubCreate(provider: OpenAIChatCompletionsProvider): {
+  requests: Record<string, unknown>[];
+} {
+  const requests: Record<string, unknown>[] = [];
+  const inner = provider as unknown as {
+    client: {
+      chat: {
+        completions: {
+          create: (
+            params: Record<string, unknown>,
+          ) => Promise<AsyncIterable<MockChunk>>;
+        };
+      };
+    };
+  };
+  inner.client.chat.completions.create = async (params) => {
+    requests.push(params);
+    return makeStream(OK_CHUNKS);
+  };
+  return { requests };
+}
+
+describe("chat-completions output-token-limit wire key", () => {
+  test("default OpenAI path sends max_completion_tokens and omits max_tokens", async () => {
+    const provider = new OpenAIChatCompletionsProvider(
+      "test-key",
+      "test-model",
+    );
+    const { requests } = stubCreate(provider);
+
+    await provider.sendMessage(
+      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      { config: { max_tokens: 64000 } },
+    );
+
+    expect(requests).toHaveLength(1);
+    expect(requests[0].max_completion_tokens).toBe(64000);
+    expect(requests[0]).not.toHaveProperty("max_tokens");
+  });
+
+  test("OpenRouter chat-completions path sends max_tokens and omits max_completion_tokens", async () => {
+    const provider = new OpenRouterProvider("or-key", "@preset/example");
+    const { requests } = stubCreate(provider);
+
+    await provider.sendMessage(
+      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      { config: { max_tokens: 64000 } },
+    );
+
+    expect(requests).toHaveLength(1);
+    expect(requests[0].max_tokens).toBe(64000);
+    expect(requests[0]).not.toHaveProperty("max_completion_tokens");
+  });
+
+  test("OpenRouter GPT-5.2 Codex sends max_completion_tokens and omits max_tokens", async () => {
+    const provider = new OpenRouterProvider("or-key", "openai/gpt-5.2-codex");
+    const { requests } = stubCreate(provider);
+
+    await provider.sendMessage(
+      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      { config: { max_tokens: 64000 } },
+    );
+
+    expect(requests).toHaveLength(1);
+    expect(requests[0].max_completion_tokens).toBe(64000);
+    expect(requests[0]).not.toHaveProperty("max_tokens");
+  });
+
+  test("OpenRouter per-call model override uses the overridden model's wire key", async () => {
+    const provider = new OpenRouterProvider("or-key", "x-ai/grok-4.20");
+    const { requests } = stubCreate(provider);
+
+    await provider.sendMessage(
+      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      { config: { model: "openai/gpt-5.1-codex-mini", max_tokens: 64000 } },
+    );
+
+    expect(requests).toHaveLength(1);
+    expect(requests[0].max_completion_tokens).toBe(64000);
+    expect(requests[0]).not.toHaveProperty("max_tokens");
+  });
+});
+
+describe("openRouterOutputTokenLimitField", () => {
+  test.each([
+    ["@preset/example", "max_tokens"],
+    ["x-ai/grok-4.20", "max_tokens"],
+    ["openai/gpt-5.2", "max_tokens"],
+    ["openai/gpt-5.2-pro", "max_tokens"],
+    ["openai/gpt-5.3-codex", "max_tokens"],
+    ["openai/gpt-5.2-codex", "max_completion_tokens"],
+    ["openai/gpt-5.2-chat", "max_completion_tokens"],
+    ["openai/gpt-5.1-codex", "max_completion_tokens"],
+    ["openai/gpt-5.1-codex-max", "max_completion_tokens"],
+    ["openai/gpt-5.1-codex-mini", "max_completion_tokens"],
+  ] as const)("%s -> %s", (model, field) => {
+    expect(openRouterOutputTokenLimitField(model)).toBe(field);
+  });
+});