Retry timed-out provider initialization once

packages/agent-runtime/src/runtime-json-rpc.ts

@@ -48,6 +48,16 @@ export class ProviderResponseEncodeError extends Error {
   }
 }
 
+export class JsonRpcRequestTimeoutError extends Error {
+  readonly method: string;
+
+  constructor(method: string) {
+    super(`JSON-RPC request timed out: ${method}`);
+    this.name = "JsonRpcRequestTimeoutError";
+    this.method = method;
+  }
+}
+
 export interface PendingJsonRpcRequest {
   resolve: (result: unknown) => void;
   reject: (error: Error) => void;
@@ -287,7 +297,7 @@ export function sendJsonRpcRequest<TResult>(
   return new Promise((resolve, reject) => {
     const timer = setTimeout(() => {
       args.pending.delete(id);
-      reject(new Error(`JSON-RPC request timed out: ${message.method}`));
+      reject(new JsonRpcRequestTimeoutError(message.method));
     }, args.timeoutMs ?? 30_000);
     args.pending.set(id, {
       resolve: (result) => {

packages/agent-runtime/src/runtime-provider-process.ts

@@ -14,6 +14,7 @@ import { createProviderForId } from "./provider-registry.js";
 import { filterSkillRootsForProvider } from "./runtime-skill-roots.js";
 import {
   ignoredJsonRpcResultSchema,
+  JsonRpcRequestTimeoutError,
   type PendingJsonRpcRequest,
   sendJsonRpcRequest,
 } from "./runtime-json-rpc.js";
@@ -61,6 +62,7 @@ export interface RuntimeProviderProcessManagerArgs {
   env: Record<string, string> | undefined;
   getNextRequestId: () => number;
   handleStdoutLine: (args: RuntimeProviderProcessLineArgs) => void;
+  initializeTimeoutMs?: number;
   onProcessExit: AgentRuntimeOptions["onProcessExit"];
   onProviderIdentityWaitersInterrupted: (
     providerProcess: RuntimeProviderProcess,
@@ -156,63 +158,76 @@ export class RuntimeProviderProcessManager {
     if (this.processes.has(args.processKey)) return;
 
     const startPromise = (async () => {
-      const adapter = this.getAdapter(args.providerId, args.acpLaunchSpec);
-      const providerProcess = this.spawnProvider({
-        adapter,
-        processKey: args.processKey,
-        providerId: args.providerId,
-      });
-
-      try {
-        if (hasChildProcessExited(providerProcess.child)) {
-          const stderr = providerProcess.stderrChunks.join("\n").slice(0, 500);
-          throw new Error(
-            `Provider "${args.providerId}" exited during startup with ${formatChildProcessExitStatus(providerProcess.child)}` +
-              (stderr ? `\nstderr: ${stderr}` : ""),
-          );
-        }
-
-        const initCmd = adapter.buildCommandPlan({ type: "initialize" });
-        if (initCmd.kind === "request") {
-          await sendJsonRpcRequest({
-            child: providerProcess.child,
-            message: initCmd,
-            pending: providerProcess.pending,
-            getNextId: this.args.getNextRequestId,
-            resultSchema: ignoredJsonRpcResultSchema,
-          });
-        }
-
-        const providerSkillRoots = filterSkillRootsForProvider({
+      for (let attempt = 1; attempt <= 2; attempt++) {
+        const adapter = this.getAdapter(args.providerId, args.acpLaunchSpec);
+        const providerProcess = this.spawnProvider({
+          adapter,
+          processKey: args.processKey,
           providerId: args.providerId,
-          skillRoots: this.args.skillRoots,
         });
-        if (providerSkillRoots.length > 0) {
-          const skillRootsCmd = adapter.buildCommandPlan({
-            type: "skills/configure",
-            skillRoots: providerSkillRoots,
-          });
-          if (skillRootsCmd.kind === "request") {
+
+        try {
+          if (hasChildProcessExited(providerProcess.child)) {
+            const stderr = providerProcess.stderrChunks.join("\n").slice(0, 500);
+            throw new Error(
+              `Provider "${args.providerId}" exited during startup with ${formatChildProcessExitStatus(providerProcess.child)}` +
+                (stderr ? `\nstderr: ${stderr}` : ""),
+            );
+          }
+
+          const initCmd = adapter.buildCommandPlan({ type: "initialize" });
+          if (initCmd.kind === "request") {
             await sendJsonRpcRequest({
               child: providerProcess.child,
-              message: skillRootsCmd,
+              message: initCmd,
               pending: providerProcess.pending,
               getNextId: this.args.getNextRequestId,
               resultSchema: ignoredJsonRpcResultSchema,
+              timeoutMs: this.args.initializeTimeoutMs,
             });
           }
-        }
-      } catch (startupError) {
-        await this.cleanupFailedStartup({
-          processKey: args.processKey,
-          providerId: args.providerId,
-          providerProcess,
-          startupError:
+
+          const providerSkillRoots = filterSkillRootsForProvider({
+            providerId: args.providerId,
+            skillRoots: this.args.skillRoots,
+          });
+          if (providerSkillRoots.length > 0) {
+            const skillRootsCmd = adapter.buildCommandPlan({
+              type: "skills/configure",
+              skillRoots: providerSkillRoots,
+            });
+            if (skillRootsCmd.kind === "request") {
+              await sendJsonRpcRequest({
+                child: providerProcess.child,
+                message: skillRootsCmd,
+                pending: providerProcess.pending,
+                getNextId: this.args.getNextRequestId,
+                resultSchema: ignoredJsonRpcResultSchema,
+              });
+            }
+          }
+          return;
+        } catch (startupError) {
+          const error =
             startupError instanceof Error
               ? startupError
-              : new Error(String(startupError)),
-        });
-        throw startupError;
+              : new Error(String(startupError));
+          await this.cleanupFailedStartup({
+            processKey: args.processKey,
+            providerId: args.providerId,
+            providerProcess,
+            startupError: error,
+          });
+          if (
+            attempt === 1 &&
+            !this.shuttingDown &&
+            error instanceof JsonRpcRequestTimeoutError &&
+            error.method === "initialize"
+          ) {
+            continue;
+          }
+          throw error;
+        }
       }
     })();
 

packages/agent-runtime/src/runtime.process-lifecycle.test.ts

@@ -9,7 +9,10 @@ import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { turnScope, type ThreadEvent } from "@bb/domain";
-import type { ProviderAdapter } from "./provider-adapter.js";
+import type {
+  ProviderAdapter,
+  ProviderAdapterFactory,
+} from "./provider-adapter.js";
 import { createAgentRuntimeWithAdapters } from "./runtime.js";
 import { RuntimeProviderProcessManager } from "./runtime-provider-process.js";
 import { RuntimeThreadIdentityRegistry } from "./runtime-thread-identity.js";
@@ -25,8 +28,10 @@ import type { AgentRuntimeOptions } from "./types.js";
 import type { ProviderRuntimeEvent } from "./runtime-json-rpc.js";
 
 interface CreateProviderProcessManagerArgs {
+  adapterFactory?: ProviderAdapterFactory;
   adapterProcessEnv?: Record<string, string>;
   env?: Record<string, string>;
+  initializeTimeoutMs?: number;
   onStderr?: NonNullable<AgentRuntimeOptions["onStderr"]>;
   onProcessExit: NonNullable<AgentRuntimeOptions["onProcessExit"]>;
   scriptPath: string;
@@ -65,8 +70,10 @@ describe("createAgentRuntime process lifecycle", () => {
     let nextRequestId = 1;
     return new RuntimeProviderProcessManager({
       additionalWorkspaceWriteRoots: [],
-      adapterFactory: () =>
-        createNoopInitializeAdapter(args.scriptPath, args.adapterProcessEnv),
+      adapterFactory:
+        args.adapterFactory ??
+        (() =>
+          createNoopInitializeAdapter(args.scriptPath, args.adapterProcessEnv)),
       bridgeBundleDir: undefined,
       captureThreadExitState: (threadId) => ({
         activeTurnId: null,
@@ -78,7 +85,17 @@ describe("createAgentRuntime process lifecycle", () => {
         identityRegistry.createProviderState({ providerId }),
       env: args.env,
       getNextRequestId: () => nextRequestId++,
-      handleStdoutLine: () => undefined,
+      handleStdoutLine: ({ line, providerProcess }) => {
+        const message: unknown = JSON.parse(line);
+        if (!isJsonRecord(message)) return;
+        const id = message.id;
+        if (typeof id !== "string" && typeof id !== "number") return;
+        const pending = providerProcess.pending.get(id);
+        if (!pending || !("result" in message)) return;
+        providerProcess.pending.delete(id);
+        pending.resolve(message.result);
+      },
+      initializeTimeoutMs: args.initializeTimeoutMs,
       onProcessExit: args.onProcessExit,
       onProviderIdentityWaitersInterrupted: (providerProcess) =>
         identityRegistry.resolvePendingIdentityWaiters(
@@ -1264,6 +1281,65 @@ rl.on("line", (line) => {
     await runtime.shutdown();
   });
 
+  it("retries provider startup once when initialize times out", async () => {
+    const attemptsPath = join(tmpDir, "initialize-timeout-attempts.txt");
+    const logPath = join(tmpDir, "initialize-timeout-log.txt");
+    const initializeTimeoutScript = join(tmpDir, "initialize-timeout.cjs");
+    writeFileSync(
+      initializeTimeoutScript,
+      `const fs = require("fs");
+      const readline = require("readline");
+      const attemptsPath = process.argv[2];
+      const logPath = process.argv[3];
+      const previousAttempts = fs.existsSync(attemptsPath)
+        ? Number(fs.readFileSync(attemptsPath, "utf8"))
+        : 0;
+      const attempt = previousAttempts + 1;
+      fs.writeFileSync(attemptsPath, String(attempt));
+      fs.appendFileSync(logPath, "spawn:" + attempt + "\\n");
+      setInterval(() => {}, 1000);
+      const rl = readline.createInterface({ input: process.stdin });
+      rl.on("line", (line) => {
+        const msg = JSON.parse(line);
+        if (msg.method !== "initialize") return;
+        fs.appendFileSync(logPath, "initialize:" + attempt + "\\n");
+        if (attempt === 1) return;
+        process.stdout.write(JSON.stringify({ jsonrpc: "2.0", id: msg.id, result: {} }) + "\\n");
+      });`,
+    );
+    const manager = createProviderProcessManager({
+      adapterFactory: () => {
+        const adapter = createFakeAdapter(scriptPath);
+        return {
+          ...adapter,
+          process: {
+            command: "node",
+            args: [initializeTimeoutScript, attemptsPath, logPath],
+          },
+        };
+      },
+      initializeTimeoutMs: 500,
+      onProcessExit: vi.fn(),
+      scriptPath,
+      workspacePath: tmpDir,
+    });
+
+    try {
+      await manager.ensureProvider({
+        processKey: "fake",
+        providerId: "fake",
+      });
+      expect(readLogLines(logPath)).toEqual([
+        "spawn:1",
+        "initialize:1",
+        "spawn:2",
+        "initialize:2",
+      ]);
+    } finally {
+      await manager.shutdown();
+    }
+  });
+
   it("removes the cached provider and retries when startup skill configuration fails", async () => {
     const attemptsPath = join(tmpDir, "startup-config-attempts.txt");
     const logPath = join(tmpDir, "startup-config-log.txt");