From e4c4f955b35ac59aec7490cebe679c59a26c64d3 Mon Sep 17 00:00:00 2001 From: Gustavo Madeira Santana Date: Fri, 17 Apr 2026 02:13:19 -0400 Subject: [PATCH] Tests: restore context-engine usage proof --- .../run/attempt.prompt-helpers.ts | 20 +++++++-- .../pi-embedded-runner/run/attempt.test.ts | 44 +++++++++++++++++++ src/agents/pi-embedded-runner/run/attempt.ts | 9 ++-- 3 files changed, 65 insertions(+), 8 deletions(-) diff --git a/src/agents/pi-embedded-runner/run/attempt.prompt-helpers.ts b/src/agents/pi-embedded-runner/run/attempt.prompt-helpers.ts index 6aaa470ee64..2cd45b0c975 100644 --- a/src/agents/pi-embedded-runner/run/attempt.prompt-helpers.ts +++ b/src/agents/pi-embedded-runner/run/attempt.prompt-helpers.ts @@ -14,6 +14,7 @@ import { resolveHeartbeatPromptForSystemPrompt } from "../../heartbeat-system-pr import { buildActiveMusicGenerationTaskPromptContextForSession } from "../../music-generation-task-status.js"; import { prependSystemPromptAdditionAfterCacheBoundary } from "../../system-prompt-cache-boundary.js"; import { resolveEffectiveToolFsWorkspaceOnly } from "../../tool-fs-policy.js"; +import { derivePromptTokens, type NormalizedUsage } from "../../usage.js"; import { buildActiveVideoGenerationTaskPromptContextForSession } from "../../video-generation-task-status.js"; import { buildEmbeddedCompactionRuntimeContext } from "../compaction-runtime-context.js"; import { log } from "../logger.js"; @@ -255,15 +256,26 @@ export function buildAfterTurnRuntimeContext(params: { ownerNumbers: params.attempt.ownerNumbers, }), ...(typeof params.tokenBudget === "number" && - Number.isFinite(params.tokenBudget) && - params.tokenBudget > 0 + Number.isFinite(params.tokenBudget) && + params.tokenBudget > 0 ? { tokenBudget: Math.floor(params.tokenBudget) } : {}), ...(typeof params.currentTokenCount === "number" && - Number.isFinite(params.currentTokenCount) && - params.currentTokenCount > 0 + Number.isFinite(params.currentTokenCount) && + params.currentTokenCount > 0 ? { currentTokenCount: Math.floor(params.currentTokenCount) } : {}), ...(params.promptCache ? { promptCache: params.promptCache } : {}), }; } + +export function buildAfterTurnRuntimeContextFromUsage( + params: Omit[0], "currentTokenCount"> & { + lastCallUsage?: NormalizedUsage; + }, +): ContextEngineRuntimeContext { + return buildAfterTurnRuntimeContext({ + ...params, + currentTokenCount: derivePromptTokens(params.lastCallUsage), + }); +} diff --git a/src/agents/pi-embedded-runner/run/attempt.test.ts b/src/agents/pi-embedded-runner/run/attempt.test.ts index 16773b6b149..b7c1a75e2de 100644 --- a/src/agents/pi-embedded-runner/run/attempt.test.ts +++ b/src/agents/pi-embedded-runner/run/attempt.test.ts @@ -7,6 +7,7 @@ import { buildAgentSystemPrompt } from "../../system-prompt.js"; import { buildContextEnginePromptCacheInfo, buildAfterTurnRuntimeContext, + buildAfterTurnRuntimeContextFromUsage, composeSystemPromptWithHookContext, decodeHtmlEntitiesInObject, mergeOrphanedTrailingUserPrompt, @@ -2881,6 +2882,49 @@ describe("buildAfterTurnRuntimeContext", () => { }); }); + it("derives afterTurn token count from the current assistant usage snapshot", () => { + const lastCallUsage = { + input: 10, + output: 5, + cacheRead: 40, + cacheWrite: 2, + total: 57, + }; + const promptCache = buildContextEnginePromptCacheInfo({ lastCallUsage }); + const legacy = buildAfterTurnRuntimeContextFromUsage({ + attempt: { + sessionKey: "agent:main:session:abc", + messageChannel: "slack", + messageProvider: "slack", + agentAccountId: "acct-1", + authProfileId: "openai:p1", + config: { plugins: { slots: { contextEngine: "lossless-claw" } } } as OpenClawConfig, + skillsSnapshot: undefined, + senderIsOwner: true, + provider: "openai-codex", + modelId: "gpt-5.4", + thinkLevel: "off", + reasoningLevel: "on", + extraSystemPrompt: "extra", + ownerNumbers: ["+15555550123"], + }, + workspaceDir: "/tmp/workspace", + agentDir: "/tmp/agent", + tokenBudget: 1050000, + lastCallUsage, + promptCache, + }); + + expect(legacy).toMatchObject({ + currentTokenCount: 52, + promptCache: { + lastCallUsage: { + total: 57, + }, + }, + }); + }); + it("preserves sender and channel routing context for scoped compaction discovery", () => { const legacy = buildAfterTurnRuntimeContext({ attempt: { diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 986a0e7b081..6fac45ebd4a 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -120,7 +120,7 @@ import { resolveTranscriptPolicy, shouldAllowProviderOwnedThinkingReplay, } from "../../transcript-policy.js"; -import { derivePromptTokens, normalizeUsage, type NormalizedUsage } from "../../usage.js"; +import { normalizeUsage, type NormalizedUsage } from "../../usage.js"; import { DEFAULT_BOOTSTRAP_FILENAME } from "../../workspace.js"; import { isRunnerAbortError } from "../abort.js"; import { isCacheTtlEligibleProvider, readLastCacheTtlTimestamp } from "../cache-ttl.js"; @@ -193,6 +193,7 @@ import { } from "./attempt.context-engine-helpers.js"; import { buildAfterTurnRuntimeContext, + buildAfterTurnRuntimeContextFromUsage, mergeOrphanedTrailingUserPrompt, prependSystemPromptAddition, resolveAttemptFsWorkspaceOnly, @@ -256,6 +257,7 @@ export { } from "./attempt.thread-helpers.js"; export { buildAfterTurnRuntimeContext, + buildAfterTurnRuntimeContextFromUsage, mergeOrphanedTrailingUserPrompt, prependSystemPromptAddition, resolveAttemptFsWorkspaceOnly, @@ -2296,13 +2298,12 @@ export async function runEmbeddedAttempt( // Let the active context engine run its post-turn lifecycle. if (params.contextEngine) { - const runtimeCurrentTokenCount = derivePromptTokens(lastCallUsage); - const afterTurnRuntimeContext = buildAfterTurnRuntimeContext({ + const afterTurnRuntimeContext = buildAfterTurnRuntimeContextFromUsage({ attempt: params, workspaceDir: effectiveWorkspace, agentDir, tokenBudget: params.contextTokenBudget, - currentTokenCount: runtimeCurrentTokenCount, + lastCallUsage, promptCache, }); await finalizeAttemptContextEngineTurn({