refactor: simplify p2 cache warming and its tests

- pi-durable: GenerationRequest drops payloadReplaced; the body chain is
  unchanged.
- p2: the warmer asks the host for the conversation's phase at each refresh
  (streaming while pi.live.run exists, idle for top-level conversations,
  nothing once the request is stale) instead of tracking afterRun, uses real
  timers, and drops the late-timer guard.
- Tests run the warmer on mocked timers in four tests plus the worker test.
This commit is contained in:
Armin Ronacher 2026-11-06 00:17:06 +01:00
parent dd8e0df4a9
commit b76afbdb2f
10 changed files with 234 additions and 445 deletions

View file

@ -1183,21 +1183,19 @@ translated phone drawer and intercept chat taps.
- **Cache warming** (`plugins/cache-warming`): pi-coding-agent's cache warmer for every conversation of a session. - **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 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 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% starts that conversation's run, replacing the previous one. Before the entry expires (90% of
of `model.promptCache[retention]`: 270 s for the 5-minute tier) the run sends the request again with `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 `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 conversation's `pi.usage` (`recordUsage`). The policy is pi-ai's, shared with pi-coding-agent: up to an hour after the
with pi-coding-agent): streaming up to an hour after the request, idle up to 30 minutes, a refresh only when its request while the run streams, 30 minutes once idle, a refresh only when its expected savings reach $0.05
expected savings reach $0.05 (continuation probability 1 while streaming, 0.15 idle). p2's additions: only (`evaluateCacheWarming`). Each refresh reads the conversation's phase: streaming while `pi.live.run` exists, else
top-level conversations (`isTopLevel`) are warmed idle, since a settled subagent or one-shot agent rarely continues; idle for top-level conversations (`isTopLevel`; a settled subagent or one-shot agent is not warmed), and nothing
a run that ended without `afterRun` (an abort during tools, a turn limit) turns idle at the next refresh (`pi.live.run` once the next request would start differently (another model, or a context that no longer starts with the sent
gone); a refresh first checks the conversation still sends the same prefix (same model, a context that extends the messages by role and timestamp: a reset, compaction or edit). Timers live in `CacheWarmingRuntime`'s native, closed by
sent messages by role and timestamp), so a reset, compaction or model change stops it; requests whose body a a reload or stop; nothing is durable, the next response starts again. While a run is armed the worker's liveness
`beforeProviderRequest` hook replaced (`payloadReplaced`) are not warmed. Timers and refreshes are a native of counts it (`cacheWarming`), so the idle stop (`P2_IDLE_MS`) waits. Setting `cacheWarming` (`"idle"` default,
`CacheWarmingRuntime`, closed by a reload or stop (nothing is durable: the next response starts again); while a `"streaming"`, `"off"`, switched live through `ExtensionSwitches`). No status UI and no decision hook yet (pi's
run is armed the worker's liveness counts it (`cacheWarming`), so the idle stop (`P2_IDLE_MS`) waits. Setting: `/session` line, `cache_warming_decision`). `test/p2/cache-warming.test.ts`.
`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, - **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. `<P2_DATA_DIR>/mcp.json` host lifetime) owns every MCP process and connection; workers never spawn one or hold a credential. `<P2_DATA_DIR>/mcp.json`
is the owner's `mcpServers` (pi-mcp's `validateMcpServerConfig`, plus p2's `shared` and `runOn`; invalid entries are is the owner's `mcpServers` (pi-mcp's `validateMcpServerConfig`, plus p2's `shared` and `runOn`; invalid entries are

View file

@ -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`) | | 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 | | 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`) | cache-warming (`afterResponse` with the sent request, `afterRun`) | | 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) |
| Wraps | `wraps: [wrapTool(tool, fn) \| wrapSection(section, fn)]` (`types.ts:360`) | replace another extension's tool or section where selected | none yet | | 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`) | | 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 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). Only cache-warming uses hooks (`afterResponse`, `afterRun`); pi's `cache_warming_decision` has no p2 counterpart. | | `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`); 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. | | `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. | | `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. | | `registerFlag` | none | Plugins get configuration only via env vars or their own files; `settings.json` is closed. |

View file

