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);