From 714adeb7f6e7172f46423e69cacf1cee204d3dd0 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 9 Apr 2026 04:51:47 +0100 Subject: [PATCH] test: make context injection coverage pure --- .../run/attempt.context-engine-helpers.ts | 40 +++ ....spawn-workspace.context-injection.test.ts | 230 +++++++----------- src/agents/pi-embedded-runner/run/attempt.ts | 30 +-- 3 files changed, 145 insertions(+), 155 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 77129f90618..d59a569d3b4 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 @@ -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; + resolveBootstrapContextForRun: () => Promise; +}): 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; diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-injection.test.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-injection.test.ts index 0af89db050f..6916872a75a 100644 --- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.context-injection.test.ts +++ b/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( diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index ece10891c6a..fd85d0850fa 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/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({