From b04696aa7b1edbcbdc51de57360f80087f08b326 Mon Sep 17 00:00:00 2001 From: Austin Kurpuis Date: Tue, 18 Aug 2026 06:43:51 -0700 Subject: [PATCH] Time out awaitJob when a launched tool never reports back CallbackReceiver/NatsJobReceiver.awaitJob() had no timeout: if a Job/ToolRun failed silently or never posted a callback, the Promise hung forever, and the chat UI appeared to hang indefinitely (the SSE heartbeat and TemporalEngine's 30-minute poll only kept the transport alive, they never surfaced an error). awaitJob now rejects with JobTimeoutError after 10 minutes of silence (configurable via opts.timeoutMs), and graph.ts labels that case distinctly from a launch failure. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01HmQ2VLfAt9s1SCXrY5mLPr --- apps/agent-orchestrator/src/agent/graph.ts | 7 ++-- .../src/callback/nats-job-receiver.ts | 14 ++++++-- .../src/callback/receiver.test.ts | 11 +++++- .../src/callback/receiver.ts | 35 ++++++++++++++++--- 4 files changed, 57 insertions(+), 10 deletions(-) diff --git a/apps/agent-orchestrator/src/agent/graph.ts b/apps/agent-orchestrator/src/agent/graph.ts index c1a2a50..46ad972 100644 --- a/apps/agent-orchestrator/src/agent/graph.ts +++ b/apps/agent-orchestrator/src/agent/graph.ts @@ -7,7 +7,7 @@ import type { AgentOrchestratorChannel, AgentTurnResult } from "../agents/nats-a import { AgentTurnFailedError, AgentTurnTimeoutError, AgentTurnTransportError } from "../agents/nats-agent-channel.js"; import type { AgentDescriptor, AgentSearchResult, AgentStore } from "../agents/types.js"; import type { PendingToolCall } from "../caller-tools/types.js"; -import type { JobResultReceiver } from "../callback/receiver.js"; +import { JobTimeoutError, type JobResultReceiver } from "../callback/receiver.js"; import type { ContainerToolLauncher } from "../k8s/container-tool-launcher.js"; import type { AgentRunLauncherPort } from "../k8s/agentrun-launcher.js"; import type { SecretKeySelector } from "../k8s/toolrun-launcher.js"; @@ -1995,7 +1995,10 @@ export function buildAgentGraph(deps: AgentGraphDeps) { // keeps that context. return { jobId, - error: `tool ${tool.id} failed to launch: ${err instanceof Error ? err.message : String(err)}`, + error: + err instanceof JobTimeoutError + ? `tool ${tool.id} timed out: ${err.message}` + : `tool ${tool.id} failed to launch: ${err instanceof Error ? err.message : String(err)}`, }; } finally { unsubscribeProgress(); diff --git a/apps/agent-orchestrator/src/callback/nats-job-receiver.ts b/apps/agent-orchestrator/src/callback/nats-job-receiver.ts index 0d0961e..159c956 100644 --- a/apps/agent-orchestrator/src/callback/nats-job-receiver.ts +++ b/apps/agent-orchestrator/src/callback/nats-job-receiver.ts @@ -1,10 +1,11 @@ import { connect, JSONCodec, type NatsConnection, type Subscription } from "nats"; import { EventSchema, NATS_RECONNECT_OPTIONS, type Event } from "@controller-agent/messaging"; -import type { JobResultReceiver, ProgressHandler } from "./receiver.js"; +import { DEFAULT_JOB_TIMEOUT_MS, JobTimeoutError, type JobResultReceiver, type ProgressHandler } from "./receiver.js"; type PendingJob = { resolve: (event: Event) => void; reject: (err: Error) => void; + timer: ReturnType; }; /** @@ -47,9 +48,15 @@ export class NatsJobReceiver implements JobResultReceiver { return receiver; } - awaitJob(jobId: string): Promise { + awaitJob(jobId: string, opts: { timeoutMs?: number } = {}): Promise { + const timeoutMs = opts.timeoutMs ?? DEFAULT_JOB_TIMEOUT_MS; return new Promise((resolve, reject) => { - this.pending.set(jobId, { resolve, reject }); + const timer = setTimeout(() => { + this.pending.delete(jobId); + this.progressHandlers.delete(jobId); + reject(new JobTimeoutError(`tool run ${jobId} did not report back within ${timeoutMs}ms`)); + }, timeoutMs); + this.pending.set(jobId, { resolve, reject, timer }); }); } @@ -85,6 +92,7 @@ export class NatsJobReceiver implements JobResultReceiver { if (event.type === "succeeded" || event.type === "failed") { const pending = this.pending.get(jobId); if (pending) { + clearTimeout(pending.timer); this.pending.delete(jobId); this.progressHandlers.delete(jobId); pending.resolve(event); diff --git a/apps/agent-orchestrator/src/callback/receiver.test.ts b/apps/agent-orchestrator/src/callback/receiver.test.ts index f0d5a0a..0308aa9 100644 --- a/apps/agent-orchestrator/src/callback/receiver.test.ts +++ b/apps/agent-orchestrator/src/callback/receiver.test.ts @@ -1,6 +1,6 @@ import { createHmac } from "node:crypto"; import { describe, expect, it } from "vitest"; -import { CallbackAuthError, CallbackReceiver, verifyAndParseCallback } from "./receiver.js"; +import { CallbackAuthError, CallbackReceiver, JobTimeoutError, verifyAndParseCallback } from "./receiver.js"; /** * Every request in this file goes over a fresh TCP connection — see the long @@ -128,4 +128,13 @@ describe("CallbackReceiver", () => { await receiver.close(); }); + + it("rejects awaitJob with JobTimeoutError if no callback arrives within timeoutMs", async () => { + const receiver = new CallbackReceiver(SECRET); + await receiver.listen(0); + + await expect(receiver.awaitJob("job-never-reports", { timeoutMs: 10 })).rejects.toThrow(JobTimeoutError); + + await receiver.close(); + }); }); diff --git a/apps/agent-orchestrator/src/callback/receiver.ts b/apps/agent-orchestrator/src/callback/receiver.ts index 455fb2f..43beb1d 100644 --- a/apps/agent-orchestrator/src/callback/receiver.ts +++ b/apps/agent-orchestrator/src/callback/receiver.ts @@ -4,6 +4,21 @@ import { EventSchema, type Event } from "@controller-agent/messaging"; export class CallbackAuthError extends Error {} +/** Thrown by `awaitJob` when no terminal event arrives within its timeout — see {@link DEFAULT_JOB_TIMEOUT_MS}. */ +export class JobTimeoutError extends Error {} + +/** + * How long `awaitJob` waits for a terminal (`succeeded`/`failed`) callback + * before giving up. Without this, a Job/ToolRun that never reports back + * (stuck in `ImagePullBackOff`, crash-looping, never scheduled, etc.) left + * `awaitJob`'s Promise pending forever — the only backstops were the + * TemporalEngine's 30-minute whole-turn poll deadline (which just returns a + * "still running" placeholder, not an error) and the SSE heartbeat (a + * transport keep-alive that never terminates a stuck turn). That hang is + * what surfaced as the chat UI appearing to hang after a tool failed to launch. + */ +export const DEFAULT_JOB_TIMEOUT_MS = 10 * 60 * 1000; + /** Handler called for each `progress` or `warning` event received for a specific job. */ export type ProgressHandler = (stage: string, message: string | undefined) => void; @@ -13,8 +28,12 @@ export type ProgressHandler = (stage: string, message: string | undefined) => vo * transport can be swapped without touching graph logic. */ export interface JobResultReceiver { - /** Returns a Promise that resolves once a terminal (`succeeded`/`failed`) event arrives. */ - awaitJob(jobId: string): Promise; + /** + * Returns a Promise that resolves once a terminal (`succeeded`/`failed`) + * event arrives, or rejects with {@link JobTimeoutError} after + * `opts.timeoutMs` (default {@link DEFAULT_JOB_TIMEOUT_MS}) of silence. + */ + awaitJob(jobId: string, opts?: { timeoutMs?: number }): Promise; /** * Registers a handler to receive `progress`/`warning` events for `jobId`. * Returns an unsubscribe function — always call it when the job is done. @@ -46,6 +65,7 @@ export function verifyAndParseCallback(rawBody: string, signatureHeader: string type PendingJob = { resolve: (event: Event) => void; reject: (err: Error) => void; + timer: ReturnType; }; /** @@ -64,9 +84,15 @@ export class CallbackReceiver { constructor(private readonly secret: string) {} - awaitJob(jobId: string): Promise { + awaitJob(jobId: string, opts: { timeoutMs?: number } = {}): Promise { + const timeoutMs = opts.timeoutMs ?? DEFAULT_JOB_TIMEOUT_MS; return new Promise((resolve, reject) => { - this.pending.set(jobId, { resolve, reject }); + const timer = setTimeout(() => { + this.pending.delete(jobId); + this.progressHandlers.delete(jobId); + reject(new JobTimeoutError(`tool run ${jobId} did not report back within ${timeoutMs}ms`)); + }, timeoutMs); + this.pending.set(jobId, { resolve, reject, timer }); }); } @@ -130,6 +156,7 @@ export class CallbackReceiver { if (jobId && (event.type === "succeeded" || event.type === "failed")) { const pending = this.pending.get(jobId); if (pending) { + clearTimeout(pending.timer); this.pending.delete(jobId); this.progressHandlers.delete(jobId); pending.resolve(event);