From 23f5cc80a4deec20cf293b5ec0ad53f13d97bb96 Mon Sep 17 00:00:00 2001 From: Dakshay Mehta <50276213+dakshaymehta@users.noreply.github.com> Date: Mon, 16 Feb 2026 14:03:40 -0800 Subject: [PATCH] Agents: wire command poll backoff into process poll --- .../bash-tools.process.poll-timeout.test.ts | 78 +++++++++++++++++++ src/agents/bash-tools.process.ts | 34 ++++++++ 2 files changed, 112 insertions(+) diff --git a/src/agents/bash-tools.process.poll-timeout.test.ts b/src/agents/bash-tools.process.poll-timeout.test.ts index 44e3bb741..a40aba434 100644 --- a/src/agents/bash-tools.process.poll-timeout.test.ts +++ b/src/agents/bash-tools.process.poll-timeout.test.ts @@ -7,9 +7,11 @@ import { resetProcessRegistryForTests, } from "./bash-process-registry.js"; import { createProcessTool } from "./bash-tools.process.js"; +import { resetDiagnosticSessionStateForTest } from "../logging/diagnostic-session-state.js"; afterEach(() => { resetProcessRegistryForTests(); + resetDiagnosticSessionStateForTest(); }); function createBackgroundSession(id: string): ProcessSession { @@ -98,3 +100,79 @@ test("process poll accepts string timeout values", async () => { vi.useRealTimers(); } }); + +test("process poll exposes adaptive retryInMs for repeated no-output polls", async () => { + const processTool = createProcessTool(); + const sessionId = "sess-retry"; + const session = createBackgroundSession(sessionId); + addSession(session); + + const poll1 = await processTool.execute("toolcall-1", { + action: "poll", + sessionId, + }); + const poll2 = await processTool.execute("toolcall-2", { + action: "poll", + sessionId, + }); + const poll3 = await processTool.execute("toolcall-3", { + action: "poll", + sessionId, + }); + const poll4 = await processTool.execute("toolcall-4", { + action: "poll", + sessionId, + }); + const poll5 = await processTool.execute("toolcall-5", { + action: "poll", + sessionId, + }); + + expect((poll1.details as { retryInMs?: number }).retryInMs).toBe(5000); + expect((poll2.details as { retryInMs?: number }).retryInMs).toBe(10000); + expect((poll3.details as { retryInMs?: number }).retryInMs).toBe(30000); + expect((poll4.details as { retryInMs?: number }).retryInMs).toBe(60000); + expect((poll5.details as { retryInMs?: number }).retryInMs).toBe(60000); +}); + +test("process poll resets retryInMs when output appears and clears on completion", async () => { + const processTool = createProcessTool(); + const sessionId = "sess-reset"; + const session = createBackgroundSession(sessionId); + addSession(session); + + const poll1 = await processTool.execute("toolcall-1", { + action: "poll", + sessionId, + }); + const poll2 = await processTool.execute("toolcall-2", { + action: "poll", + sessionId, + }); + expect((poll1.details as { retryInMs?: number }).retryInMs).toBe(5000); + expect((poll2.details as { retryInMs?: number }).retryInMs).toBe(10000); + + appendOutput(session, "stdout", "step complete\n"); + const pollWithOutput = await processTool.execute("toolcall-output", { + action: "poll", + sessionId, + }); + expect((pollWithOutput.details as { retryInMs?: number }).retryInMs).toBe(5000); + + markExited(session, 0, null, "completed"); + const pollCompleted = await processTool.execute("toolcall-completed", { + action: "poll", + sessionId, + }); + const completedDetails = pollCompleted.details as { status?: string; retryInMs?: number }; + expect(completedDetails.status).toBe("completed"); + expect(completedDetails.retryInMs).toBeUndefined(); + + const pollFinished = await processTool.execute("toolcall-finished", { + action: "poll", + sessionId, + }); + const finishedDetails = pollFinished.details as { status?: string; retryInMs?: number }; + expect(finishedDetails.status).toBe("completed"); + expect(finishedDetails.retryInMs).toBeUndefined(); +}); diff --git a/src/agents/bash-tools.process.ts b/src/agents/bash-tools.process.ts index 38a8ac357..b9378f758 100644 --- a/src/agents/bash-tools.process.ts +++ b/src/agents/bash-tools.process.ts @@ -1,8 +1,10 @@ import type { AgentTool, AgentToolResult } from "@mariozechner/pi-agent-core"; import { Type } from "@sinclair/typebox"; import { formatDurationCompact } from "../infra/format-time/format-duration.ts"; +import { getDiagnosticSessionState } from "../logging/diagnostic-session-state.js"; import { killProcessTree } from "../process/kill-tree.js"; import { getProcessSupervisor } from "../process/supervisor/index.js"; +import { recordCommandPoll, resetCommandPollCount } from "./command-poll-backoff.js"; import { type ProcessSession, deleteSession, @@ -96,6 +98,24 @@ function failText(text: string): AgentToolResult { }; } +function recordPollRetrySuggestion(sessionId: string, hasNewOutput: boolean): number | undefined { + try { + const sessionState = getDiagnosticSessionState({ sessionId }); + return recordCommandPoll(sessionState, sessionId, hasNewOutput); + } catch { + return undefined; + } +} + +function resetPollRetrySuggestion(sessionId: string): void { + try { + const sessionState = getDiagnosticSessionState({ sessionId }); + resetCommandPollCount(sessionState, sessionId); + } catch { + // Ignore diagnostics state failures for process tool behavior. + } +} + export function createProcessTool( defaults?: ProcessToolDefaults, // oxlint-disable-next-line typescript/no-explicit-any @@ -262,6 +282,7 @@ export function createProcessTool( case "poll": { if (!scopedSession) { if (scopedFinished) { + resetPollRetrySuggestion(params.sessionId); return { content: [ { @@ -287,6 +308,7 @@ export function createProcessTool( }, }; } + resetPollRetrySuggestion(params.sessionId); return failText(`No session found for ${params.sessionId}`); } if (!scopedSession.backgrounded) { @@ -320,6 +342,13 @@ export function createProcessTool( : "failed" : "running"; const output = [stdout.trimEnd(), stderr.trimEnd()].filter(Boolean).join("\n").trim(); + const hasNewOutput = output.length > 0; + const retryInMs = exited + ? undefined + : recordPollRetrySuggestion(params.sessionId, hasNewOutput); + if (exited) { + resetPollRetrySuggestion(params.sessionId); + } return { content: [ { @@ -339,6 +368,7 @@ export function createProcessTool( exitCode: exited ? exitCode : undefined, aggregated: scopedSession.aggregated, name: deriveSessionName(scopedSession.command), + ...(typeof retryInMs === "number" ? { retryInMs } : {}), }, }; } @@ -549,6 +579,7 @@ export function createProcessTool( } markExited(scopedSession, null, "SIGKILL", "failed"); } + resetPollRetrySuggestion(params.sessionId); return { content: [ { @@ -567,6 +598,7 @@ export function createProcessTool( case "clear": { if (scopedFinished) { + resetPollRetrySuggestion(params.sessionId); deleteSession(params.sessionId); return { content: [{ type: "text", text: `Cleared session ${params.sessionId}.` }], @@ -601,6 +633,7 @@ export function createProcessTool( markExited(scopedSession, null, "SIGKILL", "failed"); deleteSession(params.sessionId); } + resetPollRetrySuggestion(params.sessionId); return { content: [ { @@ -617,6 +650,7 @@ export function createProcessTool( }; } if (scopedFinished) { + resetPollRetrySuggestion(params.sessionId); deleteSession(params.sessionId); return { content: [{ type: "text", text: `Removed session ${params.sessionId}.` }],