Skip to content

Commit 200dea1

Browse files
ZengWen-DTZengWen-DT
committed
fix: surface stalled-session interruptions
- wrapToolWithAbortSignal: race tool execution against run abort signal so non-cooperative tools no longer block an aborted turn - Stale reply-operation recovery now preserves run_stalled as a user-visible interruption instead of silently dropping the draft - Both the error handler and fallback-cycle settlement check for stalled operations before the broad user-abort predicate - isReplyOperationStalled and buildStalledRunReplyPayload helpers keep stalled handling explicit and separate from abort paths Closes #103905
1 parent fceae4d commit 200dea1

8 files changed

Lines changed: 220 additions & 2 deletions

‎src/agents/agent-tools.abort.test.ts‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import { describe, expect, it } from "vitest";
2+
import { wrapToolWithAbortSignal } from "./agent-tools.abort.js";
3+
import type { AnyAgentTool } from "./agent-tools.types.js";
4+
5+
describe("wrapToolWithAbortSignal", () => {
6+
it("returns promptly on run abort while a non-cooperative tool settles later", async () => {
7+
const abort = new AbortController();
8+
let resolveTool: (() => void) | undefined;
9+
let receivedSignal: AbortSignal | undefined;
10+
const tool = {
11+
name: "slow-tool",
12+
description: "waits until released",
13+
parameters: {},
14+
execute: async (_toolCallId: string, _params: unknown, signal?: AbortSignal) => {
15+
receivedSignal = signal;
16+
await new Promise<void>((resolve) => {
17+
resolveTool = resolve;
18+
});
19+
return { content: [] };
20+
},
21+
} as unknown as AnyAgentTool;
22+
const execute = wrapToolWithAbortSignal(tool, abort.signal).execute;
23+
if (!execute) {
24+
throw new Error("Expected wrapped tool execute function");
25+
}
26+
27+
const result = execute("tool-call", {}, undefined, undefined);
28+
abort.abort();
29+
30+
await expect(result).rejects.toMatchObject({ name: "AbortError", message: "Aborted" });
31+
expect(receivedSignal?.aborted).toBe(true);
32+
33+
resolveTool?.();
34+
});
35+
});

‎src/agents/agent-tools.abort.ts‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,25 @@ function throwAbortError(): never {
1313
throw createAbortError("Aborted");
1414
}
1515

16+
function rejectOnAbort(signal: AbortSignal): { promise: Promise<never>; dispose: () => void } {
17+
let onAbort: (() => void) | undefined;
18+
const promise = new Promise<never>((_, reject) => {
19+
onAbort = () => reject(createAbortError("Aborted"));
20+
signal.addEventListener("abort", onAbort, { once: true });
21+
if (signal.aborted) {
22+
onAbort();
23+
}
24+
});
25+
return {
26+
promise,
27+
dispose: () => {
28+
if (onAbort) {
29+
signal.removeEventListener("abort", onAbort);
30+
}
31+
},
32+
};
33+
}
34+
1635
/** Wrap a tool so every execute call observes the supplied run abort signal. */
1736
export function wrapToolWithAbortSignal(
1837
tool: AnyAgentTool,
@@ -32,7 +51,17 @@ export function wrapToolWithAbortSignal(
3251
if (combinedSignal.aborted) {
3352
throwAbortError();
3453
}
35-
return await execute(toolCallId, params, combinedSignal, onUpdate);
54+
const aborted = rejectOnAbort(combinedSignal);
55+
try {
56+
// Tool cancellation is cooperative, so the handler may settle after
57+
// the caller exits. The race keeps the run responsive and observes late rejection.
58+
return await Promise.race([
59+
execute(toolCallId, params, combinedSignal, onUpdate),
60+
aborted.promise,
61+
]);
62+
} finally {
63+
aborted.dispose();
64+
}
3665
},
3766
};
3867
copyPluginToolMeta(tool, wrappedTool);

