import { afterEach, describe, expect, it } from "vitest";
import {
complete,
fauxAssistantMessage,
fauxText,
fauxThinking,
fauxToolCall,
registerFauxProvider,
stream,
Type,
} from "../src/index.js";
import type { AssistantMessageEvent, Context } from "../src/types.js";
async function collectEvents(streamResult: ReturnType<typeof stream>): Promise<AssistantMessageEvent[]> {
const events: AssistantMessageEvent[] = [];
for await (const event of streamResult) {
events.push(event);
}
return events;
}
const registrations: Array<{ unregister: () => void }> = [];
afterEach(() => {
for (const registration of registrations.splice(0)) {
registration.unregister();
}
});
describe("faux provider", () => {
it("registers a custom provider and estimates usage", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("hello world")]);
const context: Context = {
systemPrompt: "Be concise.",
messages: [{ role: "user", content: "hi there", timestamp: Date.now() }],
};
const response = await complete(registration.getModel(), context);
expect(response.content).toEqual([{ type: "text", text: "hello world" }]);
expect(response.usage.input).toBeGreaterThan(0);
expect(response.usage.output).toBeGreaterThan(0);
expect(response.usage.totalTokens).toBe(response.usage.input + response.usage.output);
expect(registration.state.callCount).toBe(1);
});
it("supports helper blocks for text, thinking, and tool calls", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([
fauxAssistantMessage([fauxThinking("think"), fauxToolCall("echo", { text: "hi" }), fauxText("done")], {
stopReason: "toolUse",
}),
]);
const response = await complete(registration.getModel(), {
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
});
expect(response.content).toEqual([
{ type: "thinking", thinking: "think" },
{ type: "toolCall", id: expect.any(String), name: "echo", arguments: { text: "hi" } },
{ type: "text", text: "done" },
]);
expect(response.stopReason).toBe("toolUse");
});
it("supports multiple models with per-model reasoning and model-aware factories", async () => {
const registration = registerFauxProvider({
models: [
{ id: "faux-fast", name: "Faux Fast", reasoning: false },
{ id: "faux-thinker", name: "Faux Thinker", reasoning: true },
],
});
registrations.push(registration);
registration.setResponses([
(_context, _options, _state, model) => fauxAssistantMessage(`${model.id}:${String(model.reasoning)}`),
(_context, _options, _state, model) => fauxAssistantMessage(`${model.id}:${String(model.reasoning)}`),
]);
expect(registration.models.map((model) => model.id)).toEqual(["faux-fast", "faux-thinker"]);
expect(registration.getModel()).toBe(registration.models[0]);
expect(registration.getModel("faux-fast")?.reasoning).toBe(false);
expect(registration.getModel("faux-thinker")?.reasoning).toBe(true);
const fast = await complete(registration.getModel("faux-fast")!, {
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
});
const thinker = await complete(registration.getModel("faux-thinker")!, {
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
});
expect(fast.content).toEqual([{ type: "text", text: "faux-fast:false" }]);
expect(thinker.content).toEqual([{ type: "text", text: "faux-thinker:true" }]);
});
it("rewrites api, provider, and model on returned messages", async () => {
const registration = registerFauxProvider({
api: "faux:test",
provider: "faux-provider",
models: [{ id: "faux-model" }],
});
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("hello")]);
const response = await complete(registration.getModel(), {
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
});
expect(response.api).toBe("faux:test");
expect(response.provider).toBe("faux-provider");
expect(response.model).toBe("faux-model");
});
it("consumes queued responses in order and errors when exhausted", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
const context: Context = {
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
};
const first = await complete(registration.getModel(), context);
const second = await complete(registration.getModel(), context);
const exhausted = await complete(registration.getModel(), context);
expect(first.content).toEqual([{ type: "text", text: "first" }]);
expect(second.content).toEqual([{ type: "text", text: "second" }]);
expect(exhausted.stopReason).toBe("error");
expect(exhausted.errorMessage).toBe("No more faux responses queued");
expect(registration.getPendingResponseCount()).toBe(0);
expect(registration.state.callCount).toBe(3);
});
it("can replace and append queued responses", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("first")]);
const context: Context = {
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
};
expect((await complete(registration.getModel(), context)).content).toEqual([{ type: "text", text: "first" }]);
expect(registration.getPendingResponseCount()).toBe(0);
registration.setResponses([fauxAssistantMessage("second")]);
expect(registration.getPendingResponseCount()).toBe(1);
expect((await complete(registration.getModel(), context)).content).toEqual([{ type: "text", text: "second" }]);
registration.appendResponses([fauxAssistantMessage("third"), fauxAssistantMessage("fourth")]);
expect(registration.getPendingResponseCount()).toBe(2);
expect((await complete(registration.getModel(), context)).content).toEqual([{ type: "text", text: "third" }]);
expect((await complete(registration.getModel(), context)).content).toEqual([{ type: "text", text: "fourth" }]);
expect(registration.getPendingResponseCount()).toBe(0);
});
it("supports async response factories", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([
async (context, _options, state) => fauxAssistantMessage(`${context.messages.length}:${state.callCount}`),
]);
const response = await complete(registration.getModel(), {
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
});
expect(response.content).toEqual([{ type: "text", text: "1:1" }]);
});
it("emits an error when a response factory throws", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([
() => {
throw new Error("boom");
},
]);
const events = await collectEvents(
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
);
expect(events).toHaveLength(1);
expect(events[0].type).toBe("error");
if (events[0].type === "error") {
expect(events[0].error.stopReason).toBe("error");
expect(events[0].error.errorMessage).toBe("boom");
}
});
it("estimates prompt and output tokens from serialized context", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("done")]);
const tool = {
name: "echo",
description: "Echo back text",
parameters: Type.Object({ text: Type.String() }),
};
const context: Context = {
systemPrompt: "sys",
messages: [
{
role: "user",
content: [
{ type: "text", text: "hello" },
{ type: "image", mimeType: "image/png", data: "abcd" },
],
timestamp: 1,
},
fauxAssistantMessage("prior"),
{
role: "toolResult",
toolCallId: "tool-1",
toolName: "echo",
content: [{ type: "text", text: "tool out" }],
isError: false,
timestamp: 2,
},
],
tools: [tool],
};
const response = await complete(registration.getModel(), context);
const promptText = [
"system:sys",
"user:hello\n[image:image/png:4]",
"assistant:prior",
"toolResult:echo\ntool out",
`tools:${JSON.stringify([tool])}`,
].join("\n\n");
const expectedPromptTokens = Math.ceil(promptText.length / 4);
const expectedOutputTokens = Math.ceil("done".length / 4);
expect(response.usage.input).toBe(expectedPromptTokens);
expect(response.usage.output).toBe(expectedOutputTokens);
expect(response.usage.cacheRead).toBe(0);
expect(response.usage.cacheWrite).toBe(0);
expect(response.usage.totalTokens).toBe(expectedPromptTokens + expectedOutputTokens);
});
it("does not share cache across sessions or requests without sessionId", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([
fauxAssistantMessage("first"),
fauxAssistantMessage("second"),
fauxAssistantMessage("third"),
]);
const context: Context = {
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
};
const first = await complete(registration.getModel(), context, {
sessionId: "session-1",
cacheRetention: "short",
});
expect(first.usage.cacheWrite).toBeGreaterThan(0);
context.messages.push(first);
context.messages.push({ role: "user", content: "follow up", timestamp: Date.now() + 1 });
const second = await complete(registration.getModel(), context, {
sessionId: "session-2",
cacheRetention: "short",
});
expect(second.usage.cacheRead).toBe(0);
expect(second.usage.cacheWrite).toBeGreaterThan(0);
const third = await complete(registration.getModel(), context);
expect(third.usage.cacheRead).toBe(0);
expect(third.usage.cacheWrite).toBe(0);
});
it("simulates prompt caching per sessionId", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
const context: Context = {
systemPrompt: "Be concise.",
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
};
const first = await complete(registration.getModel(), context, {
sessionId: "session-1",
cacheRetention: "short",
});
expect(first.usage.cacheRead).toBe(0);
expect(first.usage.cacheWrite).toBeGreaterThan(0);
context.messages.push(first);
context.messages.push({ role: "user", content: "follow up", timestamp: Date.now() + 1 });
const second = await complete(registration.getModel(), context, {
sessionId: "session-1",
cacheRetention: "short",
});
expect(second.usage.cacheRead).toBeGreaterThan(0);
expect(second.usage.input + second.usage.cacheRead).toBeGreaterThan(second.usage.input);
});
it("does not simulate caching when cacheRetention is none", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
const context: Context = {
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
};
await complete(registration.getModel(), context, { sessionId: "session-1", cacheRetention: "none" });
context.messages.push(fauxAssistantMessage("first"));
context.messages.push({ role: "user", content: "follow up", timestamp: Date.now() + 1 });
const second = await complete(registration.getModel(), context, {
sessionId: "session-1",
cacheRetention: "none",
});
expect(second.usage.cacheRead).toBe(0);
expect(second.usage.cacheWrite).toBe(0);
});
it("streams thinking, text, and partial tool call deltas", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([
fauxAssistantMessage(
[
fauxThinking("thinking text"),
fauxText("answer text"),
fauxToolCall("echo", { text: "hi", count: 12 }, { id: "tool-1" }),
],
{ stopReason: "toolUse" },
),
]);
const events: string[] = [];
const toolCallDeltas: string[] = [];
const s = stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] });
for await (const event of s) {
events.push(event.type);
if (event.type === "toolcall_delta") {
toolCallDeltas.push(event.delta);
}
}
expect(events).toContain("thinking_start");
expect(events).toContain("thinking_delta");
expect(events).toContain("text_start");
expect(events).toContain("text_delta");
expect(events).toContain("toolcall_start");
expect(events).toContain("toolcall_delta");
expect(events).toContain("toolcall_end");
expect(toolCallDeltas.length).toBeGreaterThan(1);
expect(JSON.parse(toolCallDeltas.join(""))).toEqual({ text: "hi", count: 12 });
});
it("streams an exact event order for fixed-size chunks", async () => {
const registration = registerFauxProvider({ tokenSize: { min: 1, max: 1 } });
registrations.push(registration);
registration.setResponses([
fauxAssistantMessage([fauxThinking("go"), fauxText("ok"), fauxToolCall("echo", {}, { id: "tool-1" })], {
stopReason: "toolUse",
}),
]);
const events = await collectEvents(
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
);
expect(events.map((event) => event.type)).toEqual([
"start",
"thinking_start",
"thinking_delta",
"thinking_end",
"text_start",
"text_delta",
"text_end",
"toolcall_start",
"toolcall_delta",
"toolcall_end",
"done",
]);
});
it("streams multiple tool calls in one message", async () => {
const registration = registerFauxProvider();
registrations.push(registration);
registration.setResponses([
fauxAssistantMessage(
[
fauxToolCall("echo", { text: "one" }, { id: "tool-1" }),
fauxToolCall("echo", { text: "two" }, { id: "tool-2" }),
],
{ stopReason: "toolUse" },
),
]);
const events = await collectEvents(
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
);
expect(events.filter((event) => event.type === "toolcall_start")).toHaveLength(2);
expect(events.filter((event) => event.type === "toolcall_end")).toHaveLength(2);
});
it("streams an explicit assistant error message as a terminal error", async () => {
const registration = registerFauxProvider({ tokenSize: { min: 2, max: 2 } });
registrations.push(registration);
registration.setResponses([
{
...fauxAssistantMessage("partial"),
stopReason: "error",
errorMessage: "upstream failed",
},
]);
const events = await collectEvents(
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
);
expect(events.map((event) => event.type)).toEqual(["start", "text_start", "text_delta", "text_end", "error"]);
const terminal = events[events.length - 1];
expect(terminal.type).toBe("error");
if (terminal.type === "error") {
expect(terminal.reason).toBe("error");
expect(terminal.error.stopReason).toBe("error");
expect(terminal.error.errorMessage).toBe("upstream failed");
}
});
it("streams an explicit assistant aborted message as a terminal error", async () => {
const registration = registerFauxProvider({ tokenSize: { min: 2, max: 2 } });
registrations.push(registration);
registration.setResponses([
{
...fauxAssistantMessage("partial"),
stopReason: "aborted",
errorMessage: "Request was aborted",
},
]);
const events = await collectEvents(
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
);
expect(events.map((event) => event.type)).toEqual(["start", "text_start", "text_delta", "text_end", "error"]);
const terminal = events[events.length - 1];
expect(terminal.type).toBe("error");
if (terminal.type === "error") {
expect(terminal.reason).toBe("aborted");
expect(terminal.error.stopReason).toBe("aborted");
expect(terminal.error.errorMessage).toBe("Request was aborted");
}
});
it("supports aborting before the first chunk", async () => {
const registration = registerFauxProvider({ tokensPerSecond: 50, tokenSize: { min: 3, max: 3 } });
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("abcdefghijklmnopqrstuvwxyz")]);
const controller = new AbortController();
controller.abort();
const events = await collectEvents(
stream(
registration.getModel(),
{ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] },
{ signal: controller.signal },
),
);
expect(events).toHaveLength(1);
expect(events[0].type).toBe("error");
if (events[0].type === "error") {
expect(events[0].reason).toBe("aborted");
expect(events[0].error.stopReason).toBe("aborted");
}
});
it("supports aborting mid-text stream when paced", async () => {
const registration = registerFauxProvider({ tokensPerSecond: 100, tokenSize: { min: 3, max: 3 } });
registrations.push(registration);
registration.setResponses([fauxAssistantMessage("abcdefghijklmnopqrstuvwxyz")]);
const controller = new AbortController();
const events: string[] = [];
let textDeltaCount = 0;
const s = stream(
registration.getModel(),
{ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] },
{ signal: controller.signal },
);
for await (const event of s) {
events.push(event.type);
if (event.type === "text_delta") {
textDeltaCount++;
controller.abort();
}
}
expect(textDeltaCount).toBe(1);
expect(events).toContain("text_start");
expect(events).toContain("text_delta");
expect(events).toContain("error");
expect(events).not.toContain("text_end");
});
it("supports aborting mid-thinking stream when paced", async () => {
const registration = registerFauxProvider({ tokensPerSecond: 100, tokenSize: { min: 3, max: 3 } });
registrations.push(registration);
registration.setResponses([
{
...fauxAssistantMessage("ignored"),
content: [{ type: "thinking", thinking: "abcdefghijklmnopqrstuvwxyz" }],
},
]);
const controller = new AbortController();
const events: string[] = [];
let thinkingDeltaCount = 0;
const s = stream(
registration.getModel(),
{ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] },
{ signal: controller.signal },
);
for await (const event of s) {
events.push(event.type);
if (event.type === "thinking_delta") {
thinkingDeltaCount++;
controller.abort();
}
}
expect(thinkingDeltaCount).toBe(1);
expect(events).toContain("thinking_start");
expect(events).toContain("thinking_delta");
expect(events).toContain("error");
expect(events).not.toContain("thinking_end");
});
it("supports aborting mid-toolcall stream when paced", async () => {
const registration = registerFauxProvider({ tokensPerSecond: 100, tokenSize: { min: 3, max: 3 } });
registrations.push(registration);
registration.setResponses([
{
...fauxAssistantMessage("done"),
content: [
{
type: "toolCall",
id: "tool-1",
name: "echo",
arguments: { text: "abcdefghijklmnopqrstuvwxyz", count: 123456789 },
},
],
stopReason: "toolUse",
},
]);
const controller = new AbortController();
const events: string[] = [];
let toolCallDeltaCount = 0;
const s = stream(
registration.getModel(),
{ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] },
{ signal: controller.signal },
);
for await (const event of s) {
events.push(event.type);
if (event.type === "toolcall_delta") {
toolCallDeltaCount++;
controller.abort();
}
}
expect(toolCallDeltaCount).toBe(1);
expect(events).toContain("toolcall_start");
expect(events).toContain("toolcall_delta");
expect(events).toContain("error");
expect(events).not.toContain("toolcall_end");
});
it("unregisters the provider", async () => {
const registration = registerFauxProvider();
registration.setResponses([fauxAssistantMessage("hello")]);
registration.unregister();
await expect(
complete(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
).rejects.toThrow(`No API provider registered for api: ${registration.api}`);
});
});