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

@ -20,6 +20,7 @@ import type {
KnownProvider,
Model,
ModelCost,
ModelPromptCache,
OpenAICompletionsCompat,
OpenAIResponsesCompat,
} from "../src/types.ts";
@ -931,6 +932,21 @@ function applyOpenAIExplicitPromptCacheMetadata(model: Model<Api>): void {
};
}
// Anthropic ephemeral entries have a hard five-minute lifetime; `ttl: "1h"`
// extends it to one hour. Only direct Anthropic is annotated so cache warming
// does not assume equivalent behavior through proxies.
// https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
const ANTHROPIC_PROMPT_CACHE: ModelPromptCache = { short: 300, long: 3600 };
function applyPromptCacheMetadata(model: Model<Api>): void {
if (model.provider === "anthropic" && model.api === "anthropic-messages") {
model.promptCache = ANTHROPIC_PROMPT_CACHE;
}
// Do not add OpenAI lifetimes yet. Before enabling warming for explicit
// OpenAI caches, re-evaluate it using observed expiry, replay, and billing
// behavior; a documented TTL alone does not establish full cache loss.
}
function isGemma4Model(modelId: string): boolean {
return /gemma-?4/.test(modelId.toLowerCase());
}
@ -3055,6 +3071,7 @@ async function generateModels() {
applyOpenAICompletionsTranscriptMetadata(model);
applyOpenAIResponsesTranscriptMetadata(model);
applyOpenAIExplicitPromptCacheMetadata(model);
applyPromptCacheMetadata(model);
}
applyAnthropicAllowedFallbackModelMetadata(allModels.filter(isAnthropicFallbackMetadataModel));

View file