@ -1,7 +1,6 @@
// Prompt cache warming for the session's conversations, as pi-coding-agent does it for its one session: before the // 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 // 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 // long as that is expected to pay off. The policy is pi-ai's; this file schedules it per conversation.
// conversation.
import { import {
type AssistantMessage, type AssistantMessage,
CACHE_WARMING_MAX_AGE_MS, CACHE_WARMING_MAX_AGE_MS,
@ -12,7 +11,6 @@ import {
getCacheWarmingRequestOptions, getCacheWarmingRequestOptions,
getPromptCacheTtlMs, getPromptCacheTtlMs,
isCacheWarmingReplayable, isCacheWarmingReplayable,
type Message,
type Models, type Models,
} from "@earendil-works/pi-ai"; } from "@earendil-works/pi-ai";
import type { ConversationId, GenerationRequest, GenerationResponse } from "@earendil-works/pi-durable"; import type { ConversationId, GenerationRequest, GenerationResponse } from "@earendil-works/pi-durable";
@ -22,22 +20,16 @@ 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. */ /** What the warmers read and do outside themselves; the worker facet binds it to the Harness and the settings. */
export type WarmerHost = { export type WarmerHost = {
/** The `cacheWarming` setting now. */
mode(): CacheWarmingMode; mode(): CacheWarmingMode;
now(): number; /**
/** Arms a timer and returns its cancel. */ * "streaming" while the conversation's run is under way, "idle" after it, or `undefined` when warming the
setTimer(fn: () => void, delayMs: number): () => void; * conversation is pointless: its next request will not start with `request`, or nobody is expected to continue it.
/** Whether the conversation's next request still starts with `request`: same model, a context extending its messages. */ */
isCurrent(conversationId: ConversationId, request: GenerationRequest): Promise<boolean>; phase(conversationId: ConversationId, request: GenerationRequest): Promise<CacheWarmingPhase | undefined>;
/** Whether a run is under way in the conversation. */
running(conversationId: ConversationId): Promise<boolean>;
/** Whether the conversation is warmed after its run settled: top-level conversations, not subagents. */
warmsIdle(conversationId: ConversationId): Promise<boolean>;
/** Adds a refresh's spend to the conversation's usage. */ /** Adds a refresh's spend to the conversation's usage. */
record(conversationId: ConversationId, message: AssistantMessage): Promise<void>; record(conversationId: ConversationId, message: AssistantMessage): Promise<void>;
/** Counts the conversations being warmed; a count keeps the worker process alive. */ /** Counts the conversations being warmed; a count keeps the worker process alive. */
keepAlive(delta: number): void; keepAlive(delta: number): void;
report(error: unknown): void;
}; };
type Run = { type Run = {
@ -45,20 +37,16 @@ type Run = {
readonly models: Models; readonly models: Models;
/** Prompt size of the request, as the provider reported it. */ /** Prompt size of the request, as the provider reported it. */
readonly promptTokens: number; readonly promptTokens: number;
readonly ttlMs: number;
readonly delayMs: number; readonly delayMs: number;
readonly startedAt: number; readonly startedAt: number;
readonly controller: AbortController; readonly controller: AbortController;
phase: CacheWarmingPhase; timer?: ReturnType<typeof setTimeout>;
/** 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 * One warming run per conversation, started by each response and replaced by the next one. Each refresh asks the host
* refreshes while its conversation's run streams (up to an hour after the request) and, in mode "idle", for up to 30 * for the conversation's phase and refreshes while the request is at most an hour old (30 minutes once idle) and the
* minutes after it settled, while each refresh pays off by pi-ai's `evaluateCacheWarming()`. * refresh pays off by pi-ai's `evaluateCacheWarming()`.
*/ */
export class ConversationWarmers { export class ConversationWarmers {
readonly #host: WarmerHost; readonly #host: WarmerHost;
@ -69,40 +57,26 @@ export class ConversationWarmers {
} }
/** Warm the cache entry `response`'s request wrote, replacing the conversation's run; `models` sends the refreshes. */ /** 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 { start(conversationId: ConversationId, { message, request }: GenerationResponse, models: Models): void {
this.stop(conversationId); this.stop(conversationId);
const { message, request } = response; // A polled deferred response sent no request; a failed or aborted one ended its run.
// 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 (request === undefined || message.stopReason === "error" || message.stopReason === "aborted") return;
if (this.#host.mode() === "off" || request.payloadReplaced) return; if (this.#host.mode() === "off" || !isCacheWarmingReplayable(request.model, request.options)) return;
if (!isCacheWarmingReplayable(request.model, request.options)) return;
const ttlMs = getPromptCacheTtlMs(request.model, request.options); const ttlMs = getPromptCacheTtlMs(request.model, request.options);
const delayMs = ttlMs === undefined ? undefined : getCacheWarmingDelayMs(ttlMs); const delayMs = ttlMs === undefined ? undefined : getCacheWarmingDelayMs(ttlMs);
if (ttlMs === undefined || delayMs === undefined) return; if (delayMs === undefined) return;
const usage = message.usage; const { input, cacheRead, cacheWrite } = message.usage;
const run: Run = { const run: Run = {
request, request,
models, models,
promptTokens: usage.input + usage.cacheRead + usage.cacheWrite, promptTokens: input + cacheRead + cacheWrite,
ttlMs,
delayMs, delayMs,
startedAt: this.#host.now(), startedAt: Date.now(),
controller: new AbortController(), controller: new AbortController(),
phase: "streaming",
refreshDeadlineAt: 0,
}; };
this.#runs.set(conversationId, run); this.#runs.set(conversationId, run);
this.#host.keepAlive(1); this.#host.keepAlive(1);
this.#schedule(conversationId, run); run.timer = setTimeout(() => void this.#refresh(conversationId, run), delayMs);
}
/** 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 warming the conversation; a refresh in flight is aborted. */
@ -110,7 +84,7 @@ export class ConversationWarmers {
const run = this.#runs.get(conversationId); const run = this.#runs.get(conversationId);
if (run === undefined) return; if (run === undefined) return;
this.#runs.delete(conversationId); this.#runs.delete(conversationId);
run.cancelTimer?.(); clearTimeout(run.timer);
run.controller.abort(); run.controller.abort();
this.#host.keepAlive(-1); this.#host.keepAlive(-1);
} }
@ -120,24 +94,9 @@ export class ConversationWarmers {
for (const conversationId of [...this.#runs.keys()]) this.stop(conversationId); 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<void> { async #refresh(conversationId: ConversationId, run: Run): Promise<void> {
run.cancelTimer = undefined;
try { try {
if (!(await this.#shouldRefresh(conversationId, run))) { if (!(await this.#wanted(conversationId, run))) {
if (this.#runs.get(conversationId) === run) this.stop(conversationId); if (this.#runs.get(conversationId) === run) this.stop(conversationId);
return; return;
} }
@ -152,35 +111,21 @@ export class ConversationWarmers {
await this.#host.record(conversationId, message); await this.#host.record(conversationId, message);
} catch (error) { } catch (error) {
// Warming is best effort: a failed refresh never touches the conversation, the next one may succeed. // 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 (!run.controller.signal.aborted) console.error("p2 cache-warming: refreshing failed:", error);
} }
if (this.#runs.get(conversationId) === run) this.#schedule(conversationId, run); if (this.#runs.get(conversationId) === run)
run.timer = setTimeout(() => void this.#refresh(conversationId, run), run.delayMs);
} }
/** Whether the armed refresh is still wanted: the run current, in its window, and paying off. */ /** Whether this refresh is still wanted: the run current, in its phase's window, and paying off. */
async #shouldRefresh(conversationId: ConversationId, run: Run): Promise<boolean> { async #wanted(conversationId: ConversationId, run: Run): Promise<boolean> {
const mode = this.#host.mode(); const mode = this.#host.mode();
if (mode === "off" || this.#runs.get(conversationId) !== run) return false; if (mode === "off") return false;
if (this.#host.now() > run.refreshDeadlineAt) return false; const phase = await this.#host.phase(conversationId, run.request);
// A run that ended without an answer (aborted during its tools, a turn limit) never called `afterRun`. if (phase === undefined || (phase === "idle" && mode === "streaming")) return false;
if (run.phase === "streaming" && !(await this.#host.running(conversationId))) run.phase = "idle"; const maxAge = phase === "idle" ? CACHE_WARMING_MAX_IDLE_AGE_MS : CACHE_WARMING_MAX_AGE_MS;
if (run.phase === "idle") { if (Date.now() > run.startedAt + maxAge) return false;
if (mode !== "idle" || !(await this.#host.warmsIdle(conversationId))) return false; const { action } = evaluateCacheWarming(run.request.model, run.promptTokens, phase);
if (this.#host.now() > run.startedAt + CACHE_WARMING_MAX_IDLE_AGE_MS) return false; return action === "warm" && this.#runs.get(conversationId) === run;
}
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,
);
}

View file

@ -1,12 +1,13 @@
// Cache warming in the session worker: a Durable extension whose generation hooks start, idle and stop the warming of // Cache warming in the session worker: a Durable extension whose `afterResponse` hook starts warming each
// each conversation's last request (`./warmer.ts`), bound to the Harness, the settings and the worker's liveness. // conversation's last request (`./warmer.ts`), bound to the Harness, the settings and the worker's liveness.
import type { Message } from "@earendil-works/pi-ai";
import { defineExtension, GenerationTask, hook, LiveDoc, recordUsage } from "@earendil-works/pi-durable"; import { defineExtension, GenerationTask, hook, LiveDoc, recordUsage } from "@earendil-works/pi-durable";
import { BACKGROUND_CONTEXT, defineFacet } from "@p2/app/v1"; import { BACKGROUND_CONTEXT, defineFacet } from "@p2/app/v1";
import { ExtensionSwitches } from "@p2/app/v1/node"; import { ExtensionSwitches } from "@p2/app/v1/node";
import { Extensions, harnessOf, isTopLevel, workerOf } from "@p2/app/v1/worker"; import { Extensions, harnessOf, isTopLevel, workerOf } from "@p2/app/v1/worker";
import { WorkerModels } from "../models/contract.ts"; import { WorkerModels } from "../models/contract.ts";
import { CacheWarmingRuntime } from "./contract.ts"; import { CacheWarmingRuntime } from "./contract.ts";
import { ConversationWarmers, extendsMessages } from "./warmer.ts"; import { ConversationWarmers } from "./warmer.ts";
export default defineFacet((host) => { export default defineFacet((host) => {
host.serve(CacheWarmingRuntime, (_params, self) => { host.serve(CacheWarmingRuntime, (_params, self) => {
@ -14,44 +15,35 @@ export default defineFacet((host) => {
const worker = workerOf(self); const worker = workerOf(self);
const settings = self.get(WorkerModels, {}); const settings = self.get(WorkerModels, {});
const durable = () => harness().durable; const durable = () => harness().durable;
const context = BACKGROUND_CONTEXT;
// Timers and in-flight refreshes live in this native: a reload or the worker's stop ends them, and the next // 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. // response starts warming again. Nothing of it is durable: the requests it repeats are not stored.
const warmers = self.open( const warmers = self.open(
async () => async () =>
new ConversationWarmers({ new ConversationWarmers({
mode: () => settings.settings.value.cacheWarming ?? "idle", mode: () => settings.settings.value.cacheWarming ?? "idle",
now: () => Date.now(), phase: async (conversationId, request) => {
setTimer: (fn, delayMs) => { const running = (await durable().snapshot(LiveDoc, conversationId, context))?.run !== undefined;
const timer = setTimeout(fn, delayMs); // After its run, a subagent's or one-shot agent's work is usually done; the main conversation and its
return () => clearTimeout(timer); // `/tree` branches wait for the person or for agents' reports.
if (!running && !(await isTopLevel(durable(), conversationId, context))) return undefined;
const conversation = await durable().conversation(conversationId, context);
if (conversation === undefined) return undefined;
const { model } = await conversation.agent(context);
if (model?.provider !== request.model.provider || model.modelId !== request.model.id)
return undefined;
// A reset, compaction or edit replaced what the request sent: the next request starts differently.
const current = (await conversation.context(context)).messages;
if (!startsWith(current, request.messages)) return undefined;
return running ? "streaming" : "idle";
}, },
isCurrent: async (conversationId, request) => { record: async (conversationId, { provider, model, usage }) => {
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( await durable().commit(
(tx) => (tx) => recordUsage(tx, conversationId, "models", `${provider}/${model}`, usage),
recordUsage( context,
tx,
conversationId,
"models",
`${message.provider}/${message.model}`,
message.usage,
),
BACKGROUND_CONTEXT,
); );
}, },
keepAlive: (delta) => worker().liveness.add("cacheWarming", delta), keepAlive: (delta) => worker().liveness.add("cacheWarming", delta),
report: (error) => console.error("p2 cache-warming: refreshing the prompt cache failed:", error),
}), }),
(opened) => opened.close(), (opened) => opened.close(),
); );
@ -60,14 +52,10 @@ export default defineFacet((host) => {
hooks: [ hooks: [
hook(GenerationTask, { hook(GenerationTask, {
afterResponse: (response, api) => warmers.value.start(api.conversationId, response, api.models), 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. // `cacheWarming: "off"` uninstalls the hook live; runs already warming stop at their next refresh.
self.contribute( self.contribute(
ExtensionSwitches, ExtensionSwitches,
self.derive([settings.settings], ([value]) => ({ [extension.name]: value.cacheWarming !== "off" })), self.derive([settings.settings], ([value]) => ({ [extension.name]: value.cacheWarming !== "off" })),
@ -76,3 +64,13 @@ export default defineFacet((host) => {
return {}; return {};
}); });
}); });
/** Whether `current` starts with `sent`, compared by role and timestamp. */
function startsWith(current: readonly Message[], sent: readonly Message[]): boolean {
return (
current.length >= sent.length &&
sent.every(
(message, index) => current[index]!.role === message.role && current[index]!.timestamp === message.timestamp,
)
);
}

View file

@ -1,35 +1,27 @@
// Cache warming (plugins/cache-warming): the warmers re-send a conversation's last request before its prompt cache // Cache warming (plugins/cache-warming): a conversation's last request is sent again with one output token before its
// entry expires, with one output token, for as long as pi's policy says it pays off, and stop when the request no // prompt cache entry expires, while pi's policy says it pays off. The warmers run on mocked timers with fake models;
// longer matches what the conversation sends next. Driven with a fake clock, timers and models, then end to end in a // a worker runs them end to end with the fake provider given a short cache lifetime.
// worker with the fake provider given a short cache lifetime.
import assert from "node:assert/strict"; import assert from "node:assert/strict";
import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os"; import { tmpdir } from "node:os";
import { join } from "node:path"; import { join } from "node:path";
import { test } from "node:test"; import { mock, test } from "node:test";
import type { import type {
AssistantMessage, CacheWarmingPhase,
Context, Context,
Message,
Model, Model,
Models, Models,
ModelsSimpleStreamOptions, ModelsSimpleStreamOptions,
SimpleStreamOptions,
StopReason, StopReason,
} from "@earendil-works/pi-ai"; } from "@earendil-works/pi-ai";
import { fauxAssistantMessage } from "@earendil-works/pi-ai/providers/faux"; import { fauxAssistantMessage } from "@earendil-works/pi-ai/providers/faux";
import type { import type { ConversationId, ConversationView, LiveState, UsageState } from "@earendil-works/pi-durable";
ConversationId,
ConversationView,
GenerationRequest,
LiveState,
UsageState,
} from "@earendil-works/pi-durable";
import { BACKGROUND_CONTEXT } from "@p2/app/v1"; import { BACKGROUND_CONTEXT } from "@p2/app/v1";
import { type CacheWarmingMode, ConversationWarmers, extendsMessages } from "../../plugins/cache-warming/warmer.ts"; import { type CacheWarmingMode, ConversationWarmers } from "../../plugins/cache-warming/warmer.ts";
import { Chat } from "../../plugins/chat/contract.ts"; import { Chat } from "../../plugins/chat/contract.ts";
import { Sessions } from "../../plugins/sessions/contract.ts"; import { Sessions } from "../../plugins/sessions/contract.ts";
import { validateSettings } from "../../src/runtime/settings.ts";
import { holdIn, startFixture, until } from "./helpers.ts"; import { holdIn, startFixture, until } from "./helpers.ts";
const MODEL: Model<"pi-messages"> = { const MODEL: Model<"pi-messages"> = {
@ -46,254 +38,154 @@ const MODEL: Model<"pi-messages"> = {
promptCache: { short: 300, long: 3600 }, promptCache: { short: 300, long: 3600 },
}; };
const CONVERSATION = 1 as ConversationId; const CONVERSATION = 1 as ConversationId;
const SUBAGENT = 7 as ConversationId; const REQUEST = {
const MESSAGES: Message[] = [ model: MODEL,
{ role: "user", content: "hello", timestamp: 1 }, messages: [{ role: "user" as const, content: "hello", timestamp: 1 }],
{ ...fauxAssistantMessage([]), timestamp: 2 }, options: { sessionId: "provider-session", reasoning: "high" as const },
{ role: "user", content: "go on", timestamp: 3 }, };
];
function request(overrides: Partial<GenerationRequest> = {}): GenerationRequest { /** Warmers on mocked timers; `phase` is what the host reports, `sent` the refreshes the fake models received. */
return { function setup(t: { after(fn: () => void): void }, mode: CacheWarmingMode = "idle") {
model: MODEL, mock.timers.enable({ apis: ["setTimeout", "Date"] });
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 = { const state = {
mode, mode,
now: 0, phase: "streaming" as CacheWarmingPhase | undefined,
running: new Set<ConversationId>([CONVERSATION, SUBAGENT]),
current: true,
alive: 0, alive: 0,
timers: [] as { at: number; fn: () => void; cancelled: boolean }[], sent: [] as ModelsSimpleStreamOptions[],
sent: [] as { model: Model<"pi-messages">; context: Context; options: ModelsSimpleStreamOptions }[], recorded: 0,
recorded: [] as ConversationId[],
reports: [] as unknown[],
/** Resolves the refresh in flight, when `hold` is set. */
release: undefined as (() => void) | undefined,
hold: false,
}; };
const models = { const models = {
completeSimple: async (model: Model<"pi-messages">, context: Context, options: ModelsSimpleStreamOptions) => { completeSimple: async (_model: Model<"pi-messages">, context: Context, options: ModelsSimpleStreamOptions) => {
state.sent.push({ model, context, options }); assert.deepEqual(context.messages, REQUEST.messages);
if (state.hold) state.sent.push(options);
await new Promise<void>((resolve) => { return { ...fauxAssistantMessage([]), stopReason: options.signal?.aborted ? "aborted" : "length" };
state.release = resolve;
});
const reply = fauxAssistantMessage([]);
return { ...reply, stopReason: options.signal?.aborted ? "aborted" : "length" };
}, },
} as unknown as Models; } as unknown as Models;
const warmers = new ConversationWarmers({ const warmers = new ConversationWarmers({
mode: () => state.mode, mode: () => state.mode,
now: () => state.now, phase: async () => state.phase,
setTimer: (fn, delayMs) => { record: async () => {
const timer = { at: state.now + delayMs, fn, cancelled: false }; state.recorded++;
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) => { keepAlive: (delta) => {
state.alive += 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. */ t.after(() => {
const tick = async (): Promise<boolean> => { warmers.close();
const next = state.timers.filter((timer) => !timer.cancelled).sort((a, b) => a.at - b.at)[0]; mock.timers.reset();
if (next === undefined) return false; });
next.cancelled = true; /** Starts warming after a response to a `promptTokens`-token prompt. */
state.now = next.at; const start = (
next.fn(); promptTokens: number,
// The refresh awaits the host and the models a few times. stopReason: StopReason = "toolUse",
for (let turn = 0; turn < 20; turn++) await Promise.resolve(); options: SimpleStreamOptions = REQUEST.options,
return true; ) => {
const message = fauxAssistantMessage([]);
const usage = { ...message.usage, input: 2, cacheRead: promptTokens - 1002, cacheWrite: 1000 };
warmers.start(
CONVERSATION,
{ message: { ...message, stopReason, usage }, request: { ...REQUEST, options } },
models,
);
}; };
const start = (conversationId: ConversationId, promptTokens: number, sent = request(), stop?: StopReason) => /** Advances the clock by `ms` and lets the refreshes it fired finish (they arm the next timer when done). */
warmers.start(conversationId, { message: response(promptTokens, stop), request: sent }, models); const advance = async (ms: number) => {
return { state, warmers, tick, start }; mock.timers.tick(ms);
for (let turn = 0; turn < 10; turn++) await new Promise((resolve) => setImmediate(resolve));
};
/** Advances refresh by refresh until warming stops; returns how many refreshes were sent. */
const runOut = async () => {
for (let refresh = 0; refresh < 100 && state.alive > 0; refresh++) await advance(270_000);
return state.sent.length;
};
return { state, warmers, start, advance, runOut };
} }
test("refreshes before the 5-minute entry expires, with the request's own options and one output token", async () => { test("refreshes every 270 s with the request's options and one output token, and counts the spend", async (t) => {
const { state, tick, start } = setup(); const { state, start, advance } = setup(t);
start(CONVERSATION, 600_000); start(600_000);
assert.equal(state.alive, 1); assert.equal(state.alive, 1);
assert.ok(await tick()); await advance(269_000);
assert.equal(state.now, 270_000); assert.equal(state.sent.length, 0);
await advance(1_000);
assert.equal(state.sent.length, 1); assert.equal(state.sent.length, 1);
const [refresh] = state.sent; const [options] = state.sent;
assert.equal(refresh!.model, MODEL); assert.equal(options!.sessionId, "provider-session");
assert.deepEqual(refresh!.context.messages, MESSAGES); assert.equal(options!.reasoning, "high");
assert.equal(refresh!.options.sessionId, "provider-session"); assert.equal(options!.maxTokens, 1);
assert.equal(refresh!.options.reasoning, "high"); assert.equal(options!.maxRetries, 0);
assert.equal(refresh!.options.maxTokens, 1); assert.equal(options!.disableContinuation, true);
assert.equal(refresh!.options.maxRetries, 0); assert.equal(state.recorded, 1);
assert.equal(refresh!.options.disableContinuation, true); await advance(270_000);
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); assert.equal(state.sent.length, 2);
}); });
test("a 1-hour retention refreshes after 54 minutes", async () => { test("warms up to an hour while streaming and 30 minutes while idle; mode streaming stops when idle", async (t) => {
const { state, tick, start } = setup(); const streaming = setup(t);
start(CONVERSATION, 600_000, request({ options: { cacheRetention: "long" } })); streaming.start(600_000);
assert.ok(await tick()); assert.equal(await streaming.runOut(), 13);
assert.equal(state.now, 3_240_000); mock.timers.reset();
assert.equal(state.sent.length, 1);
const idle = setup(t);
idle.state.phase = "idle";
idle.start(600_000);
assert.equal(await idle.runOut(), 6);
mock.timers.reset();
const streamingOnly = setup(t, "streaming");
streamingOnly.state.phase = "idle";
streamingOnly.start(600_000);
assert.equal(await streamingOnly.runOut(), 0);
}); });
test("streams up to an hour, idles up to 30 minutes, and stops after a settled subagent or in mode streaming", async () => { test("does not warm what is not worth it or cannot be repeated", async (t) => {
// While the run streams (a long tool call), refreshes go on until the hour after the request. const { state, warmers, start, advance } = setup(t);
const streaming = setup(); // A failed or aborted response, a model without a cache lifetime for the retention, mode off: nothing starts.
streaming.start(CONVERSATION, 600_000); start(600_000, "error");
while (await streaming.tick()); start(600_000, "aborted");
assert.equal(streaming.state.sent.length, 13); start(600_000, "toolUse", { cacheRetention: "none" });
assert.ok(streaming.state.now <= 3_600_000); state.mode = "off";
assert.equal(streaming.state.alive, 0); start(600_000);
// 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); assert.equal(state.alive, 0);
state.mode = "idle";
// 10k tokens: a refresh costs more than it is expected to save.
start(10_000);
await advance(270_000);
// 80k tokens: worth it while the run streams, not once only a 15% chance remains that it continues.
start(80_000);
await advance(270_000);
assert.equal(state.sent.length, 1);
state.phase = "idle";
await advance(270_000);
assert.equal(state.sent.length, 1);
assert.equal(state.alive, 0);
// The host says the conversation changed (reset, compaction, model) or ended (a settled subagent): it stops.
state.phase = undefined;
start(600_000);
await advance(270_000);
assert.equal(state.sent.length, 1);
assert.equal(state.alive, 0);
warmers.close();
}); });
test("the next response replaces a conversation's run; close aborts a refresh in flight", async () => { test("the next response replaces the run; close stops it", async (t) => {
const { state, warmers, tick, start } = setup(); const { state, warmers, start, advance } = setup(t);
start(CONVERSATION, 600_000); start(600_000);
start(CONVERSATION, 610_000); await advance(200_000);
start(600_000);
assert.equal(state.alive, 1); assert.equal(state.alive, 1);
assert.equal(state.timers.filter((timer) => !timer.cancelled).length, 1); // The first run's timer is gone: the refresh comes 270 s after the second response.
start(SUBAGENT, 600_000); await advance(200_000);
assert.equal(state.alive, 2); assert.equal(state.sent.length, 0);
await advance(70_000);
state.hold = true;
await tick();
assert.equal(state.sent.length, 1); assert.equal(state.sent.length, 1);
warmers.close(); warmers.close();
assert.equal(state.alive, 0); assert.equal(state.alive, 0);
assert.equal(state.sent[0]!.options.signal?.aborted, true); await advance(600_000);
state.release?.(); assert.equal(state.sent.length, 1);
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( test(
@ -332,26 +224,16 @@ test(
const chat = await holdIn(owner, Chat, { session, conversation: 1 }, "participant"); const chat = await holdIn(owner, Chat, { session, conversation: 1 }, "participant");
await chat.send({ text: "hello", requestId: "m1" }, BACKGROUND_CONTEXT); await chat.send({ text: "hello", requestId: "m1" }, BACKGROUND_CONTEXT);
const live = (view: ConversationView | null | undefined) => view?.docs["pi.live"] as LiveState | undefined; const live = (view: ConversationView | null | undefined) => view?.docs["pi.live"] as LiveState | undefined;
const usage = (view: ConversationView | null | undefined) => const tokens = (view: ConversationView | null | undefined) =>
Object.values((view?.docs["pi.usage"] as unknown as UsageState | undefined)?.models ?? {}).reduce( Object.values((view?.docs["pi.usage"] as unknown as UsageState | undefined)?.models ?? {}).reduce(
(total, model) => total + model.totalTokens, (total, model) => total + model.totalTokens,
0, 0,
); );
const answered = await until(chat.view, (view) => (view?.entries.length ?? 0) > 1 && !live(view)?.run, 30_000); const answered = await until(chat.view, (view) => (view?.entries.length ?? 0) > 1 && !live(view)?.run, 30_000);
const requests = () => const requests = () => readFileSync(requestLog, "utf8").split("\n").filter(Boolean);
readFileSync(requestLog, "utf8") // Idle now: the refreshes repeat the request exactly, and their spend is the conversation's.
.split("\n") await until(chat.view, (view) => tokens(view) > tokens(answered) && requests().length >= 3, 20_000);
.filter((line) => line.length > 0); const [sent, ...refreshes] = requests();
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); 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);
}, },
); );

View file

@ -4,7 +4,7 @@
### Breaking Changes ### 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 }) => ...`. - `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), 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 ### Added

View file

@ -3402,8 +3402,6 @@ type GenerationRequest = {
readonly messages: readonly Message[]; readonly messages: readonly Message[];
/** The request's options without the attempt's `signal`, `telemetry` and `onPayload`. */ /** The request's options without the attempt's `signal`, `telemetry` and `onPayload`. */
readonly options: SimpleStreamOptions; 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. */ /** `request` is `undefined` for a polled deferred response, which sent no request. */
@ -4937,8 +4935,8 @@ The run's inputs live in `pi.live.run`, not in the task input.
- `afterResponse` observes every terminal message, from `request` or `poll`, - `afterResponse` observes every terminal message, from `request` or `poll`,
before classification; a still deferred result is not terminal. It gets the before classification; a still deferred result is not terminal. It gets the
request the attempt sent (model, messages after `beforeRequest`, options request the attempt sent (model, messages after `beforeRequest`, options
without the attempt's signal, telemetry and payload callback, and whether a without the attempt's signal, telemetry and payload callback); a polled
`beforeProviderRequest` hook replaced the body); a polled response has none. response has none.
- Before classifying, the handler stops the partial throttle and awaits any - Before classifying, the handler stops the partial throttle and awaits any
partial commit in flight, so no stale partial lands after the outcome. The 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: terminal message is classified in one commit that also clears the partial:

View file

@ -296,19 +296,14 @@ export const GenerationTask = defineTask<GenerationInput, GenerationCheckpoint,
sessionId: await ensureProviderSessionId(runtime, context), sessionId: await ensureProviderSessionId(runtime, context),
...(thinkingLevel === "off" ? {} : { reasoning: thinkingLevel }), ...(thinkingLevel === "off" ? {} : { reasoning: thinkingLevel }),
}; };
let payloadReplaced = false;
const options: SimpleStreamOptions = { const options: SimpleStreamOptions = {
...sentOptions, ...sentOptions,
signal: runtime.signal, signal: runtime.signal,
telemetry: context.telemetry, telemetry: context.telemetry,
onPayload: async (payload) => { onPayload: (payload) => providerPayload(runtime, model, payload, context),
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 message = await streamResponse(runtime, model, messages, options, attempt, context);
const sent = { model, messages, options: sentOptions, payloadReplaced }; const sent = { model, messages, options: sentOptions };
const request = { attempt, compacted, model: ref, thinkingLevel, cutoff, messages: view.messages, sent }; const request = { attempt, compacted, model: ref, thinkingLevel, cutoff, messages: view.messages, sent };
await classify(runtime, request, message, context); await classify(runtime, request, message, context);
}, },
@ -679,23 +674,19 @@ export async function convertPartial(tx: Tx, live: Draft<LiveState>, conversatio
await appendAssistant(tx, conversationId, { ...message, stopReason: "aborted" }); await appendAssistant(tx, conversationId, { ...message, stopReason: "aborted" });
} }
/** The `beforeProviderRequest` replacement chain over one provider request body; `undefined` when no hook replaced it. */ /** The `beforeProviderRequest` replacement chain over one provider request body. */
async function providerPayload( async function providerPayload(
runtime: Runtime, runtime: Runtime,
model: Model<Api>, model: Model<Api>,
payload: unknown, payload: unknown,
context: Context, context: Context,
): Promise<unknown> { ): Promise<unknown> {
let current: { readonly payload: unknown } | undefined; let current = payload;
await runtime.hooks.each("beforeProviderRequest", async (hook) => { await runtime.hooks.each("beforeProviderRequest", async (hook) => {
const replaced = await hook( const replaced = await hook({ model, payload: current }, runtime, context);
{ model, payload: current === undefined ? payload : current.payload }, if (replaced !== undefined) current = replaced.payload;
runtime,
context,
);
if (replaced !== undefined) current = replaced;
}); });
return current?.payload; return current;
} }
/** /**

View file

@ -982,8 +982,8 @@ export interface HookApi extends DocumentReader {
export type HookResult<T> = T | undefined | Promise<T | undefined>; export type HookResult<T> = T | undefined | Promise<T | undefined>;
/** /**
* A provider request one generation attempt sent, as it was passed to `models.streamSimple()`: sending it again with the * A provider request one generation attempt sent, as it was passed to `models.streamSimple()`: sending it again repeats
* same model, messages and options repeats it, which keeps its prompt cache entry alive (cache warming). * it (without the `beforeProviderRequest` chain), which keeps its prompt cache entry alive (cache warming).
*/ */
export type GenerationRequest = { export type GenerationRequest = {
readonly model: Model<Api>; readonly model: Model<Api>;
@ -991,11 +991,6 @@ export type GenerationRequest = {
readonly messages: readonly Message[]; readonly messages: readonly Message[];
/** The request's options without the attempt's `signal`, `telemetry` and `onPayload`. */ /** The request's options without the attempt's `signal`, `telemetry` and `onPayload`. */
readonly options: SimpleStreamOptions; 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`. */ /** A terminal provider message for `afterResponse`. */

View file

@ -793,45 +793,27 @@ describe("generation hooks", () => {
await harness.close(context); await harness.close(context);
}); });
it("gives afterResponse the request each attempt sent", async () => { it("gives afterResponse the request the attempt sent", async () => {
const setup = chatSetup(); const setup = chatSetup();
addHooks(setup.registry, GenerationTask, { addHooks(setup.registry, GenerationTask, {
beforeRequest: ({ messages }) => ({ beforeRequest: ({ messages }) => ({
messages: [...messages, { role: "user", content: "injected", timestamp: 0 }], messages: [...messages, { role: "user", content: "injected", timestamp: 0 }],
}), }),
}); });
let replace = false; const requests: (GenerationRequest | undefined)[] = [];
addHooks(setup.registry, GenerationTask, { addHooks(setup.registry, GenerationTask, { afterResponse: ({ request }) => void requests.push(request) });
beforeProviderRequest: ({ payload }) => const { harness, status } = await run(setup, [DONE]);
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"); expect(status).toBe("done");
const [first, second] = requests; const [request] = requests;
expect(`${first!.model.provider}/${first!.model.id}`).toBe("faux/faux-1"); expect(`${request!.model.provider}/${request!.model.id}`).toBe("faux/faux-1");
// The messages are the ones sent: after the `beforeRequest` chain. // The messages after the `beforeRequest` chain, and the request's options without the attempt's own.
expect(textOf(first!.messages.at(-1) as never)).toBe("injected"); expect(textOf(request!.messages.at(-1) as never)).toBe("injected");
// The options are the request's own, without what belongs to the attempt. expect(typeof request!.options.sessionId).toBe("string");
expect(typeof first!.options.sessionId).toBe("string"); expect([request!.options.signal, request!.options.onPayload, request!.options.telemetry]).toEqual([
expect(first!.options.sessionId).toBe(second!.options.sessionId); undefined,
expect(first!.options.signal).toBeUndefined(); undefined,
expect(first!.options.onPayload).toBeUndefined(); undefined,
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); await harness.close(context);
}); });