‎src/auto-reply/reply/agent-runner-error-handler.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import {
3838
buildAuthProfileFailoverFailureText,
3939
buildExternalRunFailureReply,
4040
buildRateLimitCooldownMessage,
41+
buildStalledRunReplyPayload,
4142
hasBillingAttemptSummary,
4243
isNonDirectConversationContext,
4344
isPureTransientRateLimitSummary,
@@ -48,10 +49,12 @@ import {
4849
} from "./agent-runner-failure-reply.js";
4950
import type { AgentFallbackCycleState } from "./agent-runner-fallback-cycle.js";
5051
import type { AgentTurnTimingTracker } from "./agent-runner-turn-timing.js";
52+
import { drainPendingToolTasks } from "./pending-tool-task-drain.js";
5153
import { classifyProviderRequestError } from "./provider-request-error-classifier.js";
5254
import {
5355
buildRestartLifecycleReplyText,
5456
isReplyOperationRestartAbort,
57+
isReplyOperationStalled,
5558
isReplyOperationUserAbort,
5659
resolveRestartLifecycleError,
5760
} from "./reply-operation-abort.js";
@@ -229,6 +232,13 @@ export async function handleAgentExecutionError(params: {
229232
isTransientHttpError(message) ||
230233
(isFailoverError(err) && (err.reason === "timeout" || err.reason === "server_error"));
231234

235+
// Stale recovery records run_stalled before aborting this operation.
236+
// Handle it before the broad user-abort predicate so the interruption remains visible.
237+
if (isReplyOperationStalled(turn.replyOperation)) {
238+
takePendingLifecycleTerminal()?.emit("error", err);
239+
await drainPendingToolTasks({ tasks: turn.pendingToolTasks, onTimeout: logVerbose });
240+
return { kind: "final", payload: buildStalledRunReplyPayload(turn.isHeartbeat) };
241+
}
232242
const replyOperationAbortAction = resolveReplyOperationAbortAction(err);
233243
if (replyOperationAbortAction) {
234244
return replyOperationAbortAction;

‎src/auto-reply/reply/agent-runner-execution-stalled-recovery.test.ts‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
import { describe, expect, it } from "vitest";
2+
import {
3+
createMinimalRunAgentTurnParams,
4+
getRunAgentTurnWithFallback,
5+
setupAgentRunnerExecutionTestState,
6+
} from "./agent-runner-execution.test-support.js";
7+
import { HEARTBEAT_EXTERNAL_RUN_FAILURE_TEXT } from "./agent-runner-failure-copy.js";
8+
import { createReplyOperation, expireStaleReplyOperation } from "./reply-run-registry.js";
9+
10+
const state = setupAgentRunnerExecutionTestState();
11+
12+
function createStalledReplyOperation(sessionId: string) {
13+
const replyOperation = createReplyOperation({
14+
sessionKey: `agent:main:${sessionId}`,
15+
sessionId,
16+
resetTriggered: false,
17+
});
18+
replyOperation.setPhase("running");
19+
return replyOperation;
20+
}
21+
22+
describe("runAgentTurnWithFallback: stalled recovery", () => {
23+
it("surfaces a stalled reply operation after the embedded run returns no payload", async () => {
24+
const replyOperation = createStalledReplyOperation("stalled-empty-reply");
25+
state.runEmbeddedAgentMock.mockImplementationOnce(async () => {
26+
expect(expireStaleReplyOperation(replyOperation, "no_activity")).toBe(true);
27+
expect(replyOperation.abortSignal.aborted).toBe(true);
28+
return { payloads: [], meta: {} };
29+
});
30+
let releaseToolTask: () => void = () => undefined;
31+
const pendingToolTask = new Promise<void>((resolve) => {
32+
releaseToolTask = resolve;
33+
});
34+
35+
const runAgentTurnWithFallback = await getRunAgentTurnWithFallback();
36+
const pending = runAgentTurnWithFallback({
37+
...createMinimalRunAgentTurnParams({ replyOperation }),
38+
pendingToolTasks: new Set([pendingToolTask]),
39+
});
40+
let settled = false;
41+
void pending.then(() => {
42+
settled = true;
43+
});
44+
await new Promise<void>((resolve) => {
45+
setImmediate(resolve);
46+
});
47+
expect(settled).toBe(false);
48+
releaseToolTask();
49+
50+
await expect(pending).resolves.toEqual({
51+
kind: "final",
52+
payload: {
53+
text: "⚠️ This turn was interrupted because it stopped making progress. Please try again.",
54+
isError: true,
55+
},
56+
});
57+
});
58+
59+
it("surfaces a stalled reply operation after its aborted embedded run rejects", async () => {
60+
const replyOperation = createStalledReplyOperation("stalled-rejected-reply");
61+
state.runEmbeddedAgentMock.mockImplementationOnce(async () => {
62+
expect(expireStaleReplyOperation(replyOperation, "no_activity")).toBe(true);
63+
expect(replyOperation.abortSignal.aborted).toBe(true);
64+
throw new Error("embedded run aborted after stale recovery");
65+
});
66+
let releaseToolTask: () => void = () => undefined;
67+
const pendingToolTask = new Promise<void>((resolve) => {
68+
releaseToolTask = resolve;
69+
});
70+
71+
const runAgentTurnWithFallback = await getRunAgentTurnWithFallback();
72+
const pending = runAgentTurnWithFallback({
73+
...createMinimalRunAgentTurnParams({ replyOperation }),
74+
pendingToolTasks: new Set([pendingToolTask]),
75+
});
76+
let settled = false;
77+
void pending.then(() => {
78+
settled = true;
79+
});
80+
await new Promise<void>((resolve) => {
81+
setImmediate(resolve);
82+
});
83+
expect(settled).toBe(false);
84+
releaseToolTask();
85+
86+
await expect(pending).resolves.toEqual({
87+
kind: "final",
88+
payload: {
89+
text: "⚠️ This turn was interrupted because it stopped making progress. Please try again.",
90+
isError: true,
91+
},
92+
});
93+
});
94+
95+
it("uses heartbeat failure copy for a stalled heartbeat operation", async () => {
96+
const replyOperation = createStalledReplyOperation("stalled-heartbeat");
97+
state.runEmbeddedAgentMock.mockImplementationOnce(async () => {
98+
expect(expireStaleReplyOperation(replyOperation, "no_activity")).toBe(true);
99+
return { payloads: [], meta: {} };
100+
});
101+
102+
const runAgentTurnWithFallback = await getRunAgentTurnWithFallback();
103+
const result = await runAgentTurnWithFallback({
104+
...createMinimalRunAgentTurnParams({ replyOperation }),
105+
isHeartbeat: true,
106+
});
107+
108+
expect(result).toEqual({
109+
kind: "final",
110+
payload: {
111+
text: HEARTBEAT_EXTERNAL_RUN_FAILURE_TEXT,
112+
isError: true,
113+
},
114+
});
115+
});
116+
});

‎src/auto-reply/reply/agent-runner-failure-copy.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@ export const GENERIC_EXTERNAL_RUN_FAILURE_TEXT =
55
export const HEARTBEAT_EXTERNAL_RUN_FAILURE_TEXT =
66
"⚠️ Heartbeat check failed before it could produce an update. The main chat session remains available.";
77

8+
export const STALLED_RUN_FAILURE_TEXT =
9+
"⚠️ This turn was interrupted because it stopped making progress. Please try again.";
10+
811
/** True when text is exactly the generic external run failure copy. */
912
function isGenericExternalRunFailureText(text: string | undefined): boolean {
1013
return text?.trim() === GENERIC_EXTERNAL_RUN_FAILURE_TEXT;

‎src/auto-reply/reply/agent-runner-failure-reply.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import type { ReplyPayload } from "../types.js";
3636
import {
3737
GENERIC_EXTERNAL_RUN_FAILURE_TEXT,
3838
HEARTBEAT_EXTERNAL_RUN_FAILURE_TEXT,
39+
STALLED_RUN_FAILURE_TEXT,
3940
} from "./agent-runner-failure-copy.js";
4041
import { classifyProviderRequestError } from "./provider-request-error-classifier.js";
4142

@@ -459,6 +460,12 @@ export function markAgentRunFailureReplyPayload<T extends ReplyPayload>(payload:
459460
return marked;
460461
}
461462

463+
export function buildStalledRunReplyPayload(isHeartbeat: boolean): ReplyPayload {
464+
return markAgentRunFailureReplyPayload({
465+
text: isHeartbeat ? HEARTBEAT_EXTERNAL_RUN_FAILURE_TEXT : STALLED_RUN_FAILURE_TEXT,
466+
});
467+
}
468+
462469
export function buildTerminalAgentRunFailureReplyPayload(params: {
463470
isHeartbeat?: boolean;
464471
sessionCtx: ExternalFailureConversationContext;

‎src/auto-reply/reply/agent-runner-fallback-settlement.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,10 @@ import { defaultRuntime } from "../../runtime.js";
1010
import { SILENT_REPLY_TOKEN } from "../tokens.js";
1111
import { resolveAgentLifecycleTerminalMetadata } from "./agent-lifecycle-terminal.js";
1212
import { buildContextOverflowRecoveryText } from "./agent-runner-context-recovery.js";
13-
import { markAgentRunFailureReplyPayload } from "./agent-runner-failure-reply.js";
13+
import {
14+
buildStalledRunReplyPayload,
15+
markAgentRunFailureReplyPayload,
16+
} from "./agent-runner-failure-reply.js";
1417
import type { AgentFallbackCandidatesResult } from "./agent-runner-fallback-candidate.js";
1518
import type {
1619
AgentFallbackCycleParams,
@@ -23,6 +26,7 @@ import {
2326
} from "./provider-request-error-classifier.js";
2427
import {
2528
isReplyOperationRestartAbort,
29+
isReplyOperationStalled,
2630
isReplyOperationUserAbort,
2731
} from "./reply-operation-abort.js";
2832

@@ -49,6 +53,16 @@ export async function settleAgentFallbackCycle(params: {
4953
? cycle.runAbortSignal?.reason
5054
: createAgentRunRestartAbortError();
5155
}
56+
// Stale recovery records run_stalled before aborting this operation.
57+
// Prefer that terminal cause so the broader abort-signal fallback does not hide its reply.
58+
if (isReplyOperationStalled(turn.replyOperation)) {
59+
settledLifecycleTerminal?.emit(
60+
"error",
61+
new Error("Reply operation expired after making no progress"),
62+
);
63+
await drainPendingToolTasks({ tasks: turn.pendingToolTasks, onTimeout: logVerbose });
64+
return { kind: "final", payload: buildStalledRunReplyPayload(turn.isHeartbeat) };
65+
}
5266
if (isReplyOperationUserAbort(turn.replyOperation)) {
5367
settledLifecycleTerminal?.emit("end", runResult);
5468
await drainPendingToolTasks({ tasks: turn.pendingToolTasks, onTimeout: logVerbose });

‎src/auto-reply/reply/reply-operation-abort.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ export function isReplyOperationRestartAbort(replyOperation?: ReplyOperation): b
2929
return abortSignal?.aborted === true && isAgentRunRestartAbortReason(abortSignal.reason);
3030
}
3131

32+
export function isReplyOperationStalled(replyOperation?: ReplyOperation): boolean {
33+
return replyOperation?.result?.kind === "failed" && replyOperation.result.code === "run_stalled";
34+
}
35+
3236
export function resolveRestartLifecycleError(
3337
error: unknown,
3438
): GatewayDrainingError | CommandLaneClearedError | undefined {

0 commit comments

Comments
 (0)