Skip to content

Commit 2962db3

Browse files
fuller-stack-devfuller-stack-dev
committed
fix context pressure preflight for tool payloads
1 parent c127334 commit 2962db3

4 files changed

Lines changed: 383 additions & 37 deletions

File tree

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

Lines changed: 30 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -2497,26 +2497,21 @@ export async function runEmbeddedAttempt(
24972497
const onMidTurnPrecheck = (request: MidTurnPrecheckRequest) => {
24982498
pendingMidTurnPrecheckRequest = request;
24992499
};
2500-
if (!activeContextEngine || activeContextEngine.info.ownsCompaction !== true) {
2501-
removeToolResultContextGuard = installToolResultContextGuard({
2502-
agent: activeSession.agent,
2503-
contextWindowTokens: contextTokenBudgetForGuard,
2504-
...(midTurnPrecheckEnabled
2505-
? {
2506-
midTurnPrecheck: {
2507-
enabled: true,
2508-
contextTokenBudget: contextTokenBudgetForGuard,
2509-
reserveTokens: () => settingsManager.getCompactionReserveTokens(),
2510-
toolResultMaxChars: toolResultMaxCharsForGuard,
2511-
getSystemPrompt: () => systemPromptText,
2512-
getPrePromptMessageCount: () => prePromptMessageCount,
2513-
onMidTurnPrecheck,
2514-
},
2515-
}
2516-
: {}),
2517-
});
2518-
} else {
2519-
removeToolResultContextGuard = installContextEngineLoopHook({
2500+
const midTurnPrecheckOptions = midTurnPrecheckEnabled
2501+
? {
2502+
midTurnPrecheck: {
2503+
enabled: true,
2504+
contextTokenBudget: contextTokenBudgetForGuard,
2505+
reserveTokens: () => settingsManager.getCompactionReserveTokens(),
2506+
toolResultMaxChars: toolResultMaxCharsForGuard,
2507+
getSystemPrompt: () => systemPromptText,
2508+
getPrePromptMessageCount: () => prePromptMessageCount,
2509+
onMidTurnPrecheck,
2510+
},
2511+
}
2512+
: {};
2513+
if (activeContextEngine?.info.ownsCompaction === true) {
2514+
const removeContextEngineLoopHook = installContextEngineLoopHook({
25202515
agent: activeSession.agent,
25212516
contextEngine: activeContextEngine,
25222517
sessionId: params.sessionId,
@@ -2547,6 +2542,21 @@ export async function runEmbeddedAttempt(
25472542
}),
25482543
}),
25492544
});
2545+
const removeGuard = installToolResultContextGuard({
2546+
agent: activeSession.agent,
2547+
contextWindowTokens: contextTokenBudgetForGuard,
2548+
...midTurnPrecheckOptions,
2549+
});
2550+
removeToolResultContextGuard = () => {
2551+
removeGuard();
2552+
removeContextEngineLoopHook();
2553+
};
2554+
} else {
2555+
removeToolResultContextGuard = installToolResultContextGuard({
2556+
agent: activeSession.agent,
2557+
contextWindowTokens: contextTokenBudgetForGuard,
2558+
...midTurnPrecheckOptions,
2559+
});
25502560
}
25512561
const removeLoopContextGuard = removeToolResultContextGuard;
25522562
const removeHistoryImagePruneContextTransform = installHistoryImagePruneContextTransform(

‎src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import "../../test-helpers/pi-coding-agent-token-mock.js";
44
import { estimateToolResultReductionPotential } from "../tool-result-truncation.js";
55

66
let PREEMPTIVE_OVERFLOW_ERROR_TEXT: typeof import("./preemptive-compaction.js").PREEMPTIVE_OVERFLOW_ERROR_TEXT;
7+
let estimateLlmBoundaryTokenPressure: typeof import("./preemptive-compaction.js").estimateLlmBoundaryTokenPressure;
78
let estimatePrePromptTokens: typeof import("./preemptive-compaction.js").estimatePrePromptTokens;
89
let formatPrePromptPrecheckLog: typeof import("./preemptive-compaction.js").formatPrePromptPrecheckLog;
910
let shouldPreemptivelyCompactBeforePrompt: typeof import("./preemptive-compaction.js").shouldPreemptivelyCompactBeforePrompt;
@@ -12,6 +13,7 @@ beforeAll(async () => {
1213
vi.resetModules();
1314
({
1415
PREEMPTIVE_OVERFLOW_ERROR_TEXT,
16+
estimateLlmBoundaryTokenPressure,
1517
estimatePrePromptTokens,
1618
formatPrePromptPrecheckLog,
1719
shouldPreemptivelyCompactBeforePrompt,
@@ -39,6 +41,32 @@ function makeToolResultMessage(...texts: string[]): AgentMessage {
3941
} as AgentMessage;
4042
}
4143

44+
function makeJsonToolResultMessage(payload: unknown): AgentMessage {
45+
return {
46+
role: "toolResult",
47+
toolCallId: `call_${timestamp}`,
48+
toolName: "json_tool",
49+
content: [{ type: "json", payload }],
50+
isError: false,
51+
timestamp: timestamp++,
52+
} as unknown as AgentMessage;
53+
}
54+
55+
function makeAssistantToolCall(args: unknown): AgentMessage {
56+
return {
57+
role: "assistant",
58+
content: [
59+
{
60+
type: "toolCall",
61+
id: `call_${timestamp}`,
62+
name: "bulk_lookup",
63+
arguments: args,
64+
},
65+
],
66+
timestamp: timestamp++,
67+
} as AgentMessage;
68+
}
69+
4270
describe("preemptive-compaction", () => {
4371
const verboseHistory =
4472
"alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu ".repeat(40);
@@ -146,6 +174,85 @@ describe("preemptive-compaction", () => {
146174
expect(result.estimatedPromptTokens).toBeGreaterThan(result.promptBudgetBeforeReserve);
147175
});
148176

177+
it("counts array/object tool-result payloads at the LLM boundary", () => {
178+
const objectPayload = {
179+
rows: Array.from({ length: 120 }, (_, index) => ({
180+
path: `/tmp/generated-${index}.txt`,
181+
body: "x".repeat(1_500),
182+
})),
183+
};
184+
const messages = [makeJsonToolResultMessage(objectPayload)];
185+
const estimatedPromptTokens = estimateLlmBoundaryTokenPressure({
186+
messages,
187+
systemPrompt: "sys",
188+
prompt: "continue",
189+
});
190+
191+
expect(estimatedPromptTokens).toBeGreaterThan(80_000);
192+
193+
const result = shouldPreemptivelyCompactBeforePrompt({
194+
messages,
195+
systemPrompt: "sys",
196+
prompt: "continue",
197+
contextTokenBudget: 96_000,
198+
reserveTokens: 20_000,
199+
});
200+
201+
expect(result.route).not.toBe("fits");
202+
expect(result.estimatedPromptTokens).toBe(estimatedPromptTokens);
203+
expect(result.overflowTokens).toBeGreaterThan(0);
204+
});
205+
206+
it("counts assistant tool-call arguments instead of trusting text-only token estimates", () => {
207+
const messages = [
208+
makeAssistantToolCall({
209+
queryPlan: "find relevant files",
210+
candidates: Array.from({ length: 100 }, (_, index) => ({
211+
path: `/repo/file-${index}.ts`,
212+
content: "z".repeat(1_000),
213+
})),
214+
}),
215+
];
216+
const estimatedPromptTokens = estimatePrePromptTokens({
217+
messages,
218+
systemPrompt: "sys",
219+
prompt: "continue",
220+
});
221+
222+
expect(estimatedPromptTokens).toBeGreaterThan(30_000);
223+
});
224+
225+
it("prechecks a regression-sized synthetic tool-heavy transcript as over budget", () => {
226+
const toolResultCharsPerMessage = Math.ceil(427_000 / 120);
227+
const generalCharsPerMessage = Math.ceil((503_000 - 427_000) / 121);
228+
const messages: AgentMessage[] = [];
229+
for (let index = 0; index < 241; index += 1) {
230+
if (index % 2 === 0) {
231+
messages.push(
232+
makeToolResultMessage(
233+
"t".repeat(toolResultCharsPerMessage),
234+
JSON.stringify({ index, payload: "p".repeat(80) }),
235+
),
236+
);
237+
} else {
238+
messages.push(makeAssistantHistory("h".repeat(generalCharsPerMessage)));
239+
}
240+
}
241+
242+
const result = shouldPreemptivelyCompactBeforePrompt({
243+
messages,
244+
systemPrompt: "system".repeat(200),
245+
prompt: "continue",
246+
contextTokenBudget: 200_000,
247+
reserveTokens: 32_000,
248+
});
249+
250+
expect(result.estimatedPromptTokens).toBeGreaterThan(200_000);
251+
expect(result.promptBudgetBeforeReserve).toBe(168_000);
252+
expect(result.route).not.toBe("fits");
253+
expect(result.overflowTokens).toBeGreaterThan(0);
254+
});
255+
149256
it("caps reserve tokens so small context models keep usable prompt budget", () => {
150257
const result = shouldPreemptivelyCompactBeforePrompt({
151258
messages: [makeAssistantHistory("short history")],

0 commit comments

Comments
 (0)