mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(coding-agent): add prompt cache warming (#9668)
This commit is contained in:
parent
7d5eb0ee3b
commit
c596d09d9c
32 changed files with 1380 additions and 69 deletions
|
|
@ -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();
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
328
packages/coding-agent/test/cache-warmer.test.ts
Normal file
328
packages/coding-agent/test/cache-warmer.test.ts
Normal 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");
|
||||
});
|
||||
});
|
||||
|
|
@ -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: {
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue