fix(orchestrate): stop nudge waits at the max-wait deadline

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();
+    }
+  });
+});

orchestrate/skills/orchestrate/scripts/tools/nudge-root.ts

@@ -175,26 +175,70 @@ async function sleep(ms: number): Promise<void> {
   await new Promise(r => setTimeout(r, ms));
 }
 
-async function waitUntilIdle(args: Args, deadline: number): Promise<void> {
+export interface NudgeClock {
+  now?: () => number;
+  sleep?: (ms: number) => Promise<void>;
+  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<void> {
+  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<void> {
@@ -213,14 +257,20 @@ async function sendWithBusyRetry(args: Args, deadline: number): Promise<void> {
       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<void> {
   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);