diff --git a/orchestrate/skills/orchestrate/scripts/__tests__/nudge-max-wait.test.ts b/orchestrate/skills/orchestrate/scripts/__tests__/nudge-max-wait.test.ts new file mode 100644 index 000000000..06f39100e --- /dev/null +++ b/orchestrate/skills/orchestrate/scripts/__tests__/nudge-max-wait.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, spyOn, test } from "bun:test"; + +import { + NudgeDeadlineError, + waitMsUntilDeadline, + waitUntilIdle, +} from "../tools/nudge-root.ts"; + +describe("nudge max-wait", () => { + test("Sleeps only until the deadline when the poll interval is longer", () => { + expect( + waitMsUntilDeadline({ now: 0, deadline: 1_000, pollMs: 15_000 }) + ).toBe(1_000); + }); + + test("Uses the poll interval when the deadline is further out", () => { + expect( + waitMsUntilDeadline({ now: 0, deadline: 60_000, pollMs: 15_000 }) + ).toBe(15_000); + }); + + test("Does not sleep once the deadline has passed", () => { + expect( + waitMsUntilDeadline({ now: 5_000, deadline: 5_000, pollMs: 15_000 }) + ).toBe(null); + expect( + waitMsUntilDeadline({ now: 5_001, deadline: 5_000, pollMs: 15_000 }) + ).toBe(null); + }); + + test("Idle wait stops at the deadline instead of sleeping a full poll", async () => { + const log = spyOn(console, "log").mockImplementation(() => {}); + let now = 0; + const sleeps: number[] = []; + try { + await expect( + waitUntilIdle( + { + agentId: "bc-root", + message: "ping", + waitIdle: true, + maxWaitMs: 1_000, + pollMs: 15_000, + }, + 1_000, + { + now: () => now, + sleep: async ms => { + sleeps.push(ms); + now += ms; + }, + listRuns: async () => ({ items: [{ status: "running" }] }), + } + ) + ).rejects.toBeInstanceOf(NudgeDeadlineError); + expect(sleeps).toEqual([1_000]); + expect(now).toBe(1_000); + } finally { + log.mockRestore(); + } + }); +}); diff --git a/orchestrate/skills/orchestrate/scripts/tools/nudge-root.ts b/orchestrate/skills/orchestrate/scripts/tools/nudge-root.ts index bf045b83e..8ce408d07 100644 --- a/orchestrate/skills/orchestrate/scripts/tools/nudge-root.ts +++ b/orchestrate/skills/orchestrate/scripts/tools/nudge-root.ts @@ -175,26 +175,70 @@ async function sleep(ms: number): Promise { await new Promise(r => setTimeout(r, ms)); } -async function waitUntilIdle(args: Args, deadline: number): Promise { +export interface NudgeClock { + now?: () => number; + sleep?: (ms: number) => Promise; + listRuns?: ( + agentId: string + ) => Promise<{ items?: Array<{ status?: string }> }>; +} + +// How long the next poll may sleep. `--max-wait` is a deadline, so a poll +// interval longer than the time left must not run past it. +export function waitMsUntilDeadline(args: { + now: number; + deadline: number; + pollMs: number; +}): number | null { + const remaining = args.deadline - args.now; + if (!(remaining > 0)) return null; + return Math.min(args.pollMs, remaining); +} + +export async function waitUntilIdle( + args: Args, + deadline: number, + clock: NudgeClock = {} +): Promise { + const now = clock.now ?? Date.now; + const sleepFn = clock.sleep ?? sleep; + const listRuns = + clock.listRuns ?? + (agentId => + Agent.listRuns(agentId, { + runtime: "cloud", + limit: 1, + })); let attempt = 0; - while (Date.now() < deadline) { + while (now() < deadline) { attempt++; try { - const runs = await Agent.listRuns(args.agentId, { - runtime: "cloud", - limit: 1, - }); - const status = runs.items[0]?.status; + const runs = await listRuns(args.agentId); + const status = runs.items?.[0]?.status; log(`idle-check ${attempt}: status=${status ?? "(no runs)"}`); if (status !== "running") return; } catch (err) { const msg = err instanceof Error ? err.message : String(err); log(`idle-check ${attempt}: Agent.listRuns failed: ${msg}`); } - await sleep(args.pollMs); + const waitMs = waitMsUntilDeadline({ + now: now(), + deadline, + pollMs: args.pollMs, + }); + if (waitMs === null) break; + await sleepFn(waitMs); + } + const message = `idle-wait exhausted after ${args.maxWaitMs / 1000}s; giving up`; + log(message); + throw new NudgeDeadlineError(message); +} + +export class NudgeDeadlineError extends Error { + constructor(message: string) { + super(message); + this.name = "NudgeDeadlineError"; } - log(`idle-wait exhausted after ${args.maxWaitMs / 1000}s; giving up`); - process.exit(4); } async function sendWithBusyRetry(args: Args, deadline: number): Promise { @@ -213,14 +257,20 @@ async function sendWithBusyRetry(args: Args, deadline: number): Promise { return; } catch (err) { if (isBusy(err)) { + const waitMs = waitMsUntilDeadline({ + now: Date.now(), + deadline, + pollMs: args.pollMs, + }); + if (waitMs === null) break; log( "attempt " + attempt + ": agent_busy; retrying in " + - args.pollMs / 1000 + + waitMs / 1000 + "s" ); - await sleep(args.pollMs); + await sleep(waitMs); continue; } const msg = err instanceof Error ? err.message : String(err); @@ -236,24 +286,31 @@ async function sendWithBusyRetry(args: Args, deadline: number): Promise { process.exit(4); } -const args = await parseArgs(process.argv.slice(2)); -const deadline = Date.now() + args.maxWaitMs; -log( - "nudging " + - args.agentId + - " (mode=" + - (args.waitIdle ? "wait-idle" : "send-now") + - ", max-wait=" + - args.maxWaitMs / 1000 + - "s, poll=" + - args.pollMs / 1000 + - "s, message=" + - args.message.length + - "ch)" -); +if (import.meta.main) { + const args = await parseArgs(process.argv.slice(2)); + const deadline = Date.now() + args.maxWaitMs; + log( + "nudging " + + args.agentId + + " (mode=" + + (args.waitIdle ? "wait-idle" : "send-now") + + ", max-wait=" + + args.maxWaitMs / 1000 + + "s, poll=" + + args.pollMs / 1000 + + "s, message=" + + args.message.length + + "ch)" + ); -if (args.waitIdle) { - await waitUntilIdle(args, deadline); + try { + if (args.waitIdle) { + await waitUntilIdle(args, deadline); + } + await sendWithBusyRetry(args, deadline); + } catch (err) { + if (err instanceof NudgeDeadlineError) process.exit(4); + throw err; + } + process.exit(0); } -await sendWithBusyRetry(args, deadline); -process.exit(0);