Agents: wire command poll backoff into process poll

This commit is contained in:
Dakshay Mehta
2026-02-16 14:03:40 -08:00
committed by Peter Steinberger
parent 054745a7e0
commit 23f5cc80a4
2 changed files with 112 additions and 0 deletions

View File

@@ -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<unknown> {
};
}
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}.` }],