import { createServer, type IncomingHttpHeaders, type Server } from "node:http";
import type { AddressInfo } from "node:net";
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest";
import { getModel, getModels } from "../src/models.js";
import { openCodeSessionHeaders } from "../src/providers/opencode-headers.js";
import { complete, completeSimple, stream } from "../src/stream.js";
import type { Api, AssistantMessage, Context, Model, SimpleStreamOptions } from "../src/types.js";
describe("Cache Retention (PI_CACHE_RETENTION)", () => {
const originalEnv = process.env.PI_CACHE_RETENTION;
beforeEach(() => {
delete process.env.PI_CACHE_RETENTION;
});
afterEach(() => {
if (originalEnv !== undefined) {
process.env.PI_CACHE_RETENTION = originalEnv;
} else {
delete process.env.PI_CACHE_RETENTION;
}
});
const context: Context = {
systemPrompt: "You are a helpful assistant.",
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
};
describe("Anthropic Provider", () => {
it.skipIf(!process.env.ANTHROPIC_API_KEY)(
"should use default cache TTL (no ttl field) when PI_CACHE_RETENTION is not set",
async () => {
const model = getModel("anthropic", "claude-haiku-4-5");
let capturedPayload: any = null;
const s = stream(model, context, {
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const _ of s) {
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.system).toBeDefined();
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral" });
},
);
it.skipIf(!process.env.ANTHROPIC_API_KEY)("should use 1h cache TTL when PI_CACHE_RETENTION=long", async () => {
process.env.PI_CACHE_RETENTION = "long";
const model = getModel("anthropic", "claude-haiku-4-5");
let capturedPayload: any = null;
const s = stream(model, context, {
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const _ of s) {
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.system).toBeDefined();
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral", ttl: "1h" });
});
it("should add ttl for non-api.anthropic.com baseUrl by default", async () => {
process.env.PI_CACHE_RETENTION = "long";
const baseModel = getModel("anthropic", "claude-haiku-4-5");
const proxyModel = {
...baseModel,
baseUrl: "https://my-proxy.example.com/v1",
};
let capturedPayload: any = null;
const { streamAnthropic } = await import("../src/providers/anthropic.js");
try {
const s = streamAnthropic(proxyModel, context, {
apiKey: "fake-key",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral", ttl: "1h" });
});
it("should omit ttl when supportsLongCacheRetention is false", async () => {
const baseModel = getModel("anthropic", "claude-haiku-4-5");
const proxyModel = {
...baseModel,
baseUrl: "https://my-proxy.example.com/v1",
compat: { supportsLongCacheRetention: false },
};
let capturedPayload: any = null;
const { streamAnthropic } = await import("../src/providers/anthropic.js");
try {
const s = streamAnthropic(proxyModel, context, {
apiKey: "fake-key",
cacheRetention: "long",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral" });
});
it("should omit cache_control when cacheRetention is none", async () => {
const baseModel = getModel("anthropic", "claude-haiku-4-5");
let capturedPayload: any = null;
const { streamAnthropic } = await import("../src/providers/anthropic.js");
try {
const s = streamAnthropic(baseModel, context, {
apiKey: "fake-key",
cacheRetention: "none",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.system[0].cache_control).toBeUndefined();
});
it("should add cache_control to string user messages", async () => {
const baseModel = getModel("anthropic", "claude-haiku-4-5");
let capturedPayload: any = null;
const { streamAnthropic } = await import("../src/providers/anthropic.js");
try {
const s = streamAnthropic(baseModel, context, {
apiKey: "fake-key",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
const lastMessage = capturedPayload.messages[capturedPayload.messages.length - 1];
expect(Array.isArray(lastMessage.content)).toBe(true);
const lastBlock = lastMessage.content[lastMessage.content.length - 1];
expect(lastBlock.cache_control).toEqual({ type: "ephemeral" });
});
it("should set 1h cache TTL when cacheRetention is long", async () => {
const baseModel = getModel("anthropic", "claude-haiku-4-5");
let capturedPayload: any = null;
const { streamAnthropic } = await import("../src/providers/anthropic.js");
try {
const s = streamAnthropic(baseModel, context, {
apiKey: "fake-key",
cacheRetention: "long",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral", ttl: "1h" });
});
});
describe("OpenAI Responses Provider", () => {
it.skipIf(!process.env.OPENAI_API_KEY)(
"should not set prompt_cache_retention when PI_CACHE_RETENTION is not set",
async () => {
const model = getModel("openai", "gpt-4o-mini");
let capturedPayload: any = null;
const s = stream(model, context, {
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const _ of s) {
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.prompt_cache_retention).toBeUndefined();
},
);
it.skipIf(!process.env.OPENAI_API_KEY)(
"should set prompt_cache_retention to 24h when PI_CACHE_RETENTION=long",
async () => {
process.env.PI_CACHE_RETENTION = "long";
const model = getModel("openai", "gpt-4o-mini");
let capturedPayload: any = null;
const s = stream(model, context, {
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const _ of s) {
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.prompt_cache_retention).toBe("24h");
},
);
it("should set prompt_cache_retention for non-api.openai.com baseUrl by default", async () => {
process.env.PI_CACHE_RETENTION = "long";
const baseModel = getModel("openai", "gpt-4o-mini");
const proxyModel = {
...baseModel,
baseUrl: "https://my-proxy.example.com/v1",
};
let capturedPayload: any = null;
const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");
try {
const s = streamOpenAIResponses(proxyModel, context, {
apiKey: "fake-key",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.prompt_cache_retention).toBe("24h");
});
it("should omit prompt_cache_retention when supportsLongCacheRetention is false", async () => {
const model = {
...getModel("openai", "gpt-4o-mini"),
compat: { supportsLongCacheRetention: false },
};
let capturedPayload: any = null;
const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");
try {
const s = streamOpenAIResponses(model, context, {
apiKey: "fake-key",
cacheRetention: "long",
sessionId: "session-compat-false",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.prompt_cache_retention).toBeUndefined();
});
it("should omit prompt_cache_key when cacheRetention is none", async () => {
const model = getModel("openai", "gpt-4o-mini");
let capturedPayload: any = null;
const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");
try {
const s = streamOpenAIResponses(model, context, {
apiKey: "fake-key",
cacheRetention: "none",
sessionId: "session-1",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.prompt_cache_key).toBeUndefined();
expect(capturedPayload.prompt_cache_retention).toBeUndefined();
});
it("should set prompt_cache_retention when cacheRetention is long", async () => {
const model = getModel("openai", "gpt-4o-mini");
let capturedPayload: any = null;
const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");
try {
const s = streamOpenAIResponses(model, context, {
apiKey: "fake-key",
cacheRetention: "long",
sessionId: "session-2",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.prompt_cache_key).toBe("session-2");
expect(capturedPayload.prompt_cache_retention).toBe("24h");
});
});
describe("OpenAI Completions Provider", () => {
function createCompletionsModel(compat?: Model<"openai-completions">["compat"]): Model<"openai-completions"> {
return {
id: "test-model",
name: "Test Model",
api: "openai-completions",
provider: "test-openai-completions",
baseUrl: "https://my-proxy.example.com/v1",
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 128000,
maxTokens: 4096,
compat,
};
}
async function capturePayload(
model: Model<"openai-completions">,
options: { cacheRetention?: "none" | "short" | "long"; sessionId?: string },
) {
let capturedPayload: any = null;
const { streamOpenAICompletions } = await import("../src/providers/openai-completions.js");
try {
const s = streamOpenAICompletions(model, context, {
apiKey: "fake-key",
...options,
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
return capturedPayload;
}
it.each([
{ name: "caching enabled", options: { sessionId: "session-1" }, key: "session-1", retention: undefined },
{
name: "cacheRetention long",
options: { cacheRetention: "long" as const, sessionId: "session-2" },
key: "session-2",
retention: "24h",
},
{
name: "cacheRetention none",
options: { cacheRetention: "none" as const, sessionId: "session-3" },
key: undefined,
retention: undefined,
},
{
name: "PI_CACHE_RETENTION=long",
options: { sessionId: "session-4" },
env: "long",
key: "session-4",
retention: "24h",
},
])("should shape direct OpenAI prompt cache fields for $name", async ({ options, env, key, retention }) => {
if (env) process.env.PI_CACHE_RETENTION = env;
const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini");
const model = { ...baseModel, api: "openai-completions" } as Model<"openai-completions">;
const payload = await capturePayload(model, options);
expect(payload.prompt_cache_key).toBe(key);
expect(payload.prompt_cache_retention).toBe(retention);
});
it("should set prompt_cache_retention for non-api.openai.com baseUrl by default", async () => {
let capturedPayload: any = null;
const { streamOpenAICompletions } = await import("../src/providers/openai-completions.js");
try {
const s = streamOpenAICompletions(createCompletionsModel(), context, {
apiKey: "fake-key",
cacheRetention: "long",
sessionId: "session-completions",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.prompt_cache_key).toBe("session-completions");
expect(capturedPayload.prompt_cache_retention).toBe("24h");
});
it("should omit prompt_cache_retention when supportsLongCacheRetention is false", async () => {
let capturedPayload: any = null;
const { streamOpenAICompletions } = await import("../src/providers/openai-completions.js");
try {
const s = streamOpenAICompletions(createCompletionsModel({ supportsLongCacheRetention: false }), context, {
apiKey: "fake-key",
cacheRetention: "long",
sessionId: "session-completions-false",
onPayload: (payload) => {
capturedPayload = payload;
},
});
for await (const event of s) {
if (event.type === "error") break;
}
} catch {
// The fake proxy request fails after the payload capture used by this assertion.
}
expect(capturedPayload).not.toBeNull();
expect(capturedPayload.prompt_cache_key).toBeUndefined();
expect(capturedPayload.prompt_cache_retention).toBeUndefined();
});
});
});
const identityContext: Context = { messages: [{ role: "user", content: "Reply OK", timestamp: 1 }] };
const identityModels = (["opencode", "opencode-go"] as const).flatMap((provider) => {
const catalog = getModels(provider);
return [...new Set(catalog.map((model) => model.api))].map((api) => catalog.find((model) => model.api === api)!);
});
let server: Server;
let fixtureBaseUrl: string;
let requests: { headers: IncomingHttpHeaders; body: Record<string, unknown> }[];
let enforceContract: boolean;
beforeAll(async () => {
server = createServer(async (req, res) => {
let raw = "";
for await (const chunk of req) raw += chunk.toString();
const body = JSON.parse(raw) as Record<string, unknown>;
requests.push({ headers: req.headers, body });
if (
enforceContract &&
(!req.headers["x-opencode-session"] || !req.headers["user-agent"]?.startsWith("prime-agent"))
) {
res.writeHead(400, { "content-type": "application/json" });
res.end(JSON.stringify({ error: { type: "invalid_request_error", message: "Missing OpenCode identity" } }));
return;
}
res.writeHead(200, { "content-type": "text/event-stream" });
const event = (type: string, data: Record<string, unknown>) =>
res.write(`event: ${type}\ndata: ${JSON.stringify({ type, ...data })}\n\n`);
if (req.url?.includes("/messages")) {
event("message_start", {
message: {
id: "fixture-message",
type: "message",
role: "assistant",
model: body.model,
content: [],
stop_reason: null,
stop_sequence: null,
usage: { input_tokens: 1, output_tokens: 0 },
},
});
event("content_block_start", { index: 0, content_block: { type: "text", text: "" } });
event("content_block_delta", { index: 0, delta: { type: "text_delta", text: "OK" } });
event("content_block_stop", { index: 0 });
event("message_delta", {
delta: { stop_reason: "end_turn", stop_sequence: null },
usage: { output_tokens: 1 },
});
event("message_stop", {});
} else if (req.url?.includes("/responses")) {
const item = {
id: "fixture-output",
type: "message",
status: "completed",
role: "assistant",
content: [{ type: "output_text", text: "OK", annotations: [] }],
};
event("response.created", { response: { id: "fixture-response", status: "in_progress", output: [] } });
event("response.output_item.added", {
output_index: 0,
item: { ...item, status: "in_progress", content: [] },
});
event("response.content_part.added", {
output_index: 0,
content_index: 0,
part: { type: "output_text", text: "", annotations: [] },
});
event("response.output_text.delta", { output_index: 0, content_index: 0, delta: "OK" });
event("response.output_item.done", { output_index: 0, item });
event("response.completed", {
response: {
id: "fixture-response",
status: "completed",
output: [item],
usage: { input_tokens: 1, output_tokens: 1 },
},
});
} else if (req.url?.includes(":streamGenerateContent")) {
res.write(
`data: ${JSON.stringify({ candidates: [{ content: { role: "model", parts: [{ text: "OK" }] }, finishReason: "STOP" }] })}\n\n`,
);
} else {
res.write(
`data: ${JSON.stringify({ id: "fixture-completion", choices: [{ index: 0, delta: { role: "assistant", content: "OK" }, finish_reason: "stop" }] })}\n\n`,
);
res.write("data: [DONE]\n\n");
}
res.end();
});
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
fixtureBaseUrl = `http://127.0.0.1:${(server.address() as AddressInfo).port}/v1`;
});
afterAll(async () => {
server.closeAllConnections();
await new Promise<void>((resolve, reject) => server.close((error) => (error ? reject(error) : resolve())));
});
beforeEach(() => {
requests = [];
enforceContract = true;
});
function expectSuccess(result: AssistantMessage) {
expect(result.stopReason, result.errorMessage).toBe("stop");
expect(
result.content
.filter((block) => block.type === "text")
.map((block) => block.text)
.join(""),
).toBe("OK");
}
describe.each(identityModels)("request identity headers: $provider/$api", (catalogModel) => {
function request(options: SimpleStreamOptions = {}, model: Model<Api> = catalogModel, simple = true) {
return (simple ? completeSimple : complete)({ ...model, baseUrl: fixtureBaseUrl }, identityContext, {
apiKey: "fixture-not-a-real-key",
sessionId: "conversation-a",
maxTokens: 512,
transport: "sse",
signal: AbortSignal.timeout(5000),
...options,
});
}
it.each([true, false])("identifies repeated and distinct conversations (simple=%s)", async (simple) => {
for (const sessionId of ["conversation-a", "conversation-a", "conversation-b"]) {
const result = await request({ sessionId }, catalogModel, simple);
const headers = requests.at(-1)?.headers;
expect({ session: headers?.["x-opencode-session"], userAgent: headers?.["user-agent"] }).toEqual({
session: sessionId,
userAgent: expect.stringMatching(/^prime-agent(?:\/|$)/),
});
expectSuccess(result);
}
});
it.each(["none", "short", "long"] as const)(
"identifies conversations with cacheRetention=%s",
async (cacheRetention) => {
const result = await request({ cacheRetention });
expect(requests.at(-1)?.headers["x-opencode-session"]).toBe("conversation-a");
expectSuccess(result);
if (cacheRetention === "none") {
expect(requests.at(-1)?.body.prompt_cache_key).toBeUndefined();
expect(requests.at(-1)?.headers.session_id).toBeUndefined();
}
},
);
it("lets model headers and then request headers override generated identity regardless of casing", async () => {
const model = {
...catalogModel,
headers: {
"User-Agent": "prime-agent-model/1",
"X-OpenCode-Session": "model-conversation",
"X-Fixture": "model",
},
};
expectSuccess(await request({}, model));
expect(requests.at(-1)?.headers["user-agent"]).toBe("prime-agent-model/1");
expect(requests.at(-1)?.headers["x-opencode-session"]).toBe("model-conversation");
expectSuccess(
await request(
{
headers: {
"user-agent": "prime-agent-request/1",
"x-opencode-session": "request-conversation",
"x-fixture": "request",
},
},
model,
),
);
expect(requests.at(-1)?.headers["user-agent"]).toBe("prime-agent-request/1");
expect(requests.at(-1)?.headers["x-opencode-session"]).toBe("request-conversation");
expect(requests.at(-1)?.headers["x-fixture"]).toBe("request");
expect(model.headers["X-Fixture"]).toBe("model");
});
it.each([undefined, ""])("does not invent a conversation when sessionId=%s", async (sessionId) => {
enforceContract = false;
expectSuccess(await request({ sessionId }));
expect(requests.at(-1)?.headers["x-opencode-session"]).toBeUndefined();
expect(requests.at(-1)?.headers["user-agent"]).toMatch(/^prime-agent(?:\/|$)/);
});
it("does not add OpenCode identity for unrelated providers", async () => {
enforceContract = false;
expectSuccess(await request({}, { ...catalogModel, provider: "fixture-provider" }));
expect(requests.at(-1)?.headers["x-opencode-session"]).toBeUndefined();
expect(requests.at(-1)?.headers["user-agent"]).not.toMatch(/^prime-agent(?:\/|$)/);
});
});
describe("session affinity headers", () => {
function affinityModel(): Model<"openai-completions"> {
return {
id: "test-model",
name: "Test Model",
api: "openai-completions",
provider: "test-openai-completions",
baseUrl: fixtureBaseUrl,
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 128000,
maxTokens: 4096,
compat: { sendSessionAffinityHeaders: true },
};
}
const affinityHeaders = ["session_id", "x-client-request-id", "x-session-affinity"] as const;
it.each([
{ name: "sends affinity headers when compat enables them", options: {}, expected: "session-affinity" },
{ name: "omits affinity headers when cacheRetention is none", options: { cacheRetention: "none" as const } },
{
name: "lets explicit headers override generated affinity headers",
options: {
headers: {
session_id: "override",
"x-client-request-id": "override",
"x-session-affinity": "override",
},
},
expected: "override",
},
])("$name", async ({ options, expected }) => {
enforceContract = false;
await completeSimple(affinityModel(), identityContext, {
apiKey: "fixture-not-a-real-key",
sessionId: "session-affinity",
transport: "sse",
signal: AbortSignal.timeout(5000),
...options,
});
const headers = requests.at(-1)?.headers ?? {};
for (const header of affinityHeaders) {
expect(headers[header]).toBe(expected);
}
});
});
describe("openCodeSessionHeaders", () => {
it("names the session for OpenCode providers only", () => {
for (const provider of ["opencode", "opencode-go"]) {
expect(openCodeSessionHeaders(provider, "session-1")).toEqual({ "x-opencode-session": "session-1" });
expect(openCodeSessionHeaders(provider, undefined)).toEqual({});
}
for (const provider of ["anthropic", "prime-inference", ""]) {
expect(openCodeSessionHeaders(provider, "session-1")).toEqual({});
}
});
});