From 65b84ea1b46c845951ff752684c32bb3304d71d8 Mon Sep 17 00:00:00 2001 From: Iury Souza Date: Sun, 26 Jul 2026 11:57:19 +0200 Subject: [PATCH] feat(cache): integrate predictor footer status Advertise priority 200 to the optional pi-ext footer host while retaining native setStatus behavior. Keep predictions through failures and clear only after a successful response on the matching lane. --- packages/pi-cache-hit-predictor/README.md | 18 +- packages/pi-cache-hit-predictor/index.ts | 265 ++++++++++-------- .../pi-cache-hit-predictor/src/footer-slot.ts | 27 ++ .../tests/extension.test.ts | 204 +++++++++----- 4 files changed, 323 insertions(+), 191 deletions(-) create mode 100644 packages/pi-cache-hit-predictor/src/footer-slot.ts diff --git a/packages/pi-cache-hit-predictor/README.md b/packages/pi-cache-hit-predictor/README.md index 9f17b77..b975059 100644 --- a/packages/pi-cache-hit-predictor/README.md +++ b/packages/pi-cache-hit-predictor/README.md @@ -1,23 +1,25 @@ -# @howaboua/pi-cache-hit-predictor +# @iurysza/pi-cache-hit-predictor -Shows an inline cache-hit prediction when you switch Pi models or reasoning levels. +Shows a cache-hit prediction when you switch Pi models or reasoning levels. ```text Cache hit prediction · gpt-5.6-sol · low: ~27k / ~104k (26%) cached ``` -The prediction is a UI-only transcript notification. It is not sent to the model and does not change the prompt. Back-to-back predictions update the same line while you cycle through models or reasoning levels. +The prediction is UI-only. It is not sent to the model and does not change the prompt. It remains visible through aborts and failed requests, then clears after the first successful response on the predicted provider/API/model/reasoning lane. + +The package uses Pi's native `setStatus()` API when installed alone. If `@iurysza/pi-ext` is also installed, it advertises priority metadata so pi-ext can place the same status in its bounded auxiliary footer line. ## Install ```bash -pi install npm:@howaboua/pi-cache-hit-predictor +pi install npm:@iurysza/pi-cache-hit-predictor ``` Try it for one session: ```bash -pi -e npm:@howaboua/pi-cache-hit-predictor +pi -e npm:@iurysza/pi-cache-hit-predictor ``` ## How it works @@ -31,8 +33,8 @@ This is an estimate, not provider preflight data. Provider expiry or eviction, c ## Local development ```bash -bun install -bun run check -bun run pack:dry +npm install +npm run check +npm pack --dry-run pi -e ./index.ts ``` diff --git a/packages/pi-cache-hit-predictor/index.ts b/packages/pi-cache-hit-predictor/index.ts index 20fa82a..58edcb7 100644 --- a/packages/pi-cache-hit-predictor/index.ts +++ b/packages/pi-cache-hit-predictor/index.ts @@ -1,132 +1,169 @@ import type { - ExtensionAPI, - ExtensionContext, + ExtensionAPI, + ExtensionContext, } from "@earendil-works/pi-coding-agent"; import { - type CacheLane, - type CachePrediction, - predictCacheHit, - recordAssistantUsage, - scanCacheHistory, + type CacheLane, + type CachePrediction, + predictCacheHit, + recordAssistantUsage, + scanCacheHistory, } from "./src/predictor.js"; +import { createFooterSlotRegistration } from "./src/footer-slot.js"; + +const STATUS_KEY = "pi-cache-hit-predictor"; interface ModelIdentity { - provider: string; - api: string; - id: string; + provider: string; + api: string; + id: string; } function formatTokens(tokens: number): string { - if (tokens < 1_000) return Math.round(tokens).toString(); - if (tokens < 1_000_000) { - return `${(tokens / 1_000).toFixed(tokens < 10_000 ? 1 : 0)}k`; - } - return `${(tokens / 1_000_000).toFixed(tokens < 10_000_000 ? 1 : 0)}m`; + if (tokens < 1_000) return Math.round(tokens).toString(); + if (tokens < 1_000_000) { + return `${(tokens / 1_000).toFixed(tokens < 10_000 ? 1 : 0)}k`; + } + return `${(tokens / 1_000_000).toFixed(tokens < 10_000_000 ? 1 : 0)}m`; } function predictionText(prediction: CachePrediction): string { - const lane = `${prediction.lane.model} · ${prediction.lane.thinkingLevel}`; - if (!prediction.hasLaneHistory) { - const prompt = prediction.currentPromptTokens - ? ` of ~${formatTokens(prediction.currentPromptTokens)}` - : ""; - return `Cache hit prediction · ${lane}: cold lane (0%${prompt})`; - } - - if (prediction.currentPromptTokens === null || prediction.percent === null) { - return `Cache hit prediction · ${lane}: ~${formatTokens(prediction.estimatedCacheTokens)} cached`; - } - return `Cache hit prediction · ${lane}: ~${formatTokens(prediction.estimatedCacheTokens)} / ~${formatTokens(prediction.currentPromptTokens)} (${Math.round(prediction.percent)}%) cached`; + const lane = `${prediction.lane.model} · ${prediction.lane.thinkingLevel}`; + if (!prediction.hasLaneHistory) { + const prompt = prediction.currentPromptTokens + ? ` of ~${formatTokens(prediction.currentPromptTokens)}` + : ""; + return `Cache hit prediction · ${lane}: cold lane (0%${prompt})`; + } + + if (prediction.currentPromptTokens === null || prediction.percent === null) { + return `Cache hit prediction · ${lane}: ~${formatTokens(prediction.estimatedCacheTokens)} cached`; + } + return `Cache hit prediction · ${lane}: ~${formatTokens(prediction.estimatedCacheTokens)} / ~${formatTokens(prediction.currentPromptTokens)} (${Math.round(prediction.percent)}%) cached`; } function laneFor(model: ModelIdentity, thinkingLevel: string): CacheLane { - return { - provider: model.provider, - api: model.api, - model: model.id, - thinkingLevel, - }; + return { + provider: model.provider, + api: model.api, + model: model.id, + thinkingLevel, + }; +} + +function sameLane(left: CacheLane, right: CacheLane): boolean { + return left.provider === right.provider + && left.api === right.api + && left.model === right.model + && left.thinkingLevel === right.thinkingLevel; } -export default function (pi: ExtensionAPI) { - let history = scanCacheHistory([]); - let pendingPredictionTimer: ReturnType | undefined; - - const rebuild = (ctx: ExtensionContext) => { - history = scanCacheHistory( - ctx.sessionManager.getBranch(), - pi.getThinkingLevel(), - ); - }; - - const appendPrediction = ( - ctx: ExtensionContext, - model: ModelIdentity, - thinkingLevel: string, - ) => { - if (ctx.mode !== "tui") return; - const contextTokens = ctx.getContextUsage()?.tokens ?? null; - const prediction = predictCacheHit( - history, - laneFor(model, thinkingLevel), - contextTokens, - ); - ctx.ui.notify(predictionText(prediction), "info"); - }; - - const schedulePrediction = ( - ctx: ExtensionContext, - model: ModelIdentity, - thinkingLevel: string, - ) => { - if (pendingPredictionTimer) clearTimeout(pendingPredictionTimer); - pendingPredictionTimer = setTimeout(() => { - pendingPredictionTimer = undefined; - appendPrediction(ctx, model, thinkingLevel); - }, 0); - }; - - pi.on("session_start", async (_event, ctx) => rebuild(ctx)); - pi.on("session_tree", async (_event, ctx) => rebuild(ctx)); - pi.on("session_compact", async (_event, ctx) => rebuild(ctx)); - - pi.on("message_end", async (event, ctx) => { - if (event.message.role !== "assistant") return; - const selected = ctx.model; - recordAssistantUsage( - history, - event.message, - laneFor( - selected?.provider === event.message.provider && - selected.id === event.message.model && - selected.api === event.message.api - ? selected - : { - provider: event.message.provider, - api: event.message.api, - id: event.message.model, - }, - pi.getThinkingLevel(), - ), - ); - }); - - pi.on("thinking_level_select", async (event, ctx) => { - if (event.level === event.previousLevel || !ctx.model) return; - schedulePrediction(ctx, ctx.model, event.level); - }); - - pi.on("model_select", async (event, ctx) => { - if (event.source === "restore" || !event.previousModel) { - if (pendingPredictionTimer) clearTimeout(pendingPredictionTimer); - pendingPredictionTimer = undefined; - return; - } - schedulePrediction(ctx, event.model, pi.getThinkingLevel()); - }); - - pi.on("session_shutdown", async () => { - if (pendingPredictionTimer) clearTimeout(pendingPredictionTimer); - pendingPredictionTimer = undefined; - }); +export default function cacheHitPredictor(pi: ExtensionAPI) { + const footerSlot = createFooterSlotRegistration(pi.events, STATUS_KEY, 200); + let history = scanCacheHistory([]); + let pendingPredictionTimer: ReturnType | undefined; + let displayedLane: CacheLane | undefined; + + const rebuild = (ctx: ExtensionContext) => { + history = scanCacheHistory( + ctx.sessionManager.getBranch(), + pi.getThinkingLevel(), + ); + }; + + const clearPrediction = (ctx: ExtensionContext) => { + displayedLane = undefined; + ctx.ui.setStatus(STATUS_KEY, undefined); + }; + + const appendPrediction = ( + ctx: ExtensionContext, + model: ModelIdentity, + thinkingLevel: string, + ) => { + if (ctx.mode !== "tui") return; + const contextTokens = ctx.getContextUsage()?.tokens ?? null; + const prediction = predictCacheHit( + history, + laneFor(model, thinkingLevel), + contextTokens, + ); + displayedLane = prediction.lane; + ctx.ui.setStatus(STATUS_KEY, predictionText(prediction)); + }; + + const schedulePrediction = ( + ctx: ExtensionContext, + model: ModelIdentity, + thinkingLevel: string, + ) => { + if (pendingPredictionTimer) clearTimeout(pendingPredictionTimer); + pendingPredictionTimer = setTimeout(() => { + pendingPredictionTimer = undefined; + appendPrediction(ctx, model, thinkingLevel); + }, 0); + }; + + pi.on("session_start", async (_event, ctx) => { + footerSlot.register(); + clearPrediction(ctx); + rebuild(ctx); + }); + pi.on("session_tree", async (_event, ctx) => { + clearPrediction(ctx); + rebuild(ctx); + }); + pi.on("session_compact", async (_event, ctx) => { + clearPrediction(ctx); + rebuild(ctx); + }); + + pi.on("message_end", async (event, ctx) => { + if (event.message.role !== "assistant") return; + const selected = ctx.model; + recordAssistantUsage( + history, + event.message, + laneFor( + selected?.provider === event.message.provider + && selected.id === event.message.model + && selected.api === event.message.api + ? selected + : { + provider: event.message.provider, + api: event.message.api, + id: event.message.model, + }, + pi.getThinkingLevel(), + ), + ); + }); + + pi.on("thinking_level_select", async (event, ctx) => { + if (event.level === event.previousLevel || !ctx.model) return; + schedulePrediction(ctx, ctx.model, event.level); + }); + + pi.on("model_select", async (event, ctx) => { + if (event.source === "restore" || !event.previousModel) { + if (pendingPredictionTimer) clearTimeout(pendingPredictionTimer); + pendingPredictionTimer = undefined; + clearPrediction(ctx); + return; + } + schedulePrediction(ctx, event.model, pi.getThinkingLevel()); + }); + + pi.on("after_provider_response", async (event, ctx) => { + if (event.status < 200 || event.status >= 300 || !displayedLane || !ctx.model) return; + const destinationLane = laneFor(ctx.model, pi.getThinkingLevel()); + if (sameLane(displayedLane, destinationLane)) clearPrediction(ctx); + }); + + pi.on("session_shutdown", async (_event, ctx) => { + if (pendingPredictionTimer) clearTimeout(pendingPredictionTimer); + pendingPredictionTimer = undefined; + clearPrediction(ctx); + footerSlot.dispose(); + }); } diff --git a/packages/pi-cache-hit-predictor/src/footer-slot.ts b/packages/pi-cache-hit-predictor/src/footer-slot.ts new file mode 100644 index 0000000..a0eca03 --- /dev/null +++ b/packages/pi-cache-hit-predictor/src/footer-slot.ts @@ -0,0 +1,27 @@ +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; + +const PROTOCOL_VERSION = 1 as const; +const HOST_READY = "@iurysza/pi-ext/footer-slot/ready/v1"; +const REGISTER = "@iurysza/pi-ext/footer-slot/register/v1"; +const UNREGISTER = "@iurysza/pi-ext/footer-slot/unregister/v1"; + +export function createFooterSlotRegistration( + events: ExtensionAPI["events"], + id: string, + priority: number, +) { + const payload = { protocolVersion: PROTOCOL_VERSION, id, priority }; + const register = () => events.emit(REGISTER, payload); + const disposeReady = events.on(HOST_READY, (data) => { + if ((data as { protocolVersion?: unknown } | undefined)?.protocolVersion === PROTOCOL_VERSION) register(); + }); + register(); + + return { + register, + dispose() { + events.emit(UNREGISTER, { protocolVersion: PROTOCOL_VERSION, id }); + disposeReady(); + }, + }; +} diff --git a/packages/pi-cache-hit-predictor/tests/extension.test.ts b/packages/pi-cache-hit-predictor/tests/extension.test.ts index 0be5a31..b27df29 100644 --- a/packages/pi-cache-hit-predictor/tests/extension.test.ts +++ b/packages/pi-cache-hit-predictor/tests/extension.test.ts @@ -7,63 +7,58 @@ import type { } from "@earendil-works/pi-coding-agent"; import cacheHitPredictor from "../index.js"; -test("coalesces a model clamp into one inline notification", async () => { - const handlers = new Map unknown>(); - const notifications: string[] = []; - const pi = { - on(event: string, handler: (event: never, ctx: ExtensionContext) => unknown) { - handlers.set(event, handler); - }, - getThinkingLevel: () => "high", - } as unknown as ExtensionAPI; +const oldModel = { + provider: "openai", + api: "openai-responses", + id: "gpt-old", + name: "Old", + baseUrl: "https://api.openai.com/v1", + reasoning: true, + input: ["text"], + cost: { input: 1, output: 1, cacheRead: 0.1, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 10_000, +} as const; +const newModel = { ...oldModel, id: "gpt-new", name: "New" }; - cacheHitPredictor(pi); - - const branch = [ - { - type: "thinking_level_change", - id: "00000001", - parentId: null, - timestamp: new Date(1_000).toISOString(), - thinkingLevel: "low", - }, - { - type: "message", - id: "00000002", - parentId: "00000001", - timestamp: new Date(2_000).toISOString(), - message: { - role: "assistant", - content: [], - api: "openai-responses", - provider: "openai", - model: "gpt-old", - usage: { - input: 17_000, - output: 10, - cacheRead: 8_000, - cacheWrite: 0, - totalTokens: 25_010, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - }, - stopReason: "stop", - timestamp: 2_000, +const branch = [ + { + type: "thinking_level_change", + id: "00000001", + parentId: null, + timestamp: new Date(1_000).toISOString(), + thinkingLevel: "low", + }, + { + type: "message", + id: "00000002", + parentId: "00000001", + timestamp: new Date(2_000).toISOString(), + message: { + role: "assistant", + content: [], + api: "openai-responses", + provider: "openai", + model: "gpt-old", + usage: { + input: 17_000, + output: 10, + cacheRead: 8_000, + cacheWrite: 0, + totalTokens: 25_010, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, + stopReason: "stop", + timestamp: 2_000, }, - ] as SessionEntry[]; - const oldModel = { - provider: "openai", - api: "openai-responses", - id: "gpt-old", - name: "Old", - baseUrl: "https://api.openai.com/v1", - reasoning: true, - input: ["text"], - cost: { input: 1, output: 1, cacheRead: 0.1, cacheWrite: 0 }, - contextWindow: 200_000, - maxTokens: 10_000, - } as const; - const newModel = { ...oldModel, id: "gpt-new", name: "New" }; + }, +] as SessionEntry[]; + +function createHarness() { + const handlers = new Map unknown>(); + const busHandlers = new Map void>>(); + const busEvents: Array<{ channel: string; data: unknown }> = []; + const statuses: Array = []; const ctx = { mode: "tui", model: newModel, @@ -71,24 +66,95 @@ test("coalesces a model clamp into one inline notification", async () => { sessionManager: { getBranch: () => branch }, modelRegistry: { find: () => undefined }, ui: { - notify(message: string) { - notifications.push(message); + setStatus(key: string, value: string | undefined) { + assert.equal(key, "pi-cache-hit-predictor"); + statuses.push(value); }, }, } as unknown as ExtensionContext; - - await handlers.get("session_start")?.({} as never, ctx); - await handlers.get("thinking_level_select")?.( - { level: "high", previousLevel: "low" } as never, - ctx, - ); - await handlers.get("model_select")?.( - { model: newModel, previousModel: oldModel, source: "set" } as never, + const pi = { + events: { + emit(channel: string, data: unknown) { + busEvents.push({ channel, data }); + for (const handler of busHandlers.get(channel) ?? []) handler(data); + }, + on(channel: string, handler: (data: unknown) => void) { + const channelHandlers = busHandlers.get(channel) ?? new Set(); + channelHandlers.add(handler); + busHandlers.set(channel, channelHandlers); + return () => channelHandlers.delete(handler); + }, + }, + on(event: string, handler: (event: never, ctx: ExtensionContext) => unknown) { + handlers.set(event, handler); + }, + getThinkingLevel: () => "high", + } as unknown as ExtensionAPI; + cacheHitPredictor(pi); + return { + pi, ctx, - ); - assert.equal(notifications.length, 0); - await new Promise((resolve) => setTimeout(resolve, 5)); - assert.deepEqual(notifications, [ + statuses, + busEvents, + async fire(event: string, data: unknown = {}) { + await handlers.get(event)?.(data as never, ctx); + }, + }; +} + +const waitForPrediction = () => new Promise((resolve) => setTimeout(resolve, 5)); + +test("coalesces a model clamp into one footer status", async () => { + const harness = createHarness(); + await harness.fire("session_start"); + await harness.fire("thinking_level_select", { level: "high", previousLevel: "low" }); + await harness.fire("model_select", { + model: newModel, + previousModel: oldModel, + source: "set", + }); + await waitForPrediction(); + assert.equal(harness.statuses.filter(Boolean).length, 1); + assert.equal( + harness.statuses.at(-1), "Cache hit prediction · gpt-new · high: cold lane (0% of ~100k)", - ]); + ); +}); + +test("clears only after a successful response on the predicted lane", async () => { + const harness = createHarness(); + await harness.fire("model_select", { + model: newModel, + previousModel: oldModel, + source: "set", + }); + await waitForPrediction(); + const prediction = harness.statuses.at(-1); + + await harness.fire("after_provider_response", { status: 500, headers: {} }); + assert.equal(harness.statuses.at(-1), prediction); + + harness.ctx.model = oldModel as unknown as ExtensionContext["model"]; + await harness.fire("after_provider_response", { status: 200, headers: {} }); + assert.equal(harness.statuses.at(-1), prediction); + + harness.ctx.model = newModel as unknown as ExtensionContext["model"]; + await harness.fire("after_provider_response", { status: 204, headers: {} }); + assert.equal(harness.statuses.at(-1), undefined); +}); + +test("registers priority metadata and cleans up at shutdown", async () => { + const harness = createHarness(); + const registrations = () => harness.busEvents.filter(({ channel }) => channel.endsWith("/register/v1")); + assert.deepEqual(registrations().at(-1)?.data, { + protocolVersion: 1, + id: "pi-cache-hit-predictor", + priority: 200, + }); + harness.pi.events.emit("@iurysza/pi-ext/footer-slot/ready/v1", { protocolVersion: 1 }); + assert.equal(registrations().length, 2); + + await harness.fire("session_shutdown"); + assert.equal(harness.statuses.at(-1), undefined); + assert.ok(harness.busEvents.some(({ channel }) => channel.endsWith("/unregister/v1"))); }); -- 2.51.2