@ -107,6 +107,12 @@ export interface ThinkingBudgets {
// Base options all providers share
export type CacheRetention = "none" | "short" | "long";
/**
* Best-effort prompt cache lifetime in seconds for each retention tier a request can ask for.
* A missing tier means the lifetime is unknown; pi does not warm such caches.
*/
export type ModelPromptCache = Partial<Record<Exclude<CacheRetention, "none">, number>>;
export type Transport = "sse" | "websocket" | "websocket-cached" | "auto";
/** Provider-scoped environment overrides. Values take precedence over process.env. */
@ -963,6 +969,8 @@ export interface Model<TApi extends Api> {
thinkingLevelMap?: ThinkingLevelMap;
input: ("text" | "image")[];
cost: ModelCost;
/** Prompt cache lifetimes per retention tier. Unset when the provider's cache behavior is unknown. */
promptCache?: ModelPromptCache;
contextWindow: number;
maxTokens: number;
/** Default sampling parameters for this model. See {@link StreamOptions.samplingParams}; per-request keys override these. */

View file

@ -240,6 +240,13 @@ describe("Cache Retention (PI_CACHE_RETENTION)", () => {
});
describe("OpenAI Responses Provider", () => {
it.each(["gpt-5.6-luna", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-6-astra"] as const)(
"does not enable cache warming from the documented TTL alone for %s",
(modelId) => {
expect(getModel("openai", modelId).promptCache).toBeUndefined();
},
);
it.skipIf(!process.env.OPENAI_API_KEY)(
"should not set prompt_cache_retention when PI_CACHE_RETENTION is not set",
async () => {

View file

@ -733,6 +733,9 @@ interface ProviderModelConfig {
cacheWrite: number;
};
/** Best-effort prompt cache lifetime in seconds per retention tier. Unset disables cache warming. */
promptCache?: { short?: number; long?: number };
/** Maximum context window size in tokens. */
contextWindow: number;

View file

@ -738,6 +738,25 @@ pi.on("after_provider_response", (event, ctx) => {
Header availability depends on provider and transport. Providers that abstract HTTP responses may not expose headers.
#### cache_warming_decision
Fired before each prompt-cache refresh with pi's decision filled in. The event carries only pi's cost estimates; use `ctx.model`, `ctx.isIdle()`, and `ctx.getContextUsage()` for everything else.
```typescript
pi.on("cache_warming_decision", (event, ctx) => {
// event.warmCost: price of this refresh
// event.missCost: extra price of the next request if the entry is lost
// event.continuationProbability: pi's estimate that a request arrives in time
// event.action: "warm" | "stop", pi's decision
if (ctx.model?.provider === "my-provider") {
return { action: "stop" };
}
});
```
Return `{ action: "warm" }` or `{ action: "stop" }` to override; the last handler that returns an action wins. `"stop"` ends warming until the next real request.
### Model Events
#### model_select

View file

@ -9,6 +9,7 @@ Add custom providers and models (Ollama, vLLM, LM Studio, proxies) via `~/.pi/ag
- [Supported APIs](#supported-apis)
- [Provider Configuration](#provider-configuration)
- [Model Configuration](#model-configuration)
- [Prompt Cache Lifetimes](#prompt-cache-lifetimes)
- [Overriding Built-in Providers](#overriding-built-in-providers)
- [Per-model Overrides](#per-model-overrides)
- [Anthropic Messages Compatibility](#anthropic-messages-compatibility)
@ -208,6 +209,7 @@ If your command is slow, expensive, rate-limited, or should keep using a previou
| `maxTokens` | No | `16384` | Maximum output tokens |
| `samplingParams` | No | omitted | Sampling parameters merged verbatim into every request body (see below) |
| `cost` | No | all zeros | Per-million-token rates with optional request-wide input pricing tiers |
| `promptCache` | No | omitted | Best-effort prompt cache lifetime in seconds per retention tier (see below) |
| `compat` | No | provider `compat` | Provider compatibility overrides. Merged with provider-level `compat` when both are set. |
A cost tier supplies a complete alternate rate set and applies to the full request when total input usage (`input + cacheRead + cacheWrite`) exceeds `inputTokensAbove`. When multiple tiers match, the highest threshold wins.
@ -236,6 +238,19 @@ Current behavior:
- `/model`, `--list-models`, and the interactive footer display entries by model `id`.
- The configured `name` is used for model matching and secondary model detail text. It does not replace the footer/status-bar model id.
### Prompt Cache Lifetimes
`promptCache` states how long the provider keeps a prompt cache entry alive for each retention tier pi can request (`short` is the default tier; `long` is used when `PI_CACHE_RETENTION=long`). Values are seconds and are estimates: providers publish ranges, so pick the conservative end.
```json
{
"id": "claude-sonnet-5",
"promptCache": { "short": 300, "long": 3600 }
}
```
The built-in catalog fills this in for direct Anthropic (5 min / 1 h). Other providers, including direct OpenAI, have no built-in lifetime until their cache-expiry and replay behavior has been validated for warming. A model without a value for the tier a request used is never warmed; custom models and provider overrides can opt in when the backing cache behavior is known. See [Cache Warming](settings.md#cache-warming).
### Sampling Parameters
`samplingParams` is a free-form object merged verbatim into every request body for the model, after the fields pi sets itself, so its keys win. Use it to send sampling parameters pi does not model — including server-specific ones like llama.cpp's `min_p` or vLLM's `top_k`:
@ -359,7 +374,23 @@ Use `modelOverrides` to customize built-in models and matching extension-registe
}
```
`modelOverrides` supports these fields per model: `name`, `reasoning`, `thinkingLevelMap`, `input`, `cost` (partial), `contextWindow`, `maxTokens`, `samplingParams` (merged per key), `headers`, `compat`.
`modelOverrides` supports these fields per model: `name`, `reasoning`, `thinkingLevelMap`, `input`, `cost` (partial), `promptCache` (merged per tier), `contextWindow`, `maxTokens`, `samplingParams` (merged per key), `headers`, `compat`.
Use a `promptCache` override to enable cache warming through a proxy whose backing cache you know, for example OpenRouter routed to Anthropic:
```json
{
"providers": {
"openrouter": {
"modelOverrides": {
"anthropic/claude-sonnet-4": {
"promptCache": { "short": 300 }
}
}
}
}
}
```
Direct OpenAI GPT-5.6 Sol, Terra, and Luna default to a `272000` context window so requests remain within OpenAI's short-context pricing tier. To opt into OpenAI's 1.05M context window, increase it for each model you use:

View file

@ -256,6 +256,16 @@ Emitted when the user changes the thinking/reasoning level.
{"type":"thinking_level_change","id":"e5f6g7h8","parentId":"d4e5f6g7","timestamp":"2024-12-03T14:06:00.000Z","thinkingLevel":"high"}
```
### UsageEntry
Records model-attributed usage that is not an assistant message and does not participate in LLM context. `kind` is an arbitrary string identifying the operation; for example, cache warming uses `"cache_warm"`.
```json
{"type":"usage","id":"f6g7h8i9","parentId":"e5f6g7h8","timestamp":"2024-12-03T14:08:00.000Z","kind":"cache_warm","provider":"anthropic","model":"claude-sonnet-4-5","usage":{"input":0,"output":0,"cacheRead":50000,"cacheWrite":0,"totalTokens":50000,"cost":{"input":0,"output":0,"cacheRead":0.015,"cacheWrite":0,"total":0.015}}}
```
Usage entries contribute to session token and cost totals. Pi hides them from the conversation tree. Consumers should treat unknown `kind` values as normal usage rather than rejecting them.
### CompactionEntry
Created when context is compacted. Stores a summary of earlier messages and a complete system prompt/tool checkpoint.
@ -364,7 +374,7 @@ Entries normally form one tree, but navigation APIs can create multiple roots:
- `compaction` -> complete system checkpoint followed by `compactionSummary`
- `branch_summary` -> `branchSummary`
- `custom_message` -> `CustomMessage`
- `custom` -> no context message
- `usage` and `custom` -> no context message
The compaction summary replaces entries before `firstKeptEntryId`. Pre-compaction system messages are folded into the complete checkpoint rather than replayed from the retained range. Retained non-system entries and all entries after the compaction remain available to the LLM.
@ -391,6 +401,9 @@ for (const line of lines) {
case "branch_summary":
console.log(`[${entry.id}] Branch from ${entry.fromId}`);
break;
case "usage":
console.log(`[${entry.id}] Usage (${entry.kind}): ${entry.usage.totalTokens} tokens`);
break;
case "custom":
console.log(`[${entry.id}] Custom (${entry.customType}): ${JSON.stringify(entry.data)}`);
break;
@ -435,6 +448,7 @@ Key methods for working with sessions programmatically.
- `appendMessage(message)` - Add message
- `appendThinkingLevelChange(level)` - Record thinking change
- `appendModelChange(provider, modelId)` - Record model change
- `appendUsage(kind, provider, model, usage)` - Record model-attributed usage outside the conversation
- `appendCompaction(summary, firstKeptEntryId, tokensBefore, details?, fromHook?, usage?)` - Add compaction
- `appendCustomEntry(customType, data?)` - Extension state (not in context)
- `appendSessionInfo(name)` - Set session display name

View file

@ -32,8 +32,31 @@ Use `/trust` in interactive mode to save a project trust decision for future ses
| `defaultThinkingLevel` | string | - | Startup thinking level (saved with Ctrl+S in `/thinking`, or edited manually): `"off"`, `"minimal"`, `"low"`, `"medium"`, `"high"`, `"xhigh"`, `"max"` |
| `modelThinkingLevels` | object | - | Per-model startup thinking levels keyed by `"provider/modelId"`; configure from `/settings` → Default thinking level per model or edit manually |
| `hideThinkingBlock` | boolean | `false` | Hide thinking blocks in output |
| `showCacheMissNotices` | boolean | `false` | Show transcript notices for significant prompt-cache misses, compaction or branch-summary usage, and provider recovery diagnostics such as dropped Anthropic thinking blocks |
| `showCacheMissNotices` | boolean | `false` | Show transcript notices for significant prompt-cache misses, successful cache-warming usage, compaction or branch-summary usage, and provider recovery diagnostics such as dropped Anthropic thinking blocks |
| `thinkingBudgets` | object | - | Custom token budgets per thinking level. Anthropic, Google, and Bedrock use these natively. OpenAI-compatible models use them when `compat.thinkingTokenBudgetField` (or `supportsThinkingTokenBudget`) is set. |
| `cacheWarming` | string | `"streaming"` | Prompt cache-warming mode: `"off"`, `"streaming"`, or `"idle"`. Global setting only. |
#### Cache Warming
Providers drop a prompt cache entry after a period of inactivity, so the first request after a pause pays full input price again. Cache warming re-sends the last request with a one-token output budget shortly before expiry:
- `"off"` disables warming.
- `"streaming"` protects expensive prefixes during long tool executions and stops as soon as the agent settles.
- `"idle"` also considers refreshes while waiting for your next prompt, using a fixed 15% continuation probability measured from real usage.
```json
{
"cacheWarming": "idle"
}
```
A refresh is sent only when the expected avoided cache-miss cost, minus the cost of the refresh, leaves at least $0.05 of expected savings. Active agent runs use 100% continuation probability. `/session` shows the next decision, continuation probability, expected savings, threshold, and estimated costs. When cache miss notices are enabled, each successful refresh appears in the transcript with its cost; notices identify extension overrides.
Warming stops when the context changes (model switch, compaction, branch navigation). Idle warming stops no later than 30 minutes after the last real provider request; warming during an active agent run stops after 60 minutes. Extensions can override each decision through the [`cache_warming_decision`](extensions.md#cache_warming_decision) event.
Each refresh is billed as a cache read of the full context plus one output token. Usage and cost show up in session totals but never enter model context. Pi schedules candidates at 90% of the cache lifetime while leaving at least ten seconds before expiry.
Warming needs a known cache lifetime for the model and the retention tier the request used (`short`, or `long` with `PI_CACHE_RETENTION=long`). The built-in catalog carries lifetimes for direct Anthropic; custom models and other providers can declare theirs with `promptCache` in `models.json` (see [Prompt Cache Lifetimes](models.md#prompt-cache-lifetimes)). Claude models that use budget-based rather than adaptive thinking are skipped while thinking is on, because Anthropic derives the thinking budget from `max_tokens` and keys the message cache on it, so a one-token request cannot reproduce the entry.
#### thinkingBudgets

View file

@ -55,6 +55,7 @@ import { normalizeToolResultImages } from "../utils/tool-result-images.ts";
import { formatNoApiKeyFoundMessage, formatNoModelSelectedMessage } from "./auth-guidance.ts";
import { type BashResult, executeBashWithOperations } from "./bash-executor.ts";
import { generateBugReportSummary } from "./bug-report.ts";
import type { CacheWarmer, CacheWarmingStatus } from "./cache-warmer.ts";
import {
type CompactionPreparation,
type CompactionResult,
@ -106,7 +107,7 @@ import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.t
import { exportSessionToJsonl } from "./session-export.ts";
import type { BranchSummaryEntry, CompactionEntry, SessionEntry, SessionManager } from "./session-manager.ts";
import { getLatestCompactionEntry } from "./session-manager.ts";
import type { SettingsManager } from "./settings-manager.ts";
import type { CacheWarmingMode, SettingsManager } from "./settings-manager.ts";
import type { SlashCommandInfo } from "./slash-commands.ts";
import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.ts";
import {
@ -218,6 +219,8 @@ export interface AgentSessionConfig {
customTools?: ToolDefinition[];
/** Canonical model/auth runtime used by coding-agent internals. */
modelRuntime: ModelRuntime;
/** Keeps the prompt cache entry of the last session request warm. */
cacheWarmer?: Pick<CacheWarmer, "cancel" | "status" | "onAgentSettled" | "onModeChanged" | "onWarmed">;
/** Initial active built-in tool names. Default: [read, bash, edit, write] */
initialActiveToolNames?: string[];
/** Optional allowlist of tool names. When provided, only these tool names are exposed. */
@ -374,6 +377,7 @@ export class AgentSession {
private _extensionErrorUnsubscriber?: () => void;
private _modelRuntime: ModelRuntime;
private _cacheWarmer?: Pick<CacheWarmer, "cancel" | "status" | "onAgentSettled" | "onModeChanged" | "onWarmed">;
// Tool registry for extension getTools/setTools
private _toolRegistry: Map<string, AgentTool> = new Map();
@ -394,6 +398,10 @@ export class AgentSession {
this._customTools = config.customTools ?? [];
this._cwd = config.cwd;
this._modelRuntime = config.modelRuntime;
this._cacheWarmer = config.cacheWarmer;
if (this._cacheWarmer) {
this._cacheWarmer.onWarmed = (entry) => this._emit({ type: "entry_appended", entry });
}
this._extensionRunnerRef = config.extensionRunnerRef;
this._initialActiveToolNames = config.initialActiveToolNames;
this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined;
@ -652,6 +660,7 @@ export class AgentSession {
}
private async _emitAgentSettled(): Promise<void> {
this._cacheWarmer?.onAgentSettled();
this._isAgentRunActive = false;
try {
await this._extensionRunner.emit({ type: "agent_settled" });
@ -921,6 +930,10 @@ export class AgentSession {
);
this._disconnectFromAgent();
this._eventListeners = [];
if (this._cacheWarmer) {
this._cacheWarmer.onWarmed = undefined;
this._cacheWarmer.cancel();
}
cleanupSessionResources(this.sessionId);
}
@ -933,6 +946,17 @@ export class AgentSession {
return this.agent.state;
}
/** Current cache-warming state and the policy inputs that produced it. */
get cacheWarmingStatus(): CacheWarmingStatus | undefined {
return this._cacheWarmer?.status;
}
/** Persist the cache-warming mode and immediately reconcile active warming. */
setCacheWarmingMode(mode: CacheWarmingMode): void {
this.settingsManager.setCacheWarmingMode(mode);
this._cacheWarmer?.onModeChanged();
}
/** Current model (may be undefined if not yet selected) */
get model(): Model<any> | undefined {
return this.agent.state.model;
@ -3461,7 +3485,9 @@ export class AgentSession {
const usageTotals = createUsageTotals();
for (const entry of this.sessionManager.getEntries()) {
if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) {
if (entry.type === "usage") {
addUsageToTotals(usageTotals, entry.usage);
} else if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) {
addUsageToTotals(usageTotals, entry.usage);
}
if (entry.type !== "message") continue;

View file

@ -117,7 +117,17 @@ function scan(
prev = undefined;
continue;
}
if (entry.type === "message" && entry.message.role === "assistant") {
if (entry.type === "usage" && entry.kind === "cache_warm") {
const promptTokens = entry.usage.input + entry.usage.cacheRead + entry.usage.cacheWrite;
if (promptTokens > 0) {
prev = {
promptTokens,
modelKey: `${entry.provider}/${entry.model}`,
timestamp: Date.parse(entry.timestamp),
reportedCache: true,
};
}
} else if (entry.type === "message" && entry.message.role === "assistant") {
const miss = detectMiss(prev, entry.message, models);
if (miss) {
totals.missedTokens += miss.missedTokens;

View file

@ -0,0 +1,437 @@
import {
type Api,
type Context,
calculateCost,
type Model,
type ModelsSimpleStreamOptions,
type SimpleStreamOptions,
type Usage,
} from "@earendil-works/pi-ai";
import { getProviderEnvValue } from "@earendil-works/pi-ai/utils/provider-env";
import type { ModelRuntime } from "./model-runtime.ts";
import type { SessionEntry, SessionManager, UsageEntry } from "./session-manager.ts";
import type { CacheWarmingMode } from "./settings-manager.ts";
/** Streaming warming never continues past this long after the real request that started it. */
const MAX_WARMING_AGE_MS = 60 * 60_000;
/** Idle warming uses a shorter horizon because continuation estimates become less reliable with age. */
const MAX_IDLE_WARMING_AGE_MS = 30 * 60_000;
/** A refresh is sent only when it is expected to save at least this many dollars. */
const CACHE_WARMING_MINIMUM_EXPECTED_SAVINGS = 0.05;
/**
* Chance that a real request arrives before the cache entry expires while the
* agent sits idle. Measured from our own usage; per-session estimates were not
* better than this constant.
*/
const IDLE_CONTINUATION_PROBABILITY = 0.15;
/** Refresh at 90% of the TTL while preserving at least ten seconds of margin. */
export function getCacheWarmingDelayMs(ttlMs: number): number | undefined {
if (ttlMs <= 10_000) return undefined;
return Math.max(1, Math.floor(Math.min(ttlMs * 0.9, ttlMs - 10_000)));
}
/**
* Lifetime of the prompt cache entry a request writes, from the model's
* `promptCache` tier for the retention the request used. Undefined when the
* model has no lifetime for that tier or caching is off.
*/
export function getPromptCacheTtlMs(model: Model<Api>, options: SimpleStreamOptions | undefined): number | undefined {
const retention =
options?.cacheRetention ??
(getProviderEnvValue("PI_CACHE_RETENTION", options?.env) === "long" ? "long" : "short");
if (retention === "none") return undefined;
const seconds = model.promptCache?.[retention];
return seconds === undefined ? undefined : seconds * 1000;
}
/**
* Whether replaying the request with a one-token output cap leaves its cache
* entry untouched. Anthropic's budget-based thinking (Claude models without
* adaptive thinking) derives `budget_tokens` from `max_tokens`; the replay
* would get a different budget, which Anthropic keys the message cache on,
* and the model could still think for thousands of tokens.
*/
export function isReplayable(model: Model<Api>, options: SimpleStreamOptions | undefined): boolean {
if (!options?.reasoning || model.api !== "anthropic-messages") return true;
return (model as Model<"anthropic-messages">).compat?.forceAdaptiveThinking === true;
}
/** Prompt size of the most recent real request on the branch, as reported by the provider. */
function lastPromptTokens(entries: SessionEntry[]): number {
for (let index = entries.length - 1; index >= 0; index--) {
const entry = entries[index];
if (entry.type === "message" && entry.message.role === "assistant") {
const usage = entry.message.usage;
return usage.input + usage.cacheRead + usage.cacheWrite;
}
}
return 0;
}
function price(
model: Model<Api>,
tokens: Partial<Pick<Usage, "input" | "output" | "cacheRead" | "cacheWrite">>,
): number {
const usage: Usage = {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
...tokens,
};
return calculateCost(model, usage).total;
}
export type CacheWarmingAction = "warm" | "stop";
/** Inputs and outcome of one warm-or-stop decision, as shown by `/session`. */
export interface CacheWarmingDecision {
/** "streaming" while the agent run that sent the request is still active. */
phase: "streaming" | "idle";
/** Price of this refresh: a cache read of the prompt plus one output token. */
warmCost: number;
/** Extra price of the next real request if the cache entry is lost. */
missCost: number;
/** Estimated chance that a real request arrives before the entry expires. */
continuationProbability: number;
/** `continuationProbability * missCost - warmCost`. */
expectedSavings: number;
/** False when the prompt size or the model's prices are unknown. */
economicsAvailable: boolean;
/** Pi's decision: "warm" when `expectedSavings` is at least $0.05. */
action: CacheWarmingAction;
}
/**
* Fired before each refresh with pi's decision filled in. Everything else an
* extension might want (model, idle state, context size) is on the context.
*/
export interface CacheWarmingDecisionEvent
extends Pick<CacheWarmingDecision, "warmCost" | "missCost" | "continuationProbability" | "action"> {
type: "cache_warming_decision";
}
export interface CacheWarmingDecisionEventResult {
/** Override whether this refresh is sent. "stop" ends warming until the next real request. */
action?: CacheWarmingAction;
}
export interface CacheWarmingStatus {
/** "scheduled": a refresh timer is armed; "refreshing": a warm request is in flight. */
state: "inactive" | "scheduled" | "refreshing";
/** Why nothing is scheduled. */
reason?: string;
nextWarmAt?: number;
/** The pending decision, or the decision that stopped warming. */
decision?: CacheWarmingDecision;
/** True when an extension changed `decision.action`. */
extensionOverride?: boolean;
}
/** The request whose prompt cache entry should be kept warm, exactly as it was sent. */
export interface CacheWarmRequest {
model: Model<Api>;
context: Context;
options: ModelsSimpleStreamOptions;
}
interface ActiveRun extends CacheWarmRequest {
/** False once the session's model or messages no longer match the request. */
isCurrent: () => boolean;
delayMs: number;
startedAt: number;
controller: AbortController;
phase: "streaming" | "idle";
nextWarmAt: number;
/** Set while a refresh that an extension forced is in flight. */
extensionOverride: boolean;
timer?: ReturnType<typeof setTimeout>;
}
/**
* Keeps one prompt cache entry alive by re-sending its request with a
* one-token output cap before the entry expires. `start` replaces any
* previous run; warm requests never extend the fixed safety windows.
*/
export class CacheWarmer {
private run?: ActiveRun;
private inactive: CacheWarmingStatus;
private readonly models: Pick<ModelRuntime, "streamSimple">;
private readonly sessionManager: Pick<SessionManager, "appendUsage" | "getBranch">;
private readonly getMode: () => CacheWarmingMode;
/** Lets extensions override `event.action`; failures fall back to pi's decision. */
private readonly decide: (event: CacheWarmingDecisionEvent) => Promise<CacheWarmingAction>;
/** Called with the persisted usage entry after each successful refresh. */
onWarmed?: (entry: UsageEntry) => void;
constructor(
models: Pick<ModelRuntime, "streamSimple">,
sessionManager: Pick<SessionManager, "appendUsage" | "getBranch">,
getMode: () => CacheWarmingMode,
decide: (event: CacheWarmingDecisionEvent) => Promise<CacheWarmingAction> = async (event) => event.action,
) {
this.models = models;
this.sessionManager = sessionManager;
this.getMode = getMode;
this.decide = decide;
this.inactive = { state: "inactive", reason: "waiting for first request" };
}
get status(): CacheWarmingStatus {
if (this.getMode() === "off") return { state: "inactive", reason: "cache warming disabled" };
const run = this.run;
if (!run) return this.inactive;
if (!run.isCurrent()) return { state: "inactive", reason: "conversation context changed" };
const decision = this.evaluate(run);
const refreshing = run.timer === undefined;
if (!decision.economicsAvailable && !refreshing) {
return { state: "inactive", reason: "cache economics unavailable" };
}
return {
state: refreshing ? "refreshing" : "scheduled",
nextWarmAt: run.nextWarmAt,
decision,
extensionOverride: run.extensionOverride,
};
}
/** Keep the prompt cache entry written by `request` warm while `isCurrent` holds. */
start(request: CacheWarmRequest, isCurrent: () => boolean): void {
this.clearRun();
const mode = this.getMode();
if (mode === "off") {
this.stop("cache warming disabled");
return;
}
if (!isReplayable(request.model, request.options)) {
this.stop("request cannot be replayed safely");
return;
}
const ttlMs = getPromptCacheTtlMs(request.model, request.options);
if (ttlMs === undefined) {
this.stop(
request.options.cacheRetention === "none"
? "request disabled prompt caching"
: "cache lifetime unavailable",
);
return;
}
const delayMs = getCacheWarmingDelayMs(ttlMs);
if (delayMs === undefined) {
this.stop("cache lifetime unavailable");
return;
}
this.run = {
...request,
isCurrent,
delayMs,
startedAt: Date.now(),
controller: new AbortController(),
phase: "streaming",
nextWarmAt: 0,
extensionOverride: false,
};
this.schedule(this.run);
}
onAgentSettled(): void {
const run = this.run;
if (!run) return;
if (this.getMode() === "streaming") {
this.stop("agent run settled");
return;
}
run.phase = "idle";
const deadline = run.startedAt + MAX_IDLE_WARMING_AGE_MS;
if (run.nextWarmAt > deadline || Date.now() >= deadline) {
this.stop("30-minute idle safety limit reached");
}
}
/** Reconcile an active run after the persisted warming mode changes. */
onModeChanged(): void {
const run = this.run;
if (!run) return;
const reason = this.getModeStopReason(run);
if (reason) this.stop(reason);
}
cancel(): void {
this.stop("inactive");
}
private clearRun(): void {
const run = this.run;
if (!run) return;
this.run = undefined;
if (run.timer) clearTimeout(run.timer);
run.controller.abort();
}
private stop(reason: string, stopped?: Pick<CacheWarmingStatus, "decision" | "extensionOverride">): void {
this.clearRun();
this.inactive = { state: "inactive", reason, ...stopped };
}
private schedule(run: ActiveRun): void {
run.extensionOverride = false;
run.nextWarmAt = Date.now() + run.delayMs;
const deadline = run.startedAt + (run.phase === "idle" ? MAX_IDLE_WARMING_AGE_MS : MAX_WARMING_AGE_MS);
if (run.nextWarmAt > deadline || Date.now() >= deadline) {
this.stop(run.phase === "idle" ? "30-minute idle safety limit reached" : "one-hour safety limit reached");
return;
}
run.timer = setTimeout(() => void this.refresh(run), Math.max(0, run.nextWarmAt - Date.now()));
run.timer.unref?.();
}
private async refresh(run: ActiveRun): Promise<void> {
run.timer = undefined;
if (!this.validateRun(run)) return;
const decision = this.evaluate(run);
const { warmCost, missCost, continuationProbability } = decision;
let action = decision.action;
try {
action = await this.decide({
type: "cache_warming_decision",
warmCost,
missCost,
continuationProbability,
action,
});
} catch {
// Extension failures fall back to pi's own decision.
}
if (!this.validateRun(run)) return;
const extensionOverride = action !== decision.action;
if (action === "stop") {
const reason = extensionOverride
? "stopped by extension"
: decision.economicsAvailable
? "expected savings below threshold"
: "cache economics unavailable";
this.stop(reason, { decision, extensionOverride });
return;
}
run.extensionOverride = extensionOverride;
try {
const message = await this.models
.streamSimple(run.model, run.context, {
...run.options,
maxTokens: 1,
maxRetries: 0,
signal: run.controller.signal,
})
.result();
if (!this.validateRun(run)) return;
if (message.stopReason !== "error" && message.stopReason !== "aborted") {
const entry = this.sessionManager.appendUsage(
"cache_warm",
message.provider,
message.responseModel ?? message.model,
message.usage,
extensionOverride ? "extension override" : undefined,
);
this.onWarmed?.(entry);
}
} catch {
// Cache warming is best-effort and must not affect the active agent run.
}
if (this.run === run) this.schedule(run);
}
private validateRun(run: ActiveRun): boolean {
if (this.run !== run) return false;
const reason = this.getModeStopReason(run) ?? (!run.isCurrent() ? "conversation context changed" : undefined);
if (!reason) return true;
this.stop(reason);
return false;
}
private getModeStopReason(run: ActiveRun): string | undefined {
const mode = this.getMode();
if (mode === "off") return "cache warming disabled";
if (mode === "streaming" && run.phase === "idle") return "agent run settled";
return undefined;
}
private evaluate(run: ActiveRun): CacheWarmingDecision {
const model = run.model;
const promptTokens = lastPromptTokens(this.sessionManager.getBranch());
const cacheHitCost = price(model, { cacheRead: promptTokens });
const cacheMissCost = price(
model,
model.cost.cacheWrite > 0 ? { cacheWrite: promptTokens } : { input: promptTokens },
);
const warmCost = price(model, { cacheRead: promptTokens, output: 1 });
const missCost = Math.max(0, cacheMissCost - cacheHitCost);
const continuationProbability = run.phase === "idle" ? IDLE_CONTINUATION_PROBABILITY : 1;
const economicsAvailable = promptTokens > 0 && (cacheHitCost > 0 || cacheMissCost > 0);
const expectedSavings = continuationProbability * missCost - warmCost;
return {
phase: run.phase,
warmCost,
missCost,
continuationProbability,
expectedSavings,
economicsAvailable,
action: expectedSavings >= CACHE_WARMING_MINIMUM_EXPECTED_SAVINGS ? "warm" : "stop",
};
}
}
function formatDollars(value: number): string {
return value < 0 ? `-$${Math.abs(value).toFixed(3)}` : `$${value.toFixed(3)}`;
}
function formatCacheWarmingEconomics(decision: CacheWarmingDecision): string {
if (!decision.economicsAvailable) return "cache economics unavailable";
const probability = Math.round(decision.continuationProbability * 100);
const probabilityText =
decision.phase === "streaming"
? `${probability}% continuation probability while agent is running`
: `${probability}% continuation probability`;
const comparison = decision.action === "warm" ? ">=" : "<";
return `${probabilityText}, expected savings ${formatDollars(decision.expectedSavings)} ${comparison} $${CACHE_WARMING_MINIMUM_EXPECTED_SAVINGS.toFixed(3)}`;
}
function formatCacheWarmingDecisionTime(nextWarmAt: number | undefined, now: number): string {
if (nextWarmAt === undefined || nextWarmAt <= now) return "Decision now";
let remainingSeconds = Math.ceil((nextWarmAt - now) / 1000);
const hours = Math.floor(remainingSeconds / 3600);
remainingSeconds %= 3600;
const minutes = Math.floor(remainingSeconds / 60);
const seconds = remainingSeconds % 60;
const parts: string[] = [];
if (hours > 0) parts.push(`${hours}h`);
if (minutes > 0) parts.push(`${minutes}m`);
if (seconds > 0 || parts.length === 0) parts.push(`${seconds}s`);
return `Decision in ${parts.join(" ")}`;
}
/** One-line status for `/session`. */
export function formatCacheWarmingStatus(status: CacheWarmingStatus, now = Date.now()): string {
const decision = status.decision;
// A decision is attached once pi (or an extension) acted on it; "inactive"
// without one never got that far.
if (!decision || (status.state === "inactive" && !decision.economicsAvailable && !status.extensionOverride)) {
return `Inactive (${status.reason ?? "unknown reason"})`;
}
const details = status.extensionOverride
? `extension override, ${formatCacheWarmingEconomics(decision)}`
: `${formatCacheWarmingEconomics(decision)} -> ${decision.action}`;
if (status.state === "inactive") return `Stopped (${details})`;
if (status.state === "refreshing") return `Warming cache (${details})`;
return `${formatCacheWarmingDecisionTime(status.nextWarmAt, now)} (${details})`;
}
/** One-line transcript text for persisted cache-warming usage. */
export function formatCacheWarmingUsage(entry: UsageEntry): string {
const note = entry.note ? ` (${entry.note})` : "";
const cost = entry.usage.cost.total.toFixed(6).replace(/(\.\d{3}\d*?)0+$/, "$1");
return `Cache warmed${note}: $${cost}`;
}

View file

@ -40,6 +40,8 @@ export type {
BeforeProviderRequestEvent,
BeforeProviderRequestEventResult,
BuildSystemPromptOptions,
CacheWarmingDecisionEvent,
CacheWarmingDecisionEventResult,
// Context
CompactOptions,
// Events - Agent

View file

@ -6,6 +6,7 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { ImageContent, Model, Provider, ProviderHeaders } from "@earendil-works/pi-ai";
import type { KeyId } from "@earendil-works/pi-tui";
import { type Theme, theme } from "../../modes/interactive/theme/theme.ts";
import type { CacheWarmingAction } from "../cache-warmer.ts";
import type { ResourceDiagnostic } from "../diagnostics.ts";
import type { KeybindingsConfig } from "../keybindings.ts";
import type { ModelRegistry } from "../model-registry.ts";
@ -22,6 +23,8 @@ import type {
BeforeAgentStartEventResult,
BeforeProviderHeadersEvent,
BeforeProviderRequestEvent,
CacheWarmingDecisionEvent,
CacheWarmingDecisionEventResult,
CompactOptions,
ContextEvent,
ContextEventResult,
@ -161,6 +164,7 @@ type RunnerEmitEvent = Exclude<
| ToolResultEvent
| UserBashEvent
| ContextEvent
| CacheWarmingDecisionEvent
| BeforeProviderRequestEvent
| BeforeProviderHeadersEvent
| BeforeAgentStartEvent
@ -913,6 +917,30 @@ export class ExtensionRunner {
return result as RunnerEmitResult<TEvent>;
}
/** Returns the event's own action unless a handler overrides it; the last override wins. */
async emitCacheWarmingDecision(event: CacheWarmingDecisionEvent): Promise<CacheWarmingAction> {
const ctx = this.createContext();
let action = event.action;
for (const { ext, handlers } of snapshotEventHandlers(this.extensions, event.type)) {
for (const handler of handlers) {
try {
const result = (await handler(event, ctx)) as CacheWarmingDecisionEventResult | undefined;
if (result?.action !== undefined) action = result.action;
} catch (err) {
this.emitError({
extensionPath: ext.path,
event: event.type,
error: err instanceof Error ? err.message : String(err),
stack: err instanceof Error ? err.stack : undefined,
});
}
}
}
return action;
}
async emitMessageEnd(event: MessageEndEvent): Promise<AgentMessage | undefined> {
const ctx = this.createContext();
let currentMessage = event.message;

View file

@ -47,6 +47,7 @@ import type {
import type { Static, TSchema } from "typebox";
import type { Theme } from "../../modes/interactive/theme/theme.ts";
import type { BashResult } from "../bash-executor.ts";
import type { CacheWarmingDecisionEvent, CacheWarmingDecisionEventResult } from "../cache-warmer.ts";
import type { CompactionPreparation, CompactionResult } from "../compaction/index.ts";
import type { EventBus } from "../event-bus.ts";
import type { ExecOptions, ExecResult } from "../exec.ts";
@ -1088,6 +1089,7 @@ export type ExtensionEvent =
| ResourcesDiscoverEvent
| SessionEvent
| ContextEvent
| CacheWarmingDecisionEvent
| BeforeProviderRequestEvent
| BeforeProviderHeadersEvent
| AfterProviderResponseEvent
@ -1122,6 +1124,8 @@ export interface ContextEventResult {
export type BeforeProviderRequestEventResult = unknown;
export type { CacheWarmingDecisionEvent, CacheWarmingDecisionEventResult } from "../cache-warmer.ts";
export interface ToolCallEventResult {
/** Block tool execution. To modify arguments, mutate `event.input` in place instead. */
block?: boolean;
@ -1287,6 +1291,10 @@ export interface ExtensionAPI {
): () => void;
on(event: "session_tree", handler: ExtensionHandler<SessionTreeEvent>): () => void;
on(event: "context", handler: ExtensionHandler<ContextEvent, ContextEventResult>): () => void;
on(
event: "cache_warming_decision",
handler: ExtensionHandler<CacheWarmingDecisionEvent, CacheWarmingDecisionEventResult>,
): () => void;
on(
event: "before_provider_request",
handler: ExtensionHandler<BeforeProviderRequestEvent, BeforeProviderRequestEventResult>,
@ -1597,6 +1605,8 @@ export interface ProviderModelConfig {
input: ("text" | "image")[];
/** Per-million-token cost rates and optional request-wide input pricing tiers. */
cost: Model<Api>["cost"];
/** Best-effort prompt cache lifetime in seconds per retention tier. Unset disables cache warming. */
promptCache?: Model<Api>["promptCache"];
/** Maximum context window size in tokens. */
contextWindow: number;
/** Maximum output tokens. */

View file

@ -26,6 +26,7 @@ export {
createAgentSessionServices,
} from "./agent-session-services.ts";
export { type BashExecutorOptions, type BashResult, executeBashWithOperations } from "./bash-executor.ts";
export type { CacheWarmingDecision, CacheWarmingStatus } from "./cache-warmer.ts";
export type { CompactionResult } from "./compaction/index.ts";
export { createEventBus, type EventBus, type EventBusController } from "./event-bus.ts";
export { areExperimentalFeaturesEnabled } from "./experimental.ts";
@ -39,6 +40,8 @@ export {
type BeforeAgentStartEvent,
type BeforeAgentStartEventResult,
type BuildSystemPromptOptions,
type CacheWarmingDecisionEvent,
type CacheWarmingDecisionEventResult,
type ContextEvent,
defineTool,
discoverAndLoadExtensions,

View file

@ -136,6 +136,10 @@ const ModelCostSchema = Type.Object({
...ModelCostRatesSchema,
tiers: Type.Optional(Type.Array(ModelCostTierSchema)),
});
const ModelPromptCacheSchema = Type.Object({
short: Type.Optional(Type.Number({ exclusiveMinimum: 0 })),
long: Type.Optional(Type.Number({ exclusiveMinimum: 0 })),
});
const AnthropicMessagesCompatSchema = Type.Object({
supportsEagerToolInputStreaming: Type.Optional(Type.Boolean()),
@ -174,6 +178,7 @@ const ModelDefinitionSchema = Type.Object({
thinkingLevelMap: Type.Optional(ThinkingLevelMapSchema),
input: Type.Optional(Type.Array(Type.Union([Type.Literal("text"), Type.Literal("image")]))),
cost: Type.Optional(ModelCostSchema),
promptCache: Type.Optional(ModelPromptCacheSchema),
contextWindow: Type.Optional(Type.Number()),
maxTokens: Type.Optional(Type.Number()),
samplingParams: Type.Optional(Type.Record(Type.String(), Type.Unknown())),
@ -195,6 +200,7 @@ const ModelOverrideSchema = Type.Object({
tiers: Type.Optional(Type.Array(ModelCostTierSchema)),
}),
),
promptCache: Type.Optional(ModelPromptCacheSchema),
contextWindow: Type.Optional(Type.Number()),
maxTokens: Type.Optional(Type.Number()),
samplingParams: Type.Optional(Type.Record(Type.String(), Type.Unknown())),

View file

@ -65,6 +65,7 @@ export interface ProviderConfigInput {
thinkingLevelMap?: Model<Api>["thinkingLevelMap"];
input: ("text" | "image")[];
cost: Model<Api>["cost"];
promptCache?: Model<Api>["promptCache"];
contextWindow: number;
maxTokens: number;
samplingParams?: Record<string, unknown>;
@ -122,6 +123,7 @@ function applyModelOverride(model: Model<Api>, override: ModelsJsonModelOverride
tiers: override.cost.tiers ?? model.cost.tiers,
}
: model.cost,
promptCache: override.promptCache ? { ...model.promptCache, ...override.promptCache } : model.promptCache,
contextWindow: override.contextWindow ?? model.contextWindow,
maxTokens: override.maxTokens ?? model.maxTokens,
samplingParams: override.samplingParams
@ -161,6 +163,7 @@ function modelFromJson(
thinkingLevelMap: definition.thinkingLevelMap,
input: (definition.input ?? ["text"]) as ("text" | "image")[],
cost: definition.cost ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
promptCache: definition.promptCache,
contextWindow: definition.contextWindow ?? 128000,
maxTokens: definition.maxTokens ?? 16384,
samplingParams: definition.samplingParams,

View file

@ -1,10 +1,12 @@
import { join } from "node:path";
import { Agent, type AgentMessage, setDefaultStreamFn, type ThinkingLevel } from "@earendil-works/pi-agent-core";
import type { ModelsSimpleStreamOptions } from "@earendil-works/pi-ai";
import { clampThinkingLevel, type Message, type Model, streamSimple } from "@earendil-works/pi-ai/compat";
import { getAgentDir } from "../config.ts";
import { resolvePath } from "../utils/paths.ts";
import { AgentSession } from "./agent-session.ts";
import { formatNoModelsAvailableMessage } from "./auth-guidance.ts";
import { CacheWarmer } from "./cache-warmer.ts";
import { DEFAULT_THINKING_LEVEL } from "./defaults.ts";
import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefinition } from "./extensions/index.ts";
import { convertToLlm } from "./messages.ts";
@ -302,6 +304,66 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
};
const extensionRunnerRef: { current?: ExtensionRunner } = {};
const cacheWarmer = new CacheWarmer(
modelRuntime,
sessionManager,
() => settingsManager.getCacheWarmingMode(),
async (event) => extensionRunnerRef.current?.emitCacheWarmingDecision(event) ?? event.action,
);
const buildRequestOptions = (
requestModel: Model<any>,
options: ModelsSimpleStreamOptions = {},
): ModelsSimpleStreamOptions => {
const providerRetrySettings = settingsManager.getProviderRetrySettings();
const httpIdleTimeoutMs = settingsManager.getHttpIdleTimeoutMs();
const effectiveTimeoutMs = httpIdleTimeoutMs === 0 ? 2147483647 : httpIdleTimeoutMs;
const headerRunner = extensionRunnerRef.current;
return {
...options,
timeoutMs: options.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs,
websocketConnectTimeoutMs: options.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs(),
maxRetries: options.maxRetries ?? providerRetrySettings.maxRetries,
maxRetryDelayMs: options.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs,
transformHeaders: async (requestHeaders) => {
const headers = mergeProviderAttributionHeaders(
requestModel,
settingsManager,
options.sessionId,
requestHeaders,
);
return headerRunner?.hasHandlers("before_provider_headers")
? headerRunner.emitBeforeProviderHeaders(headers ?? {})
: (headers ?? {});
},
};
};
const cacheContextIsCurrent = (requestModel: Model<any>) => {
const messages = agent.state.messages;
return () => {
const currentModel = agent.state.model;
const currentMessages = agent.state.messages;
return (
currentModel.provider === requestModel.provider &&
currentModel.id === requestModel.id &&
messages.length <= currentMessages.length &&
messages.every((message, index) => currentMessages[index] === message)
);
};
};
const transformProviderPayload = async (payload: unknown) => {
const runner = extensionRunnerRef.current;
if (!runner?.hasHandlers("before_provider_request")) return payload;
return runner.emitBeforeProviderRequest(payload);
};
const handleProviderResponse: NonNullable<ModelsSimpleStreamOptions["onResponse"]> = async (response) => {
const runner = extensionRunnerRef.current;
if (!runner?.hasHandlers("after_provider_response")) return;
await runner.emit({
type: "after_provider_response",
status: response.status,
headers: response.headers,
});
};
agent = new Agent({
initialState: {
@ -312,52 +374,19 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
},
convertToLlm: convertToLlmWithBlockImages,
streamFn: async (model, context, options) => {
const providerRetrySettings = settingsManager.getProviderRetrySettings();
const httpIdleTimeoutMs = settingsManager.getHttpIdleTimeoutMs();
// SDKs treat timeout=0 as 0ms (immediate timeout), not "no timeout".
// Use max int32 to effectively disable the timeout.
const effectiveTimeoutMs = httpIdleTimeoutMs === 0 ? 2147483647 : httpIdleTimeoutMs;
const timeoutMs = options?.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs;
const websocketConnectTimeoutMs =
options?.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs();
const headerRunner = extensionRunnerRef.current;
return modelRuntime.streamSimple(model, context, {
...options,
timeoutMs,
websocketConnectTimeoutMs,
maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries,
maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs,
transformHeaders: async (requestHeaders) => {
const headers = mergeProviderAttributionHeaders(
model,
settingsManager,
options?.sessionId,
requestHeaders,
);
return headerRunner?.hasHandlers("before_provider_headers")
? headerRunner.emitBeforeProviderHeaders(headers ?? {})
: (headers ?? {});
},
});
},
onPayload: async (payload, _model) => {
const runner = extensionRunnerRef.current;
if (!runner?.hasHandlers("before_provider_request")) {
return payload;
const requestOptions = buildRequestOptions(model, options);
// Compaction and summaries use their own routing ids; only session requests
// replace the cache entry, so warming restarts from them. Keep warming while
// the current transcript still extends the request's prefix. Agent state may
// shallow-copy the messages array or refresh the model object without changing
// the provider request, so top-level object identity is not a valid cache key.
if (options?.sessionId === sessionManager.getSessionId()) {
cacheWarmer.start({ model, context, options: requestOptions }, cacheContextIsCurrent(model));
}
return runner.emitBeforeProviderRequest(payload);
},
onResponse: async (response, _model) => {
const runner = extensionRunnerRef.current;
if (!runner?.hasHandlers("after_provider_response")) {
return;
}
await runner.emit({
type: "after_provider_response",
status: response.status,
headers: response.headers,
});
return modelRuntime.streamSimple(model, context, requestOptions);
},
onPayload: transformProviderPayload,
onResponse: handleProviderResponse,
sessionId: sessionManager.getSessionId(),
transformContext: async (messages) => {
const runner = extensionRunnerRef.current;
@ -394,12 +423,14 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
resourceLoader,
customTools: options.customTools,
modelRuntime,
cacheWarmer,
initialActiveToolNames,
allowedToolNames,
excludedToolNames,
extensionRunnerRef,
sessionStartEvent: options.sessionStartEvent,
});
const extensionsResult = resourceLoader.getExtensions();
return {

View file

@ -73,6 +73,17 @@ export interface ModelChangeEntry extends SessionEntryBase {
modelId: string;
}
export interface UsageEntry extends SessionEntryBase {
type: "usage";
/** Arbitrary usage category, such as "cache_warm". */
kind: string;
provider: string;
model: string;
usage: Usage;
/** Optional human-readable qualifier for usage notices. */
note?: string;
}
export interface CompactionEntry<T = unknown> extends SessionEntryBase {
type: "compaction";
summary: string;
@ -154,6 +165,7 @@ export type SessionEntry =
| SessionMessageEntry
| ThinkingLevelChangeEntry
| ModelChangeEntry
| UsageEntry
| CompactionEntry
| BranchSummaryEntry
| CustomEntry
@ -1118,6 +1130,23 @@ export class SessionManager {
return entry.id;
}
/** Append model-attributed usage that does not participate in LLM context. Returns the appended entry. */
appendUsage(kind: string, provider: string, model: string, usage: Usage, note?: string): UsageEntry {
const entry: UsageEntry = {
type: "usage",
id: generateId(this.byId),
parentId: this.leafId,
timestamp: new Date().toISOString(),
kind,
provider,
model,
usage,
...(note ? { note } : {}),
};
this._appendEntry(entry);
return entry;
}
/** Append a compaction summary as child of current leaf, then advance leaf. Returns entry id. */
appendCompaction<T = unknown>(
summary: string,

View file

@ -73,6 +73,10 @@ export interface ThinkingBudgetsSettings {
export type MermaidRenderingMode = "off" | "final" | "streaming";
/** Cache-warming profile. "idle" also warms between agent runs. */
export const CACHE_WARMING_MODES = ["off", "streaming", "idle"] as const;
export type CacheWarmingMode = (typeof CACHE_WARMING_MODES)[number];
export interface MarkdownSettings {
codeBlockIndent?: string; // default: " "
mermaid?: MermaidRenderingMode; // default: "streaming"
@ -150,6 +154,7 @@ export interface Settings {
sessionDir?: string; // Custom session storage directory (same format as --session-dir CLI flag)
httpProxy?: string; // Proxy URL applied as HTTP_PROXY and HTTPS_PROXY for Pi-managed HTTP clients
httpIdleTimeoutMs?: number; // HTTP header/body idle timeout in milliseconds; 0 disables it
cacheWarming?: CacheWarmingMode; // default: "streaming"; global only because each refresh costs money
websocketConnectTimeoutMs?: number; // WebSocket connect/open handshake timeout in milliseconds; 0 disables it
tuiMode?: TuiMode; // default: "regular"
fullscreenExitOutput?: FullscreenExitOutput; // default: "transcript"; no effect in regular TUI mode
@ -946,6 +951,18 @@ export class SettingsManager {
this.save();
}
/** Read from global settings only because warming costs money. */
getCacheWarmingMode(): CacheWarmingMode {
const mode = this.globalSettings.cacheWarming;
return mode !== undefined && CACHE_WARMING_MODES.includes(mode) ? mode : "streaming";
}
setCacheWarmingMode(mode: CacheWarmingMode): void {
this.globalSettings.cacheWarming = mode;
this.markModified("cacheWarming");
this.save();
}
getProviderRetrySettings(): { timeoutMs?: number; maxRetries?: number; maxRetryDelayMs: number } {
return {
timeoutMs: this.settings.retry?.provider?.timeoutMs,

View file

@ -33,7 +33,7 @@ export interface UsageCostBreakdownEntry {
tokens: number;
}
/** Group attributable assistant usage by model and all other usage into a separate bucket. */
/** Group model-attributed usage by model and all other usage into a separate bucket. */
export function getUsageCostBreakdown(entries: SessionEntry[]): UsageCostBreakdownEntry[] {
const totalsByKey = new Map<string, UsageTotals>();
@ -43,6 +43,9 @@ export function getUsageCostBreakdown(entries: SessionEntry[]): UsageCostBreakdo
if (entry.type === "message" && entry.message.role === "assistant") {
key = `${entry.message.provider}/${entry.message.responseModel ?? entry.message.model}`;
usage = entry.message.usage;
} else if (entry.type === "usage") {
key = `${entry.provider}/${entry.model}`;
usage = entry.usage;
} else if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.usage) {
key = "Tools/summaries";
usage = entry.message.usage;

View file

@ -24,6 +24,7 @@ export {
type SessionStats,
} from "./core/agent-session.ts";
export { readStoredCredential } from "./core/auth-storage.ts";
export type { CacheWarmingDecision, CacheWarmingStatus } from "./core/cache-warmer.ts";
// Compaction
export {
type BranchPreparation,
@ -66,6 +67,8 @@ export type {
BeforeProviderRequestEvent,
BeforeProviderRequestEventResult,
BuildSystemPromptOptions,
CacheWarmingDecisionEvent,
CacheWarmingDecisionEventResult,
CompactOptions,
ContextEvent,
ContextEventResult,
@ -271,6 +274,7 @@ export {
type ThinkingLevelChangeEntry,
} from "./core/session-manager.ts";
export {
type CacheWarmingMode,
type CompactionModelOverride,
type CompactionSettings,
type DefaultProjectTrust,

View file

@ -89,7 +89,9 @@ export class FooterComponent implements Component {
let latestCacheHitRate: number | undefined;
for (const entry of this.session.sessionManager.getEntries()) {
if (entry.type === "message" && entry.message.role === "assistant") {
if (entry.type === "usage") {
addUsageToTotals(usageTotals, entry.usage);
} else if (entry.type === "message" && entry.message.role === "assistant") {
addUsageToTotals(usageTotals, entry.message.usage);
const latestPromptTokens =

View file

@ -12,12 +12,14 @@ import {
Text,
} from "@earendil-works/pi-tui";
import { formatHttpIdleTimeoutMs, HTTP_IDLE_TIMEOUT_CHOICES } from "../../../core/http-dispatcher.ts";
import type {
DefaultProjectTrust,
FullscreenExitOutput,
MermaidRenderingMode,
TuiMode,
WarningSettings,
import {
CACHE_WARMING_MODES,
type CacheWarmingMode,
type DefaultProjectTrust,
type FullscreenExitOutput,
type MermaidRenderingMode,
type TuiMode,
type WarningSettings,
} from "../../../core/settings-manager.ts";
import { getSettingsListTheme, parseAutoThemeSetting, type TerminalTheme, theme } from "../theme/theme.ts";
import { DynamicBorder } from "./dynamic-border.ts";
@ -60,6 +62,7 @@ export interface SettingsConfig {
followUpMode: "all" | "one-at-a-time";
transport: Transport;
httpIdleTimeoutMs: number;
cacheWarmingMode: CacheWarmingMode;
thinkingLevel: ThinkingLevel;
availableThinkingLevels: ThinkingLevel[];
modelThinkingLevels: Record<string, ThinkingLevel>;
@ -99,6 +102,7 @@ export interface SettingsCallbacks {
onFollowUpModeChange: (mode: "all" | "one-at-a-time") => void;
onTransportChange: (transport: Transport) => void;
onHttpIdleTimeoutMsChange: (timeoutMs: number) => void;
onCacheWarmingModeChange: (mode: CacheWarmingMode) => void;
onModelThinkingLevelChange: (provider: string, modelId: string, level: ThinkingLevel) => void;
onModelThinkingLevelRemove: (provider: string, modelId: string) => void;
onThemeChange: (theme: string) => void;
@ -495,6 +499,14 @@ export class SettingsSelectorComponent extends Container {
currentValue: formatHttpIdleTimeoutMs(config.httpIdleTimeoutMs),
values: HTTP_IDLE_TIMEOUT_CHOICES.map((choice) => choice.label),
},
{
id: "cache-warming-mode",
label: "Cache warming",
description:
"off; streaming while the agent runs; idle also between runs while continuation stays profitable",
currentValue: config.cacheWarmingMode,
values: [...CACHE_WARMING_MODES],
},
{
id: "hide-thinking",
label: "Hide thinking",
@ -863,6 +875,9 @@ export class SettingsSelectorComponent extends Container {
}
break;
}
case "cache-warming-mode":
callbacks.onCacheWarmingModeChange(newValue as CacheWarmingMode);
break;
case "hide-thinking":
callbacks.onHideThinkingBlockChange(newValue === "true");
break;

View file

@ -338,6 +338,7 @@ class TreeList implements Component {
this.filteredNodes = this.flatNodes.filter((flatNode) => {
const entry = flatNode.node.entry;
if (entry.type === "usage") return false;
const isCurrentLeaf = entry.id === this.currentLeafId;
// Skip assistant messages with only tool calls (no text) unless error/aborted

View file

@ -65,6 +65,7 @@ import {
computeCacheWaste,
detectCacheMiss,
} from "../../core/cache-stats.ts";
import { formatCacheWarmingStatus, formatCacheWarmingUsage } from "../../core/cache-warmer.ts";
import { recordCrash, takeUnnotifiedCrash } from "../../core/crash-log.ts";
import { DEFAULT_THINKING_LEVEL, THINKING_LEVEL_OPTIONS } from "../../core/defaults.ts";
import type {
@ -94,7 +95,12 @@ import { CredentialSynchronizationError } from "../../core/model-runtime.ts";
import { DefaultPackageManager } from "../../core/package-manager.ts";
import type { ResourceDiagnostic } from "../../core/resource-loader.ts";
import { formatMissingSessionCwdPrompt, MissingSessionCwdError } from "../../core/session-cwd.ts";
import { type SessionEntry, SessionManager, sessionEntryToContextMessages } from "../../core/session-manager.ts";
import {
type SessionEntry,
SessionManager,
sessionEntryToContextMessages,
type UsageEntry,
} from "../../core/session-manager.ts";
import type { FullscreenExitOutput, TuiMode } from "../../core/settings-manager.ts";
import { BUILTIN_SLASH_COMMANDS } from "../../core/slash-commands.ts";
import type { SourceInfo } from "../../core/source-info.ts";
@ -234,7 +240,7 @@ type CompactionCostNotice = {
usage: Usage;
};
type RenderSessionItem = AgentMessage | Extract<SessionEntry, { type: "custom" }> | CompactionCostNotice;
type RenderSessionItem = AgentMessage | Extract<SessionEntry, { type: "custom" | "usage" }> | CompactionCostNotice;
function isCustomSessionEntry(item: RenderSessionItem): item is Extract<SessionEntry, { type: "custom" }> {
return "type" in item && item.type === "custom";
@ -244,6 +250,10 @@ function isCompactionCostNotice(item: RenderSessionItem): item is CompactionCost
return "type" in item && item.type === "compaction_cost";
}
function isUsageSessionEntry(item: RenderSessionItem): item is Extract<SessionEntry, { type: "usage" }> {
return "type" in item && item.type === "usage";
}
const DEAD_TERMINAL_ERROR_CODES = new Set(["EIO", "EPIPE", "ENOTCONN"]);
function isDeadTerminalError(error: unknown): boolean {
@ -3264,6 +3274,9 @@ export class InteractiveMode {
if (event.entry.type === "custom") {
this.addCustomEntryToChat(event.entry);
this.ui.requestRender();
} else if (event.entry.type === "usage" && event.entry.kind === "cache_warm") {
this.addCacheWarmingUsage(event.entry);
this.ui.requestRender();
}
break;
@ -3754,8 +3767,8 @@ export class InteractiveMode {
): void {
this.pendingTools.clear();
const renderedPendingTools = new Map<string, ToolExecutionComponent>();
// Cache-miss notices are not persisted; re-derive them from the full entry
// list and re-inject them after the assistant messages that paid for them.
// Cache misses are not persisted, unlike successful cache-warming usage.
// Re-derive them and inject them after the assistant messages that paid for them.
const cacheMisses = this.settingsManager.getShowCacheMissNotices()
? collectCacheMisses(this.sessionManager.getEntries(), this.session.modelRuntime)
: new Map<AssistantMessage, CacheMiss>();
@ -3770,6 +3783,10 @@ export class InteractiveMode {
this.addCustomEntryToChat(item);
continue;
}
if (isUsageSessionEntry(item)) {
this.addCacheWarmingUsage(item);
continue;
}
if (isCompactionCostNotice(item)) {
this.addCompactionCostNotice(item);
continue;
@ -3848,7 +3865,7 @@ export class InteractiveMode {
options: { updateFooter?: boolean; populateHistory?: boolean } = {},
): void {
const items = entries.flatMap((entry): RenderSessionItem[] => {
if (entry.type === "custom") {
if (entry.type === "custom" || (entry.type === "usage" && entry.kind === "cache_warm")) {
return [entry];
}
const messages = sessionEntryToContextMessages(entry);
@ -3860,6 +3877,12 @@ export class InteractiveMode {
this.renderSessionItems(items, options);
}
private addCacheWarmingUsage(entry: UsageEntry): void {
if (!this.settingsManager.getShowCacheMissNotices()) return;
this.chatContainer.addChild(new Spacer(1));
this.chatContainer.addChild(new Text(theme.fg("dim", formatCacheWarmingUsage(entry)), 1, 0));
}
/**
* Render billing usage for a compaction or branch summary. The notice is derived
* from persisted summary usage and is not stored as a separate session entry.
@ -4655,6 +4678,7 @@ export class InteractiveMode {
followUpMode: this.session.followUpMode,
transport: this.settingsManager.getTransport(),
httpIdleTimeoutMs: this.settingsManager.getHttpIdleTimeoutMs(),
cacheWarmingMode: this.settingsManager.getCacheWarmingMode(),
thinkingLevel: this.settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL,
availableThinkingLevels: [...THINKING_LEVEL_OPTIONS],
modelThinkingLevels: this.settingsManager.getAllModelThinkingLevels(),
@ -4728,6 +4752,10 @@ export class InteractiveMode {
configureHttpDispatcher(timeoutMs);
this.showStatus(`HTTP idle timeout: ${formatHttpIdleTimeoutMs(timeoutMs)}`);
},
onCacheWarmingModeChange: (mode) => {
this.session.setCacheWarmingMode(mode);
this.showStatus(`Cache warming: ${mode}`);
},
onModelThinkingLevelChange: (provider, modelId, level) => {
this.settingsManager.setModelThinkingLevel(provider, modelId, level);
// If the override is for the current model, apply it to the session too
@ -6351,6 +6379,16 @@ export class InteractiveMode {
info += `${theme.fg("dim", "Output:")} ${stats.tokens.output.toLocaleString()}\n`;
info += `${theme.fg("dim", "Total:")} ${stats.tokens.total.toLocaleString()}\n`;
const cacheWarmingStatus = this.session.cacheWarmingStatus;
info += `\n${theme.bold("Cache Warming")}\n`;
info += `${theme.fg("dim", "Mode:")} ${this.settingsManager.getCacheWarmingMode()}\n`;
info += `${theme.fg("dim", "Status:")} ${cacheWarmingStatus ? formatCacheWarmingStatus(cacheWarmingStatus) : "Inactive (cache warming unavailable)"}\n`;
const decision = cacheWarmingStatus?.decision;
if (decision?.economicsAvailable) {
info += `${theme.fg("dim", "Cache miss penalty:")} $${decision.missCost.toFixed(3)}\n`;
info += `${theme.fg("dim", "Refresh cost:")} $${decision.warmCost.toFixed(3)}\n`;
}
if (stats.cost > 0 || cacheWaste.missedTokens > 0) {
info += `\n${theme.bold("Cost")}\n`;
info += `${theme.fg("dim", "Total:")} $${stats.cost.toFixed(3)}`;

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;