mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
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
parent
dd8e0df4a9
commit
b76afbdb2f
10 changed files with 234 additions and 445 deletions
|
|
@ -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.
|
||||
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
|
||||
starts that conversation's run, replacing the previous one. 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`.
|
||||
conversation's `pi.usage` (`recordUsage`). The policy is pi-ai's, shared with pi-coding-agent: up to an hour after the
|
||||
request while the run streams, 30 minutes once idle, a refresh only when its expected savings reach $0.05
|
||||
(`evaluateCacheWarming`). Each refresh reads the conversation's phase: streaming while `pi.live.run` exists, else
|
||||
idle for top-level conversations (`isTopLevel`; a settled subagent or one-shot agent is not warmed), and nothing
|
||||
once the next request would start differently (another model, or a context that no longer starts with the sent
|
||||
messages by role and timestamp: a reset, compaction or edit). Timers live in `CacheWarmingRuntime`'s native, 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. `<P2_DATA_DIR>/mcp.json`
|
||||
is the owner's `mcpServers` (pi-mcp's `validateMcpServerConfig`, plus p2's `shared` and `runOn`; invalid entries are
|
||||
|
|
|
|||
|
|
@ -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`) | 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 |
|
||||
| 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). 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. |
|
||||
| `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. |
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
// 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.
|
||||
// long as that is expected to pay off. The policy is pi-ai's; this file schedules it per conversation.
|
||||
import {
|
||||
type AssistantMessage,
|
||||
CACHE_WARMING_MAX_AGE_MS,
|
||||
|
|
@ -12,7 +11,6 @@ import {
|
|||
getCacheWarmingRequestOptions,
|
||||
getPromptCacheTtlMs,
|
||||
isCacheWarmingReplayable,
|
||||
type Message,
|
||||
type Models,
|
||||
} from "@earendil-works/pi-ai";
|
||||
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. */
|
||||
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<boolean>;
|
||||
/** 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>;
|
||||
/**
|
||||
* "streaming" while the conversation's run is under way, "idle" after it, or `undefined` when warming the
|
||||
* conversation is pointless: its next request will not start with `request`, or nobody is expected to continue it.
|
||||
*/
|
||||
phase(conversationId: ConversationId, request: GenerationRequest): Promise<CacheWarmingPhase | undefined>;
|
||||
/** Adds a refresh's spend to the conversation's usage. */
|
||||
record(conversationId: ConversationId, message: AssistantMessage): Promise<void>;
|
||||
/** Counts the conversations being warmed; a count keeps the worker process alive. */
|
||||
keepAlive(delta: number): void;
|
||||
report(error: unknown): void;
|
||||
};
|
||||
|
||||
type Run = {
|
||||
|
|
@ -45,20 +37,16 @@ type Run = {
|
|||
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;
|
||||
timer?: ReturnType<typeof setTimeout>;
|
||||
};
|
||||
|
||||
/**
|
||||
* 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()`.
|
||||
* One warming run per conversation, started by each response and replaced by the next one. Each refresh asks the host
|
||||
* for the conversation's phase and refreshes while the request is at most an hour old (30 minutes once idle) and the
|
||||
* refresh pays off by pi-ai's `evaluateCacheWarming()`.
|
||||
*/
|
||||
export class ConversationWarmers {
|
||||
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. */
|
||||
start(conversationId: ConversationId, response: GenerationResponse, models: Models): void {
|
||||
start(conversationId: ConversationId, { message, request }: 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.
|
||||
// A polled deferred response sent no request; a failed or aborted one ended its run.
|
||||
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;
|
||||
if (this.#host.mode() === "off" || !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;
|
||||
if (delayMs === undefined) return;
|
||||
const { input, cacheRead, cacheWrite } = message.usage;
|
||||
const run: Run = {
|
||||
request,
|
||||
models,
|
||||
promptTokens: usage.input + usage.cacheRead + usage.cacheWrite,
|
||||
ttlMs,
|
||||
promptTokens: input + cacheRead + cacheWrite,
|
||||
delayMs,
|
||||
startedAt: this.#host.now(),
|
||||
startedAt: Date.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);
|
||||
run.timer = setTimeout(() => void this.#refresh(conversationId, run), delayMs);
|
||||
}
|
||||
|
||||
/** Stop warming the conversation; a refresh in flight is aborted. */
|
||||
|
|
@ -110,7 +84,7 @@ export class ConversationWarmers {
|
|||
const run = this.#runs.get(conversationId);
|
||||
if (run === undefined) return;
|
||||
this.#runs.delete(conversationId);
|
||||
run.cancelTimer?.();
|
||||
clearTimeout(run.timer);
|
||||
run.controller.abort();
|
||||
this.#host.keepAlive(-1);
|
||||
}
|
||||
|
|
@ -120,24 +94,9 @@ export class ConversationWarmers {
|
|||
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> {
|
||||
run.cancelTimer = undefined;
|
||||
try {
|
||||
if (!(await this.#shouldRefresh(conversationId, run))) {
|
||||
if (!(await this.#wanted(conversationId, run))) {
|
||||
if (this.#runs.get(conversationId) === run) this.stop(conversationId);
|
||||
return;
|
||||
}
|
||||
|
|
@ -152,35 +111,21 @@ export class ConversationWarmers {
|
|||
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 (!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. */
|
||||
async #shouldRefresh(conversationId: ConversationId, run: Run): Promise<boolean> {
|
||||
/** Whether this refresh is still wanted: the run current, in its phase's window, and paying off. */
|
||||
async #wanted(conversationId: ConversationId, run: Run): Promise<boolean> {
|
||||
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;
|
||||
if (mode === "off") return false;
|
||||
const phase = await this.#host.phase(conversationId, run.request);
|
||||
if (phase === undefined || (phase === "idle" && mode === "streaming")) return false;
|
||||
const maxAge = phase === "idle" ? CACHE_WARMING_MAX_IDLE_AGE_MS : CACHE_WARMING_MAX_AGE_MS;
|
||||
if (Date.now() > run.startedAt + maxAge) return false;
|
||||
const { action } = evaluateCacheWarming(run.request.model, run.promptTokens, phase);
|
||||
return 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,
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,12 +1,13 @@
|
|||
// 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.
|
||||
// Cache warming in the session worker: a Durable extension whose `afterResponse` hook starts warming each
|
||||
// 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 { 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";
|
||||
import { ConversationWarmers } from "./warmer.ts";
|
||||
|
||||
export default defineFacet((host) => {
|
||||
host.serve(CacheWarmingRuntime, (_params, self) => {
|
||||
|
|
@ -14,44 +15,35 @@ export default defineFacet((host) => {
|
|||
const worker = workerOf(self);
|
||||
const settings = self.get(WorkerModels, {});
|
||||
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
|
||||
// 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
|
||||
phase: async (conversationId, request) => {
|
||||
const running = (await durable().snapshot(LiveDoc, conversationId, context))?.run !== undefined;
|
||||
// After its run, a subagent's or one-shot agent's 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) => {
|
||||
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";
|
||||
},
|
||||
record: async (conversationId, { provider, model, usage }) => {
|
||||
await durable().commit(
|
||||
(tx) =>
|
||||
recordUsage(
|
||||
tx,
|
||||
conversationId,
|
||||
"models",
|
||||
`${message.provider}/${message.model}`,
|
||||
message.usage,
|
||||
),
|
||||
BACKGROUND_CONTEXT,
|
||||
(tx) => recordUsage(tx, conversationId, "models", `${provider}/${model}`, usage),
|
||||
context,
|
||||
);
|
||||
},
|
||||
keepAlive: (delta) => worker().liveness.add("cacheWarming", delta),
|
||||
report: (error) => console.error("p2 cache-warming: refreshing the prompt cache failed:", error),
|
||||
}),
|
||||
(opened) => opened.close(),
|
||||
);
|
||||
|
|
@ -60,14 +52,10 @@ export default defineFacet((host) => {
|
|||
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.
|
||||
// `cacheWarming: "off"` uninstalls the hook live; runs already warming stop at their next refresh.
|
||||
self.contribute(
|
||||
ExtensionSwitches,
|
||||
self.derive([settings.settings], ([value]) => ({ [extension.name]: value.cacheWarming !== "off" })),
|
||||
|
|
@ -76,3 +64,13 @@ export default defineFacet((host) => {
|
|||
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,
|
||||
)
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,35 +1,27 @@
|
|||
// 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.
|
||||
// Cache warming (plugins/cache-warming): a conversation's last request is sent again with one output token before its
|
||||
// prompt cache entry expires, while pi's policy says it pays off. The warmers run on mocked timers with fake models;
|
||||
// a worker runs them end to end 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 { mock, test } from "node:test";
|
||||
import type {
|
||||
AssistantMessage,
|
||||
CacheWarmingPhase,
|
||||
Context,
|
||||
Message,
|
||||
Model,
|
||||
Models,
|
||||
ModelsSimpleStreamOptions,
|
||||
SimpleStreamOptions,
|
||||
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 type { ConversationId, ConversationView, 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 { type CacheWarmingMode, ConversationWarmers } 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"> = {
|
||||
|
|
@ -46,254 +38,154 @@ const MODEL: Model<"pi-messages"> = {
|
|||
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> = {}): GenerationRequest {
|
||||
return {
|
||||
const REQUEST = {
|
||||
model: MODEL,
|
||||
messages: MESSAGES,
|
||||
options: { sessionId: "provider-session", reasoning: "high" },
|
||||
payloadReplaced: false,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
messages: [{ role: "user" as const, content: "hello", timestamp: 1 }],
|
||||
options: { sessionId: "provider-session", reasoning: "high" as const },
|
||||
};
|
||||
|
||||
/** 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") {
|
||||
/** Warmers on mocked timers; `phase` is what the host reports, `sent` the refreshes the fake models received. */
|
||||
function setup(t: { after(fn: () => void): void }, mode: CacheWarmingMode = "idle") {
|
||||
mock.timers.enable({ apis: ["setTimeout", "Date"] });
|
||||
const state = {
|
||||
mode,
|
||||
now: 0,
|
||||
running: new Set<ConversationId>([CONVERSATION, SUBAGENT]),
|
||||
current: true,
|
||||
phase: "streaming" as CacheWarmingPhase | undefined,
|
||||
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,
|
||||
sent: [] as ModelsSimpleStreamOptions[],
|
||||
recorded: 0,
|
||||
};
|
||||
const models = {
|
||||
completeSimple: async (model: Model<"pi-messages">, context: Context, options: ModelsSimpleStreamOptions) => {
|
||||
state.sent.push({ model, context, options });
|
||||
if (state.hold)
|
||||
await new Promise<void>((resolve) => {
|
||||
state.release = resolve;
|
||||
});
|
||||
const reply = fauxAssistantMessage([]);
|
||||
return { ...reply, stopReason: options.signal?.aborted ? "aborted" : "length" };
|
||||
completeSimple: async (_model: Model<"pi-messages">, context: Context, options: ModelsSimpleStreamOptions) => {
|
||||
assert.deepEqual(context.messages, REQUEST.messages);
|
||||
state.sent.push(options);
|
||||
return { ...fauxAssistantMessage([]), 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;
|
||||
};
|
||||
phase: async () => state.phase,
|
||||
record: async () => {
|
||||
state.recorded++;
|
||||
},
|
||||
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<boolean> => {
|
||||
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;
|
||||
t.after(() => {
|
||||
warmers.close();
|
||||
mock.timers.reset();
|
||||
});
|
||||
/** Starts warming after a response to a `promptTokens`-token prompt. */
|
||||
const start = (
|
||||
promptTokens: number,
|
||||
stopReason: StopReason = "toolUse",
|
||||
options: SimpleStreamOptions = REQUEST.options,
|
||||
) => {
|
||||
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) =>
|
||||
warmers.start(conversationId, { message: response(promptTokens, stop), request: sent }, models);
|
||||
return { state, warmers, tick, start };
|
||||
/** Advances the clock by `ms` and lets the refreshes it fired finish (they arm the next timer when done). */
|
||||
const advance = async (ms: number) => {
|
||||
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 () => {
|
||||
const { state, tick, start } = setup();
|
||||
start(CONVERSATION, 600_000);
|
||||
test("refreshes every 270 s with the request's options and one output token, and counts the spend", async (t) => {
|
||||
const { state, start, advance } = setup(t);
|
||||
start(600_000);
|
||||
assert.equal(state.alive, 1);
|
||||
assert.ok(await tick());
|
||||
assert.equal(state.now, 270_000);
|
||||
await advance(269_000);
|
||||
assert.equal(state.sent.length, 0);
|
||||
await advance(1_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);
|
||||
const [options] = state.sent;
|
||||
assert.equal(options!.sessionId, "provider-session");
|
||||
assert.equal(options!.reasoning, "high");
|
||||
assert.equal(options!.maxTokens, 1);
|
||||
assert.equal(options!.maxRetries, 0);
|
||||
assert.equal(options!.disableContinuation, true);
|
||||
assert.equal(state.recorded, 1);
|
||||
await advance(270_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("warms up to an hour while streaming and 30 minutes while idle; mode streaming stops when idle", async (t) => {
|
||||
const streaming = setup(t);
|
||||
streaming.start(600_000);
|
||||
assert.equal(await streaming.runOut(), 13);
|
||||
mock.timers.reset();
|
||||
|
||||
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 () => {
|
||||
// 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);
|
||||
test("does not warm what is not worth it or cannot be repeated", async (t) => {
|
||||
const { state, warmers, start, advance } = setup(t);
|
||||
// A failed or aborted response, a model without a cache lifetime for the retention, mode off: nothing starts.
|
||||
start(600_000, "error");
|
||||
start(600_000, "aborted");
|
||||
start(600_000, "toolUse", { cacheRetention: "none" });
|
||||
state.mode = "off";
|
||||
start(600_000);
|
||||
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 () => {
|
||||
const { state, warmers, tick, start } = setup();
|
||||
start(CONVERSATION, 600_000);
|
||||
start(CONVERSATION, 610_000);
|
||||
test("the next response replaces the run; close stops it", async (t) => {
|
||||
const { state, warmers, start, advance } = setup(t);
|
||||
start(600_000);
|
||||
await advance(200_000);
|
||||
start(600_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();
|
||||
// The first run's timer is gone: the refresh comes 270 s after the second response.
|
||||
await advance(200_000);
|
||||
assert.equal(state.sent.length, 0);
|
||||
await advance(70_000);
|
||||
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/);
|
||||
await advance(600_000);
|
||||
assert.equal(state.sent.length, 1);
|
||||
});
|
||||
|
||||
test(
|
||||
|
|
@ -332,26 +224,16 @@ test(
|
|||
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) =>
|
||||
const tokens = (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}`);
|
||||
const requests = () => readFileSync(requestLog, "utf8").split("\n").filter(Boolean);
|
||||
// Idle now: the refreshes repeat the request exactly, and their spend is the conversation's.
|
||||
await until(chat.view, (view) => tokens(view) > tokens(answered) && requests().length >= 3, 20_000);
|
||||
const [sent, ...refreshes] = requests();
|
||||
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);
|
||||
},
|
||||
);
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
|
||||
### 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
|
||||
|
||||
|
|
|
|||
|
|
@ -3402,8 +3402,6 @@ type GenerationRequest = {
|
|||
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. */
|
||||
|
|
@ -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`,
|
||||
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.
|
||||
without the attempt's signal, telemetry and payload callback); 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:
|
||||
|
|
|
|||
|
|
@ -296,19 +296,14 @@ export const GenerationTask = defineTask<GenerationInput, GenerationCheckpoint,
|
|||
sessionId: await ensureProviderSessionId(runtime, context),
|
||||
...(thinkingLevel === "off" ? {} : { reasoning: thinkingLevel }),
|
||||
};
|
||||
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;
|
||||
},
|
||||
onPayload: (payload) => providerPayload(runtime, model, payload, 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 };
|
||||
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" });
|
||||
}
|
||||
|
||||
/** 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(
|
||||
runtime: Runtime,
|
||||
model: Model<Api>,
|
||||
payload: unknown,
|
||||
context: Context,
|
||||
): Promise<unknown> {
|
||||
let current: { readonly payload: unknown } | undefined;
|
||||
let current = payload;
|
||||
await runtime.hooks.each("beforeProviderRequest", async (hook) => {
|
||||
const replaced = await hook(
|
||||
{ model, payload: current === undefined ? payload : current.payload },
|
||||
runtime,
|
||||
context,
|
||||
);
|
||||
if (replaced !== undefined) current = replaced;
|
||||
const replaced = await hook({ model, payload: current }, runtime, context);
|
||||
if (replaced !== undefined) current = replaced.payload;
|
||||
});
|
||||
return current?.payload;
|
||||
return current;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -982,8 +982,8 @@ export interface HookApi extends DocumentReader {
|
|||
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
|
||||
* same model, messages and options repeats it, which keeps its prompt cache entry alive (cache warming).
|
||||
* A provider request one generation attempt sent, as it was passed to `models.streamSimple()`: sending it again repeats
|
||||
* it (without the `beforeProviderRequest` chain), which keeps its prompt cache entry alive (cache warming).
|
||||
*/
|
||||
export type GenerationRequest = {
|
||||
readonly model: Model<Api>;
|
||||
|
|
@ -991,11 +991,6 @@ export type GenerationRequest = {
|
|||
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`. */
|
||||
|
|
|
|||
|
|
@ -793,45 +793,27 @@ describe("generation hooks", () => {
|
|||
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();
|
||||
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]);
|
||||
const requests: (GenerationRequest | undefined)[] = [];
|
||||
addHooks(setup.registry, GenerationTask, { afterResponse: ({ request }) => void requests.push(request) });
|
||||
const { harness, status } = await run(setup, [DONE]);
|
||||
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]);
|
||||
const [request] = requests;
|
||||
expect(`${request!.model.provider}/${request!.model.id}`).toBe("faux/faux-1");
|
||||
// The messages after the `beforeRequest` chain, and the request's options without the attempt's own.
|
||||
expect(textOf(request!.messages.at(-1) as never)).toBe("injected");
|
||||
expect(typeof request!.options.sessionId).toBe("string");
|
||||
expect([request!.options.signal, request!.options.onPayload, request!.options.telemetry]).toEqual([
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
]);
|
||||
await harness.close(context);
|
||||
});
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue