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 70f06a69097..4cd179894bc 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 @@ -5,7 +5,6 @@ import { clearMemoryPluginState, registerMemoryPromptSection, } from "../../../plugins/memory-state.js"; -import { derivePromptTokens } from "../../usage.js"; import { type AttemptContextEngine, buildLoopPromptCacheInfo, @@ -16,8 +15,8 @@ import { resolvePromptCacheTouchTimestamp, runAttemptContextEngineBootstrap, } from "./attempt.context-engine-helpers.js"; -import { buildAfterTurnRuntimeContext } from "./attempt.prompt-helpers.js"; import { + cleanupTempPaths, createContextEngineBootstrapAndAssemble, expectCalledWithSessionKey, getHoisted, @@ -113,6 +112,7 @@ async function finalizeTurn( describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { const sessionKey = "agent:main:discord:channel:test-ctx-engine"; + const tempPaths: string[] = []; beforeEach(() => { resetEmbeddedAttemptHarness(); clearMemoryPluginState(); @@ -120,6 +120,7 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { }); afterEach(async () => { + await cleanupTempPaths(tempPaths); clearMemoryPluginState(); vi.restoreAllMocks(); }); @@ -481,74 +482,6 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { ); }); - it("derives deferred maintenance currentTokenCount from prompt-only usage", async () => { - const afterTurn = vi.fn( - async (_params: { - runtimeContext?: { - currentTokenCount?: number; - promptCache?: { lastCallUsage?: { total?: number } }; - }; - }) => {}, - ); - - const messagesSnapshot = [ - seedMessage, - { - role: "assistant", - content: "done", - timestamp: 2, - usage: { - input: 10, - output: 5, - cacheRead: 40, - cacheWrite: 2, - total: 57, - }, - } as unknown as AgentMessage, - ]; - const promptCache = buildLoopPromptCacheInfo({ - messagesSnapshot, - prePromptMessageCount: 1, - }); - - await finalizeTurn(sessionKey, createTestContextEngine({ afterTurn }), { - messagesSnapshot, - prePromptMessageCount: 1, - runtimeContext: buildAfterTurnRuntimeContext({ - attempt: { - sessionKey, - config: {} as never, - skillsSnapshot: undefined, - senderIsOwner: true, - provider: "openai", - modelId: "gpt-test", - thinkLevel: "off", - reasoningLevel: undefined, - extraSystemPrompt: undefined, - ownerNumbers: undefined, - }, - workspaceDir: "/tmp/workspace", - agentDir: "/tmp/agent", - tokenBudget: 2048, - currentTokenCount: derivePromptTokens(promptCache?.lastCallUsage), - promptCache, - }), - }); - - expect(afterTurn).toHaveBeenCalledWith( - expect.objectContaining({ - runtimeContext: expect.objectContaining({ - currentTokenCount: 52, - promptCache: expect.objectContaining({ - lastCallUsage: expect.objectContaining({ - total: 57, - }), - }), - }), - }), - ); - }); - it("skips maintenance when ingestBatch fails", async () => { const { bootstrap, assemble } = createContextEngineBootstrapAndAssemble(); const ingestBatch = vi.fn(async () => { diff --git a/src/agents/pi-embedded-runner/run/attempt.test.ts b/src/agents/pi-embedded-runner/run/attempt.test.ts index 49af9a330b0..16773b6b149 100644 --- a/src/agents/pi-embedded-runner/run/attempt.test.ts +++ b/src/agents/pi-embedded-runner/run/attempt.test.ts @@ -5,6 +5,7 @@ import { appendBootstrapPromptWarning } from "../../bootstrap-budget.js"; import { SYSTEM_PROMPT_CACHE_BOUNDARY } from "../../system-prompt-cache-boundary.js"; import { buildAgentSystemPrompt } from "../../system-prompt.js"; import { + buildContextEnginePromptCacheInfo, buildAfterTurnRuntimeContext, composeSystemPromptWithHookContext, decodeHtmlEntitiesInObject, @@ -444,9 +445,7 @@ describe("resolveUnknownToolGuardThreshold", () => { it("falls back to the default threshold when the override is non-positive", () => { expect(resolveUnknownToolGuardThreshold({ unknownToolThreshold: 0 })).toBe(10); expect(resolveUnknownToolGuardThreshold({ unknownToolThreshold: -5 })).toBe(10); - expect( - resolveUnknownToolGuardThreshold({ unknownToolThreshold: Number.NaN }), - ).toBe(10); + expect(resolveUnknownToolGuardThreshold({ unknownToolThreshold: Number.NaN })).toBe(10); }); it("floors fractional overrides", () => { @@ -1739,9 +1738,11 @@ describe("wrapStreamFnSanitizeMalformedToolCalls", () => { ); const wrapped = wrapStreamFnSanitizeMalformedToolCalls(baseFn as never, new Set(["read"])); - const stream = wrapped({ api: "google-gemini" } as never, { messages } as never, {} as never) as - | FakeWrappedStream - | Promise; + const stream = wrapped( + { api: "google-gemini" } as never, + { messages } as never, + {} as never, + ) as FakeWrappedStream | Promise; await Promise.resolve(stream); expect(baseFn).toHaveBeenCalledTimes(1); @@ -2831,6 +2832,15 @@ describe("buildAfterTurnRuntimeContext", () => { }); }); it("includes resolved auth profile fields for context-engine afterTurn compaction", () => { + const promptCache = buildContextEnginePromptCacheInfo({ + lastCallUsage: { + input: 10, + output: 5, + cacheRead: 40, + cacheWrite: 2, + total: 57, + }, + }); const legacy = buildAfterTurnRuntimeContext({ attempt: { sessionKey: "agent:main:session:abc", @@ -2851,7 +2861,8 @@ describe("buildAfterTurnRuntimeContext", () => { workspaceDir: "/tmp/workspace", agentDir: "/tmp/agent", tokenBudget: 1050000, - currentTokenCount: 232393, + currentTokenCount: 52, + promptCache, }); expect(legacy).toMatchObject({ @@ -2861,7 +2872,12 @@ describe("buildAfterTurnRuntimeContext", () => { workspaceDir: "/tmp/workspace", agentDir: "/tmp/agent", tokenBudget: 1050000, - currentTokenCount: 232393, + currentTokenCount: 52, + promptCache: { + lastCallUsage: { + total: 57, + }, + }, }); }); diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 37ac21a1b08..986a0e7b081 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -180,6 +180,7 @@ import { import { splitSdkTools } from "../tool-split.js"; import { mapThinkingLevel } from "../utils.js"; import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.js"; +export { buildContextEnginePromptCacheInfo } from "./attempt.context-engine-helpers.js"; import { assembleAttemptContextEngine, buildLoopPromptCacheInfo,