@vellumai/assistant 0.9.0 → 0.10.0-dev.202606192226.8ed8a6a

package/src/providers/anthropic/client.ts

@@ -18,6 +18,7 @@ import type {
   Provider,
   ProviderResponse,
   SendMessageOptions,
+  ToolDefinition,
 } from "../types.js";
 import {
   ContextOverflowError,
@@ -823,181 +824,7 @@ export class AnthropicProvider implements Provider {
     // edge LB, NAT idle) — only the latter should be retried.
     let innerTimeoutSignal: AbortSignal | undefined;
     try {
-      const formatted = messages
-        .map((m) => {
-          // Track whether an unknown block was dropped during filtering
-          let droppedUnknownBlock = false;
-
-          const content = m.content
-            .map((block) => {
-              const result = this.toAnthropicBlockSafe(block);
-              if (result == null) {
-                droppedUnknownBlock = true;
-              }
-              return result;
-            })
-            .filter(
-              (block): block is Anthropic.ContentBlockParam => block != null,
-            )
-            .filter(
-              (block) =>
-                !(
-                  block.type === "text" &&
-                  !(block as { text?: string }).text?.trim()
-                ),
-            );
-
-          // Preserve assistant turns that would otherwise become empty after filtering
-          // unknown block types (e.g. ui_surface). Dropping these messages can violate
-          // Anthropic's role alternation requirement.
-          if (
-            content.length === 0 &&
-            m.role === "assistant" &&
-            droppedUnknownBlock
-          ) {
-            return {
-              role: m.role as "assistant",
-              content: [
-                { type: "text" as const, text: PLACEHOLDER_BLOCKS_OMITTED },
-              ],
-            };
-          }
-
-          return {
-            role: m.role,
-            content,
-          } as Anthropic.MessageParam;
-        })
-        .reduce<Anthropic.MessageParam[]>((acc, m) => {
-          if (m.content.length > 0) {
-            acc.push(m);
-            return acc;
-          }
-          // Dropping an empty assistant message between two user messages (or vice
-          // versa) would create consecutive same-role messages, violating
-          // Anthropic's role alternation requirement. Inject a placeholder instead.
-          const prev = acc[acc.length - 1];
-          if (m.role === "assistant" && prev && prev.role !== "assistant") {
-            acc.push({
-              role: "assistant" as const,
-              content: [
-                { type: "text" as const, text: PLACEHOLDER_EMPTY_TURN },
-              ],
-            });
-          }
-          return acc;
-        }, []);
-
-      // Post-processing: merge consecutive same-role messages that violate
-      // Anthropic's strict user/assistant alternation requirement. These can
-      // arise from:
-      //   - Dropping empty messages in the reduce above (placeholder-adjacent)
-      //   - History reconstruction artifacts that bypass repairHistory
-      //
-      // Walk backwards so splice indices stay valid. After a merge+splice
-      // the element that was at i+1 shifts to i, potentially creating a
-      // new adjacent pair — bump i back up to recheck that position.
-      {
-        let i = formatted.length - 1;
-        while (i > 0 && i < formatted.length) {
-          if (formatted[i].role !== formatted[i - 1].role) {
-            i--;
-            continue;
-          }
-
-          const iContent = (
-            Array.isArray(formatted[i].content) ? formatted[i].content : []
-          ) as Anthropic.ContentBlockParam[];
-          const prevContent = (
-            Array.isArray(formatted[i - 1].content)
-              ? formatted[i - 1].content
-              : []
-          ) as Anthropic.ContentBlockParam[];
-          const isPlaceholder = (c: Anthropic.ContentBlockParam[]): boolean => {
-            if (
-              c.length !== 1 ||
-              typeof c[0] === "string" ||
-              c[0].type !== "text"
-            )
-              return false;
-            const text = (c[0] as { text?: string }).text;
-            return typeof text === "string" && isPlaceholderSentinelText(text);
-          };
-
-          if (isPlaceholder(iContent)) {
-            formatted.splice(i, 1);
-            // Removed the later element. The new formatted[i] (formerly
-            // i+1) may now be same-role as i-1, so decrement once to
-            // recheck from the correct position.
-            i--;
-          } else if (isPlaceholder(prevContent)) {
-            formatted.splice(i - 1, 1);
-            // Removed the earlier element — everything shifted down by 1.
-            // The element that was at i is now at i-1. Decrement so the
-            // next iteration compares the new i-1 with i-2 (or exits if
-            // i-1 is 0).
-            i--;
-          } else {
-            // Neither is a placeholder — merge content blocks into the
-            // earlier message and remove the later one. Skip the merge
-            // when either message carries tool_use or tool_result blocks;
-            // those require structural alternation for ensureToolPairing
-            // to inject the correct synthetic results downstream.
-            const hasToolBlock = (c: Anthropic.ContentBlockParam[]): boolean =>
-              c.some(
-                (b) =>
-                  typeof b !== "string" &&
-                  (b.type === "tool_use" || b.type === "tool_result"),
-              );
-            if (!hasToolBlock(prevContent) && !hasToolBlock(iContent)) {
-              formatted[i - 1] = {
-                ...formatted[i - 1],
-                content: [...prevContent, ...iContent],
-              };
-              formatted.splice(i, 1);
-              // Clamp i to the new last index — the splice may have put
-              // us past the end. If there's a new element at i (formerly
-              // i+1), it will be rechecked against the merged i-1.
-              if (i >= formatted.length) {
-                i = formatted.length - 1;
-              }
-            } else {
-              // Can't merge (tool blocks present) — leave for
-              // ensureToolPairing which handles tool_use/tool_result
-              // alternation in its own forward walk.
-              i--;
-            }
-          }
-        }
-      }
-
-      // Strip thinking/redacted_thinking blocks from completed historical
-      // assistant turns. Anthropic only requires these blocks for active
-      // tool-use continuation (the tail span where assistant tool_use is
-      // followed by user tool_result). Replaying stale thinking blocks from
-      // earlier turns causes 400 errors when the signature is no longer
-      // valid (e.g. after a provider/model/profile switch).
-      const activeToolUseStart = findActiveToolUseContinuationStart(formatted);
-      for (let i = 0; i < activeToolUseStart; i++) {
-        const msg = formatted[i];
-        if (msg.role !== "assistant" || !Array.isArray(msg.content)) continue;
-        const stripped = (msg.content as Anthropic.ContentBlockParam[]).filter(
-          (b) =>
-            typeof b === "string" ||
-            (b.type !== "thinking" && b.type !== "redacted_thinking"),
-        );
-        if (stripped.length === 0) {
-          stripped.push({
-            type: "text" as const,
-            text: PLACEHOLDER_BLOCKS_OMITTED,
-          });
-        }
-        formatted[i] = { ...msg, content: stripped };
-      }
-
-      sentMessages = ensureToolPairing(
-        repairOrphanedServerToolBlocks(formatted),
-      );
+      sentMessages = this.buildSentMessages(messages);
       const {
         effort,
         speed,
@@ -1681,6 +1508,234 @@ export class AnthropicProvider implements Provider {
     }
   }
 
+  /**
+   * Exact prompt-token count via Anthropic's `/v1/messages/count_tokens`
+   * endpoint — the real tokenizer, no inference. Serializes `messages` /
+   * `systemPrompt` / `tools` the same way {@link sendMessage} does (so the
+   * count tracks what the next call would actually send), minus the
+   * `cache_control` breakpoints, which don't affect token counts.
+   *
+   * The serialization here is intentionally simpler than `sendMessage`'s
+   * (no role-alternation merge / placeholder injection): on a pathological
+   * history `count_tokens` may reject the request, which surfaces as a thrown
+   * error the caller turns into a local-estimator fallback. The common path —
+   * a well-formed history — counts exactly.
+   */
+  async countInputTokens(
+    messages: Message[],
+    systemPrompt: string,
+    tools?: ToolDefinition[],
+  ): Promise<number> {
+    const sentMessages = this.buildSentMessages(messages);
+
+    const system = systemPrompt
+      ? systemPrompt
+          .split(SYSTEM_PROMPT_CACHE_BOUNDARY)
+          .filter((text) => text.length > 0)
+          .map((text) => ({ type: "text" as const, text }))
+      : [];
+    const toolsParam =
+      tools && tools.length > 0
+        ? tools.map((t) => ({
+            name: t.name,
+            description: t.description,
+            input_schema: t.input_schema as Anthropic.Tool["input_schema"],
+          }))
+        : undefined;
+
+    const res = await this.client.messages.countTokens({
+      model: this.model,
+      messages: sentMessages,
+      ...(system.length > 0 ? { system } : {}),
+      ...(toolsParam ? { tools: toolsParam } : {}),
+    });
+    return res.input_tokens;
+  }
+
+  /**
+   * Serialize internal `Message[]` into the Anthropic `MessageParam[]` the
+   * Messages API expects: drop unknown/empty blocks, preserve role
+   * alternation (placeholder injection + same-role merge), strip stale
+   * thinking blocks outside the active tool-use continuation, and repair
+   * tool_use/tool_result pairing. Shared by {@link sendMessage} and
+   * {@link countInputTokens} so both send the identical message payload —
+   * `cache_control` breakpoints (which don't affect token counts) are the
+   * only thing layered on top in `sendMessage`.
+   */
+  private buildSentMessages(messages: Message[]): Anthropic.MessageParam[] {
+    const formatted = messages
+      .map((m) => {
+        // Track whether an unknown block was dropped during filtering
+        let droppedUnknownBlock = false;
+
+        const content = m.content
+          .map((block) => {
+            const result = this.toAnthropicBlockSafe(block);
+            if (result == null) {
+              droppedUnknownBlock = true;
+            }
+            return result;
+          })
+          .filter(
+            (block): block is Anthropic.ContentBlockParam => block != null,
+          )
+          .filter(
+            (block) =>
+              !(
+                block.type === "text" &&
+                !(block as { text?: string }).text?.trim()
+              ),
+          );
+
+        // Preserve assistant turns that would otherwise become empty after filtering
+        // unknown block types (e.g. ui_surface). Dropping these messages can violate
+        // Anthropic's role alternation requirement.
+        if (
+          content.length === 0 &&
+          m.role === "assistant" &&
+          droppedUnknownBlock
+        ) {
+          return {
+            role: m.role as "assistant",
+            content: [
+              { type: "text" as const, text: PLACEHOLDER_BLOCKS_OMITTED },
+            ],
+          };
+        }
+
+        return {
+          role: m.role,
+          content,
+        } as Anthropic.MessageParam;
+      })
+      .reduce<Anthropic.MessageParam[]>((acc, m) => {
+        if (m.content.length > 0) {
+          acc.push(m);
+          return acc;
+        }
+        // Dropping an empty assistant message between two user messages (or vice
+        // versa) would create consecutive same-role messages, violating
+        // Anthropic's role alternation requirement. Inject a placeholder instead.
+        const prev = acc[acc.length - 1];
+        if (m.role === "assistant" && prev && prev.role !== "assistant") {
+          acc.push({
+            role: "assistant" as const,
+            content: [{ type: "text" as const, text: PLACEHOLDER_EMPTY_TURN }],
+          });
+        }
+        return acc;
+      }, []);
+
+    // Post-processing: merge consecutive same-role messages that violate
+    // Anthropic's strict user/assistant alternation requirement. These can
+    // arise from:
+    //   - Dropping empty messages in the reduce above (placeholder-adjacent)
+    //   - History reconstruction artifacts that bypass repairHistory
+    //
+    // Walk backwards so splice indices stay valid. After a merge+splice
+    // the element that was at i+1 shifts to i, potentially creating a
+    // new adjacent pair — bump i back up to recheck that position.
+    {
+      let i = formatted.length - 1;
+      while (i > 0 && i < formatted.length) {
+        if (formatted[i].role !== formatted[i - 1].role) {
+          i--;
+          continue;
+        }
+
+        const iContent = (
+          Array.isArray(formatted[i].content) ? formatted[i].content : []
+        ) as Anthropic.ContentBlockParam[];
+        const prevContent = (
+          Array.isArray(formatted[i - 1].content)
+            ? formatted[i - 1].content
+            : []
+        ) as Anthropic.ContentBlockParam[];
+        const isPlaceholder = (c: Anthropic.ContentBlockParam[]): boolean => {
+          if (
+            c.length !== 1 ||
+            typeof c[0] === "string" ||
+            c[0].type !== "text"
+          )
+            return false;
+          const text = (c[0] as { text?: string }).text;
+          return typeof text === "string" && isPlaceholderSentinelText(text);
+        };
+
+        if (isPlaceholder(iContent)) {
+          formatted.splice(i, 1);
+          // Removed the later element. The new formatted[i] (formerly
+          // i+1) may now be same-role as i-1, so decrement once to
+          // recheck from the correct position.
+          i--;
+        } else if (isPlaceholder(prevContent)) {
+          formatted.splice(i - 1, 1);
+          // Removed the earlier element — everything shifted down by 1.
+          // The element that was at i is now at i-1. Decrement so the
+          // next iteration compares the new i-1 with i-2 (or exits if
+          // i-1 is 0).
+          i--;
+        } else {
+          // Neither is a placeholder — merge content blocks into the
+          // earlier message and remove the later one. Skip the merge
+          // when either message carries tool_use or tool_result blocks;
+          // those require structural alternation for ensureToolPairing
+          // to inject the correct synthetic results downstream.
+          const hasToolBlock = (c: Anthropic.ContentBlockParam[]): boolean =>
+            c.some(
+              (b) =>
+                typeof b !== "string" &&
+                (b.type === "tool_use" || b.type === "tool_result"),
+            );
+          if (!hasToolBlock(prevContent) && !hasToolBlock(iContent)) {
+            formatted[i - 1] = {
+              ...formatted[i - 1],
+              content: [...prevContent, ...iContent],
+            };
+            formatted.splice(i, 1);
+            // Clamp i to the new last index — the splice may have put
+            // us past the end. If there's a new element at i (formerly
+            // i+1), it will be rechecked against the merged i-1.
+            if (i >= formatted.length) {
+              i = formatted.length - 1;
+            }
+          } else {
+            // Can't merge (tool blocks present) — leave for
+            // ensureToolPairing which handles tool_use/tool_result
+            // alternation in its own forward walk.
+            i--;
+          }
+        }
+      }
+    }
+
+    // Strip thinking/redacted_thinking blocks from completed historical
+    // assistant turns. Anthropic only requires these blocks for active
+    // tool-use continuation (the tail span where assistant tool_use is
+    // followed by user tool_result). Replaying stale thinking blocks from
+    // earlier turns causes 400 errors when the signature is no longer
+    // valid (e.g. after a provider/model/profile switch).
+    const activeToolUseStart = findActiveToolUseContinuationStart(formatted);
+    for (let i = 0; i < activeToolUseStart; i++) {
+      const msg = formatted[i];
+      if (msg.role !== "assistant" || !Array.isArray(msg.content)) continue;
+      const stripped = (msg.content as Anthropic.ContentBlockParam[]).filter(
+        (b) =>
+          typeof b === "string" ||
+          (b.type !== "thinking" && b.type !== "redacted_thinking"),
+      );
+      if (stripped.length === 0) {
+        stripped.push({
+          type: "text" as const,
+          text: PLACEHOLDER_BLOCKS_OMITTED,
+        });
+      }
+      formatted[i] = { ...msg, content: stripped };
+    }
+
+    return ensureToolPairing(repairOrphanedServerToolBlocks(formatted));
+  }
+
   /**
    * Convert a content block to Anthropic format, returning null for unknown
    * block types instead of throwing.  Unknown types (e.g. ui_surface stored

package/src/providers/atlascloud/client.ts

@@ -0,0 +1,85 @@
+import OpenAI from "openai";
+
+import { getLogger } from "../../util/logger.js";
+import { OpenAIChatCompletionsProvider } from "../openai/chat-completions-provider.js";
+
+const log = getLogger("atlascloud-client");
+
+/** Validation-specific timeout (10s) so a stalled network doesn't block key submission. */
+const VALIDATION_TIMEOUT_MS = 10_000;
+
+export interface AtlasCloudProviderOptions {
+  apiKey?: string;
+  streamTimeoutMs?: number;
+}
+
+/**
+ * Atlas Cloud exposes a single OpenAI-compatible endpoint. Unlike MiniMax it
+ * has no regional fallback host, so validation targets this URL only.
+ */
+const DEFAULT_ATLASCLOUD_BASE_URL = "https://api.atlascloud.ai/v1";
+
+/**
+ * Validate an Atlas Cloud API key by listing models against the OpenAI-compatible
+ * endpoint. Definitive auth failures (401/403) reject the key; transient errors
+ * (429, 5xx, network) allow the key to be stored so a flaky network doesn't
+ * block setup.
+ */
+export async function validateAtlasCloudApiKey(
+  apiKey: string,
+): Promise<{ valid: true } | { valid: false; reason: string }> {
+  try {
+    const client = new OpenAI({
+      apiKey,
+      baseURL: DEFAULT_ATLASCLOUD_BASE_URL,
+      timeout: VALIDATION_TIMEOUT_MS,
+      maxRetries: 0,
+    });
+    await client.models.list();
+    return { valid: true };
+  } catch (error) {
+    if (error instanceof OpenAI.APIError) {
+      if (error.status === 401) {
+        return { valid: false, reason: "API key is invalid or expired." };
+      }
+      if (error.status === 403) {
+        return {
+          valid: false,
+          reason: `Atlas Cloud API error (${error.status}): ${error.message}`,
+        };
+      }
+      // Transient errors (429, 5xx, etc.) — allow the key to be stored.
+      log.warn(
+        { status: error.status },
+        "Atlas Cloud API returned a transient error during key validation — allowing key",
+      );
+      return { valid: true };
+    }
+    // Network errors — allow the key to be stored.
+    log.warn(
+      { error: error instanceof Error ? error.message : String(error) },
+      "Network error during Atlas Cloud key validation — allowing key",
+    );
+    return { valid: true };
+  }
+}
+
+export class AtlasCloudProvider extends OpenAIChatCompletionsProvider {
+  constructor(
+    apiKey: string,
+    model: string,
+    options: AtlasCloudProviderOptions = {},
+  ) {
+    super(apiKey, model, {
+      baseURL: DEFAULT_ATLASCLOUD_BASE_URL,
+      providerName: "atlascloud",
+      providerLabel: "Atlas Cloud",
+      streamTimeoutMs: options.streamTimeoutMs,
+      // Atlas Cloud hosts reasoning models (e.g. DeepSeek) that emit chain of
+      // thought via `reasoning_content` rather than inline `<think>` tags. The
+      // base provider parses this field into thinking blocks and replays it on
+      // multi-turn requests.
+      assistantReasoningField: "reasoning_content",
+    });
+  }
+}

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

@@ -60,6 +60,12 @@ export class CallSiteRoutingProvider implements Provider {
     return this._activeProviderContext.getStore() ?? this.defaultProvider.name;
   }
 
+  // Forward the optional token-counting endpoint from the default provider —
+  // the same one whose `tokenEstimationProvider` this wrapper surfaces, and
+  // the provider that handles the main agent turn that `/compact` sizes
+  // against. Per-call connection routing only affects `sendMessage`.
+  public readonly countInputTokens?: NonNullable<Provider["countInputTokens"]>;
+
   constructor(
     private readonly defaultProvider: Provider,
     /**
@@ -88,6 +94,10 @@ export class CallSiteRoutingProvider implements Provider {
     ) => Promise<Provider | null>,
   ) {
     this.tokenEstimationProvider = defaultProvider.tokenEstimationProvider;
+    if (defaultProvider.countInputTokens) {
+      this.countInputTokens =
+        defaultProvider.countInputTokens.bind(defaultProvider);
+    }
   }
 
   async sendMessage(

package/src/providers/fetch-provider-catalog.ts

@@ -0,0 +1,85 @@
+/**
+ * Canonical catalog of built-in web-fetch providers.
+ *
+ * This is the single source of truth that drives the config-schema enum
+ * (`VALID_WEB_FETCH_PROVIDERS` in `assistant/src/config/schemas/services.ts`)
+ * and the client picker mirror in
+ * `clients/web/src/assistant/generated/web-fetch-provider-catalog.gen.ts`.
+ *
+ * Mirrors `search-provider-catalog.ts`. Where web-search distinguishes
+ * `managed` (platform proxy) from `byok` (user key), web-fetch has no managed
+ * proxy yet, so the two kinds are:
+ *
+ *   - `builtin` — the daemon's own HTTP fetch + extract path (`default`). No
+ *     key, always available, the default.
+ *   - `byok`    — an external provider that needs a user-supplied key (e.g.
+ *     `firecrawl`, which scrapes via its hosted API and returns clean
+ *     markdown, including for JavaScript-rendered pages the builtin fetcher
+ *     can't see).
+ *
+ * BYOK fetch providers intentionally reuse the SAME bare-name credential as
+ * their search counterpart (e.g. `firecrawl` → `FIRECRAWL_API_KEY`), so a
+ * single stored key powers both `web_search` and `web_fetch`. The key is
+ * already registered in `API_KEY_PROVIDERS` via `SEARCH_PROVIDER_CATALOG`;
+ * this catalog deliberately does NOT re-register it.
+ */
+
+export type FetchProviderKind = "builtin" | "byok";
+
+export interface FetchProviderCatalogEntry {
+  /** Stable provider identifier. Matches config values. */
+  readonly id: string;
+  /** Short display name used by picker UIs. */
+  readonly displayName: string;
+  /** Optional long display name for prose contexts. */
+  readonly displayNameLong?: string;
+  /** `builtin` (no key) or `byok` (user-supplied key). */
+  readonly kind: FetchProviderKind;
+  /** Placeholder shown in the API-key input. BYOK providers only. */
+  readonly apiKeyPrefix?: string;
+  /** Environment variable name carrying the API key. BYOK providers only. */
+  readonly envVar?: string;
+  /** Secret-catalog key (the bare provider name accepted by
+   *  `getProviderKeyAsync`). BYOK providers only. */
+  readonly secretKey?: string;
+  /** Privacy-policy URL surfaced in marketing data-sharing docs.
+   *  BYOK providers only. */
+  readonly privacyPolicyUrl?: string;
+}
+
+export const FETCH_PROVIDER_CATALOG: readonly FetchProviderCatalogEntry[] = [
+  {
+    id: "default",
+    displayName: "Built-in",
+    displayNameLong: "Built-in fetcher",
+    kind: "builtin",
+  },
+  {
+    id: "firecrawl",
+    displayName: "Firecrawl",
+    kind: "byok",
+    apiKeyPrefix: "fc-...",
+    envVar: "FIRECRAWL_API_KEY",
+    secretKey: "firecrawl",
+    privacyPolicyUrl: "https://www.firecrawl.dev/privacy-policy",
+  },
+];
+
+/** Provider ids accepted by the web-fetch config schema. */
+export const FETCH_PROVIDER_IDS: readonly string[] =
+  FETCH_PROVIDER_CATALOG.map((p) => p.id);
+
+/** Catalog entries that require a user-supplied API key. */
+export const BYOK_FETCH_PROVIDERS: readonly FetchProviderCatalogEntry[] =
+  FETCH_PROVIDER_CATALOG.filter((p) => p.kind === "byok");
+
+/** BYOK fetch-provider ids. */
+export const BYOK_FETCH_PROVIDER_IDS: readonly string[] =
+  BYOK_FETCH_PROVIDERS.map((p) => p.id);
+
+/** Look up a single catalog entry by id. Returns `undefined` if unknown. */
+export function getFetchProvider(
+  id: string,
+): FetchProviderCatalogEntry | undefined {
+  return FETCH_PROVIDER_CATALOG.find((p) => p.id === id);
+}

package/src/providers/gemini/client.ts

@@ -17,6 +17,7 @@ import type {
   Provider,
   ProviderResponse,
   SendMessageOptions,
+  ToolDefinition,
 } from "../types.js";
 import {
   ContextOverflowError,
@@ -518,6 +519,48 @@ export class GeminiProvider implements Provider {
     }
   }
 
+  /**
+   * Exact prompt-token count via Gemini's `models.countTokens` — the real
+   * tokenizer, no generation. Mirrors {@link sendMessage}'s composition
+   * (contents + `systemInstruction` + tool function declarations) so the count
+   * tracks what a real call would send. Throws when the endpoint returns no
+   * `totalTokens`, so the caller falls back to the local estimate.
+   */
+  async countInputTokens(
+    messages: Message[],
+    systemPrompt: string,
+    tools?: ToolDefinition[],
+  ): Promise<number> {
+    const contents = this.toGeminiContents(messages, this.model);
+    const config: genai.CountTokensConfig = {};
+    if (systemPrompt) {
+      config.systemInstruction = systemPrompt.replaceAll(
+        SYSTEM_PROMPT_CACHE_BOUNDARY,
+        "\n\n",
+      );
+    }
+    if (tools && tools.length > 0) {
+      config.tools = [
+        {
+          functionDeclarations: tools.map((t) => ({
+            name: t.name,
+            description: t.description,
+            parametersJsonSchema: t.input_schema,
+          })),
+        },
+      ];
+    }
+    const res = await this.client.models.countTokens({
+      model: this.model,
+      contents,
+      ...(Object.keys(config).length > 0 ? { config } : {}),
+    });
+    if (typeof res.totalTokens !== "number") {
+      throw new Error("Gemini countTokens returned no totalTokens");
+    }
+    return res.totalTokens;
+  }
+
   /** Convert neutral messages to Gemini Content[] format. */
   private toGeminiContents(
     messages: Message[],

package/src/providers/inference/adapter-factory.ts

@@ -19,6 +19,7 @@
  */
 
 import { AnthropicProvider } from "../anthropic/client.js";
+import { AtlasCloudProvider } from "../atlascloud/client.js";
 import { FireworksProvider } from "../fireworks/client.js";
 import { GeminiProvider } from "../gemini/client.js";
 import { MinimaxProvider } from "../minimax/client.js";
@@ -114,6 +115,8 @@ const ADAPTER_FACTORIES: Record<string, AdapterFactory> = {
     }),
   minimax: ({ apiKey, model, streamTimeoutMs }) =>
     new MinimaxProvider(apiKey, model, { streamTimeoutMs }),
+  atlascloud: ({ apiKey, model, streamTimeoutMs }) =>
+    new AtlasCloudProvider(apiKey, model, { streamTimeoutMs }),
 };
 
 /**