import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { buildAPIRequest, callServerProxy, buildSystemBlocks, prepareWithCaching, trimOpenAIContext, trimGeminiContext } from "./api-client.js";
import type { APIRequestConfig } from "./types.js";
// ============================================================================
// buildAPIRequest
// ============================================================================
describe("buildAPIRequest", () => {
// ---------- main profile ----------
describe("main profile", () => {
it("returns compact + clear betas for opus-4-6", () => {
const config = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
});
expect(config.betas).toContain("context-management-2025-06-27");
expect(config.betas).toContain("compact-2026-01-12");
});
it("returns clear beta only for sonnet", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "main",
});
expect(config.betas).toContain("context-management-2025-06-27");
expect(config.betas).not.toContain("compact-2026-01-12");
});
it("includes compact edit at 120K and clear edit at 80K for opus-4-6", () => {
const config = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
});
const compactEdit = config.contextManagement.edits.find(
(e) => e.type === "compact_20260112",
);
const clearEdit = config.contextManagement.edits.find(
(e) => e.type === "clear_tool_uses_20250919",
);
expect(compactEdit).toBeDefined();
expect((compactEdit!.trigger as any).value).toBe(120_000);
expect(clearEdit).toBeDefined();
expect((clearEdit!.trigger as any).value).toBe(80_000);
expect((clearEdit!.keep as any).value).toBe(3);
});
it("includes only clear edit for non-opus models", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "main",
});
expect(config.contextManagement.edits).toHaveLength(1);
expect(config.contextManagement.edits[0].type).toBe("clear_tool_uses_20250919");
});
it("defaults maxTokens to 16384", () => {
const config = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
});
expect(config.maxTokens).toBe(16384);
});
it("respects maxOutputTokens override", () => {
const config = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
maxOutputTokens: 4096,
});
expect(config.maxTokens).toBe(4096);
});
it("caps maxOutputTokens at model maximum", () => {
const config = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
maxOutputTokens: 99999,
});
expect(config.maxTokens).toBe(16384);
});
});
// ---------- subagent profile ----------
describe("subagent profile", () => {
it("uses clear at 60K/keep 2 for Anthropic models", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "subagent",
});
expect(config.contextManagement.edits).toHaveLength(1);
const edit = config.contextManagement.edits[0];
expect(edit.type).toBe("clear_tool_uses_20250919");
expect((edit.trigger as any).value).toBe(60_000);
expect((edit.keep as any).value).toBe(2);
});
it("defaults maxTokens to 8192 (subagent cap)", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "subagent",
});
expect(config.maxTokens).toBe(8192);
});
it("allows explicit maxOutputTokens to override subagent default", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "subagent",
maxOutputTokens: 4096,
});
expect(config.maxTokens).toBe(4096);
});
it("has context-management beta for Anthropic models", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "subagent",
});
expect(config.betas).toContain("context-management-2025-06-27");
});
});
// ---------- teammate profile ----------
describe("teammate profile", () => {
it("uses same context management as main for opus-4-6", () => {
const main = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
});
const teammate = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "teammate",
});
expect(teammate.betas).toEqual(main.betas);
expect(teammate.contextManagement.edits).toEqual(main.contextManagement.edits);
});
it("defaults maxTokens to 16384 (not subagent cap)", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "teammate",
});
expect(config.maxTokens).toBe(16384);
});
});
// ---------- thinking config ----------
describe("thinking integration", () => {
it("does not include thinking when disabled", () => {
const config = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
thinkingEnabled: false,
});
expect(config.thinking).toBeUndefined();
});
it("defaults thinkingEnabled to false", () => {
const config = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
});
expect(config.thinking).toBeUndefined();
});
it("includes adaptive thinking beta for opus-4-6 when enabled", () => {
const config = buildAPIRequest({
model: "claude-opus-4-6",
contextProfile: "main",
thinkingEnabled: true,
});
expect(config.betas).toContain("adaptive-thinking-2026-01-28");
expect(config.thinking).toBeDefined();
expect(config.thinking!.type).toBe("adaptive");
expect(config.thinking!.budget_tokens).toBeUndefined();
});
it("includes interleaved thinking beta for sonnet when enabled", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "main",
thinkingEnabled: true,
});
expect(config.betas).toContain("interleaved-thinking-2025-05-14");
expect(config.thinking).toBeDefined();
expect(config.thinking!.type).toBe("enabled");
expect(config.thinking!.budget_tokens).toBeDefined();
});
it("caps thinking budget_tokens below maxTokens", () => {
const config = buildAPIRequest({
model: "claude-sonnet-4-20250514",
contextProfile: "main",
thinkingEnabled: true,
maxOutputTokens: 16384,
});
// budget_tokens should be < maxTokens
expect(config.thinking!.budget_tokens!).toBeLessThan(config.maxTokens);
});
});
// ---------- Gemini (non-Anthropic) models ----------
describe("non-Anthropic models (Gemini)", () => {
it("returns empty betas for gemini models on main profile", () => {
const config = buildAPIRequest({
model: "gemini-3-pro-preview",
contextProfile: "main",
});
expect(config.betas).toEqual([]);
});
it("returns empty edits for gemini models on main profile", () => {
const config = buildAPIRequest({
model: "gemini-3-pro-preview",
contextProfile: "main",
});
expect(config.contextManagement.edits).toEqual([]);
});
it("returns empty betas for gemini models on subagent profile", () => {
const config = buildAPIRequest({
model: "gemini-2.5-flash",
contextProfile: "subagent",
});
expect(config.betas).toEqual([]);
expect(config.contextManagement.edits).toEqual([]);
});
it("disables thinking for gemini models even when enabled", () => {
const config = buildAPIRequest({
model: "gemini-3-pro-preview",
contextProfile: "main",
thinkingEnabled: true,
});
expect(config.thinking).toBeUndefined();
});
});
});
// ============================================================================
// callServerProxy
// ============================================================================
describe("callServerProxy", () => {
let originalFetch: typeof globalThis.fetch;
beforeEach(() => {
originalFetch = globalThis.fetch;
});
afterEach(() => {
globalThis.fetch = originalFetch;
});
function mockFetch(impl: (...args: any[]) => Promise<Response>) {
globalThis.fetch = vi.fn(impl) as any;
}
function makeMockStream(): ReadableStream<Uint8Array> {
return new ReadableStream({
start(controller) {
controller.enqueue(new TextEncoder().encode("data: hello\n\n"));
controller.close();
},
});
}
const baseConfig = {
proxyUrl: "https://whale-agent.fly.dev/proxy",
token: "test-jwt-token",
model: "claude-opus-4-6",
system: [{ type: "text", text: "You are helpful." }],
messages: [{ role: "user", content: "Hello" }],
tools: [{ name: "bash", description: "Run commands" }],
apiConfig: {
betas: ["context-management-2025-06-27"],
contextManagement: {
edits: [
{
type: "clear_tool_uses_20250919",
trigger: { type: "input_tokens", value: 80_000 },
keep: { type: "tool_uses", value: 3 },
},
],
},
maxTokens: 16384,
} as APIRequestConfig,
};
// ---------- Happy path ----------
it("returns a ReadableStream on successful 200 response", async () => {
const stream = makeMockStream();
mockFetch(async () => new Response(stream, { status: 200 }));
const result = await callServerProxy({ ...baseConfig });
expect(result).toBeInstanceOf(ReadableStream);
});
it("sends correct headers (Content-Type and Authorization)", async () => {
const stream = makeMockStream();
mockFetch(async (_url: string, opts: any) => {
expect(opts.headers["Content-Type"]).toBe("application/json");
expect(opts.headers["Authorization"]).toBe("Bearer test-jwt-token");
return new Response(stream, { status: 200 });
});
await callServerProxy({ ...baseConfig });
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
});
it("sends correct body fields to the proxy", async () => {
const stream = makeMockStream();
let capturedBody: any;
mockFetch(async (_url: string, opts: any) => {
capturedBody = JSON.parse(opts.body);
return new Response(stream, { status: 200 });
});
await callServerProxy({ ...baseConfig });
expect(capturedBody.mode).toBe("proxy");
expect(capturedBody.model).toBe("claude-opus-4-6");
expect(capturedBody.stream).toBe(true);
expect(capturedBody.max_tokens).toBe(16384);
expect(capturedBody.messages).toEqual(baseConfig.messages);
expect(capturedBody.system).toEqual(baseConfig.system);
expect(capturedBody.tools).toEqual(baseConfig.tools);
expect(capturedBody.betas).toEqual(["context-management-2025-06-27"]);
expect(capturedBody.context_management).toBeDefined();
expect(capturedBody.context_management.edits).toHaveLength(1);
});
it("includes thinking in body when apiConfig has thinking", async () => {
const stream = makeMockStream();
let capturedBody: any;
mockFetch(async (_url: string, opts: any) => {
capturedBody = JSON.parse(opts.body);
return new Response(stream, { status: 200 });
});
await callServerProxy({
...baseConfig,
apiConfig: {
...baseConfig.apiConfig,
thinking: { type: "adaptive" },
},
});
expect(capturedBody.thinking).toEqual({ type: "adaptive" });
});
it("omits thinking from body when apiConfig has no thinking", async () => {
const stream = makeMockStream();
let capturedBody: any;
mockFetch(async (_url: string, opts: any) => {
capturedBody = JSON.parse(opts.body);
return new Response(stream, { status: 200 });
});
await callServerProxy({ ...baseConfig });
expect(capturedBody.thinking).toBeUndefined();
});
it("omits context_management when edits array is empty", async () => {
const stream = makeMockStream();
let capturedBody: any;
mockFetch(async (_url: string, opts: any) => {
capturedBody = JSON.parse(opts.body);
return new Response(stream, { status: 200 });
});
await callServerProxy({
...baseConfig,
apiConfig: {
...baseConfig.apiConfig,
contextManagement: { edits: [] },
},
});
expect(capturedBody.context_management).toBeUndefined();
});
it("includes store_id in body when provided", async () => {
const stream = makeMockStream();
let capturedBody: any;
mockFetch(async (_url: string, opts: any) => {
capturedBody = JSON.parse(opts.body);
return new Response(stream, { status: 200 });
});
await callServerProxy({
...baseConfig,
storeId: "store-uuid-123",
});
expect(capturedBody.store_id).toBe("store-uuid-123");
});
it("omits store_id from body when not provided", async () => {
const stream = makeMockStream();
let capturedBody: any;
mockFetch(async (_url: string, opts: any) => {
capturedBody = JSON.parse(opts.body);
return new Response(stream, { status: 200 });
});
await callServerProxy({ ...baseConfig });
expect(capturedBody.store_id).toBeUndefined();
});
// ---------- Error handling ----------
it("throws on non-OK response with status and error body", async () => {
mockFetch(
async () => new Response("Bad Request: invalid model", { status: 400 }),
);
await expect(callServerProxy({ ...baseConfig })).rejects.toThrow(
"Proxy error (400): Bad Request: invalid model",
);
});
it("throws on 401 without retrying", async () => {
mockFetch(
async () => new Response("Unauthorized", { status: 401 }),
);
await expect(callServerProxy({ ...baseConfig })).rejects.toThrow(
"Proxy error (401)",
);
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
});
it("throws on null response body even with 200 status", async () => {
mockFetch(async () => {
// Response with ok=true but null body
return new Response(null, { status: 204 });
});
// 204 is not ok in the range check — actually Response(null, {status: 204}) has ok=true
// but body is null. Let's check:
// Actually Response with status 204 has ok=true. The code checks `response.ok && response.body`.
// A Response(null) has body === null, so this will fall to the error path.
await expect(callServerProxy({ ...baseConfig })).rejects.toThrow(
"Proxy error (204)",
);
});
// ---------- Retry behavior ----------
it("retries on 429 status with exponential backoff", async () => {
let callCount = 0;
mockFetch(async () => {
callCount++;
if (callCount <= 2) {
return new Response("Rate limited", { status: 429 });
}
return new Response(makeMockStream(), { status: 200 });
});
const result = await callServerProxy({ ...baseConfig });
expect(result).toBeInstanceOf(ReadableStream);
expect(callCount).toBe(3);
});
it("retries on 500 status", async () => {
let callCount = 0;
mockFetch(async () => {
callCount++;
if (callCount === 1) {
return new Response("Internal Server Error", { status: 500 });
}
return new Response(makeMockStream(), { status: 200 });
});
const result = await callServerProxy({ ...baseConfig });
expect(result).toBeInstanceOf(ReadableStream);
expect(callCount).toBe(2);
});
it("retries on 529 status (overloaded)", async () => {
let callCount = 0;
mockFetch(async () => {
callCount++;
if (callCount === 1) {
return new Response("Overloaded", { status: 529 });
}
return new Response(makeMockStream(), { status: 200 });
});
const result = await callServerProxy({ ...baseConfig });
expect(result).toBeInstanceOf(ReadableStream);
expect(callCount).toBe(2);
});
it("throws after MAX_RETRIES (3) retries exhausted", async () => {
mockFetch(
async () => new Response("Rate limited", { status: 429 }),
);
await expect(callServerProxy({ ...baseConfig })).rejects.toThrow(
"Proxy error (429)",
);
// 1 initial + 3 retries = 4 total attempts
expect(globalThis.fetch).toHaveBeenCalledTimes(4);
}, 15_000);
it("calls onRetry callback on each retry", async () => {
let callCount = 0;
mockFetch(async () => {
callCount++;
if (callCount <= 2) {
return new Response("Rate limited", { status: 429 });
}
return new Response(makeMockStream(), { status: 200 });
});
const onRetry = vi.fn();
await callServerProxy({ ...baseConfig, onRetry });
expect(onRetry).toHaveBeenCalledTimes(2);
expect(onRetry).toHaveBeenNthCalledWith(1, 1, 3, expect.stringContaining("429"));
expect(onRetry).toHaveBeenNthCalledWith(2, 2, 3, expect.stringContaining("429"));
});
// ---------- Fallback model ----------
it("switches to fallback model on last retry attempt", async () => {
let callCount = 0;
let lastModel: string | undefined;
mockFetch(async (_url: string, opts: any) => {
callCount++;
lastModel = JSON.parse(opts.body).model;
if (callCount <= 3) {
return new Response("Rate limited", { status: 429 });
}
return new Response(makeMockStream(), { status: 200 });
});
const onFallback = vi.fn();
await callServerProxy({
...baseConfig,
fallbackModel: "claude-sonnet-4-20250514",
onFallback,
});
expect(onFallback).toHaveBeenCalledWith(
"claude-opus-4-6",
"claude-sonnet-4-20250514",
);
expect(lastModel).toBe("claude-sonnet-4-20250514");
expect(callCount).toBe(4);
}, 15_000);
it("does not call onFallback when no fallbackModel is set", async () => {
let callCount = 0;
mockFetch(async () => {
callCount++;
if (callCount <= 2) {
return new Response("Rate limited", { status: 429 });
}
return new Response(makeMockStream(), { status: 200 });
});
const onFallback = vi.fn();
await callServerProxy({ ...baseConfig, onFallback });
expect(onFallback).not.toHaveBeenCalled();
});
// ---------- Abort signal ----------
it("throws immediately when signal is already aborted", async () => {
const controller = new AbortController();
controller.abort();
mockFetch(async () => {
throw new DOMException("Aborted", "AbortError");
});
await expect(
callServerProxy({ ...baseConfig, signal: controller.signal }),
).rejects.toThrow();
// Should not retry on abort
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
});
it("does not retry when signal is aborted", async () => {
const controller = new AbortController();
let callCount = 0;
mockFetch(async () => {
callCount++;
controller.abort();
throw new DOMException("Aborted", "AbortError");
});
await expect(
callServerProxy({ ...baseConfig, signal: controller.signal }),
).rejects.toThrow();
expect(callCount).toBe(1);
});
// ---------- Timeout ----------
it("applies timeout when timeoutMs is set and no signal provided", async () => {
let receivedSignal: AbortSignal | undefined;
mockFetch(async (_url: string, opts: any) => {
receivedSignal = opts.signal;
return new Response(makeMockStream(), { status: 200 });
});
await callServerProxy({
...baseConfig,
timeoutMs: 5000,
});
// The fetch should have received an AbortSignal from the timeout controller
expect(receivedSignal).toBeDefined();
expect(receivedSignal).toBeInstanceOf(AbortSignal);
});
it("does not create timeout controller when signal is already provided", async () => {
const controller = new AbortController();
let receivedSignal: AbortSignal | undefined;
mockFetch(async (_url: string, opts: any) => {
receivedSignal = opts.signal;
return new Response(makeMockStream(), { status: 200 });
});
await callServerProxy({
...baseConfig,
signal: controller.signal,
timeoutMs: 5000,
});
// Should use the user-provided signal, not a timeout-created one
expect(receivedSignal).toBe(controller.signal);
});
// ---------- Non-retryable errors ----------
it("does not retry on 400 Bad Request", async () => {
mockFetch(
async () => new Response("Bad Request", { status: 400 }),
);
await expect(callServerProxy({ ...baseConfig })).rejects.toThrow(
"Proxy error (400)",
);
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
});
it("does not retry on 403 Forbidden", async () => {
mockFetch(
async () => new Response("Forbidden", { status: 403 }),
);
await expect(callServerProxy({ ...baseConfig })).rejects.toThrow(
"Proxy error (403)",
);
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
});
it("does not retry on 404 Not Found", async () => {
mockFetch(
async () => new Response("Not Found", { status: 404 }),
);
await expect(callServerProxy({ ...baseConfig })).rejects.toThrow(
"Proxy error (404)",
);
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
});
// ---------- Network errors (retryable by message) ----------
it("retries on network error with 'timeout' in message", async () => {
let callCount = 0;
mockFetch(async () => {
callCount++;
if (callCount === 1) {
throw Object.assign(new Error("Request timeout"), { status: undefined });
}
return new Response(makeMockStream(), { status: 200 });
});
const result = await callServerProxy({ ...baseConfig });
expect(result).toBeInstanceOf(ReadableStream);
expect(callCount).toBe(2);
});
it("retries on error with 'overloaded' in message", async () => {
let callCount = 0;
mockFetch(async () => {
callCount++;
if (callCount === 1) {
throw new Error("API is overloaded");
}
return new Response(makeMockStream(), { status: 200 });
});
const result = await callServerProxy({ ...baseConfig });
expect(result).toBeInstanceOf(ReadableStream);
expect(callCount).toBe(2);
});
});
// ============================================================================
// buildSystemBlocks
// ============================================================================
describe("buildSystemBlocks", () => {
it("returns system prompt with cache_control when caching enabled", () => {
const blocks = buildSystemBlocks("You are helpful.", undefined, true);
expect(blocks).toHaveLength(1);
expect(blocks[0]).toEqual({
type: "text",
text: "You are helpful.",
cache_control: { type: "ephemeral" },
});
});
it("returns system prompt without cache_control when caching disabled", () => {
const blocks = buildSystemBlocks("You are helpful.", undefined, false);
expect(blocks).toHaveLength(1);
expect(blocks[0]).toEqual({
type: "text",
text: "You are helpful.",
});
});
it("defaults to caching enabled", () => {
const blocks = buildSystemBlocks("Hello.");
expect(blocks[0]).toHaveProperty("cache_control", { type: "ephemeral" });
});
it("appends cost context as a second block", () => {
const blocks = buildSystemBlocks("System prompt.", "Session cost: $0.05");
expect(blocks).toHaveLength(2);
expect(blocks[0].text).toBe("System prompt.");
expect(blocks[1]).toEqual({
type: "text",
text: "Session cost: $0.05",
});
});
it("cost context block does not have cache_control (dynamic content)", () => {
const blocks = buildSystemBlocks("System.", "Cost: $1.00", true);
expect(blocks[1]).not.toHaveProperty("cache_control");
});
it("omits cost context block when undefined", () => {
const blocks = buildSystemBlocks("System.", undefined, true);
expect(blocks).toHaveLength(1);
});
it("omits cost context block when empty string (falsy)", () => {
const blocks = buildSystemBlocks("System.", "", true);
expect(blocks).toHaveLength(1);
});
});
// ============================================================================
// prepareWithCaching
// ============================================================================
describe("prepareWithCaching", () => {
it("delegates to addPromptCaching and returns cached tools and messages", () => {
const tools = [{ name: "tool1" }, { name: "tool2" }];
const messages = [
{ role: "user", content: "first" },
{ role: "assistant", content: "second" },
{ role: "user", content: "third" },
];
const result = prepareWithCaching(tools, messages);
// Last tool should have cache_control
expect(result.tools[result.tools.length - 1]).toHaveProperty(
"cache_control",
{ type: "ephemeral" },
);
// First tool should not
expect(result.tools[0]).not.toHaveProperty("cache_control");
});
it("does not mutate original arrays", () => {
const tools = [{ name: "t1" }];
const messages = [
{ role: "user", content: "a" },
{ role: "assistant", content: "b" },
];
prepareWithCaching(tools, messages);
expect(tools[0]).not.toHaveProperty("cache_control");
expect(messages[0]).toEqual({ role: "user", content: "a" });
});
it("handles empty tools and messages", () => {
const result = prepareWithCaching([], []);
expect(result.tools).toEqual([]);
expect(result.messages).toEqual([]);
});
});
// ============================================================================
// trimOpenAIContext
// ============================================================================
describe("trimOpenAIContext", () => {
function makeMessages(toolResultCount: number): Array<Record<string, unknown>> {
const msgs: Array<Record<string, unknown>> = [];
for (let i = 0; i < toolResultCount; i++) {
msgs.push({
role: "assistant",
content: [{ type: "tool_use", id: `t${i}`, name: "bash", input: {} }],
});
msgs.push({
role: "user",
content: [{ type: "tool_result", tool_use_id: `t${i}`, content: `Result ${i}` }],
});
}
return msgs;
}
it("returns same reference when below threshold", () => {
const msgs = makeMessages(10);
const result = trimOpenAIContext(msgs, 100_000); // 100K < 150K default threshold
expect(result).toBe(msgs); // exact same reference
});
it("returns same reference when above threshold but few tool results", () => {
const msgs = makeMessages(3); // 3 tool results <= keepRecent=5
const result = trimOpenAIContext(msgs, 200_000);
expect(result).toBe(msgs);
});
it("trims oldest tool results when above threshold", () => {
const msgs = makeMessages(10);
const result = trimOpenAIContext(msgs, 200_000, 150_000, 5);
expect(result).not.toBe(msgs); // new array
// Should have 5 trimmed (oldest) and 5 kept
let trimmedCount = 0;
let keptCount = 0;
for (const msg of result) {
if (Array.isArray(msg.content)) {
for (const block of msg.content as Array<Record<string, unknown>>) {
if (block.type === "tool_result") {
if (block.content === "[trimmed]") trimmedCount++;
else keptCount++;
}
}
}
}
expect(trimmedCount).toBe(5);
expect(keptCount).toBe(5);
});
it("preserves most recent tool results (not oldest)", () => {
const msgs = makeMessages(8);
const result = trimOpenAIContext(msgs, 200_000, 150_000, 5);
// Last 5 tool results should be preserved
const toolResults: Array<Record<string, unknown>> = [];
for (const msg of result) {
if (Array.isArray(msg.content)) {
for (const block of msg.content as Array<Record<string, unknown>>) {
if (block.type === "tool_result") {
toolResults.push(block);
}
}
}
}
// First 3 should be trimmed, last 5 should be kept
expect(toolResults[0].content).toBe("[trimmed]");
expect(toolResults[1].content).toBe("[trimmed]");
expect(toolResults[2].content).toBe("[trimmed]");
expect(toolResults[3].content).toBe("Result 3");
expect(toolResults[7].content).toBe("Result 7");
});
it("respects custom threshold and keepRecent", () => {
const msgs = makeMessages(6);
const result = trimOpenAIContext(msgs, 100_000, 50_000, 2); // lower threshold, keep 2
let trimmedCount = 0;
for (const msg of result) {
if (Array.isArray(msg.content)) {
for (const block of msg.content as Array<Record<string, unknown>>) {
if (block.type === "tool_result" && block.content === "[trimmed]") trimmedCount++;
}
}
}
expect(trimmedCount).toBe(4); // 6 - 2 = 4 trimmed
});
it("handles messages with non-array content", () => {
const msgs = [
{ role: "user", content: "Hello" },
...makeMessages(6),
];
// Should not crash on string content
const result = trimOpenAIContext(msgs, 200_000, 150_000, 3);
expect(result[0].content).toBe("Hello");
});
});
// ============================================================================
// trimGeminiContext
// ============================================================================
describe("trimGeminiContext", () => {
function makeMessages(toolResultCount: number): Array<Record<string, unknown>> {
const msgs: Array<Record<string, unknown>> = [];
for (let i = 0; i < toolResultCount; i++) {
msgs.push({
role: "assistant",
content: [{ type: "tool_use", id: `t${i}`, name: "bash", input: {} }],
});
msgs.push({
role: "user",
content: [{ type: "tool_result", tool_use_id: `t${i}`, content: `Result ${i}` }],
});
}
return msgs;
}
it("returns same reference when below 800K threshold", () => {
const msgs = makeMessages(10);
const result = trimGeminiContext(msgs, 700_000);
expect(result).toBe(msgs);
});
it("trims when above 800K threshold", () => {
const msgs = makeMessages(10);
const result = trimGeminiContext(msgs, 900_000);
expect(result).not.toBe(msgs);
let trimmedCount = 0;
for (const msg of result) {
if (Array.isArray(msg.content)) {
for (const block of msg.content as Array<Record<string, unknown>>) {
if (block.type === "tool_result" && block.content === "[trimmed]") trimmedCount++;
}
}
}
expect(trimmedCount).toBe(5); // 10 - 5 (keepRecent default) = 5
});
});
