diff --git a/apps/p2/README.md b/apps/p2/README.md index 72dbf311d..5dedac57e 100644 --- a/apps/p2/README.md +++ b/apps/p2/README.md @@ -188,6 +188,13 @@ restart, and in-flight requests keep their configuration. Set the default model **Make default for new sessions** in the model picker. `ownerName` (1 to 40 characters, no control characters, `<` or `>`; default "Owner") is the name your messages carry from then on; messages already written keep theirs. +`cacheWarming` (`"idle"` by default, `"streaming"` or `"off"`) keeps prompt caches warm as pi does: before the cache +entry of a conversation's last request expires (5 minutes for Anthropic and Claude on Radius), p2 sends the request +again with a one-token answer, while that is cheaper than writing the cache again. `"streaming"` warms while a run +streams (up to an hour, e.g. during a long tool call); `"idle"` also warms the main conversation for up to 30 minutes +after its run answered. The refreshes count in the conversation's cost. `stream.cacheRetention: "long"` writes 1-hour +cache entries instead (at twice the input price). + ### Voice mode Voice mode talks to the agent hands-free: the microphone button in the composer (owner only) streams your speech to a diff --git a/apps/p2/docs/handoff.md b/apps/p2/docs/handoff.md index 230f9be81..443b29bc0 100644 --- a/apps/p2/docs/handoff.md +++ b/apps/p2/docs/handoff.md @@ -95,7 +95,7 @@ verifying an agent's files on their own before committing. Node 24: `~/node-v24. `tui*.ts` (launcher, sandbox, filter, grants), `dev*.ts` (dev supervisor). - **Runtime** (`src/runtime`): shared contracts and helpers (principals, worker and Harness, Durable registry, extension switches, input entries, settings, logs, fetch). -- **Plugins**: chat, chat-mode, sessions, subagents, agent (one-shot agents), codemode, coding-tools, prompt, +- **Plugins**: chat, chat-mode, sessions, subagents, agent (one-shot agents), codemode, cache-warming, coding-tools, prompt, models, login, search, status, drafts, files, send-file, artifacts, workspace, file-index, terminals, envs, mcp, ports, radius, sharing, devices, proxy-marker, logs, dev-tools, powershell, web-shell. Before writing one, read Chord's tutorial (`packages/chord/docs/tutorial/`) and plugin guide (`packages/chord/docs/plugin-authoring.md`), then an diff --git a/apps/p2/docs/p2-notes.md b/apps/p2/docs/p2-notes.md index ef4b93d55..4bf9d62b7 100644 --- a/apps/p2/docs/p2-notes.md +++ b/apps/p2/docs/p2-notes.md @@ -1180,6 +1180,24 @@ translated phone drawer and intercept chat taps. reload during a script, SIGKILL and graceful restarts mid-call (finished calls do not rerun; the interrupted unsafe `bash` call returns an error the script catches); the removal mid-script (prompt error, turn ends, the script's thread goes); `codemode-plugin.test.ts` the sandbox across reload/removal. +- **Cache warming** (`plugins/cache-warming`): pi-coding-agent's cache warmer for every conversation of a session. + The worker facet registers a hook-only Durable extension (`cache-warming`): `afterResponse` gets the request the + attempt sent (`GenerationRequest`: model, messages after `beforeRequest`, options without the attempt's signal) and + starts that conversation's run, replacing the previous one; `afterRun` turns it idle. Before the entry expires (90% + of `model.promptCache[retention]`: 270 s for the 5-minute tier) the run sends the request again with + `getCacheWarmingRequestOptions` (one output token, no retries, `disableContinuation`) and adds the spend to the + conversation's `pi.usage` (`recordUsage`, the model's key). The policy is pi-ai's (`evaluateCacheWarming`, shared + with pi-coding-agent): streaming up to an hour after the request, idle up to 30 minutes, a refresh only when its + expected savings reach $0.05 (continuation probability 1 while streaming, 0.15 idle). p2's additions: only + top-level conversations (`isTopLevel`) are warmed idle, since a settled subagent or one-shot agent rarely continues; + a run that ended without `afterRun` (an abort during tools, a turn limit) turns idle at the next refresh (`pi.live.run` + gone); a refresh first checks the conversation still sends the same prefix (same model, a context that extends the + sent messages by role and timestamp), so a reset, compaction or model change stops it; requests whose body a + `beforeProviderRequest` hook replaced (`payloadReplaced`) are not warmed. Timers and refreshes are a native of + `CacheWarmingRuntime`, closed by a reload or stop (nothing is durable: the next response starts again); while a + run is armed the worker's liveness counts it (`cacheWarming`), so the idle stop (`P2_IDLE_MS`) waits. Setting: + `cacheWarming` (`"idle"` default, `"streaming"`, `"off"`, switched live through `ExtensionSwitches`). No status UI + and no decision hook yet (pi's `/session` line, `cache_warming_decision`). `test/p2/cache-warming.test.ts`. - **MCP servers** (`plugins/mcp`, plan `docs/mcp-durable.md`, WP6-8): the server's `McpManager` (local, host lifetime) owns every MCP process and connection; workers never spawn one or hold a credential. `/mcp.json` is the owner's `mcpServers` (pi-mcp's `validateMcpServerConfig`, plus p2's `shared` and `runOn`; invalid entries are diff --git a/apps/p2/docs/plugin-hooks.md b/apps/p2/docs/plugin-hooks.md index 6ed7af76f..a037eb695 100644 --- a/apps/p2/docs/plugin-hooks.md +++ b/apps/p2/docs/plugin-hooks.md @@ -155,7 +155,7 @@ through `durableRegistry` (`src/runtime/durable.ts:41`). | --- | --- | --- | --- | | Tools | `tools` (`defineTool`, `types.ts:293`) | `parameters`, `execute(args, api, ctx)` returning `output`/`structuredOutput`/`details`; `replay: "safe" \| "unsafe"` (rerun after a crash, `:298`); `callers: ["model" \| "tools"]` (`:317`); `structuredOutputSchema` (`:311`); `deferred` (`:322`); `executionMode`, `prepareArguments`, `outputLimits`; `api.executeTool` for nested calls (`:281`) | coding-tools, artifacts, send-file, ports, conversations, sessions, subagents, dev-tools, codemode, powershell; `agent` is `callers: ["tools"]`, replay-safe with a schema (`plugins/agent/tool.ts:53-57`) | | Prompt | `sections` (`section(key, render, {tag})`, `types.ts:349`) | rendered before each request, in extension order | prompt (`plugins/prompt/worker.ts:185-191`), coding-tools, ports, conversations, dev-tools | -| Hooks | `hooks: [hook(Task, handlers)]` | `GenerationTask`: `beforeRequest`, `afterResponse`, `onYield`, `afterTools` (`types.ts:757`); `ToolTask`: `beforeTool`, `afterTool` (`:778`); `CompactionTask`: `beforeCompact` (`:795`); handlers get `HookApi` (doc reads, memos, models, `:745`) | none yet | +| Hooks | `hooks: [hook(Task, handlers)]` | `GenerationTask`: `beforeRequest`, `afterResponse`, `onYield`, `afterTools` (`types.ts:757`); `ToolTask`: `beforeTool`, `afterTool` (`:778`); `CompactionTask`: `beforeCompact` (`:795`); handlers get `HookApi` (doc reads, memos, models, `:745`) | cache-warming (`afterResponse` with the sent request, `afterRun`) | | Wraps | `wraps: [wrapTool(tool, fn) \| wrapSection(section, fn)]` (`types.ts:360`) | replace another extension's tool or section where selected | none yet | | Tasks | `tasks: [defineTask(…)]` | durable background work, resolved by name | subagents (`Anchor`, `Reporter`, `plugins/subagents/tools.ts:72,118,473`) | @@ -404,7 +404,7 @@ Reference: `packages/coding-agent/docs/extensions.md`. | pi 1.0 | p2 today | Gap | | --- | --- | --- | -| `pi.on()` lifecycle: `session_start`/`shutdown`, `before_agent_start`, `context`, `message_end`, `tool_call`/`tool_result`, `turn_end`, `agent_before_settle`, `session_before_compact`, `user_bash`, `input`, `provider_stream_event`, `cache_warming_decision` | Chord resources for start/shutdown; Durable hooks cover `context` (`beforeRequest`), `message_end` (notify-only `afterResponse`), `tool_call`/`tool_result`, `turn_end` (`afterTools`), settle (`onYield`), compaction; `SendRouters` covers `Chat.send` input only | No input hook for peer/subagent/agent inputs; no `user_bash` (shell is hard-coded in chat); no stream events; no whole-prompt override (sections and `wrapSection` instead). No plugin uses hooks yet, so they are untested in p2. | +| `pi.on()` lifecycle: `session_start`/`shutdown`, `before_agent_start`, `context`, `message_end`, `tool_call`/`tool_result`, `turn_end`, `agent_before_settle`, `session_before_compact`, `user_bash`, `input`, `provider_stream_event`, `cache_warming_decision` | Chord resources for start/shutdown; Durable hooks cover `context` (`beforeRequest`), `message_end` (notify-only `afterResponse`), `tool_call`/`tool_result`, `turn_end` (`afterTools`), settle (`onYield`), compaction; `SendRouters` covers `Chat.send` input only | No input hook for peer/subagent/agent inputs; no `user_bash` (shell is hard-coded in chat); no stream events; no whole-prompt override (sections and `wrapSection` instead). Only cache-warming uses hooks (`afterResponse`, `afterRun`); pi's `cache_warming_decision` has no p2 counterpart. | | `registerTool` with `exposure`, `namespace`, `annotations`, `prepareLoadout`, `reachesTools` | `callers`, `deferred`, replay policy | No annotations (no approval policy), no namespaces; `tool_search` not installed. | | `registerCommand`, `registerShortcut` | TUI only | No web commands or shortcuts. | | `registerFlag` | none | Plugins get configuration only via env vars or their own files; `settings.json` is closed. | diff --git a/apps/p2/plugins/cache-warming/contract.ts b/apps/p2/plugins/cache-warming/contract.ts new file mode 100644 index 000000000..deea9e50e --- /dev/null +++ b/apps/p2/plugins/cache-warming/contract.ts @@ -0,0 +1,8 @@ +// The worker's cache warmer (host-internal): pi's prompt cache warming for every conversation of the session. +import { local } from "@p2/app/v1"; + +/** + * Keeps the prompt cache entries of the session's conversations warm (`./warmer.ts`): one per worker, started with + * it; a reload or a stop closes its timers and in-flight refreshes, and the next real request starts warming again. + */ +export const CacheWarmingRuntime = local("p2.cache-warming-runtime", { version: 1, params: {}, members: {} }); diff --git a/apps/p2/plugins/cache-warming/package.json b/apps/p2/plugins/cache-warming/package.json new file mode 100644 index 000000000..7f1273ff6 --- /dev/null +++ b/apps/p2/plugins/cache-warming/package.json @@ -0,0 +1,5 @@ +{ + "name": "cache-warming", + "private": true, + "type": "module" +} diff --git a/apps/p2/plugins/cache-warming/warmer.ts b/apps/p2/plugins/cache-warming/warmer.ts new file mode 100644 index 000000000..9dbe34f45 --- /dev/null +++ b/apps/p2/plugins/cache-warming/warmer.ts @@ -0,0 +1,186 @@ +// Prompt cache warming for the session's conversations, as pi-coding-agent does it for its one session: before the +// cache entry a conversation's last request wrote expires, send that request again with a one-token output cap, for as +// long as that is expected to pay off. The policy (when, whether, how) is pi-ai's; this file schedules it per +// conversation. +import { + type AssistantMessage, + CACHE_WARMING_MAX_AGE_MS, + CACHE_WARMING_MAX_IDLE_AGE_MS, + type CacheWarmingPhase, + evaluateCacheWarming, + getCacheWarmingDelayMs, + getCacheWarmingRequestOptions, + getPromptCacheTtlMs, + isCacheWarmingReplayable, + type Message, + type Models, +} from "@earendil-works/pi-ai"; +import type { ConversationId, GenerationRequest, GenerationResponse } from "@earendil-works/pi-durable"; + +/** The `cacheWarming` setting: "idle" also warms after a run settled, "streaming" only while it runs. */ +export type CacheWarmingMode = "off" | "streaming" | "idle"; + +/** What the warmers read and do outside themselves; the worker facet binds it to the Harness and the settings. */ +export type WarmerHost = { + /** The `cacheWarming` setting now. */ + mode(): CacheWarmingMode; + now(): number; + /** Arms a timer and returns its cancel. */ + setTimer(fn: () => void, delayMs: number): () => void; + /** Whether the conversation's next request still starts with `request`: same model, a context extending its messages. */ + isCurrent(conversationId: ConversationId, request: GenerationRequest): Promise; + /** Whether a run is under way in the conversation. */ + running(conversationId: ConversationId): Promise; + /** Whether the conversation is warmed after its run settled: top-level conversations, not subagents. */ + warmsIdle(conversationId: ConversationId): Promise; + /** Adds a refresh's spend to the conversation's usage. */ + record(conversationId: ConversationId, message: AssistantMessage): Promise; + /** Counts the conversations being warmed; a count keeps the worker process alive. */ + keepAlive(delta: number): void; + report(error: unknown): void; +}; + +type Run = { + readonly request: GenerationRequest; + readonly models: Models; + /** Prompt size of the request, as the provider reported it. */ + readonly promptTokens: number; + readonly ttlMs: number; + readonly delayMs: number; + readonly startedAt: number; + readonly controller: AbortController; + phase: CacheWarmingPhase; + /** Latest time the armed refresh may still be sent, leaving half the planned margin before expiry. */ + refreshDeadlineAt: number; + cancelTimer?: () => void; +}; + +/** + * One warming run per conversation, started by each response (`afterResponse`) and replaced by the next one. A run + * refreshes while its conversation's run streams (up to an hour after the request) and, in mode "idle", for up to 30 + * minutes after it settled, while each refresh pays off by pi-ai's `evaluateCacheWarming()`. + */ +export class ConversationWarmers { + readonly #host: WarmerHost; + readonly #runs = new Map(); + + constructor(host: WarmerHost) { + this.#host = host; + } + + /** Warm the cache entry `response`'s request wrote, replacing the conversation's run; `models` sends the refreshes. */ + start(conversationId: ConversationId, response: GenerationResponse, models: Models): void { + this.stop(conversationId); + const { message, request } = response; + // A polled deferred response sent no request; a failed or aborted one wrote no entry worth keeping, and its run + // is over. + if (request === undefined || message.stopReason === "error" || message.stopReason === "aborted") return; + if (this.#host.mode() === "off" || request.payloadReplaced) return; + if (!isCacheWarmingReplayable(request.model, request.options)) return; + const ttlMs = getPromptCacheTtlMs(request.model, request.options); + const delayMs = ttlMs === undefined ? undefined : getCacheWarmingDelayMs(ttlMs); + if (ttlMs === undefined || delayMs === undefined) return; + const usage = message.usage; + const run: Run = { + request, + models, + promptTokens: usage.input + usage.cacheRead + usage.cacheWrite, + ttlMs, + delayMs, + startedAt: this.#host.now(), + controller: new AbortController(), + phase: "streaming", + refreshDeadlineAt: 0, + }; + this.#runs.set(conversationId, run); + this.#host.keepAlive(1); + this.#schedule(conversationId, run); + } + + /** The conversation's run answered (`afterRun`): warm on as idle, or stop in mode "streaming". */ + settled(conversationId: ConversationId): void { + const run = this.#runs.get(conversationId); + if (run === undefined) return; + if (this.#host.mode() === "idle") run.phase = "idle"; + else this.stop(conversationId); + } + + /** Stop warming the conversation; a refresh in flight is aborted. */ + stop(conversationId: ConversationId): void { + const run = this.#runs.get(conversationId); + if (run === undefined) return; + this.#runs.delete(conversationId); + run.cancelTimer?.(); + run.controller.abort(); + this.#host.keepAlive(-1); + } + + /** Stop every run (a reload or the worker's stop). */ + close(): void { + for (const conversationId of [...this.#runs.keys()]) this.stop(conversationId); + } + + #schedule(conversationId: ConversationId, run: Run): void { + const now = this.#host.now(); + const nextAt = now + run.delayMs; + const maxAge = run.phase === "idle" ? CACHE_WARMING_MAX_IDLE_AGE_MS : CACHE_WARMING_MAX_AGE_MS; + if (nextAt > run.startedAt + maxAge) { + this.stop(conversationId); + return; + } + // A timer can fire late after sleep or a blocked event loop; a refresh past this point would likely be a + // full-price cache write, not a warm. + run.refreshDeadlineAt = nextAt + Math.floor((run.ttlMs - run.delayMs) / 2); + run.cancelTimer = this.#host.setTimer(() => void this.#refresh(conversationId, run), run.delayMs); + } + + async #refresh(conversationId: ConversationId, run: Run): Promise { + run.cancelTimer = undefined; + try { + if (!(await this.#shouldRefresh(conversationId, run))) { + if (this.#runs.get(conversationId) === run) this.stop(conversationId); + return; + } + const { model, messages, options } = run.request; + const message = await run.models.completeSimple( + model, + { messages: [...messages] }, + getCacheWarmingRequestOptions(options, run.controller.signal), + ); + if (this.#runs.get(conversationId) !== run) return; + if (message.stopReason !== "error" && message.stopReason !== "aborted") + await this.#host.record(conversationId, message); + } catch (error) { + // Warming is best effort: a failed refresh never touches the conversation, the next one may succeed. + if (!run.controller.signal.aborted) this.#host.report(error); + } + if (this.#runs.get(conversationId) === run) this.#schedule(conversationId, run); + } + + /** Whether the armed refresh is still wanted: the run current, in its window, and paying off. */ + async #shouldRefresh(conversationId: ConversationId, run: Run): Promise { + const mode = this.#host.mode(); + if (mode === "off" || this.#runs.get(conversationId) !== run) return false; + if (this.#host.now() > run.refreshDeadlineAt) return false; + // A run that ended without an answer (aborted during its tools, a turn limit) never called `afterRun`. + if (run.phase === "streaming" && !(await this.#host.running(conversationId))) run.phase = "idle"; + if (run.phase === "idle") { + if (mode !== "idle" || !(await this.#host.warmsIdle(conversationId))) return false; + if (this.#host.now() > run.startedAt + CACHE_WARMING_MAX_IDLE_AGE_MS) return false; + } + if (!(await this.#host.isCurrent(conversationId, run.request))) return false; + const decision = evaluateCacheWarming(run.request.model, run.promptTokens, run.phase); + return decision.action === "warm" && this.#runs.get(conversationId) === run; + } +} + +/** + * Whether `current` (a conversation's model context now) starts with `sent` (the messages a request sent): the same + * messages by role and timestamp, so no reset, compaction or edit replaced any of them. + */ +export function extendsMessages(current: readonly Message[], sent: readonly Message[]): boolean { + if (current.length < sent.length) return false; + return sent.every( + (message, index) => current[index]!.role === message.role && current[index]!.timestamp === message.timestamp, + ); +} diff --git a/apps/p2/plugins/cache-warming/worker.ts b/apps/p2/plugins/cache-warming/worker.ts new file mode 100644 index 000000000..45ba17788 --- /dev/null +++ b/apps/p2/plugins/cache-warming/worker.ts @@ -0,0 +1,78 @@ +// Cache warming in the session worker: a Durable extension whose generation hooks start, idle and stop the warming of +// each conversation's last request (`./warmer.ts`), bound to the Harness, the settings and the worker's liveness. +import { defineExtension, GenerationTask, hook, LiveDoc, recordUsage } from "@earendil-works/pi-durable"; +import { BACKGROUND_CONTEXT, defineFacet } from "@p2/app/v1"; +import { ExtensionSwitches } from "@p2/app/v1/node"; +import { Extensions, harnessOf, isTopLevel, workerOf } from "@p2/app/v1/worker"; +import { WorkerModels } from "../models/contract.ts"; +import { CacheWarmingRuntime } from "./contract.ts"; +import { ConversationWarmers, extendsMessages } from "./warmer.ts"; + +export default defineFacet((host) => { + host.serve(CacheWarmingRuntime, (_params, self) => { + const harness = harnessOf(self); + const worker = workerOf(self); + const settings = self.get(WorkerModels, {}); + const durable = () => harness().durable; + // Timers and in-flight refreshes live in this native: a reload or the worker's stop ends them, and the next + // response starts warming again. Nothing of it is durable: the requests it repeats are not stored. + const warmers = self.open( + async () => + new ConversationWarmers({ + mode: () => settings.settings.value.cacheWarming ?? "idle", + now: () => Date.now(), + setTimer: (fn, delayMs) => { + const timer = setTimeout(fn, delayMs); + return () => clearTimeout(timer); + }, + isCurrent: async (conversationId, request) => { + const conversation = await durable().conversation(conversationId, BACKGROUND_CONTEXT); + if (conversation === undefined) return false; + const { model } = await conversation.agent(BACKGROUND_CONTEXT); + if (model?.provider !== request.model.provider || model.modelId !== request.model.id) return false; + return extendsMessages((await conversation.context(BACKGROUND_CONTEXT)).messages, request.messages); + }, + running: async (conversationId) => + (await durable().snapshot(LiveDoc, conversationId, BACKGROUND_CONTEXT))?.run !== undefined, + // After a subagent's or one-shot agent's run its work is usually done; the main conversation and its + // `/tree` branches wait for the person or for agents' reports. + warmsIdle: (conversationId) => isTopLevel(durable(), conversationId, BACKGROUND_CONTEXT), + record: async (conversationId, message) => { + await durable().commit( + (tx) => + recordUsage( + tx, + conversationId, + "models", + `${message.provider}/${message.model}`, + message.usage, + ), + BACKGROUND_CONTEXT, + ); + }, + keepAlive: (delta) => worker().liveness.add("cacheWarming", delta), + report: (error) => console.error("p2 cache-warming: refreshing the prompt cache failed:", error), + }), + (opened) => opened.close(), + ); + const extension = defineExtension({ + name: "cache-warming", + hooks: [ + hook(GenerationTask, { + afterResponse: (response, api) => warmers.value.start(api.conversationId, response, api.models), + afterRun: (_run, api) => { + warmers.value.settled(api.conversationId); + return undefined; + }, + }), + ], + }); + // `cacheWarming: "off"` uninstalls the hooks live; runs already warming stop at their next refresh. + self.contribute( + ExtensionSwitches, + self.derive([settings.settings], ([value]) => ({ [extension.name]: value.cacheWarming !== "off" })), + ); + self.contribute(Extensions, { [extension.name]: extension }); + return {}; + }); +}); diff --git a/apps/p2/src/runtime/settings.ts b/apps/p2/src/runtime/settings.ts index 09e89e723..17376aed2 100644 --- a/apps/p2/src/runtime/settings.ts +++ b/apps/p2/src/runtime/settings.ts @@ -47,6 +47,12 @@ const schema = obj({ * Default on in p2 (Pi Mobile: off); `false` uninstalls it live from every worker. */ codemode: opt(bool()), + /** + * pi's `cacheWarming`, kept by the cache-warming plugin: "idle" (default) keeps a conversation's prompt cache entry + * warm while its run streams and for up to 30 minutes after it settled, "streaming" only while it streams, "off" + * never. + */ + cacheWarming: opt(StringEnum(["off", "streaming", "idle"] as const)), /** * p2 addition: the owner's name on the entries the owner writes from now on (default "Owner"); entries already * written keep the name they were written with. No control characters, `<` or `>`. diff --git a/apps/p2/src/runtime/worker.ts b/apps/p2/src/runtime/worker.ts index 0382db37c..354ed8376 100644 --- a/apps/p2/src/runtime/worker.ts +++ b/apps/p2/src/runtime/worker.ts @@ -181,13 +181,13 @@ export const HarnessEnvironments = keyedRegistry("p2.ha /** * What keeps the worker process alive besides held references (Chord's own count): its Durable tasks, calls under way - * (`call`), viewers plugins count themselves, and the relocation fence. A JSON watchable of those counts - * (`workerMain`'s `busy`): the worker is idle when every count is zero. + * (`call`), viewers plugins count themselves, conversations whose prompt cache is kept warm, and the relocation fence. + * A JSON watchable of those counts (`workerMain`'s `busy`): the worker is idle when every count is zero. */ export class Liveness implements Watchable>> { #relocating = false; #stopping = false; - #counts: Readonly> = { activeCalls: 0, viewers: 0, durableTasks: 0 }; + #counts: Readonly> = { activeCalls: 0, viewers: 0, durableTasks: 0, cacheWarming: 0 }; readonly #listeners = new Set<() => void>(); get value(): Readonly> { return this.#counts; @@ -225,7 +225,7 @@ export class Liveness implements Watchable>> { this.#counts = { ...this.#counts, [kind]: count }; for (const listener of [...this.#listeners]) listener(); } - add(kind: "activeCalls" | "viewers", delta: number): void { + add(kind: "activeCalls" | "viewers" | "cacheWarming", delta: number): void { this.#set(kind, (this.#counts[kind] ?? 0) + delta); } setDurableTasks(count: number): void { diff --git a/apps/p2/test/p2/cache-warming.test.ts b/apps/p2/test/p2/cache-warming.test.ts new file mode 100644 index 000000000..9c92d7b77 --- /dev/null +++ b/apps/p2/test/p2/cache-warming.test.ts @@ -0,0 +1,357 @@ +// Cache warming (plugins/cache-warming): the warmers re-send a conversation's last request before its prompt cache +// entry expires, with one output token, for as long as pi's policy says it pays off, and stop when the request no +// longer matches what the conversation sends next. Driven with a fake clock, timers and models, then end to end in a +// worker with the fake provider given a short cache lifetime. + +import assert from "node:assert/strict"; +import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { test } from "node:test"; +import type { + AssistantMessage, + Context, + Message, + Model, + Models, + ModelsSimpleStreamOptions, + StopReason, +} from "@earendil-works/pi-ai"; +import { fauxAssistantMessage } from "@earendil-works/pi-ai/providers/faux"; +import type { + ConversationId, + ConversationView, + GenerationRequest, + LiveState, + UsageState, +} from "@earendil-works/pi-durable"; +import { BACKGROUND_CONTEXT } from "@p2/app/v1"; +import { type CacheWarmingMode, ConversationWarmers, extendsMessages } from "../../plugins/cache-warming/warmer.ts"; +import { Chat } from "../../plugins/chat/contract.ts"; +import { Sessions } from "../../plugins/sessions/contract.ts"; +import { validateSettings } from "../../src/runtime/settings.ts"; +import { holdIn, startFixture, until } from "./helpers.ts"; + +const MODEL: Model<"pi-messages"> = { + id: "claude-opus-5-5", + name: "Claude Opus 5.5", + api: "pi-messages", + provider: "radius", + baseUrl: "https://radius.example/v1", + reasoning: true, + input: ["text"], + cost: { input: 4, output: 20, cacheRead: 0.2, cacheWrite: 5 }, + contextWindow: 1_000_000, + maxTokens: 128_000, + promptCache: { short: 300, long: 3600 }, +}; +const CONVERSATION = 1 as ConversationId; +const SUBAGENT = 7 as ConversationId; +const MESSAGES: Message[] = [ + { role: "user", content: "hello", timestamp: 1 }, + { ...fauxAssistantMessage([]), timestamp: 2 }, + { role: "user", content: "go on", timestamp: 3 }, +]; + +function request(overrides: Partial = {}): GenerationRequest { + return { + model: MODEL, + messages: MESSAGES, + options: { sessionId: "provider-session", reasoning: "high" }, + payloadReplaced: false, + ...overrides, + }; +} + +/** A response whose prompt had `promptTokens` tokens, mostly read from the cache. */ +function response(promptTokens: number, stopReason: StopReason = "toolUse"): AssistantMessage { + const message = fauxAssistantMessage([]); + return { + ...message, + stopReason, + usage: { ...message.usage, input: 2, cacheRead: promptTokens - 1002, cacheWrite: 1000 }, + }; +} + +/** A fake host: a manual clock and timers, models that record refreshes, and the conversation state as variables. */ +function setup(mode: CacheWarmingMode = "idle") { + const state = { + mode, + now: 0, + running: new Set([CONVERSATION, SUBAGENT]), + current: true, + alive: 0, + timers: [] as { at: number; fn: () => void; cancelled: boolean }[], + sent: [] as { model: Model<"pi-messages">; context: Context; options: ModelsSimpleStreamOptions }[], + recorded: [] as ConversationId[], + reports: [] as unknown[], + /** Resolves the refresh in flight, when `hold` is set. */ + release: undefined as (() => void) | undefined, + hold: false, + }; + const models = { + completeSimple: async (model: Model<"pi-messages">, context: Context, options: ModelsSimpleStreamOptions) => { + state.sent.push({ model, context, options }); + if (state.hold) + await new Promise((resolve) => { + state.release = resolve; + }); + const reply = fauxAssistantMessage([]); + return { ...reply, stopReason: options.signal?.aborted ? "aborted" : "length" }; + }, + } as unknown as Models; + const warmers = new ConversationWarmers({ + mode: () => state.mode, + now: () => state.now, + setTimer: (fn, delayMs) => { + const timer = { at: state.now + delayMs, fn, cancelled: false }; + state.timers.push(timer); + return () => { + timer.cancelled = true; + }; + }, + isCurrent: async () => state.current, + running: async (conversationId) => state.running.has(conversationId), + warmsIdle: async (conversationId) => conversationId === CONVERSATION, + record: async (conversationId) => void state.recorded.push(conversationId), + keepAlive: (delta) => { + state.alive += delta; + }, + report: (error) => void state.reports.push(error), + }); + /** Moves the clock to the next armed timer and runs it; false when none is armed. */ + const tick = async (): Promise => { + const next = state.timers.filter((timer) => !timer.cancelled).sort((a, b) => a.at - b.at)[0]; + if (next === undefined) return false; + next.cancelled = true; + state.now = next.at; + next.fn(); + // The refresh awaits the host and the models a few times. + for (let turn = 0; turn < 20; turn++) await Promise.resolve(); + return true; + }; + const start = (conversationId: ConversationId, promptTokens: number, sent = request(), stop?: StopReason) => + warmers.start(conversationId, { message: response(promptTokens, stop), request: sent }, models); + return { state, warmers, tick, start }; +} + +test("refreshes before the 5-minute entry expires, with the request's own options and one output token", async () => { + const { state, tick, start } = setup(); + start(CONVERSATION, 600_000); + assert.equal(state.alive, 1); + assert.ok(await tick()); + assert.equal(state.now, 270_000); + assert.equal(state.sent.length, 1); + const [refresh] = state.sent; + assert.equal(refresh!.model, MODEL); + assert.deepEqual(refresh!.context.messages, MESSAGES); + assert.equal(refresh!.options.sessionId, "provider-session"); + assert.equal(refresh!.options.reasoning, "high"); + assert.equal(refresh!.options.maxTokens, 1); + assert.equal(refresh!.options.maxRetries, 0); + assert.equal(refresh!.options.disableContinuation, true); + assert.deepEqual(state.recorded, [CONVERSATION]); + // The next refresh is armed again one delay later. + assert.ok(await tick()); + assert.equal(state.now, 540_000); + assert.equal(state.sent.length, 2); +}); + +test("a 1-hour retention refreshes after 54 minutes", async () => { + const { state, tick, start } = setup(); + start(CONVERSATION, 600_000, request({ options: { cacheRetention: "long" } })); + assert.ok(await tick()); + assert.equal(state.now, 3_240_000); + assert.equal(state.sent.length, 1); +}); + +test("streams up to an hour, idles up to 30 minutes, and stops after a settled subagent or in mode streaming", async () => { + // While the run streams (a long tool call), refreshes go on until the hour after the request. + const streaming = setup(); + streaming.start(CONVERSATION, 600_000); + while (await streaming.tick()); + assert.equal(streaming.state.sent.length, 13); + assert.ok(streaming.state.now <= 3_600_000); + assert.equal(streaming.state.alive, 0); + + // After the run answered, the main conversation is warmed for 30 minutes. + const idle = setup(); + idle.start(CONVERSATION, 600_000); + idle.warmers.settled(CONVERSATION); + idle.state.running.delete(CONVERSATION); + while (await idle.tick()); + assert.equal(idle.state.sent.length, 6); + assert.equal(idle.state.alive, 0); + + // A subagent is not warmed once its run settled, also when it ended without an answer (no `afterRun`). + const subagent = setup(); + subagent.start(SUBAGENT, 600_000); + subagent.state.running.delete(SUBAGENT); + while (await subagent.tick()); + assert.equal(subagent.state.sent.length, 0); + assert.equal(subagent.state.alive, 0); + + // Mode "streaming" stops when the run settles. + const streamingOnly = setup("streaming"); + streamingOnly.start(CONVERSATION, 600_000); + streamingOnly.warmers.settled(CONVERSATION); + assert.equal(streamingOnly.state.alive, 0); + assert.equal(await streamingOnly.tick(), false); +}); + +test("warms only what pays off: small prompts are not refreshed, idle needs larger ones than streaming", async () => { + const small = setup(); + small.start(CONVERSATION, 10_000); + await small.tick(); + assert.equal(small.state.sent.length, 0); + assert.equal(small.state.alive, 0); + + // 80k tokens: worth a refresh while the run streams, not when only a 15% chance remains that it continues. + const medium = setup(); + medium.start(CONVERSATION, 80_000); + await medium.tick(); + assert.equal(medium.state.sent.length, 1); + medium.warmers.settled(CONVERSATION); + medium.state.running.delete(CONVERSATION); + await medium.tick(); + assert.equal(medium.state.sent.length, 1); + assert.equal(medium.state.alive, 0); +}); + +test("does not warm what it cannot repeat: failed responses, replaced bodies, changed contexts, mode off", async () => { + for (const stop of ["error", "aborted"] as const) { + const failed = setup(); + failed.start(CONVERSATION, 600_000, request(), stop); + assert.equal(failed.state.alive, 0); + } + const replaced = setup(); + replaced.start(CONVERSATION, 600_000, request({ payloadReplaced: true })); + assert.equal(replaced.state.alive, 0); + + const off = setup("off"); + off.start(CONVERSATION, 600_000); + assert.equal(off.state.alive, 0); + + // A reset, compaction or model change before the refresh: nothing is sent, warming stops. + const changed = setup(); + changed.start(CONVERSATION, 600_000); + changed.state.current = false; + await changed.tick(); + assert.equal(changed.state.sent.length, 0); + assert.equal(changed.state.alive, 0); + + // Switched off while warming: the next refresh stops instead. + const switched = setup(); + switched.start(CONVERSATION, 600_000); + switched.state.mode = "off"; + await switched.tick(); + assert.equal(switched.state.sent.length, 0); + assert.equal(switched.state.alive, 0); +}); + +test("a timer that fires too late stops instead of paying for a full cache write", async () => { + const { state, tick, start } = setup(); + start(CONVERSATION, 600_000); + // The process slept: the timer fires after the entry has (nearly) expired. + state.timers[0]!.at = 299_000; + await tick(); + assert.equal(state.sent.length, 0); + assert.equal(state.alive, 0); +}); + +test("the next response replaces a conversation's run; close aborts a refresh in flight", async () => { + const { state, warmers, tick, start } = setup(); + start(CONVERSATION, 600_000); + start(CONVERSATION, 610_000); + assert.equal(state.alive, 1); + assert.equal(state.timers.filter((timer) => !timer.cancelled).length, 1); + start(SUBAGENT, 600_000); + assert.equal(state.alive, 2); + + state.hold = true; + await tick(); + assert.equal(state.sent.length, 1); + warmers.close(); + assert.equal(state.alive, 0); + assert.equal(state.sent[0]!.options.signal?.aborted, true); + state.release?.(); + for (let turn = 0; turn < 20; turn++) await Promise.resolve(); + // The aborted refresh records nothing and arms nothing. + assert.deepEqual(state.recorded, []); + assert.equal(await tick(), false); + assert.deepEqual(state.reports, []); +}); + +test("a context extends the sent messages when it starts with them, by role and timestamp", () => { + const next: Message = { role: "user", content: "more", timestamp: 4 }; + assert.equal(extendsMessages([...MESSAGES, next], MESSAGES), true); + assert.equal(extendsMessages(MESSAGES, MESSAGES), true); + // After a reset or compaction, the context starts with other messages. + assert.equal(extendsMessages([{ role: "user", content: "summary", timestamp: 9 }, next], MESSAGES), false); + assert.equal(extendsMessages(MESSAGES.slice(0, 2), MESSAGES), false); +}); + +test("settings.json accepts pi's cacheWarming modes", () => { + assert.deepEqual(validateSettings({ cacheWarming: "streaming" }), { cacheWarming: "streaming" }); + assert.throws(() => validateSettings({ cacheWarming: "always" }), /Invalid settings/); +}); + +test( + "a worker keeps the main conversation's cache warm after its run and counts the refreshes' spend", + { timeout: 90_000 }, + async (t) => { + const scratch = mkdtempSync(join(tmpdir(), "p2-cache-warming-")); + const requestLog = join(scratch, "requests.jsonl"); + writeFileSync(requestLog, ""); + const fixture = await startFixture({ workerEnv: { P2_FAKE_REQUESTS: requestLog } }); + t.after(async () => { + await fixture.close(); + rmSync(scratch, { recursive: true, force: true }); + }); + // The fake model as a cached, expensive one, before the session's worker reads models.json: an 11-second + // lifetime refreshes after one second, and every refresh pays off even for its small prompts. + writeFileSync( + join(fixture.dataDir, "models.json"), + JSON.stringify({ + providers: { + faux: { + modelOverrides: { + "faux-1": { + cost: { input: 4000, output: 20000, cacheRead: 200, cacheWrite: 5000 }, + promptCache: { short: 11 }, + }, + }, + }, + }, + }), + ); + const owner = fixture.browser(); + const session = ( + await (await holdIn(owner, Sessions, {}, "owner")).create({ cwd: fixture.workspace }, BACKGROUND_CONTEXT) + ).id; + const chat = await holdIn(owner, Chat, { session, conversation: 1 }, "participant"); + await chat.send({ text: "hello", requestId: "m1" }, BACKGROUND_CONTEXT); + const live = (view: ConversationView | null | undefined) => view?.docs["pi.live"] as LiveState | undefined; + const usage = (view: ConversationView | null | undefined) => + Object.values((view?.docs["pi.usage"] as unknown as UsageState | undefined)?.models ?? {}).reduce( + (total, model) => total + model.totalTokens, + 0, + ); + const answered = await until(chat.view, (view) => (view?.entries.length ?? 0) > 1 && !live(view)?.run, 30_000); + const requests = () => + readFileSync(requestLog, "utf8") + .split("\n") + .filter((line) => line.length > 0); + const [sent] = requests(); + assert.ok(sent); + // Idle now: each refresh repeats the request exactly. + const deadline = Date.now() + 20_000; + while (requests().length < 3 && Date.now() < deadline) await new Promise((resolve) => setTimeout(resolve, 200)); + const refreshes = requests().slice(1); + assert.ok(refreshes.length >= 2, `expected refreshes, got ${refreshes.length}`); + for (const refresh of refreshes) assert.equal(refresh, sent); + // Their spend counts as the conversation's. + const before = usage(answered); + await until(chat.view, (view) => usage(view) > before, 10_000); + }, +); diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 5625dfceb..e05d05677 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -6,6 +6,7 @@ - Anthropic requests with a `sessionId` (and `cacheRetention` not `"none"`) now use message threads (`message-threads-2026-08-12` beta): the server continues from the previous reply and pi sends only the new messages instead of the whole conversation. When anything earlier in the request changed, or the server can't continue the thread, the full request is sent. Token counts and prompt caching are unchanged. On by default only for the `anthropic` provider on `api.anthropic.com`; set `compat.supportsMessageThreads` to change that, for example for gateways that forward the `thread` field. Call `cleanupSessionResources(sessionId)` when a session ends to free the per-session state ([#11545](https://github.com/earendil-works/pi/pull/11545)) - Added `StreamOptions.disableContinuation`, which keeps a request from using or replacing server-side continuation state such as Anthropic message threads. Set it on requests that repeat an earlier one, such as cache warming ([#11545](https://github.com/earendil-works/pi/pull/11545)) +- Added the prompt cache warming policy as shared helpers, so every host warms the same way: `getPromptCacheTtlMs()`, `getCacheWarmingDelayMs()`, `isCacheWarmingReplayable()`, `getCacheWarmingRequestOptions()`, `evaluateCacheWarming()` and the `CACHE_WARMING_*` constants (moved from pi-coding-agent's cache warmer) - Added `@earendil-works/pi-ai/lazy-oauth` (`registerLazyOAuthFlows()`) and `@earendil-works/pi-ai/lazy-bedrock-provider` (`registerLazyBedrockProvider()`, built on the new `setBedrockProviderLoader()`) for Node bundles that split at dynamic imports: the OAuth login flows and Bedrock load through literal dynamic imports, so the bundler keeps them inside the bundle as chunks that load on first use, instead of through the variable-specifier imports that load pi-ai's own files at runtime. Nothing changes without calling them. ### Changed diff --git a/packages/ai/src/cache-warming.ts b/packages/ai/src/cache-warming.ts new file mode 100644 index 000000000..3675a56b2 --- /dev/null +++ b/packages/ai/src/cache-warming.ts @@ -0,0 +1,123 @@ +// Prompt cache warming policy shared by pi's hosts: when a refresh is due, whether a request can be replayed, and +// whether a refresh pays for itself. Hosts own the scheduling; this module owns the numbers. +import { calculateCost } from "./models.ts"; +import type { Api, Model, SimpleStreamOptions, Usage } from "./types.ts"; +import { getProviderEnvValue } from "./utils/provider-env.ts"; + +/** Whether a refresh is sent ("warm") or warming ends until the next real request ("stop"). */ +export type CacheWarmingAction = "warm" | "stop"; + +/** "streaming" while the agent run that sent the request is still active; "idle" after it settled. */ +export type CacheWarmingPhase = "streaming" | "idle"; + +/** Inputs and outcome of one warm-or-stop decision. */ +export interface CacheWarmingDecision { + phase: CacheWarmingPhase; + /** 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; + /** "warm" when `expectedSavings` is at least `CACHE_WARMING_MINIMUM_EXPECTED_SAVINGS`. */ + action: CacheWarmingAction; +} + +/** Streaming warming never continues past this long after the real request that started it. */ +export const CACHE_WARMING_MAX_AGE_MS = 60 * 60_000; + +/** Idle warming uses a shorter horizon because continuation estimates become less reliable with age. */ +export const CACHE_WARMING_MAX_IDLE_AGE_MS = 30 * 60_000; + +/** A refresh is sent only when it is expected to save at least this many dollars. */ +export 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. + */ +export const CACHE_WARMING_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, 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 isCacheWarmingReplayable(model: Model, options: SimpleStreamOptions | undefined): boolean { + if (!options?.reasoning || model.api !== "anthropic-messages") return true; + return (model as Model<"anthropic-messages">).compat?.forceAdaptiveThinking === true; +} + +/** + * The options of a refresh: the request's own, capped at one output token, without retries, and without replacing the + * conversation's provider thread state (the refresh repeats the last request). + */ +export function getCacheWarmingRequestOptions(options: T, signal: AbortSignal): T { + return { ...options, maxTokens: 1, disableContinuation: true, maxRetries: 0, signal }; +} + +/** Decide whether refreshing the cache entry of a `promptTokens`-token request pays for itself in `phase`. */ +export function evaluateCacheWarming( + model: Model, + promptTokens: number, + phase: CacheWarmingPhase, +): CacheWarmingDecision { + 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 = phase === "idle" ? CACHE_WARMING_IDLE_CONTINUATION_PROBABILITY : 1; + const economicsAvailable = promptTokens > 0 && (cacheHitCost > 0 || cacheMissCost > 0); + const expectedSavings = continuationProbability * missCost - warmCost; + return { + phase, + warmCost, + missCost, + continuationProbability, + expectedSavings, + economicsAvailable, + action: expectedSavings >= CACHE_WARMING_MINIMUM_EXPECTED_SAVINGS ? "warm" : "stop", + }; +} + +function price( + model: Model, + tokens: Partial>, +): 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; +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 0c9197ae6..7645e8925 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -19,6 +19,7 @@ export * from "./auth/context.ts"; export * from "./auth/credential-store.ts"; export * from "./auth/helpers.ts"; export * from "./auth/types.ts"; +export * from "./cache-warming.ts"; export type { OAuthAuthInfo, OAuthDeviceCodeInfo, diff --git a/packages/coding-agent/src/core/cache-warmer.ts b/packages/coding-agent/src/core/cache-warmer.ts index 6af5f02fa..97176f14e 100644 --- a/packages/coding-agent/src/core/cache-warmer.ts +++ b/packages/coding-agent/src/core/cache-warmer.ts @@ -1,61 +1,24 @@ import { type Api, + CACHE_WARMING_MAX_AGE_MS, + CACHE_WARMING_MAX_IDLE_AGE_MS, + CACHE_WARMING_MINIMUM_EXPECTED_SAVINGS, + type CacheWarmingAction, + type CacheWarmingDecision, type Context, - calculateCost, + evaluateCacheWarming, + getCacheWarmingDelayMs, + getCacheWarmingRequestOptions, + getPromptCacheTtlMs, + isCacheWarmingReplayable, 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, 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, options: SimpleStreamOptions | undefined): boolean { - if (!options?.reasoning || model.api !== "anthropic-messages") return true; - return (model as Model<"anthropic-messages">).compat?.forceAdaptiveThinking === true; -} +export type { CacheWarmingAction, CacheWarmingDecision }; /** Prompt size of the most recent real request on the branch, as reported by the provider. */ function lastPromptTokens(entries: SessionEntry[]): number { @@ -69,42 +32,6 @@ function lastPromptTokens(entries: SessionEntry[]): number { return 0; } -function price( - model: Model, - tokens: Partial>, -): 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. @@ -209,7 +136,7 @@ export class CacheWarmer { this.stop("cache warming disabled"); return; } - if (!isReplayable(request.model, request.options)) { + if (!isCacheWarmingReplayable(request.model, request.options)) { this.stop("request cannot be replayed safely"); return; } @@ -250,7 +177,7 @@ export class CacheWarmer { return; } run.phase = "idle"; - const deadline = run.startedAt + MAX_IDLE_WARMING_AGE_MS; + const deadline = run.startedAt + CACHE_WARMING_MAX_IDLE_AGE_MS; if (run.nextWarmAt > deadline || Date.now() >= deadline) { this.stop("30-minute idle safety limit reached"); } @@ -288,7 +215,8 @@ export class CacheWarmer { // the planned pre-expiry margin for that delay and request dispatch; a // late refresh is likely a full-price cache write, not a cache warm. run.refreshDeadlineAt = run.nextWarmAt + Math.floor((run.ttlMs - run.delayMs) / 2); - const deadline = run.startedAt + (run.phase === "idle" ? MAX_IDLE_WARMING_AGE_MS : MAX_WARMING_AGE_MS); + const deadline = + run.startedAt + (run.phase === "idle" ? CACHE_WARMING_MAX_IDLE_AGE_MS : CACHE_WARMING_MAX_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; @@ -330,14 +258,7 @@ export class CacheWarmer { run.extensionOverride = extensionOverride; try { const message = await this.models - .streamSimple(run.model, run.context, { - ...run.options, - maxTokens: 1, - // The warm request repeats the last request; it must not replace the conversation's thread state. - disableContinuation: true, - maxRetries: 0, - signal: run.controller.signal, - }) + .streamSimple(run.model, run.context, getCacheWarmingRequestOptions(run.options, run.controller.signal)) .result(); if (!this.validateRun(run)) return; if (message.stopReason !== "error" && message.stopReason !== "aborted") { @@ -378,27 +299,7 @@ export class CacheWarmer { } 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", - }; + return evaluateCacheWarming(run.model, lastPromptTokens(this.sessionManager.getBranch()), run.phase); } } diff --git a/packages/coding-agent/test/cache-warmer.test.ts b/packages/coding-agent/test/cache-warmer.test.ts index 1a7bf1790..492347099 100644 --- a/packages/coding-agent/test/cache-warmer.test.ts +++ b/packages/coding-agent/test/cache-warmer.test.ts @@ -2,6 +2,9 @@ import { type Api, type AssistantMessage, type AssistantMessageEventStream, + getCacheWarmingDelayMs, + getPromptCacheTtlMs, + isCacheWarmingReplayable, type Model, type ModelsSimpleStreamOptions, normalizeContext, @@ -18,9 +21,6 @@ import { 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"; @@ -137,10 +137,10 @@ describe("cache warming", () => { [270_000, 50_000, undefined], ); expect([ - isReplayable(budgetModel, { reasoning: "medium" }), - isReplayable(budgetModel, undefined), - isReplayable(adaptiveModel, { reasoning: "medium" }), - isReplayable(openaiModel, { reasoning: "medium" }), + isCacheWarmingReplayable(budgetModel, { reasoning: "medium" }), + isCacheWarmingReplayable(budgetModel, undefined), + isCacheWarmingReplayable(adaptiveModel, { reasoning: "medium" }), + isCacheWarmingReplayable(openaiModel, { reasoning: "medium" }), ]).toEqual([false, true, true, true]); }); diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index 9ef9732d9..940c262e9 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -2,8 +2,13 @@ ## [Unreleased] +### Breaking Changes + +- `GenerationHooks.afterResponse` takes a `GenerationResponse` (`{ message, request }`) instead of the bare message. `request` is the provider request the attempt sent (`GenerationRequest`: the model, the messages after `beforeRequest`, the options without the attempt's signal, telemetry and payload callback, and `payloadReplaced`), so a hook can send it again, e.g. to keep its prompt cache entry warm; it is `undefined` for a polled deferred response. Migration: `afterResponse: (message) => ...` becomes `afterResponse: ({ message }) => ...`. + ### Added +- `recordUsage(tx, conversationId, bucket, key, usage)` is exported, so an extension can add spend it caused outside a generation (such as cache warming) to a conversation's `pi.usage`. - `Tx.submission(id)` reads a committed submission record by ID, so a tool can find the entries of its run's inputs (`pi.live.run.inputs`) without scanning the transcript. ### Changed diff --git a/packages/durable/README.md b/packages/durable/README.md index 49727c747..5fc3328c6 100644 --- a/packages/durable/README.md +++ b/packages/durable/README.md @@ -659,7 +659,7 @@ const Guard = defineExtension({ }); ``` -- **Generation:** `beforeRequest` (replace the messages of one request), `beforeProviderRequest` (replace the provider request body of one request, e.g. to set a provider-specific field), `afterResponse`, `onYield` (continue the run with another user message), `afterTools` (runs once a round's tools are done), and `afterRun` (tasks to run before the run ends, below). +- **Generation:** `beforeRequest` (replace the messages of one request), `beforeProviderRequest` (replace the provider request body of one request, e.g. to set a provider-specific field), `afterResponse` (the response with the request that produced it, e.g. to repeat the request and keep its prompt cache warm), `onYield` (continue the run with another user message), `afterTools` (runs once a round's tools are done), and `afterRun` (tasks to run before the run ends, below). - **Tools:** `beforeTool` (block or rewrite arguments) and `afterTool` (replace the result). `afterTool` sees every result a tool task writes, also the error results the Harness writes itself (unavailable, invalid, blocked, interrupted, restart limit, aborted), told apart by their diagnostic code; past `settings.restartLimit` it no longer runs. On an abort it runs inside the abort handler, so a slow hook delays `abort()`. To limit a hook to some conversations, select its extension only there, for example with `configure({ extensions: { add: [Guard] } })`. diff --git a/packages/durable/docs/spec.md b/packages/durable/docs/spec.md index fcae714a3..b883dc6e4 100644 --- a/packages/durable/docs/spec.md +++ b/packages/durable/docs/spec.md @@ -3382,8 +3382,8 @@ interface GenerationHooks { api: HookApi, context: Context, ): HookResult<{ readonly payload: unknown }>; - /** Every terminal provider message, before classification. */ - afterResponse(message: AssistantMessage, api: HookApi, context: Context): void | Promise; + /** Every terminal provider message, before classification, with the request that produced it. */ + afterResponse(response: GenerationResponse, api: HookApi, context: Context): void | Promise; /** A final answer; `continue` appends a user message and continues the run. */ onYield(answer: AssistantMessage, api: HookApi, context: Context): HookResult<{ readonly continue: UserInput }>; /** After every tool of the round is terminal; `results` are the round's result entries in call order. */ @@ -3392,6 +3392,26 @@ interface GenerationHooks { afterRun(run: RunSummary, api: HookApi, context: Context): HookResult<{ readonly tasks: readonly RunEndTask[] }>; } +/** + * A provider request one attempt sent, as passed to `models.streamSimple()`: sending + * it again repeats it, which keeps its prompt cache entry alive (cache warming). + */ +type GenerationRequest = { + readonly model: Model; + /** The model context after the `beforeRequest` chain. */ + readonly messages: readonly Message[]; + /** The request's options without the attempt's `signal`, `telemetry` and `onPayload`. */ + readonly options: SimpleStreamOptions; + /** Whether a `beforeProviderRequest` hook replaced the body; a resend without the chain differs. */ + readonly payloadReplaced: boolean; +}; + +/** `request` is `undefined` for a polled deferred response, which sent no request. */ +type GenerationResponse = { + readonly message: AssistantMessage; + readonly request: GenerationRequest | undefined; +}; + /** How a run ended, for `afterRun`. */ type RunSummary = { /** A final answer, or a tool round whose every call asked to `terminate`. */ @@ -4915,7 +4935,10 @@ The run's inputs live in `pi.live.run`, not in the task input. stream options, and reruns `beforeRequest`; after a crash the attempt is charged instead, as above. - `afterResponse` observes every terminal message, from `request` or `poll`, - before classification; a still deferred result is not terminal. + before classification; a still deferred result is not terminal. It gets the + request the attempt sent (model, messages after `beforeRequest`, options + without the attempt's signal, telemetry and payload callback, and whether a + `beforeProviderRequest` hook replaced the body); a polled response has none. - Before classifying, the handler stops the partial throttle and awaits any partial commit in flight, so no stale partial lands after the outcome. The terminal message is classified in one commit that also clears the partial: diff --git a/packages/durable/src/harness/generation.ts b/packages/durable/src/harness/generation.ts index a7eff02e7..deec550b9 100644 --- a/packages/durable/src/harness/generation.ts +++ b/packages/durable/src/harness/generation.ts @@ -56,6 +56,7 @@ import type { ContextView, ConversationStreamOptions, GenerationHooks, + GenerationRequest, ModelRef, PromptInput, RunEndTask, @@ -153,6 +154,8 @@ type Request = { readonly messages?: readonly Message[]; /** Set when the message came from polling, so a still deferred result polls strictly later. */ readonly pollAt?: number; + /** The provider request the attempt sent; absent for a polled deferred response. */ + readonly sent?: GenerationRequest; }; const DEFAULT_POLL_AFTER_MS = 5000; @@ -288,16 +291,25 @@ export const GenerationTask = defineTask providerPayload(runtime, model, payload, context), + }; + let payloadReplaced = false; + const options: SimpleStreamOptions = { + ...sentOptions, + signal: runtime.signal, + telemetry: context.telemetry, + onPayload: async (payload) => { + const replaced = await providerPayload(runtime, model, payload, context); + if (replaced !== undefined) payloadReplaced = true; + return replaced ?? payload; + }, }; const message = await streamResponse(runtime, model, messages, options, attempt, context); - const request = { attempt, compacted, model: ref, thinkingLevel, cutoff, messages: view.messages }; + const sent = { model, messages, options: sentOptions, payloadReplaced }; + const request = { attempt, compacted, model: ref, thinkingLevel, cutoff, messages: view.messages, sent }; await classify(runtime, request, message, context); }, /** After the backoff's sleep; a run before `until` (an early wake, or a record of an earlier version) sleeps again. */ @@ -667,19 +679,23 @@ export async function convertPartial(tx: Tx, live: Draft, conversatio await appendAssistant(tx, conversationId, { ...message, stopReason: "aborted" }); } -/** The `beforeProviderRequest` replacement chain over one provider request body. */ +/** The `beforeProviderRequest` replacement chain over one provider request body; `undefined` when no hook replaced it. */ async function providerPayload( runtime: Runtime, model: Model, payload: unknown, context: Context, ): Promise { - let current = payload; + let current: { readonly payload: unknown } | undefined; await runtime.hooks.each("beforeProviderRequest", async (hook) => { - const replaced = await hook({ model, payload: current }, runtime, context); - if (replaced !== undefined) current = replaced.payload; + const replaced = await hook( + { model, payload: current === undefined ? payload : current.payload }, + runtime, + context, + ); + if (replaced !== undefined) current = replaced; }); - return current; + return current?.payload; } /** @@ -791,7 +807,8 @@ async function classify( }, context); return; } - await runtime.hooks.each("afterResponse", (hook) => hook(message, runtime, context)); + const response = { message, request: request.sent }; + await runtime.hooks.each("afterResponse", (hook) => hook(response, runtime, context)); const calls = message.content.filter((content): content is ToolCall => content.type === "toolCall"); if (message.stopReason === "toolUse" && calls.length > 0) { return startToolRound(runtime, request, message, calls, context); diff --git a/packages/durable/src/harness/types.ts b/packages/durable/src/harness/types.ts index ee875ded8..64344f909 100644 --- a/packages/durable/src/harness/types.ts +++ b/packages/durable/src/harness/types.ts @@ -10,6 +10,7 @@ import type { Model, Models, ModelThinkingLevel, + SimpleStreamOptions, ThinkingBudgets, Tool, ToolCall, @@ -980,6 +981,30 @@ export interface HookApi extends DocumentReader { export type HookResult = T | undefined | Promise; +/** + * A provider request one generation attempt sent, as it was passed to `models.streamSimple()`: sending it again with the + * same model, messages and options repeats it, which keeps its prompt cache entry alive (cache warming). + */ +export type GenerationRequest = { + readonly model: Model; + /** The model context after the `beforeRequest` chain. */ + readonly messages: readonly Message[]; + /** The request's options without the attempt's `signal`, `telemetry` and `onPayload`. */ + readonly options: SimpleStreamOptions; + /** + * Whether a `beforeProviderRequest` hook replaced the provider request body. A resend without that chain then + * sends a different body. + */ + readonly payloadReplaced: boolean; +}; + +/** A terminal provider message for `afterResponse`. */ +export type GenerationResponse = { + readonly message: AssistantMessage; + /** The request that produced it; `undefined` for a polled deferred response, which sent no request. */ + readonly request: GenerationRequest | undefined; +}; + /** Hooks of the built-in generation task. */ export interface GenerationHooks { /** Before every request attempt, including recovery; the result is used for that request only. */ @@ -998,8 +1023,8 @@ export interface GenerationHooks { api: HookApi, context: Context, ): HookResult<{ readonly payload: unknown }>; - /** Every terminal provider message, before classification. */ - afterResponse(message: AssistantMessage, api: HookApi, context: Context): void | Promise; + /** Every terminal provider message, before classification, with the request that produced it. */ + afterResponse(response: GenerationResponse, api: HookApi, context: Context): void | Promise; /** A final answer; the first `continue` appends a user message and continues the run. */ onYield(answer: AssistantMessage, api: HookApi, context: Context): HookResult<{ readonly continue: UserInput }>; /** After every tool of the round is terminal; `results` are the round's result entries in call order. */ diff --git a/packages/durable/src/index.ts b/packages/durable/src/index.ts index 1a3babbee..633dfaf1f 100644 --- a/packages/durable/src/index.ts +++ b/packages/durable/src/index.ts @@ -96,6 +96,8 @@ export type { EnvTarget, Extension, GenerationHooks, + GenerationRequest, + GenerationResponse, HarnessInspection, HarnessOptions, HarnessSettings, @@ -134,7 +136,7 @@ export type { Wake, Wrap, } from "./harness/types.ts"; -export { UsageDoc, type UsageState } from "./harness/usage.ts"; +export { recordUsage, UsageDoc, type UsageState } from "./harness/usage.ts"; export type { ConversationView, ConversationWindow } from "./harness/view.ts"; export { createSession } from "./session/session.ts"; export { MemoryStorage } from "./storage/memory.ts"; diff --git a/packages/durable/test/harness-tools.test.ts b/packages/durable/test/harness-tools.test.ts index f16205378..bc73e71c5 100644 --- a/packages/durable/test/harness-tools.test.ts +++ b/packages/durable/test/harness-tools.test.ts @@ -20,6 +20,7 @@ import { defineTask, defineTool, type EntryRecord, + type GenerationRequest, GenerationTask, type Harness, hook, @@ -736,7 +737,7 @@ describe("generation hooks", () => { }); const responses: string[] = []; addHooks(setup.registry, GenerationTask, { - afterResponse: (message) => void responses.push(textOf(message) ?? ""), + afterResponse: ({ message }) => void responses.push(textOf(message) ?? ""), }); let yields = 0; addHooks(setup.registry, GenerationTask, { @@ -792,6 +793,48 @@ describe("generation hooks", () => { await harness.close(context); }); + it("gives afterResponse the request each attempt sent", async () => { + const setup = chatSetup(); + addHooks(setup.registry, GenerationTask, { + beforeRequest: ({ messages }) => ({ + messages: [...messages, { role: "user", content: "injected", timestamp: 0 }], + }), + }); + let replace = false; + addHooks(setup.registry, GenerationTask, { + beforeProviderRequest: ({ payload }) => + replace ? { payload: { ...(payload as object), changed: true } } : undefined, + }); + const requests: GenerationRequest[] = []; + addHooks(setup.registry, GenerationTask, { + afterResponse: ({ request }) => void requests.push(request!), + }); + let yields = 0; + addHooks(setup.registry, GenerationTask, { + onYield: () => (yields++ === 0 ? { continue: "again" } : undefined), + }); + const respond: FauxResponseStep = async (_context, options, _state, model) => { + await options!.onPayload!({ body: 1 }, model); + replace = true; + return DONE; + }; + const { harness, status } = await run(setup, [respond, respond]); + expect(status).toBe("done"); + const [first, second] = requests; + expect(`${first!.model.provider}/${first!.model.id}`).toBe("faux/faux-1"); + // The messages are the ones sent: after the `beforeRequest` chain. + expect(textOf(first!.messages.at(-1) as never)).toBe("injected"); + // The options are the request's own, without what belongs to the attempt. + expect(typeof first!.options.sessionId).toBe("string"); + expect(first!.options.sessionId).toBe(second!.options.sessionId); + expect(first!.options.signal).toBeUndefined(); + expect(first!.options.onPayload).toBeUndefined(); + expect(first!.options.telemetry).toBeUndefined(); + // Only the second request's body was replaced by a hook. + expect([first!.payloadReplaced, second!.payloadReplaced]).toEqual([false, true]); + await harness.close(context); + }); + it("reports a throwing beforeProviderRequest hook and continues the chain", async () => { const setup = chatSetup(); addHooks( @@ -851,13 +894,16 @@ describe("generation hooks", () => { const setup = chatSetup({ deferred: { pendingFetches: 1, pollAfterMs: 1 } }); const observed: string[] = []; addHooks(setup.registry, GenerationTask, { - afterResponse: (message) => void observed.push(`${message.stopReason}:${textOf(message) ?? ""}`), + afterResponse: ({ message, request }) => + void observed.push( + `${message.stopReason}:${textOf(message) ?? ""}:${request === undefined ? "polled" : "sent"}`, + ), }); setup.settings.stream = { deferred: true }; const result = await run(setup, [fauxAssistantMessage([fauxText("late")])]); expect(result.status).toBe("done"); // The still-deferred results are not terminal. - expect(observed).toEqual(["stop:late"]); + expect(observed).toEqual(["stop:late:polled"]); await result.harness.close(context); });