import { afterEach, describe, expect, it } from "vitest";
import {
complete,
fauxAssistantMessage,
fauxText,
fauxThinking,
fauxToolCall,
registerFauxProvider,
stream,
Type,
} from "../src/index.js";
import type { AssistantMessageEvent, Context, ProviderStreamOptions } 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();
}
});
function register(options?: Parameters<typeof registerFauxProvider>[0]) {
const registration = registerFauxProvider(options);
registrations.push(registration);
return registration;
}
const userTurn = (): Context => ({ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] });
describe("faux provider", () => {
it("consumes queued responses in order and errors when exhausted", async () => {
const registration = register();
registration.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
const context = userTurn();
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("emits an error when a response factory throws", async () => {
const registration = register();
registration.setResponses([
() => {
throw new Error("boom");
},
]);
const events = await collectEvents(stream(registration.getModel(), userTurn()));
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 = register();
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);
});
// Two sequential requests over a growing context; only a repeated sessionId with a real
// retention window may report a cache hit on the second request.
it.each([
{
name: "reuses the prefix for a repeated sessionId",
first: { sessionId: "session-1", cacheRetention: "short" } as ProviderStreamOptions,
second: { sessionId: "session-1", cacheRetention: "short" } as ProviderStreamOptions,
cacheHit: true,
cacheWrite: true,
},
{
name: "does not share the cache across sessions",
first: { sessionId: "session-1", cacheRetention: "short" } as ProviderStreamOptions,
second: { sessionId: "session-2", cacheRetention: "short" } as ProviderStreamOptions,
cacheHit: false,
cacheWrite: true,
},
{
name: "does not cache without a sessionId",
first: undefined,
second: undefined,
cacheHit: false,
cacheWrite: false,
},
{
name: "does not cache when cacheRetention is none",
first: { sessionId: "session-1", cacheRetention: "none" } as ProviderStreamOptions,
second: { sessionId: "session-1", cacheRetention: "none" } as ProviderStreamOptions,
cacheHit: false,
cacheWrite: false,
},
])("simulates prompt caching: $name", async ({ first, second, cacheHit, cacheWrite }) => {
const registration = register();
registration.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
const context: Context = {
systemPrompt: "Be concise.",
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
};
const firstResponse = await complete(registration.getModel(), context, first);
context.messages.push(firstResponse);
context.messages.push({ role: "user", content: "follow up", timestamp: Date.now() + 1 });
const secondResponse = await complete(registration.getModel(), context, second);
if (cacheWrite) {
expect(firstResponse.usage.cacheWrite).toBeGreaterThan(0);
} else {
expect(firstResponse.usage.cacheWrite).toBe(0);
}
expect(firstResponse.usage.cacheRead).toBe(0);
if (cacheHit) {
expect(secondResponse.usage.cacheRead).toBeGreaterThan(0);
} else {
expect(secondResponse.usage.cacheRead).toBe(0);
}
});
it("streams an exact event order for fixed-size chunks", async () => {
const registration = register({ tokenSize: { min: 1, max: 1 } });
registration.setResponses([
fauxAssistantMessage([fauxThinking("go"), fauxText("ok"), fauxToolCall("echo", {}, { id: "tool-1" })], {
stopReason: "toolUse",
}),
]);
const events = await collectEvents(stream(registration.getModel(), userTurn()));
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("reassembles tool call arguments from partial deltas", async () => {
const registration = register();
registration.setResponses([
fauxAssistantMessage([fauxToolCall("echo", { text: "hi", count: 12 }, { id: "tool-1" })], {
stopReason: "toolUse",
}),
]);
const deltas: string[] = [];
for await (const event of stream(registration.getModel(), userTurn())) {
if (event.type === "toolcall_delta") deltas.push(event.delta);
}
expect(deltas.length).toBeGreaterThan(1);
expect(JSON.parse(deltas.join(""))).toEqual({ text: "hi", count: 12 });
});
it("streams multiple tool calls in one message", async () => {
const registration = register();
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(), userTurn()));
expect(events.filter((event) => event.type === "toolcall_start")).toHaveLength(2);
expect(events.filter((event) => event.type === "toolcall_end")).toHaveLength(2);
});
// An explicit terminal stopReason must flush the partial text and then end on a single error event.
it.each([
{ stopReason: "error" as const, errorMessage: "upstream failed", reason: "error" },
{ stopReason: "aborted" as const, errorMessage: "Request was aborted", reason: "aborted" },
])("streams an explicit $stopReason message as a terminal error", async ({ stopReason, errorMessage, reason }) => {
const registration = register({ tokenSize: { min: 2, max: 2 } });
registration.setResponses([{ ...fauxAssistantMessage("partial"), stopReason, errorMessage }]);
const events = await collectEvents(stream(registration.getModel(), userTurn()));
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(reason);
expect(terminal.error.stopReason).toBe(stopReason);
expect(terminal.error.errorMessage).toBe(errorMessage);
}
});
it("supports aborting before the first chunk", async () => {
const registration = register({ tokensPerSecond: 50, tokenSize: { min: 3, max: 3 } });
registration.setResponses([fauxAssistantMessage("abcdefghijklmnopqrstuvwxyz")]);
const controller = new AbortController();
controller.abort();
const events = await collectEvents(stream(registration.getModel(), userTurn(), { 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");
}
});
// Aborting on the first delta of a paced stream must stop the block without emitting its *_end event.
it.each([
{
kind: "text",
response: fauxAssistantMessage("abcdefghijklmnopqrstuvwxyz"),
deltaType: "text_delta" as const,
startType: "text_start",
endType: "text_end",
},
{
kind: "thinking",
response: {
...fauxAssistantMessage("ignored"),
content: [{ type: "thinking" as const, thinking: "abcdefghijklmnopqrstuvwxyz" }],
},
deltaType: "thinking_delta" as const,
startType: "thinking_start",
endType: "thinking_end",
},
{
kind: "toolcall",
response: {
...fauxAssistantMessage("done"),
content: [
{
type: "toolCall" as const,
id: "tool-1",
name: "echo",
arguments: { text: "abcdefghijklmnopqrstuvwxyz", count: 123456789 },
},
],
stopReason: "toolUse" as const,
},
deltaType: "toolcall_delta" as const,
startType: "toolcall_start",
endType: "toolcall_end",
},
])("supports aborting mid-$kind stream when paced", async ({ response, deltaType, startType, endType }) => {
const registration = register({ tokensPerSecond: 100, tokenSize: { min: 3, max: 3 } });
registration.setResponses([response]);
const controller = new AbortController();
const events: string[] = [];
let deltaCount = 0;
for await (const event of stream(registration.getModel(), userTurn(), { signal: controller.signal })) {
events.push(event.type);
if (event.type === deltaType) {
deltaCount++;
controller.abort();
}
}
expect(deltaCount).toBe(1);
expect(events).toContain(startType);
expect(events).toContain(deltaType);
expect(events).toContain("error");
expect(events).not.toContain(endType);
});
});