@vellumai/assistant 0.7.0 → 0.7.2

package/src/providers/__tests__/retry-callsite.test.ts

@@ -122,6 +122,117 @@ describe("RetryProvider — callSite resolution", () => {
     expect(config.modelIntent).toBeUndefined();
   });
 
+  test("attaches sanitized stable attribution headers only when enabled", async () => {
+    setLlmConfig({
+      default: {
+        provider: "anthropic",
+        model: "claude-default",
+      },
+      profiles: {
+        "conversation-profile": {
+          model: "claude-profile",
+          source: "user",
+        },
+      },
+      callSites: {
+        memoryRetrieval: {
+          provider: "openai",
+        },
+      },
+    });
+
+    let seen: SendMessageOptions | undefined;
+    const wrapped = new RetryProvider(
+      makeProvider("openai", (options) => {
+        seen = options;
+      }),
+      { forwardUsageAttributionHeaders: true },
+    );
+
+    await wrapped.sendMessage(DUMMY_MESSAGES, undefined, undefined, {
+      config: {
+        callSite: "memoryRetrieval",
+        overrideProfile: "conversation-profile",
+      },
+    });
+
+    const config = seen?.config as Record<string, unknown>;
+    expect(config.usageAttributionHeaders).toEqual({
+      "X-Vellum-LLM-Call-Site": "memoryRetrieval",
+      "X-Vellum-Inference-Profile": "conversation-profile",
+      "X-Vellum-Inference-Profile-Source": "conversation",
+      "X-Vellum-Resolved-Provider": "openai",
+      "X-Vellum-Resolved-Model": "claude-profile",
+    });
+    expect(
+      (config.usageAttributionHeaders as Record<string, string>)[
+        "X-Vellum-LLM-Call-Site-Label"
+      ],
+    ).toBeUndefined();
+  });
+
+  test("omits attribution headers by default for direct provider transports", async () => {
+    setLlmConfig({
+      default: {
+        provider: "openai",
+        model: "gpt-default",
+      },
+      callSites: {
+        memoryRetrieval: {
+          provider: "openai",
+        },
+      },
+    });
+
+    let seen: SendMessageOptions | undefined;
+    const wrapped = new RetryProvider(
+      makeProvider("openai", (options) => {
+        seen = options;
+      }),
+    );
+
+    await wrapped.sendMessage(DUMMY_MESSAGES, undefined, undefined, {
+      config: { callSite: "memoryRetrieval" },
+    });
+
+    const config = seen?.config as Record<string, unknown>;
+    expect(config.usageAttributionHeaders).toBeUndefined();
+    expect(config.callSite).toBeUndefined();
+  });
+
+  test("omits profile source attribution header when no profile is applied", async () => {
+    setLlmConfig({
+      default: {
+        provider: "openai",
+        model: "gpt-default",
+      },
+      callSites: {
+        memoryRetrieval: {
+          provider: "openai",
+        },
+      },
+    });
+
+    let seen: SendMessageOptions | undefined;
+    const wrapped = new RetryProvider(
+      makeProvider("openai", (options) => {
+        seen = options;
+      }),
+      { forwardUsageAttributionHeaders: true },
+    );
+
+    await wrapped.sendMessage(DUMMY_MESSAGES, undefined, undefined, {
+      config: { callSite: "memoryRetrieval" },
+    });
+
+    const config = seen?.config as Record<string, unknown>;
+    expect(config.usageAttributionHeaders).toEqual({
+      "X-Vellum-LLM-Call-Site": "memoryRetrieval",
+      "X-Vellum-Resolved-Provider": "openai",
+      "X-Vellum-Resolved-Model": "gpt-default",
+    });
+  });
+
   test("falls back to llm.default when llm.callSites[id] is absent", async () => {
     setLlmConfig({
       default: {
@@ -185,7 +296,7 @@ describe("RetryProvider — callSite resolution", () => {
     // fields with `{type:"invalid_request_error", message:"contextWindow:
     // Extra inputs are not permitted"}`. Provider routing is handled by
     // CallSiteRoutingProvider; contextWindow is consumed by the agent loop
-    // directly from `config.llm.default.contextWindow.*`.
+    // from the effective per-call-site/profile context resolver.
     expect(config.contextWindow).toBeUndefined();
     expect(config.provider).toBeUndefined();
   });
@@ -386,6 +497,32 @@ describe("RetryProvider — no callSite (pre-resolved config passes through)", (
     expect(config.model).not.toBe("MUST-NOT-LEAK");
     expect(config.model).not.toBe("ALSO-MUST-NOT-LEAK");
   });
+
+  test("does not forward caller-supplied attribution headers without callSite", async () => {
+    setLlmConfig({
+      default: { provider: "anthropic", model: "MUST-NOT-LEAK" },
+    });
+
+    let seen: SendMessageOptions | undefined;
+    const wrapped = new RetryProvider(
+      makeProvider("anthropic", (options) => {
+        seen = options;
+      }),
+    );
+
+    await wrapped.sendMessage(DUMMY_MESSAGES, undefined, undefined, {
+      config: {
+        model: "explicit-model",
+        usageAttributionHeaders: {
+          "X-Vellum-LLM-Call-Site": "injected",
+        },
+      },
+    });
+
+    const config = seen?.config as Record<string, unknown>;
+    expect(config.model).toBe("explicit-model");
+    expect(config.usageAttributionHeaders).toBeUndefined();
+  });
 });
 
 // ── getConfiguredProvider — call-site routing ──────────────────────────────

package/src/providers/anthropic/client.ts

@@ -866,6 +866,7 @@ export class AnthropicProvider implements Provider {
         output_config,
         cacheTtl: _cacheTtl,
         max_tokens: callerMaxTokens,
+        usageAttributionHeaders,
         ...restConfig
       } = (config ?? {}) as Record<string, unknown> & {
         // "xhigh" is an intermediate tier between "high" and "max" supported
@@ -877,9 +878,10 @@ export class AnthropicProvider implements Provider {
         effort?: "none" | "low" | "medium" | "high" | "xhigh" | "max";
         speed?: "standard" | "fast";
         output_config?: Record<string, unknown>;
+        usageAttributionHeaders?: Record<string, string>;
       };
-      // Haiku does not support the effort / output_config parameter,
-      // extended cache TTL betas, 1M context, or 64K output tokens.
+      // Haiku does not support the effort / output_config parameter or
+      // extended cache TTL betas.
       // Determine the effective model (per-call override or provider default)
       // and gate features accordingly.
       const effectiveModel =
@@ -991,24 +993,48 @@ export class AnthropicProvider implements Provider {
       // TTL is appropriate. Walk backwards to find the last user message
       // with a real text block (skipping tool_result-only messages and
       // synthetic continuation placeholders injected by ensureToolPairing).
-      let turnStartIdx = -1;
-      for (let i = sentMessages.length - 1; i >= 0; i--) {
-        const msg = sentMessages[i];
-        if (msg.role !== "user" || !Array.isArray(msg.content)) continue;
-        const hasText = msg.content.some(
-          (b) =>
-            typeof b !== "string" &&
-            b.type === "text" &&
-            b.text !== SYNTHETIC_CONTINUATION_TEXT,
-        );
-        if (!hasText) continue;
-        const lastBlock = msg.content[msg.content.length - 1];
+      const msgs = sentMessages;
+      const findUserTextMsgIdx = (startIdx: number): number => {
+        for (let i = startIdx; i >= 0; i--) {
+          const msg = msgs[i];
+          if (msg.role !== "user" || !Array.isArray(msg.content)) continue;
+          const hasText = msg.content.some(
+            (b) =>
+              typeof b !== "string" &&
+              b.type === "text" &&
+              b.text !== SYNTHETIC_CONTINUATION_TEXT,
+          );
+          if (hasText) return i;
+        }
+        return -1;
+      };
+      const applyCacheControlToLastBlock = (msgIdx: number): void => {
+        const content = msgs[msgIdx].content;
+        if (!Array.isArray(content) || content.length === 0) return;
+        const lastBlock = content[content.length - 1];
         if (typeof lastBlock !== "string") {
           (lastBlock as unknown as Record<string, unknown>).cache_control =
             cacheControl;
         }
-        turnStartIdx = i;
-        break;
+      };
+      const turnStartIdx = findUserTextMsgIdx(msgs.length - 1);
+      if (turnStartIdx >= 0) applyCacheControlToLastBlock(turnStartIdx);
+
+      // Previous-turn anchor: when this request is the first of a new turn
+      // (turn-start is the very last message — no tool-use loop yet), also
+      // place a 1h breakpoint on the *previous* turn-starting user message.
+      // Anthropic only matches the cache at cache_control points present in
+      // the current request, so without this anchor the breakpoint slides
+      // forward each new user turn and the prior cached prefix becomes
+      // unreachable — forcing a full re-creation of history (200K+
+      // cache_creation tokens per new turn). Skipped during tool-use loops
+      // where the current turn-start already covers the same prefix and a
+      // second anchor would blow the 4-breakpoint budget.
+      let prevTurnAnchorIdx = -1;
+      if (turnStartIdx === msgs.length - 1 && turnStartIdx > 0) {
+        prevTurnAnchorIdx = findUserTextMsgIdx(turnStartIdx - 1);
+        if (prevTurnAnchorIdx >= 0)
+          applyCacheControlToLastBlock(prevTurnAnchorIdx);
       }
 
       // Advancing tail: place a short-lived 5m cache breakpoint on the last
@@ -1045,19 +1071,21 @@ export class AnthropicProvider implements Provider {
       }
 
       // Enforce Anthropic API maximum of 4 cache_control blocks.
-      // When the system prompt boundary splits into 2 cached blocks AND
-      // tools + turn-start + advancing-tail breakpoints are all present,
-      // we'd have 5.  Drop the static system block's breakpoint — it's
-      // small (<1K tokens) so the re-read cost is negligible, while the
-      // dynamic block (workspace context) rarely changes mid-session and
-      // benefits more from caching.
-      const hasTailBreakpoint = tailBreakpointApplied;
+      // With the system prompt boundary split into 2 cached blocks AND
+      // tools + turn-start + (tail OR prev-turn-anchor), we'd have 5.
+      // Drop the static system block's breakpoint — it's small (<1K
+      // tokens) so the re-read cost is negligible, while the dynamic
+      // block (workspace context) rarely changes mid-session and
+      // benefits more from caching. Tail and prev-turn-anchor are
+      // mutually exclusive (prev-turn-anchor only fires when turn-start
+      // is the last message, which is the exact condition that suppresses
+      // the tail), so we never exceed 5.
       const hasToolCacheBreakpoint =
         params.tools?.some(
           (t) => "cache_control" in t && t.cache_control != null,
         ) ?? false;
       if (
-        hasTailBreakpoint &&
+        (tailBreakpointApplied || prevTurnAnchorIdx >= 0) &&
         Array.isArray(params.system) &&
         params.system.length === 2 &&
         hasToolCacheBreakpoint
@@ -1091,15 +1119,13 @@ export class AnthropicProvider implements Provider {
         finalMessage(): Promise<Anthropic.Message>;
       }
 
-      // Fast mode: use the beta endpoint with speed: "fast" for Opus models (4.6, 4.7)
+      // Fast mode: use the beta endpoint with speed: "fast" for Opus models (4.6, 4.7).
       const useFastMode = speed === "fast" && effectiveModel.includes("opus");
 
-      // Collect required betas: extended cache TTL for 1h system prompt caching,
-      // 1M context window, and fast-mode when applicable.
-      // Haiku doesn't support the extended cache TTL or 1M context betas.
-      const betas: string[] = isHaiku
-        ? []
-        : ["extended-cache-ttl-2025-04-11", "context-1m-2025-08-07"];
+      // Collect request betas that are still explicit transport features.
+      // Current long-context Anthropic models expose their larger windows by
+      // model capability in the catalog/resolver, not by this generic header.
+      const betas: string[] = isHaiku ? [] : ["extended-cache-ttl-2025-04-11"];
       if (useFastMode) {
         betas.push("fast-mode-2026-02-01");
       }
@@ -1114,7 +1140,12 @@ export class AnthropicProvider implements Provider {
                 betas,
               } as Anthropic.Beta.Messages.MessageCreateParamsNonStreaming &
                 Anthropic.Beta.Messages.MessageCreateParamsStreaming,
-              { signal: timeoutSignal },
+              {
+                signal: timeoutSignal,
+                ...(usageAttributionHeaders
+                  ? { headers: usageAttributionHeaders }
+                  : {}),
+              },
             ) as unknown as UnifiedStream)
           : betas.length > 0
             ? (this.client.beta.messages.stream(
@@ -1123,10 +1154,18 @@ export class AnthropicProvider implements Provider {
                   betas,
                 } as Anthropic.Beta.Messages.MessageCreateParamsNonStreaming &
                   Anthropic.Beta.Messages.MessageCreateParamsStreaming,
-                { signal: timeoutSignal },
+                {
+                  signal: timeoutSignal,
+                  ...(usageAttributionHeaders
+                    ? { headers: usageAttributionHeaders }
+                    : {}),
+                },
               ) as unknown as UnifiedStream)
             : (this.client.messages.stream(params, {
                 signal: timeoutSignal,
+                ...(usageAttributionHeaders
+                  ? { headers: usageAttributionHeaders }
+                  : {}),
               }) as unknown as UnifiedStream);
 
         // Buffer streaming text until it's clear the accumulated text isn't

package/src/providers/call-site-routing.ts

@@ -19,6 +19,8 @@
  * stable identity.
  */
 
+import { AsyncLocalStorage } from "node:async_hooks";
+
 import { resolveCallSiteConfig } from "../config/llm-resolver.js";
 import { getConfig } from "../config/loader.js";
 import type {
@@ -30,14 +32,28 @@ import type {
 } from "./types.js";
 
 export class CallSiteRoutingProvider implements Provider {
-  public readonly name: string;
   public readonly tokenEstimationProvider?: string;
 
+  // Per-call async context that tracks which provider is currently executing.
+  // Using AsyncLocalStorage instead of a plain instance field means concurrent
+  // sendMessage calls (e.g. the main agent turn and a title-generation call
+  // both in-flight at the same time on the same provider instance) each see
+  // their own value — no clobbering, no premature clear.
+  //
+  // During sendMessage, emitLlmCallStartedIfNeeded reads provider.name on the
+  // first text_delta (before the response completes). The getter below returns
+  // the async-context value so streaming trace events carry the routed
+  // provider's name, not the default's.
+  private readonly _activeProviderContext = new AsyncLocalStorage<string>();
+
+  get name(): string {
+    return this._activeProviderContext.getStore() ?? this.defaultProvider.name;
+  }
+
   constructor(
     private readonly defaultProvider: Provider,
     private readonly getProviderByName: (name: string) => Provider | undefined,
   ) {
-    this.name = defaultProvider.name;
     this.tokenEstimationProvider = defaultProvider.tokenEstimationProvider;
   }
 
@@ -48,7 +64,30 @@ export class CallSiteRoutingProvider implements Provider {
     options?: SendMessageOptions,
   ): Promise<ProviderResponse> {
     const target = this.selectProvider(options);
-    return target.sendMessage(messages, tools, systemPrompt, options);
+    const isRouted = target !== this.defaultProvider;
+
+    const doSend = async (): Promise<ProviderResponse> => {
+      const response = await target.sendMessage(
+        messages,
+        tools,
+        systemPrompt,
+        options,
+      );
+      // Also stamp actualProvider on the response so that handleUsage /
+      // llm_call_finished (which read event.actualProvider, not provider.name)
+      // attribute the call to the right provider.
+      if (isRouted && response.actualProvider == null) {
+        return { ...response, actualProvider: target.name };
+      }
+      return response;
+    };
+
+    // Run inside the async context so that any code reading provider.name
+    // during streaming (e.g. emitLlmCallStartedIfNeeded on text_delta) sees
+    // the routed provider's name for this specific call, not the default.
+    return isRouted
+      ? this._activeProviderContext.run(target.name, doSend)
+      : doSend();
   }
 
   /**

package/src/providers/gemini/client.ts

@@ -30,10 +30,20 @@ const GEMINI_CONTEXT_OVERFLOW_TOKEN_PATTERNS =
 const GEMINI_3_UNSIGNED_TOOL_CALL_THOUGHT_SIGNATURE =
   "context_engineering_is_the_way_to_go";
 
-export function isGemini3Model(model: string): boolean {
+function isGemini3Model(model: string): boolean {
   return model.startsWith("gemini-3") || model.startsWith("models/gemini-3");
 }
 
+function stripGeminiHttpOptions(
+  config: genai.GenerateContentConfig,
+): genai.GenerateContentConfig {
+  const { httpOptions: _httpOptions, ...rest } =
+    config as genai.GenerateContentConfig & {
+      httpOptions?: unknown;
+    };
+  return rest;
+}
+
 /**
  * Detect Gemini's context-overflow signals on an `ApiError`. Gemini surfaces
  * this condition via its "RESOURCE_EXHAUSTED" category. The Gemini SDK's
@@ -162,6 +172,9 @@ export class GeminiProvider implements Provider {
     const configObj = config as Record<string, unknown> | undefined;
     const maxTokens = configObj?.max_tokens as number | undefined;
     const modelOverride = configObj?.model as string | undefined;
+    const usageAttributionHeaders = configObj?.usageAttributionHeaders as
+      | Record<string, string>
+      | undefined;
     const activeModel = modelOverride ?? this.model;
 
     try {
@@ -193,6 +206,9 @@ export class GeminiProvider implements Provider {
       const { signal: timeoutSignal, cleanup: cleanupTimeout } =
         createStreamTimeout(this.streamTimeoutMs, signal);
       geminiConfig.abortSignal = timeoutSignal;
+      if (usageAttributionHeaders) {
+        geminiConfig.httpOptions = { headers: usageAttributionHeaders };
+      }
 
       // Accumulate from streaming chunks
       let fullText = "";
@@ -296,7 +312,7 @@ export class GeminiProvider implements Provider {
       const rawRequest = {
         model: activeModel,
         contents: geminiContents,
-        config: geminiConfig,
+        config: stripGeminiHttpOptions(geminiConfig),
       };
       const rawResponse = {
         model: responseModel,

package/src/providers/managed-proxy/context.ts

@@ -92,8 +92,3 @@ export async function managedFallbackEnabledFor(
   if (!meta?.managed) return false;
   return await hasManagedProxyPrereqs();
 }
-
-/** @internal Test-only: reset the cached managed-proxy-enabled flag. */
-export function _resetManagedProxyEnabledCache(): void {
-  _managedProxyEnabled = false;
-}