From 9765a071e82f0dd15aadd2dd6d56de929c045d89 Mon Sep 17 00:00:00 2001 From: Eva Date: Wed, 29 Apr 2026 16:09:11 +0700 Subject: [PATCH] fix(context-engine): honor assembled prompt authority --- ...mpt.spawn-workspace.context-engine.test.ts | 67 +++++++++++++++++++ src/agents/pi-embedded-runner/run/attempt.ts | 6 +- .../run/preemptive-compaction.test.ts | 2 +- src/context-engine/types.ts | 5 ++ 4 files changed, 78 insertions(+), 2 deletions(-) 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 74c65599e66..250d06412d2 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 @@ -322,6 +322,73 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { ); }); + it("uses assembled context as the default precheck authority", async () => { + let sawPrompt = false; + const hugeHistory = "large raw history ".repeat(25_000); + + const result = await createContextEngineAttemptRunner({ + contextEngine: createTestContextEngine({ + assemble: async () => ({ + messages: [ + { role: "user", content: "small assembled context", timestamp: 1 }, + ] as AgentMessage[], + estimatedTokens: 8, + }), + }), + sessionKey, + tempPaths, + sessionMessages: [{ role: "user", content: hugeHistory, timestamp: 1 }] as AgentMessage[], + attemptOverrides: { + contextTokenBudget: 500, + }, + sessionPrompt: async (session) => { + sawPrompt = true; + session.messages = [ + ...session.messages, + { role: "assistant", content: "done", timestamp: 2 }, + ]; + }, + }); + + expect(sawPrompt).toBe(true); + expect(result.promptError).toBeNull(); + expect(result.promptErrorSource).toBeNull(); + }); + + it("honors context engines that opt into preassembly overflow authority", async () => { + let sawPrompt = false; + const hugeHistory = "large raw history ".repeat(25_000); + + const result = await createContextEngineAttemptRunner({ + contextEngine: createTestContextEngine({ + assemble: async () => ({ + messages: [ + { role: "user", content: "small assembled context", timestamp: 1 }, + ] as AgentMessage[], + estimatedTokens: 8, + promptAuthority: "preassembly_may_overflow", + }), + }), + sessionKey, + tempPaths, + sessionMessages: [{ role: "user", content: hugeHistory, timestamp: 1 }] as AgentMessage[], + attemptOverrides: { + contextTokenBudget: 500, + }, + sessionPrompt: async (session) => { + sawPrompt = true; + session.messages = [ + ...session.messages, + { role: "assistant", content: "done", timestamp: 2 }, + ]; + }, + }); + + expect(sawPrompt).toBe(false); + expect(result.promptErrorSource).toBe("precheck"); + expect(result.preflightRecovery?.route).toBe("compact_only"); + }); + it("keeps gateway model runs independent from agent context and session history", async () => { const bootstrap = vi.fn(async () => ({ bootstrapped: true })); const assemble = vi.fn(async ({ messages }: { messages: AgentMessage[] }) => ({ diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index a473e3f45e3..b655034ed55 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -1527,6 +1527,7 @@ export async function runEmbeddedAttempt( } let prePromptMessageCount = activeSession.messages.length; let unwindowedContextEngineMessagesForPrecheck: AgentMessage[] | undefined; + let contextEnginePromptAuthority: "assembled" | "preassembly_may_overflow" = "assembled"; abortSessionForYield = () => { yieldAbortSettled = Promise.resolve(activeSession.abort()); }; @@ -2089,6 +2090,7 @@ export async function runEmbeddedAttempt( if (assembled.messages !== activeSession.messages) { activeSession.agent.state.messages = assembled.messages; } + contextEnginePromptAuthority = assembled.promptAuthority ?? "assembled"; if (assembled.systemPromptAddition) { systemPromptText = prependSystemPromptAddition({ systemPrompt: systemPromptText, @@ -2760,7 +2762,9 @@ export async function runEmbeddedAttempt( const preemptiveCompaction = shouldPreemptivelyCompactBeforePrompt({ messages: activeSession.messages, - unwindowedMessages: unwindowedContextEngineMessagesForPrecheck, + ...(contextEnginePromptAuthority === "preassembly_may_overflow" + ? { unwindowedMessages: unwindowedContextEngineMessagesForPrecheck } + : {}), systemPrompt: systemPromptText, prompt: effectivePrompt, contextTokenBudget, diff --git a/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts b/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts index 2cc4cb218da..67bdb2cf383 100644 --- a/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts +++ b/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts @@ -93,7 +93,7 @@ describe("preemptive-compaction", () => { expect(result.estimatedPromptTokens).toBeLessThan(result.promptBudgetBeforeReserve); }); - it("uses the larger unwindowed message estimate when context engine assembly windows history", () => { + it("uses the larger unwindowed message estimate when explicitly provided", () => { const result = shouldPreemptivelyCompactBeforePrompt({ messages: [makeAssistantHistory("small assembled window")], unwindowedMessages: [makeAssistantHistory(verboseHistory.repeat(4))], diff --git a/src/context-engine/types.ts b/src/context-engine/types.ts index 8a7d6e8b6f0..29f132066f4 100644 --- a/src/context-engine/types.ts +++ b/src/context-engine/types.ts @@ -8,6 +8,11 @@ export type AssembleResult = { messages: AgentMessage[]; /** Estimated total tokens in assembled context */ estimatedTokens: number; + /** + * Declares whether the assembled messages are the authoritative prompt for + * overflow prechecks. Defaults to "assembled". + */ + promptAuthority?: "assembled" | "preassembly_may_overflow"; /** Optional context-engine-provided instructions prepended to the runtime system prompt */ systemPromptAddition?: string; };