diff --git a/src/agents/pi-embedded-runner.guard.waitforidle-before-flush.test.ts b/src/agents/pi-embedded-runner.guard.waitforidle-before-flush.test.ts index 207e721ac81..b627c222bd7 100644 --- a/src/agents/pi-embedded-runner.guard.waitforidle-before-flush.test.ts +++ b/src/agents/pi-embedded-runner.guard.waitforidle-before-flush.test.ts @@ -1,7 +1,11 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import { SessionManager } from "@mariozechner/pi-coding-agent"; import { afterEach, describe, expect, it, vi } from "vitest"; -import { flushPendingToolResultsAfterIdle } from "./pi-embedded-runner/wait-for-idle-before-flush.js"; +import { + flushPendingToolResultsAfterIdle, + PROBE_WAIT_FOR_IDLE_TIMEOUT_MS, + resolvePostRunIdleFlushTimeoutMs, +} from "./pi-embedded-runner/wait-for-idle-before-flush.js"; import { guardSessionManager } from "./session-tool-result-guard-wrapper.js"; function assistantToolCall(id: string): AgentMessage { @@ -138,4 +142,20 @@ describe("flushPendingToolResultsAfterIdle", () => { }); expect(vi.getTimerCount()).toBe(0); }); + + it("uses a short best-effort idle wait for probe sessions", () => { + expect( + resolvePostRunIdleFlushTimeoutMs({ + sessionId: "probe-xai-123", + timeoutMs: 30_000, + }), + ).toBe(PROBE_WAIT_FOR_IDLE_TIMEOUT_MS); + expect( + resolvePostRunIdleFlushTimeoutMs({ + sessionId: "session:test", + timeoutMs: 30_000, + }), + ).toBe(30_000); + expect(resolvePostRunIdleFlushTimeoutMs({ sessionId: "session:test" })).toBeUndefined(); + }); }); diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index d785218f819..74f228839c9 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -138,7 +138,10 @@ import { collectAllowedToolNames } from "../tool-name-allowlist.js"; import { installToolResultContextGuard } from "../tool-result-context-guard.js"; import { splitSdkTools } from "../tool-split.js"; import { describeUnknownError, mapThinkingLevel } from "../utils.js"; -import { flushPendingToolResultsAfterIdle } from "../wait-for-idle-before-flush.js"; +import { + flushPendingToolResultsAfterIdle, + resolvePostRunIdleFlushTimeoutMs, +} from "../wait-for-idle-before-flush.js"; import { waitForCompactionRetryWithAggregateTimeout } from "./compaction-retry-aggregate-timeout.js"; import { resolveRunTimeoutDuringCompaction, @@ -2451,6 +2454,9 @@ export async function runEmbeddedAttempt( await flushPendingToolResultsAfterIdle({ agent: activeSession?.agent, sessionManager, + timeoutMs: resolvePostRunIdleFlushTimeoutMs({ + sessionId: params.sessionId, + }), clearPendingOnTimeout: true, }); activeSession.dispose(); @@ -3197,6 +3203,9 @@ export async function runEmbeddedAttempt( await flushPendingToolResultsAfterIdle({ agent: session?.agent, sessionManager, + timeoutMs: resolvePostRunIdleFlushTimeoutMs({ + sessionId: params.sessionId, + }), clearPendingOnTimeout: true, }); session?.dispose(); diff --git a/src/agents/pi-embedded-runner/wait-for-idle-before-flush.ts b/src/agents/pi-embedded-runner/wait-for-idle-before-flush.ts index 71b661aadb7..a0a0ed54aa0 100644 --- a/src/agents/pi-embedded-runner/wait-for-idle-before-flush.ts +++ b/src/agents/pi-embedded-runner/wait-for-idle-before-flush.ts @@ -8,6 +8,17 @@ type ToolResultFlushManager = { }; export const DEFAULT_WAIT_FOR_IDLE_TIMEOUT_MS = 30_000; +export const PROBE_WAIT_FOR_IDLE_TIMEOUT_MS = 50; + +export function resolvePostRunIdleFlushTimeoutMs(opts: { + sessionId?: string; + timeoutMs?: number; +}): number | undefined { + if (opts.sessionId?.startsWith("probe-")) { + return PROBE_WAIT_FOR_IDLE_TIMEOUT_MS; + } + return opts.timeoutMs; +} async function waitForAgentIdleBestEffort( agent: IdleAwareAgent | null | undefined,