clanker-agent/packages/ai/test/tokens.test.ts

import { describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { stream } from "../src/stream.js";
import type { Api, Context, Model, StreamOptions } from "../src/types.js";

type StreamOptionsWithExtras = StreamOptions & Record<string, unknown>;

import {
  hasAzureOpenAICredentials,
  resolveAzureDeploymentName,
} from "./azure-utils.js";
import { hasBedrockCredentials } from "./bedrock-utils.js";
import { resolveApiKey } from "./oauth.js";

// Resolve OAuth tokens at module level (async, runs before tests)
const oauthTokens = await Promise.all([
  resolveApiKey("anthropic"),
  resolveApiKey("github-copilot"),
  resolveApiKey("google-gemini-cli"),
  resolveApiKey("google-antigravity"),
  resolveApiKey("openai-codex"),
]);
const [
  anthropicOAuthToken,
  githubCopilotToken,
  geminiCliToken,
  antigravityToken,
  openaiCodexToken,
] = oauthTokens;

async function testTokensOnAbort<TApi extends Api>(
  llm: Model<TApi>,
  options: StreamOptionsWithExtras = {},
) {
  const context: Context = {
    messages: [
      {
        role: "user",
        content:
          "Write a long poem with 20 stanzas about the beauty of nature.",
        timestamp: Date.now(),
      },
    ],
    systemPrompt: "You are a helpful assistant.",
  };

  const controller = new AbortController();
  const response = stream(llm, context, {
    ...options,
    signal: controller.signal,
  });

  let abortFired = false;
  let text = "";
  for await (const event of response) {
    if (
      !abortFired &&
      (event.type === "text_delta" || event.type === "thinking_delta")
    ) {
      text += event.delta;
      if (text.length >= 1000) {
        abortFired = true;
        controller.abort();
      }
    }
  }

  const msg = await response.result();

  expect(msg.stopReason).toBe("aborted");

  // OpenAI providers, OpenAI Codex, Gemini CLI, zai, Amazon Bedrock, and the GPT-OSS model on Antigravity only send usage in the final chunk,
  // so when aborted they have no token stats. Anthropic and Google send usage information early in the stream.
  // MiniMax reports input tokens but not output tokens when aborted.
  if (
    llm.api === "openai-completions" ||
    llm.api === "mistral-conversations" ||
    llm.api === "openai-responses" ||
    llm.api === "azure-openai-responses" ||
    llm.api === "openai-codex-responses" ||
    llm.provider === "google-gemini-cli" ||
    llm.provider === "zai" ||
    llm.provider === "amazon-bedrock" ||
    llm.provider === "vercel-ai-gateway" ||
    (llm.provider === "google-antigravity" && llm.id.includes("gpt-oss"))
  ) {
    expect(msg.usage.input).toBe(0);
    expect(msg.usage.output).toBe(0);
  } else if (llm.provider === "minimax") {
    // MiniMax reports input tokens early but output tokens only in final chunk
    expect(msg.usage.input).toBeGreaterThan(0);
    expect(msg.usage.output).toBe(0);
  } else {
    expect(msg.usage.input).toBeGreaterThan(0);
    expect(msg.usage.output).toBeGreaterThan(0);

    // Some providers (Antigravity, Copilot) have zero cost rates
    if (llm.cost.input > 0) {
      expect(msg.usage.cost.input).toBeGreaterThan(0);
      expect(msg.usage.cost.total).toBeGreaterThan(0);
    }
  }
}

describe("Token Statistics on Abort", () => {
  describe.skipIf(!process.env.GEMINI_API_KEY)("Google Provider", () => {
    const llm = getModel("google", "gemini-2.5-flash");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm, { thinking: { enabled: true } });
      },
    );
  });

  describe.skipIf(!process.env.OPENAI_API_KEY)(
    "OpenAI Completions Provider",
    () => {
      const { compat: _compat, ...baseModel } = getModel(
        "openai",
        "gpt-4o-mini",
      )!;
      void _compat;
      const llm: Model<"openai-completions"> = {
        ...baseModel,
        api: "openai-completions",
      };

      it(
        "should include token stats when aborted mid-stream",
        { retry: 3, timeout: 30000 },
        async () => {
          await testTokensOnAbort(llm);
        },
      );
    },
  );

  describe.skipIf(!process.env.OPENAI_API_KEY)(
    "OpenAI Responses Provider",
    () => {
      const llm = getModel("openai", "gpt-5-mini");

      it(
        "should include token stats when aborted mid-stream",
        { retry: 3, timeout: 30000 },
        async () => {
          await testTokensOnAbort(llm);
        },
      );
    },
  );

  describe.skipIf(!hasAzureOpenAICredentials())(
    "Azure OpenAI Responses Provider",
    () => {
      const llm = getModel("azure-openai-responses", "gpt-4o-mini");
      const azureDeploymentName = resolveAzureDeploymentName(llm.id);
      const azureOptions = azureDeploymentName ? { azureDeploymentName } : {};

      it(
        "should include token stats when aborted mid-stream",
        { retry: 3, timeout: 30000 },
        async () => {
          await testTokensOnAbort(llm, azureOptions);
        },
      );
    },
  );

  describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic Provider", () => {
    const llm = getModel("anthropic", "claude-3-5-haiku-20241022");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.XAI_API_KEY)("xAI Provider", () => {
    const llm = getModel("xai", "grok-3-fast");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.GROQ_API_KEY)("Groq Provider", () => {
    const llm = getModel("groq", "openai/gpt-oss-20b");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.CEREBRAS_API_KEY)("Cerebras Provider", () => {
    const llm = getModel("cerebras", "gpt-oss-120b");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.HF_TOKEN)("Hugging Face Provider", () => {
    const llm = getModel("huggingface", "moonshotai/Kimi-K2.5");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.ZAI_API_KEY)("zAI Provider", () => {
    const llm = getModel("zai", "glm-4.5-flash");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.MISTRAL_API_KEY)("Mistral Provider", () => {
    const llm = getModel("mistral", "devstral-medium-latest");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.MINIMAX_API_KEY)("MiniMax Provider", () => {
    const llm = getModel("minimax", "MiniMax-M2.1");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider", () => {
    const llm = getModel("kimi-coding", "kimi-k2-thinking");

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });

  describe.skipIf(!process.env.AI_GATEWAY_API_KEY)(
    "Vercel AI Gateway Provider",
    () => {
      const llm = getModel("vercel-ai-gateway", "google/gemini-2.5-flash");

      it(
        "should include token stats when aborted mid-stream",
        { retry: 3, timeout: 30000 },
        async () => {
          await testTokensOnAbort(llm);
        },
      );
    },
  );

  // =========================================================================
  // OAuth-based providers (credentials from ~/.companion/agent/oauth.json)
  // =========================================================================

  describe("Anthropic OAuth Provider", () => {
    const llm = getModel("anthropic", "claude-3-5-haiku-20241022");

    it.skipIf(!anthropicOAuthToken)(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm, { apiKey: anthropicOAuthToken });
      },
    );
  });

  describe("GitHub Copilot Provider", () => {
    it.skipIf(!githubCopilotToken)(
      "gpt-4o - should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        const llm = getModel("github-copilot", "gpt-4o");
        await testTokensOnAbort(llm, { apiKey: githubCopilotToken });
      },
    );

    it.skipIf(!githubCopilotToken)(
      "claude-sonnet-4 - should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        const llm = getModel("github-copilot", "claude-sonnet-4");
        await testTokensOnAbort(llm, { apiKey: githubCopilotToken });
      },
    );
  });

  describe("Google Gemini CLI Provider", () => {
    it.skipIf(!geminiCliToken)(
      "gemini-2.5-flash - should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
        await testTokensOnAbort(llm, { apiKey: geminiCliToken });
      },
    );
  });

  describe("Google Antigravity Provider", () => {
    it.skipIf(!antigravityToken)(
      "gemini-3-flash - should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        const llm = getModel("google-antigravity", "gemini-3-flash");
        await testTokensOnAbort(llm, { apiKey: antigravityToken });
      },
    );

    it.skipIf(!antigravityToken)(
      "claude-sonnet-4-5 - should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        const llm = getModel("google-antigravity", "claude-sonnet-4-5");
        await testTokensOnAbort(llm, { apiKey: antigravityToken });
      },
    );

    it.skipIf(!antigravityToken)(
      "gpt-oss-120b-medium - should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        const llm = getModel("google-antigravity", "gpt-oss-120b-medium");
        await testTokensOnAbort(llm, { apiKey: antigravityToken });
      },
    );
  });

  describe("OpenAI Codex Provider", () => {
    it.skipIf(!openaiCodexToken)(
      "gpt-5.2-codex - should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        const llm = getModel("openai-codex", "gpt-5.2-codex");
        await testTokensOnAbort(llm, { apiKey: openaiCodexToken });
      },
    );
  });

  describe.skipIf(!hasBedrockCredentials())("Amazon Bedrock Provider", () => {
    const llm = getModel(
      "amazon-bedrock",
      "global.anthropic.claude-sonnet-4-5-20250929-v1:0",
    );

    it(
      "should include token stats when aborted mid-stream",
      { retry: 3, timeout: 30000 },
      async () => {
        await testTokensOnAbort(llm);
      },
    );
  });
});