fix(langsmith): bound test bodies and retry on fresh worker to stop CI hangs

contrib/langsmith/src/__tests__/helpers.ts

@@ -105,6 +105,22 @@ export class InMemoryRunCollector {
     this.flushCount = 0;
   }
 
+  /** Copy of the current state; lets the harness roll back a retried attempt. */
+  snapshot(): { createOrder: string[]; byId: Map<string, CollectedRun>; flushCount: number } {
+    return { createOrder: [...this.createOrder], byId: new Map(this.byId), flushCount: this.flushCount };
+  }
+
+  /** Restore a {@link snapshot}, discarding anything recorded since. */
+  restore(state: { createOrder: string[]; byId: Map<string, CollectedRun>; flushCount: number }): void {
+    this.createOrder.length = 0;
+    this.createOrder.push(...state.createOrder);
+    this.byId.clear();
+    for (const [id, run] of state.byId) {
+      this.byId.set(id, run);
+    }
+    this.flushCount = state.flushCount;
+  }
+
   /** View this collector as a LangSmith client for the plugin's `client` option. */
   asClient(): LangSmithClient {
     return this as unknown as LangSmithClient;
@@ -198,37 +214,130 @@ function getBundle(plugin: LangSmithPlugin, workflowsPath: string, optionsKey: s
   return bundle;
 }
 
+// Bound on one body attempt, and how many fresh workers to try. On a loaded CI
+// machine the dev server + worker pair can permanently fail to deliver a
+// workflow's *first* workflow task: the task stays SCHEDULED at attempt 1
+// forever (normal-queue first tasks have no schedule-to-start timeout), while a
+// fresh poller on the same queue receives it instantly. Left alone, the case
+// promise never settles, AVA's 120s inactivity watchdog fires, and its SIGTERM
+// is swallowed by the SDK Runtime's shutdown handler — wedging the whole suite
+// until the CI job timeout. Bounding each attempt and retrying on a fresh
+// worker + task queue converts that hang into (at worst) a visible failure and
+// (in practice) a recovered pass.
+//
+// Calibration (from CI run 31544937973, linux-arm Node 24 leg):
+//  - 30s per attempt is latency headroom, not the recovery lever. Even on the
+//    slowest, contended runners a fresh worker reaches RUNNING in <300ms and
+//    healthy bodies finish in single-digit seconds; stalled attempts show zero
+//    activity for the entire window, and an unmitigated stall never recovers
+//    (20+ minute hung jobs). Waiting longer per attempt therefore cannot help;
+//    only a fresh worker can.
+//  - Stalls are correlated in time: that leg saw one case stall on 2/3
+//    attempts then recover, and another stall on 3/3 back-to-back attempts
+//    (~90s window). Six attempts sample a ~3 minute window — double the worst
+//    observed sequence — while still failing loudly within minutes if the
+//    degradation persists.
+//  - Attempts are not capped by AVA's 120s inactivity watchdog: AVA debounces
+//    that timer on every stateChange record carrying a testFile, which
+//    includes worker-stdout/stderr chunks (ava 5.3.1 lib/fork.js tags all
+//    records with testFile; lib/api.js debounces on any of them). The retry
+//    warning below is therefore load-bearing: it guarantees output every
+//    ≤~30s while attempts continue, so the watchdog only fires for a genuine
+//    silent wedge (its backstop role, unchanged).
+const BODY_STALL_TIMEOUT_MS = 30_000;
+const MAX_BODY_ATTEMPTS = 6;
+
+/** A body attempt exceeded {@link BODY_STALL_TIMEOUT_MS}; the harness retries on a fresh worker. */
+class HarnessStallError extends Error {
+  constructor(taskQueue: string, attempt: number) {
+    super(
+      `Test body did not settle within ${BODY_STALL_TIMEOUT_MS}ms on task queue ${taskQueue} ` +
+        `(attempt ${attempt}/${MAX_BODY_ATTEMPTS}); assuming the first-workflow-task delivery stall`
+    );
+  }
+}
+
+/** Terminate workflows a stalled attempt left running so a retry can reuse their workflow ids. */
+async function terminateLeakedWorkflows(env: TestWorkflowEnvironment, taskQueue: string): Promise<void> {
+  try {
+    const leaked = env.client.workflow.list({
+      query: `TaskQueue = '${taskQueue}' AND ExecutionStatus = 'Running'`,
+    });
+    for await (const wf of leaked) {
+      try {
+        await env.client.workflow.getHandle(wf.workflowId, wf.runId).terminate('stalled harness attempt cleanup');
+      } catch {
+        /* already closed */
+      }
+    }
+  } catch {
+    /* best-effort: the retry runs on a fresh task queue regardless */
+  }
+}
+
 export async function withTracingWorker<T>(args: HarnessArgs<T>): Promise<T> {
   const privateEnv = sharedEnv ? undefined : await TestWorkflowEnvironment.createLocal();
   const env = sharedEnv ?? privateEnv!;
   try {
-    const taskQueue = args.taskQueue ?? `langsmith-test-${randomUUID()}`;
     const plugin = new LangSmithPlugin({ ...args.options, client: args.collector.asClient() });
 
     const { workflowsPath = WORKFLOWS_PATH, ...restWorkerOpts } = args.workerOptions ?? {};
     const workflowBundle = await getBundle(plugin, workflowsPath, JSON.stringify(args.options ?? {}));
 
-    const worker = await Worker.create({
-      connection: env.nativeConnection,
-      namespace: env.namespace,
-      taskQueue,
-      workflowBundle,
-      activities: args.activities,
-      plugins: [plugin],
-      // Avoid waiting for the default 10s sticky execution timeout on worker
-      // transition: these short-lived per-case workers can otherwise stall a
-      // full 10s on task redelivery and, on loaded CI, blow the 120s AVA cap.
-      stickyQueueScheduleToStartTimeout: '1s',
-      ...restWorkerOpts,
-    });
+    // Roll the collector back on retry so assertions never see a stalled
+    // attempt's partial emissions.
+    const preAttemptState = args.collector.snapshot();
 
-    const client = new Client({
-      connection: env.connection,
-      namespace: env.namespace,
-      plugins: [plugin],
-    });
+    for (let attempt = 1; ; attempt++) {
+      const taskQueue = args.taskQueue ?? `langsmith-test-${randomUUID()}`;
+
+      const worker = await Worker.create({
+        connection: env.nativeConnection,
+        namespace: env.namespace,
+        taskQueue,
+        workflowBundle,
+        activities: args.activities,
+        plugins: [plugin],
+        // Avoid waiting for the default 10s sticky execution timeout on worker
+        // transition: these short-lived per-case workers can otherwise stall a
+        // full 10s on task redelivery and, on loaded CI, blow the 120s AVA cap.
+        stickyQueueScheduleToStartTimeout: '1s',
+        ...restWorkerOpts,
+      });
 
-    return await worker.runUntil(args.body({ client, taskQueue, env }));
+      const client = new Client({
+        connection: env.connection,
+        namespace: env.namespace,
+        plugins: [plugin],
+      });
+
+      try {
+        // Deferred body (function form): the worker is polling before the body
+        // starts any workflow.
+        return await worker.runUntil(async () => {
+          let stallTimer: ReturnType<typeof setTimeout> | undefined;
+          const stall = new Promise<never>((_, reject) => {
+            stallTimer = setTimeout(() => reject(new HarnessStallError(taskQueue, attempt)), BODY_STALL_TIMEOUT_MS);
+          });
+          try {
+            return await Promise.race([args.body({ client, taskQueue, env }), stall]);
+          } finally {
+            clearTimeout(stallTimer);
+          }
+        });
+      } catch (err) {
+        if (!(err instanceof HarnessStallError) || attempt >= MAX_BODY_ATTEMPTS) {
+          throw err;
+        }
+        // Load-bearing, do not remove: surfaces in the archived test log so CI
+        // stalls stay diagnosable, AND (as worker stdout) debounces AVA's
+        // inactivity watchdog so continued attempts can't trip it — see the
+        // calibration notes on MAX_BODY_ATTEMPTS.
+        console.warn(`withTracingWorker: ${err.message}; retrying on a fresh worker`);
+        await terminateLeakedWorkflows(env, taskQueue);
+        args.collector.restore(preAttemptState);
+      }
+    }
   } finally {
     if (privateEnv) await privateEnv.teardown();
   }