fix(omp): honor configured timeout while waiting for ready (#4143)

* fix(omp): honor configured timeout while waiting for ready

* fix(omp): increase ready timeout to 20 seconds

* docs: keep OMP timeout policy in one place

---------

Co-authored-by: Mohamed Boudra <boudra.moha@gmail.com>

docs/custom-providers.md

@@ -404,7 +404,7 @@ Custom OMP profiles should extend `omp`. They inherit the OMP adapter's `rpc-ui`
 }
 ```
 
-`params.sessionDir` is used only for importing sessions that were started outside Paseo. If `command` or XDG env vars move OMP's state directory, set `params.sessionDir` to the resulting OMP JSONL session directory; launching and resuming still go through the configured command. `params.rpcTimeoutMs` overrides the 60-second OMP control-plane RPC deadline.
+`params.sessionDir` is used only for importing sessions that were started outside Paseo. If `command` or XDG env vars move OMP's state directory, set `params.sessionDir` to the resulting OMP JSONL session directory; launching and resuming still go through the configured command. OMP waits 20 seconds for its initial `ready` frame and 60 seconds for later control-plane RPCs by default. `params.rpcTimeoutMs` overrides both deadlines.
 
 For other providers that keep Pi's `--mode rpc` API but write sessions somewhere else, extend `pi`, replace the command, and provide the JSONL session directory:
 

docs/providers.md

@@ -67,7 +67,7 @@ Pi model records expose input capabilities through `model.input`. Only send raw
 
 Pi MCP support depends on the open-source `pi-mcp-adapter` extension being loaded for the agent cwd. Probe with Pi RPC `get_commands`; the adapter registers an extension command named `mcp` (often with `sourceInfo.source` containing `pi-mcp-adapter`). When Paseo injects MCP servers into Pi, write a per-agent MCP config and pass it with `--mcp-config` instead of modifying user or project MCP files. Because that flag replaces the Pi global config layer, preserve the existing `<Pi agent dir>/mcp.json` in the generated file before overlaying injected servers. For local HTTP servers such as Paseo's own `/mcp/agents` endpoint, explicitly disable adapter OAuth (`auth: false`, `oauth: false`) in the generated config.
 
-Pi and OMP control-plane RPCs wait 60 seconds by default. Override the provider's `params.rpcTimeoutMs` when extension or MCP startup on a slow host needs a different deadline. Timeout errors name the pending RPC phase and report both elapsed time and the configured deadline. This setting does not govern long-running Pi compaction or Pi extension UI results.
+Pi control-plane RPCs wait 60 seconds by default. Override `params.rpcTimeoutMs` when extension or MCP startup on a slow host needs more time. Timeout errors name the pending RPC phase and report both elapsed time and the configured deadline. This setting does not govern long-running Pi compaction or Pi extension UI results. See [OMP profiles and Pi-compatible forks](custom-providers.md#omp-profiles-and-pi-compatible-forks) for OMP startup and RPC deadlines.
 
 Pi import discovery reads Pi's persisted JSONL session files because Pi RPC does not expose a recent-session listing command. Resume and full history hydration still go through `pi --mode rpc` using the session file as `nativeHandle`.
 

packages/server/src/server/agent/providers/omp/agent.test.ts

@@ -7,11 +7,15 @@ import type { OmpUsagePollScheduler } from "./usage-poller.js";
 import { resolveOmpProviderParams } from "./provider-config.js";
 import { OmpHarness } from "./test-utils/omp-harness.js";
 
-test("OMP RPC timeout defaults to 60 seconds and accepts an override", () => {
-  expect(resolveOmpProviderParams({}).runtimeProviderParams.rpcTimeoutMs).toBe(60_000);
-  expect(
-    resolveOmpProviderParams({ rpcTimeoutMs: 90_000 }).runtimeProviderParams.rpcTimeoutMs,
-  ).toBe(90_000);
+test("OMP ready timeout defaults to 20 seconds and RPC timeout overrides both", () => {
+  expect(resolveOmpProviderParams({}).runtimeProviderParams).toMatchObject({
+    readyTimeoutMs: 20_000,
+    rpcTimeoutMs: 60_000,
+  });
+  expect(resolveOmpProviderParams({ rpcTimeoutMs: 90_000 }).runtimeProviderParams).toMatchObject({
+    readyTimeoutMs: 90_000,
+    rpcTimeoutMs: 90_000,
+  });
 });
 
 class ManualIdleScheduler implements OmpProviderIdleScheduler {

packages/server/src/server/agent/providers/omp/agent.ts

@@ -828,14 +828,15 @@ function buildExtensionUiResponse(
 function createRuntime(
   logger: Logger,
   runtimeSettings: ProviderRuntimeSettings | undefined,
-  requestTimeoutMs: number,
+  providerParams: OmpRuntimeProviderParams,
 ): OmpRuntime {
   return new OmpCliRuntime({
     logger,
     runtimeSettings,
     command: ["omp"],
     commandsRpcName: "get_available_commands",
-    requestTimeoutMs,
+    readyTimeoutMs: providerParams.readyTimeoutMs,
+    requestTimeoutMs: providerParams.rpcTimeoutMs,
   });
 }
 
@@ -2216,8 +2217,7 @@ export class OmpAgentClient implements AgentClient {
     this.noTurnScheduler = options.noTurnScheduler;
     this.usagePollScheduler = options.usagePollScheduler;
     this.runtime =
-      options.runtime ??
-      createRuntime(options.logger, runtimeSettings, this.providerParams.rpcTimeoutMs);
+      options.runtime ?? createRuntime(options.logger, runtimeSettings, this.providerParams);
   }
 
   private async configureNativePaseoTools(

packages/server/src/server/agent/providers/omp/cli-runtime.ts

@@ -47,6 +47,7 @@ export interface OmpCliRuntimeOptions {
   runtimeSettings?: ProviderRuntimeSettings;
   command?: [string, ...string[]];
   commandsRpcName?: "get_available_commands";
+  readyTimeoutMs?: number;
   requestTimeoutMs?: number;
   spawnProcess?: (launch: OmpRuntimeLaunch) => ChildProcessWithoutNullStreams;
 }
@@ -87,7 +88,10 @@ export class OmpCliRuntime implements OmpRuntime {
     const handleAbort = () => void process.close(input.signal?.reason).catch(() => undefined);
     input.signal?.addEventListener("abort", handleAbort, { once: true });
     try {
-      await establishOmpProtocol(process, this.options.logger, this.options.requestTimeoutMs);
+      await establishOmpProtocol(process, this.options.logger, {
+        readyTimeoutMs: this.options.readyTimeoutMs,
+        requestTimeoutMs: this.options.requestTimeoutMs,
+      });
       input.signal?.throwIfAborted();
       return new OmpCliRuntimeSession(process, this.commandsRpcName);
     } catch (error) {

packages/server/src/server/agent/providers/omp/protocol-session.test.ts

@@ -1,5 +1,5 @@
 import pino from "pino";
-import { describe, expect, it } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import type { JsonlRpcExit } from "../jsonl-rpc-process.js";
 import { establishOmpProtocol, type OmpProtocolTransport } from "./protocol-session.js";
 
@@ -44,6 +44,10 @@ const V2_READY = {
   maxReassembledFrameBytes: 64 * 1024 * 1024,
 };
 
+afterEach(() => {
+  vi.useRealTimers();
+});
+
 describe("establishOmpProtocol", () => {
   it("owns readiness and v2 negotiation", async () => {
     const harness = transportHarness();
@@ -59,6 +63,43 @@ describe("establishOmpProtocol", () => {
     expect(harness.exitSubscribed()).toBe(false);
   });
 
+  it("uses the configured RPC timeout while waiting for a cold OMP", async () => {
+    vi.useFakeTimers();
+    const harness = transportHarness();
+    const negotiation = establishOmpProtocol(harness.transport, pino({ level: "silent" }), {
+      readyTimeoutMs: 60_000,
+      requestTimeoutMs: 60_000,
+    });
+    const outcome = negotiation.then(
+      () => "resolved",
+      () => "rejected",
+    );
+
+    await vi.advanceTimersByTimeAsync(12_000);
+    harness.emitReady(V2_READY);
+
+    await expect(outcome).resolves.toBe("resolved");
+  });
+
+  it("waits 20 seconds for OMP to become ready", async () => {
+    vi.useFakeTimers();
+    const harness = transportHarness();
+    const negotiation = establishOmpProtocol(harness.transport, pino({ level: "silent" }));
+    const outcome = negotiation.then(
+      () => "resolved",
+      (error: unknown) =>
+        error instanceof Error ? error.message : "rejected with a non-Error value",
+    );
+
+    await vi.advanceTimersByTimeAsync(19_999);
+    expect(harness.exitSubscribed()).toBe(true);
+    expect(await Promise.race([outcome, Promise.resolve("pending")])).toBe("pending");
+
+    await vi.advanceTimersByTimeAsync(1);
+    await expect(outcome).resolves.toBe("Timed out waiting for OMP to become ready");
+    expect(harness.exitSubscribed()).toBe(false);
+  });
+
   it("keeps protocol v1 when the ready frame has no matching capability", async () => {
     const harness = transportHarness();
     const negotiation = establishOmpProtocol(harness.transport, pino({ level: "silent" }));

packages/server/src/server/agent/providers/omp/protocol-session.ts

@@ -5,7 +5,7 @@ import {
   type JsonlRpcExit,
 } from "../jsonl-rpc-process.js";
 
-const OMP_READY_TIMEOUT_MS = 10_000;
+const OMP_READY_TIMEOUT_MS = 20_000;
 
 export interface OmpProtocolTransport {
   onMessage(callback: (message: Record<string, unknown>) => void): () => void;
@@ -13,12 +13,19 @@ export interface OmpProtocolTransport {
   request(command: Record<string, unknown>, timeoutMs: number | null): Promise<unknown>;
 }
 
+export interface OmpProtocolTimeouts {
+  readyTimeoutMs?: number;
+  requestTimeoutMs?: number;
+}
+
 export async function establishOmpProtocol(
   transport: OmpProtocolTransport,
   logger: Logger,
-  requestTimeoutMs = JSONL_RPC_DEFAULT_TIMEOUT_MS,
+  timeouts: OmpProtocolTimeouts = {},
 ): Promise<void> {
-  const ready = await waitForReady(transport);
+  const readyTimeoutMs = timeouts.readyTimeoutMs ?? OMP_READY_TIMEOUT_MS;
+  const requestTimeoutMs = timeouts.requestTimeoutMs ?? JSONL_RPC_DEFAULT_TIMEOUT_MS;
+  const ready = await waitForReady(transport, readyTimeoutMs);
   if (!supportsJsonlRpcProtocolV2(ready)) return;
   const response = (await transport.request(
     { type: "negotiate_protocol", protocolVersion: 2 },
@@ -28,14 +35,17 @@ export async function establishOmpProtocol(
   logger.debug({}, "Negotiated OMP RPC protocol v2 (chunked frame transport)");
 }
 
-function waitForReady(transport: OmpProtocolTransport): Promise<Record<string, unknown>> {
+function waitForReady(
+  transport: OmpProtocolTransport,
+  requestTimeoutMs: number,
+): Promise<Record<string, unknown>> {
   return new Promise((resolve, reject) => {
     let settled = false;
     let unsubscribeMessage = (): void => {};
     let unsubscribeExit = (): void => {};
     const timer = setTimeout(
       () => finish(new Error("Timed out waiting for OMP to become ready")),
-      OMP_READY_TIMEOUT_MS,
+      requestTimeoutMs,
     );
     const finish = (result: Record<string, unknown> | Error): void => {
       if (settled) return;

packages/server/src/server/agent/providers/omp/provider-config.ts

@@ -8,6 +8,7 @@ import type { ProviderRuntimeSettings } from "../../provider-launch-config.js";
 
 const OMP_SESSION_DIR = "~/.omp/agent/sessions";
 const DEFAULT_OMP_MODE_ID = "full";
+const DEFAULT_OMP_READY_TIMEOUT_MS = 20_000;
 const DEFAULT_OMP_RPC_TIMEOUT_MS = 60_000;
 
 export const MIN_SUPPORTED_OMP_VERSION = "16.3.9";
@@ -16,7 +17,7 @@ export { OMP_MODES };
 export const OmpProviderParamsSchema = z
   .object({
     sessionDir: z.string().min(1).optional(),
-    rpcTimeoutMs: z.number().int().positive().default(DEFAULT_OMP_RPC_TIMEOUT_MS),
+    rpcTimeoutMs: z.number().int().positive().optional(),
     smolModel: z.string().min(1).optional(),
     slowModel: z.string().min(1).optional(),
     planModel: z.string().min(1).optional(),
@@ -25,6 +26,7 @@ export const OmpProviderParamsSchema = z
 
 export interface OmpRuntimeProviderParams {
   sessionDir: string;
+  readyTimeoutMs: number;
   rpcTimeoutMs: number;
 }
 
@@ -132,10 +134,12 @@ export function resolveOmpProviderParams(providerParams: unknown): {
   modelRoleParams: OmpModelRoleParams;
 } {
   const params = OmpProviderParamsSchema.parse(providerParams ?? {});
+  const configuredRpcTimeoutMs = params.rpcTimeoutMs;
   return {
     runtimeProviderParams: {
       sessionDir: params.sessionDir ?? OMP_SESSION_DIR,
-      rpcTimeoutMs: params.rpcTimeoutMs,
+      readyTimeoutMs: configuredRpcTimeoutMs ?? DEFAULT_OMP_READY_TIMEOUT_MS,
+      rpcTimeoutMs: configuredRpcTimeoutMs ?? DEFAULT_OMP_RPC_TIMEOUT_MS,
     },
     modelRoleParams: {
       ...(params.smolModel ? { smolModel: params.smolModel } : {}),