From ce28073970f1a25191fb9f97c6bbce4e5bd0012e Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 9 Apr 2026 04:58:36 +0100 Subject: [PATCH] test: move context-engine cache coverage to helpers --- .../run/attempt.context-engine-helpers.ts | 58 ++++++ ...mpt.spawn-workspace.context-engine.test.ts | 165 ++++-------------- src/agents/pi-embedded-runner/run/attempt.ts | 56 +----- 3 files changed, 97 insertions(+), 182 deletions(-) diff --git a/src/agents/pi-embedded-runner/run/attempt.context-engine-helpers.ts b/src/agents/pi-embedded-runner/run/attempt.context-engine-helpers.ts index d59a569d3b4..e0e2ca048dc 100644 --- a/src/agents/pi-embedded-runner/run/attempt.context-engine-helpers.ts +++ b/src/agents/pi-embedded-runner/run/attempt.context-engine-helpers.ts @@ -1,6 +1,9 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { MemoryCitationsMode } from "../../../config/types.memory.js"; import type { ContextEngine, ContextEngineRuntimeContext } from "../../../context-engine/types.js"; +import type { NormalizedUsage } from "../../usage.js"; +import type { PromptCacheChange } from "../prompt-cache-observability.js"; +import type { EmbeddedRunAttemptResult } from "./types.js"; export type AttemptContextEngine = ContextEngine; @@ -44,6 +47,61 @@ export async function resolveAttemptBootstrapContext< }; } +export function buildContextEnginePromptCacheInfo(params: { + retention?: "none" | "short" | "long"; + lastCallUsage?: NormalizedUsage; + observation?: + | { + broke: boolean; + previousCacheRead?: number; + cacheRead?: number; + changes?: PromptCacheChange[] | null; + } + | undefined; + lastCacheTouchAt?: number | null; +}): EmbeddedRunAttemptResult["promptCache"] { + const promptCache: NonNullable = {}; + if (params.retention) { + promptCache.retention = params.retention; + } + if (params.lastCallUsage) { + promptCache.lastCallUsage = { ...params.lastCallUsage }; + } + if (params.observation) { + promptCache.observation = { + broke: params.observation.broke, + ...(typeof params.observation.previousCacheRead === "number" + ? { previousCacheRead: params.observation.previousCacheRead } + : {}), + ...(typeof params.observation.cacheRead === "number" + ? { cacheRead: params.observation.cacheRead } + : {}), + ...(params.observation.changes && params.observation.changes.length > 0 + ? { + changes: params.observation.changes.map((change) => ({ + code: change.code, + detail: change.detail, + })), + } + : {}), + }; + } + if (typeof params.lastCacheTouchAt === "number" && Number.isFinite(params.lastCacheTouchAt)) { + promptCache.lastCacheTouchAt = params.lastCacheTouchAt; + } + return Object.keys(promptCache).length > 0 ? promptCache : undefined; +} + +export function findCurrentAttemptAssistantMessage(params: { + messagesSnapshot: AgentMessage[]; + prePromptMessageCount: number; +}): AgentMessage | undefined { + return params.messagesSnapshot + .slice(Math.max(0, params.prePromptMessageCount)) + .toReversed() + .find((message) => message.role === "assistant"); +} + export async function runAttemptContextEngineBootstrap(params: { hadSessionFile: boolean; contextEngine?: AttemptContextEngine; diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-engine.test.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-engine.test.ts index 2a7b7053d07..2f45a46c1b7 100644 --- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-engine.test.ts +++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-engine.test.ts @@ -8,17 +8,15 @@ import { import { type AttemptContextEngine, assembleAttemptContextEngine, + buildContextEnginePromptCacheInfo, + findCurrentAttemptAssistantMessage, finalizeAttemptContextEngineTurn, runAttemptContextEngineBootstrap, } from "./attempt.context-engine-helpers.js"; import { - cacheTtlEligibleModel, - cleanupTempPaths, - createContextEngineAttemptRunner, createContextEngineBootstrapAndAssemble, expectCalledWithSessionKey, getHoisted, - type MutableSession, resetEmbeddedAttemptHarness, } from "./attempt.spawn-workspace.test-support.js"; import { @@ -32,27 +30,6 @@ const sessionFile = "/tmp/session.jsonl"; const seedMessage = { role: "user", content: "seed", timestamp: 1 } as AgentMessage; const doneMessage = { role: "assistant", content: "done", timestamp: 2 } as unknown as AgentMessage; type AfterTurnPromptCacheCall = { runtimeContext?: { promptCache?: Record } }; -type AfterTurnUnknownPromptCacheCall = { runtimeContext?: { promptCache?: unknown } }; - -function appendAssistantWithUsage(usage: { - input?: number; - output?: number; - cacheRead?: number; - cacheWrite?: number; - total?: number; -}) { - return async (session: MutableSession, _prompt: string, _options?: { images?: unknown[] }) => { - session.messages = [ - ...session.messages, - { - role: "assistant", - content: "done", - timestamp: 2, - usage, - } as unknown as AgentMessage, - ]; - }; -} function createTestContextEngine(params: Partial): AttemptContextEngine { return { @@ -132,8 +109,6 @@ async function finalizeTurn( describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { const sessionKey = "agent:main:discord:channel:test-ctx-engine"; - const tempPaths: string[] = []; - beforeEach(() => { resetEmbeddedAttemptHarness(); clearMemoryPluginState(); @@ -143,7 +118,6 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { afterEach(async () => { clearMemoryPluginState(); vi.restoreAllMocks(); - await cleanupTempPaths(tempPaths); }); it("forwards sessionKey to bootstrap, assemble, and afterTurn", async () => { @@ -332,43 +306,20 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { ); }); - it("passes prompt-cache retention, last-call usage, and cache-touch metadata to afterTurn", async () => { - const afterTurn = vi.fn(async (_params: AfterTurnPromptCacheCall) => {}); - - await createContextEngineAttemptRunner({ - contextEngine: { - assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }), - afterTurn, - }, - attemptOverrides: { - config: { - agents: { - defaults: { - contextPruning: { - mode: "cache-ttl", - }, - }, - }, + it("builds prompt-cache retention, last-call usage, and cache-touch metadata", () => { + expect( + buildContextEnginePromptCacheInfo({ + retention: "short", + lastCallUsage: { + input: 10, + output: 5, + cacheRead: 40, + cacheWrite: 2, + total: 57, }, - provider: "anthropic", - modelId: "claude-sonnet-4-5", - model: cacheTtlEligibleModel, - }, - sessionPrompt: appendAssistantWithUsage({ - input: 10, - output: 5, - cacheRead: 40, - cacheWrite: 2, - total: 57, + lastCacheTouchAt: 123, }), - sessionKey, - tempPaths, - }); - - const afterTurnCall = afterTurn.mock.calls.at(0)?.[0]; - const runtimeContext = afterTurnCall?.runtimeContext; - - expect(runtimeContext?.promptCache).toEqual( + ).toEqual( expect.objectContaining({ retention: "short", lastCallUsage: { @@ -378,80 +329,38 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { cacheWrite: 2, total: 57, }, - lastCacheTouchAt: expect.any(Number), + lastCacheTouchAt: 123, }), ); }); - it("omits prompt-cache metadata from afterTurn when no cache data is available", async () => { - const afterTurn = vi.fn(async (_params: AfterTurnUnknownPromptCacheCall) => {}); - - await createContextEngineAttemptRunner({ - contextEngine: { - assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }), - afterTurn, - }, - sessionKey, - tempPaths, - }); - - const afterTurnCall = afterTurn.mock.calls.at(0)?.[0]; - const runtimeContext = afterTurnCall?.runtimeContext; - - expect(runtimeContext?.promptCache).toBeUndefined(); + it("omits prompt-cache metadata when no cache data is available", () => { + expect(buildContextEnginePromptCacheInfo({})).toBeUndefined(); }); - it("does not reuse a prior turn's usage when the current attempt exits before a new assistant", async () => { - const afterTurn = vi.fn(async (_params: AfterTurnPromptCacheCall) => {}); - - await createContextEngineAttemptRunner({ - contextEngine: { - assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }), - afterTurn, + it("does not reuse a prior turn's usage when the current attempt has no assistant", () => { + const priorAssistant = { + role: "assistant", + content: "prior turn", + timestamp: 2, + usage: { + input: 99, + output: 7, + cacheRead: 1234, + total: 1340, }, - attemptOverrides: { - config: { - agents: { - defaults: { - contextPruning: { - mode: "cache-ttl", - }, - }, - }, - }, - provider: "anthropic", - modelId: "claude-sonnet-4-5", - model: cacheTtlEligibleModel, - contextTokenBudget: 1, - prompt: "force-preflight-overflow", - }, - sessionMessages: [ - seedMessage, - { - role: "assistant", - content: "prior turn", - timestamp: 2, - usage: { - input: 99, - output: 7, - cacheRead: 1234, - total: 1340, - }, - } as unknown as AgentMessage, - ], - sessionKey, - tempPaths, + } as unknown as AgentMessage; + const currentAttemptAssistant = findCurrentAttemptAssistantMessage({ + messagesSnapshot: [seedMessage, priorAssistant], + prePromptMessageCount: 2, + }); + const promptCache = buildContextEnginePromptCacheInfo({ + retention: "short", + lastCallUsage: (currentAttemptAssistant as { usage?: undefined } | undefined)?.usage, }); - const afterTurnCall = afterTurn.mock.calls.at(0)?.[0]; - const promptCache = afterTurnCall?.runtimeContext?.promptCache; - - expect(promptCache).toEqual( - expect.objectContaining({ - retention: "short", - }), - ); - expect(promptCache?.lastCallUsage).toBeUndefined(); + expect(currentAttemptAssistant).toBeUndefined(); + expect(promptCache).toEqual({ retention: "short" }); }); it("threads prompt-cache break observations into afterTurn", async () => { diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index fd85d0850fa..47927f0a036 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -160,6 +160,8 @@ import { mapThinkingLevel } from "../utils.js"; import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.js"; import { assembleAttemptContextEngine, + buildContextEnginePromptCacheInfo, + findCurrentAttemptAssistantMessage, finalizeAttemptContextEngineTurn, resolveAttemptBootstrapContext, runAttemptContextEngineBootstrap, @@ -247,60 +249,6 @@ export { wrapOllamaCompatNumCtx, } from "../../../plugin-sdk/ollama-runtime.js"; -function buildContextEnginePromptCacheInfo(params: { - retention?: "none" | "short" | "long"; - lastCallUsage?: NormalizedUsage; - observation?: - | { - broke: boolean; - previousCacheRead?: number; - cacheRead?: number; - changes?: PromptCacheChange[] | null; - } - | undefined; - lastCacheTouchAt?: number | null; -}): EmbeddedRunAttemptResult["promptCache"] { - const promptCache: NonNullable = {}; - if (params.retention) { - promptCache.retention = params.retention; - } - if (params.lastCallUsage) { - promptCache.lastCallUsage = { ...params.lastCallUsage }; - } - if (params.observation) { - promptCache.observation = { - broke: params.observation.broke, - ...(typeof params.observation.previousCacheRead === "number" - ? { previousCacheRead: params.observation.previousCacheRead } - : {}), - ...(typeof params.observation.cacheRead === "number" - ? { cacheRead: params.observation.cacheRead } - : {}), - ...(params.observation.changes && params.observation.changes.length > 0 - ? { - changes: params.observation.changes.map((change) => ({ - code: change.code, - detail: change.detail, - })), - } - : {}), - }; - } - if (typeof params.lastCacheTouchAt === "number" && Number.isFinite(params.lastCacheTouchAt)) { - promptCache.lastCacheTouchAt = params.lastCacheTouchAt; - } - return Object.keys(promptCache).length > 0 ? promptCache : undefined; -} - -function findCurrentAttemptAssistantMessage(params: { - messagesSnapshot: AgentMessage[]; - prePromptMessageCount: number; -}): AgentMessage | undefined { - return params.messagesSnapshot - .slice(Math.max(0, params.prePromptMessageCount)) - .toReversed() - .find((message) => message.role === "assistant"); -} export { decodeHtmlEntitiesInObject, wrapStreamFnRepairMalformedToolCallArguments,