feat(coding-agent): add prompt cache warming (#9668)

This commit is contained in:
Armin Ronacher 2026-09-19 21:18:15 +02:00 • committed by GitHub
parent 7d5eb0ee3b
commit c596d09d9c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
32 changed files with 1380 additions and 69 deletions

View file

@ -208,6 +208,39 @@ describe("AgentSession.getSessionStats", () => {
}
});
it("includes cache-warming usage exactly once without adding messages", async () => {
const { session, sessionManager } = await createSession();
try {
sessionManager.appendUsage(
"cache_warm",
"anthropic",
model.id,
{
input: 2,
output: 1,
cacheRead: 97,
cacheWrite: 0,
totalTokens: 100,
cost: { input: 0.001, output: 0.002, cacheRead: 0.007, cacheWrite: 0, total: 0.01 },
},
"extension override",
);
const [entry] = sessionManager.getEntries();
expect(entry).toMatchObject({ type: "usage", kind: "cache_warm", note: "extension override" });
const stats = session.getSessionStats();
expect(stats.tokens).toEqual({ input: 2, output: 1, cacheRead: 97, cacheWrite: 0, total: 100 });
expect(stats.totalMessages).toBe(0);
expect(sessionManager.buildSessionContext().messages).toEqual([]);
expect(getUsageCostBreakdown(sessionManager.getEntries())).toEqual([
{ key: `anthropic/${model.id}`, cost: 0.01, tokens: 100 },
]);
} finally {
session.dispose();
}
});
it("includes tool result usage in session totals", async () => {
const { session, sessionManager } = await createSession();

View file

@ -46,6 +46,26 @@ function entry(message: AssistantMessage): SessionEntry {
return { type: "message", id: "x", parentId: null, timestamp: "", message } as SessionEntry;
}
function usageEntry(kind: string, timestamp: number): SessionEntry {
return {
type: "usage",
id: `usage-${kind}`,
parentId: null,
timestamp: new Date(timestamp).toISOString(),
kind,
provider: "test",
model: "test-model",
usage: {
input: 0,
output: 0,
cacheRead: 100_000,
cacheWrite: 0,
totalTokens: 100_000,
cost: zeroCost,
},
};
}
// Turn 1: fresh 100k cache write at $3.75/M
const turn1 = assistant({ cacheWrite: 100_000, cost: { cacheWrite: 0.375 }, timestamp: 0 });
// Turn 2: healthy, everything read back at $0.30/M
@ -127,6 +147,20 @@ describe("detectCacheMiss", () => {
expect(miss?.modelChanged).toBe(true);
});
it("uses only cache-warm usage entries as cache refreshes", () => {
const missMessage = assistant({ cacheWrite: 110_000, cost: { cacheWrite: 0.4125 }, timestamp: 600_000 });
const afterCacheWarm = detectCacheMiss([entry(turn1), usageEntry("cache_warm", 500_000)], missMessage, models);
const afterOtherUsage = detectCacheMiss(
[entry(turn1), usageEntry("custom_operation", 500_000)],
missMessage,
models,
);
expect(afterCacheWarm?.idleMs).toBe(100_000);
expect(afterOtherUsage?.idleMs).toBe(600_000);
});
it("returns undefined for healthy turns", () => {
const healthy = assistant({
cacheRead: 105_000,

View file

@ -0,0 +1,328 @@
import {
type Api,
type AssistantMessage,
type AssistantMessageEventStream,
type Model,
type ModelsSimpleStreamOptions,
normalizeContext,
type Usage,
} from "@earendil-works/pi-ai";
import { getBuiltinModel } from "@earendil-works/pi-ai/providers/all";
import { afterEach, describe, expect, it, vi } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.ts";
import {
CacheWarmer,
type CacheWarmingAction,
type CacheWarmingDecision,
type CacheWarmingDecisionEvent,
type CacheWarmRequest,
formatCacheWarmingStatus,
formatCacheWarmingUsage,
getCacheWarmingDelayMs,
getPromptCacheTtlMs,
isReplayable,
} from "../src/core/cache-warmer.ts";
import { createEventBus } from "../src/core/event-bus.ts";
import { createExtensionRuntime, loadExtensionFromFactory } from "../src/core/extensions/loader.ts";
import { ExtensionRunner } from "../src/core/extensions/runner.ts";
import type { ExtensionFactory } from "../src/core/extensions/types.ts";
import { type SessionEntry, SessionManager, type UsageEntry } from "../src/core/session-manager.ts";
import type { CacheWarmingMode } from "../src/core/settings-manager.ts";
import { createInMemoryModelRegistry } from "./model-runtime-test-utils.ts";
const adaptiveModel: Model<Api> = {
...getBuiltinModel("anthropic", "claude-opus-4-6"),
promptCache: { short: 300, long: 3600 },
};
const budgetModel: Model<Api> = {
...getBuiltinModel("anthropic", "claude-sonnet-4-5"),
promptCache: { short: 300, long: 3600 },
};
const openaiModel: Model<Api> = {
...getBuiltinModel("openai", "gpt-5"),
promptCache: { short: 300, long: 86_400 },
};
const unknownModel: Model<Api> = { ...adaptiveModel, promptCache: undefined };
const warmUsage: Usage = {
input: 0,
output: 1,
cacheRead: 100,
cacheWrite: 0,
totalTokens: 101,
cost: { input: 0, output: 0, cacheRead: 0.01, cacheWrite: 0, total: 0.01 },
};
function response(model: Model<Api>, stopReason: AssistantMessage["stopReason"] = "length"): AssistantMessage {
return {
role: "assistant",
content: [],
api: model.api,
provider: model.provider,
model: model.id,
usage: warmUsage,
stopReason,
timestamp: 0,
};
}
function branchWithPrompt(promptTokens: number): SessionEntry[] {
return [
{
type: "message",
id: "a",
parentId: null,
timestamp: new Date(0).toISOString(),
message: {
...response(adaptiveModel),
usage: { ...warmUsage, output: 10, cacheRead: promptTokens, totalTokens: promptTokens + 10 },
},
},
];
}
function fakeRuntime(
options: {
result?: (model: Model<Api>) => Promise<AssistantMessage>;
decide?: (event: CacheWarmingDecisionEvent) => CacheWarmingAction;
mode?: CacheWarmingMode;
branch?: SessionEntry[];
} = {},
) {
const calls: Array<{ model: Model<Api>; options: ModelsSimpleStreamOptions | undefined }> = [];
const events: CacheWarmingDecisionEvent[] = [];
const warmedEntries: UsageEntry[] = [];
const usageManager = SessionManager.inMemory();
const appendUsage = vi.fn(usageManager.appendUsage.bind(usageManager));
const state = { mode: options.mode ?? "idle", branch: options.branch ?? branchWithPrompt(100_000) };
const warmer = new CacheWarmer(
{
streamSimple: (model, _context, streamOptions) => {
calls.push({ model, options: streamOptions });
return {
result: () => (options.result ?? (async (m: Model<Api>) => response(m)))(model),
} as unknown as AssistantMessageEventStream;
},
},
{ appendUsage, getBranch: () => state.branch },
() => state.mode,
async (event) => {
events.push(event);
return options.decide?.(event) ?? event.action;
},
);
warmer.onWarmed = (entry) => warmedEntries.push(entry);
return { warmer, calls, events, warmedEntries, appendUsage, state };
}
function request(model: Model<Api> = adaptiveModel, options: ModelsSimpleStreamOptions = {}): CacheWarmRequest {
return { model, context: normalizeContext({ messages: [] }), options };
}
const current = () => true;
afterEach(() => vi.useRealTimers());
describe("cache warming", () => {
it("derives eligibility and timing from retention and provider behavior", () => {
expect([
getPromptCacheTtlMs(adaptiveModel, undefined),
getPromptCacheTtlMs(adaptiveModel, { cacheRetention: "long" }),
getPromptCacheTtlMs(adaptiveModel, { cacheRetention: "none" }),
getPromptCacheTtlMs(adaptiveModel, { env: { PI_CACHE_RETENTION: "long" } }),
getPromptCacheTtlMs(openaiModel, { cacheRetention: "long" }),
getPromptCacheTtlMs(unknownModel, undefined),
]).toEqual([300_000, 3_600_000, undefined, 3_600_000, 86_400_000, undefined]);
expect([getCacheWarmingDelayMs(300_000), getCacheWarmingDelayMs(60_000), getCacheWarmingDelayMs(10_000)]).toEqual(
[270_000, 50_000, undefined],
);
expect([
isReplayable(budgetModel, { reasoning: "medium" }),
isReplayable(budgetModel, undefined),
isReplayable(adaptiveModel, { reasoning: "medium" }),
isReplayable(openaiModel, { reasoning: "medium" }),
]).toEqual([false, true, true, true]);
});
it("replays profitable requests and preserves options across repeated refreshes", async () => {
vi.useFakeTimers();
const { warmer, calls, events, appendUsage, warmedEntries } = fakeRuntime();
const signal = new AbortController().signal;
const transformHeaders = async () => ({});
warmer.start(request(adaptiveModel, { reasoning: "high", signal, sessionId: "s", transformHeaders }), current);
await vi.advanceTimersByTimeAsync(270_000);
expect(calls[0]).toMatchObject({
model: adaptiveModel,
options: { reasoning: "high", sessionId: "s", transformHeaders, maxTokens: 1, maxRetries: 0 },
});
expect(calls[0].options?.signal).not.toBe(signal);
expect(events[0]).toMatchObject({
type: "cache_warming_decision",
continuationProbability: 1,
action: "warm",
});
expect(events[0].missCost).toBeCloseTo(0.575);
expect(events[0].warmCost).toBeCloseTo(0.050025);
expect(appendUsage).toHaveBeenCalledWith(
"cache_warm",
adaptiveModel.provider,
adaptiveModel.id,
warmUsage,
undefined,
);
expect(warmedEntries).toEqual([appendUsage.mock.results[0]?.value]);
await vi.advanceTimersByTimeAsync(270_000);
expect(calls).toHaveLength(2);
warmer.cancel();
});
it("applies economic decisions and extension overrides", async () => {
vi.useFakeTimers();
const unprofitable = fakeRuntime({ branch: branchWithPrompt(5_000) });
unprofitable.warmer.start(request(), current);
await vi.advanceTimersByTimeAsync(270_000);
expect(unprofitable.calls).toHaveLength(0);
expect(unprofitable.warmer.status).toMatchObject({
state: "inactive",
decision: { action: "stop", economicsAvailable: true },
extensionOverride: false,
});
const forced = fakeRuntime({ branch: branchWithPrompt(5_000), decide: () => "warm" });
forced.warmer.start(request(), current);
await vi.advanceTimersByTimeAsync(270_000);
expect(forced.calls).toHaveLength(1);
expect(forced.warmedEntries[0]?.note).toBe("extension override");
forced.warmer.cancel();
const vetoed = fakeRuntime({ decide: () => "stop" });
vetoed.warmer.start(request(), current);
await vi.advanceTimersByTimeAsync(270_000);
expect(vetoed.calls).toHaveLength(0);
expect(vetoed.warmer.status).toMatchObject({ state: "inactive", extensionOverride: true });
const unavailable = fakeRuntime({ branch: branchWithPrompt(0) });
unavailable.warmer.start(request(), current);
expect(unavailable.warmer.status).toMatchObject({
state: "inactive",
reason: "cache economics unavailable",
});
await vi.advanceTimersByTimeAsync(270_000);
expect(unavailable.calls).toHaveLength(0);
});
it("stops for unsupported requests, context changes, and mode changes", async () => {
vi.useFakeTimers();
const unsupported = fakeRuntime();
unsupported.state.mode = "off";
unsupported.warmer.start(request(), current);
expect(unsupported.warmer.status.reason).toBe("cache warming disabled");
unsupported.state.mode = "idle";
unsupported.warmer.start(request(unknownModel), current);
expect(unsupported.warmer.status.reason).toBe("cache lifetime unavailable");
unsupported.warmer.start(request(budgetModel, { reasoning: "high" }), current);
expect(unsupported.warmer.status.reason).toBe("request cannot be replayed safely");
let stillCurrent = true;
unsupported.warmer.start(request(), () => stillCurrent);
stillCurrent = false;
expect(unsupported.warmer.status.reason).toBe("conversation context changed");
await vi.advanceTimersByTimeAsync(270_000);
expect(unsupported.calls).toHaveLength(0);
unsupported.warmer.start(request(), current);
unsupported.state.mode = "off";
await vi.advanceTimersByTimeAsync(270_000);
expect(unsupported.calls).toHaveLength(0);
const streaming = fakeRuntime({ mode: "streaming", branch: branchWithPrompt(400_000) });
streaming.warmer.start(request(), current);
streaming.warmer.onAgentSettled();
expect(streaming.warmer.status.reason).toBe("agent run settled");
});
it("aborts replaced requests and does not record failed refreshes", async () => {
vi.useFakeTimers();
let release!: () => void;
const pending = fakeRuntime({
result: (model) =>
new Promise((resolve) => {
release = () => resolve(response(model));
}),
});
pending.warmer.start(request(), current);
await vi.advanceTimersByTimeAsync(270_000);
pending.warmer.start(request(), current);
expect(pending.calls[0].options?.signal?.aborted).toBe(true);
release();
pending.warmer.cancel();
await vi.advanceTimersByTimeAsync(600_000);
expect(pending.calls).toHaveLength(1);
const failed = fakeRuntime({ result: async (model) => response(model, "error") });
failed.warmer.start(request(), current);
await vi.advanceTimersByTimeAsync(270_000);
expect(failed.appendUsage).not.toHaveBeenCalled();
failed.warmer.cancel();
});
it("formats status and usage entries", () => {
const decision: CacheWarmingDecision = {
phase: "idle",
warmCost: 0.013,
missCost: 0.621,
continuationProbability: 0.6,
expectedSavings: 0.36,
economicsAvailable: true,
action: "warm",
};
expect(formatCacheWarmingStatus({ state: "scheduled", nextWarmAt: 222_000, decision }, 0)).toBe(
"Decision in 3m 42s (60% continuation probability, expected savings $0.360 >= $0.050 -> warm)",
);
const usage = {
...warmUsage,
cost: { input: 0.00004, output: 0.00005, cacheRead: 0.02940725, cacheWrite: 0, total: 0.02949725 },
};
const entry = SessionManager.inMemory().appendUsage(
"cache_warm",
adaptiveModel.provider,
adaptiveModel.id,
usage,
"extension override",
);
expect(formatCacheWarmingUsage(entry)).toBe("Cache warmed (extension override): $0.029497");
});
});
describe("ExtensionRunner.emitCacheWarmingDecision", () => {
it("uses the last extension override", async () => {
const runtime = createExtensionRuntime();
const eventBus = createEventBus();
const factories: ExtensionFactory[] = [
(pi) => {
pi.on("cache_warming_decision", () => ({ action: "warm" }));
},
(pi) => {
pi.on("cache_warming_decision", () => ({ action: "stop" }));
},
];
const extensions = [];
for (const factory of factories) {
extensions.push(await loadExtensionFromFactory(factory, process.cwd(), eventBus, runtime));
}
const modelRegistry = await createInMemoryModelRegistry(AuthStorage.inMemory());
const runner = new ExtensionRunner(extensions, runtime, process.cwd(), SessionManager.inMemory(), modelRegistry);
const event: CacheWarmingDecisionEvent = {
type: "cache_warming_decision",
warmCost: 0.05,
missCost: 0.5,
continuationProbability: 0.15,
action: "warm",
};
expect(await runner.emitCacheWarmingDecision(event)).toBe("stop");
});
});

View file

@ -787,6 +787,34 @@ describe("ModelRegistry", () => {
expect(opus?.samplingParams).toBeUndefined();
});
test("custom model and model override carry prompt cache lifetimes", async () => {
writeRawModelsJson({
openrouter: {
baseUrl: "https://my-proxy.example.com/v1",
api: "openai-completions",
models: [{ id: "custom/cached-model", promptCache: { short: 120 } }],
modelOverrides: {
"anthropic/claude-sonnet-4": { promptCache: { short: 300 } },
},
},
anthropic: {
modelOverrides: {
"claude-sonnet-4-6": { promptCache: { long: 1800 } },
},
},
});
const registry = await createModelRegistry(authStorage, modelsJsonPath);
const openrouter = getModelsForProvider(registry, "openrouter");
expect(registry.getError()).toBeUndefined();
expect(openrouter.find((m) => m.id === "custom/cached-model")?.promptCache).toEqual({ short: 120 });
expect(openrouter.find((m) => m.id === "anthropic/claude-sonnet-4")?.promptCache).toEqual({ short: 300 });
expect(openrouter.find((m) => m.id === "anthropic/claude-opus-4")?.promptCache).toBeUndefined();
// Overrides merge per tier with the built-in catalog.
expect(registry.find("anthropic", "claude-sonnet-4-6")?.promptCache).toEqual({ short: 300, long: 1800 });
});
test("model override with compat.openRouterRouting", async () => {
writeRawModelsJson({
openrouter: {

View file

@ -51,9 +51,8 @@ describe("createAgentSession stream options", () => {
};
}
function createDoneStream(api: Api) {
const stream = createAssistantMessageEventStream();
const message: AssistantMessage = {
function createDoneMessage(api: Api, promptTokens = 0): AssistantMessage {
return {
role: "assistant",
content: [{ type: "text", text: "ok" }],
api,
@ -62,15 +61,19 @@ describe("createAgentSession stream options", () => {
usage: {
input: 0,
output: 0,
cacheRead: 0,
cacheRead: promptTokens,
cacheWrite: 0,
totalTokens: 0,
totalTokens: promptTokens,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: Date.now(),
};
stream.end(message);
}
function createDoneStream(api: Api, promptTokens = 0) {
const stream = createAssistantMessageEventStream();
stream.end(createDoneMessage(api, promptTokens));
return stream;
}
@ -123,6 +126,80 @@ describe("createAgentSession stream options", () => {
}
}
async function createCacheWarmingSession(populate?: (manager: SessionManager, model: Model<Api>) => void) {
const model: Model<Api> = {
...createModel("anthropic-messages"),
cost: { input: 10, output: 50, cacheRead: 0.25, cacheWrite: 12.5 },
promptCache: { short: 300 },
};
const authStorage = AuthStorage.create(join(agentDir, "auth.json"));
await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "test-api-key" }));
const modelRegistry = await createModelRegistry(authStorage, join(agentDir, "models.json"));
let providerCalls = 0;
modelRegistry.registerProvider(model.provider, {
api: model.api,
streamSimple: () => {
providerCalls++;
return createDoneStream(model.api, 100_000);
},
});
const sessionManager = SessionManager.inMemory(cwd);
populate?.(sessionManager, model);
const { session } = await createAgentSession({
cwd,
agentDir,
model,
modelRuntime: getModelRuntime(modelRegistry),
settingsManager: SettingsManager.inMemory({ cacheWarming: "idle" }),
sessionManager,
});
return {
session,
providerCalls: () => providerCalls,
dispose: () => {
session.dispose();
modelRegistry.unregisterProvider(model.provider);
},
};
}
it("schedules cache warming after a completed session request", async () => {
const fixture = await createCacheWarmingSession();
try {
await fixture.session.prompt("test");
expect(fixture.session.cacheWarmingStatus?.nextWarmAt).toBeGreaterThan(Date.now());
// Equivalent shallow copies remain current, but removing the request prefix does not.
fixture.session.agent.state.messages = [...fixture.session.agent.state.messages];
fixture.session.agent.state.model = { ...fixture.session.agent.state.model };
expect(fixture.session.cacheWarmingStatus?.nextWarmAt).toBeGreaterThan(Date.now());
fixture.session.agent.state.messages = fixture.session.agent.state.messages.slice(1);
expect(fixture.session.cacheWarmingStatus?.reason).toBe("conversation context changed");
} finally {
fixture.dispose();
}
});
it("waits for the next request instead of restoring cache warming", async () => {
const fixture = await createCacheWarmingSession((manager, model) => {
manager.appendModelChange(model.provider, model.id);
manager.appendThinkingLevelChange("off");
manager.appendMessage({ role: "user", content: "test", timestamp: Date.now() - 60_000 });
const assistant = { ...createDoneMessage(model.api, 100_000), timestamp: Date.now() - 59_000 };
manager.appendMessage(assistant);
manager.appendUsage("cache_warm", model.provider, model.id, assistant.usage);
});
try {
expect(fixture.providerCalls()).toBe(0);
expect(fixture.session.cacheWarmingStatus).toEqual({
state: "inactive",
reason: "waiting for first request",
});
} finally {
fixture.dispose();
}
});
it("forwards httpIdleTimeoutMs as timeoutMs for OpenAI Codex", async () => {
const options = await captureStreamOptions("openai-codex-responses", { httpIdleTimeoutMs: 1234 });

View file

@ -393,6 +393,30 @@ describe("SettingsManager", () => {
});
});
describe("cacheWarming", () => {
it("defaults to streaming and ignores project settings", () => {
expect(SettingsManager.create(projectDir, agentDir).getCacheWarmingMode()).toBe("streaming");
writeFileSync(join(projectDir, ".pi", "settings.json"), JSON.stringify({ cacheWarming: "idle" }));
expect(SettingsManager.create(projectDir, agentDir).getCacheWarmingMode()).toBe("streaming");
writeFileSync(join(agentDir, "settings.json"), JSON.stringify({ cacheWarming: "idle" }));
expect(SettingsManager.create(projectDir, agentDir).getCacheWarmingMode()).toBe("idle");
writeFileSync(join(agentDir, "settings.json"), JSON.stringify({ cacheWarming: "bogus" }));
expect(SettingsManager.create(projectDir, agentDir).getCacheWarmingMode()).toBe("streaming");
});
it("persists the mode globally", async () => {
const manager = SettingsManager.create(projectDir, agentDir);
manager.setCacheWarmingMode("off");
await manager.flush();
expect(SettingsManager.create(projectDir, agentDir).getCacheWarmingMode()).toBe("off");
expect(JSON.parse(readFileSync(join(agentDir, "settings.json"), "utf8"))).toEqual({ cacheWarming: "off" });
});
});
describe("externalEditor", () => {
const originalVisual = process.env.VISUAL;
const originalEditor = process.env.EDITOR;