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 ba23e83c6c06..aa66a4662f73 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 @@ -211,6 +211,32 @@ describe("runEmbeddedAttempt context engine sessionKey forwarding", () => { vi.restoreAllMocks(); }); + it("flushes block replies again after compaction retry wait resolves", async () => { + const order: string[] = []; + let flushCount = 0; + const onBlockReplyFlush = vi.fn(async () => { + flushCount += 1; + order.push(`flush-${flushCount}`); + }); + hoisted.waitForCompactionRetryWithAggregateTimeoutMock.mockImplementation(async () => { + order.push("retry-wait"); + return { timedOut: false }; + }); + + await createContextEngineAttemptRunner({ + contextEngine: createContextEngineBootstrapAndAssemble(), + sessionKey, + tempPaths, + attemptOverrides: { + onBlockReplyFlush, + }, + }); + + expect(onBlockReplyFlush).toHaveBeenCalledTimes(2); + expect(hoisted.waitForCompactionRetryWithAggregateTimeoutMock).toHaveBeenCalledTimes(1); + expect(order).toEqual(["flush-1", "retry-wait", "flush-2"]); + }); + it("enables Tool Search controls for embedded PI runs when configured", async () => { await createContextEngineAttemptRunner({ contextEngine: { diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts index ff60234e6cf2..41bfd9ad70a9 100644 --- a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts +++ b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts @@ -32,6 +32,8 @@ type AcquireSessionWriteLockFn = typeof import("../../session-write-lock.js").acquireSessionWriteLock; type ShouldPreemptivelyCompactBeforePromptFn = typeof import("./preemptive-compaction.js").shouldPreemptivelyCompactBeforePrompt; +type WaitForCompactionRetryWithAggregateTimeoutFn = + typeof import("./compaction-retry-aggregate-timeout.js").waitForCompactionRetryWithAggregateTimeout; type SubscriptionMock = ReturnType; type UnknownMock = Mock<(...args: unknown[]) => unknown>; @@ -92,6 +94,7 @@ type AttemptSpawnWorkspaceHoisted = { (sessionKey: string | undefined, config: unknown) => number | undefined >; limitHistoryTurnsMock: Mock<(messages: T, limit: number | undefined) => T>; + waitForCompactionRetryWithAggregateTimeoutMock: Mock; preemptiveCompactionCalls: Parameters[0][]; systemPromptOverrideTexts: string[]; sessionManager: SessionManagerMocks; @@ -191,6 +194,10 @@ const hoisted = vi.hoisted((): AttemptSpawnWorkspaceHoisted => { const limitHistoryTurnsMock = vi.fn<(messages: T, limit: number | undefined) => T>( (messages) => messages, ); + const waitForCompactionRetryWithAggregateTimeoutMock = + vi.fn(async () => ({ + timedOut: false, + })); const preemptiveCompactionCalls: Parameters[0][] = []; const systemPromptOverrideTexts: string[] = []; const sessionManager = { @@ -234,6 +241,7 @@ const hoisted = vi.hoisted((): AttemptSpawnWorkspaceHoisted => { detectAndLoadPromptImagesMock, getHistoryLimitFromSessionKeyMock, limitHistoryTurnsMock, + waitForCompactionRetryWithAggregateTimeoutMock, preemptiveCompactionCalls, systemPromptOverrideTexts, sessionManager, @@ -790,10 +798,9 @@ vi.mock("../utils.js", () => ({ })); vi.mock("./compaction-retry-aggregate-timeout.js", () => ({ - waitForCompactionRetryWithAggregateTimeout: async () => ({ - timedOut: false, - aborted: false, - }), + waitForCompactionRetryWithAggregateTimeout: ( + ...args: Parameters + ) => hoisted.waitForCompactionRetryWithAggregateTimeoutMock(...args), })); vi.mock("./compaction-timeout.js", () => ({ @@ -975,6 +982,9 @@ export function resetEmbeddedAttemptHarness( hoisted.runContextEngineMaintenanceMock.mockReset().mockResolvedValue(undefined); hoisted.getHistoryLimitFromSessionKeyMock.mockReset().mockReturnValue(undefined); hoisted.limitHistoryTurnsMock.mockReset().mockImplementation((messages) => messages); + hoisted.waitForCompactionRetryWithAggregateTimeoutMock + .mockReset() + .mockResolvedValue({ timedOut: false }); hoisted.preemptiveCompactionCalls.length = 0; hoisted.systemPromptOverrideTexts.length = 0; hoisted.sessionManager.getLeafEntry.mockReset().mockReturnValue(null); diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 5b70a28df78f..22795bdd0907 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -4483,6 +4483,11 @@ export async function runEmbeddedAttempt( `proceeding with pre-compaction state runId=${params.runId} sessionId=${params.sessionId}`, ); } + } else if (onBlockReplyFlush) { + // Retry-generated blocks can still be draining when the compaction + // retry wait resolves; this second drain is idempotent when no new + // blocks were produced. + await onBlockReplyFlush(); } } catch (err) { if (isRunnerAbortError(err)) {