@vellumai/assistant 0.12.0 → 0.12.1-dev.202609141913.5e16607

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-provider-reasoning.test.ts

@@ -934,6 +934,55 @@ describe("thinking-mode tool_choice rejection fallback", () => {
     expect(text?.text).toBe("ok");
   });
 
+  test("retries once when Kimi rejects a specified tool_choice in thinking mode", async () => {
+    const { provider, requests } = stubProviderWithErrors(
+      [
+        rejection(
+          "tool_choice 'specified' is incompatible with thinking enabled",
+        ),
+      ],
+      OK_CHUNKS,
+    );
+
+    const response = await provider.sendMessage(
+      [{ role: "user", content: [{ type: "text", text: "hi" }] }],
+      {
+        tools: [
+          {
+            name: "select_pages",
+            description: "Pick relevant memory pages",
+            input_schema: { type: "object", properties: {} },
+          },
+        ],
+        config: {
+          tool_choice: { type: "tool", name: "select_pages" },
+          effort: "high",
+        },
+      },
+    );
+
+    expect(requests).toHaveLength(2);
+    const first = requests[0] as {
+      tool_choice?: { type: string; function: { name: string } };
+      reasoning_effort?: string;
+    };
+    const second = requests[1] as {
+      tool_choice?: unknown;
+      reasoning_effort?: string;
+    };
+    expect(first.tool_choice).toEqual({
+      type: "function",
+      function: { name: "select_pages" },
+    });
+    expect(first.reasoning_effort).toBe("high");
+    expect(second.tool_choice).toBeUndefined();
+    expect(second.reasoning_effort).toBe("high");
+    const text = response.content.find((b) => b.type === "text") as
+      | { type: "text"; text: string }
+      | undefined;
+    expect(text?.text).toBe("ok");
+  });
+
   test("retries once for an OpenRouter-wrapped thinking-mode tool_choice rejection", async () => {
     const wrapped = new OpenAI.APIError(
       400,
@@ -1141,33 +1190,80 @@ describe("unknown assistant reasoning field rejection fallback", () => {
     },
   ];
 
-  test("retries once without reasoning_content when a strict schema rejects it", async () => {
+  test.each([
+    [
+      "a strict schema",
+      "Additional properties are not allowed ('reasoning_content' was unexpected)",
+    ],
+    ["Groq", "property 'reasoning_content' is unsupported"],
+  ])(
+    "retries once without reasoning_content when %s rejects it",
+    async (_upstream, message) => {
+      const { provider, requests } = stubProviderWithErrors(
+        [rejection(message)],
+        OK_CHUNKS,
+        { assistantReasoningField: "reasoning_content" },
+      );
+
+      const response = await provider.sendMessage(thinkingHistory);
+
+      expect(requests).toHaveLength(2);
+      const first = requests[0] as {
+        messages: Array<{ reasoning_content?: string; content: string | null }>;
+      };
+      const second = requests[1] as {
+        messages: Array<{ reasoning_content?: string; content: string | null }>;
+      };
+      expect(first.messages[0].reasoning_content).toBe("hidden chain");
+      expect(second.messages[0].reasoning_content).toBeUndefined();
+      expect(second.messages[0].content).toBe("answer");
+      const text = response.content.find((b) => b.type === "text") as
+        | { type: "text"; text: string }
+        | undefined;
+      expect(text?.text).toBe("ok");
+    },
+  );
+
+  test("preserves assistant tool calls when stripping an unsupported reasoning_content", async () => {
     const { provider, requests } = stubProviderWithErrors(
-      [
-        rejection(
-          "Additional properties are not allowed ('reasoning_content' was unexpected)",
-        ),
-      ],
+      [rejection("property 'reasoning_content' is unsupported")],
       OK_CHUNKS,
       { assistantReasoningField: "reasoning_content" },
     );
 
-    const response = await provider.sendMessage(thinkingHistory);
+    await provider.sendMessage([
+      {
+        role: "assistant",
+        content: [
+          { type: "thinking", thinking: "hidden chain", signature: "" },
+          {
+            type: "tool_use",
+            id: "call_1",
+            name: "lookup",
+            input: { q: "x" },
+          },
+        ],
+      },
+      {
+        role: "user",
+        content: [
+          { type: "tool_result", tool_use_id: "call_1", content: "found" },
+        ],
+      },
+    ]);
 
     expect(requests).toHaveLength(2);
-    const first = requests[0] as {
-      messages: Array<{ reasoning_content?: string; content: string | null }>;
-    };
     const second = requests[1] as {
-      messages: Array<{ reasoning_content?: string; content: string | null }>;
+      messages: Array<{
+        role: string;
+        reasoning_content?: string;
+        tool_calls?: Array<{ id: string; function: { name: string } }>;
+      }>;
     };
-    expect(first.messages[0].reasoning_content).toBe("hidden chain");
-    expect(second.messages[0].reasoning_content).toBeUndefined();
-    expect(second.messages[0].content).toBe("answer");
-    const text = response.content.find((b) => b.type === "text") as
-      | { type: "text"; text: string }
-      | undefined;
-    expect(text?.text).toBe("ok");
+    const assistant = second.messages.find((m) => m.role === "assistant");
+    expect(assistant?.reasoning_content).toBeUndefined();
+    expect(assistant?.tool_calls?.[0]?.id).toBe("call_1");
+    expect(assistant?.tool_calls?.[0]?.function.name).toBe("lookup");
   });
 
   test("does not strip reasoning_content on a must-be-passed-back error", async () => {
@@ -1189,6 +1285,38 @@ describe("unknown assistant reasoning field rejection fallback", () => {
     ).toBe("hidden chain");
   });
 
+  test("does not strip reasoning_content on a 429 output-token limit", async () => {
+    const { provider, requests } = stubProviderWithErrors(
+      [
+        rejection(
+          "Rate limit reached for model `qwen/qwen3-32b` in organization `org_x` on tokens per minute (TPM): Limit 1000, Requested 4000.",
+          429,
+        ),
+      ],
+      OK_CHUNKS,
+      { assistantReasoningField: "reasoning_content" },
+    );
+
+    await expect(provider.sendMessage(thinkingHistory)).rejects.toThrow();
+    expect(requests).toHaveLength(1);
+  });
+
+  test("does not strip reasoning_content on a 413 input-token limit", async () => {
+    const { provider, requests } = stubProviderWithErrors(
+      [
+        rejection(
+          "Request too large for model `qwen/qwen3-32b` in organization `org_x` on tokens per minute (TPM): Limit 7000, Requested 24936.",
+          413,
+        ),
+      ],
+      OK_CHUNKS,
+      { assistantReasoningField: "reasoning_content" },
+    );
+
+    await expect(provider.sendMessage(thinkingHistory)).rejects.toThrow();
+    expect(requests).toHaveLength(1);
+  });
+
   test("does not retry unknown-field 500s", async () => {
     const { provider, requests } = stubProviderWithErrors(
       [