diff --git a/src/agents/pi-embedded-runner/run.overflow-compaction.loop.test.ts b/src/agents/pi-embedded-runner/run.overflow-compaction.loop.test.ts index 448f9daa7c1..0e7e6870b84 100644 --- a/src/agents/pi-embedded-runner/run.overflow-compaction.loop.test.ts +++ b/src/agents/pi-embedded-runner/run.overflow-compaction.loop.test.ts @@ -253,6 +253,41 @@ describe("overflow compaction in run loop", () => { expect(result.meta.error).toBeUndefined(); }); + it("continues from the transcript after mid-turn precheck truncation handled the overflow", async () => { + mockedRunEmbeddedAttempt + .mockResolvedValueOnce( + makeAttemptResult({ + promptError: null, + preflightRecovery: { + route: "truncate_tool_results_only", + source: "mid-turn", + handled: true, + truncatedCount: 2, + }, + }), + ) + .mockResolvedValueOnce(makeAttemptResult({ promptError: null })); + + const result = await runEmbeddedPiAgent(baseParams); + + expect(mockedCompactDirect).not.toHaveBeenCalled(); + expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2); + expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + prompt: expect.stringContaining("Continue from the current transcript"), + }), + ); + expect(mockedRunEmbeddedAttempt).not.toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ prompt: baseParams.prompt }), + ); + expect(mockedLog.info).toHaveBeenCalledWith( + expect.stringContaining("retrying from current transcript"), + ); + expect(result.meta.error).toBeUndefined(); + }); + it("falls back to compaction when early truncate-only recovery does not help", async () => { mockedRunEmbeddedAttempt .mockResolvedValueOnce( @@ -286,6 +321,44 @@ describe("overflow compaction in run loop", () => { expect(result.meta.error).toBeUndefined(); }); + it("continues from the transcript after mid-turn precheck compaction", async () => { + mockedRunEmbeddedAttempt + .mockResolvedValueOnce( + makeAttemptResult({ + promptError: makeOverflowError( + "Context overflow: prompt too large for the model (mid-turn precheck).", + ), + promptErrorSource: "precheck", + preflightRecovery: { route: "compact_only", source: "mid-turn" }, + }), + ) + .mockResolvedValueOnce(makeAttemptResult({ promptError: null })); + + mockedCompactDirect.mockResolvedValueOnce( + makeCompactionSuccess({ + summary: "Compacted after mid-turn precheck", + firstKeptEntryId: "entry-8", + tokensBefore: 155000, + }), + ); + + const result = await runEmbeddedPiAgent(baseParams); + + expect(mockedCompactDirect).toHaveBeenCalledTimes(1); + expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2); + expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + prompt: expect.stringContaining("Continue from the current transcript"), + }), + ); + expect(mockedRunEmbeddedAttempt).not.toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ prompt: baseParams.prompt }), + ); + expect(result.meta.error).toBeUndefined(); + }); + it("runs post-compaction tool-result truncation before retry for mixed precheck routes", async () => { mockedRunEmbeddedAttempt .mockResolvedValueOnce( diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index c179c327596..cf0e586193d 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -163,6 +163,8 @@ type ApiKeyInfo = ResolvedProviderAuth; const MAX_SAME_MODEL_IDLE_TIMEOUT_RETRIES = 1; const EMBEDDED_RUN_LANE_TIMEOUT_GRACE_MS = 30_000; +const MID_TURN_PRECHECK_CONTINUATION_PROMPT = + "Continue from the current transcript after the latest tool result. Do not repeat the original user request, and do not rerun completed tools unless the transcript shows they are still needed."; type EmbeddedRunAttemptForRunner = Awaited>; function resolveEmbeddedRunLaneTimeoutMs(timeoutMs: number): number | undefined { @@ -759,6 +761,7 @@ export async function runEmbeddedPiAgent( let planningOnlyRetryInstruction: string | null = null; let reasoningOnlyRetryInstruction: string | null = null; let emptyResponseRetryInstruction: string | null = null; + let nextAttemptPromptOverride: string | null = null; const ackExecutionFastPathInstruction = resolveAckExecutionFastPathInstruction({ provider, modelId, @@ -963,7 +966,9 @@ export async function runEmbeddedPiAgent( await fs.mkdir(resolvedWorkspace, { recursive: true }); const basePrompt = - provider === "anthropic" ? scrubAnthropicRefusalMagic(params.prompt) : params.prompt; + nextAttemptPromptOverride ?? + (provider === "anthropic" ? scrubAnthropicRefusalMagic(params.prompt) : params.prompt); + nextAttemptPromptOverride = null; const promptAdditions = [ ackExecutionFastPathInstruction, planningOnlyRetryInstruction, @@ -1201,10 +1206,15 @@ export async function runEmbeddedPiAgent( (attempt.toolMetas?.length ?? 0) === 0 && (attempt.assistantTexts?.length ?? 0) === 0; if (preflightRecovery?.handled) { + const retryingFromTranscript = preflightRecovery.source === "mid-turn"; log.info( `[context-overflow-precheck] early recovery route=${preflightRecovery.route} ` + - `completed for ${provider}/${modelId}; retrying prompt`, + `completed for ${provider}/${modelId}; ` + + (retryingFromTranscript ? "retrying from current transcript" : "retrying prompt"), ); + if (retryingFromTranscript) { + nextAttemptPromptOverride = MID_TURN_PRECHECK_CONTINUATION_PROMPT; + } continue; } const requestedSelection = shouldSwitchToLiveModel({ @@ -1385,6 +1395,9 @@ export async function runEmbeddedPiAgent( log.warn( `context overflow persisted after in-attempt compaction (attempt ${overflowCompactionAttempts}/${MAX_OVERFLOW_COMPACTION_ATTEMPTS}); retrying prompt without additional compaction for ${provider}/${modelId}`, ); + if (preflightRecovery?.source === "mid-turn") { + nextAttemptPromptOverride = MID_TURN_PRECHECK_CONTINUATION_PROMPT; + } continue; } // Attempt explicit overflow compaction only when this attempt did not @@ -1512,6 +1525,9 @@ export async function runEmbeddedPiAgent( } autoCompactionCount += 1; log.info(`auto-compaction succeeded for ${provider}/${modelId}; retrying prompt`); + if (preflightRecovery?.source === "mid-turn") { + nextAttemptPromptOverride = MID_TURN_PRECHECK_CONTINUATION_PROMPT; + } continue; } log.warn( @@ -1550,6 +1566,9 @@ export async function runEmbeddedPiAgent( log.info( `[context-overflow-recovery] Truncated ${truncResult.truncatedCount} tool result(s); retrying prompt`, ); + if (preflightRecovery?.source === "mid-turn") { + nextAttemptPromptOverride = MID_TURN_PRECHECK_CONTINUATION_PROMPT; + } continue; } log.warn( 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 c0f43a1a752..74c65599e66 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 @@ -903,7 +903,7 @@ describe("runEmbeddedAttempt context engine mid-turn precheck integration", () = }); expect(result.promptErrorSource).toBe("precheck"); - expect(result.preflightRecovery).toEqual({ route: "compact_only" }); + expect(result.preflightRecovery).toEqual({ route: "compact_only", source: "mid-turn" }); expect(result.messagesSnapshot).toEqual([seedMessage]); }); }); diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index def2558a43a..36f482994a0 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -2357,6 +2357,7 @@ export async function runEmbeddedAttempt( if (truncationResult.truncated) { preflightRecovery = { route: "truncate_tool_results_only", + source: "mid-turn", handled: true, truncatedCount: truncationResult.truncatedCount, }; @@ -2367,7 +2368,7 @@ export async function runEmbeddedAttempt( `handled=true truncatedCount=${truncationResult.truncatedCount}`, ); } else { - preflightRecovery = { route: "compact_only" }; + preflightRecovery = { route: "compact_only", source: "mid-turn" }; promptError = new Error(PREEMPTIVE_OVERFLOW_ERROR_TEXT); promptErrorSource = "precheck"; logMidTurnPrecheck( @@ -2376,7 +2377,7 @@ export async function runEmbeddedAttempt( ); } } else { - preflightRecovery = { route: request.route }; + preflightRecovery = { route: request.route, source: "mid-turn" }; promptError = new Error(PREEMPTIVE_OVERFLOW_ERROR_TEXT); promptErrorSource = "precheck"; logMidTurnPrecheck(request.route); diff --git a/src/agents/pi-embedded-runner/run/types.ts b/src/agents/pi-embedded-runner/run/types.ts index 7e88fa19d28..6f75d152bb6 100644 --- a/src/agents/pi-embedded-runner/run/types.ts +++ b/src/agents/pi-embedded-runner/run/types.ts @@ -68,11 +68,13 @@ export type EmbeddedRunAttemptResult = { preflightRecovery?: | { route: Exclude; + source?: "mid-turn"; handled: true; truncatedCount?: number; } | { route: Exclude; + source?: "mid-turn"; handled?: false; }; sessionIdUsed: string;