test: make context injection coverage pure

src/agents/pi-embedded-runner/run/attempt.context-engine-helpers.ts

@@ -4,6 +4,46 @@ import type { ContextEngine, ContextEngineRuntimeContext } from "../../../contex
 
 export type AttemptContextEngine = ContextEngine;
 
+export type AttemptBootstrapContext = {
+  bootstrapFiles: unknown[];
+  contextFiles: unknown[];
+};
+
+export async function resolveAttemptBootstrapContext<
+  TContext extends AttemptBootstrapContext,
+>(params: {
+  contextInjectionMode: "always" | "continuation-skip";
+  bootstrapContextMode?: string;
+  bootstrapContextRunKind?: string;
+  sessionFile: string;
+  hasCompletedBootstrapTurn: (sessionFile: string) => Promise<boolean>;
+  resolveBootstrapContextForRun: () => Promise<TContext>;
+}): Promise<
+  TContext & {
+    isContinuationTurn: boolean;
+    shouldRecordCompletedBootstrapTurn: boolean;
+  }
+> {
+  const isContinuationTurn =
+    params.contextInjectionMode === "continuation-skip" &&
+    params.bootstrapContextRunKind !== "heartbeat" &&
+    (await params.hasCompletedBootstrapTurn(params.sessionFile));
+  const shouldRecordCompletedBootstrapTurn =
+    !isContinuationTurn &&
+    params.bootstrapContextMode !== "lightweight" &&
+    params.bootstrapContextRunKind !== "heartbeat";
+
+  const context = isContinuationTurn
+    ? ({ bootstrapFiles: [], contextFiles: [] } as unknown as TContext)
+    : await params.resolveBootstrapContextForRun();
+
+  return {
+    ...context,
+    isContinuationTurn,
+    shouldRecordCompletedBootstrapTurn,
+  };
+}
+
 export async function runAttemptContextEngineBootstrap(params: {
   hadSessionFile: boolean;
   contextEngine?: AttemptContextEngine;

src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-injection.test.ts

@@ -1,189 +1,139 @@
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { describe, expect, it, vi } from "vitest";
+import { filterHeartbeatPairs } from "../../../auto-reply/heartbeat-filter.js";
 import { HEARTBEAT_PROMPT } from "../../../auto-reply/heartbeat.js";
 import { limitHistoryTurns } from "../history.js";
 import {
-  cleanupTempPaths,
-  createContextEngineAttemptRunner,
-  getHoisted,
-  resetEmbeddedAttemptHarness,
-} from "./attempt.spawn-workspace.test-support.js";
+  assembleAttemptContextEngine,
+  type AttemptContextEngine,
+  resolveAttemptBootstrapContext,
+} from "./attempt.context-engine-helpers.js";
 
-const hoisted = getHoisted();
+async function resolveBootstrapContext(params: {
+  contextInjectionMode?: "always" | "continuation-skip";
+  bootstrapContextMode?: string;
+  bootstrapContextRunKind?: string;
+  completed?: boolean;
+  resolver?: () => Promise<{ bootstrapFiles: unknown[]; contextFiles: unknown[] }>;
+}) {
+  const hasCompletedBootstrapTurn = vi.fn(async () => params.completed ?? false);
+  const resolveBootstrapContextForRun =
+    params.resolver ??
+    vi.fn(async () => ({
+      bootstrapFiles: [],
+      contextFiles: [],
+    }));
 
-describe("runEmbeddedAttempt context injection", () => {
-  const tempPaths: string[] = [];
-
-  beforeEach(() => {
-    resetEmbeddedAttemptHarness();
+  const result = await resolveAttemptBootstrapContext({
+    contextInjectionMode: params.contextInjectionMode ?? "always",
+    bootstrapContextMode: params.bootstrapContextMode ?? "full",
+    bootstrapContextRunKind: params.bootstrapContextRunKind ?? "default",
+    sessionFile: "/tmp/session.jsonl",
+    hasCompletedBootstrapTurn,
+    resolveBootstrapContextForRun,
   });
 
-  afterEach(async () => {
-    await cleanupTempPaths(tempPaths);
-  });
+  return { result, hasCompletedBootstrapTurn, resolveBootstrapContextForRun };
+}
 
+describe("embedded attempt context injection", () => {
   it("skips bootstrap reinjection on safe continuation turns when configured", async () => {
-    hoisted.resolveContextInjectionModeMock.mockReturnValue("continuation-skip");
-    hoisted.hasCompletedBootstrapTurnMock.mockResolvedValue(true);
+    const { result, hasCompletedBootstrapTurn, resolveBootstrapContextForRun } =
+      await resolveBootstrapContext({
+        contextInjectionMode: "continuation-skip",
+        completed: true,
+      });
 
-    await createContextEngineAttemptRunner({
-      contextEngine: {
-        assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }),
-      },
-      sessionKey: "agent:main",
-      tempPaths,
-    });
-
-    expect(hoisted.hasCompletedBootstrapTurnMock).toHaveBeenCalled();
-    expect(hoisted.resolveBootstrapContextForRunMock).not.toHaveBeenCalled();
-  });
-
-  it("checks continuation state only after taking the session lock", async () => {
-    hoisted.resolveContextInjectionModeMock.mockReturnValue("continuation-skip");
-    hoisted.hasCompletedBootstrapTurnMock.mockResolvedValue(true);
-
-    await createContextEngineAttemptRunner({
-      contextEngine: {
-        assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }),
-      },
-      sessionKey: "agent:main",
-      tempPaths,
-    });
-
-    expect(hoisted.acquireSessionWriteLockMock).toHaveBeenCalled();
-    expect(hoisted.hasCompletedBootstrapTurnMock).toHaveBeenCalled();
-    const lockCallOrder = hoisted.acquireSessionWriteLockMock.mock.invocationCallOrder[0];
-    const continuationCallOrder = hoisted.hasCompletedBootstrapTurnMock.mock.invocationCallOrder[0];
-    expect(lockCallOrder).toBeLessThan(continuationCallOrder);
+    expect(result.isContinuationTurn).toBe(true);
+    expect(result.bootstrapFiles).toEqual([]);
+    expect(result.contextFiles).toEqual([]);
+    expect(hasCompletedBootstrapTurn).toHaveBeenCalledWith("/tmp/session.jsonl");
+    expect(resolveBootstrapContextForRun).not.toHaveBeenCalled();
   });
 
   it("still resolves bootstrap context when continuation-skip has no completed assistant turn yet", async () => {
-    hoisted.resolveContextInjectionModeMock.mockReturnValue("continuation-skip");
-    hoisted.hasCompletedBootstrapTurnMock.mockResolvedValue(false);
+    const resolver = vi.fn(async () => ({
+      bootstrapFiles: [{ name: "AGENTS.md" }],
+      contextFiles: [{ path: "AGENTS.md" }],
+    }));
 
-    await createContextEngineAttemptRunner({
-      contextEngine: {
-        assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }),
-      },
-      sessionKey: "agent:main",
-      tempPaths,
+    const { result } = await resolveBootstrapContext({
+      contextInjectionMode: "continuation-skip",
+      completed: false,
+      resolver,
     });
 
-    expect(hoisted.resolveBootstrapContextForRunMock).toHaveBeenCalledTimes(1);
+    expect(result.isContinuationTurn).toBe(false);
+    expect(result.bootstrapFiles).toEqual([{ name: "AGENTS.md" }]);
+    expect(result.contextFiles).toEqual([{ path: "AGENTS.md" }]);
+    expect(resolver).toHaveBeenCalledTimes(1);
   });
 
   it("never skips heartbeat bootstrap filtering", async () => {
-    hoisted.resolveContextInjectionModeMock.mockReturnValue("continuation-skip");
-    hoisted.hasCompletedBootstrapTurnMock.mockResolvedValue(true);
-
-    await createContextEngineAttemptRunner({
-      contextEngine: {
-        assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }),
-      },
-      attemptOverrides: {
+    const { result, hasCompletedBootstrapTurn, resolveBootstrapContextForRun } =
+      await resolveBootstrapContext({
+        contextInjectionMode: "continuation-skip",
         bootstrapContextMode: "lightweight",
         bootstrapContextRunKind: "heartbeat",
-      },
-      sessionKey: "agent:main:heartbeat:test",
-      tempPaths,
-    });
+        completed: true,
+      });
 
-    expect(hoisted.hasCompletedBootstrapTurnMock).not.toHaveBeenCalled();
-    expect(hoisted.resolveBootstrapContextForRunMock).toHaveBeenCalledWith(
-      expect.objectContaining({
-        contextMode: "lightweight",
-        runKind: "heartbeat",
-      }),
-    );
+    expect(result.isContinuationTurn).toBe(false);
+    expect(result.shouldRecordCompletedBootstrapTurn).toBe(false);
+    expect(hasCompletedBootstrapTurn).not.toHaveBeenCalled();
+    expect(resolveBootstrapContextForRun).toHaveBeenCalledTimes(1);
   });
 
   it("runs full bootstrap injection after a successful non-heartbeat turn", async () => {
-    hoisted.resolveBootstrapContextForRunMock.mockResolvedValue({
-      bootstrapFiles: [
-        {
-          name: "AGENTS.md",
-          path: "AGENTS.md",
-          content: "bootstrap context",
-          missing: false,
-        },
-      ],
-      contextFiles: [
-        {
-          path: "AGENTS.md",
-          content: "bootstrap context",
-        },
-      ],
+    const resolver = vi.fn(async () => ({
+      bootstrapFiles: [{ name: "AGENTS.md", content: "bootstrap context" }],
+      contextFiles: [{ path: "AGENTS.md", content: "bootstrap context" }],
+    }));
+
+    const { result } = await resolveBootstrapContext({
+      bootstrapContextMode: "full",
+      bootstrapContextRunKind: "default",
+      resolver,
     });
 
-    const result = await createContextEngineAttemptRunner({
-      contextEngine: {
-        assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }),
-      },
-      attemptOverrides: {
-        bootstrapContextMode: "full",
-        bootstrapContextRunKind: "default",
-      },
-      sessionKey: "agent:main",
-      tempPaths,
-    });
-
-    expect(result.promptError).toBeNull();
-    expect(hoisted.resolveBootstrapContextForRunMock).toHaveBeenCalledWith(
-      expect.objectContaining({
-        contextMode: "full",
-        runKind: "default",
-      }),
-    );
+    expect(result.shouldRecordCompletedBootstrapTurn).toBe(true);
+    expect(result.bootstrapFiles).toEqual([{ name: "AGENTS.md", content: "bootstrap context" }]);
   });
 
   it("does not record full bootstrap completion for heartbeat runs", async () => {
-    await createContextEngineAttemptRunner({
-      contextEngine: {
-        assemble: async ({ messages }) => ({ messages, estimatedTokens: 1 }),
-      },
-      attemptOverrides: {
-        bootstrapContextMode: "lightweight",
-        bootstrapContextRunKind: "heartbeat",
-      },
-      sessionKey: "agent:main:heartbeat:test",
-      tempPaths,
+    const { result } = await resolveBootstrapContext({
+      bootstrapContextMode: "lightweight",
+      bootstrapContextRunKind: "heartbeat",
     });
 
-    expect(hoisted.sessionManager.appendCustomEntry).not.toHaveBeenCalledWith(
-      "openclaw:bootstrap-context:full",
-      expect.anything(),
-    );
+    expect(result.shouldRecordCompletedBootstrapTurn).toBe(false);
   });
 
   it("filters no-op heartbeat pairs before history limiting and context-engine assembly", async () => {
-    hoisted.getDmHistoryLimitFromSessionKeyMock.mockReturnValue(1);
-    hoisted.limitHistoryTurnsMock.mockImplementation(
-      (messages: unknown, limit: number | undefined) =>
-        limitHistoryTurns(messages as AgentMessage[], limit),
-    );
     const assemble = vi.fn(async ({ messages }: { messages: AgentMessage[] }) => ({
       messages,
       estimatedTokens: 1,
     }));
     const sessionMessages: AgentMessage[] = [
-      { role: "user", content: "real question", timestamp: 1 } as unknown as AgentMessage,
+      { role: "user", content: "real question", timestamp: 1 } as AgentMessage,
       { role: "assistant", content: "real answer", timestamp: 2 } as unknown as AgentMessage,
-      { role: "user", content: HEARTBEAT_PROMPT, timestamp: 3 } as unknown as AgentMessage,
+      { role: "user", content: HEARTBEAT_PROMPT, timestamp: 3 } as AgentMessage,
       { role: "assistant", content: "HEARTBEAT_OK", timestamp: 4 } as unknown as AgentMessage,
     ];
 
-    await createContextEngineAttemptRunner({
-      contextEngine: { assemble },
-      attemptOverrides: {
-        config: {
-          agents: {
-            list: [{ id: "main", heartbeat: {} }],
-          },
-        },
-      },
+    const heartbeatFiltered = filterHeartbeatPairs(sessionMessages, undefined, HEARTBEAT_PROMPT);
+    const limited = limitHistoryTurns(heartbeatFiltered, 1);
+    await assembleAttemptContextEngine({
+      contextEngine: {
+        info: { id: "test", name: "Test", version: "0.0.1" },
+        ingest: async () => ({ ingested: true }),
+        compact: async () => ({ ok: false, compacted: false, reason: "unused" }),
+        assemble,
+      } satisfies AttemptContextEngine,
+      sessionId: "session",
       sessionKey: "agent:main:discord:dm:test-user",
-      sessionMessages,
-      tempPaths,
+      messages: limited,
+      modelId: "gpt-test",
     });
 
     expect(assemble).toHaveBeenCalledWith(

src/agents/pi-embedded-runner/run/attempt.ts

@@ -161,6 +161,7 @@ import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.
 import {
   assembleAttemptContextEngine,
   finalizeAttemptContextEngineTurn,
+  resolveAttemptBootstrapContext,
   runAttemptContextEngineBootstrap,
 } from "./attempt.context-engine-helpers.js";
 import {
@@ -451,20 +452,18 @@ export async function runEmbeddedAttempt(
 
     const sessionLabel = params.sessionKey ?? params.sessionId;
     const contextInjectionMode = resolveContextInjectionMode(params.config);
-    const isContinuationTurn =
-      contextInjectionMode === "continuation-skip" &&
-      params.bootstrapContextRunKind !== "heartbeat" &&
-      (await hasCompletedBootstrapTurn(params.sessionFile));
-    const shouldRecordCompletedBootstrapTurn =
-      !isContinuationTurn &&
-      params.bootstrapContextMode !== "lightweight" &&
-      params.bootstrapContextRunKind !== "heartbeat";
-    const { bootstrapFiles: hookAdjustedBootstrapFiles, contextFiles } = isContinuationTurn
-      ? {
-          bootstrapFiles: [],
-          contextFiles: [],
-        }
-      : await resolveBootstrapContextForRun({
+    const {
+      bootstrapFiles: hookAdjustedBootstrapFiles,
+      contextFiles,
+      shouldRecordCompletedBootstrapTurn,
+    } = await resolveAttemptBootstrapContext({
+      contextInjectionMode,
+      bootstrapContextMode: params.bootstrapContextMode,
+      bootstrapContextRunKind: params.bootstrapContextRunKind,
+      sessionFile: params.sessionFile,
+      hasCompletedBootstrapTurn,
+      resolveBootstrapContextForRun: async () =>
+        await resolveBootstrapContextForRun({
           workspaceDir: effectiveWorkspace,
           config: params.config,
           sessionKey: params.sessionKey,
@@ -472,7 +471,8 @@ export async function runEmbeddedAttempt(
           warn: makeBootstrapWarn({ sessionLabel, warn: (message) => log.warn(message) }),
           contextMode: params.bootstrapContextMode,
           runKind: params.bootstrapContextRunKind,
-        });
+        }),
+    });
     const bootstrapMaxChars = resolveBootstrapMaxChars(params.config);
     const bootstrapTotalMaxChars = resolveBootstrapTotalMaxChars(params.config);
     const bootstrapAnalysis = analyzeBootstrapBudget({