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";
import { getKimiCodingTestModel } from "./kimi-test-model.js";
import { getZaiTestModel } from "./zai-test-model.js";
type StreamOptionsWithExtras = StreamOptions & Record<string, unknown>;
import { hasAzureOpenAICredentials, resolveAzureDeploymentName } from "./azure-utils.js";
import { hasBedrockCredentials } from "./bedrock-utils.js";
import { hasCloudflareWorkersAICredentials } from "./cloudflare-utils.js";
import { resolveApiKey } from "./oauth.js";
const oauthTokens = await Promise.all([
resolveApiKey("anthropic"),
resolveApiKey("github-copilot"),
resolveApiKey("openai-codex"),
]);
const [anthropicOAuthToken, githubCopilotToken, 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");
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 === "zai" ||
llm.provider === "amazon-bedrock" ||
llm.provider === "vercel-ai-gateway"
) {
expect(msg.usage.input).toBe(0);
expect(msg.usage.output).toBe(0);
} else if (llm.provider === "minimax") {
expect(msg.usage.input).toBe(0);
expect(msg.usage.output).toBe(0);
} else if (llm.provider === "kimi-coding") {
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);
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.4-mini");
it("should include token stats when aborted mid-stream", { retry: 3, timeout: 30000 }, async () => {
await testTokensOnAbort(llm, { reasoningEffort: "low" });
});
});
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-sonnet-4-6");
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-4.3");
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(!hasCloudflareWorkersAICredentials())("Cloudflare Workers AI Provider", () => {
const llm = getModel("cloudflare-workers-ai", "@cf/moonshotai/kimi-k2.6");
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 = getZaiTestModel();
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.7");
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 = getKimiCodingTestModel();
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);
});
});
describe.skipIf(!process.env.XIAOMI_API_KEY)("Xiaomi MiMo (API billing) Provider", () => {
const llm = getModel("xiaomi", "mimo-v2.5-pro");
// Xiaomi only reports this streaming usage at message_stop, after an abort.
it.skip("should include token stats when aborted mid-stream", { retry: 3, timeout: 30000 }, async () => {
await testTokensOnAbort(llm);
});
});
describe.skipIf(!process.env.XIAOMI_TOKEN_PLAN_CN_API_KEY)("Xiaomi MiMo Token Plan (CN) Provider", () => {
const llm = getModel("xiaomi-token-plan-cn", "mimo-v2.5-pro");
it.skip("should include token stats when aborted mid-stream", { retry: 3, timeout: 30000 }, async () => {
await testTokensOnAbort(llm);
});
});
describe.skipIf(!process.env.XIAOMI_TOKEN_PLAN_AMS_API_KEY)("Xiaomi MiMo Token Plan (AMS) Provider", () => {
const llm = getModel("xiaomi-token-plan-ams", "mimo-v2.5-pro");
it.skip("should include token stats when aborted mid-stream", { retry: 3, timeout: 30000 }, async () => {
await testTokensOnAbort(llm);
});
});
describe.skipIf(!process.env.XIAOMI_TOKEN_PLAN_SGP_API_KEY)("Xiaomi MiMo Token Plan (SGP) Provider", () => {
const llm = getModel("xiaomi-token-plan-sgp", "mimo-v2.5-pro");
it.skip("should include token stats when aborted mid-stream", { retry: 3, timeout: 30000 }, async () => {
await testTokensOnAbort(llm);
});
});
describe("Anthropic OAuth Provider", () => {
const llm = getModel("anthropic", "claude-sonnet-4-6");
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-5-mini - should include token stats when aborted mid-stream",
{ retry: 3, timeout: 30000 },
async () => {
const llm = getModel("github-copilot", "gpt-5-mini");
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.5");
await testTokensOnAbort(llm, { apiKey: githubCopilotToken });
},
);
});
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);
});
});
});