test: move context-engine cache coverage to helpers

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<EmbeddedRunAttemptResult["promptCache"]> = {};
+  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;

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<string, unknown> } };
-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>): 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 () => {

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<EmbeddedRunAttemptResult["promptCache"]> = {};
-  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,