Merge 710eedef4f53278c889d36ab9776e7886d7187e9 into 598f1826d8b2bc969aace2c6459824737667218c

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();
+  });
 });

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();

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,