Skip to content

Commit 1c5e35f

Browse files
authored
refactor(gateway): split agent runtime helpers (#106544)
* refactor(gateway): split agent runtime helpers * fix(gateway): keep task status type internal
1 parent 70feb13 commit 1c5e35f

4 files changed

Lines changed: 772 additions & 710 deletions

File tree

Lines changed: 261 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,261 @@
1+
import { ErrorCodes, errorShape } from "../../../packages/gateway-protocol/src/index.js";
2+
import {
3+
buildAgentRunTerminalOutcome,
4+
type AgentRunTerminalOutcome,
5+
} from "../../agents/agent-run-terminal-outcome.js";
6+
import { isAgentRunRestartAbortReason } from "../../agents/run-termination.js";
7+
import {
8+
normalizeAgentRunTimeoutPhase,
9+
normalizeProviderStarted,
10+
} from "../../agents/run-timeout-attribution.js";
11+
import { agentCommandFromIngress } from "../../commands/agent.js";
12+
import { isAbortError } from "../../infra/abort-signal.js";
13+
import { clearAgentRunContext } from "../../infra/agent-events.js";
14+
import { readErrorName } from "../../infra/errors.js";
15+
import { defaultRuntime } from "../../runtime.js";
16+
import { createRunningTaskRun } from "../../tasks/detached-task-runtime.js";
17+
import { normalizeDeliveryContext } from "../../utils/delivery-context.shared.js";
18+
import type { ChatAbortControllerEntry } from "../chat-abort.js";
19+
import { formatForLog } from "../ws-log.js";
20+
import { setGatewayDedupeEntries } from "./agent-dedupe.js";
21+
import {
22+
resolveFailedTrackedAgentTaskStatus,
23+
tryFinalizeTrackedAgentTask,
24+
type GatewayAgentTaskTrackingMode,
25+
} from "./agent-task-tracking.js";
26+
import type { GatewayRequestContext, GatewayRequestHandlerOptions } from "./types.js";
27+
28+
function readAgentRunTimeoutAttribution(meta: unknown) {
29+
const record =
30+
meta && typeof meta === "object" && !Array.isArray(meta)
31+
? (meta as Record<string, unknown>)
32+
: undefined;
33+
return {
34+
timeoutPhase: normalizeAgentRunTimeoutPhase(record?.timeoutPhase),
35+
providerStarted: normalizeProviderStarted(record?.providerStarted),
36+
};
37+
}
38+
39+
function isGatewayAbortSignalReason(reason: unknown): boolean {
40+
return reason === undefined || isAbortError(reason) || readErrorName(reason) === "TimeoutError";
41+
}
42+
43+
function isGatewayAgentAbortRejection(error: unknown, signal: AbortSignal): boolean {
44+
if (!signal.aborted) {
45+
return false;
46+
}
47+
if (isAgentRunRestartAbortReason(signal.reason)) {
48+
return true;
49+
}
50+
if (readErrorName(signal.reason) === "TimeoutError") {
51+
return true;
52+
}
53+
if (!isGatewayAbortSignalReason(signal.reason)) {
54+
return false;
55+
}
56+
return isAbortError(error) || readErrorName(error) === "TimeoutError";
57+
}
58+
59+
function resolveGatewayAgentAbortStopReason(signal: AbortSignal): "restart" | "rpc" | "timeout" {
60+
if (isAgentRunRestartAbortReason(signal.reason)) {
61+
return "restart";
62+
}
63+
return readErrorName(signal.reason) === "TimeoutError" ? "timeout" : "rpc";
64+
}
65+
66+
export function resolveAbortedAgentStopReason(entry?: ChatAbortControllerEntry): string {
67+
return entry?.abortStopReason?.trim() || "rpc";
68+
}
69+
70+
export function deleteGatewayDedupeEntries(params: {
71+
dedupe: GatewayRequestContext["dedupe"];
72+
keys: readonly string[];
73+
}) {
74+
for (const key of params.keys) {
75+
params.dedupe.delete(key);
76+
}
77+
}
78+
79+
export function dispatchAgentRunFromGateway(params: {
80+
ingressOpts: Parameters<typeof agentCommandFromIngress>[0];
81+
runId: string;
82+
dedupeKeys: readonly string[];
83+
/**
84+
* Controller whose signal is wired into `ingressOpts.abortSignal`. Used on
85+
* completion to drop the matching `chatAbortControllers` entry without
86+
* touching a same-runId entry owned by a concurrent chat.send.
87+
*/
88+
abortController: AbortController;
89+
cleanupAbortController: () => void;
90+
respond: GatewayRequestHandlerOptions["respond"];
91+
context: GatewayRequestHandlerOptions["context"];
92+
taskTrackingMode: Exclude<GatewayAgentTaskTrackingMode, "plugin_subagent">;
93+
onSettled?: (outcome: {
94+
terminalOutcome: AgentRunTerminalOutcome;
95+
onRecovered?: () => void;
96+
}) => Promise<boolean> | boolean;
97+
}) {
98+
const shouldTrackTask = params.taskTrackingMode === "cli";
99+
let taskTracked = false;
100+
if (shouldTrackTask) {
101+
try {
102+
taskTracked = Boolean(
103+
createRunningTaskRun({
104+
runtime: "cli",
105+
sourceId: params.runId,
106+
ownerKey: params.ingressOpts.sessionKey,
107+
scopeKind: "session",
108+
requesterOrigin: normalizeDeliveryContext({
109+
channel: params.ingressOpts.channel,
110+
to: params.ingressOpts.to,
111+
accountId: params.ingressOpts.accountId,
112+
threadId: params.ingressOpts.threadId,
113+
}),
114+
childSessionKey: params.ingressOpts.sessionKey,
115+
runId: params.runId,
116+
task: params.ingressOpts.message,
117+
deliveryStatus: "not_applicable",
118+
startedAt: Date.now(),
119+
}),
120+
);
121+
} catch (err) {
122+
// Best-effort only: background task tracking must not block agent runs.
123+
// Still surface the swallowed error so non-transient tracking failures stay observable.
124+
params.context.logGateway.warn(
125+
`failed to start tracked agent task ${params.runId}: ${formatForLog(err)}`,
126+
);
127+
}
128+
}
129+
const settle = async (outcome: {
130+
terminalOutcome: AgentRunTerminalOutcome;
131+
onRecovered?: () => void;
132+
}): Promise<boolean> => {
133+
try {
134+
return (await params.onSettled?.(outcome)) ?? true;
135+
} catch (error) {
136+
params.context.logGateway.warn(
137+
`failed to settle agent continuation ${params.runId}: ${formatForLog(error)}`,
138+
);
139+
return false;
140+
}
141+
};
142+
void agentCommandFromIngress(params.ingressOpts, defaultRuntime, params.context.deps)
143+
.then(async (result) => {
144+
const aborted = result?.meta?.aborted === true;
145+
const timeoutAttribution = readAgentRunTimeoutAttribution(result?.meta);
146+
if (taskTracked) {
147+
tryFinalizeTrackedAgentTask({
148+
runId: params.runId,
149+
status: aborted ? "timed_out" : "succeeded",
150+
terminalSummary: aborted ? "aborted" : "completed",
151+
log: params.context.logGateway,
152+
});
153+
}
154+
const payload = {
155+
runId: params.runId,
156+
status: aborted ? ("timeout" as const) : ("ok" as const),
157+
summary: aborted ? "aborted" : "completed",
158+
...(aborted ? { stopReason: result?.meta?.stopReason ?? "rpc" } : {}),
159+
...(aborted && timeoutAttribution.timeoutPhase
160+
? { timeoutPhase: timeoutAttribution.timeoutPhase }
161+
: {}),
162+
...(aborted && timeoutAttribution.providerStarted !== undefined
163+
? { providerStarted: timeoutAttribution.providerStarted }
164+
: {}),
165+
result,
166+
};
167+
const terminalOutcome = buildAgentRunTerminalOutcome({
168+
status:
169+
aborted || result?.meta?.stopReason === "timeout" || timeoutAttribution.timeoutPhase
170+
? "timeout"
171+
: result?.meta?.error || result?.meta?.stopReason === "error"
172+
? "error"
173+
: "ok",
174+
error: result?.meta?.error,
175+
stopReason: result?.meta?.stopReason,
176+
livenessState: result?.meta?.livenessState,
177+
timeoutPhase: timeoutAttribution.timeoutPhase,
178+
providerStarted: timeoutAttribution.providerStarted,
179+
});
180+
const persistTerminalDedupe = () => {
181+
setGatewayDedupeEntries({
182+
dedupe: params.context.dedupe,
183+
keys: params.dedupeKeys,
184+
entry: {
185+
ts: Date.now(),
186+
ok: true,
187+
payload,
188+
},
189+
});
190+
};
191+
const settled = await settle({ terminalOutcome, onRecovered: persistTerminalDedupe });
192+
if (!settled) {
193+
const summary = "failed to persist cron continuation settlement";
194+
const error = errorShape(ErrorCodes.UNAVAILABLE, summary);
195+
const failedPayload = { runId: params.runId, status: "error" as const, summary };
196+
setGatewayDedupeEntries({
197+
dedupe: params.context.dedupe,
198+
keys: params.dedupeKeys,
199+
entry: { ts: Date.now(), ok: false, payload: failedPayload, error },
200+
});
201+
params.respond(false, failedPayload, error, { runId: params.runId, error: summary });
202+
return;
203+
}
204+
persistTerminalDedupe();
205+
// Send a second res frame (same id) so TS clients with expectFinal can wait.
206+
// Swift clients will typically treat the first res as the result and ignore this.
207+
params.respond(true, payload, undefined, { runId: params.runId });
208+
})
209+
.catch(async (err: unknown) => {
210+
const aborted = isGatewayAgentAbortRejection(err, params.abortController.signal);
211+
const renderedErr = formatForLog(err);
212+
if (taskTracked) {
213+
tryFinalizeTrackedAgentTask({
214+
runId: params.runId,
215+
status: aborted ? "timed_out" : resolveFailedTrackedAgentTaskStatus(err),
216+
error: renderedErr,
217+
terminalSummary: renderedErr,
218+
log: params.context.logGateway,
219+
});
220+
}
221+
const error = errorShape(ErrorCodes.UNAVAILABLE, renderedErr);
222+
const stopReason = resolveGatewayAgentAbortStopReason(params.abortController.signal);
223+
const terminalOutcome = buildAgentRunTerminalOutcome({
224+
status: aborted ? "timeout" : "error",
225+
error: renderedErr,
226+
stopReason,
227+
timeoutPhase: aborted ? "gateway_draining" : undefined,
228+
});
229+
const payload = {
230+
runId: params.runId,
231+
status: aborted ? ("timeout" as const) : ("error" as const),
232+
summary: aborted ? "aborted" : renderedErr,
233+
...(aborted ? { stopReason, timeoutPhase: "gateway_draining" as const } : {}),
234+
};
235+
const persistTerminalDedupe = (settlementPersisted: boolean) => {
236+
setGatewayDedupeEntries({
237+
dedupe: params.context.dedupe,
238+
keys: params.dedupeKeys,
239+
entry: {
240+
ts: Date.now(),
241+
ok: aborted && settlementPersisted,
242+
payload,
243+
...(aborted ? {} : { error }),
244+
},
245+
});
246+
};
247+
const settled = await settle({
248+
terminalOutcome,
249+
onRecovered: () => persistTerminalDedupe(true),
250+
});
251+
persistTerminalDedupe(settled);
252+
params.respond(aborted && settled, payload, aborted && settled ? undefined : error, {
253+
runId: params.runId,
254+
...(aborted ? {} : { error: formatForLog(err) }),
255+
});
256+
})
257+
.finally(() => {
258+
clearAgentRunContext(params.runId, params.ingressOpts.lifecycleGeneration);
259+
params.cleanupAbortController();
260+
});
261+
}

0 commit comments

Comments
 (0)